[fix]: guard configure_local_attention pending LSE port (SR PR scope)

Addresses Oracle deep-review finding (MF1): the existing _local_window_attention math sums per-block SDPA outputs with a plain '+', but upstream daVinci-MagiHuman _flash_attn_with_correction reconstructs the full softmax via log-sum-exp correction across overlapping key ranges. The two are not numerically equivalent.

On this branch the path is unreachable by default (local_attn_layers=() for both base and DMD-2-distill variants) and the parity tests pass local_attn_handler=None, so this is dead code today. But it is also broken-code-on-activation: the moment the SR-1080p variant PR calls configure_local_attention([...]), it would silently diverge from upstream.

Raise NotImplementedError when any local-attn layer is requested, with a message that points the SR activation PR at the upstream LSE math (daVinci-MagiHuman/inference/model/dit/dit_module.py::_flash_attn_with_correction). The empty-list case (the only one used by this PR) is unchanged. Per the decomposition contract, every intermediate state is dead code, not broken code.

fastvideo/configs/models/dits/__init__.py

@@ -5,6 +5,7 @@ from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
+from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
@@ -13,5 +14,5 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
     "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
-    "StableAudioConfig"
+    "MagiHumanVideoConfig", "StableAudioConfig"
 ]

fastvideo/configs/models/dits/magi_human.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Architecture / model config for the daVinci-MagiHuman DiT.
+
+The MagiHuman base DiT is a 15B-parameter single-stream transformer that
+jointly denoises video, audio, and text tokens in one flat sequence. Layout
+details verified against GAIR/daVinci-MagiHuman's base/ shards (2026-04-24).
+
+This file captures only configuration. The module implementation lives in
+fastvideo/models/dits/magi_human.py and the pipeline wiring in
+fastvideo/pipelines/basic/magi_human/.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def _is_block_layer(n: str, m) -> bool:
+    # Match "block.layers.<idx>" — the FSDP shard boundary for MagiHuman.
+    parts = n.split(".")
+    return (len(parts) >= 3 and parts[0] == "block" and parts[1] == "layers" and str.isdigit(parts[2]))
+
+
+@dataclass
+class MagiHumanArchConfig(DiTArchConfig):
+    """MagiHuman base DiT architecture constants.
+
+    **Scope contract:** fields here must match the `transformer/config.json`
+    emitted by `scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py`
+    1:1, and both are sourced from the upstream Python reference
+    `inference/common/config.py::ModelConfig` (the HF root `config.json`
+    is empty so the Python source is canonical). Pipeline-level knobs
+    (VAE stride, fps, num_inference_steps, CFG scales, flow_shift,
+    t5_gemma_target_length) and data-proxy knobs (coords_style,
+    frame_receptive_field, ref_audio_offset, text_offset) live on
+    `MagiHumanBaseConfig`, NOT here.
+
+    `param_names_mapping` is intentionally empty: the FastVideo implementation
+    keeps the same module tree as the reference (`adapter.*`,
+    `block.layers.<i>.*`, `final_linear_{video,audio}.*`,
+    `final_norm_{video,audio}.*`), so converted weights load directly.
+    """
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_block_layer])
+
+    # No renames needed — the FastVideo module mirrors the reference names.
+    param_names_mapping: dict = field(default_factory=dict)
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    # --- transformer shape ---
+    num_layers: int = 40
+    hidden_size: int = 5120
+    head_dim: int = 128
+    num_query_groups: int = 8  # num_heads_kv (GQA)
+
+    # --- modality channels ---
+    # video_in_channels = z_dim (48) * patch_size product (1*2*2=4), so the
+    # embedder receives 192 per token. text_in_channels is T5Gemma-9B's
+    # encoder hidden size.
+    video_in_channels: int = 192
+    audio_in_channels: int = 64
+    text_in_channels: int = 3584
+
+    # --- block-level architecture switches ---
+    # Sandwich MoE: first and last 4 layers have per-modality experts
+    # (video/audio/text), middle layers share a single set of weights.
+    mm_layers: tuple[int, ...] = (0, 1, 2, 3, 36, 37, 38, 39)
+    local_attn_layers: tuple[int, ...] = ()
+    gelu7_layers: tuple[int, ...] = (0, 1, 2, 3)
+    post_norm_layers: tuple[int, ...] = ()
+    enable_attn_gating: bool = True
+    activation_type: str = "swiglu7"
+
+    # --- DiT patching (upstream `ModelConfig`-equivalent; NOT the VAE
+    #     stride, which is pipeline-level). ---
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    spatial_rope_interpolation: str = "extra"
+
+    # --- TReAD (token routing + early drop). Flattened from the upstream
+    #     nested `tread_config` dict so it round-trips through
+    #     `update_model_arch` cleanly. ---
+    tread_selection_rate: float = 0.5
+    tread_start_layer_idx: int = 2
+    tread_end_layer_idx: int = 25
+
+    # --- derived fields (populated in __post_init__) ---
+    num_attention_heads: int = 0  # hidden_size / head_dim
+    num_heads_kv: int = 0  # == num_query_groups
+    in_channels: int = 0  # mirror of video_in_channels (FastVideo contract)
+    out_channels: int = 0  # mirror of video_in_channels
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.num_attention_heads = self.hidden_size // self.head_dim
+        self.num_heads_kv = self.num_query_groups
+        self.in_channels = self.video_in_channels
+        self.out_channels = self.video_in_channels
+        # num_channels_latents is the VAE latent z_dim (48 for Wan 2.2 TI2V-5B).
+        # We don't declare z_dim on the arch config (it's a VAE property),
+        # but we still set num_channels_latents for the BaseDiT contract.
+        self.num_channels_latents = 48
+
+
+@dataclass
+class MagiHumanVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=MagiHumanArchConfig)
+
+    prefix: str = "magi_human"

fastvideo/models/dits/magi_human.py

@@ -0,0 +1,927 @@
+# SPDX-License-Identifier: Apache-2.0
+"""daVinci-MagiHuman DiT (base variant).
+
+Ported from https://github.com/GAIR-NLP/daVinci-MagiHuman
+(inference/model/dit/dit_module.py, ~950 lines in the reference).
+
+Architecture summary (verified against GAIR/daVinci-MagiHuman/base/ weights):
+
+  - 40 transformer layers, hidden 5120, head_dim 128.
+  - GQA with 40 query heads and 8 KV heads.
+  - Multi-modality "sandwich": layers 0..3 and 36..39 use 3-way modality
+    experts (video/audio/text) packed inside each linear as
+    weight[..., out * 3, in]. Middle layers share a single expert.
+  - Per-head attention gating: the QKV projection emits an extra
+    num_heads_q channels that are sigmoid-gated onto the attention output.
+  - Activation is GELU7 on layers 0..3 (non-gated, intermediate=4*hidden)
+    and SwiGLU7 elsewhere (gated, intermediate=int(hidden*4*2/3)//4*4).
+  - Position encoding is an element-wise Fourier embedding over 9-column
+    coords (t,h,w + original TxHxW + reference TxHxW), not a standard
+    1D/3D RoPE.
+  - Forward takes a flat concatenated token stream (video first, then
+    audio, then text) plus a modality map; the internal ModalityDispatcher
+    permutes by modality before each linear so per-expert chunks line up.
+
+Deviations from the "use fastvideo.layers primitives everywhere" guideline
+in the add-model skill:
+
+  - The packed-expert linears store weight as [out * num_experts, in].
+    FastVideo's ReplicatedLinear does not model this layout; we use raw
+    nn.Parameter with a small wrapper below. This is deliberate and scoped
+    to this DiT: ReplicatedLinear still handles the adapter.* embedders
+    and final_linear_{video,audio} (single-expert) here.
+  - Self-attention is full-sequence and crosses modalities inside the flat
+    concat stream; DistributedAttention assumes a clean spatial-sequence
+    layout, so for the first port we use torch SDPA. Multi-GPU sequence
+    parallelism is a follow-up.
+  - torch.compile via magi_compiler is replaced with a plain nn.Module.
+
+For the full history and shape-by-shape verification notes, see
+.claude/skills/add-model/SKILL.md and the scaffold PR description.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.attention import LocalAttention
+from fastvideo.configs.models.dits.magi_human import (
+    MagiHumanArchConfig,
+    MagiHumanVideoConfig,
+)
+from fastvideo.layers.rotary_embedding import _apply_rotary_emb
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.platforms import AttentionBackendEnum
+
+
+# ---------------------------------------------------------------------------
+# Enums
+# ---------------------------------------------------------------------------
+
+
+class Modality(IntEnum):
+    VIDEO = 0
+    AUDIO = 1
+    TEXT = 2
+
+
+# ---------------------------------------------------------------------------
+# Activations
+# ---------------------------------------------------------------------------
+
+
+def swiglu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
+    """Gated swish-GLU with OpenAI-OSS-style limits and +1 linear bias."""
+    in_dtype = x.dtype
+    x = x.to(torch.float32)
+    x_glu, x_linear = x[..., ::2], x[..., 1::2]
+    x_glu = x_glu.clamp(max=limit)
+    x_linear = x_linear.clamp(min=-limit, max=limit)
+    out_glu = x_glu * torch.sigmoid(alpha * x_glu)
+    return (out_glu * (x_linear + 1)).to(in_dtype)
+
+
+def gelu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
+    in_dtype = x.dtype
+    x = x.to(torch.float32).clamp(max=limit)
+    return (x * torch.sigmoid(alpha * x)).to(in_dtype)
+
+
+# ---------------------------------------------------------------------------
+# Modality dispatcher
+# ---------------------------------------------------------------------------
+
+
+class ModalityDispatcher:
+    """Permute a flat token stream so same-modality tokens are contiguous.
+
+    The DiT's multi-expert linears apply a different weight chunk per modality.
+    Instead of carrying a branch inside each Linear, we pre-permute tokens so
+    each chunk sees a contiguous slice, then un-permute before computing
+    RoPE/attention across the full sequence.
+    """
+
+    def __init__(self, modality_mapping: torch.Tensor, num_modalities: int):
+        self.modality_mapping = modality_mapping
+        self.num_modalities = num_modalities
+        self.permute_mapping = torch.argsort(modality_mapping)
+        self.inv_permute_mapping = torch.argsort(self.permute_mapping)
+        permuted = modality_mapping[self.permute_mapping]
+        self.group_size = torch.bincount(permuted, minlength=num_modalities).to(torch.int32)
+        self.group_size_cpu: list[int] = [int(x) for x in self.group_size.cpu().tolist()]
+
+    def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]:
+        return list(torch.split(x, self.group_size_cpu, dim=0))
+
+    def undispatch(self, *chunks: torch.Tensor) -> torch.Tensor:
+        return torch.cat(chunks, dim=0)
+
+    @staticmethod
+    def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor:
+        return x[permute_mapping]
+
+    @staticmethod
+    def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor:
+        return x[inv_permute_mapping]
+
+
+# ---------------------------------------------------------------------------
+# Norms, rotary embed
+# ---------------------------------------------------------------------------
+
+
+class MultiModalityRMSNorm(nn.Module):
+    """RMSNorm with optional per-modality scale.
+
+    When num_modality == 1, behaves identically to a standard RMSNorm with
+    weight initialized to zero (effective weight is 1 + weight, hence the
+    learnable +1 offset baked into the forward path). When num_modality > 1,
+    the weight tensor packs per-modality scales along its flat axis and the
+    dispatcher selects the right chunk per modality.
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6, num_modality: int = 1):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.num_modality = num_modality
+        # Always stored in fp32; matches the reference initialization.
+        self.weight = nn.Parameter(torch.zeros(dim * num_modality, dtype=torch.float32))
+
+    def _rms(self, x: torch.Tensor) -> torch.Tensor:
+        t = x.float()
+        return t * torch.rsqrt(t.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        modality_dispatcher: Optional[ModalityDispatcher] = None,
+    ) -> torch.Tensor:
+        original_dtype = x.dtype
+        t = self._rms(x)
+        if self.num_modality == 1:
+            return (t * (self.weight + 1)).to(original_dtype)
+        assert modality_dispatcher is not None, (
+            "MultiModalityRMSNorm with num_modality>1 requires a dispatcher"
+        )
+        weight_chunks = self.weight.chunk(self.num_modality, dim=0)
+        parts = modality_dispatcher.dispatch(t)
+        for i in range(self.num_modality):
+            parts[i] = parts[i] * (weight_chunks[i] + 1)
+        return modality_dispatcher.undispatch(*parts).to(original_dtype)
+
+
+def _freq_bands(num_bands: int, temperature: float = 10000.0) -> torch.Tensor:
+    exp = torch.arange(0, num_bands, 1, dtype=torch.int64).float() / num_bands
+    return 1.0 / (temperature ** exp)
+
+
+class ElementWiseFourierEmbed(nn.Module):
+    """Element-wise Fourier embedding over 9-column coords (t, h, w, T, H, W,
+    ref_T, ref_H, ref_W). Produces a per-token positional embedding that
+    acts as the RoPE angle input for attention.
+
+    Weight: `bands` of shape `[dim // 8]` (fixed at init via freq_bands).
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        temperature: float = 10000.0,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        self.dim = dim
+        bands = _freq_bands(dim // 8, temperature=temperature).to(dtype)
+        # `register_buffer` so state_dict keeps it, matching upstream naming.
+        self.register_buffer("bands", bands)
+
+    def forward(self, coords: torch.Tensor) -> torch.Tensor:
+        # coords: [L, 9] = (t, h, w, T, H, W, ref_T, ref_H, ref_W)
+        coords_xyz = coords[:, :3]
+        sizes = coords[:, 3:6]
+        refs = coords[:, 6:9]
+
+        scales = (refs - 1) / (sizes - 1)
+        scales[(refs == 1) & (sizes == 1)] = 1
+        # Center H and W (leave time uncentered).
+        centers = (sizes - 1) / 2
+        centers[:, 0] = 0
+        coords_xyz = coords_xyz - centers
+
+        proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands  # [L, 3, B]
+        sin_proj = proj.sin()
+        cos_proj = proj.cos()
+        return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)
+
+
+# ---------------------------------------------------------------------------
+# Packed-expert linear
+# ---------------------------------------------------------------------------
+
+
+class PackedExpertLinear(nn.Module):
+    """Linear where the weight is packed per-modality along the output axis.
+
+    Shapes:
+        weight: [out_features * num_experts, in_features]
+        bias:   [out_features * num_experts]  (optional)
+
+    When `num_experts == 1`, behaves exactly like `nn.Linear`. When
+    `num_experts > 1`, `forward` dispatches the input via the supplied
+    `ModalityDispatcher`, applies the per-modality weight/bias chunk, and
+    gathers the outputs in original order.
+
+    Why not use `ReplicatedLinear`? Because the packed-expert layout is not
+    what ReplicatedLinear (or any other fastvideo.layers.linear) is wired
+    for. Using raw `nn.Parameter` keeps weight loading trivial (names map
+    1:1 to the upstream checkpoint) and avoids quantization-path assumptions
+    that don't match this layout.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        num_experts: int = 1,
+        bias: bool = False,
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_experts = num_experts
+        self.use_bias = bias
+        self.weight = nn.Parameter(
+            torch.zeros(out_features * num_experts, in_features, dtype=dtype)
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                torch.zeros(out_features * num_experts, dtype=dtype)
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        modality_dispatcher: Optional[ModalityDispatcher] = None,
+    ) -> torch.Tensor:
+        if self.num_experts == 1:
+            return F.linear(x, self.weight, self.bias)
+        assert modality_dispatcher is not None, (
+            "PackedExpertLinear with num_experts>1 requires a dispatcher"
+        )
+        parts = modality_dispatcher.dispatch(x)
+        w_chunks = self.weight.chunk(self.num_experts, dim=0)
+        b_chunks = (
+            self.bias.chunk(self.num_experts, dim=0)
+            if self.bias is not None else [None] * self.num_experts
+        )
+        for i in range(self.num_experts):
+            parts[i] = F.linear(parts[i], w_chunks[i], b_chunks[i])
+        return modality_dispatcher.undispatch(*parts)
+
+
+# ---------------------------------------------------------------------------
+# Attention & MLP
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class AttentionSubConfig:
+    hidden_size: int
+    num_heads_q: int
+    num_heads_kv: int
+    head_dim: int
+    num_modality: int
+    enable_attn_gating: bool
+    use_local_attn: bool = False
+    frame_receptive_field: int = 11
+
+
+class MagiAttention(nn.Module):
+    """Self-attention with GQA + optional per-head sigmoid gating."""
+
+    def __init__(self, cfg: AttentionSubConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.gating_size = cfg.num_heads_q if cfg.enable_attn_gating else 0
+        qkv_out = (
+            cfg.num_heads_q * cfg.head_dim
+            + 2 * cfg.num_heads_kv * cfg.head_dim
+            + self.gating_size
+        )
+        self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality)
+        self.linear_qkv = PackedExpertLinear(
+            cfg.hidden_size, qkv_out, num_experts=cfg.num_modality, bias=False,
+        )
+        self.linear_proj = PackedExpertLinear(
+            cfg.num_heads_q * cfg.head_dim, cfg.hidden_size,
+            num_experts=cfg.num_modality, bias=False,
+        )
+        self.q_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality)
+        self.k_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality)
+
+        self.q_size = cfg.num_heads_q * cfg.head_dim
+        self.kv_size = cfg.num_heads_kv * cfg.head_dim
+
+        self.attn = LocalAttention(
+            num_heads=cfg.num_heads_q,
+            head_size=cfg.head_dim,
+            num_kv_heads=cfg.num_heads_kv,
+            causal=False,
+            supported_attention_backends=(
+                AttentionBackendEnum.FLASH_ATTN,
+                AttentionBackendEnum.TORCH_SDPA,
+            ),
+        )
+
+    def configure_local_attention(
+        self,
+        *,
+        enabled: bool,
+        frame_receptive_field: int = 11,
+    ) -> None:
+        self.cfg.use_local_attn = enabled
+        self.cfg.frame_receptive_field = frame_receptive_field
+
+    def _sdpa(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        """Run SDPA on [L, H, D] tensors and return [L, Hq, D]."""
+        if q.numel() == 0:
+            return q.new_empty(q.shape)
+        out = F.scaled_dot_product_attention(
+            q.transpose(0, 1).unsqueeze(0).contiguous(),
+            k.transpose(0, 1).unsqueeze(0).contiguous(),
+            v.transpose(0, 1).unsqueeze(0).contiguous(),
+            enable_gqa=self.cfg.num_heads_q != self.cfg.num_heads_kv,
+        )
+        return out.squeeze(0).transpose(0, 1).contiguous()
+
+    def _local_window_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        *,
+        num_video_tokens: int,
+        num_frames: int,
+    ) -> torch.Tensor:
+        """Approximate upstream FFAHandler block accumulation with SDPA.
+
+        SR-1080p's reference kernel sums three independently-normalized
+        attention contributions:
+
+        * video frame queries -> local-window video keys;
+        * all video queries -> all audio+text keys;
+        * all audio+text queries -> full sequence keys.
+
+        This method mirrors that accumulator semantics with ordinary SDPA
+        slices. It is intentionally scoped to single-process inference; layers
+        without ``use_local_attn`` keep the existing full LocalAttention path.
+        """
+        if num_frames <= 0 or num_video_tokens <= 0:
+            return self._sdpa(q, k, v)
+        if num_video_tokens % num_frames != 0:
+            raise ValueError(
+                f"MagiHuman local attention expects video tokens divisible by "
+                f"frames, got {num_video_tokens=} and {num_frames=}."
+            )
+
+        token_per_frame = num_video_tokens // num_frames
+        out = torch.zeros(
+            q.shape[0],
+            self.cfg.num_heads_q,
+            self.cfg.head_dim,
+            device=q.device,
+            dtype=q.dtype,
+        )
+        rf = int(self.cfg.frame_receptive_field)
+
+        q_video = q[:num_video_tokens]
+        k_video = k[:num_video_tokens]
+        v_video = v[:num_video_tokens]
+        for frame_idx in range(num_frames):
+            q_start = frame_idx * token_per_frame
+            q_end = q_start + token_per_frame
+            k_start = max(0, (frame_idx - rf) * token_per_frame)
+            k_end = min(num_video_tokens, (frame_idx + rf + 1) * token_per_frame)
+            out[q_start:q_end] = self._sdpa(
+                q_video[q_start:q_end],
+                k_video[k_start:k_end],
+                v_video[k_start:k_end],
+            )
+
+        if num_video_tokens < q.shape[0]:
+            k_at = k[num_video_tokens:]
+            v_at = v[num_video_tokens:]
+            out[:num_video_tokens] = out[:num_video_tokens] + self._sdpa(
+                q[:num_video_tokens],
+                k_at,
+                v_at,
+            )
+            out[num_video_tokens:] = self._sdpa(
+                q[num_video_tokens:],
+                k,
+                v,
+            )
+        return out
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rope: torch.Tensor,
+        permute_mapping: torch.Tensor,
+        inv_permute_mapping: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher,
+        num_video_tokens: int | None = None,
+        num_frames: int | None = None,
+    ) -> torch.Tensor:
+        orig_dtype = self.linear_qkv.weight.dtype
+        h = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(orig_dtype)
+        qkv = self.linear_qkv(h, modality_dispatcher=modality_dispatcher).float()
+        q, k, v, g = torch.split(
+            qkv, [self.q_size, self.kv_size, self.kv_size, self.gating_size], dim=-1,
+        )
+        q = q.view(-1, self.cfg.num_heads_q, self.cfg.head_dim)
+        k = k.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim)
+        v = v.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim)
+        g = g.view(-1, self.cfg.num_heads_q, 1) if self.gating_size else None
+
+        q = self.q_norm(q, modality_dispatcher=modality_dispatcher)
+        k = self.k_norm(k, modality_dispatcher=modality_dispatcher)
+
+        # Un-permute before RoPE + attention so positional order reflects
+        # the original (video, audio, text) concat — matches reference.
+        q = ModalityDispatcher.inv_permute(q, inv_permute_mapping)
+        k = ModalityDispatcher.inv_permute(k, inv_permute_mapping)
+        v = ModalityDispatcher.inv_permute(v, inv_permute_mapping)
+        if g is not None:
+            g = ModalityDispatcher.inv_permute(g, inv_permute_mapping)
+
+        # Element-wise Fourier embed packs sin/cos of 3 axes into a single
+        # `rope` tensor. Match reference's split:
+        #   sin_emb, cos_emb = rope.tensor_split(2, -1)
+        # Reference passes (cos_emb, sin_emb) but splits sin first — replicated
+        # exactly so weight parity holds. Partial RoPE: rope dim is
+        # 6 * (head_dim // 8) = 96 < head_dim (128), so the trailing 32
+        # head_dim positions stay unrotated, matching the reference.
+        sin_emb, cos_emb = rope.tensor_split(2, -1)
+        rot_dim = cos_emb.shape[-1] * 2
+        q_rot = _apply_rotary_emb(q[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True)
+        k_rot = _apply_rotary_emb(k[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True)
+        if rot_dim < q.shape[-1]:
+            q = torch.cat([q_rot, q[..., rot_dim:]], dim=-1)
+            k = torch.cat([k_rot, k[..., rot_dim:]], dim=-1)
+        else:
+            q, k = q_rot, k_rot
+
+        # Run SDPA via FastVideo's LocalAttention so the backend selection
+        # (SDPA / FlashAttn / SLA / SageAttn) flows through the standard
+        # configurable path. GQA is handled inside the SDPA backend via
+        # `enable_gqa=True` when num_heads_q != num_heads_kv, so we no
+        # longer need the manual `repeat_interleave` here.
+        # Attention math runs at orig_dtype (bf16 in production and in the
+        # parity test, since PackedExpertLinear's default is bf16, matching
+        # upstream BaseLinear at dit_module.py:330). The gating multiply
+        # promotes back to fp32 implicitly via PyTorch's type-promotion
+        # rules: bf16_attn_out * sigmoid(fp32_g) -> fp32, mirroring upstream
+        # dit_module.py:649.
+        q = q.to(orig_dtype)
+        k = k.to(orig_dtype)
+        v = v.to(orig_dtype)
+        if self.cfg.use_local_attn:
+            if num_video_tokens is None or num_frames is None:
+                raise ValueError("MagiHuman local attention requires video token/frame metadata.")
+            out = self._local_window_attention(
+                q,
+                k,
+                v,
+                num_video_tokens=num_video_tokens,
+                num_frames=num_frames,
+            )
+        else:
+            out = self.attn(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0)).squeeze(0)
+
+        out = ModalityDispatcher.permute(out, permute_mapping)
+        if g is not None:
+            g = ModalityDispatcher.permute(g, permute_mapping)
+            out = out * torch.sigmoid(g)
+        out = out.reshape(-1, self.cfg.num_heads_q * self.cfg.head_dim).to(orig_dtype)
+        return self.linear_proj(out, modality_dispatcher=modality_dispatcher)
+
+
+@dataclass
+class MLPSubConfig:
+    hidden_size: int
+    intermediate_size: int
+    activation: str              # "swiglu7" or "gelu7"
+    num_modality: int
+    gated: bool
+
+
+class MagiMLP(nn.Module):
+    def __init__(self, cfg: MLPSubConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality)
+        up_out = cfg.intermediate_size * 2 if cfg.gated else cfg.intermediate_size
+        self.up_gate_proj = PackedExpertLinear(
+            cfg.hidden_size, up_out, num_experts=cfg.num_modality, bias=False,
+        )
+        self.down_proj = PackedExpertLinear(
+            cfg.intermediate_size, cfg.hidden_size,
+            num_experts=cfg.num_modality, bias=False,
+        )
+        self._act = swiglu7 if cfg.activation == "swiglu7" else gelu7
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher,
+    ) -> torch.Tensor:
+        orig_dtype = self.up_gate_proj.weight.dtype
+        x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(orig_dtype)
+        x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).float()
+        x = self._act(x).to(orig_dtype)
+        x = self.down_proj(x, modality_dispatcher=modality_dispatcher).float()
+        return x
+
+
+class MagiTransformerLayer(nn.Module):
+    def __init__(self, arch: MagiHumanArchConfig, layer_idx: int):
+        super().__init__()
+        num_modality = 3 if layer_idx in arch.mm_layers else 1
+        self.post_norm = layer_idx in arch.post_norm_layers
+        self.layer_idx = layer_idx
+
+        self.attention = MagiAttention(AttentionSubConfig(
+            hidden_size=arch.hidden_size,
+            num_heads_q=arch.num_attention_heads,
+            num_heads_kv=arch.num_heads_kv,
+            head_dim=arch.head_dim,
+            num_modality=num_modality,
+            enable_attn_gating=arch.enable_attn_gating,
+            use_local_attn=layer_idx in arch.local_attn_layers,
+        ))
+
+        is_gelu7 = layer_idx in arch.gelu7_layers
+        if is_gelu7:
+            intermediate = arch.hidden_size * 4
+            gated = False
+            activation = "gelu7"
+        else:
+            intermediate = (arch.hidden_size * 4 * 2 // 3) // 4 * 4
+            gated = True
+            activation = "swiglu7"
+
+        self.mlp = MagiMLP(MLPSubConfig(
+            hidden_size=arch.hidden_size,
+            intermediate_size=intermediate,
+            activation=activation,
+            num_modality=num_modality,
+            gated=gated,
+        ))
+
+        if self.post_norm:
+            self.attn_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality)
+            self.mlp_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rope: torch.Tensor,
+        permute_mapping: torch.Tensor,
+        inv_permute_mapping: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher,
+        num_video_tokens: int | None = None,
+        num_frames: int | None = None,
+    ) -> torch.Tensor:
+        attn_out = self.attention(
+            hidden_states, rope, permute_mapping, inv_permute_mapping, modality_dispatcher,
+            num_video_tokens=num_video_tokens,
+            num_frames=num_frames,
+        )
+        if self.post_norm:
+            attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher)
+        hidden_states = hidden_states + attn_out
+
+        mlp_out = self.mlp(hidden_states, modality_dispatcher=modality_dispatcher)
+        if self.post_norm:
+            mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher)
+        return hidden_states + mlp_out
+
+
+# ---------------------------------------------------------------------------
+# Adapter (per-modality embedders + Fourier RoPE producer)
+# ---------------------------------------------------------------------------
+
+
+class MagiAdapter(nn.Module):
+    def __init__(self, arch: MagiHumanArchConfig):
+        super().__init__()
+        # Embedders stay in fp32 to match the reference dtype exactly.
+        self.video_embedder = nn.Linear(
+            arch.video_in_channels, arch.hidden_size, bias=True, dtype=torch.float32,
+        )
+        self.text_embedder = nn.Linear(
+            arch.text_in_channels, arch.hidden_size, bias=True, dtype=torch.float32,
+        )
+        self.audio_embedder = nn.Linear(
+            arch.audio_in_channels, arch.hidden_size, bias=True, dtype=torch.float32,
+        )
+        self.rope = ElementWiseFourierEmbed(arch.head_dim)
+        # RoPE cache: coords_mapping is the same tensor object across timesteps
+        # in the denoising loop. The key combines (id(t), data_ptr, _version,
+        # shape, dtype, device): id(t) catches Python-object identity (so a
+        # new tensor allocated at a previously-freed data_ptr does not collide),
+        # _version catches in-place mutation of the same tensor, and the
+        # remaining fields are belt-and-suspenders. We hold a strong reference
+        # to the source tensor so its id() stays unique while the cache is hot.
+        self._cached_rope: Optional[torch.Tensor] = None
+        self._cached_rope_key: Optional[tuple] = None
+        self._cached_rope_src: Optional[torch.Tensor] = None
+
+    def _rope_cache_key(self, t: torch.Tensor) -> tuple:
+        return (id(t), t.data_ptr(), t._version, t.shape, t.dtype, t.device)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        coords_mapping: torch.Tensor,
+        video_mask: torch.Tensor,
+        audio_mask: torch.Tensor,
+        text_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        key = self._rope_cache_key(coords_mapping)
+        if key != self._cached_rope_key:
+            self._cached_rope = self.rope(coords_mapping)
+            self._cached_rope_key = key
+            self._cached_rope_src = coords_mapping
+        rope = self._cached_rope
+        # Embedder dtypes may differ from x's dtype when FastVideo's FSDP
+        # loader casts all weights to `pipeline_config.precision` (bf16).
+        # Match the weight dtype per modality.
+        v_w = self.video_embedder.weight
+        a_w = self.audio_embedder.weight
+        t_w = self.text_embedder.weight
+        out = torch.zeros(
+            x.shape[0], self.video_embedder.out_features,
+            device=x.device, dtype=v_w.dtype,
+        )
+        out[text_mask] = self.text_embedder(
+            x[text_mask, : self.text_embedder.in_features].to(t_w.dtype)
+        ).to(out.dtype)
+        out[audio_mask] = self.audio_embedder(
+            x[audio_mask, : self.audio_embedder.in_features].to(a_w.dtype)
+        ).to(out.dtype)
+        out[video_mask] = self.video_embedder(
+            x[video_mask, : self.video_embedder.in_features].to(v_w.dtype)
+        ).to(out.dtype)
+        return out, rope
+
+
+# ---------------------------------------------------------------------------
+# Top-level DiT
+# ---------------------------------------------------------------------------
+
+
+class _TransformerBlock(nn.Module):
+    """Thin ModuleList wrapper to keep the 'block.layers.<i>' state_dict
+    naming identical to the upstream checkpoint (which uses a magi_compile
+    decorator producing `block.layers.<i>.*`)."""
+
+    def __init__(self, arch: MagiHumanArchConfig):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            MagiTransformerLayer(arch, i) for i in range(arch.num_layers)
+        ])
+
+    def configure_local_attention(
+        self,
+        local_attn_layers: tuple[int, ...],
+        frame_receptive_field: int = 11,
+    ) -> None:
+        enabled_layers = set(local_attn_layers)
+        for idx, layer in enumerate(self.layers):
+            layer.attention.configure_local_attention(
+                enabled=idx in enabled_layers,
+                frame_receptive_field=frame_receptive_field,
+            )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: torch.Tensor,
+        permute_mapping: torch.Tensor,
+        inv_permute_mapping: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher,
+        num_video_tokens: int | None = None,
+        num_frames: int | None = None,
+    ) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(
+                x,
+                rope,
+                permute_mapping,
+                inv_permute_mapping,
+                modality_dispatcher,
+                num_video_tokens=num_video_tokens,
+                num_frames=num_frames,
+            )
+        return x
+
+
+_CFG = MagiHumanVideoConfig()
+
+
+class MagiHumanDiT(BaseDiT):
+    """Top-level DiT for daVinci-MagiHuman (base).
+
+    Forward signature mirrors the reference `DiTModel.forward`: it takes a
+    flat token stream, its per-token coords and modality mapping, and
+    returns per-modality outputs packed into a max-channel-width tensor.
+
+    This scaffold is single-GPU only; the `ulysses_scheduler().dispatch(...)`
+    sequence-parallel wrapping in the reference has no equivalent here yet.
+
+    Dtype contract (the production loader / converter MUST preserve this):
+      * PackedExpertLinear params (all packed-expert qkv/proj/up/down) — bf16.
+      * MagiAdapter embedders, final_norm_{video,audio},
+        final_linear_{video,audio}, MultiModalityRMSNorm.weight — fp32.
+      * ElementWiseFourierEmbed.bands buffer — fp32.
+
+    The cross-layer residual stream stays in fp32; only the per-block linear
+    interior runs at bf16 (matching upstream BaseLinear's dtype boundary).
+    A loader that uniformly casts every tensor to one dtype (e.g.
+    `dit_precision=bf16`) WILL silently break numerical parity across the
+    40-layer accumulator. The conversion script in
+    `scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py` and
+    the pipeline-level loader configuration are responsible for honoring
+    this contract.
+    """
+
+    # BaseDiT requires these class attrs. Source them from the config so
+    # they stay in sync with MagiHumanVideoConfig edits.
+    _fsdp_shard_conditions = _CFG._fsdp_shard_conditions
+    _compile_conditions = _CFG._compile_conditions
+    _supported_attention_backends = _CFG._supported_attention_backends
+    param_names_mapping = _CFG.param_names_mapping
+    reverse_param_names_mapping = _CFG.reverse_param_names_mapping
+    lora_param_names_mapping = _CFG.lora_param_names_mapping
+
+    def __init__(self, config: MagiHumanVideoConfig, hf_config: dict | None = None, **kwargs):
+        super().__init__(config=config, hf_config=hf_config or {})
+        arch: MagiHumanArchConfig = getattr(config, "arch_config", config)
+        self.arch = arch
+
+        # BaseDiT contract instance vars.
+        self.hidden_size = arch.hidden_size
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.num_channels_latents
+
+        self.adapter = MagiAdapter(arch)
+        self.block = _TransformerBlock(arch)
+        self.final_norm_video = MultiModalityRMSNorm(arch.hidden_size)
+        self.final_norm_audio = MultiModalityRMSNorm(arch.hidden_size)
+        self.final_linear_video = nn.Linear(
+            arch.hidden_size, arch.video_in_channels, bias=False, dtype=torch.float32,
+        )
+        self.final_linear_audio = nn.Linear(
+            arch.hidden_size, arch.audio_in_channels, bias=False, dtype=torch.float32,
+        )
+        # Dispatcher + mask cache: modality_mapping is the same tensor object
+        # across all timesteps in the denoising loop. See _modality_cache_key
+        # below for the multi-field key design (id + data_ptr + _version +
+        # shape + dtype + device) that defends against CUDA allocator pointer
+        # reuse and in-place mutation. Strong references to the source tensors
+        # are held so id() stays unique while the cache is hot.
+        self._cached_dispatcher: Optional[ModalityDispatcher] = None
+        self._cached_video_mask: Optional[torch.Tensor] = None
+        self._cached_audio_mask: Optional[torch.Tensor] = None
+        self._cached_text_mask: Optional[torch.Tensor] = None
+        self._cached_modality_key: Optional[tuple] = None
+        self._cached_modality_src: Optional[torch.Tensor] = None
+        # num_frames is keyed on (coords_mapping_key, num_video_tokens) since
+        # the value reads coords_mapping[:num_video_tokens, 0].max() — both
+        # source inputs must match for the cached number to remain valid.
+        self._cached_num_video_tokens: int = 0
+        self._cached_num_frames: int = 0
+        self._cached_coords_key: Optional[tuple] = None
+        self._cached_coords_src: Optional[torch.Tensor] = None
+
+    def configure_local_attention(
+        self,
+        local_attn_layers: tuple[int, ...] | list[int],
+        frame_receptive_field: int = 11,
+    ) -> None:
+        layers = tuple(int(layer) for layer in local_attn_layers)
+        if layers:
+            # The current `_local_window_attention` accumulates per-block SDPA
+            # outputs with a plain sum, but upstream
+            # `daVinci-MagiHuman/inference/model/dit/dit_module.py::
+            # _flash_attn_with_correction` reconstructs the equivalent full
+            # softmax via log-sum-exp correction across overlapping key
+            # ranges. The two are NOT numerically equivalent, so enabling
+            # any local-attn layer on this DiT (used by the SR-1080p variant
+            # in the upstream stack) would silently diverge from upstream
+            # output. The activation PR for the SR variant must port the
+            # LSE-corrected math and add a parity test before re-enabling
+            # this path.
+            raise NotImplementedError(
+                "MagiHuman local-window attention is not yet ported. Upstream "
+                "uses log-sum-exp-corrected block accumulation; this DiT "
+                "currently has only a plain-sum approximation that does not "
+                "match upstream output. The SR-1080p activation PR will port "
+                "and parity-test this path. For the base / DMD-2-distill "
+                "variants leave `local_attn_layers=[]` (the default)."
+            )
+        self.arch.local_attn_layers = layers
+        self.block.configure_local_attention(layers, frame_receptive_field)
+
+    def _modality_cache_key(self, t: torch.Tensor) -> tuple:
+        return (id(t), t.data_ptr(), t._version, t.shape, t.dtype, t.device)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        coords_mapping: torch.Tensor,
+        modality_mapping: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x:                [L, max(V_ch, A_ch, T_ch)]
+            coords_mapping:   [L, 9]
+            modality_mapping: [L]  (int in {VIDEO, AUDIO, TEXT})
+        Returns:
+            out: [L, max(V_ch, A_ch)] with video channels in video slots and
+                 audio channels in audio slots; text slots are zero.
+        """
+        key = self._modality_cache_key(modality_mapping)
+        if key != self._cached_modality_key:
+            self._cached_dispatcher = ModalityDispatcher(modality_mapping, num_modalities=3)
+            self._cached_video_mask = modality_mapping == Modality.VIDEO
+            self._cached_audio_mask = modality_mapping == Modality.AUDIO
+            self._cached_text_mask = modality_mapping == Modality.TEXT
+            self._cached_num_video_tokens = int(self._cached_video_mask.sum().item())
+            self._cached_modality_key = key
+            self._cached_modality_src = modality_mapping
+        dispatcher = self._cached_dispatcher
+        video_mask = self._cached_video_mask
+        audio_mask = self._cached_audio_mask
+        text_mask = self._cached_text_mask
+        num_video_tokens = self._cached_num_video_tokens
+
+        coords_key = (self._modality_cache_key(coords_mapping), num_video_tokens)
+        if coords_key != self._cached_coords_key:
+            if num_video_tokens:
+                self._cached_num_frames = int(
+                    coords_mapping[:num_video_tokens, 0].max().item()) + 1
+            else:
+                self._cached_num_frames = 0
+            self._cached_coords_key = coords_key
+            self._cached_coords_src = coords_mapping
+        num_frames = self._cached_num_frames
+
+        x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask)
+        # Keep the residual stream in adapter dtype (fp32) entering the block.
+        # Upstream daVinci-MagiHuman dit_module.py:923 casts to params_dtype,
+        # which is fp32 by default; each layer's pre_norm.to(bf16) handles
+        # the bf16 internal-compute boundary, and linear_proj outputs bf16
+        # which gets promoted back to fp32 by the residual addition. Casting
+        # the residual to bf16 here degrades the cross-layer accumulator and
+        # compounds visibly over 40 layers in pipeline parity.
+        x = ModalityDispatcher.permute(x, dispatcher.permute_mapping)
+
+        x = self.block(
+            x, rope,
+            permute_mapping=dispatcher.permute_mapping,
+            inv_permute_mapping=dispatcher.inv_permute_mapping,
+            modality_dispatcher=dispatcher,
+            num_video_tokens=num_video_tokens,
+            num_frames=num_frames,
+        )
+        x = ModalityDispatcher.inv_permute(x, dispatcher.inv_permute_mapping)
+
+        x_video = x[video_mask].to(self.final_norm_video.weight.dtype)
+        x_video = self.final_norm_video(x_video)
+        x_video = self.final_linear_video(x_video)
+
+        x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype)
+        x_audio = self.final_norm_audio(x_audio)
+        x_audio = self.final_linear_audio(x_audio)
+
+        max_ch = max(self.arch.video_in_channels, self.arch.audio_in_channels)
+        out = torch.zeros(x.shape[0], max_ch, device=x.device, dtype=x.dtype)
+        out[video_mask, : self.arch.video_in_channels] = x_video.to(out.dtype)
+        out[audio_mask, : self.arch.audio_in_channels] = x_audio.to(out.dtype)
+        return out
+
+
+EntryClass = MagiHumanDiT

tests/local_tests/helpers/magi_human_upstream.py

@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stubs to make upstream daVinci-MagiHuman code importable in-process.
+
+The upstream DiT (daVinci-MagiHuman/inference/model/dit/dit_module.py) hard-
+imports SandAI's internal `magi_compiler` + a distributed-runtime init
+that requires `torchrun`. Neither is available in a single-process
+parity test. This module installs the minimum stubs to let the upstream
+DiT load and run on a single GPU with cp_world_size == 1 (which makes
+Ulysses's scatter/gather a no-op).
+
+Use:
+    from tests.local_tests.helpers.magi_human_upstream import (
+        install_stubs, load_upstream_dit,
+    )
+    install_stubs()
+    model = load_upstream_dit(base_shard_dir, device=torch.device("cuda"))
+"""
+from __future__ import annotations
+
+import importlib
+import sys
+import types
+from pathlib import Path
+from typing import Callable
+
+
+# ---------------------------------------------------------------------------
+# magi_compiler stubs — identity decorators.
+# ---------------------------------------------------------------------------
+
+
+def _install_magi_compiler_stub() -> None:
+    """Stub magi_compiler + register the `torch.ops.infra.*` ops upstream
+    calls via `torch.ops`.
+
+    Upstream code decorates plain Python fns with
+    `@magi_register_custom_op(name="infra::flash_attn_func", ...)` and
+    then calls them as `torch.ops.infra.flash_attn_func(...)`. Our stub
+    decorator has to both (a) preserve the decorated fn for direct call
+    sites and (b) register the fn under the advertised torch.ops
+    namespace so `torch.ops.infra.*` resolves.
+
+    For parity testing we route `infra::flash_attn_func` through
+    `F.scaled_dot_product_attention`, matching the FastVideo DiT's
+    kernel choice so drift measured in this test is architectural,
+    not kernel-dependent.
+    """
+    if "magi_compiler" in sys.modules:
+        return
+
+    import torch
+    import torch.nn.functional as F
+
+    pkg = types.ModuleType("magi_compiler")
+
+    def magi_compile(config_patch=None):
+        def decorator(cls_or_fn):
+            return cls_or_fn
+        return decorator
+
+    # Create one Library per (namespace, schema) pair. Track by namespace
+    # so we don't define the same op twice on re-import.
+    _libs: dict[str, torch.library.Library] = {}
+    _defined: set[tuple[str, str]] = set()
+
+    def _sdpa_flash_attn_func(q, k, v):
+        # Upstream shape: [batch=1, L, H, D]. SDPA expects [B, H, L, D]
+        # and no native GQA; expand K/V to match Q heads.
+        num_heads_q = q.shape[2]
+        num_heads_kv = k.shape[2]
+        if num_heads_q != num_heads_kv:
+            assert num_heads_q % num_heads_kv == 0
+            repeat = num_heads_q // num_heads_kv
+            k = k.repeat_interleave(repeat, dim=2)
+            v = v.repeat_interleave(repeat, dim=2)
+        q2 = q.transpose(1, 2).contiguous()
+        k2 = k.transpose(1, 2).contiguous()
+        v2 = v.transpose(1, 2).contiguous()
+        out = F.scaled_dot_product_attention(q2, k2, v2)
+        return out.transpose(1, 2).contiguous()
+
+    def _sdpa_segments(q, k, v, q_ranges, k_ranges):
+        # Upstream flex op shape: [L, H, D]. FFA accumulates each block's
+        # independently normalized attention output into the destination query
+        # slice. This SDPA fallback mirrors the accumulator semantics for
+        # SR-1080p parity tests without requiring SandAI's MagiAttention wheel.
+        out = torch.zeros(
+            q.shape[0],
+            q.shape[1],
+            q.shape[2],
+            dtype=q.dtype,
+            device=q.device,
+        )
+        num_heads_q = q.shape[1]
+        num_heads_kv = k.shape[1]
+        for q_range, k_range in zip(q_ranges.tolist(), k_ranges.tolist()):
+            qs, qe = int(q_range[0]), int(q_range[1])
+            ks, ke = int(k_range[0]), int(k_range[1])
+            q_block = q[qs:qe]
+            k_block = k[ks:ke]
+            v_block = v[ks:ke]
+            if num_heads_q != num_heads_kv:
+                assert num_heads_q % num_heads_kv == 0
+                repeat = num_heads_q // num_heads_kv
+                k_block = k_block.repeat_interleave(repeat, dim=1)
+                v_block = v_block.repeat_interleave(repeat, dim=1)
+            block_out = F.scaled_dot_product_attention(
+                q_block.transpose(0, 1).unsqueeze(0).contiguous(),
+                k_block.transpose(0, 1).unsqueeze(0).contiguous(),
+                v_block.transpose(0, 1).unsqueeze(0).contiguous(),
+            )
+            out[qs:qe] += block_out.squeeze(0).transpose(0, 1).contiguous()
+        lse = torch.empty((q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device)
+        return out, lse
+
+    def magi_register_custom_op(name=None, mutates_args=(), infer_output_meta_fn=None, is_subgraph_boundary=False, **kwargs):
+        def decorator(fn):
+            if not name:
+                return fn
+            namespace, op_name = name.split("::", 1)
+            if namespace not in _libs:
+                _libs[namespace] = torch.library.Library(namespace, "FRAGMENT")
+            if (namespace, op_name) in _defined:
+                # Already registered in a previous test run — reuse.
+                return fn
+            # Route known ops through SDPA; leave unknown ones as direct fn.
+            returns = "(Tensor, Tensor)" if op_name == "flex_flash_attn_func" else "Tensor"
+            schema_name = f"{op_name}({_infer_schema(fn)}) -> {returns}"
+            try:
+                _libs[namespace].define(schema_name)
+            except Exception:
+                pass
+            if op_name == "flash_attn_func":
+                torch.library.impl(
+                    _libs[namespace], op_name, "CUDA"
+                )(_sdpa_flash_attn_func)
+                torch.library.impl(
+                    _libs[namespace], op_name, "CPU"
+                )(_sdpa_flash_attn_func)
+            elif op_name == "flex_flash_attn_func":
+                torch.library.impl(
+                    _libs[namespace], op_name, "CUDA"
+                )(_sdpa_segments)
+            else:
+                # For ops we don't care about (compile-only wrappers), the
+                # Python fn path inside the module body is used directly —
+                # we just need `torch.ops.<ns>.<op>` to exist so module-
+                # load-time attribute lookups succeed.
+                torch.library.impl(
+                    _libs[namespace], op_name, "CUDA"
+                )(fn)
+            _defined.add((namespace, op_name))
+            return fn
+        return decorator
+
+    pkg.magi_compile = magi_compile
+    sys.modules["magi_compiler"] = pkg
+
+    api = types.ModuleType("magi_compiler.api")
+    api.magi_register_custom_op = magi_register_custom_op
+    sys.modules["magi_compiler.api"] = api
+    pkg.api = api
+
+    config_mod = types.ModuleType("magi_compiler.config")
+
+    class CompileConfig:
+        class offload_config:  # pragma: no cover - pass-through
+            gpu_resident_weight_ratio = 1.0
+    config_mod.CompileConfig = CompileConfig
+    sys.modules["magi_compiler.config"] = config_mod
+    pkg.config = config_mod
+
+
+def _infer_schema(fn) -> str:
+    """Return a minimal torch.library schema string for the given fn.
+
+    For our stub we just need *something* parseable; all real ops we
+    care about take `(q, k, v)` or `(q, k, v, q_ranges, k_ranges)` or
+    variants. Use generic `Tensor a, Tensor b, ...` arg names.
+    """
+    import inspect
+    sig = inspect.signature(fn)
+    parts = []
+    for i, name in enumerate(sig.parameters):
+        arg_name = name if name.isidentifier() else f"a{i}"
+        parts.append(f"Tensor {arg_name}")
+    return ", ".join(parts)
+
+
+# ---------------------------------------------------------------------------
+# distributed / CP stubs — single-GPU, cp_world_size == 1.
+# ---------------------------------------------------------------------------
+
+
+def _install_distributed_stubs() -> None:
+    """Monkey-patch upstream distributed + parallelism modules for cp=1."""
+    # inference.infra.distributed.*
+    # The real module requires NCCL / parallel_state to be initialized
+    # from torchrun; here we short-circuit the handful of getters the DiT
+    # actually calls.
+    import inference.infra.distributed as dist_mod
+
+    dist_mod.get_cp_world_size = lambda: 1
+    dist_mod.get_cp_group = lambda: None
+    dist_mod.get_cp_rank = lambda: 0
+    dist_mod.get_tp_rank = lambda: 0
+    dist_mod.get_pp_rank = lambda: 0
+
+    # inference.infra.parallelism.*
+    # At cp_world_size=1, scatter/gather are trivially no-ops.
+    import inference.infra.parallelism.gather_scatter_primitive as gs
+
+    def _scatter_noop(x, cp_split_sizes, group=None):
+        return x
+
+    def _gather_noop(x, cp_split_sizes, group=None):
+        return x
+
+    gs.scatter_to_context_parallel_region = _scatter_noop
+    gs.gather_from_context_parallel_region = _gather_noop
+
+    # Re-import ulysses_scheduler with patched scatter/gather in place.
+    import inference.infra.parallelism.ulysses_scheduler as us
+    us.scatter_to_context_parallel_region = _scatter_noop
+    us.gather_from_context_parallel_region = _gather_noop
+    us.get_cp_world_size = lambda: 1
+    us.get_cp_group = lambda: None
+
+    # all-to-all primitives used by flash_attn_with_cp. At cp=1 they are
+    # entered only as a no-op path (the `if cp_world_size > 1` branch is
+    # skipped), so no stubs needed there.
+
+
+def install_stubs() -> None:
+    """Install all stubs. Idempotent."""
+    _install_magi_compiler_stub()
+    repo_root = Path(__file__).resolve().parents[3]
+    upstream = repo_root / "daVinci-MagiHuman"
+    path_s = str(upstream)
+    if path_s not in sys.path:
+        sys.path.insert(0, path_s)
+    # Reload inference.* after sys.path mutation so it picks up the real
+    # upstream package (not a stale one).
+    for name in list(sys.modules):
+        if name == "inference" or name.startswith("inference."):
+            del sys.modules[name]
+    import inference  # noqa: F401
+    _install_distributed_stubs()
+
+
+# ---------------------------------------------------------------------------
+# Upstream DiTModel loader — instantiate + load base shards.
+# ---------------------------------------------------------------------------
+
+
+def _base_arch_dict() -> dict:
+    """Return the upstream `ModelConfig`-equivalent dict for the base variant.
+    Matches `inference/common/config.py::ModelConfig` defaults for base.
+    """
+    import torch
+    return dict(
+        num_layers=40,
+        hidden_size=5120,
+        head_dim=128,
+        num_query_groups=8,
+        video_in_channels=48 * 4,
+        audio_in_channels=64,
+        text_in_channels=3584,
+        checkpoint_qk_layernorm_rope=False,
+        params_dtype=torch.float32,
+        tread_config=dict(
+            selection_rate=0.5, start_layer_idx=2, end_layer_idx=25,
+        ),
+        mm_layers=[0, 1, 2, 3, 36, 37, 38, 39],
+        local_attn_layers=[],
+        enable_attn_gating=True,
+        activation_type="swiglu7",
+        gelu7_layers=[0, 1, 2, 3],
+        # derived
+        num_heads_q=40,
+        num_heads_kv=8,
+        post_norm_layers=[],
+    )
+
+
+def load_upstream_dit(base_shard_dir, device=None, dtype=None, local_attn_layers=None):
+    """Instantiate upstream `DiTModel` and load the base shards into it.
+
+    Args:
+        base_shard_dir: path to `base/` (contains `model-0000*-of-00007.safetensors`
+            and `model.safetensors.index.json`).
+        device: torch device (default cuda if available).
+        dtype: dtype cast (default: leave checkpoint dtypes as-is).
+
+    Returns:
+        An upstream `DiTModel` in `.eval()` mode with weights loaded.
+    """
+    import glob
+    import json
+    import types as _types
+
+    import torch
+    from safetensors.torch import load_file
+
+    from inference.common.config import ModelConfig  # upstream pydantic class
+    from inference.model.dit.dit_module import DiTModel
+
+    arch_dict = _base_arch_dict()
+    if local_attn_layers is not None:
+        arch_dict["local_attn_layers"] = list(local_attn_layers)
+    # ModelConfig is a pydantic BaseModel — build via kwargs.
+    model_config = ModelConfig(**arch_dict)
+
+    model = DiTModel(model_config=model_config)
+
+    # Load all base shards into a single state dict.
+    base_shard_dir = Path(base_shard_dir)
+    shard_paths = sorted(base_shard_dir.glob("*.safetensors"))
+    state = {}
+    for p in shard_paths:
+        state.update(load_file(str(p)))
+
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    if missing:
+        raise RuntimeError(f"Upstream DiT missing {len(missing)} keys: {missing[:5]}")
+    if unexpected:
+        raise RuntimeError(f"Upstream DiT unexpected {len(unexpected)} keys: {unexpected[:5]}")
+
+    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device=device)
+    if dtype is not None:
+        model = model.to(dtype=dtype)
+    model.eval()
+    return model

tests/local_tests/magi_human/test_magi_human_distill_parity.py

@@ -0,0 +1,256 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DiT parity test for the daVinci-MagiHuman DMD-2 distilled checkpoint.
+
+The distill variant has the SAME architecture as the base model (same 40
+layers, same hidden_size, same mm_layers / gelu7_layers / local_attn_layers,
+same head_dim and num_query_groups; see
+`daVinci-MagiHuman/inference/common/config.py:ModelConfig`). Only the
+weights differ: distill is trained for 8-step DMD-2 inference without CFG.
+
+This test mirrors `test_magi_human_parity.py::test_magi_human_dit_parity`
+exactly, just pointing at the `distill/` subfolder of GAIR/daVinci-MagiHuman
+and the matching `converted_weights/magi_human_distill/`.
+
+Skips cleanly when:
+  * `daVinci-MagiHuman/` clone is absent
+  * GAIR/daVinci-MagiHuman distill shards are not locally available
+  * Converted distill weights have not been produced yet
+  * CUDA is unavailable
+"""
+from __future__ import annotations
+
+import gc
+import glob
+import os
+from pathlib import Path
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.forward_context import set_forward_context
+
+
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+
+def _find_distill_shard_dir() -> Path | None:
+    """Return the local path to GAIR/daVinci-MagiHuman/distill/ shards or None."""
+    override = os.getenv("MAGI_HUMAN_DISTILL_SHARD_DIR")
+    if override:
+        p = Path(override)
+        return p if p.is_dir() else None
+    try:
+        from huggingface_hub import snapshot_download
+        snap = snapshot_download(
+            repo_id="GAIR/daVinci-MagiHuman",
+            allow_patterns=[
+                "distill/*.safetensors",
+                "distill/model.safetensors.index.json",
+            ],
+        )
+        candidate = Path(snap) / "distill"
+        if candidate.is_dir() and any(candidate.glob("*.safetensors")):
+            return candidate
+        return None
+    except Exception:
+        return None
+
+
+def _cleanup_gpu() -> None:
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="MagiHuman distill DiT parity requires CUDA.",
+)
+def test_magi_human_distill_dit_parity():
+    repo_root = Path(__file__).resolve().parents[3]
+    upstream_src = repo_root / "daVinci-MagiHuman"
+    if not upstream_src.exists():
+        pytest.skip(
+            "Upstream daVinci-MagiHuman/ clone missing. Run "
+            "`git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git`"
+        )
+
+    distill_shard_dir = _find_distill_shard_dir()
+    if distill_shard_dir is None or not distill_shard_dir.is_dir():
+        pytest.skip(
+            "GAIR/daVinci-MagiHuman distill/ shards not available locally. "
+            "Set MAGI_HUMAN_DISTILL_SHARD_DIR or run the conversion once to "
+            "populate the HF cache."
+        )
+
+    converted_dir = Path(os.getenv(
+        "MAGI_HUMAN_DISTILL_DIFFUSERS_PATH",
+        repo_root / "converted_weights" / "magi_human_distill",
+    ))
+    transformer_dir = converted_dir / "transformer"
+    if not transformer_dir.is_dir():
+        pytest.skip(
+            f"Converted distill transformer dir missing at {transformer_dir}. Run "
+            f"scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py "
+            f"--subfolder distill --cast-bf16 first."
+        )
+
+    from tests.local_tests.helpers.magi_human_upstream import (
+        install_stubs,
+        load_upstream_dit,
+    )
+    install_stubs()
+
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+
+    z_dim = 48
+    pT, pH, pW = 1, 2, 2
+    lat_T, lat_H, lat_W = 2, 6, 6
+    video_latent = torch.randn(
+        (1, z_dim, lat_T, lat_H, lat_W),
+        dtype=torch.float32, device=device,
+    )
+    num_video_tokens = (lat_T // pT) * (lat_H // pH) * (lat_W // pW)
+    num_audio_tokens = 4
+    num_text_tokens = 8
+    audio_latent = torch.randn(
+        (1, num_audio_tokens, 64),
+        dtype=torch.float32, device=device,
+    )
+    text_feat = torch.randn(
+        (1, num_text_tokens, 3584),
+        dtype=torch.float32, device=device,
+    )
+
+    # The pipeline stages module ships in a later step of the PR #1280
+    # decomposition (4/8). On this branch (3/8 — DiT only) the import fails
+    # cleanly; the test is dead code until the pipeline lands.
+    try:
+        from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
+            build_packed_inputs,
+        )
+    except ImportError:
+        pytest.skip(
+            "fastvideo.pipelines.basic.magi_human.stages.latent_preparation "
+            "is not present on this branch yet (activates in step 4/8 of the "
+            "PR #1280 decomposition). DiT parity will run once the pipeline "
+            "lands."
+        )
+    from fastvideo.models.dits.magi_human import Modality  # noqa: F401
+
+    x, coords, mm = build_packed_inputs(
+        video_latent=video_latent,
+        audio_latent=audio_latent,
+        audio_feat_len=num_audio_tokens,
+        txt_feat=text_feat,
+        txt_feat_len=num_text_tokens,
+        patch_size=(pT, pH, pW),
+        coords_style="v2",
+    )
+    assert x.shape[0] == num_video_tokens + num_audio_tokens + num_text_tokens
+
+    total_tokens = x.shape[0]
+
+    # Distill arch is identical to base; load_upstream_dit's _base_arch_dict
+    # describes both since they share num_layers / hidden_size / mm_layers
+    # / etc. Just point at the distill shards.
+    print("Loading upstream distill DiTModel from distill shards...")
+    upstream_model = load_upstream_dit(
+        distill_shard_dir,
+        device=device,
+        dtype=None,
+    )
+
+    from inference.common import VarlenHandler
+    cu = torch.tensor([0, total_tokens], dtype=torch.int32, device=device)
+    varlen = VarlenHandler(
+        cu_seqlens_q=cu,
+        cu_seqlens_k=cu,
+        max_seqlen_q=total_tokens,
+        max_seqlen_k=total_tokens,
+    )
+
+    print("Running upstream distill forward...")
+    with torch.inference_mode():
+        ref_out = upstream_model(
+            x=x.clone(),
+            coords_mapping=coords.clone(),
+            modality_mapping=mm.clone(),
+            varlen_handler=varlen,
+            local_attn_handler=None,
+        ).detach().float().cpu()
+
+    del upstream_model
+    _cleanup_gpu()
+
+    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
+    from fastvideo.models.dits.magi_human import MagiHumanDiT
+    from safetensors.torch import load_file
+    print("Loading FastVideo MagiHumanDiT from converted distill transformer/...")
+    fv_cfg = MagiHumanVideoConfig()
+    fv_model = MagiHumanDiT(fv_cfg)
+
+    fv_state = {}
+    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
+        fv_state.update(load_file(shard))
+    missing, unexpected = fv_model.load_state_dict(fv_state, strict=False)
+    assert not missing, f"FastVideo distill DiT missing {len(missing)} keys: {missing[:5]}"
+    assert not unexpected, f"FastVideo distill DiT unexpected {len(unexpected)} keys: {unexpected[:5]}"
+
+    fv_model = fv_model.to(device=device)
+    fv_model.eval()
+
+    print("Running FastVideo distill forward...")
+    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
+        fv_out = fv_model(x.clone(), coords.clone(), mm.clone()).detach().float().cpu()
+
+    print(
+        f"ref sum={ref_out.sum().item():.4f} "
+        f"abs_mean={ref_out.abs().mean().item():.4f} "
+        f"shape={tuple(ref_out.shape)}"
+    )
+    print(
+        f"fv  sum={fv_out.sum().item():.4f} "
+        f"abs_mean={fv_out.abs().mean().item():.4f} "
+        f"shape={tuple(fv_out.shape)}"
+    )
+    diff = (ref_out - fv_out).abs()
+    print(
+        f"diff max={diff.max().item():.6f} "
+        f"mean={diff.mean().item():.6f} "
+        f"median={diff.median().item():.6f}"
+    )
+
+    ref_video = ref_out[:num_video_tokens]
+    fv_video = fv_out[:num_video_tokens]
+    ref_audio = ref_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64]
+    fv_audio = fv_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64]
+    ref_text = ref_out[num_video_tokens + num_audio_tokens:]
+    fv_text = fv_out[num_video_tokens + num_audio_tokens:]
+    video_diff = (ref_video - fv_video).abs()
+    audio_diff = (ref_audio - fv_audio).abs()
+    text_diff = (ref_text - fv_text).abs()
+    print(
+        f"video  ref_abs={ref_video.abs().mean():.4f} "
+        f"diff_max={video_diff.max():.4f} diff_mean={video_diff.mean():.4f}"
+    )
+    print(
+        f"audio  ref_abs={ref_audio.abs().mean():.4f} "
+        f"diff_max={audio_diff.max():.4f} diff_mean={audio_diff.mean():.4f}"
+    )
+    print(
+        f"text   ref_abs={ref_text.abs().mean():.4f} "
+        f"diff_max={text_diff.max():.4f} diff_mean={text_diff.mean():.4f}"
+    )
+
+    assert ref_out.shape == fv_out.shape
+    assert_close(fv_text, ref_text, atol=1e-6, rtol=1e-6)
+    # Same tolerance as base DiT parity (post-Wave 14b dtype-boundary fixes,
+    # both DiTs are bit-exact via shared upstream BaseLinear bf16 path).
+    assert_close(fv_out, ref_out, atol=0.03, rtol=0.01)
+    ref_abs = ref_out.abs().mean().item()
+    fv_abs = fv_out.abs().mean().item()
+    rel = abs(ref_abs - fv_abs) / max(ref_abs, 1e-6)
+    assert rel < 0.05, f"abs_mean drift {rel:.3%} > 5%"

tests/local_tests/magi_human/test_magi_human_parity.py

@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical parity test: FastVideo MagiHumanDiT vs upstream DiTModel.
+
+Loads both models from the **same converted base checkpoint** and runs them
+on identical small inputs. Asserts closeness on the joint video+audio
+output tensor.
+
+What this catches (that the preflight test does NOT):
+  - Silent weight-name mismatches that `strict=False` loading would hide.
+  - Wrong modality-expert chunking inside `PackedExpertLinear`.
+  - RoPE sin/cos ordering flipped.
+  - Per-head gating dtype / split order.
+  - swiglu7 / gelu7 off-by-one on the `+1` linear bias.
+
+Skips cleanly when:
+  - `daVinci-MagiHuman/` clone is absent (no upstream source).
+  - GAIR/daVinci-MagiHuman base shards are not available locally.
+  - CUDA is unavailable.
+
+Tolerance: `atol=5e-3, rtol=5e-3` on bf16 forward paths. The FastVideo
+attention path uses `F.scaled_dot_product_attention` while upstream uses
+`flash_attn_func`; both accumulate in bf16 but via different kernels, so
+small drift is expected and bounded.
+"""
+from __future__ import annotations
+
+import gc
+import os
+from pathlib import Path
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.forward_context import set_forward_context
+
+
+# Force TORCH_SDPA for FastVideo so the attention kernel is deterministic.
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
+
+
+def _find_base_shard_dir() -> Path | None:
+    """Return the local path to GAIR/daVinci-MagiHuman/base/ with shards present, or None."""
+    override = os.getenv("MAGI_HUMAN_BASE_SHARD_DIR")
+    if override:
+        p = Path(override)
+        return p if p.is_dir() else None
+    try:
+        from huggingface_hub import snapshot_download
+        snap = snapshot_download(
+            repo_id="GAIR/daVinci-MagiHuman",
+            allow_patterns=[
+                "base/*.safetensors",
+                "base/model.safetensors.index.json",
+            ],
+        )
+        candidate = Path(snap) / "base"
+        if candidate.is_dir() and any(candidate.glob("*.safetensors")):
+            return candidate
+        return None
+    except Exception:
+        return None
+
+
+def _cleanup_gpu() -> None:
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="MagiHuman DiT parity requires CUDA.",
+)
+def test_magi_human_dit_parity():
+    repo_root = Path(__file__).resolve().parents[3]
+    upstream_src = repo_root / "daVinci-MagiHuman"
+    if not upstream_src.exists():
+        pytest.skip(
+            "Upstream daVinci-MagiHuman/ clone missing. Run "
+            "`git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git`"
+        )
+
+    base_shard_dir = _find_base_shard_dir()
+    if base_shard_dir is None or not base_shard_dir.is_dir():
+        pytest.skip(
+            "GAIR/daVinci-MagiHuman base/ shards not available locally. "
+            "Set MAGI_HUMAN_BASE_SHARD_DIR or run the conversion once to "
+            "populate the HF cache."
+        )
+
+    converted_dir = Path(os.getenv(
+        "MAGI_HUMAN_DIFFUSERS_PATH",
+        repo_root / "converted_weights" / "magi_human_base",
+    ))
+    transformer_dir = converted_dir / "transformer"
+    if not transformer_dir.is_dir():
+        pytest.skip(
+            f"Converted transformer dir missing at {transformer_dir}. Run "
+            f"scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py first."
+        )
+
+    # Add upstream to sys.path and install compiler/distributed stubs.
+    from tests.local_tests.helpers.magi_human_upstream import (
+        install_stubs,
+        load_upstream_dit,
+    )
+    install_stubs()
+
+    # --- Shared inputs (deliberately small) ---
+    device = torch.device("cuda:0")
+    torch.manual_seed(0)
+
+    # Mirror MagiDataProxy.process_input for a tiny frame:
+    # video_latent: [1, z_dim, T, H, W], T=2, H=6, W=6, z_dim=48
+    # -> video tokens: (T/pT)*(H/pH)*(W/pW) with patch=(1,2,2) = 2*3*3 = 18
+    # audio tokens: 4
+    # text tokens: 8
+    # max channel width = 192 (video)
+    z_dim = 48
+    pT, pH, pW = 1, 2, 2
+    lat_T, lat_H, lat_W = 2, 6, 6
+    video_latent = torch.randn(
+        (1, z_dim, lat_T, lat_H, lat_W),
+        dtype=torch.float32, device=device,
+    )
+    num_video_tokens = (lat_T // pT) * (lat_H // pH) * (lat_W // pW)  # 18
+    num_audio_tokens = 4
+    num_text_tokens = 8
+    audio_latent = torch.randn(
+        (1, num_audio_tokens, 64),
+        dtype=torch.float32, device=device,
+    )
+    text_feat = torch.randn(
+        (1, num_text_tokens, 3584),
+        dtype=torch.float32, device=device,
+    )
+
+    # --- Build the packed inputs the DiT consumes ---
+    # The pipeline stages module ships in a later step of the PR #1280
+    # decomposition (4/8). On this branch (3/8 — DiT only) the import fails
+    # cleanly; the test is dead code until the pipeline lands.
+    try:
+        from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
+            build_packed_inputs,
+        )
+    except ImportError:
+        pytest.skip(
+            "fastvideo.pipelines.basic.magi_human.stages.latent_preparation "
+            "is not present on this branch yet (activates in step 4/8 of the "
+            "PR #1280 decomposition). DiT parity will run once the pipeline "
+            "lands."
+        )
+    from fastvideo.models.dits.magi_human import Modality  # noqa: F401
+
+    x, coords, mm = build_packed_inputs(
+        video_latent=video_latent,
+        audio_latent=audio_latent,
+        audio_feat_len=num_audio_tokens,
+        txt_feat=text_feat,
+        txt_feat_len=num_text_tokens,
+        patch_size=(pT, pH, pW),
+        coords_style="v2",
+    )
+    assert x.shape[0] == num_video_tokens + num_audio_tokens + num_text_tokens
+
+    total_tokens = x.shape[0]
+
+    # --- Load upstream DiT first (so we know weights round-trip cleanly).
+    #     Upstream reads raw base/ shards; we keep it in bf16 for speed
+    #     and because that matches the FastVideo side after FSDP load.
+    print("Loading upstream DiTModel from base shards...")
+    upstream_model = load_upstream_dit(
+        base_shard_dir,
+        device=device,
+        dtype=None,  # keep checkpoint dtypes (fp32 for norms, bf16 for matmuls)
+    )
+
+    # --- VarlenHandler for upstream (batch=1, total_tokens).
+    from inference.common import VarlenHandler
+    cu = torch.tensor([0, total_tokens], dtype=torch.int32, device=device)
+    varlen = VarlenHandler(
+        cu_seqlens_q=cu,
+        cu_seqlens_k=cu,
+        max_seqlen_q=total_tokens,
+        max_seqlen_k=total_tokens,
+    )
+
+    # --- Forward upstream and capture output. ---
+    print("Running upstream forward...")
+    with torch.inference_mode():
+        ref_out = upstream_model(
+            x=x.clone(),
+            coords_mapping=coords.clone(),
+            modality_mapping=mm.clone(),
+            varlen_handler=varlen,
+            local_attn_handler=None,  # local_attn_layers=[] for base
+        ).detach().float().cpu()
+
+    # Free upstream model before loading FastVideo (saves ~30 GB on GPU).
+    del upstream_model
+    _cleanup_gpu()
+
+    # --- Load FastVideo MagiHumanDiT from the converted transformer/ ---
+    from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
+    from fastvideo.models.dits.magi_human import MagiHumanDiT
+    from safetensors.torch import load_file
+    import glob
+    print("Loading FastVideo MagiHumanDiT from converted transformer/...")
+    fv_cfg = MagiHumanVideoConfig()
+    fv_model = MagiHumanDiT(fv_cfg)
+
+    fv_state = {}
+    for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))):
+        fv_state.update(load_file(shard))
+    missing, unexpected = fv_model.load_state_dict(fv_state, strict=False)
+    assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}"
+    assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}"
+
+    fv_model = fv_model.to(device=device)
+    fv_model.eval()
+
+    print("Running FastVideo forward...")
+    with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
+        fv_out = fv_model(x.clone(), coords.clone(), mm.clone()).detach().float().cpu()
+
+    # Global stats
+    print(
+        f"ref sum={ref_out.sum().item():.4f} "
+        f"abs_mean={ref_out.abs().mean().item():.4f} "
+        f"shape={tuple(ref_out.shape)}"
+    )
+    print(
+        f"fv  sum={fv_out.sum().item():.4f} "
+        f"abs_mean={fv_out.abs().mean().item():.4f} "
+        f"shape={tuple(fv_out.shape)}"
+    )
+    diff = (ref_out - fv_out).abs()
+    print(
+        f"diff max={diff.max().item():.6f} "
+        f"mean={diff.mean().item():.6f} "
+        f"median={diff.median().item():.6f}"
+    )
+
+    # Per-modality diagnostic (video, audio, text). Text rows are zero-
+    # padded on both sides; video and audio should carry comparable
+    # abs_mean.
+    ref_video = ref_out[:num_video_tokens]
+    fv_video = fv_out[:num_video_tokens]
+    ref_audio = ref_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64]
+    fv_audio = fv_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64]
+    ref_text = ref_out[num_video_tokens + num_audio_tokens:]
+    fv_text = fv_out[num_video_tokens + num_audio_tokens:]
+    video_diff = (ref_video - fv_video).abs()
+    audio_diff = (ref_audio - fv_audio).abs()
+    text_diff = (ref_text - fv_text).abs()
+    print(
+        f"video  ref_abs={ref_video.abs().mean():.4f} "
+        f"diff_max={video_diff.max():.4f} diff_mean={video_diff.mean():.4f}"
+    )
+    print(
+        f"audio  ref_abs={ref_audio.abs().mean():.4f} "
+        f"diff_max={audio_diff.max():.4f} diff_mean={audio_diff.mean():.4f}"
+    )
+    print(
+        f"text   ref_abs={ref_text.abs().mean():.4f} "
+        f"diff_max={text_diff.max():.4f} diff_mean={text_diff.mean():.4f}"
+    )
+
+    # --- Assertions ---
+    assert ref_out.shape == fv_out.shape, (
+        f"shape mismatch: ref={ref_out.shape} fv={fv_out.shape}"
+    )
+
+    # Text rows are zero-padded on both sides — must match exactly.
+    assert_close(fv_text, ref_text, atol=1e-6, rtol=1e-6)
+
+    # Video + audio: bf16 single-forward DiT noise floor is ~1e-3 to
+    # 5e-3 per element. atol=0.03 catches gross structural bugs
+    # (permutation flips, sign inversions, wrong modality dispatch,
+    # missing sub-layers) while leaving 6-10x margin over actual bf16
+    # noise. Observed diff_max=0.057 will FAIL — that is the bug
+    # surfacing and is the intended spec for downstream root-cause
+    # investigation.
+    assert_close(fv_out, ref_out, atol=0.03, rtol=0.01)
+
+    # Sanity: mean magnitudes should match within 5%. A gross bug
+    # (e.g. dropping a modality branch) would show up here.
+    ref_abs = ref_out.abs().mean().item()
+    fv_abs = fv_out.abs().mean().item()
+    rel = abs(ref_abs - fv_abs) / max(ref_abs, 1e-6)
+    assert rel < 0.05, f"abs_mean drift {rel:.3%} > 5% — possible structural bug"