[Model] Add native MAGI-2 Preview diffusion support (#5918)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

vllm_omni/diffusion/models/magi2/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""SandAI MAGI-2 Preview diffusion support."""
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = ["Magi2Pipeline", "get_magi2_post_process_func"]
+
+
+def __getattr__(name: str) -> Any:
+    if name in __all__:
+        from .pipeline_magi2 import Magi2Pipeline, get_magi2_post_process_func
+
+        return {
+            "Magi2Pipeline": Magi2Pipeline,
+            "get_magi2_post_process_func": get_magi2_post_process_func,
+        }[name]
+    raise AttributeError(name)

vllm_omni/diffusion/models/magi2/attention.py

@@ -0,0 +1,293 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""Native MAGI-2 packed attention with sinks and Ulysses exchange.
+
+The sink and context-parallel math is adapted from SandAI's Apache-2.0
+MAGI-2 preview implementation.  This version uses vLLM's bundled
+FlashAttention extension and vLLM-Omni's existing Ulysses process group; the
+PyTorch path is an exact, portable oracle for small tests.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from functools import cache
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.utils.fa import (
+    resolve_vllm_flash_attn_version,
+    vllm_flash_attn_varlen_with_lse,
+)
+
+from .parallel import (
+    Magi2ParallelGroup,
+    get_magi2_ulysses_group,
+    scatter_heads_gather_seqlen,
+    scatter_seqlen_gather_heads,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@cache
+def _resolve_flash_attn_version() -> int:
+    """Apply MAGI-2's operator override to the shared FA2/FA3/FA4 resolver."""
+
+    requested = os.environ.get("MAGI2_FLASH_ATTN_VERSION")
+    version = resolve_vllm_flash_attn_version(requested)
+    logger.info("MAGI-2 selected FlashAttention %d", version)
+    return version
+
+
+@dataclass(frozen=True)
+class VarlenHandler:
+    """Packed-sequence metadata consumed by MAGI-2 attention."""
+
+    cu_seqlens_q: torch.Tensor | None
+    cu_seqlens_k: torch.Tensor | None
+    max_seqlen_q: int | None = None
+    max_seqlen_k: int | None = None
+
+    def resolved(self, q_tokens: int, k_tokens: int) -> tuple[torch.Tensor, torch.Tensor, int, int]:
+        if self.cu_seqlens_q is None:
+            cu_q = torch.tensor([0, q_tokens], device="cpu", dtype=torch.int32)
+        else:
+            cu_q = self.cu_seqlens_q
+        if self.cu_seqlens_k is None:
+            cu_k = torch.tensor([0, k_tokens], device="cpu", dtype=torch.int32)
+        else:
+            cu_k = self.cu_seqlens_k
+        max_q = self.max_seqlen_q
+        max_k = self.max_seqlen_k
+        if max_q is None:
+            max_q = int(torch.diff(cu_q).max().item()) if cu_q.numel() > 1 else 0
+        if max_k is None:
+            max_k = int(torch.diff(cu_k).max().item()) if cu_k.numel() > 1 else 0
+        return cu_q, cu_k, max_q, max_k
+
+
+def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
+    if not interleaved:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+    x1, x2 = x[..., ::2], x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    interleaved: bool = False,
+) -> torch.Tensor:
+    """Apply the released element-wise RoPE layout to ``[...,H,D]``."""
+
+    rotary_dim = cos.shape[-1] * 2
+    if rotary_dim > x.shape[-1]:
+        raise ValueError(f"RoPE dimension {rotary_dim} exceeds head dimension {x.shape[-1]}")
+    if interleaved:
+        cos = cos.unsqueeze(-2).repeat_interleave(2, dim=-1)
+        sin = sin.unsqueeze(-2).repeat_interleave(2, dim=-1)
+    else:
+        cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2)
+        sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2)
+    rotated = x[..., :rotary_dim] * cos + rotate_half(x[..., :rotary_dim], interleaved) * sin
+    return torch.cat((rotated, x[..., rotary_dim:]), dim=-1)
+
+
+def correct_out_lse_with_sink(
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    sink: torch.Tensor | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Add zero-valued attention sinks to an already-computed softmax.
+
+    FlashAttention returns ``out[T,H,D]`` and conventionally ``lse[H,T]``.
+    A MAGI sink contains additional logits ``[num_sink,H]`` whose values are
+    zero vectors, so only the denominator and LSE change.
+    """
+
+    if sink is None or sink.numel() == 0:
+        return out, lse
+    if out.ndim != 3 or lse.ndim != 2 or sink.ndim != 2:
+        raise ValueError(
+            "expected out[T,H,D], lse[H,T], sink[num_sink,H], got "
+            f"{tuple(out.shape)}, {tuple(lse.shape)}, {tuple(sink.shape)}"
+        )
+    old_lse = lse.float().transpose(0, 1)
+    sink_lse = torch.logsumexp(sink.float(), dim=0).unsqueeze(0)
+    if old_lse.shape[-1] != sink_lse.shape[-1]:
+        raise ValueError("attention sink and FlashAttention head counts differ")
+    new_lse = torch.logaddexp(old_lse, sink_lse)
+    delta = old_lse - new_lse
+    delta = torch.where(torch.isfinite(delta), delta, torch.full_like(delta, -torch.inf))
+    corrected = out * torch.exp(delta).unsqueeze(-1).to(out.dtype)
+    return corrected, new_lse.transpose(0, 1).contiguous()
+
+
+def _repeat_kv_heads(tensor: torch.Tensor, query_heads: int) -> torch.Tensor:
+    kv_heads = tensor.shape[1]
+    if kv_heads == query_heads:
+        return tensor
+    if query_heads % kv_heads:
+        raise ValueError(f"query heads {query_heads} must be divisible by KV heads {kv_heads}")
+    return tensor.repeat_interleave(query_heads // kv_heads, dim=1)
+
+
+def torch_varlen_attention_with_sink(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    softcap: float = -1.0,
+    sink: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Reference packed attention, including GQA and sink logits."""
+
+    if q.ndim != 3 or k.ndim != 3 or v.ndim != 3:
+        raise ValueError("packed attention expects q/k/v shaped [tokens,heads,dim]")
+    if cu_seqlens_q.numel() != cu_seqlens_k.numel():
+        raise ValueError("query and key cumulative-length arrays must contain the same batch count")
+    output = torch.empty_like(q)
+    scale = q.shape[-1] ** -0.5
+    for batch_idx in range(cu_seqlens_q.numel() - 1):
+        q_start, q_end = (int(v) for v in cu_seqlens_q[batch_idx : batch_idx + 2].tolist())
+        k_start, k_end = (int(v) for v in cu_seqlens_k[batch_idx : batch_idx + 2].tolist())
+        q_part = q[q_start:q_end].float()
+        k_part = _repeat_kv_heads(k[k_start:k_end], q.shape[1]).float()
+        v_part = _repeat_kv_heads(v[k_start:k_end], q.shape[1]).float()
+        scores = torch.einsum("qhd,khd->hqk", q_part, k_part) * scale
+        if softcap > 0:
+            scores = softcap * torch.tanh(scores / softcap)
+        if sink is not None and sink.numel() > 0:
+            sink_scores = sink.float().transpose(0, 1).unsqueeze(1).expand(-1, q_part.shape[0], -1)
+            probabilities = torch.softmax(torch.cat((scores, sink_scores), dim=-1), dim=-1)[..., : k_part.shape[0]]
+        else:
+            probabilities = torch.softmax(scores, dim=-1)
+        output[q_start:q_end] = torch.einsum("hqk,khd->qhd", probabilities, v_part).to(output.dtype)
+    return output
+
+
+def packed_attention_with_sink(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    varlen: VarlenHandler,
+    *,
+    softcap: float = -1.0,
+    sink: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run packed attention on one rank after Ulysses head exchange."""
+
+    cu_q, cu_k, max_q, max_k = varlen.resolved(q.shape[0], k.shape[0])
+    cu_q = cu_q.to(device=q.device, dtype=torch.int32).contiguous()
+    cu_k = cu_k.to(device=q.device, dtype=torch.int32).contiguous()
+    if q.is_cuda:
+        out, lse = vllm_flash_attn_varlen_with_lse(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            max_seqlen_q=max_q,
+            max_seqlen_k=max_k,
+            softcap=softcap,
+            deterministic=os.environ.get("MAGI2_DETERMINISTIC", "0") == "1",
+            fa_version=_resolve_flash_attn_version(),
+        )
+        return correct_out_lse_with_sink(out, lse, sink)[0]
+    return torch_varlen_attention_with_sink(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_q,
+        cu_seqlens_k=cu_k,
+        softcap=softcap,
+        sink=sink,
+    )
+
+
+def ulysses_packed_attention_with_sink(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    varlen: VarlenHandler,
+    split_sizes: list[int] | torch.Tensor,
+    *,
+    softcap: float = -1.0,
+    sink: torch.Tensor | None = None,
+    group: Magi2ParallelGroup | None = None,
+) -> torch.Tensor:
+    """MAGI-2 attention with overlapping Ulysses CP/head exchange."""
+
+    group = group or get_magi2_ulysses_group()
+    if isinstance(split_sizes, torch.Tensor):
+        split_sizes = [int(v) for v in split_sizes.detach().cpu().tolist()]
+    if group.world_size > 1:
+        q, k, v = scatter_heads_gather_seqlen((q, k, v), split_sizes, group)
+        if sink is not None:
+            if sink.shape[-1] % group.world_size:
+                raise ValueError("attention sink heads must divide across Ulysses ranks")
+            sink = torch.chunk(sink, group.world_size, dim=-1)[group.rank].contiguous()
+    output = packed_attention_with_sink(q, k, v, varlen, softcap=softcap, sink=sink)
+    if group.world_size > 1:
+        output = scatter_seqlen_gather_heads(output.contiguous(), split_sizes, group)
+        assert isinstance(output, torch.Tensor)
+    return output
+
+
+class Magi2PackedAttentionKernel(nn.Module):
+    """Model kernel plugged into the shared diffusion Attention contract."""
+
+    def __init__(self, softcap: float) -> None:
+        super().__init__()
+        self.softcap = softcap
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        if attn_metadata is None:
+            raise ValueError("MAGI-2 packed attention requires attention metadata")
+        varlen = attn_metadata.extra.get("magi2_varlen")
+        split_sizes = attn_metadata.extra.get("magi2_split_sizes")
+        sink = attn_metadata.extra.get("magi2_sink")
+        if not isinstance(varlen, VarlenHandler):
+            raise TypeError("magi2_varlen must be a VarlenHandler")
+        if not isinstance(split_sizes, (list, torch.Tensor)):
+            raise TypeError("magi2_split_sizes must be a list or tensor")
+        if sink is not None and not isinstance(sink, torch.Tensor):
+            raise TypeError("magi2_sink must be a tensor or None")
+        return ulysses_packed_attention_with_sink(
+            query,
+            key,
+            value,
+            varlen,
+            split_sizes,
+            softcap=self.softcap,
+            sink=sink,
+        )
+
+
+__all__ = [
+    "VarlenHandler",
+    "Magi2PackedAttentionKernel",
+    "apply_rotary_emb",
+    "correct_out_lse_with_sink",
+    "packed_attention_with_sink",
+    "torch_varlen_attention_with_sink",
+    "ulysses_packed_attention_with_sink",
+]

vllm_omni/diffusion/models/magi2/audio_decoder.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Native Stable Audio Open decoder used by MAGI-2.
+
+The MAGI checkpoint embeds Stable Audio's Oobleck VAE under the
+``pretransform.model`` prefix and uses the sequential module names from
+``stable-audio-tools``. This module converts those names to Diffusers'
+``AutoencoderOobleck`` layout and loads only the decoder tensors.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+from diffusers import AutoencoderOobleck
+from safetensors import safe_open
+
+_CHECKPOINT_PREFIX = "pretransform.model.decoder."
+_BLOCK_COMPONENTS = {
+    0: "snake1",
+    1: "conv_t1",
+}
+_RESIDUAL_COMPONENTS = {
+    0: "snake1",
+    1: "conv1",
+    2: "snake2",
+    3: "conv2",
+}
+
+
+def convert_stable_audio_decoder_key(key: str) -> str | None:
+    """Map one Stable Audio decoder key to Diffusers' Oobleck decoder.
+
+    Unrelated checkpoint tensors return ``None``. A key inside the decoder
+    namespace that does not match the released Stable Audio layout raises so a
+    changed checkpoint cannot be loaded partially by accident.
+    """
+
+    if not key.startswith(_CHECKPOINT_PREFIX):
+        return None
+    suffix = key[len(_CHECKPOINT_PREFIX) :]
+    parts = suffix.split(".")
+    if len(parts) < 3 or parts[0] != "layers":
+        raise ValueError(f"Unsupported Stable Audio decoder key: {key}")
+
+    try:
+        layer = int(parts[1])
+    except ValueError as exc:
+        raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
+    tail = parts[2:]
+
+    if layer == 0 and len(tail) == 1:
+        return f"conv1.{tail[0]}"
+    if 1 <= layer <= 5:
+        block = layer - 1
+        if len(tail) == 3 and tail[0] == "layers":
+            try:
+                component = _BLOCK_COMPONENTS[int(tail[1])]
+            except (KeyError, ValueError) as exc:
+                raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
+            return f"block.{block}.{component}.{tail[2]}"
+        if len(tail) == 5 and tail[0] == "layers" and tail[2] == "layers":
+            try:
+                residual = int(tail[1]) - 1
+                component = _RESIDUAL_COMPONENTS[int(tail[3])]
+            except (KeyError, ValueError) as exc:
+                raise ValueError(f"Unsupported Stable Audio decoder key: {key}") from exc
+            if residual not in (1, 2, 3):
+                raise ValueError(f"Unsupported Stable Audio decoder key: {key}")
+            return f"block.{block}.res_unit{residual}.{component}.{tail[4]}"
+    if layer == 6 and len(tail) == 1:
+        return f"snake1.{tail[0]}"
+    if layer == 7 and len(tail) == 1:
+        return f"conv2.{tail[0]}"
+    raise ValueError(f"Unsupported Stable Audio decoder key: {key}")
+
+
+def convert_stable_audio_decoder_state_dict(
+    state_dict: Mapping[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    """Convert and filter a Stable Audio checkpoint state dictionary."""
+
+    converted: dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        target_key = convert_stable_audio_decoder_key(key)
+        if target_key is None:
+            continue
+        if target_key in converted:
+            raise ValueError(f"Duplicate converted Stable Audio key: {target_key}")
+        if target_key.endswith((".alpha", ".beta")):
+            value = value.reshape(1, -1, 1)
+        converted[target_key] = value
+    return converted
+
+
+def _load_oobleck_config(model_path: Path) -> tuple[dict[str, Any], int]:
+    config_path = model_path / "model_config.json"
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Stable Audio config does not exist: {config_path}")
+    with config_path.open(encoding="utf-8") as file:
+        full_config = json.load(file)
+
+    try:
+        pretransform = full_config["model"]["pretransform"]
+        config = pretransform["config"]
+        decoder = config["decoder"]
+        decoder_config = decoder["config"]
+        sample_rate = int(full_config["sample_rate"])
+    except (KeyError, TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid Stable Audio model config: {config_path}") from exc
+
+    if pretransform.get("type") != "autoencoder" or decoder.get("type") != "oobleck":
+        raise ValueError("MAGI-2 requires an Oobleck Stable Audio pretransform")
+    if decoder_config.get("use_nearest_upsample", False):
+        raise ValueError("Diffusers Oobleck does not support nearest-neighbor decoder upsampling")
+    if decoder_config.get("final_tanh", True):
+        raise ValueError("MAGI-2 requires the Stable Audio decoder with final_tanh=false")
+
+    strides = list(decoder_config["strides"])
+    downsampling_ratio = int(config["downsampling_ratio"])
+    if int(torch.tensor(strides).prod().item()) != downsampling_ratio:
+        raise ValueError("Stable Audio strides do not match downsampling_ratio")
+
+    kwargs = {
+        "encoder_hidden_size": int(config["encoder"]["config"]["channels"]),
+        "downsampling_ratios": strides,
+        "channel_multiples": list(decoder_config["c_mults"]),
+        "decoder_channels": int(decoder_config["channels"]),
+        "decoder_input_channels": int(decoder_config["latent_dim"]),
+        "audio_channels": int(decoder_config["out_channels"]),
+        "sampling_rate": sample_rate,
+    }
+    return kwargs, downsampling_ratio
+
+
+class Magi2AudioDecoder(nn.Module):
+    """Decode MAGI-2 audio latents with the bundled Stable Audio VAE weights."""
+
+    def __init__(
+        self,
+        model_path: str | Path,
+        *,
+        device: str | torch.device = "cpu",
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        super().__init__()
+        model_path = Path(model_path)
+        if not model_path.is_dir():
+            raise ValueError(f"Stable Audio model path must be a local directory: {model_path}")
+
+        oobleck_kwargs, self.downsampling_ratio = _load_oobleck_config(model_path)
+        self.sample_rate = int(oobleck_kwargs["sampling_rate"])
+        # Construct the Diffusers implementation for architectural compatibility,
+        # then retain only the decoder because MAGI-2 has no audio input encoder.
+        autoencoder = AutoencoderOobleck(**oobleck_kwargs)
+        self.decoder = autoencoder.decoder.to(device=device, dtype=dtype)
+        del autoencoder
+
+        weights_path = model_path / "model.safetensors"
+        if not weights_path.is_file():
+            raise FileNotFoundError(f"Stable Audio weights do not exist: {weights_path}")
+        selected: dict[str, torch.Tensor] = {}
+        with safe_open(weights_path, framework="pt", device=str(device)) as checkpoint:
+            for key in checkpoint.keys():
+                if key.startswith(_CHECKPOINT_PREFIX):
+                    selected[key] = checkpoint.get_tensor(key)
+        converted = convert_stable_audio_decoder_state_dict(selected)
+        converted = {key: value.to(dtype=dtype) for key, value in converted.items()}
+        self.decoder.load_state_dict(converted, strict=True)
+        self.eval().requires_grad_(False)
+
+    @property
+    def latent_fps(self) -> float:
+        return self.sample_rate / self.downsampling_ratio
+
+    @torch.inference_mode()
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        parameter = next(self.decoder.parameters())
+        return self.decoder(latents.to(device=parameter.device, dtype=parameter.dtype))
+
+    forward = decode
+
+
+__all__ = [
+    "Magi2AudioDecoder",
+    "convert_stable_audio_decoder_key",
+    "convert_stable_audio_decoder_state_dict",
+]

vllm_omni/diffusion/models/magi2/conditioning.py

@@ -0,0 +1,321 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Native MAGI-2 prompt conditioning.
+
+Adapted from SandAI's MAGI-2 Preview inference implementation. The model
+loading remains in-tree and uses the Transformers Qwen3.5 implementation; no
+MAGI-2 reference checkout is imported at runtime.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from tokenizers import Regex, pre_tokenizers
+from transformers import AutoTokenizer
+
+
+def _strip_empty(obj: Any) -> Any:
+    if isinstance(obj, dict):
+        return {
+            key: value
+            for key, value in ((key, _strip_empty(value)) for key, value in obj.items())
+            if value is not None and value != [] and value != {}
+        }
+    if isinstance(obj, list):
+        return [_strip_empty(value) for value in obj if value is not None]
+    return obj
+
+
+def _one_line(value: Any) -> str:
+    if value is None:
+        return ""
+    if not isinstance(value, str):
+        value = str(value)
+    return " ".join(value.split())
+
+
+def _as_dict(value: Any) -> dict[str, Any]:
+    return value if isinstance(value, dict) else {}
+
+
+def _as_list(value: Any) -> list[Any]:
+    if value is None:
+        return []
+    return value if isinstance(value, list) else [value]
+
+
+def _pick(values: dict[str, Any], keys: Iterable[str]) -> list[tuple[str, Any]]:
+    return [(key, values[key]) for key in keys if key in values and values[key] not in (None, [], {})]
+
+
+def json_to_compact_markdown(raw_json: str | dict[str, Any]) -> str:
+    """Convert MAGI's structured prompt JSON into its training-time text form."""
+
+    if isinstance(raw_json, str):
+        obj = json.loads(raw_json.strip())
+    else:
+        obj = raw_json
+
+    obj = _strip_empty(obj)
+    if not isinstance(obj, dict) or "global_layer" not in obj:
+        return raw_json if isinstance(raw_json, str) else json.dumps(raw_json, ensure_ascii=False)
+
+    global_layer = _as_dict(obj.get("global_layer"))
+    dynamic_layer = _as_dict(obj.get("dynamic_layer"))
+    reference_layer = obj.get("reference_layer", [])
+
+    lines: list[str] = []
+    context = _one_line(global_layer.get("context"))
+    description = _one_line(global_layer.get("description"))
+    if context or description:
+        lines.append(f"context: {context}" if context else "context")
+        if description:
+            lines.append(description)
+
+    aesthetics = _as_dict(global_layer.get("aesthetics"))
+    aesthetic_parts = []
+    for label, key in (("style", "style"), ("mood", "mood_atmosphere"), ("color", "color_scheme")):
+        value = _one_line(aesthetics.get(key))
+        if value:
+            aesthetic_parts.append(f"{label}={value}")
+    if aesthetic_parts:
+        lines.append("aesthetics: " + "; ".join(aesthetic_parts))
+
+    audio = _as_dict(global_layer.get("audio_baseline"))
+    dialogue = _as_dict(audio.get("dialogue"))
+    audio_parts = []
+    language = _one_line(dialogue.get("language"))
+    speakers = _as_list(dialogue.get("speaker_tags"))
+    ambience = _one_line(audio.get("ambience"))
+    if language:
+        audio_parts.append(f"language={language}")
+    if speakers:
+        audio_parts.append("speakers=" + ",".join(map(_one_line, speakers)))
+    if ambience:
+        audio_parts.append(f"ambience={ambience}")
+    if audio_parts:
+        lines.append("audio: " + "; ".join(audio_parts))
+
+    subjects = global_layer.get("alive_subjects_static")
+    if isinstance(subjects, list) and subjects:
+        lines.append("subjects:")
+        for subject in subjects:
+            if not isinstance(subject, dict):
+                continue
+            sid = _one_line(subject.get("subject_id"))
+            desc = _one_line(subject.get("description"))
+            position = _one_line(subject.get("position"))
+            orientation = _one_line(subject.get("orientation"))
+            attrs = []
+            for key, value in _pick(
+                _as_dict(subject.get("visual_attributes")),
+                ["gender", "age_group", "ethnicity", "clothing", "appearance_details"],
+            ):
+                value = _one_line(value)
+                if value:
+                    attrs.append(f"{key}={value}")
+            row = " - " + " | ".join(part for part in (sid, desc) if part)
+            extra = "; ".join(part for part in (position, orientation) if part)
+            if extra:
+                row += f" ({extra})"
+            if attrs:
+                row += " :: " + "; ".join(attrs)
+            lines.append(row)
+
+    objects = global_layer.get("objects_static")
+    if isinstance(objects, list) and objects:
+        lines.append("objects:")
+        for obj_item in objects:
+            if not isinstance(obj_item, dict):
+                continue
+            oid = _one_line(obj_item.get("object_id"))
+            desc = _one_line(obj_item.get("description"))
+            shape = _one_line(obj_item.get("shape_and_color"))
+            position = _one_line(obj_item.get("position"))
+            row = " - " + " | ".join(part for part in (oid, desc) if part)
+            details = "; ".join(part for part in (shape, position) if part)
+            if details:
+                row += " :: " + details
+            lines.append(row)
+
+    segments = dynamic_layer.get("timeline_segments")
+    if isinstance(segments, list) and segments:
+        lines.append("timeline:")
+        for segment in segments:
+            if not isinstance(segment, dict):
+                continue
+            basic = _as_dict(segment.get("segment_basic_info"))
+            timestamp = _one_line(basic.get("timestamp_range"))
+            desc = _one_line(basic.get("segment_description"))
+            head = f" - {timestamp}" if timestamp else " -"
+            if desc:
+                head += f" {desc}"
+            lines.append(head.rstrip())
+
+            segment_audio = _as_dict(segment.get("audio"))
+            for dialogue_line in segment_audio.get("dialogue_lines", []) or []:
+                if not isinstance(dialogue_line, dict):
+                    continue
+                speaker = _one_line(dialogue_line.get("speaker"))
+                text = _one_line(dialogue_line.get("text"))
+                timestamp = _one_line(dialogue_line.get("timestamp"))
+                if text:
+                    line = (f"   - dialogue {speaker}: " if speaker else "   - dialogue: ") + text
+                    if timestamp:
+                        line += f" ({timestamp})"
+                    lines.append(line)
+
+            for alive in segment.get("alive_subjects", []) or []:
+                if not isinstance(alive, dict):
+                    continue
+                sid = _one_line(alive.get("subject_id"))
+                action = _as_dict(alive.get("action"))
+                parts = [
+                    _one_line(action.get("primary_action")),
+                    _one_line(action.get("interaction")),
+                    _one_line(action.get("facial_expression")),
+                ]
+                parts = [part for part in parts if part]
+                if parts:
+                    lines.append(f"   - action {sid}: " + "; ".join(parts))
+
+            for obj_item in segment.get("objects", []) or []:
+                if not isinstance(obj_item, dict):
+                    continue
+                oid = _one_line(obj_item.get("object_id"))
+                state = _as_dict(obj_item.get("dynamic_state"))
+                parts = [_one_line(state.get("state_change")), _one_line(state.get("motion_detail"))]
+                parts = [part for part in parts if part]
+                if parts:
+                    lines.append(f"   - objects {oid}: " + "; ".join(parts))
+
+    if isinstance(reference_layer, list):
+        lines.extend(str(line) for line in reference_layer if str(line).strip())
+    return "\n".join(line for line in lines if line.strip())
+
+
+def normalize_prompt(prompt: str) -> str:
+    """Normalize structured prompts while leaving ordinary text byte-for-byte intact."""
+
+    try:
+        parsed = json.loads(prompt)
+    except (json.JSONDecodeError, TypeError):
+        return prompt
+    if isinstance(parsed, dict | list):
+        return json_to_compact_markdown(prompt)
+    return prompt
+
+
+class Magi2Qwen35TextEncoder(nn.Module):
+    """Qwen3.5-27B feature extractor used by MAGI-2 Preview."""
+
+    def __init__(
+        self,
+        model_path: str,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        max_length: int = 7000,
+        skip_layer: int = 2,
+        local_files_only: bool = True,
+    ) -> None:
+        super().__init__()
+        self.max_length = max_length
+        self.skip_layer = skip_layer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            padding_side="right",
+            local_files_only=local_files_only,
+        )
+        cjk_split = pre_tokenizers.Split(
+            pattern=Regex(
+                r"([\u1100-\u11ff\u2e80-\ua4cf\ua840-\uD7AF"
+                r"\uF900-\uFAFF\uFE30-\uFE4F\uFF65-\uFFDC\U00020000-\U0002FFFF])"
+            ),
+            behavior="isolated",
+        )
+        original = self.tokenizer.backend_tokenizer.pre_tokenizer
+        self.tokenizer.backend_tokenizer.pre_tokenizer = pre_tokenizers.Sequence([cjk_split, original])
+
+        from transformers import Qwen3_5TextModel
+
+        self.text_model = Qwen3_5TextModel.from_pretrained(
+            model_path,
+            torch_dtype=dtype,
+            local_files_only=local_files_only,
+        )
+        self.text_model.eval().requires_grad_(False)
+
+    def _tokenize(self, prompt: str, *, offsets: bool = False):
+        return self.tokenizer(
+            [normalize_prompt(prompt)],
+            return_tensors="pt",
+            padding="longest",
+            return_offsets_mapping=offsets,
+            max_length=self.max_length,
+            truncation=True,
+        )
+
+    def get_target_token_indices(self, prompt: str, target: str | None) -> list[int] | None:
+        if not target:
+            return None
+        normalized = normalize_prompt(prompt)
+        offsets = self._tokenize(prompt, offsets=True)["offset_mapping"][0]
+        start = normalized.find(target)
+        if start < 0:
+            return None
+        end = start + len(target)
+        indices = []
+        for index, (token_start, token_end) in enumerate(offsets.tolist()):
+            if token_start == 0 and token_end == 0 and index != 0:
+                continue
+            if max(token_start, start) < min(token_end, end):
+                indices.append(index)
+        return indices
+
+    def pool_figure_tokens(
+        self,
+        prompt: str,
+        targets: list[str],
+        text_features: torch.Tensor,
+    ) -> torch.Tensor:
+        """Mean-pool the token features spelling each ``<Figure N>`` marker."""
+
+        embeddings = []
+        for target in targets:
+            indices = self.get_target_token_indices(prompt, target)
+            if indices:
+                embeddings.append(text_features[0, indices].mean(dim=0).clone())
+            else:
+                embeddings.append(text_features.new_zeros(text_features.shape[-1]))
+        return torch.stack(embeddings)
+
+    @torch.inference_mode()
+    def encode(self, prompt: str) -> torch.Tensor:
+        inputs = self._tokenize(prompt)
+        device = next(self.text_model.parameters()).device
+        inputs = inputs.to(device)
+        outputs = self.text_model(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        if self.skip_layer == 0:
+            return outputs.last_hidden_state
+        return outputs.hidden_states[-(self.skip_layer + 1)]
+
+    forward = encode
+
+
+__all__ = [
+    "Magi2Qwen35TextEncoder",
+    "json_to_compact_markdown",
+    "normalize_prompt",
+]

vllm_omni/diffusion/models/magi2/configuration_magi2.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""Native configuration for the released MAGI-2 Preview checkpoint.
+
+The defaults mirror ``sand-ai/MAGI-2-preview``.  Keeping the architecture in
+vLLM-Omni makes model construction independent of SandAI's Python package and
+also gives tests a small-config entry point.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import torch
+
+
+@dataclass(frozen=True)
+class Magi2MHCConfig:
+    enabled: bool = True
+    num_streams: int = 4
+    alpha_init: float = 0.01
+    sinkhorn_iterations: int = 20
+    sinkhorn_epsilon: float = 1e-12
+
+
+@dataclass(frozen=True)
+class Magi2MoEConfig:
+    num_heads: int = 12
+    num_experts: int = 256
+    top_k: int = 6
+    expert_intermediate_size: int = 1280
+    shared_expert_intermediate_size: int = 1280
+    modality_shared_expert_intermediate_size: int = 1280
+    score_function: str = "sigmoid"
+    normalize_routing_weights: bool = True
+    routing_scale: float = 4.9
+    layers: tuple[int, ...] = tuple(range(2, 38))
+
+    def __post_init__(self) -> None:
+        if self.num_heads <= 0 or self.num_experts <= 0:
+            raise ValueError("MAGI-2 MoE head and expert counts must be positive")
+        if not 0 < self.top_k <= self.num_experts:
+            raise ValueError("MAGI-2 MoE top_k must be in [1, num_experts]")
+        if self.score_function != "sigmoid":
+            raise ValueError("The released MAGI-2 checkpoint uses sigmoid routing")
+
+
+@dataclass(frozen=True)
+class Magi2PreviewConfig:
+    num_layers: int = 40
+    hidden_size: int = 3072
+    head_dim: int = 128
+    num_query_groups: int = 24
+    video_in_channels: int = 48
+    audio_in_channels: int = 64
+    text_in_channels: int = 5120
+    intermediate_factor: float = 4.0
+    multimodal_layers: tuple[int, ...] = (0, 1, 38, 39)
+    params_dtype: torch.dtype = torch.bfloat16
+    attention_softcap: float = -1.0
+    attention_sink_tokens: int = 1
+    attention_gating: bool = True
+    mhc: Magi2MHCConfig = field(default_factory=Magi2MHCConfig)
+    moe: Magi2MoEConfig = field(default_factory=Magi2MoEConfig)
+
+    @property
+    def num_attention_heads(self) -> int:
+        if self.hidden_size % self.head_dim:
+            raise ValueError("hidden_size must be divisible by head_dim")
+        return self.hidden_size // self.head_dim
+
+    @property
+    def num_heads_q(self) -> int:
+        return self.num_attention_heads
+
+    @property
+    def num_heads_kv(self) -> int:
+        return self.num_query_groups
+
+    @property
+    def virtual_width(self) -> int:
+        return self.hidden_size * self.mhc.num_streams
+
+    def validate(self) -> None:
+        if not self.mhc.enabled:
+            raise ValueError("The released MAGI-2 Preview checkpoint requires mHC")
+        if not self.attention_gating:
+            raise ValueError("The released MAGI-2 Preview checkpoint requires attention gating")
+        if self.attention_sink_tokens < 1:
+            raise ValueError("The released MAGI-2 Preview checkpoint requires an attention sink")
+        if self.hidden_size % self.moe.num_heads:
+            raise ValueError("hidden_size must be divisible by the number of MoE heads")
+        layer_ids = set(range(self.num_layers))
+        if not set(self.multimodal_layers).issubset(layer_ids):
+            raise ValueError("multimodal layer index is outside the transformer")
+        if not set(self.moe.layers).issubset(layer_ids):
+            raise ValueError("MoE layer index is outside the transformer")
+
+
+@dataclass(frozen=True)
+class Magi2GenerationConfig:
+    duration_seconds: float = 10.0
+    fps: float = 12.5
+    output_frames: int = 125
+    preview_steps: int = 100
+    shift: float = 7.0
+    video_guidance_scale: float = 5.0
+    audio_guidance_scale: float = 7.0
+    audio_latent_fps: float = 25.0
+    video_vae_stride: tuple[int, int, int] = (8, 16, 16)
+    patch_size: tuple[int, int, int] = (1, 1, 1)
+    video_latent_channels: int = 48
+    audio_latent_channels: int = 64
+    audio_sample_rate: int = 44_100
+
+
+MAGI2_PREVIEW_CONFIG = Magi2PreviewConfig()
+MAGI2_GENERATION_CONFIG = Magi2GenerationConfig()

vllm_omni/diffusion/models/magi2/layers.py

@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""Checkpoint-compatible native layers for MAGI-2 Preview.
+
+The grouped-linear, normalization, Fourier, and mHC formulas are adapted from
+SandAI's Apache-2.0 preview implementation.  They have been modified to remove
+MagiCompiler and external Triton runtime dependencies.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .parallel import Magi2ParallelGroup, get_magi2_tp_group
+
+
+def swiglu7(
+    x: torch.Tensor,
+    alpha: float = 1.702,
+    limit: float = 7.0,
+    out_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Released GPT-OSS-style clamped SwiGLU activation."""
+
+    out_dtype = x.dtype if out_dtype is None else out_dtype
+    x = x.to(torch.float32)
+    gate, linear = x[..., ::2], x[..., 1::2]
+    gate = gate.clamp(max=limit)
+    linear = linear.clamp(min=-limit, max=limit)
+    return (gate * torch.sigmoid(alpha * gate) * (linear + 1.0)).to(out_dtype)
+
+
+class ModalityDispatcher:
+    """Precompute stable modality grouping and inverse permutation metadata."""
+
+    def __init__(self, modality_mapping: torch.Tensor, num_modalities: int) -> None:
+        if modality_mapping.ndim != 1:
+            raise ValueError("modality mapping must be a one-dimensional tensor")
+        self.modality_mapping = modality_mapping
+        self.num_modalities = num_modalities
+        self.permute_mapping = torch.argsort(modality_mapping, stable=True)
+        self.inv_permute_mapping = torch.argsort(self.permute_mapping)
+        self.permuted_modality_mapping = modality_mapping.index_select(0, self.permute_mapping)
+        self.group_size = torch.bincount(self.permuted_modality_mapping.long(), minlength=num_modalities).to(
+            torch.int32
+        )
+        self.group_size_cpu = [int(value) for value in self.group_size.cpu().tolist()]
+        self.cu_group_sizes = F.pad(torch.cumsum(self.group_size, dim=0), (1, 0))
+
+    def dispatch(self, tensor: torch.Tensor) -> list[torch.Tensor]:
+        return list(torch.split(tensor, self.group_size_cpu, dim=0))
+
+    @staticmethod
+    def undispatch(*groups: torch.Tensor) -> torch.Tensor:
+        return torch.cat(groups, dim=0)
+
+    def permute(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor.index_select(0, self.permute_mapping)
+
+    def inverse_permute(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor.index_select(0, self.inv_permute_mapping)
+
+
+class Magi2GroupedLinear(nn.Module):
+    """Modality-grouped linear with the released flattened weight layout."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        *,
+        num_experts: int = 1,
+        bias: bool = False,
+        dtype: torch.dtype | None = None,
+        device: torch.device | str | None = None,
+        parallel_mode: str | None = None,
+        qkv_splits: tuple[int, int, int] | None = None,
+        tp_group: Magi2ParallelGroup | None = None,
+    ) -> None:
+        super().__init__()
+        if parallel_mode not in (None, "column", "row"):
+            raise ValueError(f"unknown MAGI-2 linear parallel mode {parallel_mode!r}")
+        if qkv_splits is not None and parallel_mode != "column":
+            raise ValueError("segmented QKV slicing requires column parallelism")
+        if qkv_splits is not None and sum(qkv_splits) != out_features:
+            raise ValueError("QKV splits must sum to out_features")
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_experts = num_experts
+        self.parallel_mode = parallel_mode
+        self.qkv_splits = qkv_splits
+        self.tp_group = tp_group or get_magi2_tp_group()
+        self.local_in_features = in_features
+        self.local_out_features = out_features
+        if parallel_mode == "column" and self.tp_group.world_size > 1:
+            split_dims = qkv_splits or (out_features,)
+            if any(size % self.tp_group.world_size for size in split_dims):
+                raise ValueError(
+                    f"column-parallel output dimensions {split_dims} must divide TP={self.tp_group.world_size}"
+                )
+            self.local_out_features = sum(size // self.tp_group.world_size for size in split_dims)
+        elif parallel_mode == "row" and self.tp_group.world_size > 1:
+            if in_features % self.tp_group.world_size:
+                raise ValueError(f"row-parallel input {in_features} must divide TP={self.tp_group.world_size}")
+            self.local_in_features = in_features // self.tp_group.world_size
+
+        self.weight = nn.Parameter(
+            torch.empty(
+                num_experts * self.local_out_features,
+                self.local_in_features,
+                dtype=dtype,
+                device=device,
+            )
+        )
+        if bias:
+            bias_features = self.local_out_features if parallel_mode == "column" else out_features
+            self.bias = nn.Parameter(torch.empty(num_experts * bias_features, dtype=dtype, device=device))
+        else:
+            self.register_parameter("bias", None)
+        if self.tp_group.world_size > 1 and parallel_mode is not None:
+            self.weight.checkpoint_weight_transform = self.shard_checkpoint_weight
+            if self.bias is not None and parallel_mode == "column":
+                self.bias.checkpoint_weight_transform = self.shard_checkpoint_bias
+
+    def _column_slices(self) -> tuple[tuple[int, int], ...]:
+        splits = self.qkv_splits or (self.out_features,)
+        offsets: list[tuple[int, int]] = []
+        base = 0
+        for size in splits:
+            local = size // self.tp_group.world_size
+            start = base + self.tp_group.rank * local
+            offsets.append((start, start + local))
+            base += size
+        return tuple(offsets)
+
+    def shard_checkpoint_weight(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
+        """Convert the released grouped weight into this TP rank's shard."""
+
+        if tuple(checkpoint_tensor.shape) == tuple(self.weight.shape):
+            return checkpoint_tensor
+        expected = (self.num_experts * self.out_features, self.in_features)
+        if tuple(checkpoint_tensor.shape) != expected:
+            raise ValueError(
+                f"grouped linear checkpoint has shape {tuple(checkpoint_tensor.shape)}, expected {expected}"
+            )
+        grouped = checkpoint_tensor.view(self.num_experts, self.out_features, self.in_features)
+        if self.parallel_mode == "column":
+            shards = [grouped[:, start:end] for start, end in self._column_slices()]
+            return torch.cat(shards, dim=1).reshape(self.num_experts * self.local_out_features, self.in_features)
+        if self.parallel_mode == "row":
+            start = self.tp_group.rank * self.local_in_features
+            return grouped[:, :, start : start + self.local_in_features].reshape(
+                self.num_experts * self.out_features,
+                self.local_in_features,
+            )
+        return checkpoint_tensor
+
+    def shard_checkpoint_bias(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
+        if self.bias is None:
+            raise ValueError("cannot load bias into a bias-free grouped linear")
+        if tuple(checkpoint_tensor.shape) == tuple(self.bias.shape):
+            return checkpoint_tensor
+        expected = (self.num_experts * self.out_features,)
+        if tuple(checkpoint_tensor.shape) != expected:
+            raise ValueError(f"grouped linear bias has shape {tuple(checkpoint_tensor.shape)}, expected {expected}")
+        grouped = checkpoint_tensor.view(self.num_experts, self.out_features)
+        shards = [grouped[:, start:end] for start, end in self._column_slices()]
+        return torch.cat(shards, dim=1).reshape(-1)
+
+    def forward(
+        self,
+        tensor: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher | None = None,
+    ) -> torch.Tensor:
+        weight = self.weight.view(self.num_experts, self.local_out_features, self.local_in_features)
+        bias_features = self.local_out_features if self.parallel_mode == "column" else self.out_features
+        bias = self.bias.view(self.num_experts, bias_features) if self.bias is not None else None
+        linear_bias = None if self.parallel_mode == "row" else bias
+        if self.num_experts == 1:
+            output = F.linear(tensor, weight[0], None if linear_bias is None else linear_bias[0])
+        else:
+            if modality_dispatcher is None:
+                raise ValueError("modality_dispatcher is required for grouped linear")
+            if modality_dispatcher.num_modalities != self.num_experts:
+                raise ValueError("grouped linear expert count does not match modality dispatcher")
+            inputs = modality_dispatcher.dispatch(tensor)
+            outputs = [
+                F.linear(part, weight[index], None if linear_bias is None else linear_bias[index])
+                for index, part in enumerate(inputs)
+            ]
+            output = torch.cat(outputs, dim=0)
+
+        if self.parallel_mode == "row" and self.tp_group.world_size > 1:
+            dist.all_reduce(output, group=self.tp_group.group)
+        if self.parallel_mode == "row" and bias is not None:
+            if self.num_experts == 1:
+                output = output + bias[0]
+            else:
+                assert modality_dispatcher is not None
+                outputs = [part + bias[index] for index, part in enumerate(modality_dispatcher.dispatch(output))]
+                output = torch.cat(outputs, dim=0)
+        return output
+
+
+def make_grouped_linear(
+    in_features: int,
+    out_features: int,
+    *,
+    num_experts: int = 1,
+    bias: bool = False,
+    dtype: torch.dtype | None = None,
+    parallel_mode: str | None = None,
+    qkv_splits: tuple[int, int, int] | None = None,
+) -> Magi2GroupedLinear:
+    return Magi2GroupedLinear(
+        in_features,
+        out_features,
+        num_experts=num_experts,
+        bias=bias,
+        dtype=dtype,
+        parallel_mode=parallel_mode,
+        qkv_splits=qkv_splits,
+    )
+
+
+class MultiModalityRMSNorm(nn.Module):
+    """RMSNorm with independent modality and mHC-stream scales."""
+
+    def __init__(
+        self,
+        dim: int,
+        *,
+        eps: float = 1e-6,
+        num_modality: int = 1,
+        num_patterns: int = 1,
+        out_dtype: torch.dtype | None = None,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.num_modality = num_modality
+        self.num_patterns = num_patterns
+        self.out_dtype = out_dtype
+        self.weight = nn.Parameter(torch.zeros(num_patterns * dim * num_modality, dtype=torch.float32, device=device))
+
+    def forward(
+        self,
+        tensor: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher | None = None,
+    ) -> torch.Tensor:
+        original_dtype = tensor.dtype
+        normalized = tensor.float()
+        normalized = normalized * torch.rsqrt(normalized.square().mean(dim=-1, keepdim=True) + self.eps)
+        if self.num_modality == 1:
+            weight = self.weight.view(self.num_patterns, self.dim) + 1.0
+            result = normalized * weight
+        else:
+            if modality_dispatcher is None:
+                raise ValueError("modality_dispatcher is required for multimodal RMSNorm")
+            inputs = modality_dispatcher.dispatch(normalized)
+            weights = self.weight.view(self.num_modality, self.num_patterns, self.dim)
+            result = modality_dispatcher.undispatch(
+                *(part * (weights[index] + 1.0) for index, part in enumerate(inputs))
+            )
+        return result.to(self.out_dtype or original_dtype)
+
+
+def _frequency_bands(
+    num_bands: int,
+    *,
+    temperature: float,
+    device: torch.device | str | None,
+) -> torch.Tensor:
+    exponent = torch.arange(num_bands, dtype=torch.float32, device=device) / num_bands
+    return 1.0 / temperature**exponent
+
+
+class ElementWiseFourierEmbed(nn.Module):
+    """Nine-coordinate Fourier embedding used as MAGI's element-wise RoPE."""
+
+    def __init__(
+        self,
+        dim: int,
+        *,
+        temperature: float = 10000.0,
+        learnable: bool = False,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.temperature = temperature
+        bands = _frequency_bands(dim // 8, temperature=temperature, device=device).to(dtype)
+        if learnable:
+            self.bands = nn.Parameter(bands)
+        else:
+            # ``bands`` is part of the released checkpoint.  Keep it in the
+            # parameter namespace so the generic DLO mmap loader can bind it
+            # without first materializing a full CPU transformer.  It remains
+            # immutable during inference.
+            self.bands = nn.Parameter(bands, requires_grad=False)
+
+    def forward(self, coords: torch.Tensor) -> torch.Tensor:
+        if coords.ndim != 2 or coords.shape[-1] != 9:
+            raise ValueError("MAGI coordinates must have shape [tokens,9]")
+        xyz = coords[:, :3]
+        sizes = coords[:, 3:6]
+        references = coords[:, 6:9]
+        scales = (references - 1) / (sizes - 1)
+        scales = torch.where((references == 1) & (sizes == 1), torch.ones_like(scales), scales)
+        if not torch.isfinite(scales).all():
+            raise ValueError("invalid MAGI coordinate scale")
+        centers = (sizes - 1) / 2
+        centers = centers.clone()
+        centers[:, 0] = 0
+        projection = (xyz - centers).unsqueeze(-1) * scales.unsqueeze(-1) * self.bands
+        return torch.cat((projection.sin(), projection.cos()), dim=1).flatten(1)
+
+
+MHCTensorTuple = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+
+
+def sinkhorn_knopp(matrix_logits: torch.Tensor, iterations: int, epsilon: float) -> torch.Tensor:
+    matrix = torch.exp(matrix_logits - matrix_logits.amax(dim=(-2, -1), keepdim=True))
+    for _ in range(iterations):
+        matrix = matrix / (matrix.sum(dim=-2, keepdim=True) + epsilon)
+        matrix = matrix / (matrix.sum(dim=-1, keepdim=True) + epsilon)
+    return matrix
+
+
+class MHCHandler:
+    """Exact four-stream manifold-constrained hyper-connection math."""
+
+    def __init__(
+        self,
+        num_streams: int,
+        hidden_size: int,
+        *,
+        sinkhorn_iterations: int = 20,
+        sinkhorn_epsilon: float = 1e-12,
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        self.num_streams = num_streams
+        self.hidden_size = hidden_size
+        self.sinkhorn_iterations = sinkhorn_iterations
+        self.sinkhorn_epsilon = sinkhorn_epsilon
+        self.dtype = dtype
+        self.matmul_scale = 1.0 / math.sqrt(float(num_streams * hidden_size))
+
+    def flatten(self, tensor: torch.Tensor) -> torch.Tensor:
+        self._check_multi(tensor)
+        return tensor.view(tensor.shape[0], -1)
+
+    def compute_logits(
+        self,
+        flattened: torch.Tensor,
+        norm: Callable[[torch.Tensor], torch.Tensor],
+        phi_fused: torch.Tensor,
+    ) -> MHCTensorTuple:
+        if flattened.ndim != 2 or flattened.shape[-1] != self.num_streams * self.hidden_size:
+            raise ValueError("invalid flattened mHC shape")
+        fused = norm(flattened).to(self.dtype) @ phi_fused
+        pre, post, residual = torch.split(
+            fused,
+            (self.num_streams, self.num_streams, self.num_streams**2),
+            dim=-1,
+        )
+        return pre, post, residual.view(-1, self.num_streams, self.num_streams)
+
+    def apply_pre(
+        self,
+        streams: torch.Tensor,
+        alpha_bias_logits: MHCTensorTuple,
+        *,
+        out_dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
+        self._check_multi(streams)
+        alpha, bias, logits = alpha_bias_logits
+        coefficients = torch.sigmoid(alpha * self.matmul_scale * logits + bias.unsqueeze(0))
+        return torch.einsum("tn,tnc->tc", coefficients.to(out_dtype or streams.dtype), streams)
+
+    def compute_post_residual(
+        self,
+        post: MHCTensorTuple,
+        residual: MHCTensorTuple,
+        *,
+        out_dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        alpha_post, bias_post, post_logits = post
+        alpha_residual, bias_residual, residual_logits = residual
+        post_coefficients = 2.0 * torch.sigmoid(alpha_post * self.matmul_scale * post_logits + bias_post.unsqueeze(0))
+        residual_matrix = sinkhorn_knopp(
+            alpha_residual * self.matmul_scale * residual_logits.float() + bias_residual.unsqueeze(0).float(),
+            self.sinkhorn_iterations,
+            self.sinkhorn_epsilon,
+        )
+        return post_coefficients.to(out_dtype), residual_matrix.to(out_dtype)
+
+    def hyper_connect(
+        self,
+        residual_streams: torch.Tensor,
+        branch_output: torch.Tensor,
+        post_coefficients: torch.Tensor,
+        residual_matrix: torch.Tensor,
+    ) -> torch.Tensor:
+        self._check_multi(residual_streams)
+        if branch_output.ndim != 2 or branch_output.shape[-1] != self.hidden_size:
+            raise ValueError("invalid mHC branch-output shape")
+        branch = torch.einsum("tn,tc->tnc", post_coefficients, branch_output)
+        mixed = torch.einsum("tij,tjc->tic", residual_matrix, residual_streams)
+        return mixed + branch
+
+    def _check_multi(self, tensor: torch.Tensor) -> None:
+        expected = (self.num_streams, self.hidden_size)
+        if tensor.ndim != 3 or tensor.shape[1:] != expected:
+            raise ValueError(f"expected mHC tensor [tokens,{expected[0]},{expected[1]}], got {tuple(tensor.shape)}")
+
+
+__all__ = [
+    "ElementWiseFourierEmbed",
+    "MHCHandler",
+    "Magi2GroupedLinear",
+    "ModalityDispatcher",
+    "MultiModalityRMSNorm",
+    "make_grouped_linear",
+    "sinkhorn_knopp",
+    "swiglu7",
+]

vllm_omni/diffusion/models/magi2/mh_moe.py

@@ -0,0 +1,547 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2025-2026 SandAI. All Rights Reserved.
+
+"""Native multi-head MoE used by MAGI-2 Preview.
+
+Adapted from SandAI's Apache-2.0 ``flash_mh_moe`` implementation and modified
+to use vLLM's existing expert-parallel group.  MAGI's routing is unusual: each
+of twelve 256-wide hidden-state heads independently selects experts from its
+own 256-expert bank.  It is therefore not representable by vLLM's conventional
+whole-token :class:`FusedMoE` primitive.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.triton_utils import tl, triton
+
+from vllm_omni.platforms import current_omni_platform
+
+from .parallel import Magi2ParallelGroup, ep_dispatch, ep_undispatch, get_magi2_ep_group
+
+RoutingScore = Literal["softmax", "sigmoid"]
+
+
+def swiglu7_pair(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    """Released clamped SwiGLU7 expert activation, evaluated in fp32."""
+
+    dtype = gate.dtype
+    gate = gate.float().clamp(max=7.0)
+    up = up.float().clamp(min=-7.0, max=7.0)
+    return (gate * torch.sigmoid(1.702 * gate) * (up + 1.0)).to(dtype)
+
+
+def compute_topk_probs_and_indices(
+    router_logits: torch.Tensor,
+    top_k: int,
+    *,
+    score_func: RoutingScore = "sigmoid",
+    expert_bias: torch.Tensor | None = None,
+    route_norm: bool = True,
+    norm_eps: float = 1e-12,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Route independently for every ``[head, token]`` pair.
+
+    The auxiliary-free bias affects expert selection but deliberately does not
+    affect the returned routing probability, matching the training recipe.
+    """
+
+    if router_logits.ndim != 3:
+        raise ValueError("router_logits must be [heads,tokens,experts]")
+    if not 0 < top_k <= router_logits.shape[-1]:
+        raise ValueError("top_k must be in [1, num_experts]")
+    if score_func == "sigmoid":
+        router_scores = torch.sigmoid(router_logits)
+    elif score_func == "softmax":
+        router_scores = torch.softmax(router_logits, dim=-1)
+    else:
+        raise ValueError(f"unsupported routing score function {score_func!r}")
+    selection_scores = router_scores
+    if expert_bias is not None:
+        selection_scores = selection_scores + expert_bias.view(router_logits.shape[0], 1, -1)
+    # Keep the reference's default sorted=True behavior.  Besides defining the
+    # route order for ties, this also fixes the reduction order used by the
+    # following L1 normalization.
+    topk_indices = torch.topk(selection_scores, top_k, dim=-1).indices
+    topk_probs = router_scores.gather(-1, topk_indices)
+    if route_norm:
+        topk_probs = F.normalize(topk_probs, p=1, dim=-1, eps=norm_eps)
+    return topk_probs, topk_indices
+
+
+def global_sort_routes(
+    topk_probs: torch.Tensor,
+    topk_indices: torch.Tensor,
+    num_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Convert per-head routes into a stable flattened-expert CSR layout."""
+
+    if topk_probs.shape != topk_indices.shape or topk_indices.ndim != 3:
+        raise ValueError("top-k probabilities and indices must have the same [H,S,K] shape")
+    heads, sequence, top_k = topk_indices.shape
+    device = topk_indices.device
+    head_offset = torch.arange(heads, device=device).view(heads, 1, 1) * num_experts
+    flattened_experts = (topk_indices + head_offset).reshape(-1)
+    flat_probs = topk_probs.reshape(-1)
+    flat_tokens = torch.arange(sequence, device=device).view(1, sequence, 1).expand(heads, sequence, top_k).reshape(-1)
+    order = flattened_experts.argsort(stable=True)
+    gather_ids = flat_tokens[order].to(torch.int32)
+    sorted_probs = flat_probs[order].float()
+    counts = torch.bincount(flattened_experts, minlength=heads * num_experts)
+    offsets = torch.zeros(heads * num_experts + 1, device=device, dtype=torch.long)
+    offsets[1:] = counts.cumsum(0)
+    return gather_ids, sorted_probs, offsets
+
+
+def torch_mh_moe_forward(
+    x: torch.Tensor,
+    gather_ids: torch.Tensor,
+    probs: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    w_gate: torch.Tensor,
+    w_up: torch.Tensor,
+    w_down: torch.Tensor,
+) -> torch.Tensor:
+    """Small-shape correctness oracle for the fused expert kernel."""
+
+    if x.ndim != 3:
+        raise ValueError("multi-head MoE input must be [tokens,heads,head_dim]")
+    output = torch.zeros_like(x)
+    experts_per_head = (expert_offsets.numel() - 1) // x.shape[1]
+    for flat_expert in range(expert_offsets.numel() - 1):
+        begin = int(expert_offsets[flat_expert].item())
+        end = int(expert_offsets[flat_expert + 1].item())
+        if begin == end:
+            continue
+        head = flat_expert // experts_per_head
+        token_ids = gather_ids[begin:end].long()
+        expert_input = x.index_select(0, token_ids)[:, head]
+        gate = expert_input @ w_gate[flat_expert]
+        up = expert_input @ w_up[flat_expert]
+        hidden = swiglu7_pair(gate, up)
+        expert_output = hidden @ w_down[flat_expert]
+        expert_output = expert_output * probs[begin:end, None].to(expert_output.dtype)
+        output[:, head].index_add_(0, token_ids, expert_output)
+    return output
+
+
+_SWIGLU7_ALPHA = tl.constexpr(1.702)
+_SWIGLU7_LIMIT = tl.constexpr(7.0)
+_SWIGLU7_BIAS = tl.constexpr(1.0)
+
+
+@triton.jit
+def _swiglu7_kernel(gate, up, out_dtype: tl.constexpr):
+    gate_clamped = tl.minimum(gate, _SWIGLU7_LIMIT)
+    up_clamped = tl.maximum(tl.minimum(up, _SWIGLU7_LIMIT), -_SWIGLU7_LIMIT)
+    sigmoid = tl.sigmoid(_SWIGLU7_ALPHA * gate_clamped)
+    swish = gate_clamped * sigmoid
+    return (swish * (up_clamped + _SWIGLU7_BIAS)).to(out_dtype)
+
+
+@triton.jit
+def _binary_search_expert(
+    cumulative_tiles,
+    tile_id,
+    num_experts: tl.constexpr,
+    log2_num_experts: tl.constexpr,
+):
+    lo = 0
+    hi = num_experts
+    for _ in tl.static_range(0, log2_num_experts + 1):
+        mid = (lo + hi + 1) // 2
+        below = tl.load(cumulative_tiles + mid) <= tile_id
+        lo = tl.where(below, mid, lo)
+        hi = tl.where(below, hi, mid - 1)
+    return lo
+
+
+@triton.jit
+def _mh_moe_kernel(
+    x_ptr,
+    wg_ptr,
+    wu_ptr,
+    wd_ptr,
+    y_ptr,
+    gather_ids_ptr,
+    probs_ptr,
+    expert_offsets_ptr,
+    cumulative_tiles_ptr,
+    stride_x_s,
+    stride_x_h,
+    stride_x_dh,
+    stride_wg_e,
+    stride_wg_dh,
+    stride_wg_de,
+    stride_wu_e,
+    stride_wu_dh,
+    stride_wu_de,
+    stride_wd_e,
+    stride_wd_de,
+    stride_wd_dh,
+    stride_y_s,
+    stride_y_h,
+    stride_y_dh,
+    d_head: tl.constexpr,
+    d_expert: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_flat_experts: tl.constexpr,
+    log2_num_experts: tl.constexpr,
+    block_t: tl.constexpr,
+    block_dh: tl.constexpr,
+    block_de: tl.constexpr,
+    acc_dtype: tl.constexpr = tl.float32,
+    deterministic: tl.constexpr = False,
+):
+    tile_id = tl.program_id(0)
+    total_tiles = tl.load(cumulative_tiles_ptr + num_flat_experts)
+    if tile_id >= total_tiles:
+        return
+
+    expert = _binary_search_expert(cumulative_tiles_ptr, tile_id, num_flat_experts, log2_num_experts)
+    expert_i64 = expert.to(tl.int64)
+    head = expert // (num_flat_experts // num_heads)
+    tile_in_expert = tile_id - tl.load(cumulative_tiles_ptr + expert)
+    token_start = tl.load(expert_offsets_ptr + expert) + tile_in_expert * block_t
+    expert_end = tl.load(expert_offsets_ptr + expert + 1)
+    count = tl.minimum(token_start + block_t, expert_end) - token_start
+
+    dh_block_offsets = tl.arange(0, block_dh)
+    de_block_offsets = tl.arange(0, block_de)
+    token_offsets = tl.arange(0, block_t)
+    dh_offsets = tl.arange(0, d_head)
+
+    token_positions = token_start + token_offsets
+    token_mask = token_offsets < count
+    # Token indices fit in int32, but multiplying a large packed-batch index by
+    # the hidden-width stride does not. Promote before computing element offsets.
+    gather_ids = tl.load(gather_ids_ptr + token_positions, mask=token_mask, other=0).to(tl.int64)
+    probabilities = tl.load(probs_ptr + token_positions, mask=token_mask, other=0.0)
+    x_base = gather_ids * stride_x_s + head * stride_x_h
+    output_acc = tl.zeros([block_t, d_head], dtype=acc_dtype)
+
+    for de_start in tl.range(0, d_expert, block_de):
+        de_offsets = de_start + de_block_offsets
+        gate_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
+        up_acc = tl.zeros([block_t, block_de], dtype=acc_dtype)
+        for dh_start in tl.static_range(0, d_head, block_dh):
+            local_dh = dh_start + dh_block_offsets
+            x_block = tl.load(
+                x_ptr + x_base[:, None] + local_dh[None, :] * stride_x_dh,
+                mask=token_mask[:, None],
+                other=0.0,
+            )
+            wg = tl.load(
+                wg_ptr
+                + expert_i64 * stride_wg_e
+                + local_dh[:, None] * stride_wg_dh
+                + de_offsets[None, :] * stride_wg_de
+            )
+            wu = tl.load(
+                wu_ptr
+                + expert_i64 * stride_wu_e
+                + local_dh[:, None] * stride_wu_dh
+                + de_offsets[None, :] * stride_wu_de
+            )
+            gate_acc += tl.dot(x_block, wg)
+            up_acc += tl.dot(x_block, wu)
+        hidden = _swiglu7_kernel(gate_acc, up_acc, wd_ptr.dtype.element_ty)
+        down = tl.load(
+            wd_ptr + expert_i64 * stride_wd_e + de_offsets[:, None] * stride_wd_de + dh_offsets[None, :] * stride_wd_dh
+        )
+        output_acc += tl.dot(hidden, down)
+    output_acc = output_acc * probabilities[:, None]
+
+    if deterministic:
+        output_ptrs = y_ptr + token_positions[:, None] * stride_y_s + dh_offsets[None, :] * stride_y_dh
+        tl.store(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
+    else:
+        output_base = gather_ids * stride_y_s + head * stride_y_h
+        output_ptrs = y_ptr + output_base[:, None] + dh_offsets[None, :] * stride_y_dh
+        tl.atomic_add(output_ptrs, output_acc.to(y_ptr.dtype.element_ty), mask=token_mask[:, None])
+
+
+def _deterministic_scatter(
+    sorted_output: torch.Tensor,
+    reference: torch.Tensor,
+    gather_ids: torch.Tensor,
+    expert_offsets: torch.Tensor,
+) -> torch.Tensor:
+    num_flat_experts = expert_offsets.numel() - 1
+    experts_per_head = num_flat_experts // reference.shape[1]
+    expert_lengths = torch.diff(expert_offsets)
+    head_values = torch.arange(num_flat_experts, device=gather_ids.device) // experts_per_head
+    head_ids = torch.repeat_interleave(head_values, expert_lengths)
+    scatter_ids = gather_ids.long() * reference.shape[1] + head_ids.long()
+    output = torch.zeros_like(reference).view(-1, reference.shape[-1])
+    output.scatter_add_(0, scatter_ids[:, None].expand_as(sorted_output), sorted_output.to(output.dtype))
+    return output.view_as(reference)
+
+
+def _select_block_config() -> tuple[int, int, int, int, int]:
+    """Return the reference kernel config, capped for pre-Blackwell GPUs."""
+
+    capability = current_omni_platform.get_device_capability()
+    if capability is not None and capability.major >= 10:  # Blackwell
+        return (128, 64, 32, 2, 8)
+    # BLOCK_T=128 needs 122,880 bytes of shared memory and is not safe on the
+    # qualified L20X path.  This is the reference kernel's portable config.
+    return (64, 64, 32, 2, 4)
+
+
+def triton_mh_moe_forward(
+    x: torch.Tensor,
+    gather_ids: torch.Tensor,
+    probs: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    w_gate: torch.Tensor,
+    w_up: torch.Tensor,
+    w_down: torch.Tensor,
+    *,
+    deterministic: bool = False,
+) -> torch.Tensor:
+    """Fused gather/expert/scatter kernel for released MAGI dimensions."""
+
+    routed_tokens = gather_ids.numel()
+    if routed_tokens == 0:
+        return torch.zeros_like(x)
+    d_head, d_expert = x.shape[-1], w_down.shape[1]
+    block_t, block_dh, block_de, num_stages, num_warps = _select_block_config()
+    if d_head % block_dh or d_expert % block_de:
+        return torch_mh_moe_forward(x, gather_ids, probs, expert_offsets, w_gate, w_up, w_down)
+
+    if deterministic:
+        output = torch.empty((routed_tokens, 1, d_head), device=x.device, dtype=x.dtype)
+    else:
+        output = torch.zeros_like(x)
+    num_flat_experts = expert_offsets.numel() - 1
+    expert_tiles = (torch.diff(expert_offsets) + block_t - 1) // block_t
+    cumulative_tiles = torch.cat(
+        (torch.zeros(1, dtype=torch.int32, device=x.device), expert_tiles.cumsum(0, dtype=torch.int32))
+    )
+    # Match the reference launch bound.  Empty/excess programs return after
+    # comparing against ``cumulative_tiles[-1]`` inside the kernel.
+    grid = ((routed_tokens + block_t - 1) // block_t + num_flat_experts,)
+    log2_experts = max(1, math.ceil(math.log2(max(num_flat_experts, 1) + 1)))
+    _mh_moe_kernel[grid](
+        x,
+        w_gate,
+        w_up,
+        w_down,
+        output,
+        gather_ids,
+        probs,
+        expert_offsets,
+        cumulative_tiles,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        w_gate.stride(0),
+        w_gate.stride(1),
+        w_gate.stride(2),
+        w_up.stride(0),
+        w_up.stride(1),
+        w_up.stride(2),
+        w_down.stride(0),
+        w_down.stride(1),
+        w_down.stride(2),
+        output.stride(0),
+        output.stride(1),
+        output.stride(2),
+        d_head,
+        d_expert,
+        x.shape[1],
+        num_flat_experts,
+        log2_experts,
+        block_t,
+        block_dh,
+        block_de,
+        tl.float32,
+        deterministic,
+        num_stages=num_stages,
+        num_warps=num_warps,
+    )
+    if deterministic:
+        return _deterministic_scatter(output.view(routed_tokens, d_head), x, gather_ids, expert_offsets)
+    return output
+
+
+@dataclass(frozen=True)
+class Magi2MultiHeadMoEConfig:
+    hidden_size: int
+    num_heads: int
+    num_experts: int
+    top_k: int
+    expert_intermediate_size: int
+    params_dtype: torch.dtype
+    score_func: RoutingScore = "sigmoid"
+    route_norm: bool = True
+    route_scale: float = 1.0
+
+
+class Magi2MultiHeadMoE(nn.Module):
+    """Checkpoint-compatible MAGI-2 head-routed expert layer."""
+
+    _EP_SHARDED_PARAMETER_NAMES = frozenset(
+        {"gate", "W_gate", "W_up", "W_down", "router.expert_bias", "router.expert_bias_ema"}
+    )
+
+    def __init__(
+        self,
+        config: Magi2MultiHeadMoEConfig,
+        *,
+        ep_group: Magi2ParallelGroup | None = None,
+    ) -> None:
+        super().__init__()
+        if config.hidden_size % config.num_heads:
+            raise ValueError("hidden_size must be divisible by the number of MoE heads")
+        self.config = config
+        self.num_heads = config.num_heads
+        self.num_experts = config.num_experts
+        self.top_k = config.top_k
+        self.d_head = config.hidden_size // config.num_heads
+        self.d_expert = config.expert_intermediate_size
+        self.ep_group = ep_group or get_magi2_ep_group()
+        self.padded_num_heads = math.ceil(self.num_heads / self.ep_group.world_size) * self.ep_group.world_size
+        self.local_num_heads = self.padded_num_heads // self.ep_group.world_size
+        self.local_flatten_num_experts = self.local_num_heads * self.num_experts
+        self.ep_pad_heads = self.padded_num_heads - self.num_heads
+        self.local_head_start = self.ep_group.rank * self.local_num_heads
+        self.has_real_moe_heads = self.local_head_start < self.num_heads
+
+        self.gate = nn.Parameter(torch.empty(self.local_flatten_num_experts, self.d_head, dtype=torch.float32))
+        self.W_gate = nn.Parameter(
+            torch.empty(self.local_flatten_num_experts, self.d_head, self.d_expert, dtype=config.params_dtype)
+        )
+        self.W_up = nn.Parameter(
+            torch.empty(self.local_flatten_num_experts, self.d_head, self.d_expert, dtype=config.params_dtype)
+        )
+        self.W_down = nn.Parameter(
+            torch.empty(self.local_flatten_num_experts, self.d_expert, self.d_head, dtype=config.params_dtype)
+        )
+        self.router = nn.Module()
+        # Both tensors are released checkpoint entries.  Non-trainable
+        # Parameters let the DLO mmap path bind them on a meta-constructed
+        # model; persistent buffers are intentionally not mmap-loaded by the
+        # generic backend.
+        self.router.expert_bias = nn.Parameter(
+            torch.zeros(self.local_flatten_num_experts, dtype=torch.float32),
+            requires_grad=False,
+        )
+        self.router.expert_bias_ema = nn.Parameter(
+            torch.zeros(self.local_flatten_num_experts, dtype=torch.float32),
+            requires_grad=False,
+        )
+
+        for name in self._EP_SHARDED_PARAMETER_NAMES:
+            target: nn.Module | Magi2MultiHeadMoE = self
+            parts = name.split(".")
+            for part in parts[:-1]:
+                target = getattr(target, part)
+            parameter = getattr(target, parts[-1])
+            parameter.mmap_weight_transform = self.ep_slice
+
+    def ep_slice(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
+        """Slice flattened ``(head,expert)`` checkpoint rows for this rank."""
+
+        if checkpoint_tensor.shape[0] == self.local_flatten_num_experts:
+            return checkpoint_tensor
+        start = self.local_head_start * self.num_experts
+        end = min(start + self.local_flatten_num_experts, checkpoint_tensor.shape[0])
+        if start >= checkpoint_tensor.shape[0]:
+            return torch.zeros(
+                (self.local_flatten_num_experts, *checkpoint_tensor.shape[1:]),
+                dtype=checkpoint_tensor.dtype,
+                device=checkpoint_tensor.device,
+            )
+        local = checkpoint_tensor[start:end]
+        if local.shape[0] < self.local_flatten_num_experts:
+            # Uneven EP/head partitions require materialized zero padding;
+            # divisible production layouts keep the mmap-backed slice above.
+            padding = torch.zeros(
+                (self.local_flatten_num_experts - local.shape[0], *local.shape[1:]),
+                dtype=local.dtype,
+                device=local.device,
+            )
+            local = torch.cat((local, padding), dim=0)
+        return local
+
+    def _route(self, x_heads: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        gate = self.gate.view(self.local_num_heads, self.num_experts, self.d_head).float()
+        logits = torch.einsum("shd,hed->hse", x_heads.float(), gate)
+        bias_source = (os.environ.get("MAGI2_ROUTER_BIAS_SOURCE") or "ema").strip().lower()
+        bias_tensor = self.router.expert_bias if bias_source == "main" else self.router.expert_bias_ema
+        bias = bias_tensor.view(self.local_num_heads, self.num_experts)
+        probs, indices = compute_topk_probs_and_indices(
+            logits,
+            self.top_k,
+            score_func=self.config.score_func,
+            expert_bias=bias,
+            route_norm=self.config.route_norm,
+        )
+        return probs * self.config.route_scale, indices
+
+    def _local_forward(self, x_heads: torch.Tensor) -> torch.Tensor:
+        probabilities, indices = self._route(x_heads)
+        gather_ids, sorted_probs, offsets = global_sort_routes(probabilities, indices, self.num_experts)
+        if x_heads.is_cuda:
+            return triton_mh_moe_forward(
+                x_heads,
+                gather_ids,
+                sorted_probs,
+                offsets,
+                self.W_gate,
+                self.W_up,
+                self.W_down,
+                deterministic=os.environ.get("MAGI2_DETERMINISTIC", "0") == "1",
+            )
+        return torch_mh_moe_forward(x_heads, gather_ids, sorted_probs, offsets, self.W_gate, self.W_up, self.W_down)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.ep_group.world_size > 1 and self.ep_group.replicated_sequence:
+            # TP column-parallel ``split_linear`` already emits exactly this
+            # rank's contiguous MoE-head slice.  Compute it once and leave it
+            # sharded for the row-parallel ``merge_linear``; no token dispatch
+            # or head all-gather belongs on the true TP path.
+            local_hidden_size = self.local_num_heads * self.d_head
+            if x.shape[-1] != local_hidden_size:
+                raise ValueError(f"TP-local MAGI MoE input has width {x.shape[-1]}, expected {local_hidden_size}")
+            local = x.view(-1, self.local_num_heads, self.d_head)
+            output = self._local_forward(local) if self.has_real_moe_heads else torch.zeros_like(local)
+            return output.reshape(-1, local_hidden_size)
+
+        x_heads = x.view(-1, self.num_heads, self.d_head)
+        if self.ep_pad_heads:
+            padding = x_heads.new_zeros((x_heads.shape[0], self.ep_pad_heads, self.d_head))
+            x_heads = torch.cat((x_heads, padding), dim=1)
+        sequence_split_sizes: list[int] | None = None
+        if self.ep_group.world_size > 1:
+            local_size = torch.tensor([x_heads.shape[0]], dtype=torch.int64, device=x_heads.device)
+            gathered_sizes = [torch.empty_like(local_size) for _ in range(self.ep_group.world_size)]
+            torch.distributed.all_gather(gathered_sizes, local_size, group=self.ep_group.group)
+            sequence_split_sizes = [int(size.item()) for size in gathered_sizes]
+            x_heads = ep_dispatch(x_heads, self.ep_group, sequence_split_sizes)
+        output = self._local_forward(x_heads) if self.has_real_moe_heads else torch.zeros_like(x_heads)
+        if self.ep_group.world_size > 1:
+            output = ep_undispatch(output, self.ep_group, sequence_split_sizes)
+        if self.ep_pad_heads:
+            output = output[:, : self.num_heads]
+        return output.reshape(-1, self.num_heads * self.d_head)
+
+
+__all__ = [
+    "Magi2MultiHeadMoE",
+    "Magi2MultiHeadMoEConfig",
+    "compute_topk_probs_and_indices",
+    "global_sort_routes",
+    "torch_mh_moe_forward",
+    "triton_mh_moe_forward",
+]

vllm_omni/diffusion/models/magi2/modeling_magi2.py

@@ -0,0 +1,739 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""Native MAGI-2 Preview diffusion transformer.
+
+The architecture and mHC/MoE equations are adapted from SandAI's Apache-2.0
+MAGI-2 Preview implementation.  This modified vLLM-Omni version removes the
+external ``inference`` package, MagiCompiler, and SandAI process manager while
+preserving the released checkpoint module names and tensor layouts.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+from enum import IntEnum
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+
+from .attention import Magi2PackedAttentionKernel, VarlenHandler, apply_rotary_emb
+from .configuration_magi2 import Magi2PreviewConfig
+from .layers import (
+    ElementWiseFourierEmbed,
+    MHCHandler,
+    ModalityDispatcher,
+    MultiModalityRMSNorm,
+    make_grouped_linear,
+    swiglu7,
+)
+from .mh_moe import Magi2MultiHeadMoE, Magi2MultiHeadMoEConfig
+from .parallel import Magi2SequenceDispatcher
+
+
+class Modality(IntEnum):
+    VIDEO = 0
+    AUDIO = 1
+    TEXT = 2
+    TIME = 3
+
+
+class Magi2Attention(nn.Module):
+    def __init__(
+        self,
+        config: Magi2PreviewConfig,
+        *,
+        num_modality: int,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.num_modality = num_modality
+        self.head_dim = config.head_dim
+        global_q_size = config.num_heads_q * self.head_dim
+        global_kv_size = config.num_heads_kv * self.head_dim
+
+        self.pre_norm = MultiModalityRMSNorm(
+            config.hidden_size,
+            num_modality=num_modality,
+        )
+        self.linear_g = make_grouped_linear(
+            config.hidden_size,
+            config.num_heads_q,
+            num_experts=num_modality,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="column",
+        )
+        self.linear_qkv = make_grouped_linear(
+            config.hidden_size,
+            global_q_size + 2 * global_kv_size,
+            num_experts=num_modality,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="column",
+            qkv_splits=(global_q_size, global_kv_size, global_kv_size),
+        )
+        self.linear_proj = make_grouped_linear(
+            global_q_size,
+            config.hidden_size,
+            num_experts=num_modality,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="row",
+        )
+        self.tp_group = self.linear_qkv.tp_group
+        self.num_heads_q = config.num_heads_q // self.tp_group.world_size
+        self.num_heads_kv = config.num_heads_kv // self.tp_group.world_size
+        self.q_size = self.num_heads_q * self.head_dim
+        self.kv_size = self.num_heads_kv * self.head_dim
+        self.sinks = nn.Parameter(torch.empty(config.attention_sink_tokens, self.num_heads_q, dtype=torch.float32))
+        if self.tp_group.world_size > 1:
+            self.sinks.checkpoint_weight_transform = self._shard_sinks
+        self.q_norm = MultiModalityRMSNorm(
+            self.head_dim,
+            num_modality=num_modality,
+            out_dtype=torch.float32,
+        )
+        self.k_norm = MultiModalityRMSNorm(
+            self.head_dim,
+            num_modality=num_modality,
+            out_dtype=torch.float32,
+        )
+        # MAGI-2's packed CFG sequences, learned sink logits, softcap, and
+        # uneven Ulysses splits require a model kernel. Route it through the
+        # framework Attention layer so compile/dispatch ownership remains
+        # shared while the model kernel owns its specialized communication.
+        self.packed_attention = Attention(
+            num_heads=self.num_heads_q,
+            num_kv_heads=self.num_heads_kv,
+            head_size=self.head_dim,
+            causal=False,
+            softmax_scale=self.head_dim**-0.5,
+            qkv_layout="THD",
+            skip_sequence_parallel=True,
+            disable_kv_quant=True,
+            custom_attention=Magi2PackedAttentionKernel(config.attention_softcap),
+        )
+
+    def _shard_sinks(self, checkpoint_tensor: torch.Tensor) -> torch.Tensor:
+        if tuple(checkpoint_tensor.shape) == tuple(self.sinks.shape):
+            return checkpoint_tensor
+        expected = (self.config.attention_sink_tokens, self.config.num_heads_q)
+        if tuple(checkpoint_tensor.shape) != expected:
+            raise ValueError(f"attention sinks have shape {tuple(checkpoint_tensor.shape)}, expected {expected}")
+        start = self.tp_group.rank * self.num_heads_q
+        return checkpoint_tensor[:, start : start + self.num_heads_q]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rope: torch.Tensor,
+        varlen_handler: VarlenHandler,
+        modality_dispatcher: ModalityDispatcher,
+        cp_split_sizes: list[int] | torch.Tensor,
+    ) -> torch.Tensor:
+        normalized = self.pre_norm(hidden_states, modality_dispatcher)
+        gates = self.linear_g(normalized, modality_dispatcher)
+        qkv = self.linear_qkv(normalized, modality_dispatcher)
+        q, k, v = torch.split(qkv, (self.q_size, self.kv_size, self.kv_size), dim=-1)
+        q = q.view(-1, self.num_heads_q, self.head_dim)
+        k = k.view(-1, self.num_heads_kv, self.head_dim)
+        v = v.view(-1, self.num_heads_kv, self.head_dim)
+        gates = gates.view(-1, self.num_heads_q, 1)
+
+        q = self.q_norm(q, modality_dispatcher)
+        k = self.k_norm(k, modality_dispatcher)
+        q = modality_dispatcher.inverse_permute(q).unsqueeze(0)
+        k = modality_dispatcher.inverse_permute(k).unsqueeze(0)
+        v = modality_dispatcher.inverse_permute(v).unsqueeze(0)
+        sin, cos = rope.tensor_split(2, dim=-1)
+        # The released implementation normalizes Q/K in fp32, applies RoPE,
+        # then converts Q/K/V to the checkpoint dtype at the attention
+        # boundary.  This conversion is also required by FlashAttention.
+        q = apply_rotary_emb(q, cos, sin).squeeze(0).to(self.config.params_dtype)
+        k = apply_rotary_emb(k, cos, sin).squeeze(0).to(self.config.params_dtype)
+        v = v.squeeze(0).to(self.config.params_dtype)
+
+        output = self.packed_attention(
+            q,
+            k,
+            v,
+            AttentionMetadata(
+                extra={
+                    "magi2_varlen": varlen_handler,
+                    "magi2_split_sizes": cp_split_sizes,
+                    "magi2_sink": self.sinks,
+                }
+            ),
+        )
+        output = modality_dispatcher.permute(output)
+        output = output * torch.sigmoid(gates)
+        output = output.reshape(-1, self.q_size).to(self.config.params_dtype)
+        return self.linear_proj(output, modality_dispatcher)
+
+
+class Magi2MLP(nn.Module):
+    def __init__(
+        self,
+        config: Magi2PreviewConfig,
+        *,
+        num_modality: int,
+    ) -> None:
+        super().__init__()
+        intermediate_size = int(config.hidden_size * config.intermediate_factor * 2 / 3) // 128 * 128
+        # Tiny unit-test configurations can be narrower than the release's
+        # 128-channel rounding unit while retaining the same formula family.
+        if intermediate_size == 0:
+            intermediate_size = max(1, int(config.hidden_size * config.intermediate_factor * 2 / 3))
+        self.intermediate_size = intermediate_size
+        self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality)
+        self.up_gate_proj = make_grouped_linear(
+            config.hidden_size,
+            2 * intermediate_size,
+            num_experts=num_modality,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="column",
+        )
+        self.down_proj = make_grouped_linear(
+            intermediate_size,
+            config.hidden_size,
+            num_experts=num_modality,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="row",
+        )
+
+    def forward(self, hidden_states: torch.Tensor, dispatcher: ModalityDispatcher) -> torch.Tensor:
+        hidden_states = self.pre_norm(hidden_states, dispatcher)
+        hidden_states = self.up_gate_proj(hidden_states, dispatcher)
+        hidden_states = swiglu7(hidden_states)
+        return self.down_proj(hidden_states, dispatcher)
+
+
+class Magi2MultiHeadMoELayer(nn.Module):
+    def __init__(self, config: Magi2PreviewConfig) -> None:
+        super().__init__()
+        moe = config.moe
+        self.config = config
+        self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=3)
+        self.split_linear = make_grouped_linear(
+            config.hidden_size,
+            config.hidden_size,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="column",
+        )
+        self.moe_mlp = Magi2MultiHeadMoE(
+            Magi2MultiHeadMoEConfig(
+                hidden_size=config.hidden_size,
+                num_heads=moe.num_heads,
+                num_experts=moe.num_experts,
+                top_k=moe.top_k,
+                expert_intermediate_size=moe.expert_intermediate_size,
+                params_dtype=config.params_dtype,
+                score_func=moe.score_function,  # type: ignore[arg-type]
+                route_norm=moe.normalize_routing_weights,
+                route_scale=moe.routing_scale,
+            )
+        )
+        self.merge_linear = make_grouped_linear(
+            config.hidden_size,
+            config.hidden_size,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="row",
+        )
+        self.shared_expert_fc1 = make_grouped_linear(
+            config.hidden_size,
+            2 * moe.shared_expert_intermediate_size,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="column",
+        )
+        self.shared_expert_fc2 = make_grouped_linear(
+            moe.shared_expert_intermediate_size,
+            config.hidden_size,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="row",
+        )
+        self.modality_specific_shared_expert_fc1 = make_grouped_linear(
+            config.hidden_size,
+            2 * moe.modality_shared_expert_intermediate_size,
+            num_experts=3,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="column",
+        )
+        self.modality_specific_shared_expert_fc2 = make_grouped_linear(
+            moe.modality_shared_expert_intermediate_size,
+            config.hidden_size,
+            num_experts=3,
+            bias=False,
+            dtype=config.params_dtype,
+            parallel_mode="row",
+        )
+        tp_size = self.split_linear.tp_group.world_size
+        self.local_shared_expert_intermediate_size = moe.shared_expert_intermediate_size // tp_size
+        self.local_modality_shared_expert_intermediate_size = moe.modality_shared_expert_intermediate_size // tp_size
+
+    def _shared_experts(
+        self,
+        normalized: torch.Tensor,
+        dispatcher: ModalityDispatcher,
+    ) -> torch.Tensor:
+        shared = self.shared_expert_fc1(normalized)
+        modality = self.modality_specific_shared_expert_fc1(normalized, dispatcher)
+        activated = swiglu7(torch.cat((shared, modality), dim=-1))
+        shared, modality = activated.split(
+            (
+                self.local_shared_expert_intermediate_size,
+                self.local_modality_shared_expert_intermediate_size,
+            ),
+            dim=-1,
+        )
+        return self.shared_expert_fc2(shared) + self.modality_specific_shared_expert_fc2(
+            modality.contiguous(), dispatcher
+        )
+
+    def forward(self, hidden_states: torch.Tensor, dispatcher: ModalityDispatcher) -> torch.Tensor:
+        normalized = self.pre_norm(hidden_states, dispatcher)
+        routed = self.split_linear(normalized)
+        routed = self.moe_mlp(routed)
+        routed = self.merge_linear(routed)
+        return routed + self._shared_experts(normalized, dispatcher)
+
+
+class Magi2PreAdapter(nn.Module):
+    def __init__(self, config: Magi2PreviewConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.adapter_dim = config.virtual_width
+        self.video_embedder = nn.Linear(
+            config.video_in_channels,
+            self.adapter_dim,
+            bias=True,
+            dtype=torch.float32,
+        )
+        self.text_embedder = nn.Linear(
+            config.text_in_channels,
+            self.adapter_dim,
+            bias=True,
+            dtype=torch.float32,
+        )
+        self.audio_embedder = nn.Linear(
+            config.audio_in_channels,
+            self.adapter_dim,
+            bias=True,
+            dtype=torch.float32,
+        )
+        self.rope = ElementWiseFourierEmbed(config.head_dim, learnable=False)
+
+    def forward(
+        self,
+        packed: torch.Tensor,
+        coords_mapping: torch.Tensor,
+        video_indices: torch.Tensor,
+        audio_indices: torch.Tensor,
+        text_indices: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        rope = self.rope(coords_mapping)
+        output = torch.zeros(
+            packed.shape[0],
+            self.adapter_dim,
+            dtype=torch.float32,
+            device=packed.device,
+        )
+        if text_indices.numel():
+            output.index_copy_(
+                0,
+                text_indices,
+                self.text_embedder(packed.index_select(0, text_indices)[:, : self.config.text_in_channels].float()),
+            )
+        if audio_indices.numel():
+            output.index_copy_(
+                0,
+                audio_indices,
+                self.audio_embedder(packed.index_select(0, audio_indices)[:, : self.config.audio_in_channels].float()),
+            )
+        if video_indices.numel():
+            output.index_copy_(
+                0,
+                video_indices,
+                self.video_embedder(packed.index_select(0, video_indices)[:, : self.config.video_in_channels].float()),
+            )
+        return output, rope
+
+
+class Magi2PostAdapter(nn.Module):
+    def __init__(self, config: Magi2PreviewConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.adapter_dim = config.virtual_width
+        self.final_norm_video = MultiModalityRMSNorm(self.adapter_dim)
+        self.final_norm_audio = MultiModalityRMSNorm(self.adapter_dim)
+        self.final_linear_video = nn.Linear(
+            self.adapter_dim,
+            config.video_in_channels,
+            bias=False,
+            dtype=torch.float32,
+        )
+        self.final_linear_audio = nn.Linear(
+            self.adapter_dim,
+            config.audio_in_channels,
+            bias=False,
+            dtype=torch.float32,
+        )
+        self.final_out_dim = max(config.video_in_channels, config.audio_in_channels)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        video_indices: torch.Tensor,
+        audio_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        output = torch.zeros(
+            hidden_states.shape[0],
+            self.final_out_dim,
+            dtype=torch.float32,
+            device=hidden_states.device,
+        )
+        if video_indices.numel():
+            video = self.final_norm_video(hidden_states.index_select(0, video_indices).float())
+            video = self.final_linear_video(video).float()
+            output[:, : self.config.video_in_channels].index_copy_(0, video_indices, video)
+        if audio_indices.numel():
+            audio = self.final_norm_audio(hidden_states.index_select(0, audio_indices).float())
+            audio = self.final_linear_audio(audio).float()
+            output[:, : self.config.audio_in_channels].index_copy_(0, audio_indices, audio)
+        return output
+
+
+class Magi2TransformerLayer(nn.Module):
+    def __init__(self, config: Magi2PreviewConfig, layer_index: int) -> None:
+        super().__init__()
+        num_modality = 3 if layer_index in config.multimodal_layers else 1
+        self.config = config
+        self.attention = Magi2Attention(config, num_modality=num_modality)
+        self.mlp: nn.Module
+        if layer_index in config.moe.layers:
+            self.mlp = Magi2MultiHeadMoELayer(config)
+        else:
+            self.mlp = Magi2MLP(config, num_modality=num_modality)
+        self._init_mhc(num_modality)
+
+    def _init_mhc(self, num_modality: int) -> None:
+        streams = self.config.mhc.num_streams
+        hidden = self.config.hidden_size
+        alpha = self.config.mhc.alpha_init
+        for branch in ("attn", "mlp"):
+            for coefficient in ("pre", "post", "res"):
+                self.register_parameter(
+                    f"mhc_alpha_{coefficient}_{branch}",
+                    nn.Parameter(torch.full((1,), alpha, dtype=torch.float32)),
+                )
+            self.register_parameter(
+                f"mhc_bias_pre_{branch}",
+                nn.Parameter(torch.empty(streams, dtype=torch.float32)),
+            )
+            self.register_parameter(
+                f"mhc_bias_post_{branch}",
+                nn.Parameter(torch.empty(streams, dtype=torch.float32)),
+            )
+            self.register_parameter(
+                f"mhc_bias_res_{branch}",
+                nn.Parameter(torch.empty(streams, streams, dtype=torch.float32)),
+            )
+            self.register_parameter(
+                f"mhc_phi_fused_{branch}",
+                nn.Parameter(
+                    torch.empty(
+                        streams * hidden,
+                        streams + streams + streams * streams,
+                        dtype=torch.float32,
+                    )
+                ),
+            )
+        self.mhc_norm = MultiModalityRMSNorm(
+            streams * hidden,
+            num_modality=num_modality,
+            out_dtype=torch.float32,
+        )
+        self.mhc_handler = MHCHandler(
+            streams,
+            hidden,
+            sinkhorn_iterations=self.config.mhc.sinkhorn_iterations,
+            sinkhorn_epsilon=self.config.mhc.sinkhorn_epsilon,
+        )
+
+    def _branch_logits(
+        self,
+        streams: torch.Tensor,
+        branch: str,
+        dispatcher: ModalityDispatcher,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.mhc_handler.compute_logits(
+            self.mhc_handler.flatten(streams),
+            partial(self.mhc_norm, modality_dispatcher=dispatcher),
+            getattr(self, f"mhc_phi_fused_{branch}"),
+        )
+
+    def _branch_input(
+        self,
+        streams: torch.Tensor,
+        branch: str,
+        logits: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        pre_logits, _, _ = logits
+        return self.mhc_handler.apply_pre(
+            streams,
+            (
+                getattr(self, f"mhc_alpha_pre_{branch}"),
+                getattr(self, f"mhc_bias_pre_{branch}"),
+                pre_logits,
+            ),
+            out_dtype=streams.dtype,
+        )
+
+    def _connect(
+        self,
+        streams: torch.Tensor,
+        output: torch.Tensor,
+        branch: str,
+        logits: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        _, post_logits, residual_logits = logits
+        post, residual = self.mhc_handler.compute_post_residual(
+            (
+                getattr(self, f"mhc_alpha_post_{branch}"),
+                getattr(self, f"mhc_bias_post_{branch}"),
+                post_logits,
+            ),
+            (
+                getattr(self, f"mhc_alpha_res_{branch}"),
+                getattr(self, f"mhc_bias_res_{branch}"),
+                residual_logits,
+            ),
+            out_dtype=streams.dtype,
+        )
+        return self.mhc_handler.hyper_connect(streams, output, post, residual)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rope: torch.Tensor,
+        varlen_handler: VarlenHandler,
+        modality_dispatcher: ModalityDispatcher,
+        cp_split_sizes: list[int] | torch.Tensor,
+    ) -> torch.Tensor:
+        streams = hidden_states.reshape(hidden_states.shape[0], self.config.mhc.num_streams, self.config.hidden_size)
+        attention_logits = self._branch_logits(streams, "attn", modality_dispatcher)
+        attention_input = self._branch_input(streams, "attn", attention_logits)
+        attention_output = self.attention(
+            attention_input,
+            rope,
+            varlen_handler,
+            modality_dispatcher,
+            cp_split_sizes,
+        )
+        streams = self._connect(streams, attention_output, "attn", attention_logits)
+        mlp_logits = self._branch_logits(streams, "mlp", modality_dispatcher)
+        mlp_input = self._branch_input(streams, "mlp", mlp_logits)
+        mlp_output = self.mlp(mlp_input, modality_dispatcher)
+        streams = self._connect(streams, mlp_output, "mlp", mlp_logits)
+        return streams.reshape(streams.shape[0], -1)
+
+
+class Magi2TransformerBlock(nn.Module):
+    def __init__(self, config: Magi2PreviewConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList(Magi2TransformerLayer(config, index) for index in range(config.num_layers))
+
+    def __iter__(self):
+        """Expose checkpoint-nested layers to the layerwise offload backends."""
+
+        return iter(self.layers)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rope: torch.Tensor,
+        varlen_handler: VarlenHandler,
+        modality_dispatcher: ModalityDispatcher,
+        cp_split_sizes: list[int] | torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states.to(self.config.params_dtype)
+        hidden_states = modality_dispatcher.permute(hidden_states)
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                rope,
+                varlen_handler,
+                modality_dispatcher,
+                cp_split_sizes,
+            )
+        return modality_dispatcher.inverse_permute(hidden_states)
+
+
+def _is_magi2_transformer_layer(_name: str, module: nn.Module) -> bool:
+    """Shard one Preview layer at a time under the shared FSDP2/HSDP path."""
+
+    return isinstance(module, Magi2TransformerLayer)
+
+
+class Magi2PreviewTransformer(nn.Module):
+    """Native preview DiT with the released checkpoint hierarchy."""
+
+    # ``block`` must remain the registered child name because it is part of the
+    # released checkpoint hierarchy.  The ``layers`` property below exposes its
+    # ModuleList through the standard layerwise-offload contract without
+    # changing the registered module names.
+    _layerwise_offload_blocks_attrs = ["layers"]
+    _hsdp_shard_conditions = [_is_magi2_transformer_layer]
+    _hsdp_preserve_parameter_dtypes = True
+    _EP_SHARDED_SUFFIXES = (
+        ".moe_mlp.gate",
+        ".moe_mlp.W_gate",
+        ".moe_mlp.W_up",
+        ".moe_mlp.W_down",
+        ".moe_mlp.router.expert_bias",
+        ".moe_mlp.router.expert_bias_ema",
+    )
+
+    def __init__(self, config: Magi2PreviewConfig | None = None) -> None:
+        super().__init__()
+        self.config = config or Magi2PreviewConfig()
+        self.config.validate()
+        self.pre_adapter = Magi2PreAdapter(self.config)
+        self.post_adapter = Magi2PostAdapter(self.config)
+        self.block = Magi2TransformerBlock(self.config)
+        # SP doubles as MAGI's MoE-head parallel axis. These modules therefore
+        # contain different checkpoint slices on each SP rank and must remain
+        # rank-local while HSDP shards the replicated parameters around them.
+        self._hsdp_ignored_modules = [
+            f"block.layers.{index}.mlp.moe_mlp"
+            for index, layer in enumerate(self.block.layers)
+            if isinstance(layer.mlp, Magi2MultiHeadMoELayer)
+        ]
+
+    @property
+    def layers(self) -> nn.ModuleList:
+        """Expose Preview layers to the shared offload block discovery API."""
+
+        return self.block.layers
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        coords_mapping: torch.Tensor,
+        modality_mapping: torch.Tensor,
+        varlen_handler: VarlenHandler,
+        time_token_sequence: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        dispatcher = Magi2SequenceDispatcher()
+        x = dispatcher.dispatch(x)
+        coords_mapping = dispatcher.dispatch(coords_mapping)
+        modality_mapping = dispatcher.dispatch(modality_mapping)
+        if time_token_sequence is not None:
+            time_token_sequence = dispatcher.dispatch(time_token_sequence)
+        assert dispatcher.split_sizes is not None
+        cp_split_sizes = dispatcher.split_sizes
+
+        time_mask = modality_mapping == int(Modality.TIME)
+        modality_mapping = torch.where(time_mask, int(Modality.TEXT), modality_mapping)
+        modality_dispatcher = ModalityDispatcher(modality_mapping, 3)
+        video_indices = torch.nonzero(modality_mapping == int(Modality.VIDEO)).flatten()
+        audio_indices = torch.nonzero(modality_mapping == int(Modality.AUDIO)).flatten()
+        text_indices = torch.nonzero(modality_mapping == int(Modality.TEXT)).flatten()
+
+        hidden_states, rope = self.pre_adapter(
+            x,
+            coords_mapping,
+            video_indices,
+            audio_indices,
+            text_indices,
+        )
+        if time_token_sequence is not None and time_token_sequence.shape[-1] > 0:
+            hidden_states[:, : time_token_sequence.shape[-1]] = time_token_sequence.to(hidden_states.dtype)
+        hidden_states = self.block(
+            hidden_states,
+            rope,
+            varlen_handler,
+            modality_dispatcher,
+            cp_split_sizes,
+        )
+        output = self.post_adapter(hidden_states, video_indices, audio_indices)
+        return dispatcher.undispatch(output)
+
+    def _moe_for_weight(self, name: str) -> Magi2MultiHeadMoE | None:
+        if ".moe_mlp." not in name:
+            return None
+        module_name = name.split(".moe_mlp.", 1)[0] + ".moe_mlp"
+        module = dict(self.named_modules()).get(module_name)
+        return module if isinstance(module, Magi2MultiHeadMoE) else None
+
+    def validate_loaded_weights(self, loaded_names: set[str]) -> None:
+        """Fail closed when the mmap loader misses a Preview tensor."""
+
+        loaded_names = {name.removeprefix("transformer.") for name in loaded_names}
+        loaded_names = {f"block.{name}" if name.startswith("layers.") else name for name in loaded_names}
+        expected = {name for name, _ in self.named_parameters()}
+        missing = expected - loaded_names
+        if missing:
+            raise ValueError(f"MAGI-2 Preview mmap loading is missing {len(missing)} weights: {sorted(missing)[:8]}")
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Strictly load and TP/MoE-slice the released Preview checkpoint.
+
+        This is the ordinary loader used by resident and rank-local DLO
+        deployments.  DLO+AllGather uses the pipeline's mmap mapping and the
+        same per-parameter transforms before orthogonal DP sharding.
+        """
+
+        targets = self.state_dict(keep_vars=True)
+        loaded: set[str] = set()
+        for raw_name, checkpoint_tensor in weights:
+            name = raw_name.removeprefix("transformer.")
+            if name not in targets:
+                raise ValueError(f"unexpected MAGI-2 preview checkpoint weight {raw_name!r}")
+            target = targets[name]
+            checkpoint_transform = getattr(target, "checkpoint_weight_transform", None)
+            moe = self._moe_for_weight(name)
+            if callable(checkpoint_transform):
+                checkpoint_tensor = checkpoint_transform(checkpoint_tensor)
+            elif moe is not None and name.endswith(self._EP_SHARDED_SUFFIXES):
+                checkpoint_tensor = moe.ep_slice(checkpoint_tensor)
+            if tuple(target.shape) != tuple(checkpoint_tensor.shape):
+                raise ValueError(
+                    f"MAGI-2 weight {name!r} has shape {tuple(checkpoint_tensor.shape)}, expected {tuple(target.shape)}"
+                )
+            with torch.no_grad():
+                target.copy_(checkpoint_tensor.to(device=target.device, dtype=target.dtype))
+            loaded.add(name)
+
+        if (os.environ.get("MAGI2_ROUTER_BIAS_SOURCE") or "ema").strip().lower() != "main":
+            for module_name, module in self.named_modules():
+                if not isinstance(module, Magi2MultiHeadMoE):
+                    continue
+                module.router.expert_bias.copy_(module.router.expert_bias_ema)
+                loaded.add(f"{module_name}.router.expert_bias")
+
+        missing = set(targets) - loaded
+        if missing:
+            raise ValueError(f"MAGI-2 preview checkpoint is missing {len(missing)} weights: {sorted(missing)[:8]}")
+        return loaded
+
+
+__all__ = [
+    "Magi2PreviewTransformer",
+    "Modality",
+    "VarlenHandler",
+]

vllm_omni/diffusion/models/magi2/parallel.py

@@ -0,0 +1,392 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2025-2026 SandAI. All Rights Reserved.
+
+"""Distributed primitives for the native MAGI-2 preview transformer.
+
+This file is adapted from SandAI's Apache-2.0 MAGI-2 preview context- and
+expert-parallel primitives.  It has been modified to reuse vLLM-Omni's
+already-initialized sequence-parallel and expert-parallel process groups.
+
+MAGI-2 supports two equivalent layouts for its two parallel views:
+
+* Ulysses context parallelism shards tokens and gathers attention heads.
+* Multi-head expert parallelism shards the *MoE head* axis, not experts.
+
+With SP-only deployment both views use the SP group.  With TP enabled,
+Ulysses keeps using SP while native column/row tensor parallelism and the MoE
+heads use TP.  This gives TP4 and TP2SP2 independent weight/head and token
+axes while preserving the released equations and checkpoint hierarchy.
+
+These helpers intentionally do not create or own process groups.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+
+@dataclass(frozen=True)
+class Magi2ParallelGroup:
+    """The small process-group surface needed by MAGI-2 kernels."""
+
+    group: dist.ProcessGroup | None
+    world_size: int
+    rank: int
+    replicated_sequence: bool = False
+
+
+def _dist_group_info(group: dist.ProcessGroup | None) -> Magi2ParallelGroup:
+    if group is None or not dist.is_available() or not dist.is_initialized():
+        return Magi2ParallelGroup(None, 1, 0)
+    return Magi2ParallelGroup(
+        group=group,
+        world_size=dist.get_world_size(group),
+        rank=dist.get_rank(group),
+    )
+
+
+def get_magi2_ulysses_group() -> Magi2ParallelGroup:
+    """Return vLLM-Omni's Ulysses subgroup, or a rank-local fallback.
+
+    MAGI-2 does not support Ring or AllGather-KV attention.  The caller is
+    expected to validate that ``ulysses_degree == sequence_parallel_size``.
+    """
+
+    if not dist.is_available() or not dist.is_initialized():
+        return Magi2ParallelGroup(None, 1, 0)
+    try:
+        from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+
+        coordinator = get_sp_group()
+    except (AssertionError, RuntimeError):
+        return Magi2ParallelGroup(None, 1, 0)
+    return Magi2ParallelGroup(
+        group=coordinator.ulysses_group,
+        world_size=coordinator.ulysses_world_size,
+        rank=coordinator.ulysses_rank,
+    )
+
+
+def get_magi2_ep_group() -> Magi2ParallelGroup:
+    """Return the process group used for MAGI's MoE-head parallelism.
+
+    TP is the explicit MoE-head axis when it is larger than one.  Otherwise
+    the released SP-only layout overlaps head parallelism with Ulysses.  Both
+    groups are initialized and owned by vLLM-Omni; MAGI creates no ad-hoc
+    process groups.
+    """
+
+    if not dist.is_available() or not dist.is_initialized():
+        return Magi2ParallelGroup(None, 1, 0)
+    try:
+        from vllm.distributed.parallel_state import get_tp_group
+
+        coordinator = get_tp_group()
+    except (AssertionError, RuntimeError):
+        return get_magi2_ulysses_group()
+    if coordinator.world_size <= 1:
+        return get_magi2_ulysses_group()
+    return Magi2ParallelGroup(
+        group=coordinator.device_group,
+        world_size=coordinator.world_size,
+        rank=coordinator.rank_in_group,
+        replicated_sequence=True,
+    )
+
+
+def get_magi2_tp_group() -> Magi2ParallelGroup:
+    """Return vLLM's tensor-parallel group without an SP fallback."""
+
+    if not dist.is_available() or not dist.is_initialized():
+        return Magi2ParallelGroup(None, 1, 0, replicated_sequence=True)
+    try:
+        from vllm.distributed.parallel_state import get_tp_group
+
+        coordinator = get_tp_group()
+    except (AssertionError, RuntimeError):
+        return Magi2ParallelGroup(None, 1, 0, replicated_sequence=True)
+    return Magi2ParallelGroup(
+        group=coordinator.device_group,
+        world_size=coordinator.world_size,
+        rank=coordinator.rank_in_group,
+        replicated_sequence=True,
+    )
+
+
+def get_magi2_replica_group(data_parallel_size: int) -> Magi2ParallelGroup:
+    """Return the complete TP x SP group for one data-parallel replica.
+
+    The pipeline uses this group for conditioning broadcasts and output-rank
+    ownership.  With DP=1 the diffusion world is exactly one TP x SP replica.
+    With DP>1, MAGI currently requires TP=1, so the existing SP group is the
+    complete per-replica group. CFG x DP is deliberately rejected by
+    ``_validate_native_topology``; if that combination is enabled in the
+    future, this helper must select a group fixed to one DP replica and include
+    its CFG ranks rather than returning the SP group alone.
+    """
+
+    if not dist.is_available() or not dist.is_initialized():
+        return Magi2ParallelGroup(None, 1, 0)
+    if data_parallel_size > 1:
+        try:
+            from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+
+            coordinator = get_sp_group()
+        except (AssertionError, RuntimeError):
+            return Magi2ParallelGroup(None, 1, 0)
+        return Magi2ParallelGroup(
+            group=coordinator.device_group,
+            world_size=coordinator.world_size,
+            rank=coordinator.rank_in_group,
+        )
+    return _dist_group_info(dist.group.WORLD)
+
+
+def balanced_split_sizes(length: int, world_size: int) -> list[int]:
+    """Split ``length`` contiguously, putting one extra token on early ranks."""
+
+    if length < 0:
+        raise ValueError(f"length must be non-negative, got {length}")
+    if world_size < 1:
+        raise ValueError(f"world_size must be positive, got {world_size}")
+    base, remainder = divmod(length, world_size)
+    return [base + int(rank < remainder) for rank in range(world_size)]
+
+
+def shard_sequence(
+    tensor: torch.Tensor,
+    split_sizes: list[int] | None = None,
+    group: Magi2ParallelGroup | None = None,
+) -> tuple[torch.Tensor, list[int]]:
+    """Return this Ulysses rank's contiguous sequence slice.
+
+    Inputs are replicated at the pipeline boundary, matching the diffusion
+    runner contract, so no collective is needed for dispatch.
+    """
+
+    group = group or get_magi2_ulysses_group()
+    split_sizes = split_sizes or balanced_split_sizes(tensor.shape[0], group.world_size)
+    if len(split_sizes) != group.world_size or sum(split_sizes) != tensor.shape[0]:
+        raise ValueError(
+            f"split_sizes={split_sizes} do not partition sequence length "
+            f"{tensor.shape[0]} across {group.world_size} ranks"
+        )
+    start = sum(split_sizes[: group.rank])
+    return tensor.narrow(0, start, split_sizes[group.rank]).contiguous(), split_sizes
+
+
+def gather_sequence(
+    tensor: torch.Tensor,
+    split_sizes: list[int],
+    group: Magi2ParallelGroup | None = None,
+) -> torch.Tensor:
+    """Gather uneven contiguous sequence shards in rank order."""
+
+    group = group or get_magi2_ulysses_group()
+    if group.world_size == 1:
+        return tensor
+    if len(split_sizes) != group.world_size:
+        raise ValueError("split_sizes length must equal the Ulysses world size")
+    if tensor.shape[0] != split_sizes[group.rank]:
+        raise ValueError(f"rank {group.rank} owns {tensor.shape[0]} tokens, expected {split_sizes[group.rank]}")
+    output_shape = (sum(split_sizes), *tensor.shape[1:])
+    output = torch.empty(output_shape, dtype=tensor.dtype, device=tensor.device)
+    chunks = list(torch.split(output, split_sizes, dim=0))
+    dist.all_gather(chunks, tensor.contiguous(), group=group.group)
+    return output
+
+
+def scatter_seqlen_gather_heads(
+    tensor: torch.Tensor,
+    split_sizes: list[int],
+    group: Magi2ParallelGroup | None = None,
+) -> torch.Tensor:
+    """Ulysses ``[sum(S_r), H, D] -> [S_rank, world*H, D]`` exchange."""
+
+    group = group or get_magi2_ulysses_group()
+    if group.world_size == 1:
+        return tensor
+    if tensor.ndim != 3 or not tensor.is_contiguous():
+        raise ValueError("Ulysses attention input must be contiguous [S,H,D]")
+    if len(split_sizes) != group.world_size or sum(split_sizes) != tensor.shape[0]:
+        raise ValueError("split_sizes must partition the input sequence")
+
+    local_tokens = split_sizes[group.rank]
+    output = torch.empty(
+        (group.world_size * local_tokens, tensor.shape[1], tensor.shape[2]),
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    dist.all_to_all_single(
+        output,
+        tensor,
+        output_split_sizes=[local_tokens] * group.world_size,
+        input_split_sizes=split_sizes,
+        group=group.group,
+    )
+    output = (
+        output.view(group.world_size, local_tokens, tensor.shape[1], tensor.shape[2])
+        .permute(1, 0, 2, 3)
+        .reshape(local_tokens, group.world_size * tensor.shape[1], tensor.shape[2])
+    )
+    return output
+
+
+def scatter_heads_gather_seqlen(
+    tensors: Iterable[torch.Tensor],
+    split_sizes: list[int],
+    group: Magi2ParallelGroup | None = None,
+) -> list[torch.Tensor]:
+    """Inverse batched Ulysses exchange for Q/K/V.
+
+    Each input is ``[S_rank, world*H_i, D]`` and each output is
+    ``[sum(S_r), H_i, D]``.  Fusing Q/K/V into one all-to-all preserves the
+    reference communication ordering and avoids three independent collectives.
+    """
+
+    group = group or get_magi2_ulysses_group()
+    tensors = list(tensors)
+    if group.world_size == 1:
+        return tensors
+    if not tensors:
+        return []
+    local_tokens = split_sizes[group.rank]
+    if len(split_sizes) != group.world_size:
+        raise ValueError("split_sizes length must equal the Ulysses world size")
+    if any(t.ndim != 3 or t.shape[0] != local_tokens for t in tensors):
+        raise ValueError("all Ulysses inputs must be [local_tokens, heads, dim]")
+
+    reshaped: list[torch.Tensor] = []
+    local_head_counts: list[int] = []
+    head_dim = tensors[0].shape[-1]
+    for tensor in tensors:
+        if tensor.shape[-1] != head_dim or tensor.shape[1] % group.world_size:
+            raise ValueError("attention heads must divide evenly across Ulysses ranks")
+        local_heads = tensor.shape[1] // group.world_size
+        local_head_counts.append(local_heads)
+        reshaped.append(
+            tensor.view(local_tokens, group.world_size, local_heads, head_dim)
+            .permute(1, 0, 2, 3)
+            .reshape(group.world_size * local_tokens, local_heads, head_dim)
+        )
+
+    fused = torch.cat(reshaped, dim=1).contiguous()
+    output = torch.empty(
+        (sum(split_sizes), fused.shape[1], head_dim),
+        dtype=fused.dtype,
+        device=fused.device,
+    )
+    dist.all_to_all_single(
+        output,
+        fused,
+        output_split_sizes=split_sizes,
+        input_split_sizes=[local_tokens] * group.world_size,
+        group=group.group,
+    )
+    return list(torch.split(output, local_head_counts, dim=1))
+
+
+def ep_dispatch(
+    tensor: torch.Tensor,
+    group: Magi2ParallelGroup | None = None,
+    sequence_split_sizes: list[int] | None = None,
+) -> torch.Tensor:
+    """Dispatch ``[S,H,D]`` so each rank evaluates a contiguous head shard."""
+
+    group = group or get_magi2_ep_group()
+    if group.world_size == 1:
+        return tensor
+    if tensor.ndim != 3 or tensor.shape[1] % group.world_size:
+        raise ValueError(
+            f"MoE head count {tensor.shape[1] if tensor.ndim >= 2 else '?'} must divide by EP size {group.world_size}"
+        )
+    sequence, heads, dim = tensor.shape
+    local_heads = heads // group.world_size
+    if sequence_split_sizes is None:
+        local_size = torch.tensor([sequence], dtype=torch.int64, device=tensor.device)
+        gathered_sizes = [torch.empty_like(local_size) for _ in range(group.world_size)]
+        dist.all_gather(gathered_sizes, local_size, group=group.group)
+        sequence_split_sizes = [int(size.item()) for size in gathered_sizes]
+    if len(sequence_split_sizes) != group.world_size or sequence_split_sizes[group.rank] != sequence:
+        raise ValueError("EP sequence split sizes do not describe the local tensor")
+    send = tensor.contiguous().view(sequence, group.world_size, local_heads, dim).permute(1, 0, 2, 3).contiguous()
+    output = torch.empty(
+        (sum(sequence_split_sizes), local_heads, dim),
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    row_width = local_heads * dim
+    dist.all_to_all_single(
+        output.view(-1),
+        send.view(-1),
+        output_split_sizes=[size * row_width for size in sequence_split_sizes],
+        input_split_sizes=[sequence * row_width] * group.world_size,
+        group=group.group,
+    )
+    return output
+
+
+def ep_undispatch(
+    tensor: torch.Tensor,
+    group: Magi2ParallelGroup | None = None,
+    sequence_split_sizes: list[int] | None = None,
+) -> torch.Tensor:
+    """Undo :func:`ep_dispatch` and restore the complete MoE head axis."""
+
+    group = group or get_magi2_ep_group()
+    if group.world_size == 1:
+        return tensor
+    if tensor.ndim != 3:
+        raise ValueError("EP-dispatched tensor must be [global_tokens,local_heads,dim]")
+    if sequence_split_sizes is None:
+        if tensor.shape[0] % group.world_size:
+            raise ValueError("uneven EP sequence requires explicit split sizes")
+        sequence_split_sizes = [tensor.shape[0] // group.world_size] * group.world_size
+    if len(sequence_split_sizes) != group.world_size or sum(sequence_split_sizes) != tensor.shape[0]:
+        raise ValueError("EP sequence split sizes do not partition the global tensor")
+    sequence = sequence_split_sizes[group.rank]
+    local_heads, dim = tensor.shape[1:]
+    send = tensor.contiguous()
+    output = torch.empty(
+        (group.world_size, sequence, local_heads, dim),
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    row_width = local_heads * dim
+    dist.all_to_all_single(
+        output.view(-1),
+        send.view(-1),
+        output_split_sizes=[sequence * row_width] * group.world_size,
+        input_split_sizes=[size * row_width for size in sequence_split_sizes],
+        group=group.group,
+    )
+    return output.permute(1, 0, 2, 3).contiguous().view(sequence, group.world_size * local_heads, dim)
+
+
+class Magi2SequenceDispatcher:
+    """Request-scoped dispatcher that enforces one consistent token split."""
+
+    def __init__(self, group: Magi2ParallelGroup | None = None) -> None:
+        self.group = group or get_magi2_ulysses_group()
+        self.split_sizes: list[int] | None = None
+
+    def dispatch(self, tensor: torch.Tensor) -> torch.Tensor:
+        split_sizes = balanced_split_sizes(tensor.shape[0], self.group.world_size)
+        if self.split_sizes is not None and split_sizes != self.split_sizes:
+            raise ValueError(
+                f"all packed inputs must share a sequence length; got {split_sizes}, expected {self.split_sizes}"
+            )
+        self.split_sizes = split_sizes
+        return shard_sequence(tensor, split_sizes, self.group)[0]
+
+    def undispatch(self, tensor: torch.Tensor) -> torch.Tensor:
+        if self.split_sizes is None:
+            raise RuntimeError("dispatch must be called before undispatch")
+        output = gather_sequence(tensor, self.split_sizes, self.group)
+        self.split_sizes = None
+        return output

vllm_omni/diffusion/models/magi2/preview_data_proxy.py

@@ -0,0 +1,691 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""Native packing helpers for the MAGI-2 Preview transformer.
+
+The released model consumes one varlen sequence per sample.  Tokens are laid
+out in VIDEO -> AUDIO -> TEXT order, followed by zero or more reference-image
+special-token/image-token pairs.  This module keeps that layout and its 9-D
+coordinate metadata independent of any particular distributed topology.
+
+The packing math is adapted from SandAI's Apache-2.0 MAGI-2 Preview inference
+implementation.  It intentionally has no dependency on that implementation at
+runtime.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import IntEnum
+from itertools import chain
+from typing import Literal
+
+import torch
+from einops import rearrange
+from torch.nn import functional as F
+
+from .attention import VarlenHandler
+
+
+class Modality(IntEnum):
+    """Checkpoint modality IDs used by the pre/post adapters."""
+
+    VIDEO = 0
+    AUDIO = 1
+    TEXT = 2
+    # TIME is remapped to TEXT by the transformer after selecting time tokens.
+    TIME = 3
+
+
+@dataclass(frozen=True)
+class Magi2PreviewDataProxyConfig:
+    """Packing configuration for the released preview checkpoint."""
+
+    t_patch_size: int = 1
+    patch_size: int = 1
+    spatial_rope_interpolation: Literal["inter", "extra"] = "extra"
+    add_time_token: bool = False
+    time_channel_dim: int = 64
+    time_aligned_rope: bool = False
+    audio_latent_fps: float = 25.0
+    time_pos_fps: float = 3.125
+    vae_first_latent_is_image: bool = True
+    video_fps: float = 25.0
+
+    def __post_init__(self) -> None:
+        if self.t_patch_size <= 0 or self.patch_size <= 0:
+            raise ValueError("MAGI-2 patch sizes must be positive")
+        if self.time_channel_dim < 0:
+            raise ValueError("time_channel_dim must be non-negative")
+
+
+@dataclass
+class ModelInput:
+    """Unpacked tensors for one native preview-transformer invocation."""
+
+    x_t: torch.Tensor
+    audio_x_t: torch.Tensor
+    audio_feat_len: torch.Tensor | Sequence[int]
+    txt_feat: torch.Tensor
+    txt_feat_len: torch.Tensor | Sequence[int]
+    t: torch.Tensor
+    ref_audio_feat: torch.Tensor | None = None
+    ref_audio_feat_len: torch.Tensor | Sequence[int] | None = None
+    ref_video_feat: torch.Tensor | None = None
+    ref_video_feat_len: torch.Tensor | Sequence[int] | None = None
+    per_token_video_t: torch.Tensor | None = None
+    per_token_audio_t: torch.Tensor | None = None
+    # Images are [B, M, C, 1, H, W], lengths are [B, M, 2] holding
+    # the latent patch-grid (H, W), and special tokens are [B, M, D].
+    ref_image_feat: torch.Tensor | None = None
+    ref_image_feat_len: torch.Tensor | None = None
+    ref_image_special_token_embedding: torch.Tensor | None = None
+
+
+def _to_int(value: int | torch.Tensor) -> int:
+    if isinstance(value, torch.Tensor):
+        return int(value.detach().reshape(-1)[0].item())
+    return int(value)
+
+
+def _length_at(value: torch.Tensor | Sequence[int], index: int) -> int:
+    if isinstance(value, torch.Tensor) and value.ndim == 0:
+        if index:
+            raise IndexError("a scalar length only describes one sample")
+        return _to_int(value)
+    return _to_int(value[index])
+
+
+def _pad_cat(tensors: list[torch.Tensor], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    if not tensors:
+        return torch.empty(0, 0, device=device, dtype=dtype)
+    max_channel = max(tensor.shape[-1] for tensor in tensors)
+    return torch.cat(
+        [F.pad(tensor.to(device=device, dtype=dtype), (0, max_channel - tensor.shape[-1])) for tensor in tensors],
+        dim=0,
+    )
+
+
+def _segment_paint(
+    values: list[int],
+    seqlens: list[int],
+    *,
+    dtype: torch.dtype,
+    device: torch.device,
+    output_size: int,
+) -> torch.Tensor:
+    mapping = torch.empty(output_size, dtype=dtype, device=device)
+    offset = 0
+    for value, seqlen in zip(values, seqlens, strict=True):
+        mapping[offset : offset + seqlen] = int(value)
+        offset += seqlen
+    if offset != output_size:
+        raise ValueError(f"segment lengths total {offset}, expected {output_size}")
+    return mapping
+
+
+def _seqlens_to_cu_seqlens(seqlens: list[int], device: torch.device | None = None) -> torch.Tensor:
+    lengths = torch.tensor(seqlens, dtype=torch.int32, device=device)
+    return F.pad(torch.cumsum(lengths, dim=0), (1, 0))
+
+
+def _len_to_list(value: torch.Tensor) -> list[int]:
+    return [int(item) for item in value.detach().to(torch.long).reshape(-1).tolist()]
+
+
+def get_coords(
+    shape: tuple[int, int, int],
+    ref_feat_shape: tuple[int, int, int],
+    offset_thw: tuple[int, int, int] = (0, 0, 0),
+    *,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.float32,
+    time_positions: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Build MAGI-2's ``(t,h,w,T,H,W,ref_T,ref_H,ref_W)`` rows."""
+
+    ori_t, ori_h, ori_w = shape
+    ref_t, ref_h, ref_w = ref_feat_shape
+    offset_t, offset_h, offset_w = offset_thw
+
+    if time_positions is None:
+        time_range = torch.arange(ori_t, device=device, dtype=dtype) + offset_t
+    else:
+        if time_positions.numel() != ori_t:
+            raise ValueError(f"received {time_positions.numel()} time positions for {ori_t} time steps")
+        time_range = time_positions.to(device=device, dtype=dtype) + offset_t
+    height_range = torch.arange(ori_h, device=device, dtype=dtype) + offset_h
+    width_range = torch.arange(ori_w, device=device, dtype=dtype) + offset_w
+    time_grid, height_grid, width_grid = torch.meshgrid(
+        time_range,
+        height_range,
+        width_range,
+        indexing="ij",
+    )
+    coords = torch.stack((time_grid, height_grid, width_grid), dim=-1).reshape(-1, 3)
+    metadata = torch.tensor(
+        (ori_t, ori_h, ori_w, ref_t, ref_h, ref_w),
+        device=device,
+        dtype=dtype,
+    )
+    return torch.cat((coords, metadata.expand(coords.shape[0], -1)), dim=-1)
+
+
+def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor:
+    """The preview checkpoint's FP32 diffusion-time channel embedding."""
+
+    position = position.to(torch.float32) * 1000.0
+    half = dim // 2
+    frequencies = torch.exp(
+        -math.log(10000) * torch.arange(0, half, dtype=torch.float32, device=position.device) / half
+    )
+    arguments = position[:, None].float() * frequencies[None]
+    embedding = torch.cat((torch.cos(arguments), torch.sin(arguments)), dim=-1)
+    if dim % 2:
+        embedding = torch.cat((embedding, torch.zeros_like(embedding[:, :1])), dim=-1)
+    return embedding
+
+
+def _patchify_3d(x: torch.Tensor, t_patch_size: int, patch_size: int) -> torch.Tensor:
+    """Non-overlapping 3-D unfold with the released channel-major order."""
+
+    kernel = (t_patch_size, patch_size, patch_size)
+    if x.ndim != 5:
+        raise ValueError(f"expected [B,C,T,H,W], got {tuple(x.shape)}")
+    if not all(size >= kernel_size for size, kernel_size in zip(x.shape[2:], kernel, strict=True)):
+        return torch.empty(
+            x.shape[0],
+            0,
+            x.shape[1] * math.prod(kernel),
+            device=x.device,
+            dtype=x.dtype,
+        )
+
+    patches = x.unfold(2, t_patch_size, t_patch_size)
+    patches = patches.unfold(3, patch_size, patch_size)
+    patches = patches.unfold(4, patch_size, patch_size)
+    # [B,C,T,H,W,pT,pH,pW] -> [B,T*H*W,C*pT*pH*pW]
+    patches = patches.permute(0, 2, 3, 4, 1, 5, 6, 7).contiguous()
+    return patches.reshape(x.shape[0], -1, x.shape[1] * math.prod(kernel))
+
+
+@dataclass
+class SingleData:
+    """Packed metadata for one sample."""
+
+    video_x_t: torch.Tensor
+    audio_x_t: torch.Tensor
+    audio_feat_len: int
+    txt_feat: torch.Tensor
+    txt_feat_len: int
+    t: int
+    h: int
+    w: int
+    patch_size: int
+    t_patch_size: int
+    spatial_rope_interpolation: Literal["inter", "extra"]
+    diffusion_t: torch.Tensor | None = None
+    per_token_video_t: torch.Tensor | None = None
+    per_token_audio_t: torch.Tensor | None = None
+    time_channel_dim: int = 0
+    vae_first_latent_is_image: bool = True
+    video_fps: float = 25.0
+    time_pos_fps: float = 3.125
+    ref_image_feats: list[torch.Tensor] | None = None
+    ref_image_feat_lens: list[list[int]] | None = None
+    ref_image_special_tokens: list[torch.Tensor] | None = None
+
+    def __post_init__(self) -> None:
+        self.video_token_num = self.video_x_t.shape[0]
+        self.origin_audio_feat_len = self.audio_x_t.shape[0]
+        self.audio_x_t = self.audio_x_t[: self.audio_feat_len]
+        self.txt_feat = self.txt_feat[: self.txt_feat_len]
+        if self.per_token_audio_t is not None:
+            self.per_token_audio_t = self.per_token_audio_t[: self.audio_feat_len]
+
+        self.ref_image_feats = self.ref_image_feats or []
+        self.ref_image_feat_lens = self.ref_image_feat_lens or []
+        self.ref_image_special_tokens = self.ref_image_special_tokens or []
+        if not (len(self.ref_image_feats) == len(self.ref_image_feat_lens) == len(self.ref_image_special_tokens)):
+            raise ValueError("reference-image features, lengths, and special tokens must have equal counts")
+        self.ref_image_token_nums = [math.prod(feat_len) for feat_len in self.ref_image_feat_lens]
+        self.ref_image_feats = [
+            feature[:token_num]
+            for feature, token_num in zip(self.ref_image_feats, self.ref_image_token_nums, strict=True)
+        ]
+        self.num_ref_images = len(self.ref_image_feats)
+        self.total_ref_image_feat_len = sum(self.ref_image_token_nums)
+        self.video_channel = self.video_x_t.shape[-1]
+        self.audio_channel = self.audio_x_t.shape[-1]
+
+    @property
+    def device(self) -> torch.device:
+        return self.video_x_t.device
+
+    @property
+    def default_dtype(self) -> torch.dtype:
+        return self.video_x_t.dtype
+
+    @property
+    def add_time_token(self) -> bool:
+        return self.diffusion_t is not None
+
+    @property
+    def total_token_num(self) -> int:
+        total = self.video_token_num + self.audio_feat_len + self.txt_feat_len
+        total += self.total_ref_image_feat_len + self.num_ref_images
+        return total + int(self.add_time_token)
+
+    @property
+    def feat_to_cat(self) -> list[torch.Tensor]:
+        tensors = [self.video_x_t, self.audio_x_t, self.txt_feat]
+        for index in range(self.num_ref_images):
+            tensors.append(self.ref_image_special_tokens[index])
+            tensors.append(self.ref_image_feats[index])
+        if self.add_time_token:
+            tensors.append(self.diffusion_t.to(device=self.device, dtype=self.default_dtype).reshape(1, 1))
+        return tensors
+
+    @property
+    def token_sequence(self) -> torch.Tensor:
+        return _pad_cat(self.feat_to_cat, self.device, self.default_dtype)
+
+    @property
+    def modality_map_seqlens(self) -> tuple[list[int], list[int]]:
+        seqlens = [self.video_token_num, self.audio_feat_len, self.txt_feat_len]
+        modalities = [Modality.VIDEO, Modality.AUDIO, Modality.TEXT]
+        for index in range(self.num_ref_images):
+            seqlens.extend((1, self.ref_image_token_nums[index]))
+            modalities.extend((Modality.TEXT, Modality.VIDEO))
+        if self.add_time_token:
+            seqlens.append(1)
+            modalities.append(Modality.TIME)
+        return seqlens, modalities
+
+    @property
+    def modality_mapping(self) -> torch.Tensor:
+        seqlens, modalities = self.modality_map_seqlens
+        return _segment_paint(
+            modalities,
+            seqlens,
+            dtype=torch.int32,
+            device=self.device,
+            output_size=self.total_token_num,
+        )
+
+    def _default_coords(
+        self,
+        shape: tuple[int, int, int],
+        ref_feat_shape: tuple[int, int, int],
+        offset_thw: tuple[int, int, int] = (0, 0, 0),
+        time_positions: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return get_coords(
+            shape,
+            ref_feat_shape,
+            offset_thw,
+            device=self.device,
+            dtype=self.default_dtype,
+            time_positions=time_positions,
+        )
+
+    @property
+    def coords_to_cat(self) -> list[torch.Tensor]:
+        t_steps = self.t // self.t_patch_size
+        h_steps = self.h // self.patch_size
+        w_steps = self.w // self.patch_size
+        if self.spatial_rope_interpolation == "inter":
+            video_h_ref, video_w_ref = 32, 32
+        else:
+            video_h_ref, video_w_ref = h_steps, w_steps
+
+        video_coords = self._default_coords(
+            (t_steps, h_steps, w_steps),
+            (t_steps, video_h_ref, video_w_ref),
+        )
+        magic_audio_ref_t = (self.audio_feat_len - 1) // 8 + 1
+        audio_coords = self._default_coords(
+            (self.audio_feat_len, 1, 1),
+            (magic_audio_ref_t // self.t_patch_size, 1, 1),
+        )
+        coords = [
+            video_coords,
+            audio_coords,
+            self._default_coords(
+                (self.txt_feat_len, 1, 1),
+                (1, 1, 1),
+                offset_thw=(-self.txt_feat_len, 0, 0),
+            ),
+        ]
+
+        for index in range(self.num_ref_images):
+            token_len = self.ref_image_token_nums[index]
+            if len(self.ref_image_feat_lens[index]) >= 2:
+                image_h, image_w = self.ref_image_feat_lens[index][:2]
+            else:
+                image_h = image_w = math.ceil(math.sqrt(token_len))
+            # Keep image context after the generated clip with a one-step gap.
+            time_offset = t_steps + 2 + index
+            coords.append(
+                torch.tensor(
+                    [[time_offset, -1, -1, 1, image_h, image_w, 1, image_h, image_w]],
+                    device=self.device,
+                    dtype=self.default_dtype,
+                )
+            )
+            coords.append(
+                self._default_coords(
+                    (1, image_h, image_w),
+                    (1, image_h, image_w),
+                    offset_thw=(time_offset, 0, 0),
+                )[:token_len]
+            )
+
+        if self.add_time_token:
+            coords.append(self._default_coords((1, 1, 1), (1, 1, 1))[:1])
+        return coords
+
+    @property
+    def coords_mapping(self) -> torch.Tensor:
+        return torch.cat(self.coords_to_cat, dim=0)
+
+    @property
+    def time_token_sequence(self) -> torch.Tensor:
+        if self.time_channel_dim == 0:
+            return torch.empty(self.total_token_num, 0, device=self.device)
+        if self.per_token_video_t is None or self.per_token_audio_t is None:
+            raise ValueError("per-token video and audio timesteps are required when time_channel_dim is non-zero")
+        parts = [
+            self.per_token_video_t.squeeze(-1),
+            self.per_token_audio_t.squeeze(-1),
+            torch.zeros(self.txt_feat_len, device=self.device),
+        ]
+        for index in range(self.num_ref_images):
+            parts.append(torch.zeros(1, device=self.device))
+            parts.append(torch.zeros(self.ref_image_token_nums[index], device=self.device))
+        if self.add_time_token:
+            parts.append(self.diffusion_t.reshape(1).to(self.device))
+        raw_t = torch.cat(parts, dim=0)
+        if self.time_channel_dim == 1:
+            return raw_t.unsqueeze(-1)
+        return sinusoidal_embedding_1d(self.time_channel_dim, raw_t)
+
+
+@dataclass
+class SimplePackedData:
+    """A concatenation of independently addressable packed samples."""
+
+    items: list[SingleData]
+
+    def __post_init__(self) -> None:
+        if not self.items:
+            raise ValueError("SimplePackedData must contain at least one item")
+        self._total_token_num_list = [item.total_token_num for item in self.items]
+        self._total_token_num_sum = sum(self._total_token_num_list)
+        self._total_token_num_max = max(self._total_token_num_list)
+        self._total_token_cu_seqlens = _seqlens_to_cu_seqlens(self._total_token_num_list)
+
+    @property
+    def device(self) -> torch.device:
+        return self.items[0].device
+
+    @property
+    def default_dtype(self) -> torch.dtype:
+        return self.items[0].default_dtype
+
+    @property
+    def token_sequence(self) -> torch.Tensor:
+        tensors = list(chain.from_iterable(item.feat_to_cat for item in self.items))
+        return _pad_cat(tensors, self.device, self.default_dtype)
+
+    @property
+    def modality_mapping(self) -> torch.Tensor:
+        seqlens: list[int] = []
+        modalities: list[int] = []
+        for item in self.items:
+            item_seqlens, item_modalities = item.modality_map_seqlens
+            seqlens.extend(item_seqlens)
+            modalities.extend(item_modalities)
+        return _segment_paint(
+            modalities,
+            seqlens,
+            dtype=torch.int32,
+            device=self.device,
+            output_size=self.total_token_num,
+        )
+
+    @property
+    def coords_mapping(self) -> torch.Tensor:
+        coords = list(chain.from_iterable(item.coords_to_cat for item in self.items))
+        return torch.cat(coords, dim=0)
+
+    @property
+    def time_token_sequence(self) -> torch.Tensor:
+        return torch.cat([item.time_token_sequence for item in self.items], dim=0)
+
+    @property
+    def total_token_num(self) -> int:
+        return self._total_token_num_sum
+
+    @property
+    def cu_seqlen(self) -> torch.Tensor:
+        return self._total_token_cu_seqlens.clone()
+
+    @property
+    def max_seqlen(self) -> int:
+        return self._total_token_num_max
+
+    def __getitem__(self, index: int) -> SingleData:
+        return self.items[index]
+
+    def depack_token_sequence(self, token_sequence: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        if token_sequence.shape[0] != self.total_token_num:
+            raise ValueError(f"model returned {token_sequence.shape[0]} tokens, expected {self.total_token_num}")
+        videos: list[torch.Tensor] = []
+        audios: list[torch.Tensor] = []
+        token_slices = torch.split(token_sequence, self._total_token_num_list, dim=0)
+        for item, token_slice in zip(self.items, token_slices, strict=True):
+            video_flat = token_slice[: item.video_token_num, : item.video_channel]
+            output_channels = item.video_channel // (item.t_patch_size * item.patch_size * item.patch_size)
+            video = rearrange(
+                video_flat,
+                "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)",
+                T=item.t // item.t_patch_size,
+                H=item.h // item.patch_size,
+                W=item.w // item.patch_size,
+                pT=item.t_patch_size,
+                pH=item.patch_size,
+                pW=item.patch_size,
+                C=output_channels,
+            ).contiguous()
+            audio = torch.zeros(
+                item.origin_audio_feat_len,
+                item.audio_channel,
+                device=token_sequence.device,
+                dtype=token_sequence.dtype,
+            )
+            audio[: item.audio_feat_len] = token_slice[
+                item.video_token_num : item.video_token_num + item.audio_feat_len,
+                : item.audio_channel,
+            ]
+            videos.append(video)
+            audios.append(audio)
+        return torch.stack(videos, dim=0), torch.stack(audios, dim=0)
+
+
+@dataclass(frozen=True)
+class PackedModelInput:
+    """Request-owned packed transformer arguments and output layout."""
+
+    token_sequence: torch.Tensor
+    coords_mapping: torch.Tensor
+    modality_mapping: torch.Tensor
+    varlen_handler: VarlenHandler
+    time_token_sequence: torch.Tensor
+    output_layout: SimplePackedData
+
+    @property
+    def model_args(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, VarlenHandler, torch.Tensor]:
+        return (
+            self.token_sequence,
+            self.coords_mapping,
+            self.modality_mapping,
+            self.varlen_handler,
+            self.time_token_sequence,
+        )
+
+
+class Magi2DataProxy:
+    """Convert dense modality tensors to/from the transformer's packed ABI.
+
+    The proxy is request-stateless: all packing state needed by
+    ``process_output`` is returned as ``PackedModelInput.output_layout`` rather
+    than stored on this shared instance, so concurrent requests cannot
+    overwrite one another.
+    """
+
+    def __init__(self, config: Magi2PreviewDataProxyConfig | None = None) -> None:
+        self.config = config or Magi2PreviewDataProxyConfig()
+        self.patch_size = self.config.patch_size
+        self.t_patch_size = self.config.t_patch_size
+
+    def img2tokens(self, x_t: torch.Tensor) -> torch.Tensor:
+        return _patchify_3d(x_t, self.t_patch_size, self.patch_size)
+
+    @staticmethod
+    def _validate_batch(data: ModelInput) -> None:
+        if data.x_t.ndim != 5:
+            raise ValueError(f"x_t must be [B,C,T,H,W], got {tuple(data.x_t.shape)}")
+        if data.audio_x_t.ndim != 3 or data.txt_feat.ndim != 3:
+            raise ValueError("audio_x_t and txt_feat must be [B,L,C]")
+        batch_size = data.x_t.shape[0]
+        if data.audio_x_t.shape[0] != batch_size or data.txt_feat.shape[0] != batch_size:
+            raise ValueError("video, audio, and text batch sizes must match")
+
+        ref_values = (
+            data.ref_image_feat,
+            data.ref_image_feat_len,
+            data.ref_image_special_token_embedding,
+        )
+        if any(value is not None for value in ref_values) and not all(value is not None for value in ref_values):
+            raise ValueError("reference-image features, lengths, and special tokens must be provided together")
+        if data.ref_image_feat is not None:
+            if data.ref_image_feat.ndim != 6:
+                raise ValueError("ref_image_feat must be [B,M,C,T,H,W]")
+            if data.ref_image_feat_len.ndim != 3 or data.ref_image_feat_len.shape[-1] < 2:
+                raise ValueError("ref_image_feat_len must be [B,M,2]")
+            if data.ref_image_special_token_embedding.ndim != 3:
+                raise ValueError("ref_image_special_token_embedding must be [B,M,D]")
+            expected = data.ref_image_feat.shape[:2]
+            if data.ref_image_feat_len.shape[:2] != expected:
+                raise ValueError("reference-image feature and length batch/image counts differ")
+            if data.ref_image_special_token_embedding.shape[:2] != expected:
+                raise ValueError("reference-image feature and special-token batch/image counts differ")
+
+    def process_input(
+        self,
+        data: ModelInput,
+    ) -> PackedModelInput:
+        self._validate_batch(data)
+        batch_size, _, video_t, video_h, video_w = data.x_t.shape
+        video_tokens = self.img2tokens(data.x_t)
+        audio_tokens = data.audio_x_t.contiguous()
+        text_tokens = data.txt_feat.contiguous()
+
+        per_token_video_tokens = None
+        per_token_audio = None
+        if self.config.time_channel_dim > 0 and data.per_token_video_t is not None:
+            per_token_video_tokens = self.img2tokens(data.per_token_video_t)[:, :, :1]
+            per_token_audio = data.per_token_audio_t
+
+        has_ref_image = data.ref_image_feat is not None
+        num_ref_images = data.ref_image_feat.shape[1] if has_ref_image else 0
+        ref_device, ref_dtype = text_tokens.device, text_tokens.dtype
+
+        items: list[SingleData] = []
+        for batch_index in range(batch_size):
+            ref_image_feats: list[torch.Tensor] = []
+            ref_image_feat_lens: list[list[int]] = []
+            ref_image_special_tokens: list[torch.Tensor] = []
+            for image_index in range(num_ref_images):
+                feature = self.img2tokens(data.ref_image_feat[batch_index, image_index].unsqueeze(0)).squeeze(0)
+                ref_image_feats.append(feature)
+                ref_image_feat_lens.append(_len_to_list(data.ref_image_feat_len[batch_index, image_index]))
+                ref_image_special_tokens.append(
+                    data.ref_image_special_token_embedding[batch_index, image_index]
+                    .to(device=ref_device, dtype=ref_dtype)
+                    .unsqueeze(0)
+                )
+
+            if isinstance(data.t, torch.Tensor) and data.t.ndim == 0:
+                diffusion_t = data.t
+            else:
+                diffusion_t = data.t[batch_index]
+            items.append(
+                SingleData(
+                    video_x_t=video_tokens[batch_index],
+                    audio_x_t=audio_tokens[batch_index],
+                    audio_feat_len=_length_at(data.audio_feat_len, batch_index),
+                    txt_feat=text_tokens[batch_index],
+                    txt_feat_len=_length_at(data.txt_feat_len, batch_index),
+                    t=video_t,
+                    h=video_h,
+                    w=video_w,
+                    patch_size=self.patch_size,
+                    t_patch_size=self.t_patch_size,
+                    spatial_rope_interpolation=self.config.spatial_rope_interpolation,
+                    diffusion_t=diffusion_t if self.config.add_time_token else None,
+                    per_token_video_t=(
+                        per_token_video_tokens[batch_index] if per_token_video_tokens is not None else None
+                    ),
+                    per_token_audio_t=(per_token_audio[batch_index] if per_token_audio is not None else None),
+                    time_channel_dim=self.config.time_channel_dim,
+                    time_pos_fps=self.config.time_pos_fps,
+                    vae_first_latent_is_image=self.config.vae_first_latent_is_image,
+                    video_fps=self.config.video_fps,
+                    ref_image_feats=ref_image_feats,
+                    ref_image_feat_lens=ref_image_feat_lens,
+                    ref_image_special_tokens=ref_image_special_tokens,
+                )
+            )
+
+        packed = SimplePackedData(items)
+        cu_seqlens = packed.cu_seqlen.to(device=data.x_t.device, dtype=torch.int32)
+        varlen_handler = VarlenHandler(
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=packed.max_seqlen,
+            max_seqlen_k=packed.max_seqlen,
+        )
+        return PackedModelInput(
+            token_sequence=packed.token_sequence,
+            coords_mapping=packed.coords_mapping,
+            modality_mapping=packed.modality_mapping,
+            varlen_handler=varlen_handler,
+            time_token_sequence=packed.time_token_sequence,
+            output_layout=packed,
+        )
+
+    @staticmethod
+    def process_output(
+        x: torch.Tensor,
+        output_layout: SimplePackedData,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return output_layout.depack_token_sequence(x)
+
+
+__all__ = [
+    "Magi2DataProxy",
+    "Magi2PreviewDataProxyConfig",
+    "Modality",
+    "ModelInput",
+    "PackedModelInput",
+    "SimplePackedData",
+    "SingleData",
+    "VarlenHandler",
+    "get_coords",
+    "sinusoidal_embedding_1d",
+]

vllm_omni/diffusion/models/magi2/sampler_magi2.py

@@ -0,0 +1,682 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""Native MAGI-2 Preview classifier-free-guidance sampling.
+
+The CFG and denoising math is adapted from SandAI's Apache-2.0 MAGI-2
+Preview inference implementation.  Model placement is deliberately not
+managed here: vLLM-Omni's loader/offloader (including DLO) owns placement,
+while this sampler only prepares tensors and invokes the already placed model.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
+from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
+
+from .preview_data_proxy import Magi2DataProxy, ModelInput
+
+
+def _pad_or_trim(
+    tensor: torch.Tensor,
+    target_size: int,
+    dim: int,
+    pad_value: float = 0.0,
+) -> tuple[torch.Tensor, int]:
+    current_size = tensor.size(dim)
+    if current_size < target_size:
+        padding_amount = target_size - current_size
+        padding = [0] * (2 * tensor.dim())
+        padding_dim_index = tensor.dim() - 1 - dim
+        padding[2 * padding_dim_index + 1] = padding_amount
+        return F.pad(tensor, tuple(padding), "constant", pad_value), current_size
+    slices = [slice(None)] * tensor.dim()
+    slices[dim] = slice(0, target_size)
+    return tensor[tuple(slices)], current_size
+
+
+@dataclass(frozen=True)
+class CFGConfig:
+    """Preview CFG controls, defaulted to the released MAGI-2 recipe."""
+
+    use_cfg_trick: bool = False
+    cfg_trick_start_frame: int = 13
+    cfg_trick_value: float = 2.0
+    use_dynamic_cfg: bool = False
+    dynamic_cfg_start_t: int = 500
+    dynamic_cfg_cutoff_value: float = 2.0
+    video_txt_guidance_scale: float = 5.0
+    audio_txt_guidance_scale: float = 7.0
+    # The released I2V config keeps image context in both CFG branches, so CFG
+    # changes text conditioning without pushing away from the input image.
+    use_ref_for_uncond: bool = True
+    use_skimmed_cfg_linear: bool = False
+    skimmed_cfg_scale: float = 3.0
+    cfg_rescale: float = 0.0
+
+    def __post_init__(self) -> None:
+        finite_values = {
+            "cfg_trick_value": self.cfg_trick_value,
+            "dynamic_cfg_cutoff_value": self.dynamic_cfg_cutoff_value,
+            "video_txt_guidance_scale": self.video_txt_guidance_scale,
+            "audio_txt_guidance_scale": self.audio_txt_guidance_scale,
+            "skimmed_cfg_scale": self.skimmed_cfg_scale,
+            "cfg_rescale": self.cfg_rescale,
+        }
+        for name, value in finite_values.items():
+            if not math.isfinite(float(value)):
+                raise ValueError(f"{name} must be finite, got {value}")
+        if self.cfg_trick_start_frame < 0:
+            raise ValueError("cfg_trick_start_frame must be non-negative")
+        if self.skimmed_cfg_scale < 0 or self.cfg_rescale < 0:
+            raise ValueError("skimmed_cfg_scale and cfg_rescale must be non-negative")
+
+
+@dataclass
+class SamplerInput:
+    video_t_list: Sequence[torch.Tensor]
+    audio_t_list: Sequence[torch.Tensor]
+    latent: torch.Tensor  # [B,C,T,H,W]
+    audio_latent: torch.Tensor  # [B,L,C]
+    txt_feat: torch.Tensor  # [B,L,D]
+    null_txt_feat: torch.Tensor  # [B,L,D]
+    ref_audio_feat: torch.Tensor | None
+    ref_video_feat: torch.Tensor | None
+    video_scheduler: SchedulerMixin
+    audio_scheduler: SchedulerMixin
+    cfg_config: CFGConfig
+    ref_image_feat: torch.Tensor | None = None  # [B,M,C,1,H,W]
+    ref_image_feat_len: torch.Tensor | None = None  # [B,M,2]
+    ref_image_special_token_embedding: torch.Tensor | None = None  # [B,M,D]
+
+
+def build_magi2_preview_schedulers(
+    num_inference_steps: int,
+    *,
+    device: torch.device | str,
+    shift: float = 7.0,
+) -> tuple[FlowUniPCMultistepScheduler, FlowUniPCMultistepScheduler]:
+    """Create independent, identically configured video/audio schedulers."""
+
+    if num_inference_steps <= 0:
+        raise ValueError("num_inference_steps must be positive")
+    if not math.isfinite(shift) or shift <= 0:
+        raise ValueError(f"shift must be positive and finite, got {shift}")
+    video_scheduler = FlowUniPCMultistepScheduler()
+    audio_scheduler = FlowUniPCMultistepScheduler()
+    video_scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
+    audio_scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
+    return video_scheduler, audio_scheduler
+
+
+class Magi2PreviewSampler(CFGParallelMixin):
+    """Run MAGI-2 Preview's joint video/audio denoising loop."""
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        data_proxy: Magi2DataProxy | None = None,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        self.model = model
+        self.data_proxy = data_proxy or Magi2DataProxy()
+        self.device = torch.device(device) if device is not None else None
+        # Retained for the reference-compatible constructor.  Model placement
+        # and parameter dtype are owned by the generic loader/offloader.
+        self.dtype = dtype
+
+    @torch.inference_mode()
+    def forward(self, model_input: ModelInput) -> tuple[torch.Tensor, torch.Tensor]:
+        packed_input = self.data_proxy.process_input(model_input)
+        model_output = self.model(*packed_input.model_args)
+        if not isinstance(model_output, torch.Tensor):
+            raise TypeError(f"MAGI-2 preview transformer must return a tensor, got {type(model_output)!r}")
+        return self.data_proxy.process_output(model_output, packed_input.output_layout)
+
+    def predict_noise(self, model_input: ModelInput) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run one conditional or unconditional branch for shared CFG dispatch."""
+
+        return self.forward(model_input)
+
+    @torch.inference_mode()
+    def sample(self, sampler_input: SamplerInput) -> tuple[torch.Tensor, torch.Tensor]:
+        video_timesteps = list(sampler_input.video_t_list)
+        audio_timesteps = list(sampler_input.audio_t_list)
+        if not video_timesteps:
+            raise ValueError("MAGI-2 sampling needs at least one timestep")
+        if len(video_timesteps) != len(audio_timesteps):
+            raise ValueError("video and audio timestep schedules must have equal lengths")
+
+        video_cfgs, audio_cfgs = self.precalculate_cfg(
+            video_timesteps,
+            sampler_input.latent.shape[2],
+            sampler_input.cfg_config,
+            device=sampler_input.latent.device,
+        )
+        latent = sampler_input.latent.clone()
+        audio_latent = sampler_input.audio_latent.clone()
+
+        for timestep, video_cfg, audio_cfg in zip(
+            video_timesteps,
+            video_cfgs,
+            audio_cfgs,
+            strict=True,
+        ):
+            model_input = self.prepare_model_input(
+                latent=latent,
+                audio_latent=audio_latent,
+                txt_feat=sampler_input.txt_feat,
+                null_txt_feat=sampler_input.null_txt_feat,
+                ref_audio_feat=sampler_input.ref_audio_feat,
+                ref_video_feat=sampler_input.ref_video_feat,
+                ref_image_feat=sampler_input.ref_image_feat,
+                ref_image_feat_len=sampler_input.ref_image_feat_len,
+                ref_image_special_token_embedding=(sampler_input.ref_image_special_token_embedding),
+                t=timestep,
+                cfg_config=sampler_input.cfg_config,
+            )
+            if get_classifier_free_guidance_world_size() > 1:
+                positive_input, negative_input = self._split_cfg_model_input(model_input)
+                guided = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=True,
+                    true_cfg_scale=1.0,
+                    positive_kwargs={"model_input": positive_input},
+                    negative_kwargs={"model_input": negative_input},
+                    cfg_normalize=False,
+                    kwargs={
+                        "video_txt_guidance_scale": video_cfg,
+                        "audio_txt_guidance_scale": audio_cfg,
+                        "cfg_config": sampler_input.cfg_config,
+                        "latent": latent,
+                        "audio_latent": audio_latent,
+                    },
+                )
+                if not isinstance(guided, tuple) or len(guided) != 2:
+                    raise RuntimeError("MAGI-2 CFG parallel combine must return video and audio predictions")
+                latent, audio_latent = self._step_guided(
+                    guided,
+                    latent,
+                    audio_latent,
+                    sampler_input.video_scheduler,
+                    sampler_input.audio_scheduler,
+                    timestep,
+                )
+            else:
+                model_pred = self.forward(model_input)
+                latent, audio_latent, _, _ = self.step(
+                    model_pred,
+                    latent,
+                    audio_latent,
+                    video_cfg,
+                    audio_cfg,
+                    sampler_input.video_scheduler,
+                    sampler_input.audio_scheduler,
+                    timestep,
+                    cfg_config=sampler_input.cfg_config,
+                )
+
+        return latent, audio_latent
+
+    @staticmethod
+    def _validate_prepare_inputs(
+        latent: torch.Tensor,
+        audio_latent: torch.Tensor,
+        txt_feat: torch.Tensor,
+        null_txt_feat: torch.Tensor,
+    ) -> None:
+        if latent.ndim != 5:
+            raise ValueError("latent must be [B,C,T,H,W]")
+        if audio_latent.ndim != 3 or txt_feat.ndim != 3 or null_txt_feat.ndim != 3:
+            raise ValueError("audio and text features must be [B,L,C]")
+        batch_size = latent.shape[0]
+        if not (audio_latent.shape[0] == txt_feat.shape[0] == null_txt_feat.shape[0] == batch_size):
+            raise ValueError("video, audio, positive-text, and negative-text batch sizes must match")
+
+    def prepare_model_input(
+        self,
+        latent: torch.Tensor,
+        audio_latent: torch.Tensor,
+        txt_feat: torch.Tensor,
+        null_txt_feat: torch.Tensor,
+        ref_audio_feat: torch.Tensor | None = None,
+        ref_video_feat: torch.Tensor | None = None,
+        ref_image_feat: torch.Tensor | None = None,
+        ref_image_feat_len: torch.Tensor | None = None,
+        ref_image_special_token_embedding: torch.Tensor | None = None,
+        t: torch.Tensor | float | None = None,
+        cfg_config: CFGConfig | None = None,
+    ) -> ModelInput:
+        self._validate_prepare_inputs(latent, audio_latent, txt_feat, null_txt_feat)
+        batch_size = latent.shape[0]
+        audio_latent_len = audio_latent.shape[1]
+        txt_feat_len = txt_feat.shape[1]
+        null_txt_feat_len = null_txt_feat.shape[1]
+        target_length = max(txt_feat_len, null_txt_feat_len)
+        txt_feat, _ = _pad_or_trim(txt_feat, target_length, 1)
+        null_txt_feat, _ = _pad_or_trim(null_txt_feat, target_length, 1)
+
+        if ref_audio_feat is None:
+            ref_audio_feat = audio_latent.new_empty(batch_size, 0, audio_latent.shape[-1])
+        if ref_audio_feat.ndim != 3 or ref_audio_feat.shape[0] != batch_size:
+            raise ValueError("ref_audio_feat must be [B,L,C] with the same batch size")
+        ref_audio_feat_len = ref_audio_feat.shape[1]
+
+        if ref_video_feat is None:
+            # Preserve the 5-D ABI without allocating or reading a full-size
+            # placeholder that the zero reference length excludes downstream.
+            ref_video_feat = latent[:, :, :0]
+            ref_video_feat_len = 0
+        else:
+            if ref_video_feat.ndim != 5 or ref_video_feat.shape[0] != batch_size:
+                raise ValueError("ref_video_feat must be [B,C,T,H,W] with the same batch size")
+            ref_video_feat_len = (ref_video_feat.shape[3] * ref_video_feat.shape[4]) // 4
+
+        ref_image_feat_cfg, ref_image_feat_len_cfg, ref_image_special_tokens_cfg = self._prepare_ref_image_cfg(
+            ref_image_feat,
+            ref_image_feat_len,
+            ref_image_special_token_embedding,
+            cfg_config,
+        )
+
+        if t is None:
+            t_value = torch.tensor(0.0, device=latent.device, dtype=latent.dtype)
+        elif isinstance(t, torch.Tensor):
+            t_value = t.reshape(-1)[0].to(device=latent.device, dtype=latent.dtype)
+        else:
+            t_value = torch.tensor(float(t), device=latent.device, dtype=latent.dtype)
+        t_normalized = t_value / 1000.0
+        cfg_batch = batch_size * 2
+        t_batch = t_normalized.expand(cfg_batch)
+
+        _, _, video_t, video_h, video_w = latent.shape
+        audio_t = audio_latent.shape[1]
+        # Keep zero-stride views here. Patchification/materialization happens
+        # once in the data proxy, so cloning the full token grids every denoise
+        # step only increases peak memory and allocator pressure.
+        per_token_video_t = t_batch.view(-1, 1, 1, 1, 1).expand(cfg_batch, 1, video_t, video_h, video_w)
+        per_token_audio_t = t_batch.view(-1, 1, 1).expand(cfg_batch, audio_t, 1)
+
+        audio_lengths = torch.full(
+            (cfg_batch,),
+            audio_latent_len,
+            device=latent.device,
+            dtype=torch.long,
+        )
+        text_lengths = torch.tensor(
+            [txt_feat_len] * batch_size + [null_txt_feat_len] * batch_size,
+            device=latent.device,
+            dtype=torch.long,
+        )
+        ref_audio_lengths = torch.full(
+            (cfg_batch,),
+            ref_audio_feat_len,
+            device=latent.device,
+            dtype=torch.long,
+        )
+        ref_video_lengths = torch.full(
+            (cfg_batch,),
+            ref_video_feat_len,
+            device=latent.device,
+            dtype=torch.long,
+        )
+
+        return ModelInput(
+            x_t=torch.cat((latent, latent), dim=0),
+            audio_x_t=torch.cat((audio_latent, audio_latent), dim=0),
+            audio_feat_len=audio_lengths,
+            txt_feat=torch.cat((txt_feat, null_txt_feat), dim=0),
+            txt_feat_len=text_lengths,
+            t=t_batch,
+            per_token_video_t=per_token_video_t,
+            per_token_audio_t=per_token_audio_t,
+            ref_audio_feat=torch.cat((ref_audio_feat, torch.zeros_like(ref_audio_feat)), dim=0),
+            ref_audio_feat_len=ref_audio_lengths,
+            ref_video_feat=torch.cat((ref_video_feat, torch.zeros_like(ref_video_feat)), dim=0),
+            ref_video_feat_len=ref_video_lengths,
+            ref_image_feat=ref_image_feat_cfg,
+            ref_image_feat_len=ref_image_feat_len_cfg,
+            ref_image_special_token_embedding=ref_image_special_tokens_cfg,
+        )
+
+    @staticmethod
+    def _split_cfg_model_input(model_input: ModelInput) -> tuple[ModelInput, ModelInput]:
+        """Split the packed positive/negative batch for CFG rank dispatch."""
+
+        batch = model_input.x_t.shape[0]
+        if batch % 2:
+            raise ValueError("MAGI-2 CFG model input must contain equal positive and negative halves")
+        half = batch // 2
+
+        def split(value: torch.Tensor | Sequence[int] | None) -> tuple[torch.Tensor | Sequence[int] | None, ...]:
+            if value is None:
+                return None, None
+            if isinstance(value, torch.Tensor):
+                if value.ndim == 0 or value.shape[0] != batch:
+                    raise ValueError("MAGI-2 CFG tensor fields must use the packed CFG batch dimension")
+                return value[:half], value[half:]
+            if len(value) != batch:
+                raise ValueError("MAGI-2 CFG sequence fields must use the packed CFG batch dimension")
+            return value[:half], value[half:]
+
+        field_pairs = {
+            name: split(getattr(model_input, name))
+            for name in (
+                "x_t",
+                "audio_x_t",
+                "audio_feat_len",
+                "txt_feat",
+                "txt_feat_len",
+                "t",
+                "ref_audio_feat",
+                "ref_audio_feat_len",
+                "ref_video_feat",
+                "ref_video_feat_len",
+                "per_token_video_t",
+                "per_token_audio_t",
+                "ref_image_feat",
+                "ref_image_feat_len",
+                "ref_image_special_token_embedding",
+            )
+        }
+        branches = []
+        for branch in range(2):
+            branches.append(ModelInput(**{name: pair[branch] for name, pair in field_pairs.items()}))
+        return branches[0], branches[1]
+
+    @staticmethod
+    def _prepare_ref_image_cfg(
+        ref_image_feat: torch.Tensor | None,
+        ref_image_feat_len: torch.Tensor | None,
+        ref_image_special_token_embedding: torch.Tensor | None,
+        cfg_config: CFGConfig | None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
+        if ref_image_feat is None:
+            if ref_image_feat_len is not None or ref_image_special_token_embedding is not None:
+                raise ValueError("reference-image features, lengths, and special tokens must be provided together")
+            return None, None, None
+        if ref_image_feat_len is None or ref_image_special_token_embedding is None:
+            raise ValueError("reference-image features, lengths, and special tokens must be provided together")
+
+        keep_ref = bool(cfg_config is not None and cfg_config.use_ref_for_uncond)
+        uncond_feature = ref_image_feat if keep_ref else torch.zeros_like(ref_image_feat)
+        return (
+            torch.cat((ref_image_feat, uncond_feature), dim=0),
+            torch.cat((ref_image_feat_len, ref_image_feat_len), dim=0),
+            torch.cat(
+                (ref_image_special_token_embedding, ref_image_special_token_embedding),
+                dim=0,
+            ),
+        )
+
+    def precalculate_cfg(
+        self,
+        t_list: Sequence[torch.Tensor],
+        latent_length: int,
+        cfg_config: CFGConfig,
+        *,
+        device: torch.device | str | None = None,
+    ) -> tuple[list[torch.Tensor], list[float]]:
+        if latent_length < 0:
+            raise ValueError("latent_length must be non-negative")
+        if device is None:
+            if self.device is not None:
+                device = self.device
+            elif t_list:
+                device = t_list[0].device
+            else:
+                device = torch.device("cpu")
+
+        if cfg_config.use_cfg_trick:
+            video_cfg_base = (
+                torch.tensor(cfg_config.video_txt_guidance_scale, device=device)
+                .expand(1, 1, latent_length, 1, 1)
+                .clone()
+            )
+            video_cfg_base[:, :, : cfg_config.cfg_trick_start_frame] = min(
+                cfg_config.cfg_trick_value,
+                cfg_config.video_txt_guidance_scale,
+            )
+        else:
+            video_cfg_base = torch.tensor([cfg_config.video_txt_guidance_scale], device=device)
+
+        all_video_cfgs: list[torch.Tensor] = []
+        all_audio_cfgs: list[float] = []
+        for timestep in t_list:
+            video_cfg = video_cfg_base.clone()
+            if cfg_config.use_dynamic_cfg:
+                timestep_value = float(timestep.detach().reshape(-1)[0].item())
+                if timestep_value < cfg_config.dynamic_cfg_start_t:
+                    video_cfg[video_cfg > cfg_config.dynamic_cfg_cutoff_value] = cfg_config.dynamic_cfg_cutoff_value
+            all_video_cfgs.append(video_cfg)
+            all_audio_cfgs.append(cfg_config.audio_txt_guidance_scale)
+        return all_video_cfgs, all_audio_cfgs
+
+    @staticmethod
+    def _cfg_scale_tensor(
+        guidance_scale: torch.Tensor | float,
+        like: torch.Tensor,
+    ) -> torch.Tensor:
+        if isinstance(guidance_scale, torch.Tensor):
+            return guidance_scale.to(device=like.device, dtype=like.dtype)
+        return torch.tensor(float(guidance_scale), device=like.device, dtype=like.dtype)
+
+    @classmethod
+    def _get_skimming_mask(
+        cls,
+        x_orig: torch.Tensor,
+        cond: torch.Tensor,
+        uncond: torch.Tensor,
+        guidance_scale: torch.Tensor | float,
+    ) -> torch.Tensor:
+        guidance_scale = cls._cfg_scale_tensor(guidance_scale, cond)
+        x_orig = x_orig.to(device=cond.device, dtype=cond.dtype)
+        denoised = x_orig - ((x_orig - uncond) + guidance_scale * ((x_orig - cond) - (x_orig - uncond)))
+        matching_pred_signs = (cond - uncond).sign() == cond.sign()
+        matching_diff_after = cond.sign() == (cond * guidance_scale - uncond * (guidance_scale - 1)).sign()
+        deviation_influence = denoised.sign() == (denoised - x_orig).sign()
+        return matching_pred_signs & matching_diff_after & deviation_influence
+
+    @classmethod
+    def _apply_skimmed_cfg_linear(
+        cls,
+        x_orig: torch.Tensor | None,
+        cond: torch.Tensor,
+        uncond: torch.Tensor,
+        guidance_scale: torch.Tensor | float,
+        cfg_config: CFGConfig | None,
+    ) -> torch.Tensor:
+        if (
+            cfg_config is None
+            or not cfg_config.use_skimmed_cfg_linear
+            or x_orig is None
+            or not torch.any(uncond).item()
+        ):
+            return uncond
+
+        guidance_scale = cls._cfg_scale_tensor(guidance_scale, cond)
+        valid_scale = guidance_scale > 1
+        if not torch.any(valid_scale).item():
+            return uncond
+        scale_delta = torch.where(
+            valid_scale,
+            guidance_scale - 1,
+            torch.ones_like(guidance_scale),
+        )
+        fallback_weight = (cfg_config.skimmed_cfg_scale - 1) / scale_delta
+        target_uncond = cond * (1 - fallback_weight) + uncond * fallback_weight
+        skim_mask = cls._get_skimming_mask(x_orig, cond, uncond, guidance_scale)
+        uncond = torch.where(skim_mask & valid_scale, target_uncond, uncond)
+
+        target_uncond = cond * (1 - fallback_weight) + uncond * fallback_weight
+        skim_mask = cls._get_skimming_mask(x_orig, uncond, cond, guidance_scale)
+        return torch.where(skim_mask & valid_scale, target_uncond, uncond)
+
+    @staticmethod
+    def _cfg_rescale_dims(tensor: torch.Tensor) -> tuple[int, ...]:
+        if tensor.ndim <= 1:
+            return ()
+        if tensor.ndim in (2, 3):
+            return (1,)
+        return tuple(range(2, tensor.ndim))
+
+    @classmethod
+    def _apply_cfg_rescale(
+        cls,
+        pos: torch.Tensor,
+        cfg: torch.Tensor,
+        cfg_config: CFGConfig | None,
+    ) -> torch.Tensor:
+        if cfg_config is None or cfg_config.cfg_rescale <= 0 or cfg.numel() == 0:
+            return cfg
+        spatial_dims = cls._cfg_rescale_dims(cfg)
+        if not spatial_dims:
+            return cfg
+        pos_std = pos.float().std(dim=spatial_dims, keepdim=True, unbiased=False)
+        cfg_std = cfg.float().std(dim=spatial_dims, keepdim=True, unbiased=False).clamp_min(1e-6)
+        factor = cfg_config.cfg_rescale * (pos_std / cfg_std) + (1 - cfg_config.cfg_rescale)
+        return cfg * factor.to(device=cfg.device, dtype=cfg.dtype)
+
+    def cfg_velocity(
+        self,
+        model_output: tuple[torch.Tensor, torch.Tensor],
+        video_txt_guidance_scale: torch.Tensor | float,
+        audio_txt_guidance_scale: torch.Tensor | float,
+        cfg_config: CFGConfig | None = None,
+        latent: torch.Tensor | None = None,
+        audio_latent: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        cond_uncond_video, cond_uncond_audio = model_output
+        if cond_uncond_video.shape[0] % 2 or cond_uncond_audio.shape[0] % 2:
+            raise ValueError("MAGI-2 CFG outputs must contain equal conditional/unconditional halves")
+        video_batch = cond_uncond_video.shape[0] // 2
+        audio_batch = cond_uncond_audio.shape[0] // 2
+        if video_batch != audio_batch:
+            raise ValueError("video and audio CFG output batch sizes differ")
+
+        if cond_uncond_video.numel() > 0:
+            cond_video = cond_uncond_video[:video_batch]
+            uncond_video = cond_uncond_video[video_batch:]
+            uncond_video = self._apply_skimmed_cfg_linear(
+                latent,
+                cond_video,
+                uncond_video,
+                video_txt_guidance_scale,
+                cfg_config,
+            )
+            cfg_video = uncond_video + video_txt_guidance_scale * (cond_video - uncond_video)
+            cfg_video = self._apply_cfg_rescale(cond_video, cfg_video, cfg_config)
+        else:
+            cfg_video = None
+
+        if cond_uncond_audio.numel() > 0:
+            cond_audio = cond_uncond_audio[:audio_batch]
+            uncond_audio = cond_uncond_audio[audio_batch:]
+            uncond_audio = self._apply_skimmed_cfg_linear(
+                audio_latent,
+                cond_audio,
+                uncond_audio,
+                audio_txt_guidance_scale,
+                cfg_config,
+            )
+            cfg_audio = uncond_audio + audio_txt_guidance_scale * (cond_audio - uncond_audio)
+            cfg_audio = self._apply_cfg_rescale(cond_audio, cfg_audio, cfg_config)
+        else:
+            cfg_audio = None
+        return cfg_video, cfg_audio
+
+    def combine_cfg_noise(
+        self,
+        positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        true_cfg_scale: float,
+        cfg_normalize: bool = False,
+        kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Preserve MAGI-2's dual-modality guidance under shared CFG dispatch."""
+
+        del true_cfg_scale, cfg_normalize
+        if not isinstance(positive_noise_pred, tuple) or not isinstance(negative_noise_pred, tuple):
+            raise TypeError("MAGI-2 CFG predictions must contain video and audio tensors")
+        if len(positive_noise_pred) != 2 or len(negative_noise_pred) != 2:
+            raise ValueError("MAGI-2 CFG predictions must contain exactly two tensors")
+        if kwargs is None:
+            raise ValueError("MAGI-2 CFG combine requires guidance and latent context")
+        cfg_video, cfg_audio = self.cfg_velocity(
+            (
+                torch.cat((positive_noise_pred[0], negative_noise_pred[0]), dim=0),
+                torch.cat((positive_noise_pred[1], negative_noise_pred[1]), dim=0),
+            ),
+            kwargs["video_txt_guidance_scale"],
+            kwargs["audio_txt_guidance_scale"],
+            kwargs.get("cfg_config"),
+            kwargs.get("latent"),
+            kwargs.get("audio_latent"),
+        )
+        if cfg_video is None or cfg_audio is None:
+            raise RuntimeError("MAGI-2 Preview must produce both video and audio CFG predictions")
+        return cfg_video, cfg_audio
+
+    @staticmethod
+    def _step_guided(
+        guided_output: tuple[torch.Tensor, torch.Tensor],
+        latent: torch.Tensor,
+        audio_latent: torch.Tensor,
+        video_scheduler: SchedulerMixin,
+        audio_scheduler: SchedulerMixin,
+        t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cfg_video, cfg_audio = guided_output
+        latent = video_scheduler.step(cfg_video, t, latent, return_dict=False)[0]
+        audio_latent = audio_scheduler.step(cfg_audio, t, audio_latent, return_dict=False)[0]
+        return latent, audio_latent
+
+    def step(
+        self,
+        model_output: tuple[torch.Tensor, torch.Tensor],
+        latent: torch.Tensor,
+        audio_latent: torch.Tensor,
+        video_txt_guidance_scale: torch.Tensor | float,
+        audio_txt_guidance_scale: torch.Tensor | float,
+        video_scheduler: SchedulerMixin,
+        audio_scheduler: SchedulerMixin,
+        t: torch.Tensor,
+        cfg_config: CFGConfig | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        cfg_video, cfg_audio = self.cfg_velocity(
+            model_output,
+            video_txt_guidance_scale,
+            audio_txt_guidance_scale,
+            cfg_config,
+            latent,
+            audio_latent,
+        )
+        if cfg_video is not None:
+            latent = video_scheduler.step(cfg_video, t, latent, return_dict=False)[0]
+        if cfg_audio is not None:
+            audio_latent = audio_scheduler.step(
+                cfg_audio,
+                t,
+                audio_latent,
+                return_dict=False,
+            )[0]
+        return latent, audio_latent, cfg_video, cfg_audio
+
+
+__all__ = [
+    "CFGConfig",
+    "Magi2PreviewSampler",
+    "SamplerInput",
+    "build_magi2_preview_schedulers",
+]

vllm_omni/diffusion/models/magi2/turbo_vae.py

@@ -0,0 +1,582 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025-2026 SandAI. All Rights Reserved.
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Native decoder for the TurboVAED checkpoint shipped with MAGI-2.
+
+Adapted from SandAI's MAGI-2 Preview inference implementation. Training-only
+modules and MagiCompiler integration are intentionally omitted.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from einops import rearrange
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+_WAN22_LATENT_MEAN = (
+    -0.2289,
+    -0.0052,
+    -0.1323,
+    -0.2339,
+    -0.2799,
+    0.0174,
+    0.1838,
+    0.1557,
+    -0.1382,
+    0.0542,
+    0.2813,
+    0.0891,
+    0.1570,
+    -0.0098,
+    0.0375,
+    -0.1825,
+    -0.2246,
+    -0.1207,
+    -0.0698,
+    0.5109,
+    0.2665,
+    -0.2108,
+    -0.2158,
+    0.2502,
+    -0.2055,
+    -0.0322,
+    0.1109,
+    0.1567,
+    -0.0729,
+    0.0899,
+    -0.2799,
+    -0.1230,
+    -0.0313,
+    -0.1649,
+    0.0117,
+    0.0723,
+    -0.2839,
+    -0.2083,
+    -0.0520,
+    0.3748,
+    0.0152,
+    0.1957,
+    0.1433,
+    -0.2944,
+    0.3573,
+    -0.0548,
+    -0.1681,
+    -0.0667,
+)
+_WAN22_LATENT_STD = (
+    0.4765,
+    1.0364,
+    0.4514,
+    1.1677,
+    0.5313,
+    0.4990,
+    0.4818,
+    0.5013,
+    0.8158,
+    1.0344,
+    0.5894,
+    1.0901,
+    0.6885,
+    0.6165,
+    0.8454,
+    0.4978,
+    0.5759,
+    0.3523,
+    0.7135,
+    0.6804,
+    0.5833,
+    1.4146,
+    0.8986,
+    0.5659,
+    0.7069,
+    0.5338,
+    0.4889,
+    0.4917,
+    0.4069,
+    0.4999,
+    0.6866,
+    0.4093,
+    0.5709,
+    0.6065,
+    0.6415,
+    0.4944,
+    0.5726,
+    1.2042,
+    0.5458,
+    1.6887,
+    0.3971,
+    1.0600,
+    0.3943,
+    0.5537,
+    0.5444,
+    0.4089,
+    0.7468,
+    0.7744,
+)
+
+
+def turbo_unpatchify(x: torch.Tensor, patch_size: int) -> torch.Tensor:
+    if patch_size == 1:
+        return x
+    if x.ndim != 5:
+        raise ValueError(f"Expected a 5-D TurboVAED output, got {tuple(x.shape)}")
+    batch, patched_channels, frames, height, width = x.shape
+    patch_area = patch_size * patch_size
+    if patched_channels % patch_area:
+        raise ValueError("TurboVAED output channels must be divisible by patch_size squared")
+    channels = patched_channels // patch_area
+    x = x.view(batch, channels, patch_size, patch_size, frames, height, width)
+    x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous()
+    return x.view(batch, channels, frames, height * patch_size, width * patch_size)
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, eps: float = 1e-8) -> None:
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        variance = x.float().pow(2).mean(1, keepdim=True)
+        return (x * torch.rsqrt(variance + self.eps)).to(dtype)
+
+
+class _CausalConv3d(nn.Module):
+    """Checkpoint-compatible non-causal 3-D convolution with edge time padding."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int, int] = 3,
+        stride: int | tuple[int, int, int] = 1,
+        *,
+        is_causal: bool = False,
+    ) -> None:
+        super().__init__()
+        if is_causal:
+            raise ValueError("The released MAGI-2 TurboVAED checkpoint is non-causal")
+        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,) * 3
+        stride = stride if isinstance(stride, tuple) else (stride,) * 3
+        self.conv = nn.Conv3d(
+            in_channels,
+            out_channels,
+            self.kernel_size,
+            stride=stride,
+            padding=(0, self.kernel_size[1] // 2, self.kernel_size[2] // 2),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        temporal_pad = (self.kernel_size[0] - 1) // 2
+        if temporal_pad:
+            x = torch.cat(
+                [
+                    x[:, :, :1].repeat(1, 1, temporal_pad, 1, 1),
+                    x,
+                    x[:, :, -1:].repeat(1, 1, temporal_pad, 1, 1),
+                ],
+                dim=2,
+            )
+        return self.conv(x)
+
+
+class _ResNetBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        *,
+        is_upsampler_modified: bool = False,
+        is_causal: bool = False,
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or in_channels
+        self.nonlinearity = nn.SiLU()
+        self.replace_nonlinearity = nn.ReLU()
+        self.is_upsampler_modified = is_upsampler_modified
+        self.norm1 = _RMSNorm()
+        self.conv1 = _CausalConv3d(in_channels, out_channels, is_causal=is_causal)
+        self.norm2 = _RMSNorm()
+        self.conv2 = _CausalConv3d(out_channels, out_channels, is_causal=is_causal)
+        self.norm3 = _RMSNorm(eps=1e-6) if in_channels != out_channels else None
+        self.conv_shortcut = (
+            _CausalConv3d(in_channels, out_channels, kernel_size=1, is_causal=is_causal)
+            if in_channels != out_channels
+            else None
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        x = self.norm1(inputs)
+        x = self.replace_nonlinearity(x) if self.is_upsampler_modified else self.nonlinearity(x)
+        x = self.conv1(x)
+        x = self.conv2(self.nonlinearity(self.norm2(x)))
+        residual = self.norm3(inputs) if self.norm3 is not None else inputs
+        if self.conv_shortcut is not None:
+            residual = self.conv_shortcut(residual)
+        return x + residual
+
+
+class _WanUpsample(nn.Upsample):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x.float()).type_as(x)
+
+
+class _WanResample(nn.Module):
+    def __init__(self, channels: int, mode: str) -> None:
+        super().__init__()
+        self.mode = mode
+        if mode in ("upsample2d", "upsample3d"):
+            self.resample = nn.Sequential(
+                _WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(channels, channels, 3, padding=1),
+            )
+        else:
+            self.resample = nn.Identity()
+        if mode == "upsample3d":
+            self.time_conv = _CausalConv3d(channels, channels * 2, (3, 1, 1))
+
+    def forward(self, x: torch.Tensor, *, is_first_chunk: bool) -> torch.Tensor:
+        batch, channels, frames, _height, _width = x.shape
+        if self.mode == "upsample3d":
+            x = self.time_conv(x)
+            x = rearrange(x, "b (split c) t h w -> b c (t split) h w", split=2)
+            if x.shape[1:3] != (channels, frames * 2):
+                raise RuntimeError("Unexpected temporal upsample shape")
+            if is_first_chunk:
+                x = x[:, :, 1:]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        return rearrange(x, "(b t) c h w -> b c t h w", b=batch)
+
+
+class _MidBlock3d(nn.Module):
+    def __init__(self, channels: int, num_layers: int, *, is_causal: bool) -> None:
+        super().__init__()
+        self.resnets = nn.ModuleList([_ResNetBlock3d(channels, is_causal=is_causal) for _ in range(num_layers)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for resnet in self.resnets:
+            x = resnet(x)
+        return x
+
+
+class _UpBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int,
+        *,
+        spatio_temporal_scale: bool,
+        spatio_only: bool,
+        is_causal: bool,
+    ) -> None:
+        super().__init__()
+        self.conv_in = (
+            _ResNetBlock3d(in_channels, out_channels, is_causal=is_causal) if in_channels != out_channels else None
+        )
+        self.upsamplers = (
+            nn.ModuleList([_WanResample(out_channels, "upsample2d" if spatio_only else "upsample3d")])
+            if spatio_temporal_scale
+            else None
+        )
+        self.resnets = nn.ModuleList(
+            [
+                _ResNetBlock3d(
+                    out_channels,
+                    is_upsampler_modified=spatio_temporal_scale,
+                    is_causal=is_causal,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(self, x: torch.Tensor, *, is_first_chunk: bool) -> torch.Tensor:
+        if self.conv_in is not None:
+            x = self.conv_in(x)
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                x = upsampler(x, is_first_chunk=is_first_chunk)
+        for resnet in self.resnets:
+            x = resnet(x)
+        return x
+
+
+class _TurboDecoder3d(nn.Module):
+    def __init__(
+        self,
+        *,
+        latent_channels: int,
+        out_channels: int,
+        block_out_channels: tuple[int, ...],
+        layers_per_block: tuple[int, ...],
+        spatio_temporal_scaling: tuple[bool, ...],
+        spatio_only: tuple[bool, ...],
+        patch_size: int,
+        decoder_causal: bool,
+        use_unpatchify: bool,
+    ) -> None:
+        super().__init__()
+        if patch_size != 2 or not use_unpatchify:
+            raise ValueError("The released MAGI-2 TurboVAED requires patch_size=2 and use_unpatchify=true")
+        self.patch_size = patch_size
+        blocks = tuple(reversed(block_out_channels))
+        scaling = tuple(reversed(spatio_temporal_scaling))
+        layers = tuple(reversed(layers_per_block))
+        spatial_only = tuple(reversed(spatio_only))
+        current = blocks[0]
+
+        self.conv_in = _CausalConv3d(latent_channels, current, is_causal=decoder_causal)
+        self.mid_block = _MidBlock3d(current, layers[0], is_causal=decoder_causal)
+        self.up_blocks = nn.ModuleList()
+        for index, output in enumerate(blocks):
+            self.up_blocks.append(
+                _UpBlock3d(
+                    current,
+                    output,
+                    layers[index + 1],
+                    spatio_temporal_scale=scaling[index],
+                    spatio_only=spatial_only[index],
+                    is_causal=decoder_causal,
+                )
+            )
+            current = output
+        self.conv_act = nn.SiLU()
+        self.conv_out = _CausalConv3d(
+            current,
+            out_channels * patch_size * patch_size,
+            is_causal=decoder_causal,
+        )
+
+    def forward(self, x: torch.Tensor, *, is_first_chunk: bool) -> torch.Tensor:
+        x = self.mid_block(self.conv_in(x))
+        for block in self.up_blocks:
+            x = block(x, is_first_chunk=is_first_chunk)
+        dtype = x.dtype
+        x = (x * torch.rsqrt(x.float().pow(2).mean(1, keepdim=True) + 1e-8)).to(dtype)
+        return turbo_unpatchify(self.conv_out(self.conv_act(x)), self.patch_size)
+
+
+def _load_turbo_config(config_path: Path) -> dict[str, Any]:
+    if not config_path.is_file():
+        raise FileNotFoundError(f"TurboVAED config does not exist: {config_path}")
+    with config_path.open(encoding="utf-8") as file:
+        config = json.load(file)
+    unsupported = {
+        "decoder_causal": config.get("decoder_causal", False),
+        "decoder_is_dw_conv": any(config.get("decoder_is_dw_conv", [])),
+        "use_unpatchify": not config.get("use_unpatchify", False),
+    }
+    enabled = [name for name, value in unsupported.items() if value]
+    if enabled:
+        raise ValueError(f"Unsupported TurboVAED configuration: {', '.join(enabled)}")
+    return config
+
+
+def extract_turbo_decoder_state_dict(checkpoint: Mapping[str, Any]) -> dict[str, torch.Tensor]:
+    """Select decoder EMA tensors and strip the training module prefix."""
+
+    state: Any = checkpoint
+    if "ema_state_dict" in state:
+        state = state["ema_state_dict"]
+    elif "state_dict" in state:
+        state = state["state_dict"]
+    if not isinstance(state, dict):
+        raise ValueError("TurboVAED checkpoint does not contain a state dictionary")
+    result = {}
+    for key, value in state.items():
+        target = key[7:] if key.startswith("module.") else key
+        if target.startswith("decoder."):
+            result[target] = value
+    if not result:
+        raise ValueError("TurboVAED checkpoint contains no decoder tensors")
+    return result
+
+
+class Magi2TurboVAEDecoder(nn.Module, DistributedVaeMixin):
+    """Load and run MAGI-2's default distilled video decoder."""
+
+    def __init__(
+        self,
+        config_path: str | Path,
+        checkpoint_path: str | Path,
+        *,
+        device: str | torch.device = "cpu",
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        super().__init__()
+        config = _load_turbo_config(Path(config_path))
+        self.z_dim = int(config["latent_channels"])
+        self.first_chunk_size = int(config["first_chunk_size"])
+        self.step_size = int(config["step_size"])
+        self.temporal_compression_ratio = int(config["temporal_compression_ratio"])
+        self.spatial_compression_ratio = int(config["spatial_compression_ratio"])
+        self.decoder = _TurboDecoder3d(
+            latent_channels=self.z_dim,
+            out_channels=int(config["out_channels"]),
+            block_out_channels=tuple(config["decoder_block_out_channels"]),
+            layers_per_block=tuple(config["decoder_layers_per_block"]),
+            spatio_temporal_scaling=tuple(config["decoder_spatio_temporal_scaling"]),
+            spatio_only=tuple(config["decoder_spatio_only"]),
+            patch_size=int(config["patch_size"]),
+            decoder_causal=bool(config["decoder_causal"]),
+            use_unpatchify=bool(config["use_unpatchify"]),
+        ).to(device=device, dtype=dtype)
+        self.register_buffer(
+            "latent_mean",
+            torch.tensor(_WAN22_LATENT_MEAN, dtype=torch.float32, device=device),
+            persistent=False,
+        )
+        self.register_buffer(
+            "latent_std",
+            torch.tensor(_WAN22_LATENT_STD, dtype=torch.float32, device=device),
+            persistent=False,
+        )
+
+        checkpoint_path = Path(checkpoint_path)
+        if not checkpoint_path.is_file():
+            raise FileNotFoundError(f"TurboVAED checkpoint does not exist: {checkpoint_path}")
+        checkpoint = torch.load(
+            checkpoint_path,
+            map_location="cpu",
+            weights_only=True,
+            mmap=True,
+        )
+        state = extract_turbo_decoder_state_dict(checkpoint)
+        self.load_state_dict(state, strict=True)
+        self.eval().requires_grad_(False)
+        self.use_tiling = False
+        self.use_slicing = False
+        if dist.is_available() and dist.is_initialized():
+            self.init_distributed()
+
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
+        if mode != "tile":
+            raise ValueError(f"MAGI-2 TurboVAE supports temporal tile parallelism only, got {mode!r}")
+        if not hasattr(self, "distributed_executor"):
+            if parallel_size > 1:
+                raise RuntimeError("TurboVAE patch parallelism requires initialized torch.distributed")
+            return
+        super().set_parallel_size(parallel_size, mode=mode)
+
+    def is_distributed_enabled(self) -> bool:
+        return hasattr(self, "distributed_executor") and super().is_distributed_enabled()
+
+    def _prepare_latent(self, z: torch.Tensor) -> tuple[torch.Tensor, int]:
+        z = z.to(device=self.latent_std.device)
+        dtype = z.dtype
+        z = z * self.latent_std.view(1, self.z_dim, 1, 1, 1) + self.latent_mean.view(1, self.z_dim, 1, 1, 1)
+        z = z.to(dtype)
+        first = self.first_chunk_size
+        step = self.step_size
+        frames = z.shape[2]
+        if frames < first:
+            padding = first - frames
+        elif (frames - first) % step:
+            padding = step - (frames - first) % step
+        else:
+            padding = 0
+        if padding:
+            z = torch.cat([z, z[:, :, -1:].repeat(1, 1, padding, 1, 1)], dim=2)
+        return z, padding
+
+    def _chunk_tasks(self, z: torch.Tensor, padding: int) -> tuple[list[TileTask], GridSpec]:
+        first = self.first_chunk_size
+        step = self.step_size
+        overlap = self.temporal_compression_ratio
+        frames = z.shape[2]
+        descriptors: list[tuple[int, int, bool, int, int]] = []
+        if frames == first:
+            descriptors.append((0, frames, True, 0, 0))
+        else:
+            descriptors.append((0, first + 1, True, 0, overlap))
+            for index in range(first, frames, step):
+                last = index + step == frames
+                descriptors.append(
+                    (
+                        index - 1,
+                        index + step if last else index + step + 1,
+                        False,
+                        overlap,
+                        0 if last else overlap,
+                    )
+                )
+
+        tasks = [
+            TileTask(
+                tile_id=chunk_index,
+                grid_coord=(chunk_index, int(is_first), crop_left, crop_right),
+                tensor=z[:, :, left:right],
+                workload=(right - left) * z.shape[3] * z.shape[4],
+            )
+            for chunk_index, (left, right, is_first, crop_left, crop_right) in enumerate(descriptors)
+        ]
+        return tasks, GridSpec(
+            split_dims=(2,),
+            grid_shape=(len(tasks),),
+            tile_spec={"padding": padding},
+            output_dtype=z.dtype,
+        )
+
+    def _decode_chunk(self, task: TileTask) -> torch.Tensor:
+        _index, is_first, crop_left, crop_right = task.grid_coord
+        assert isinstance(task.tensor, torch.Tensor)
+        output = self.decoder(task.tensor, is_first_chunk=bool(is_first))
+        right = output.shape[2] - crop_right if crop_right else output.shape[2]
+        return output[:, :, crop_left:right]
+
+    def _merge_chunks(
+        self,
+        chunks: dict[tuple[int, ...], torch.Tensor],
+        grid_spec: GridSpec,
+    ) -> torch.Tensor:
+        output = torch.cat([tensor for _coord, tensor in sorted(chunks.items())], dim=2)
+        padding = int(grid_spec.tile_spec["padding"])
+        if padding:
+            output = output[:, :, : -padding * self.temporal_compression_ratio]
+        return output
+
+    @torch.inference_mode()
+    def decode(self, z: torch.Tensor, *, output_offload: bool = False) -> torch.Tensor:
+        z, padding = self._prepare_latent(z)
+        tasks, grid_spec = self._chunk_tasks(z, padding)
+        if self.is_distributed_enabled():
+            output = self.distributed_executor.execute(
+                z,
+                DistributedOperator(
+                    split=lambda _z: (tasks, grid_spec),
+                    exec=self._decode_chunk,
+                    merge=self._merge_chunks,
+                ),
+                broadcast_result=False,
+            )
+            return output.cpu() if output_offload and output.numel() else output
+
+        chunks = []
+        for task in tasks:
+            tensor = self._decode_chunk(task)
+            chunks.append((task.grid_coord, tensor.cpu() if output_offload else tensor))
+        return self._merge_chunks(dict(chunks), grid_spec)
+
+    forward = decode
+
+
+__all__ = [
+    "Magi2TurboVAEDecoder",
+    "extract_turbo_decoder_state_dict",
+    "turbo_unpatchify",
+]