magi_human

class Modality(IntEnum):
    VIDEO = 0
    AUDIO = 1
    TEXT = 2

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)

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&#91;inv_permute_mapping&#93;</code></pre></details>
fastvideo.models.dits.magi_human.ModalityDispatcher.dispatch
dispatch(x: torch.Tensor) -> list[torch.Tensor]
Source
def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]:
    return list(torch.split(x, self.group_size_cpu, dim=0))
fastvideo.models.dits.magi_human.ModalityDispatcher.undispatch
undispatch(*chunks: torch.Tensor) -> torch.Tensor
Source
def undispatch(self, *chunks: torch.Tensor) -> torch.Tensor:
    return torch.cat(chunks, dim=0)
fastvideo.models.dits.magi_human.ModalityDispatcher.permute
permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor
Source
@staticmethod
def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor:
    return x[permute_mapping]
fastvideo.models.dits.magi_human.ModalityDispatcher.inv_permute
inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor
Source
@staticmethod
def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor:
    return x[inv_permute_mapping]
fastvideo.models.dits.magi_human.MultiModalityRMSNorm
MultiModalityRMSNorm(dim: int, eps: float = 1e-06, num_modality: int = 1)
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.
Source
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 &gt; 1,
the weight tensor packs per-modality scales along its flat axis and the
dispatcher selects the right chunk per modality.
&quot;&quot;&quot;

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) -&gt; 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&#91;ModalityDispatcher&#93; = None,
) -&gt; 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, (&quot;MultiModalityRMSNorm with num_modality&gt;1 requires a dispatcher&quot;)
    weight_chunks = self.weight.chunk(self.num_modality, dim=0)
    parts = modality_dispatcher.dispatch(t)
    for i in range(self.num_modality):
        parts&#91;i&#93; = parts&#91;i&#93; * (weight_chunks&#91;i&#93; + 1)
    return modality_dispatcher.undispatch(*parts).to(original_dtype)</code></pre></details>
fastvideo.models.dits.magi_human.MultiModalityRMSNorm.forward
forward(x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor
Source
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)
fastvideo.models.dits.magi_human.ElementWiseFourierEmbed
ElementWiseFourierEmbed(dim: int, temperature: float = 10000.0, dtype: torch.dtype = torch.float32)
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).
Source
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 `&#91;dim // 8&#93;` (fixed at init via freq_bands).
&quot;&quot;&quot;

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(&quot;bands&quot;, bands)

def forward(self, coords: torch.Tensor) -&gt; torch.Tensor:
    # coords: &#91;L, 9&#93; = (t, h, w, T, H, W, ref_T, ref_H, ref_W)
    coords_xyz = coords&#91;:, :3&#93;
    sizes = coords&#91;:, 3:6&#93;
    refs = coords&#91;:, 6:9&#93;

    scales = (refs - 1) / (sizes - 1)
    scales&#91;(refs == 1) &amp; (sizes == 1)&#93; = 1
    # Center H and W (leave time uncentered).
    centers = (sizes - 1) / 2
    centers&#91;:, 0&#93; = 0
    coords_xyz = coords_xyz - centers

    proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands  # &#91;L, 3, B&#93;
    sin_proj = proj.sin()
    cos_proj = proj.cos()
    return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)</code></pre></details>
fastvideo.models.dits.magi_human.ElementWiseFourierEmbed.forward
forward(coords: torch.Tensor) -> torch.Tensor
Source
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&#91;(refs == 1) &amp; (sizes == 1)&#93; = 1
# Center H and W (leave time uncentered).
centers = (sizes - 1) / 2
centers&#91;:, 0&#93; = 0
coords_xyz = coords_xyz - centers

proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands  # &#91;L, 3, B&#93;
sin_proj = proj.sin()
cos_proj = proj.cos()
return torch.cat((sin_proj, cos_proj), dim=1).flatten(1)</code></pre></details>
fastvideo.models.dits.magi_human.PackedExpertLinear
PackedExpertLinear(in_features: int, out_features: int, num_experts: int = 1, bias: bool = False, dtype: torch.dtype = torch.bfloat16)
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.
Source
class PackedExpertLinear(nn.Module):
    """Linear where the weight is packed per-modality along the output axis.
Shapes:
    weight: &#91;out_features * num_experts, in_features&#93;
    bias:   &#91;out_features * num_experts&#93;  (optional)

When `num_experts == 1`, behaves exactly like `nn.Linear`. When
`num_experts &gt; 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&#x27;t match this layout.
&quot;&quot;&quot;

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.empty(out_features * num_experts, in_features, dtype=dtype))
    if bias:
        self.bias = nn.Parameter(torch.empty(out_features * num_experts, dtype=dtype))
    else:
        self.register_parameter(&quot;bias&quot;, None)

def forward(
    self,
    x: torch.Tensor,
    modality_dispatcher: Optional&#91;ModalityDispatcher&#93; = None,
) -&gt; torch.Tensor:
    if self.num_experts == 1:
        return F.linear(x, self.weight, self.bias)
    assert modality_dispatcher is not None, (&quot;PackedExpertLinear with num_experts&gt;1 requires a dispatcher&quot;)
    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 &#91;None&#93; * self.num_experts)
    for i in range(self.num_experts):
        parts&#91;i&#93; = F.linear(parts&#91;i&#93;, w_chunks&#91;i&#93;, b_chunks&#91;i&#93;)
    return modality_dispatcher.undispatch(*parts)</code></pre></details>
fastvideo.models.dits.magi_human.PackedExpertLinear.forward
forward(x: torch.Tensor, modality_dispatcher: Optional[ModalityDispatcher] = None) -> torch.Tensor
Source
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)
fastvideo.models.dits.magi_human.AttentionSubConfig
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)
Source
@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
fastvideo.models.dits.magi_human.MagiAttention
MagiAttention(cfg: AttentionSubConfig)
Self-attention with GQA + optional per-head sigmoid gating.
Source
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,
) -&gt; 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) -&gt; torch.Tensor:
    &quot;&quot;&quot;Run SDPA on &#91;L, H, D&#93; tensors and return &#91;L, Hq, D&#93;.&quot;&quot;&quot;
    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,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Approximate upstream FFAHandler block accumulation with SDPA.

    SR-1080p&#x27;s reference kernel sums three independently-normalized
    attention contributions:

    * video frame queries -&gt; local-window video keys;
    * all video queries -&gt; all audio+text keys;
    * all audio+text queries -&gt; 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.
    &quot;&quot;&quot;
    if num_frames &lt;= 0 or num_video_tokens &lt;= 0:
        return self._sdpa(q, k, v)
    if num_video_tokens % num_frames != 0:
        raise ValueError(f&quot;MagiHuman local attention expects video tokens divisible by &quot;
                         f&quot;frames, got &#123;num_video_tokens=&#125; and &#123;num_frames=&#125;.&quot;)

    token_per_frame = num_video_tokens // num_frames
    out = torch.zeros(
        q.shape&#91;0&#93;,
        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&#91;:num_video_tokens&#93;
    k_video = k&#91;:num_video_tokens&#93;
    v_video = v&#91;:num_video_tokens&#93;
    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&#91;q_start:q_end&#93; = self._sdpa(
            q_video&#91;q_start:q_end&#93;,
            k_video&#91;k_start:k_end&#93;,
            v_video&#91;k_start:k_end&#93;,
        )

    if num_video_tokens &lt; q.shape&#91;0&#93;:
        k_at = k&#91;num_video_tokens:&#93;
        v_at = v&#91;num_video_tokens:&#93;
        out&#91;:num_video_tokens&#93; = out&#91;:num_video_tokens&#93; + self._sdpa(
            q&#91;:num_video_tokens&#93;,
            k_at,
            v_at,
        )
        out&#91;num_video_tokens:&#93; = self._sdpa(
            q&#91;num_video_tokens:&#93;,
            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,
) -&gt; 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,
        &#91;self.q_size, self.kv_size, self.kv_size, self.gating_size&#93;,
        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&#x27;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 &lt; 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&#91;-1&#93; * 2
    q_rot = _apply_rotary_emb(q&#91;..., :rot_dim&#93;, cos_emb, sin_emb, is_neox_style=True)
    k_rot = _apply_rotary_emb(k&#91;..., :rot_dim&#93;, cos_emb, sin_emb, is_neox_style=True)
    if rot_dim &lt; q.shape&#91;-1&#93;:
        q = torch.cat(&#91;q_rot, q&#91;..., rot_dim:&#93;&#93;, dim=-1)
        k = torch.cat(&#91;k_rot, k&#91;..., rot_dim:&#93;&#93;, dim=-1)
    else:
        q, k = q_rot, k_rot

    # Run SDPA via FastVideo&#x27;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&#x27;s default is bf16, matching
    # upstream BaseLinear at dit_module.py:330). The gating multiply
    # promotes back to fp32 implicitly via PyTorch&#x27;s type-promotion
    # rules: bf16_attn_out * sigmoid(fp32_g) -&gt; 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(&quot;MagiHuman local attention requires video token/frame metadata.&quot;)
        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)</code></pre></details>
fastvideo.models.dits.magi_human.MagiAttention.configure_local_attention
configure_local_attention(*, enabled: bool, frame_receptive_field: int = 11) -> None
Source
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
fastvideo.models.dits.magi_human.MagiAttention.forward
forward(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
Source
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&#x27;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 &lt; 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&#91;-1&#93; * 2
q_rot = _apply_rotary_emb(q&#91;..., :rot_dim&#93;, cos_emb, sin_emb, is_neox_style=True)
k_rot = _apply_rotary_emb(k&#91;..., :rot_dim&#93;, cos_emb, sin_emb, is_neox_style=True)
if rot_dim &lt; q.shape&#91;-1&#93;:
    q = torch.cat(&#91;q_rot, q&#91;..., rot_dim:&#93;&#93;, dim=-1)
    k = torch.cat(&#91;k_rot, k&#91;..., rot_dim:&#93;&#93;, dim=-1)
else:
    q, k = q_rot, k_rot

# Run SDPA via FastVideo&#x27;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&#x27;s default is bf16, matching
# upstream BaseLinear at dit_module.py:330). The gating multiply
# promotes back to fp32 implicitly via PyTorch&#x27;s type-promotion
# rules: bf16_attn_out * sigmoid(fp32_g) -&gt; 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(&quot;MagiHuman local attention requires video token/frame metadata.&quot;)
    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)</code></pre></details>
fastvideo.models.dits.magi_human.MLPSubConfig
MLPSubConfig(hidden_size: int, intermediate_size: int, activation: str, num_modality: int, gated: bool)
Source
@dataclass
class MLPSubConfig:
    hidden_size: int
    intermediate_size: int
    activation: str  # "swiglu7" or "gelu7"
    num_modality: int
    gated: bool
fastvideo.models.dits.magi_human.MagiMLP
MagiMLP(cfg: MLPSubConfig)
Source
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 == &quot;swiglu7&quot; else gelu7

def forward(
    self,
    x: torch.Tensor,
    modality_dispatcher: ModalityDispatcher,
) -&gt; 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</code></pre></details>
fastvideo.models.dits.magi_human.MagiMLP.forward
forward(x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor
Source
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
fastvideo.models.dits.magi_human.MagiTransformerLayer
MagiTransformerLayer(arch: MagiHumanArchConfig, layer_idx: int)
Source
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 = &quot;gelu7&quot;
    else:
        intermediate = (arch.hidden_size * 4 * 2 // 3) // 4 * 4
        gated = True
        activation = &quot;swiglu7&quot;

    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,
) -&gt; 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</code></pre></details>
fastvideo.models.dits.magi_human.MagiTransformerLayer.forward
forward(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
Source
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</code></pre></details>
fastvideo.models.dits.magi_human.MagiAdapter
MagiAdapter(arch: MagiHumanArchConfig)
Source
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, so data_ptr()+shape+dtype+device is a fast,
    # collision-free key that avoids recomputing the Fourier embed each step.
    self._cached_rope: Optional&#91;torch.Tensor&#93; = None
    self._cached_rope_key: Optional&#91;tuple&#93; = None

def _rope_cache_key(self, t: torch.Tensor) -&gt; tuple:
    return (t.data_ptr(), 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,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor&#93;:
    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
    rope = self._cached_rope
    # Embedder dtypes may differ from x&#x27;s dtype when FastVideo&#x27;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&#91;0&#93;,
        self.video_embedder.out_features,
        device=x.device,
        dtype=v_w.dtype,
    )
    out&#91;text_mask&#93; = self.text_embedder(x&#91;text_mask, :self.text_embedder.in_features&#93;.to(t_w.dtype)).to(out.dtype)
    out&#91;audio_mask&#93; = self.audio_embedder(x&#91;audio_mask, :self.audio_embedder.in_features&#93;.to(a_w.dtype)).to(
        out.dtype)
    out&#91;video_mask&#93; = self.video_embedder(x&#91;video_mask, :self.video_embedder.in_features&#93;.to(v_w.dtype)).to(
        out.dtype)
    return out, rope</code></pre></details>
fastvideo.models.dits.magi_human.MagiAdapter.forward
forward(x: torch.Tensor, coords_mapping: torch.Tensor, video_mask: torch.Tensor, audio_mask: torch.Tensor, text_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]
Source
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
    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
fastvideo.models.dits.magi_human.MagiHumanDiT
MagiHumanDiT(config: MagiHumanVideoConfig, hf_config: dict | None = None, **kwargs)
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.
Source
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.
&quot;&quot;&quot;

# 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 &#123;&#125;)
    arch: MagiHumanArchConfig = getattr(config, &quot;arch_config&quot;, 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; data_ptr()+shape+dtype+device
    # is a fast, collision-free key that avoids rebuilding ModalityDispatcher
    # (which calls argsort + bincount) on every forward call.
    self._cached_dispatcher: Optional&#91;ModalityDispatcher&#93; = None
    self._cached_video_mask: Optional&#91;torch.Tensor&#93; = None
    self._cached_audio_mask: Optional&#91;torch.Tensor&#93; = None
    self._cached_text_mask: Optional&#91;torch.Tensor&#93; = None
    self._cached_modality_key: Optional&#91;tuple&#93; = None

def configure_local_attention(
    self,
    local_attn_layers: tuple&#91;int, ...&#93; | list&#91;int&#93;,
    frame_receptive_field: int = 11,
) -&gt; None:
    layers = tuple(int(layer) for layer in local_attn_layers)
    self.arch.local_attn_layers = layers
    self.block.configure_local_attention(layers, frame_receptive_field)

def _modality_cache_key(self, t: torch.Tensor) -&gt; tuple:
    return (t.data_ptr(), t.shape, t.dtype, t.device)

def forward(
    self,
    x: torch.Tensor,
    coords_mapping: torch.Tensor,
    modality_mapping: torch.Tensor,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;
    Args:
        x:                &#91;L, max(V_ch, A_ch, T_ch)&#93;
        coords_mapping:   &#91;L, 9&#93;
        modality_mapping: &#91;L&#93;  (int in &#123;VIDEO, AUDIO, TEXT&#125;)
    Returns:
        out: &#91;L, max(V_ch, A_ch)&#93; with video channels in video slots and
             audio channels in audio slots; text slots are zero.
    &quot;&quot;&quot;
    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_modality_key = key
    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 = int(video_mask.sum().item())
    if num_video_tokens:
        num_frames = int(coords_mapping&#91;:num_video_tokens, 0&#93;.max().item()) + 1
    else:
        num_frames = 0

    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&#x27;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&#91;video_mask&#93;.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&#91;audio_mask&#93;.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&#91;0&#93;, max_ch, device=x.device, dtype=x.dtype)
    out&#91;video_mask, :self.arch.video_in_channels&#93; = x_video.to(out.dtype)
    out&#91;audio_mask, :self.arch.audio_in_channels&#93; = x_audio.to(out.dtype)
    return out</code></pre></details>
fastvideo.models.dits.magi_human.MagiHumanDiT.configure_local_attention
configure_local_attention(local_attn_layers: tuple[int, ...] | list[int], frame_receptive_field: int = 11) -> None
Source
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)
    self.arch.local_attn_layers = layers
    self.block.configure_local_attention(layers, frame_receptive_field)
fastvideo.models.dits.magi_human.MagiHumanDiT.forward
forward(x: torch.Tensor, coords_mapping: torch.Tensor, modality_mapping: torch.Tensor) -> torch.Tensor
Parameters

x: torch.Tensor
[L, max(V_ch, A_ch, T_ch)]
coords_mapping: torch.Tensor
[L, 9]
modality_mapping: torch.Tensor
[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.
Source
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_modality_key = key
    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 = int(video_mask.sum().item())
    if num_video_tokens:
        num_frames = int(coords_mapping[:num_video_tokens, 0].max().item()) + 1
    else:
        num_frames = 0
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&#x27;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&#91;video_mask&#93;.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&#91;audio_mask&#93;.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&#91;0&#93;, max_ch, device=x.device, dtype=x.dtype)
out&#91;video_mask, :self.arch.video_in_channels&#93; = x_video.to(out.dtype)
out&#91;audio_mask, :self.arch.audio_in_channels&#93; = x_audio.to(out.dtype)
return out</code></pre></details>