[Diffusion] MAGI-2: regional compile coverage (#7174)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>
2026-09-16 13:01:26 +08:00
1 parent a2138b830f
commit 21d86ec923
@@ -16,6 +16,7 @@ import os
 from collections.abc import Iterable
 from enum import IntEnum
 from functools import partial
+from typing import Any

 import torch
 import torch.nn as nn
@@ -28,6 +29,7 @@ from .configuration_magi2 import Magi2PreviewConfig
 from .layers import (
    ElementWiseFourierEmbed,
    MHCHandler,
+    MHCTensorTuple,
    ModalityDispatcher,
    MultiModalityRMSNorm,
    make_grouped_linear,
@@ -130,14 +132,12 @@ class Magi2Attention(nn.Module):
        start = self.tp_group.rank * self.num_heads_q
        return checkpoint_tensor[:, start : start + self.num_heads_q]

-    def forward(
+    def project(
        self,
        hidden_states: torch.Tensor,
        rope: torch.Tensor,
-        varlen_handler: VarlenHandler,
        modality_dispatcher: ModalityDispatcher,
-        cp_split_sizes: list[int] | torch.Tensor,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        normalized = self.pre_norm(hidden_states, modality_dispatcher)
        gates = self.linear_g(normalized, modality_dispatcher)
        qkv = self.linear_qkv(normalized, modality_dispatcher)
@@ -159,8 +159,17 @@ class Magi2Attention(nn.Module):
        q = apply_rotary_emb(q, cos, sin).squeeze(0).to(self.config.params_dtype)
        k = apply_rotary_emb(k, cos, sin).squeeze(0).to(self.config.params_dtype)
        v = v.squeeze(0).to(self.config.params_dtype)
+        return q, k, v, gates

-        output = self.packed_attention(
+    def attend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        varlen_handler: VarlenHandler,
+        cp_split_sizes: list[int] | torch.Tensor,
+    ) -> torch.Tensor:
+        return self.packed_attention(
            q,
            k,
            v,
@@ -172,7 +181,14 @@ class Magi2Attention(nn.Module):
                }
            ),
        )
-        output = modality_dispatcher.permute(output)
+
+    def output(
+        self,
+        attention: torch.Tensor,
+        gates: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher,
+    ) -> torch.Tensor:
+        output = modality_dispatcher.permute(attention)
        output = output * torch.sigmoid(gates)
        output = output.reshape(-1, self.q_size).to(self.config.params_dtype)
        return self.linear_proj(output, modality_dispatcher)
@@ -303,12 +319,21 @@ class Magi2MultiHeadMoELayer(nn.Module):
            modality.contiguous(), dispatcher
        )

-    def forward(self, hidden_states: torch.Tensor, dispatcher: ModalityDispatcher) -> torch.Tensor:
+    def route_input(
+        self,
+        hidden_states: torch.Tensor,
+        dispatcher: ModalityDispatcher,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
        normalized = self.pre_norm(hidden_states, dispatcher)
-        routed = self.split_linear(normalized)
-        routed = self.moe_mlp(routed)
-        routed = self.merge_linear(routed)
-        return routed + self._shared_experts(normalized, dispatcher)
+        return normalized, self.split_linear(normalized)
+
+    def combine(
+        self,
+        normalized: torch.Tensor,
+        routed: torch.Tensor,
+        dispatcher: ModalityDispatcher,
+    ) -> torch.Tensor:
+        return self.merge_linear(routed) + self._shared_experts(normalized, dispatcher)


 class Magi2PreAdapter(nn.Module):
@@ -339,37 +364,28 @@ class Magi2PreAdapter(nn.Module):
    def forward(
        self,
        packed: torch.Tensor,
-        coords_mapping: torch.Tensor,
        video_indices: torch.Tensor,
        audio_indices: torch.Tensor,
        text_indices: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        rope = self.rope(coords_mapping)
-        output = torch.zeros(
+    ) -> torch.Tensor:
+        # Every token belongs to exactly one modality, so the three copies
+        # below cover the whole buffer.  Embedding in fp32 and rounding each
+        # modality once matches the released cast at the block boundary.
+        output = torch.empty(
            packed.shape[0],
            self.adapter_dim,
-            dtype=torch.float32,
+            dtype=self.config.params_dtype,
            device=packed.device,
        )
-        if text_indices.numel():
-            output.index_copy_(
-                0,
-                text_indices,
-                self.text_embedder(packed.index_select(0, text_indices)[:, : self.config.text_in_channels].float()),
-            )
-        if audio_indices.numel():
-            output.index_copy_(
-                0,
-                audio_indices,
-                self.audio_embedder(packed.index_select(0, audio_indices)[:, : self.config.audio_in_channels].float()),
-            )
-        if video_indices.numel():
-            output.index_copy_(
-                0,
-                video_indices,
-                self.video_embedder(packed.index_select(0, video_indices)[:, : self.config.video_in_channels].float()),
-            )
-        return output, rope
+        for indices, embedder, in_channels in (
+            (text_indices, self.text_embedder, self.config.text_in_channels),
+            (audio_indices, self.audio_embedder, self.config.audio_in_channels),
+            (video_indices, self.video_embedder, self.config.video_in_channels),
+        ):
+            if indices.numel():
+                tokens = packed[:, :in_channels].index_select(0, indices).float()
+                output.index_copy_(0, indices, embedder(tokens).to(output.dtype))
+        return output


 class Magi2PostAdapter(nn.Module):
@@ -425,8 +441,10 @@ class Magi2TransformerLayer(nn.Module):
        self.mlp: nn.Module
        if layer_index in config.moe.layers:
            self.mlp = Magi2MultiHeadMoELayer(config)
+            self.region_methods = ("_attention_input", "_moe_input", "_moe_output")
        else:
            self.mlp = Magi2MLP(config, num_modality=num_modality)
+            self.region_methods = ("_attention_input", "_dense_output")
        self._init_mhc(num_modality)

    def _init_mhc(self, num_modality: int) -> None:
@@ -525,6 +543,74 @@ class Magi2TransformerLayer(nn.Module):
        )
        return self.mhc_handler.hyper_connect(streams, output, post, residual)

+    # The layer forward is split into compile regions around the two eager
+    # kernels: packed attention and the multi-head MoE.
+
+    def _attention_input(
+        self,
+        hidden_states: torch.Tensor,
+        rope: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher,
+    ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        streams = hidden_states.reshape(hidden_states.shape[0], self.config.mhc.num_streams, self.config.hidden_size)
+        attention_logits = self._branch_logits(streams, "attn", modality_dispatcher)
+        attention_input = self._branch_input(streams, "attn", attention_logits)
+        q, k, v, gates = self.attention.project(attention_input, rope, modality_dispatcher)
+        return streams, attention_logits, q, k, v, gates
+
+    def _mlp_input(
+        self,
+        streams: torch.Tensor,
+        attention: torch.Tensor,
+        gates: torch.Tensor,
+        attention_logits: MHCTensorTuple,
+        modality_dispatcher: ModalityDispatcher,
+    ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor]:
+        attention_output = self.attention.output(attention, gates, modality_dispatcher)
+        streams = self._connect(streams, attention_output, "attn", attention_logits)
+        mlp_logits = self._branch_logits(streams, "mlp", modality_dispatcher)
+        return streams, mlp_logits, self._branch_input(streams, "mlp", mlp_logits)
+
+    def _dense_output(
+        self,
+        streams: torch.Tensor,
+        attention: torch.Tensor,
+        gates: torch.Tensor,
+        attention_logits: MHCTensorTuple,
+        modality_dispatcher: ModalityDispatcher,
+    ) -> torch.Tensor:
+        streams, mlp_logits, mlp_input = self._mlp_input(
+            streams, attention, gates, attention_logits, modality_dispatcher
+        )
+        streams = self._connect(streams, self.mlp(mlp_input, modality_dispatcher), "mlp", mlp_logits)
+        return streams.reshape(streams.shape[0], -1)
+
+    def _moe_input(
+        self,
+        streams: torch.Tensor,
+        attention: torch.Tensor,
+        gates: torch.Tensor,
+        attention_logits: MHCTensorTuple,
+        modality_dispatcher: ModalityDispatcher,
+    ) -> tuple[torch.Tensor, MHCTensorTuple, torch.Tensor, torch.Tensor]:
+        streams, mlp_logits, mlp_input = self._mlp_input(
+            streams, attention, gates, attention_logits, modality_dispatcher
+        )
+        normalized, routed = self.mlp.route_input(mlp_input, modality_dispatcher)
+        return streams, mlp_logits, normalized, routed
+
+    def _moe_output(
+        self,
+        streams: torch.Tensor,
+        mlp_logits: MHCTensorTuple,
+        normalized: torch.Tensor,
+        routed: torch.Tensor,
+        modality_dispatcher: ModalityDispatcher,
+    ) -> torch.Tensor:
+        mlp_output = self.mlp.combine(normalized, routed, modality_dispatcher)
+        streams = self._connect(streams, mlp_output, "mlp", mlp_logits)
+        return streams.reshape(streams.shape[0], -1)
+
    def forward(
        self,
        hidden_states: torch.Tensor,
@@ -533,22 +619,15 @@ class Magi2TransformerLayer(nn.Module):
        modality_dispatcher: ModalityDispatcher,
        cp_split_sizes: list[int] | torch.Tensor,
    ) -> torch.Tensor:
-        streams = hidden_states.reshape(hidden_states.shape[0], self.config.mhc.num_streams, self.config.hidden_size)
-        attention_logits = self._branch_logits(streams, "attn", modality_dispatcher)
-        attention_input = self._branch_input(streams, "attn", attention_logits)
-        attention_output = self.attention(
-            attention_input,
-            rope,
-            varlen_handler,
-            modality_dispatcher,
-            cp_split_sizes,
+        streams, attention_logits, q, k, v, gates = self._attention_input(hidden_states, rope, modality_dispatcher)
+        attention = self.attention.attend(q, k, v, varlen_handler, cp_split_sizes)
+        if not isinstance(self.mlp, Magi2MultiHeadMoELayer):
+            return self._dense_output(streams, attention, gates, attention_logits, modality_dispatcher)
+        streams, mlp_logits, normalized, routed = self._moe_input(
+            streams, attention, gates, attention_logits, modality_dispatcher
        )
-        streams = self._connect(streams, attention_output, "attn", attention_logits)
-        mlp_logits = self._branch_logits(streams, "mlp", modality_dispatcher)
-        mlp_input = self._branch_input(streams, "mlp", mlp_logits)
-        mlp_output = self.mlp(mlp_input, modality_dispatcher)
-        streams = self._connect(streams, mlp_output, "mlp", mlp_logits)
-        return streams.reshape(streams.shape[0], -1)
+        routed = self.mlp.moe_mlp(routed)
+        return self._moe_output(streams, mlp_logits, normalized, routed, modality_dispatcher)


 class Magi2TransformerBlock(nn.Module):
@@ -570,7 +649,6 @@ class Magi2TransformerBlock(nn.Module):
        modality_dispatcher: ModalityDispatcher,
        cp_split_sizes: list[int] | torch.Tensor,
    ) -> torch.Tensor:
-        hidden_states = hidden_states.to(self.config.params_dtype)
        hidden_states = modality_dispatcher.permute(hidden_states)
        for layer in self.layers:
            hidden_states = layer(
@@ -630,6 +708,15 @@ class Magi2PreviewTransformer(nn.Module):

        return self.block.layers

+    def compile_regions(self, **compile_kwargs: Any) -> None:
+        """Compile the dense compute between the eager attention and MoE kernels."""
+
+        self.pre_adapter.forward = torch.compile(self.pre_adapter.forward, **compile_kwargs)
+        self.post_adapter.forward = torch.compile(self.post_adapter.forward, **compile_kwargs)
+        for layer in self.block.layers:
+            for name in layer.region_methods:
+                setattr(layer, name, torch.compile(getattr(layer, name), **compile_kwargs))
+
    def forward(
        self,
        x: torch.Tensor,
@@ -647,6 +734,7 @@ class Magi2PreviewTransformer(nn.Module):
        assert dispatcher.split_sizes is not None
        cp_split_sizes = dispatcher.split_sizes

+        rope = self.pre_adapter.rope(coords_mapping)
        time_mask = modality_mapping == int(Modality.TIME)
        modality_mapping = torch.where(time_mask, int(Modality.TEXT), modality_mapping)
        modality_dispatcher = ModalityDispatcher(modality_mapping, 3)
@@ -654,13 +742,7 @@ class Magi2PreviewTransformer(nn.Module):
        audio_indices = torch.nonzero(modality_mapping == int(Modality.AUDIO)).flatten()
        text_indices = torch.nonzero(modality_mapping == int(Modality.TEXT)).flatten()

-        hidden_states, rope = self.pre_adapter(
-            x,
-            coords_mapping,
-            video_indices,
-            audio_indices,
-            text_indices,
-        )
+        hidden_states = self.pre_adapter(x, video_indices, audio_indices, text_indices)
        if time_token_sequence is not None and time_token_sequence.shape[-1] > 0:
            hidden_states[:, : time_token_sequence.shape[-1]] = time_token_sequence.to(hidden_states.dtype)
        hidden_states = self.block(
@@ -695,6 +695,24 @@ class Magi2Pipeline(

        return self.video_decoder.module

+    def setup_compile(self) -> None:
+        """Compile the transformer regions; the attention and MoE kernels stay eager."""
+
+        granularity = self.od_config.diffusion_compile_granularity
+        if granularity != "regional":
+            logger.warning(
+                "MAGI-2 compiles the transformer regions itself; diffusion_compile_granularity=%r is ignored.",
+                granularity,
+            )
+
+        # The mHC connections chain bf16 ops that eager rounds after every op.
+        # Emulating those casts keeps the compiled rounding boundaries equal.
+        self.transformer.compile_regions(
+            fullgraph=True,
+            dynamic=self.od_config.diffusion_compile_dynamic,
+            options={"emulate_precision_casts": True},
+        )
+
    def load_weights(
        self,
        weights: Iterable[tuple[str, torch.Tensor]],
@@ -308,29 +308,11 @@ class Magi2PreviewSampler(CFGParallelMixin):
        per_token_video_t = t_batch.view(-1, 1, 1, 1, 1).expand(cfg_batch, 1, video_t, video_h, video_w)
        per_token_audio_t = t_batch.view(-1, 1, 1).expand(cfg_batch, audio_t, 1)

-        audio_lengths = torch.full(
-            (cfg_batch,),
-            audio_latent_len,
-            device=latent.device,
-            dtype=torch.long,
-        )
-        text_lengths = torch.tensor(
-            [txt_feat_len] * batch_size + [null_txt_feat_len] * batch_size,
-            device=latent.device,
-            dtype=torch.long,
-        )
-        ref_audio_lengths = torch.full(
-            (cfg_batch,),
-            ref_audio_feat_len,
-            device=latent.device,
-            dtype=torch.long,
-        )
-        ref_video_lengths = torch.full(
-            (cfg_batch,),
-            ref_video_feat_len,
-            device=latent.device,
-            dtype=torch.long,
-        )
+        # Lengths stay host integers; the data proxy reads them as Python ints.
+        audio_lengths = [audio_latent_len] * cfg_batch
+        text_lengths = [txt_feat_len] * batch_size + [null_txt_feat_len] * batch_size
+        ref_audio_lengths = [ref_audio_feat_len] * cfg_batch
+        ref_video_lengths = [ref_video_feat_len] * cfg_batch

        return ModelInput(
            x_t=torch.cat((latent, latent), dim=0),