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

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>

tests/diffusion/models/magi2/test_native_compile.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from tests.diffusion.models.magi2.test_native_packing import (
+    _longer_text_tensors,
+    _sampler_tensors,
+    _tiny_model,
+    _tiny_sampler,
+)
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def test_layer_regions_bracket_the_eager_kernels() -> None:
+    moe_layer, dense_layer = _tiny_model().block.layers
+
+    assert moe_layer.region_methods == ("_attention_input", "_moe_input", "_moe_output")
+    assert dense_layer.region_methods == ("_attention_input", "_dense_output")
+
+
+def test_compile_regions_fullgraph_matches_eager_bitwise() -> None:
+    model = _tiny_model()
+    sampler = _tiny_sampler(model)
+    steps = [
+        sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
+        for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
+    ]
+    expected = [sampler.forward(step) for step in steps]
+
+    torch._dynamo.reset()
+    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
+    first = sampler.forward(steps[0])
+    graphs = torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        second = sampler.forward(steps[1])
+
+    for (video, audio), (video_ref, audio_ref) in zip((first, second), expected, strict=True):
+        assert torch.equal(video, video_ref)
+        assert torch.equal(audio, audio_ref)
+    # pre adapter, post adapter, three MoE-layer regions, two dense-layer regions
+    assert graphs == 7
+
+
+def test_pipeline_setup_compile_forwards_config_dynamic() -> None:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
+
+    calls: list[dict[str, object]] = []
+
+    class _Transformer(torch.nn.Module):
+        def compile_regions(self, **compile_kwargs: object) -> None:
+            calls.append(compile_kwargs)
+
+    pipeline = object.__new__(Magi2Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = OmniDiffusionConfig(
+        model="sand-ai/MAGI-2-preview",
+        model_class_name="Magi2Pipeline",
+        diffusion_compile_dynamic=False,
+    )
+    pipeline.transformer = _Transformer()
+
+    pipeline.setup_compile()
+
+    assert calls == [{"fullgraph": True, "dynamic": False, "options": {"emulate_precision_casts": True}}]
+
+
+def test_pipeline_setup_compile_warns_when_granularity_is_not_regional(caplog: pytest.LogCaptureFixture) -> None:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
+
+    calls: list[dict[str, object]] = []
+
+    class _Transformer(torch.nn.Module):
+        def compile_regions(self, **compile_kwargs: object) -> None:
+            calls.append(compile_kwargs)
+
+    pipeline = object.__new__(Magi2Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = OmniDiffusionConfig(
+        model="sand-ai/MAGI-2-preview",
+        model_class_name="Magi2Pipeline",
+        diffusion_compile_granularity="full",
+    )
+    pipeline.transformer = _Transformer()
+
+    pipeline.setup_compile()
+
+    assert calls == [{"fullgraph": True, "dynamic": True, "options": {"emulate_precision_casts": True}}]
+    assert "diffusion_compile_granularity='full'" in caplog.text
+
+
+def test_compiled_graphs_are_shared_across_layers_of_one_kind() -> None:
+    model = _tiny_model(moe_layers=2)
+    sampler = _tiny_sampler(model)
+    step = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([900.0]), cfg_config=CFGConfig())
+    expected = sampler.forward(step)
+
+    torch._dynamo.reset()
+    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
+    actual = sampler.forward(step)
+
+    assert torch.equal(actual[0], expected[0])
+    # Two MoE layers and two dense layers still compile one graph per region
+    # kind: pre adapter, post adapter, three MoE regions, two dense regions.
+    assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
+
+
+def test_compiled_regions_survive_new_requests_without_recompiling() -> None:
+    model = _tiny_model()
+    sampler = _tiny_sampler(model)
+    t = torch.tensor([900.0])
+    first_request = sampler.prepare_model_input(**_sampler_tensors(0), t=t, cfg_config=CFGConfig())
+    second_request = sampler.prepare_model_input(**_sampler_tensors(1), t=t, cfg_config=CFGConfig())
+    longer_request = sampler.prepare_model_input(**_longer_text_tensors(2), t=t, cfg_config=CFGConfig())
+    expected = [sampler.forward(request) for request in (first_request, second_request, longer_request)]
+
+    torch._dynamo.reset()
+    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
+    actual = [sampler.forward(first_request)]
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        actual.append(sampler.forward(second_request))
+        actual.append(sampler.forward(longer_request))
+
+    for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
+        assert torch.equal(video, video_ref)
+        assert torch.equal(audio, audio_ref)
+    assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7

tests/diffusion/models/magi2/test_native_compile_cuda.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from tests.diffusion.models.magi2.test_native_packing import (
+    _longer_text_tensors,
+    _sampler_tensors,
+    _tiny_model,
+    _tiny_sampler,
+)
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
+
+pytestmark = [
+    pytest.mark.diffusion,
+    pytest.mark.core_model,
+    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"),
+]
+
+
+def _cuda_steps(sampler, tensors=_sampler_tensors):
+    timesteps = (torch.tensor([900.0]), torch.tensor([600.0]), torch.tensor([300.0]))
+    return [
+        sampler.prepare_model_input(
+            **{name: value.cuda() for name, value in tensors(seed).items()},
+            t=t,
+            cfg_config=CFGConfig(),
+        )
+        for seed, t in enumerate(timesteps)
+    ]
+
+
+def test_inductor_regions_match_eager_in_deterministic_mode(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("MAGI2_DETERMINISTIC", "1")
+    model = _tiny_model(params_dtype=torch.bfloat16).cuda()
+    sampler = _tiny_sampler(model)
+    steps = _cuda_steps(sampler)
+    expected = [sampler.forward(step) for step in steps]
+    new_request_step = _cuda_steps(sampler)[0]
+    longer_text_step = _cuda_steps(sampler, _longer_text_tensors)[0]
+    expected_new_request = sampler.forward(new_request_step)
+    expected_longer_text = sampler.forward(longer_text_step)
+
+    torch._dynamo.reset()
+    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+    # The pipeline compiles with inductor emulating eager's intermediate
+    # low-precision casts, so fused bf16 chains keep eager's rounding.
+    model.compile_regions(fullgraph=True, dynamic=True, options={"emulate_precision_casts": True})
+    first = sampler.forward(steps[0])
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        compiled = [first, sampler.forward(steps[1]), sampler.forward(steps[2])]
+        repeated = [sampler.forward(step) for step in steps]
+        # A second request repacks the same shapes; a third changes the
+        # text length and therefore every token count.
+        compiled_new_request = sampler.forward(new_request_step)
+        compiled_longer_text = sampler.forward(longer_text_step)
+
+    # Emulated precision casts restore eager's rounding boundaries; what is
+    # left is one-ulp transcendental noise that some Triton backends carry.
+    for (video, audio), (video_ref, audio_ref) in zip(compiled, expected, strict=True):
+        torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
+        torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
+    for (video, audio), (video_again, audio_again) in zip(compiled, repeated, strict=True):
+        assert torch.equal(video, video_again)
+        assert torch.equal(audio, audio_again)
+    for (video, audio), (video_ref, audio_ref) in (
+        (compiled_new_request, expected_new_request),
+        (compiled_longer_text, expected_longer_text),
+    ):
+        torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
+        torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
+    assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7

tests/diffusion/models/magi2/test_native_compile_distributed.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import os
+import tempfile
+from contextlib import ExitStack, contextmanager
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+from torch.distributed import init_device_mesh
+from torch.distributed.fsdp import MixedPrecisionPolicy
+
+import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dlo_module
+from tests.diffusion.models.magi2.test_native_distributed_parity import (
+    _current_group,
+    _new_groups,
+    _patched_groups,
+)
+from tests.diffusion.models.magi2.test_native_packing import _tiny_config, _tiny_model, _tiny_sampler
+from tests.diffusion.offloader.helpers import patch_offload_runtime
+from vllm_omni.diffusion.distributed.hsdp import shard_model
+from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
+from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup
+from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
+from vllm_omni.diffusion.offloader.base import OffloadConfig
+from vllm_omni.diffusion.offloader.config import OffloadStrategy
+
+_WORLD_SIZE = 4
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+@contextmanager
+def _cpu_offload_runtime():
+    with pytest.MonkeyPatch.context() as monkeypatch:
+        patch_offload_runtime(monkeypatch, dlo_module.current_omni_platform, synchronize=True)
+        yield
+
+
+def _enable_dlo(
+    model: Magi2PreviewTransformer,
+    *,
+    dp_group: dist.ProcessGroup | None,
+    dp_size: int,
+    dp_rank: int,
+) -> None:
+    pipeline = Magi2Pipeline.__new__(Magi2Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.transformer = model
+    backend = dlo_module.DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=False,
+            dp_size=dp_size,
+            dlo_use_allgather=dp_size > 1,
+        ),
+        torch.device("cpu"),
+    )
+    backend.dp_group = dp_group
+    backend.rank = dp_rank
+    backend.enable(pipeline)
+    assert backend.enabled
+
+
+def _enable_hsdp(model: Magi2PreviewTransformer) -> None:
+    mesh = init_device_mesh("cpu", mesh_shape=(1, _WORLD_SIZE), mesh_dim_names=("replicate", "shard"))
+    ignored = {
+        parameter for name in model._hsdp_ignored_modules for parameter in model.get_submodule(name).parameters()
+    }
+    shard_model(
+        model,
+        mesh=mesh,
+        mp_policy=MixedPrecisionPolicy(cast_forward_inputs=False),
+        hsdp_shard_conditions=model._hsdp_shard_conditions,
+        ignored_params=ignored,
+    )
+    model.requires_grad_(False)
+
+
+def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
+    # 14 tokens per CFG branch, so the packed sequence splits evenly across
+    # two Ulysses ranks; gloo cannot gather uneven shards.
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    return {
+        "latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
+        "audio_latent": torch.randn(1, 5, 4, generator=generator),
+        "txt_feat": torch.randn(1, 3, 4, generator=generator),
+        "null_txt_feat": torch.randn(1, 3, 4, generator=generator),
+    }
+
+
+def _steps(sampler):
+    return [
+        sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
+        for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
+    ]
+
+
+def _run(model: Magi2PreviewTransformer, *, compile_regions: bool) -> list[tuple[torch.Tensor, torch.Tensor]]:
+    sampler = _tiny_sampler(model)
+    steps = _steps(sampler)
+    if not compile_regions:
+        return [sampler.forward(step) for step in steps]
+    torch._dynamo.reset()
+    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
+    outputs = [sampler.forward(steps[0])]
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        outputs.append(sampler.forward(steps[1]))
+    return outputs
+
+
+def _assert_same(actual, expected, *, exact: bool, label: str) -> None:
+    for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
+        if exact:
+            assert torch.equal(video, video_ref), label
+            assert torch.equal(audio, audio_ref), label
+        else:
+            torch.testing.assert_close(video, video_ref, atol=2e-5, rtol=2e-4, msg=label)
+            torch.testing.assert_close(audio, audio_ref, atol=2e-5, rtol=2e-4, msg=label)
+
+
+def _distributed_worker(rank: int, rendezvous: str) -> None:
+    torch.set_num_threads(1)
+    dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=_WORLD_SIZE)
+    try:
+        singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
+        with _patched_groups(singleton, singleton):
+            oracle = _tiny_model()
+            checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()]
+            expected = _run(oracle, compile_regions=False)
+
+        sp2_groups = _new_groups(((0, 1), (2, 3)))
+        dp2_groups = _new_groups(((0, 2), (1, 3)))
+        sp_group = _current_group(rank, sp2_groups)
+        dp_ranks, dp_process_group = next((ranks, group) for ranks, group in dp2_groups if rank in ranks)
+
+        def dlo_rank_local(model: Magi2PreviewTransformer) -> None:
+            _enable_dlo(model, dp_group=None, dp_size=1, dp_rank=0)
+
+        def dlo_allgather(model: Magi2PreviewTransformer) -> None:
+            _enable_dlo(model, dp_group=dp_process_group, dp_size=len(dp_ranks), dp_rank=dp_ranks.index(rank))
+
+        variants = (
+            ("dlo_rank_local_sp2", sp_group, dlo_rank_local),
+            ("dlo_allgather_dp2sp2", sp_group, dlo_allgather),
+            ("hsdp4", singleton, _enable_hsdp),
+        )
+        for name, ulysses_group, enable in variants:
+            with ExitStack() as stack:
+                stack.enter_context(_patched_groups(singleton, ulysses_group))
+                stack.enter_context(_cpu_offload_runtime())
+                outputs = {}
+                for compile_regions in (False, True):
+                    model = Magi2PreviewTransformer(_tiny_config())
+                    assert model.load_weights(checkpoint) == set(model.state_dict())
+                    enable(model)
+                    outputs[compile_regions] = _run(model, compile_regions=compile_regions)
+            _assert_same(outputs[False], expected, exact=False, label=f"{name} eager vs single-rank oracle")
+            _assert_same(outputs[True], outputs[False], exact=True, label=f"{name} compiled vs eager")
+            dist.barrier()
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(
+    not dist.is_available() or not dist.is_gloo_available(),
+    reason="requires torch.distributed gloo",
+)
+def test_compiled_regions_match_eager_under_dlo_and_hsdp() -> None:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
+        with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
+            mp.spawn(_distributed_worker, args=(rendezvous,), nprocs=_WORLD_SIZE, join=True)

tests/diffusion/models/magi2/test_native_packing.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
+    Magi2MHCConfig,
+    Magi2MoEConfig,
+    Magi2PreviewConfig,
+)
+from vllm_omni.diffusion.models.magi2.layers import MultiModalityRMSNorm
+from vllm_omni.diffusion.models.magi2.mh_moe import Magi2MultiHeadMoE
+from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
+from vllm_omni.diffusion.models.magi2.preview_data_proxy import (
+    Magi2DataProxy,
+    Magi2PreviewDataProxyConfig,
+)
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _tiny_config(params_dtype: torch.dtype = torch.float32, *, moe_layers: int = 1) -> Magi2PreviewConfig:
+    # The leading layers are multimodal with MoE, the rest single-modality
+    # dense, so every layer belongs to one of two kinds.
+    moe_indices = tuple(range(moe_layers))
+    return Magi2PreviewConfig(
+        num_layers=2 * moe_layers,
+        hidden_size=16,
+        head_dim=8,
+        num_query_groups=2,
+        video_in_channels=4,
+        audio_in_channels=4,
+        text_in_channels=4,
+        intermediate_factor=2,
+        multimodal_layers=moe_indices,
+        params_dtype=params_dtype,
+        mhc=Magi2MHCConfig(num_streams=2),
+        moe=Magi2MoEConfig(
+            num_heads=2,
+            num_experts=4,
+            top_k=2,
+            expert_intermediate_size=8,
+            shared_expert_intermediate_size=8,
+            modality_shared_expert_intermediate_size=8,
+            layers=moe_indices,
+        ),
+    )
+
+
+def _tiny_model(
+    seed: int = 11,
+    params_dtype: torch.dtype = torch.float32,
+    *,
+    moe_layers: int = 1,
+) -> Magi2PreviewTransformer:
+    model = Magi2PreviewTransformer(_tiny_config(params_dtype, moe_layers=moe_layers))
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
+        for module in model.modules():
+            if isinstance(module, MultiModalityRMSNorm):
+                module.weight.zero_()
+            elif isinstance(module, Magi2MultiHeadMoE):
+                module.router.expert_bias.zero_()
+                module.router.expert_bias_ema.zero_()
+    return model
+
+
+def _tiny_sampler(model: torch.nn.Module) -> Magi2PreviewSampler:
+    return Magi2PreviewSampler(model, Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=8)))
+
+
+def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    return {
+        "latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
+        "audio_latent": torch.randn(1, 5, 4, generator=generator),
+        "txt_feat": torch.randn(1, 3, 4, generator=generator),
+        "null_txt_feat": torch.randn(1, 2, 4, generator=generator),
+    }
+
+
+def _longer_text_tensors(seed: int) -> dict[str, torch.Tensor]:
+    tensors = _sampler_tensors(seed)
+    generator = torch.Generator(device="cpu").manual_seed(seed + 100)
+    tensors["txt_feat"] = torch.randn(1, 4, 4, generator=generator)
+    return tensors
+
+
+def test_prepare_model_input_keeps_lengths_on_host() -> None:
+    sampler = Magi2PreviewSampler(torch.nn.Identity())
+    model_input = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([500.0]), cfg_config=CFGConfig())
+
+    assert model_input.audio_feat_len == [5, 5]
+    assert model_input.txt_feat_len == [3, 2]
+    assert model_input.ref_audio_feat_len == [0, 0]
+    assert model_input.ref_video_feat_len == [0, 0]
+
+    positive, negative = Magi2PreviewSampler._split_cfg_model_input(model_input)
+    assert positive.txt_feat_len == [3]
+    assert negative.txt_feat_len == [2]
+
+
+def test_pre_adapter_embeds_directly_in_checkpoint_dtype() -> None:
+    model = _tiny_model(params_dtype=torch.bfloat16)
+    packed = torch.randn(6, 4)
+    indices = torch.tensor([0, 1]), torch.tensor([2, 3]), torch.tensor([4, 5])
+
+    with torch.no_grad():
+        hidden = model.pre_adapter(packed, *indices)
+
+    assert hidden.dtype == torch.bfloat16
+    assert hidden.shape == (6, model.config.virtual_width)

vllm_omni/diffusion/models/magi2/modeling_magi2.py

@@ -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(

vllm_omni/diffusion/models/magi2/pipeline_magi2.py

@@ -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]],

vllm_omni/diffusion/models/magi2/sampler_magi2.py

@@ -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),