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

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

tests/diffusion/attention/test_attention_config.py

@@ -907,6 +907,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
                 forward=lambda *args: pytest.fail("unexpected SDPA fallback"),
             ),
             _assert_metadata_compatible=lambda metadata: None,
+            _has_custom_attention=False,
             _scheduler_paged_kv=False,
             paged_kv_cache_role=None,
         )

tests/diffusion/attention/test_flash_attn.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Test script for FlashAttention backend with padding handling.
@@ -29,6 +29,30 @@ HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
 flash_attn_func = fa.flash_attn_func  # noqa: N813
 
 
+@pytest.mark.parametrize(
+    ("device_major", "requested", "supported", "expected"),
+    [
+        (8, None, frozenset((2, 3, 4)), 2),
+        (9, None, frozenset((2, 3, 4)), 3),
+        (10, None, frozenset((2, 3, 4)), 4),
+        (12, None, frozenset((2, 3, 4)), 4),
+        (9, None, frozenset((2,)), 2),
+        (10, None, frozenset((2, 3)), 3),
+        (9, "2", frozenset((2, 3, 4)), 2),
+        (9, "4", frozenset((2, 3, 4)), 4),
+    ],
+)
+def test_versioned_flash_attention_selection(device_major, requested, supported, expected):
+    assert fa._choose_vllm_flash_attn_version(device_major, requested, supported) == expected
+
+
+def test_versioned_flash_attention_selection_rejects_invalid_or_unavailable_version():
+    with pytest.raises(ValueError, match="must be 2, 3, or 4"):
+        fa._choose_vllm_flash_attn_version(9, "5", frozenset((2, 3, 4)))
+    with pytest.raises(RuntimeError, match="FlashAttention 4 is unavailable"):
+        fa._choose_vllm_flash_attn_version(9, "4", frozenset((2, 3)))
+
+
 def create_attention_mask(batch_size: int, seq_len: int, valid_len: int, device: torch.device) -> torch.Tensor:
     """
     Create attention mask where first valid_len tokens are valid (1) and rest are padding (0).

tests/diffusion/cache/test_cache_backends.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Unit tests for cache backends (cache-dit and teacache).
@@ -164,7 +164,10 @@ class TestCacheDiTBackend:
             num_inference_steps=20,
             verbose=True,
         )
-        mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
+        mock_cache_dit.summary.assert_called_once_with(
+            mock_block_adapter.return_value,
+            details=True,
+        )
 
     @patch("vllm_omni.diffusion.cache.cachedit.backend.logger")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")

tests/diffusion/cache/test_cache_dit.py

@@ -59,6 +59,7 @@ def test_custom_cache_dit_enablers_are_registered_explicitly():
         "Cosmos3OmniDiffusersPipeline": cd_model_specific.enable_cache_for_cosmos3,
         "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3,
         "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2,
+        "Magi2Pipeline": cd_model_specific.enable_cache_for_magi2,
     }
 
     with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True):
@@ -90,12 +91,50 @@ def test_cosmos3_aliases_use_cosmos3_custom_cache_dit_enabler(pipeline_name: str
     assert cd_backend.CUSTOM_DIT_ENABLERS[pipeline_name] is cd_model_specific.enable_cache_for_cosmos3
 
 
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.enable_cache_for_dit")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
+def test_magi2_cache_dit_targets_only_nested_repeated_layers(mock_block_adapter, mock_enable_cache):
+    pipeline = Mock()
+    transformer_block = pipeline.transformer.block
+    layers = torch.nn.ModuleList([torch.nn.Identity()])
+    transformer_block.layers = layers
+    adapter = mock_block_adapter.return_value
+    refresh = Mock()
+    mock_enable_cache.return_value = refresh
+
+    result = cd_model_specific.enable_cache_for_magi2(pipeline, SAMPLE_CACHE_CONFIG)
+
+    mock_block_adapter.assert_called_once()
+    adapter_kwargs = mock_block_adapter.call_args.kwargs
+    assert adapter_kwargs["transformer"] is transformer_block
+    assert adapter_kwargs["blocks"] == [layers]
+    assert adapter_kwargs["has_separate_cfg"] is False
+    assert adapter_kwargs["check_forward_pattern"] is True
+    assert result.refresh is refresh
+    assert result.targets == (adapter,)
+    get_transformer = mock_enable_cache.call_args.kwargs["get_pipeline_transformer"]
+    assert get_transformer(pipeline) is transformer_block
+
+
+@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit.summary")
+@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter.is_cached", return_value=True)
+def test_cache_summary_uses_custom_nested_targets(mock_is_cached, mock_summary):
+    target = object()
+    pipeline = Mock(_cache_dit_targets=(target,))
+
+    cd_backend.cache_summary(pipeline, details=True)
+
+    mock_is_cached.assert_called_once_with(target)
+    mock_summary.assert_called_once_with(target, details=True)
+
+
 def test_cachedit_public_api_is_explicit():
     assert set(cd_backend.__all__) == {
         "BagelCachedAdapter",
         "CUSTOM_DIT_ENABLERS",
         "CacheDiTAdapterConfig",
         "CacheDiTBackend",
+        "CacheDiTEnableResult",
         "CacheDiTConfig",
         "CacheDiTRequestSpec",
         "RequestScopedCacheDiTRuntime",
@@ -213,7 +252,7 @@ def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
     current_omni_platform.is_rocm(),
     reason="vLLM ROCm custom ops lack CPU fallback",
 )
-def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
+def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group, request: pytest.FixtureRequest):
     """CacheDiT Pattern_0 treats the second positional arg as encoder_hidden_states,
     which is a collision for one of the kwargs in LTX2 since we treat the audio
     hidden states as encoder_hidden_states.
@@ -253,6 +292,7 @@ def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
     pipeline.transformer = model
     backend = CacheDiTBackend(DiffusionCacheConfig())
     backend.enable(pipeline)
+    request.addfinalizer(lambda: backend.disable(pipeline))
     backend.refresh(pipeline, num_inference_steps=5)
 
     # Wrap call_Fn_blocks in CacheDiT so that we can verify the

tests/diffusion/distributed/test_hsdp.py

@@ -12,6 +12,7 @@ from torch.distributed.tensor import DeviceMesh, DTensor
 
 from tests.helpers.runtime import get_distributed_init_method
 from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.distributed import hsdp as hsdp_module
 from vllm_omni.diffusion.distributed.hsdp import (
     HSDPInferenceConfig,
     _unshardable_parameters,
@@ -73,6 +74,52 @@ def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
     assert not isinstance(model.block.input_global_scale, DTensor)
 
 
+def test_hsdp_resolves_nested_ignored_module_and_preserves_mixed_dtypes(
+    cpu_process_group,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    class Group:
+        world_size = 1
+        rank_in_group = 0
+        device_group = dist.group.WORLD
+
+    class MixedDtypeModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layers = nn.Module()
+            self.layers.sharded = nn.Linear(2, 2, dtype=torch.bfloat16)
+            self.layers.ignored = nn.Linear(2, 2, dtype=torch.float32)
+            self._hsdp_shard_conditions = [lambda name, _module: name == "layers.sharded"]
+            self._hsdp_ignored_modules = ["layers.ignored"]
+            self._hsdp_preserve_parameter_dtypes = True
+
+    group = Group()
+    monkeypatch.setattr(hsdp_module, "get_world_group", lambda: group)
+    monkeypatch.setattr(
+        hsdp_module,
+        "_create_hsdp_mesh",
+        lambda **_kwargs: DeviceMesh("cpu", [0]),
+    )
+
+    model = MixedDtypeModel()
+    ignored_weight = model.layers.ignored.weight
+    hsdp_module.apply_hsdp_to_model(
+        model,
+        HSDPInferenceConfig(
+            enabled=True,
+            hsdp_shard_size=1,
+            param_dtype=torch.float16,
+        ),
+        target_device=torch.device("cpu"),
+    )
+
+    assert model.layers.ignored.weight is ignored_weight
+    assert not isinstance(model.layers.ignored.weight, DTensor)
+    assert model.layers.ignored.weight.dtype == torch.float32
+    assert isinstance(model.layers.sharded.weight, DTensor)
+    assert model.layers.sharded.weight.dtype == torch.bfloat16
+
+
 class TestHSDPInferenceConfig:
     """Tests for HSDPInferenceConfig dataclass."""
 

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -199,6 +199,7 @@ def _fake_ltx_parallel_attention():
     layer.attention = SimpleNamespace(forward=run("native"))
     layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
     layer.backend_pref = "FLASH_ATTN"
+    layer._has_custom_attention = False
     return layer, calls
 
 

tests/diffusion/models/magi2/test_audio_decoder.py

@@ -0,0 +1,56 @@
+# 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.audio_decoder import (
+    convert_stable_audio_decoder_key,
+    convert_stable_audio_decoder_state_dict,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+@pytest.mark.parametrize(
+    ("source", "target"),
+    [
+        ("pretransform.model.decoder.layers.0.weight_v", "conv1.weight_v"),
+        ("pretransform.model.decoder.layers.1.layers.0.alpha", "block.0.snake1.alpha"),
+        ("pretransform.model.decoder.layers.3.layers.1.weight_g", "block.2.conv_t1.weight_g"),
+        (
+            "pretransform.model.decoder.layers.5.layers.4.layers.2.beta",
+            "block.4.res_unit3.snake2.beta",
+        ),
+        (
+            "pretransform.model.decoder.layers.2.layers.2.layers.3.weight_v",
+            "block.1.res_unit1.conv2.weight_v",
+        ),
+        ("pretransform.model.decoder.layers.6.alpha", "snake1.alpha"),
+        ("pretransform.model.decoder.layers.7.weight_g", "conv2.weight_g"),
+    ],
+)
+def test_stable_audio_decoder_key_conversion(source: str, target: str):
+    assert convert_stable_audio_decoder_key(source) == target
+
+
+def test_stable_audio_decoder_conversion_filters_and_reshapes_snake():
+    alpha = torch.arange(4, dtype=torch.float32)
+    weight = torch.ones(2, 4, 7)
+    state = {
+        "pretransform.model.decoder.layers.6.alpha": alpha,
+        "pretransform.model.decoder.layers.7.weight_v": weight,
+        "model.model.transformer.layers.0.weight": torch.empty(1),
+    }
+    converted = convert_stable_audio_decoder_state_dict(state)
+    assert set(converted) == {"snake1.alpha", "conv2.weight_v"}
+    assert converted["snake1.alpha"].shape == (1, 4, 1)
+    torch.testing.assert_close(converted["snake1.alpha"].flatten(), alpha)
+    assert converted["conv2.weight_v"] is weight
+
+
+def test_unknown_decoder_key_fails_closed():
+    with pytest.raises(ValueError, match="Unsupported Stable Audio decoder key"):
+        convert_stable_audio_decoder_key("pretransform.model.decoder.future.weight")

tests/diffusion/models/magi2/test_native_distributed_parity.py

@@ -0,0 +1,229 @@
+# 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 vllm_omni.diffusion.models.magi2.attention as attention_module
+import vllm_omni.diffusion.models.magi2.layers as layers_module
+import vllm_omni.diffusion.models.magi2.mh_moe as mh_moe_module
+import vllm_omni.diffusion.models.magi2.parallel as parallel_module
+from vllm_omni.diffusion.models.magi2.attention import VarlenHandler
+from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
+    Magi2MHCConfig,
+    Magi2MoEConfig,
+    Magi2PreviewConfig,
+)
+from vllm_omni.diffusion.models.magi2.modeling_magi2 import (
+    Magi2PreviewTransformer,
+    Modality,
+)
+from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup
+
+_WORLD_SIZE = 4
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _tiny_config() -> Magi2PreviewConfig:
+    # Four attention, KV, and MoE heads make every head axis divisible by each
+    # requested TP/SP layout. The MLP and shared-expert widths are likewise
+    # divisible by TP=4, so this exercises all native row/column collectives.
+    return Magi2PreviewConfig(
+        num_layers=1,
+        hidden_size=32,
+        head_dim=8,
+        num_query_groups=4,
+        video_in_channels=4,
+        audio_in_channels=4,
+        text_in_channels=4,
+        intermediate_factor=1.5,
+        multimodal_layers=(0,),
+        params_dtype=torch.float32,
+        mhc=Magi2MHCConfig(num_streams=2),
+        moe=Magi2MoEConfig(
+            num_heads=4,
+            num_experts=2,
+            top_k=2,
+            expert_intermediate_size=8,
+            shared_expert_intermediate_size=8,
+            modality_shared_expert_intermediate_size=8,
+            layers=(0,),
+        ),
+    )
+
+
+def _initialize_model(model: Magi2PreviewTransformer) -> None:
+    generator = torch.Generator(device="cpu").manual_seed(71)
+    with torch.no_grad():
+        for name, parameter in model.named_parameters():
+            if name == "pre_adapter.rope.bands":
+                continue
+            parameter.copy_(
+                torch.randn(
+                    parameter.shape,
+                    dtype=parameter.dtype,
+                    generator=generator,
+                )
+                * 0.025
+            )
+        moe = model.block.layers[0].mlp.moe_mlp
+        moe.router.expert_bias.zero_()
+        moe.router.expert_bias_ema.zero_()
+
+
+def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, VarlenHandler]:
+    generator = torch.Generator(device="cpu").manual_seed(113)
+    packed = torch.randn(8, 4, generator=generator) * 0.2
+    coordinates = torch.tensor(
+        [
+            [0, 0, 0, 2, 2, 2, 2, 2, 2],
+            [0, 1, 1, 2, 2, 2, 2, 2, 2],
+            [0, 0, 0, 1, 2, 2, 1, 2, 2],
+            [0, 1, 1, 1, 2, 2, 1, 2, 2],
+            [1, 0, 0, 2, 2, 2, 2, 2, 2],
+            [1, 1, 1, 2, 2, 2, 2, 2, 2],
+            [0, 0, 0, 1, 1, 1, 1, 1, 1],
+            [0, 0, 0, 1, 1, 1, 1, 1, 1],
+        ],
+        dtype=torch.float32,
+    )
+    modalities = torch.tensor(
+        [
+            Modality.VIDEO,
+            Modality.AUDIO,
+            Modality.TEXT,
+            Modality.TIME,
+            Modality.VIDEO,
+            Modality.AUDIO,
+            Modality.TEXT,
+            Modality.TIME,
+        ]
+    )
+    cumulative = torch.tensor([0, packed.shape[0]], dtype=torch.int32)
+    varlen = VarlenHandler(cumulative, cumulative, packed.shape[0], packed.shape[0])
+    return packed, coordinates, modalities, varlen
+
+
+@contextmanager
+def _patched_groups(
+    tp_group: Magi2ParallelGroup,
+    sp_group: Magi2ParallelGroup,
+):
+    # TP owns native weight and MoE-head shards when enabled. SP owns Ulysses
+    # tokens and is also the MoE-head group in the released SP-only layout.
+    ep_group = tp_group if tp_group.world_size > 1 else sp_group
+    with ExitStack() as stack:
+        stack.enter_context(patch.object(layers_module, "get_magi2_tp_group", return_value=tp_group))
+        stack.enter_context(patch.object(mh_moe_module, "get_magi2_ep_group", return_value=ep_group))
+        stack.enter_context(patch.object(parallel_module, "get_magi2_ulysses_group", return_value=sp_group))
+        stack.enter_context(patch.object(attention_module, "get_magi2_ulysses_group", return_value=sp_group))
+        yield
+
+
+def _current_group(
+    rank: int,
+    rank_groups: tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...],
+) -> Magi2ParallelGroup:
+    for ranks, process_group in rank_groups:
+        if rank in ranks:
+            return Magi2ParallelGroup(
+                process_group,
+                world_size=len(ranks),
+                rank=ranks.index(rank),
+            )
+    return Magi2ParallelGroup(None, world_size=1, rank=0)
+
+
+def _new_groups(
+    rank_sets: tuple[tuple[int, ...], ...],
+) -> tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...]:
+    groups = []
+    for ranks in rank_sets:
+        process_group = dist.new_group(ranks=list(ranks), backend="gloo")
+        groups.append((ranks, process_group))
+    return tuple(groups)
+
+
+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 = Magi2PreviewTransformer(_tiny_config())
+            _initialize_model(oracle)
+            checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()]
+            inputs = _inputs()
+            with torch.no_grad():
+                expected = oracle(*inputs)
+
+        tp2_groups = _new_groups(((0, 1), (2, 3)))
+        sp2_groups = _new_groups(((0, 2), (1, 3)))
+        world_group = Magi2ParallelGroup(
+            dist.group.WORLD,
+            world_size=_WORLD_SIZE,
+            rank=rank,
+        )
+        layouts = (
+            ("tp4", world_group, singleton),
+            ("tp2sp2", _current_group(rank, tp2_groups), _current_group(rank, sp2_groups)),
+            ("sp4", singleton, world_group),
+        )
+
+        for layout_name, tp_group, sp_group in layouts:
+            tp_group = Magi2ParallelGroup(
+                tp_group.group,
+                tp_group.world_size,
+                tp_group.rank,
+                replicated_sequence=True,
+            )
+            with _patched_groups(tp_group, sp_group):
+                model = Magi2PreviewTransformer(_tiny_config())
+                assert model.load_weights(checkpoint) == set(model.state_dict())
+                with torch.no_grad():
+                    actual = model(*inputs)
+
+            max_abs_error = (actual - expected).abs().max()
+            dist.all_reduce(max_abs_error, op=dist.ReduceOp.MAX)
+            close = torch.tensor(
+                int(torch.allclose(actual, expected, atol=2e-5, rtol=2e-4)),
+                dtype=torch.int32,
+            )
+            dist.all_reduce(close, op=dist.ReduceOp.MIN)
+            if not close.item():
+                raise AssertionError(
+                    f"{layout_name} differs from the single-rank oracle; "
+                    f"global max absolute error={max_abs_error.item():.8g}"
+                )
+            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_four_rank_tp_sp_layouts_match_single_rank_oracle() -> None:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
+        mp.spawn(
+            _distributed_worker,
+            args=(rendezvous,),
+            nprocs=_WORLD_SIZE,
+            join=True,
+        )

tests/diffusion/models/magi2/test_native_dlo_allgather.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from contextlib import ExitStack, contextmanager
+from dataclasses import dataclass
+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 safetensors.torch import load_file, save_file
+
+import vllm_omni.diffusion.models.magi2.layers as layers_module
+import vllm_omni.diffusion.models.magi2.mh_moe as mh_moe_module
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    build_checkpoint_mmap_plan,
+)
+from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
+    Magi2MHCConfig,
+    Magi2MoEConfig,
+    Magi2PreviewConfig,
+)
+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.offloader.distributed_layerwise_backend import (
+    DistributedLayerwiseOffloadBackend,
+    DistributedLayerwiseOffloadHook,
+)
+
+
+@dataclass(frozen=True)
+class _WeightSourceStub:
+    model_or_path: str
+    subfolder: str
+    revision: str | None
+    prefix: str
+
+
+@dataclass(frozen=True)
+class _PipelineModulesStub:
+    dits: list[nn.Module]
+    dit_names: list[str]
+
+
+_WORLD_SIZE = 4
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _tiny_config() -> Magi2PreviewConfig:
+    return Magi2PreviewConfig(
+        # The released checkpoint has 40 layers. Keep this tiny fixture at two
+        # so it exercises the shared backend's streamable-block contract.
+        num_layers=2,
+        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=(0, 1),
+        params_dtype=torch.float32,
+        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=(0, 1),
+        ),
+    )
+
+
+@contextmanager
+def _patched_weight_groups(ep_group: Magi2ParallelGroup):
+    singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
+    with ExitStack() as stack:
+        stack.enter_context(patch.object(layers_module, "get_magi2_tp_group", return_value=singleton))
+        stack.enter_context(patch.object(mh_moe_module, "get_magi2_ep_group", return_value=ep_group))
+        yield
+
+
+def _initialize_checkpoint(model: Magi2PreviewTransformer) -> None:
+    generator = torch.Generator(device="cpu").manual_seed(211)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.copy_(
+                torch.randn(
+                    parameter.shape,
+                    dtype=parameter.dtype,
+                    generator=generator,
+                )
+                * 0.025
+            )
+        # The normal loader intentionally mirrors EMA into the runtime router
+        # bias. Keep both released checkpoint entries equal so this test can
+        # compare every reconstructed tensor, not only inference-visible ones.
+        moe = model.block.layers[0].mlp.moe_mlp
+        moe.router.expert_bias.copy_(moe.router.expert_bias_ema)
+
+
+def _write_checkpoint(checkpoint_root: str) -> None:
+    singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
+    with _patched_weight_groups(singleton):
+        source = Magi2PreviewTransformer(_tiny_config())
+        _initialize_checkpoint(source)
+
+    weights = {name: tensor.detach().cpu().contiguous() for name, tensor in source.state_dict().items()}
+    preview_dir = os.path.join(checkpoint_root, "preview")
+    os.makedirs(preview_dir)
+    weight_file = os.path.join(preview_dir, "model.safetensors")
+    save_file(weights, weight_file)
+    with open(os.path.join(preview_dir, "model.safetensors.index.json"), "w") as index:
+        json.dump(
+            {"weight_map": {name: os.path.basename(weight_file) for name in weights}},
+            index,
+        )
+
+
+def _new_groups(
+    rank_sets: tuple[tuple[int, ...], ...],
+) -> tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...]:
+    return tuple((ranks, dist.new_group(ranks=list(ranks), backend="gloo")) for ranks in rank_sets)
+
+
+def _current_group(
+    rank: int,
+    rank_groups: tuple[tuple[tuple[int, ...], dist.ProcessGroup], ...],
+) -> tuple[tuple[int, ...], dist.ProcessGroup]:
+    return next((ranks, group) for ranks, group in rank_groups if rank in ranks)
+
+
+def _load_mmap_transform_and_reconstruct(
+    checkpoint_root: str,
+    dp_group: dist.ProcessGroup,
+    dp_size: int,
+    dp_rank: int,
+    ep_group: Magi2ParallelGroup,
+) -> None:
+    checkpoint_file = os.path.join(checkpoint_root, "preview", "model.safetensors")
+    checkpoint = load_file(checkpoint_file, device="cpu")
+
+    with _patched_weight_groups(ep_group):
+        oracle = Magi2PreviewTransformer(_tiny_config())
+        assert oracle.load_weights(checkpoint.items()) == set(oracle.state_dict())
+
+        with torch.device("meta"):
+            target = Magi2PreviewTransformer(_tiny_config())
+        target.requires_grad_(False)
+        pipeline = Magi2Pipeline.__new__(Magi2Pipeline)
+        nn.Module.__init__(pipeline)
+        pipeline.checkpoint_root = checkpoint_root
+        pipeline.transformer = target
+
+        plan_result = build_checkpoint_mmap_plan(
+            pipeline,
+            dit_modules=(("transformer", target),),
+            sources=(
+                _WeightSourceStub(
+                    model_or_path=checkpoint_root,
+                    subfolder="preview",
+                    revision=None,
+                    prefix="transformer.",
+                ),
+            ),
+            model_path=checkpoint_root,
+            tensor_parallel_size=1,
+            use_hsdp=False,
+            online_quantization=False,
+        )
+        assert plan_result.fallback_reason is None
+        assert plan_result.plan is not None
+
+        backend = object.__new__(DistributedLayerwiseOffloadBackend)
+        backend.device = torch.device("cpu")
+        backend._using_rank_local_mmap = False
+        backend._mmap_transforms_by_tensor_id = {}
+        backend._load_weights_via_mmap(
+            pipeline,
+            _PipelineModulesStub(dits=[target], dit_names=["transformer"]),
+            plan_result.plan,
+        )
+
+    target_block = target.block.layers[0]
+    oracle_block = oracle.block.layers[0]
+    target_parameters = dict(target_block.named_parameters())
+    oracle_parameters = dict(oracle_block.named_parameters())
+
+    # The mmap loader binds the released full-head tensor. The DLO shard step
+    # must apply the model-declared EP transform before it computes the DP
+    # shard, so DP peers gather identical SP-local layouts in DP2SP2.
+    moe_name = "mlp.moe_mlp.gate"
+    mmap_moe = target_parameters[moe_name]
+    mmap_transform = backend._mmap_transforms_by_tensor_id[id(mmap_moe)]
+    assert callable(mmap_transform)
+    assert tuple(mmap_moe.shape) == tuple(checkpoint[f"block.layers.0.{moe_name}"].shape)
+    assert tuple(mmap_transform(mmap_moe).shape) == tuple(oracle_parameters[moe_name].shape)
+
+    # Keep this gloo regression accelerator-isolated. Pinned allocation is the
+    # final, layout-preserving operation in _shard_and_pin and has separate
+    # backend coverage; enabling it here would initialize CUDA in every worker.
+    cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin(
+        target_parameters,
+        dict(target_block.named_buffers()),
+        dp_size=dp_size,
+        rank=dp_rank,
+        pin_memory=False,
+        tensor_transforms=backend._mmap_transforms_by_tensor_id,
+    )
+    assert all(not shard.requires_grad for shard in cpu_shards.values())
+    assert all(parameter.numel() == 0 for parameter in target_block.parameters())
+
+    for dtype, local_shard in cpu_shards.items():
+        gathered = torch.empty(
+            local_shard.numel() * dp_size,
+            dtype=dtype,
+            device="cpu",
+        )
+        dist.all_gather_into_tensor(gathered, local_shard, group=dp_group)
+
+        for tensor_metadata in metadata[dtype]:
+            name = tensor_metadata["name"]
+            offset = tensor_metadata["offset"]
+            numel = tensor_metadata["numel"]
+            reconstructed = gathered[offset : offset + numel].view(tensor_metadata["shape"])
+            torch.testing.assert_close(reconstructed, oracle_parameters[name])
+
+
+def _distributed_worker(rank: int, rendezvous: str, checkpoint_root: str) -> None:
+    torch.set_num_threads(1)
+    dist.init_process_group(
+        "gloo",
+        init_method=rendezvous,
+        rank=rank,
+        world_size=_WORLD_SIZE,
+    )
+    try:
+        dp2_groups = _new_groups(((0, 2), (1, 3)))
+        sp2_groups = _new_groups(((0, 1), (2, 3)))
+
+        _load_mmap_transform_and_reconstruct(
+            checkpoint_root,
+            dist.group.WORLD,
+            dp_size=4,
+            dp_rank=rank,
+            ep_group=Magi2ParallelGroup(None, world_size=1, rank=0),
+        )
+        dist.barrier()
+
+        dp_ranks, dp_group = _current_group(rank, dp2_groups)
+        sp_ranks, sp_group = _current_group(rank, sp2_groups)
+        _load_mmap_transform_and_reconstruct(
+            checkpoint_root,
+            dp_group,
+            dp_size=2,
+            dp_rank=dp_ranks.index(rank),
+            ep_group=Magi2ParallelGroup(
+                sp_group,
+                world_size=2,
+                rank=sp_ranks.index(rank),
+            ),
+        )
+        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_dlo_mmap_transform_dp_shard_and_allgather_for_dp4_and_dp2sp2() -> None:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        _write_checkpoint(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, temp_dir),
+                nprocs=_WORLD_SIZE,
+                join=True,
+            )

tests/diffusion/models/magi2/test_native_preview.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import cache_dit
+import pytest
+import torch
+
+from vllm_omni.diffusion.cache.cachedit.model_specific import enable_cache_for_magi2
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.diffusion.models.magi2.attention import VarlenHandler
+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,
+    Modality,
+)
+from vllm_omni.diffusion.models.magi2.preview_data_proxy import (
+    Magi2DataProxy,
+    Magi2PreviewDataProxyConfig,
+    ModelInput,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _tiny_config(
+    params_dtype: torch.dtype = torch.float32,
+    *,
+    num_layers: int = 1,
+) -> Magi2PreviewConfig:
+    layer_indices = tuple(range(num_layers))
+    return Magi2PreviewConfig(
+        num_layers=num_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=layer_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=layer_indices,
+        ),
+    )
+
+
+def _initialize_tiny_model(model: Magi2PreviewTransformer, seed: int) -> None:
+    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_()
+
+
+def test_hsdp_policy_shards_individual_preview_layers():
+    model = Magi2PreviewTransformer(_tiny_config(num_layers=2))
+    condition = model._hsdp_shard_conditions[0]
+
+    assert condition("block.layers.0", model.block.layers[0])
+    assert condition("block.layers.1", model.block.layers[1])
+    assert not condition("block", model.block)
+    assert not condition("pre_adapter", model.pre_adapter)
+
+
+def test_data_proxy_keeps_output_layout_request_scoped() -> None:
+    proxy = Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=0))
+
+    def make_input(width: int) -> ModelInput:
+        return ModelInput(
+            x_t=torch.arange(2 * width, dtype=torch.float32).reshape(1, 2, 1, 1, width),
+            audio_x_t=torch.zeros(1, 1, 2),
+            audio_feat_len=torch.tensor([1]),
+            txt_feat=torch.zeros(1, 1, 2),
+            txt_feat_len=torch.tensor([1]),
+            t=torch.tensor([0.5]),
+        )
+
+    first = proxy.process_input(make_input(2))
+    second = proxy.process_input(make_input(3))
+
+    first_video, first_audio = proxy.process_output(first.token_sequence, first.output_layout)
+    second_video, second_audio = proxy.process_output(second.token_sequence, second.output_layout)
+
+    assert first_video.shape == (1, 2, 1, 1, 2)
+    assert second_video.shape == (1, 2, 1, 1, 3)
+    assert first_audio.shape == second_audio.shape == (1, 1, 2)
+
+
+def test_tiny_native_preview_runs_through_nested_cachedit_adapter() -> None:
+    model = Magi2PreviewTransformer(_tiny_config(num_layers=3))
+    _initialize_tiny_model(model, seed=11)
+    pipeline = type("Magi2TestPipeline", (), {"transformer": model})()
+    result = enable_cache_for_magi2(
+        pipeline,
+        DiffusionCacheConfig(
+            Fn_compute_blocks=1,
+            Bn_compute_blocks=0,
+            max_warmup_steps=1,
+            residual_diff_threshold=1.0,
+        ),
+    )
+    result.refresh(pipeline, 4, verbose=False)
+    packed = torch.randn(6, 4)
+    coordinates = torch.ones(6, 9)
+    modalities = torch.tensor(
+        [
+            Modality.VIDEO,
+            Modality.VIDEO,
+            Modality.AUDIO,
+            Modality.AUDIO,
+            Modality.TEXT,
+            Modality.TEXT,
+        ]
+    )
+    cumulative = torch.tensor([0, 6], dtype=torch.int32)
+    varlen = VarlenHandler(cumulative, cumulative, 6, 6)
+
+    execution_counts = [0, 0, 0]
+    hooks = [
+        layer.register_forward_hook(
+            lambda _module, _args, _output, layer_index=index: execution_counts.__setitem__(
+                layer_index, execution_counts[layer_index] + 1
+            )
+        )
+        for index, layer in enumerate(model.block.layers)
+    ]
+    try:
+        with torch.no_grad():
+            first = model(packed, coordinates, modalities, varlen)
+            second = model(packed, coordinates, modalities, varlen)
+        assert torch.isfinite(first).all()
+        assert torch.isfinite(second).all()
+        # The first layer is the configured Fn block and always executes. The
+        # zero-residual repeated input hits the cache on the second call, so
+        # the two middle layers are skipped.
+        assert execution_counts == [2, 1, 1]
+    finally:
+        for hook in hooks:
+            hook.remove()
+        cache_dit.disable_cache(result.targets[0])
+
+    assert not getattr(model.block, "_is_cached", False)
+
+
+def test_tiny_native_preview_matches_pinned_reference_golden() -> None:
+    """Full-model golden from SandAI reference f68a0f9bbccb.
+
+    The reference's unavailable compiler/FA3/MoE CUDA boundaries were replaced
+    by their eager PyTorch equations when generating this tensor.  Keeping the
+    resulting golden local makes this regression independent of that runtime.
+    """
+
+    model = Magi2PreviewTransformer(_tiny_config(torch.bfloat16))
+    with torch.no_grad():
+        trainable_parameters = (parameter for parameter in model.parameters() if parameter.requires_grad)
+        for index, parameter in enumerate(trainable_parameters):
+            values = (torch.arange(parameter.numel(), dtype=torch.float32) % 17 - 8) * 0.002 + (index % 5 - 2) * 0.0001
+            parameter.copy_(values.reshape(parameter.shape).to(parameter.dtype))
+        moe = model.block.layers[0].mlp.moe_mlp
+        moe.router.expert_bias.zero_()
+        moe.router.expert_bias_ema.zero_()
+
+    packed = torch.tensor(
+        [
+            [0.1, -0.2, 0.3, -0.4],
+            [0.5, 0.6, -0.7, 0.8],
+            [-0.9, 1.0, -0.1, 0.2],
+            [0.3, -0.4, 0.5, -0.6],
+            [0.7, 0.8, 0.9, 1.0],
+            [-0.2, 0.4, -0.6, 0.8],
+            [0.9, -0.7, 0.5, -0.3],
+        ]
+    )
+    coordinates = torch.tensor(
+        [
+            [0, 0, 0, 2, 2, 2, 2, 2, 2],
+            [1, 1, 1, 2, 2, 2, 2, 2, 2],
+            [0, 0, 0, 1, 2, 2, 1, 2, 2],
+            [0, 1, 1, 1, 2, 2, 1, 2, 2],
+            [0, 0, 0, 1, 1, 1, 1, 1, 1],
+            [0, 0, 0, 1, 1, 1, 1, 1, 1],
+            [0, 0, 0, 1, 1, 1, 1, 1, 1],
+        ],
+        dtype=torch.float32,
+    )
+    modalities = torch.tensor(
+        [
+            Modality.VIDEO,
+            Modality.VIDEO,
+            Modality.AUDIO,
+            Modality.AUDIO,
+            Modality.TEXT,
+            Modality.TEXT,
+            Modality.TIME,
+        ]
+    )
+    cumulative = torch.tensor([0, 4, 7], dtype=torch.int32)
+    time_tokens = torch.tensor(
+        [
+            [0.1, 0.2, 0.3],
+            [0.4, 0.5, 0.6],
+            [0.7, 0.8, 0.9],
+            [1.0, 0.9, 0.8],
+            [0.7, 0.6, 0.5],
+            [0.4, 0.3, 0.2],
+            [0.1, 0.0, -0.1],
+        ]
+    )
+    expected = torch.tensor(
+        [
+            [-0.093289732933, 0.133922040462, 0.149612516165, 0.069642566144],
+            [-0.118302434683, 0.129068359733, 0.156465008855, 0.084789708257],
+            [-0.114637844265, 0.137497439981, 0.159646511078, 0.088356778026],
+            [-0.104430131614, 0.146924808621, 0.152607530355, 0.092736266553],
+            [0.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0],
+        ]
+    )
+    varlen = VarlenHandler(cumulative, cumulative, 4, 4)
+
+    with torch.no_grad():
+        actual = model(packed, coordinates, modalities, varlen, time_tokens)
+
+    torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)
+
+
+def test_strict_loader_covers_all_keys_and_uses_ema_router_bias() -> None:
+    source = Magi2PreviewTransformer(_tiny_config())
+    _initialize_tiny_model(source, seed=11)
+    source_moe = source.block.layers[0].mlp.moe_mlp
+    source_moe.router.expert_bias.fill_(2.0)
+    source_moe.router.expert_bias_ema.fill_(3.0)
+    checkpoint = [(name, value.detach().clone()) for name, value in source.state_dict().items()]
+
+    target = Magi2PreviewTransformer(_tiny_config())
+    loaded = target.load_weights(checkpoint)
+
+    assert loaded == set(target.state_dict())
+    target_moe = target.block.layers[0].mlp.moe_mlp
+    torch.testing.assert_close(
+        target_moe.router.expert_bias,
+        torch.full_like(target_moe.router.expert_bias, 3.0),
+    )
+    torch.testing.assert_close(
+        target_moe.router.expert_bias_ema, torch.full_like(target_moe.router.expert_bias_ema, 3.0)
+    )

tests/diffusion/models/magi2/test_pipeline_magi2.py

@@ -0,0 +1,631 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from types import MethodType
+from typing import Any
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+from PIL import Image
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
+from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
+from vllm_omni.diffusion.models.magi2.pipeline_magi2 import (
+    MAGI2_AUDIO_SAMPLE_RATE,
+    MAGI2_MODEL_REVISION,
+    Magi2Pipeline,
+    _Magi2StagedComponent,
+    _resolve_checkpoint_root,
+    _validate_native_topology,
+)
+from vllm_omni.diffusion.models.magi2.preview_data_proxy import Magi2DataProxy
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler
+from vllm_omni.diffusion.registry import DiffusionModelRegistry
+from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
+from vllm_omni.errors import OmniClientError
+from vllm_omni.model_extras.registry import (
+    get_extra_body_params,
+    should_preserve_reference_image_size,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+@dataclass
+class _SamplingStub:
+    width: int | None = None
+    height: int | None = None
+    fps: float | None = None
+    frame_rate: float | None = None
+    resolved_frame_rate: float | None = None
+    num_frames: int = 1
+    num_inference_steps: int | None = None
+    num_outputs_per_prompt: int = 1
+    seed: int | None = 42
+    generator: torch.Generator | list[torch.Generator] | None = None
+    extra_args: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class _RequestStub:
+    prompts: list[Any]
+    sampling_params: _SamplingStub
+    num_reqs: int = 1
+
+
+@dataclass
+class _ParallelStub:
+    pipeline_parallel_size: int = 1
+    data_parallel_size: int = 1
+    tensor_parallel_size: int = 1
+    sequence_parallel_size: int = 1
+    ulysses_degree: int = 1
+    ring_degree: int = 1
+    allgather_degree: int = 1
+    cfg_parallel_size: int = 1
+    vae_patch_parallel_size: int = 1
+    text_encoder_tp_size: int = 1
+    enable_expert_parallel: bool = False
+    use_hsdp: bool = False
+
+
+@dataclass
+class _TopologyStub:
+    parallel_config: _ParallelStub = field(default_factory=_ParallelStub)
+    enable_cpu_offload: bool = False
+    enable_layerwise_offload: bool = False
+    enable_distributed_layerwise_offload: bool = False
+    dlo_use_allgather: bool = True
+    quantization_config: object = None
+    cache_backend: str = "none"
+    custom_pipeline_args: dict[str, object] = field(default_factory=dict)
+    additional_config: dict[str, object] = field(default_factory=lambda: {"magi2_allow_unsupported_topology": True})
+
+
+class _FakeNativeRuntime:
+    def __init__(self) -> None:
+        self.calls: list[dict] = []
+
+    def evaluate(self, **kwargs):
+        self.calls.append(kwargs)
+        video = np.zeros((2, 8, 12, 3), dtype=np.uint8)
+        audio = np.zeros((32, 2), dtype=np.float32)
+        return video, audio
+
+
+def _pipeline() -> tuple[Magi2Pipeline, _FakeNativeRuntime]:
+    pipe = Magi2Pipeline.__new__(Magi2Pipeline)
+    nn.Module.__init__(pipe)
+    pipe.deterministic = False
+    pipe.enable_diffusion_pipeline_profiler = False
+    runtime = _FakeNativeRuntime()
+
+    def evaluate(_self, **kwargs):
+        return runtime.evaluate(**kwargs)
+
+    pipe._evaluate_preview = MethodType(evaluate, pipe)
+    return pipe, runtime
+
+
+def _request(prompt, **sampling_overrides):
+    sampling = _SamplingStub()
+    for key, value in sampling_overrides.items():
+        setattr(sampling, key, value)
+    return _RequestStub(prompts=[prompt], sampling_params=sampling)
+
+
+def _checkpoint_tree(root) -> None:
+    files = (
+        "preview/model.safetensors.index.json",
+        "text_encoder/config.json",
+        "text_encoder/model.safetensors.index.json",
+        "vae/Wan2.2_VAE.pth",
+        "turbo_vae/TurboV3-Wan22-TinyShallow_7_7.json",
+        "turbo_vae/checkpoint.ckpt",
+        "stable-audio-open-1.0/model_config.json",
+        "stable-audio-open-1.0/model.safetensors",
+    )
+    for relative in files:
+        path = root / relative
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.touch()
+
+
+def test_official_id_detection_and_metadata():
+    assert is_diffusion_model("sand-ai/MAGI-2-preview")
+    assert is_diffusion_model("https://huggingface.co/sand-ai/MAGI-2-preview")
+    assert resolve_model_class_name("sand-ai/MAGI-2-preview") == "Magi2Pipeline"
+    assert DiffusionModelRegistry._try_load_model_cls("Magi2Pipeline") is Magi2Pipeline
+    metadata = get_diffusion_model_metadata("Magi2Pipeline")
+    assert metadata.supports_multimodal_inputs
+    assert metadata.max_multimodal_image_inputs == 1
+    assert should_preserve_reference_image_size(
+        "Magi2Pipeline",
+        model="sand-ai/MAGI-2-preview",
+    )
+    assert {"seconds", "resolution"} <= get_extra_body_params("Magi2Pipeline")
+
+
+def test_local_preview_signature_detection_without_refiner(tmp_path):
+    _checkpoint_tree(tmp_path)
+    assert is_diffusion_model(str(tmp_path))
+    assert resolve_model_class_name(str(tmp_path)) == "Magi2Pipeline"
+    config = OmniDiffusionConfig(model=str(tmp_path))
+    config.enrich_config()
+    assert config.model_class_name == "Magi2Pipeline"
+    assert config.supports_multimodal_inputs
+
+
+def test_huggingface_url_resolves_pinned_snapshot(tmp_path, monkeypatch):
+    _checkpoint_tree(tmp_path)
+    api = Mock()
+    api.snapshot_download.side_effect = lambda *, repo_id, revision: (
+        str(tmp_path)
+        if (repo_id, revision) == ("sand-ai/MAGI-2-preview", MAGI2_MODEL_REVISION)
+        else pytest.fail("unexpected snapshot request")
+    )
+    monkeypatch.setattr(
+        "vllm.transformers_utils.repo_utils.hf_api",
+        lambda: api,
+    )
+    assert _resolve_checkpoint_root(
+        "https://huggingface.co/sand-ai/MAGI-2-preview",
+        None,
+    ) == str(tmp_path.resolve())
+
+
+def test_forward_uses_native_540p_preview_defaults(monkeypatch):
+    pipe, runtime = _pipeline()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
+        lambda: False,
+    )
+    result = pipe(_request("A fox walks through snow"))
+
+    call = runtime.calls[0]
+    assert call["prompt"] == "A fox walks through snow"
+    assert call["image"] is None
+    assert (call["width"], call["height"]) == (896, 512)
+    assert call["num_inference_steps"] == 100
+    assert result.output["metadata"]["video"]["fps"] == 12.5
+    assert result.output["metadata"]["audio"]["sample_rate"] == MAGI2_AUDIO_SAMPLE_RATE
+    assert result.stage_durations["magi2_preview_e2e"] >= 0
+
+
+def test_forward_maps_272p_i2v_and_output_resize(monkeypatch):
+    pipe, runtime = _pipeline()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
+        lambda: False,
+    )
+    image = Image.new("RGB", (16, 9), "white")
+    prompt = {
+        "prompt": "The first frame begins moving",
+        "multi_modal_data": {"image": image},
+    }
+    result = pipe(
+        _request(
+            prompt,
+            width=448,
+            height=256,
+            num_inference_steps=1,
+            extra_args={
+                "output_width": 6,
+                "output_height": 4,
+            },
+        )
+    )
+
+    call = runtime.calls[0]
+    assert call["image"] is image
+    assert pipe._ensure_figure_token(call["prompt"]) == (
+        "The first frame begins moving\nreference_layer:The first frame refers to <Figure 1>"
+    )
+    assert (call["width"], call["height"]) == (448, 256)
+    assert call["num_inference_steps"] == 1
+    assert result.output["payload"]["video"].shape == (2, 4, 6, 3)
+
+
+def test_forward_rejects_conflicting_resolution_and_sampling_geometry(monkeypatch):
+    pipe, runtime = _pipeline()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
+        lambda: False,
+    )
+
+    with pytest.raises(OmniClientError, match="requires 896x512"):
+        pipe(
+            _request(
+                "A fox walks through snow",
+                width=448,
+                height=256,
+                extra_args={"resolution": "540p"},
+            )
+        )
+    assert not runtime.calls
+
+
+@pytest.mark.parametrize(
+    ("prompt", "extra_args", "message"),
+    [
+        ("", {}, "non-empty prompt"),
+        ("prompt", {"seconds": 5}, "10-second clips"),
+        ("prompt", {"duration": 5}, "10-second clips"),
+        ("prompt", {"resolution": "720p"}, "Unsupported native"),
+        ("prompt", {"resolution": "1080p"}, "Unsupported native"),
+        ("prompt", {"resolution": "540p", "use_refiner": True}, "refiner"),
+        ("prompt", {"output_width": 0, "output_height": 4}, "positive integers"),
+    ],
+)
+def test_forward_rejects_invalid_preview_requests(
+    monkeypatch,
+    prompt,
+    extra_args,
+    message,
+):
+    pipe, runtime = _pipeline()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
+        lambda: False,
+    )
+    with pytest.raises(OmniClientError, match=message):
+        pipe(_request(prompt, extra_args=extra_args))
+    assert not runtime.calls
+
+
+def test_forward_rejects_zero_inference_steps_before_generation(monkeypatch):
+    pipe, runtime = _pipeline()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
+        lambda: False,
+    )
+    with pytest.raises(OmniClientError, match="inference steps must be positive"):
+        pipe(_request("prompt", num_inference_steps=0))
+    assert not runtime.calls
+
+
+def test_forward_stops_peak_monitor_when_initial_synchronize_fails(monkeypatch):
+    pipe, runtime = _pipeline()
+    pipe.enable_diffusion_pipeline_profiler = True
+    monitor = Mock(peak_bytes=0)
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_cuda",
+        lambda: True,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
+        lambda: True,
+    )
+    monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
+    monkeypatch.setattr(torch.accelerator, "synchronize", Mock(side_effect=RuntimeError("sync failed")))
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2._PeakReservedMonitor",
+        Mock(return_value=monitor),
+    )
+
+    with pytest.raises(RuntimeError, match="sync failed"):
+        pipe(_request("A fox walks through snow"))
+
+    monitor.start.assert_called_once_with()
+    monitor.stop.assert_called_once_with()
+    assert not runtime.calls
+
+
+def test_forward_does_not_start_peak_monitor_without_profiler(monkeypatch):
+    pipe, runtime = _pipeline()
+    monitor_factory = Mock(side_effect=AssertionError("unexpected memory monitor"))
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_cuda",
+        lambda: True,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
+        lambda: True,
+    )
+    monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
+    monkeypatch.setattr(torch.accelerator, "synchronize", Mock())
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2._PeakReservedMonitor",
+        monitor_factory,
+    )
+
+    result = pipe(_request("A fox walks through snow"))
+
+    monitor_factory.assert_not_called()
+    assert result.peak_memory_mb == 0.0
+    assert len(runtime.calls) == 1
+
+
+def test_forward_rejects_multiple_images(monkeypatch):
+    pipe, _ = _pipeline()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.current_omni_platform.is_available",
+        lambda: False,
+    )
+    prompt = {
+        "prompt": "animate",
+        "multi_modal_data": {"image": [Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2))]},
+    }
+    with pytest.raises(OmniClientError, match="at most one input image"):
+        pipe(_request(prompt))
+
+
+def test_pipeline_rejects_unknown_initialization_arguments():
+    with pytest.raises(TypeError, match=r"Unexpected MAGI-2.*unknown_option"):
+        Magi2Pipeline(None, unknown_option=True)
+
+
+def _topology_config(**overrides):
+    parallel = _ParallelStub()
+    config = _TopologyStub(parallel_config=parallel)
+    for key, value in overrides.items():
+        if hasattr(parallel, key):
+            setattr(parallel, key, value)
+        else:
+            setattr(config, key, value)
+    return config
+
+
+def test_native_topology_accepts_qualified_tensor_parallelism():
+    _validate_native_topology(_topology_config(tensor_parallel_size=4))
+
+
+def test_native_topology_rejects_nondivisible_tensor_parallelism():
+    config = _topology_config(tensor_parallel_size=3)
+    with pytest.raises(ValueError, match="does not divide"):
+        _validate_native_topology(config)
+
+
+def test_native_topology_accepts_single_device_layerwise_with_cpu_staging():
+    config = _topology_config(
+        enable_layerwise_offload=True,
+        additional_config={},
+    )
+    _validate_native_topology(config)
+
+    config.enable_cpu_offload = True
+    _validate_native_topology(config)
+
+    config.enable_layerwise_offload = False
+    with pytest.raises(ValueError, match="Combine --enable-cpu-offload"):
+        _validate_native_topology(config)
+
+
+def test_native_topology_requires_dlo_rank_local_mode():
+    config = _topology_config(enable_distributed_layerwise_offload=True)
+    with pytest.raises(ValueError, match="dlo-no-use-allgather"):
+        _validate_native_topology(config)
+    config.dlo_use_allgather = False
+    _validate_native_topology(config)
+
+
+def test_native_topology_accepts_four_device_hsdp_cfg_and_vae_layouts():
+    _validate_native_topology(
+        _topology_config(
+            sequence_parallel_size=4,
+            ulysses_degree=4,
+            use_hsdp=True,
+        )
+    )
+    _validate_native_topology(
+        _topology_config(
+            sequence_parallel_size=2,
+            ulysses_degree=2,
+            cfg_parallel_size=2,
+        )
+    )
+    _validate_native_topology(
+        _topology_config(
+            sequence_parallel_size=4,
+            ulysses_degree=4,
+            vae_patch_parallel_size=4,
+        )
+    )
+
+
+def test_native_topology_rejects_cfg_data_parallel_combination():
+    with pytest.raises(ValueError, match="CFG parallelism.*DLO data parallelism"):
+        _validate_native_topology(_topology_config(cfg_parallel_size=2, data_parallel_size=2))
+
+
+def test_cfg_parallel_branch_adapter_preserves_packed_cfg_math():
+    sampler = Magi2PreviewSampler(nn.Identity(), Magi2DataProxy(), device="cpu", dtype=torch.float32)
+    latent = torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 1, 2)
+    audio_latent = torch.arange(6, dtype=torch.float32).reshape(1, 3, 2)
+    positive_text = torch.ones(1, 3, 4)
+    negative_text = torch.full((1, 2, 4), -1.0)
+    packed = sampler.prepare_model_input(
+        latent=latent,
+        audio_latent=audio_latent,
+        txt_feat=positive_text,
+        null_txt_feat=negative_text,
+        t=torch.tensor(500.0),
+        cfg_config=CFGConfig(),
+    )
+    positive, negative = sampler._split_cfg_model_input(packed)
+
+    assert positive.x_t.shape[0] == negative.x_t.shape[0] == 1
+    assert packed.ref_video_feat.shape == (2, 2, 0, 1, 2)
+    assert packed.ref_video_feat.numel() == 0
+    torch.testing.assert_close(positive.txt_feat[:, :3], positive_text)
+    torch.testing.assert_close(negative.txt_feat[:, :2], negative_text)
+
+    video_pos = torch.full_like(latent, 3.0)
+    video_neg = torch.full_like(latent, 1.0)
+    audio_pos = torch.full_like(audio_latent, 4.0)
+    audio_neg = torch.full_like(audio_latent, 2.0)
+    expected = sampler.cfg_velocity(
+        (torch.cat((video_pos, video_neg)), torch.cat((audio_pos, audio_neg))),
+        5.0,
+        7.0,
+        CFGConfig(),
+        latent,
+        audio_latent,
+    )
+    actual = sampler.combine_cfg_noise(
+        (video_pos, audio_pos),
+        (video_neg, audio_neg),
+        1.0,
+        kwargs={
+            "video_txt_guidance_scale": 5.0,
+            "audio_txt_guidance_scale": 7.0,
+            "cfg_config": CFGConfig(),
+            "latent": latent,
+            "audio_latent": audio_latent,
+        },
+    )
+    torch.testing.assert_close(actual[0], expected[0])
+    torch.testing.assert_close(actual[1], expected[1])
+
+
+def test_decode_audio_preserves_batch_for_released_preview_shape(monkeypatch):
+    class ShapeCheckingDecoder(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.seen_shape: tuple[int, ...] | None = None
+
+        def decode(self, latent: torch.Tensor) -> torch.Tensor:
+            self.seen_shape = tuple(latent.shape)
+            if self.seen_shape != (1, 64, 250):
+                raise AssertionError(f"unexpected decoder input shape {self.seen_shape}")
+            return torch.zeros(1, 2, 512_000)
+
+    decoder = ShapeCheckingDecoder()
+    stager = Mock()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.PinnedModuleStager",
+        Mock(return_value=stager),
+    )
+    component = _Magi2StagedComponent(
+        decoder,
+        torch.device("cpu"),
+        pin_memory=False,
+    )
+    pipeline = object.__new__(Magi2Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline._is_output_rank = True
+    pipeline._offload_aux_after_use = True
+    pipeline.device_str = "cpu"
+    pipeline.audio_decoder = component
+
+    def fake_resample(audio: np.ndarray, target_length: int) -> np.ndarray:
+        assert audio.shape == (512_000, 2)
+        assert target_length == 441_000
+        return np.zeros((target_length, 2), dtype=audio.dtype)
+
+    monkeypatch.setattr("scipy.signal.resample", fake_resample)
+    audio = pipeline._decode_audio(torch.zeros(1, 250, 64))
+
+    assert decoder.seen_shape == (1, 64, 250)
+    assert audio is not None
+    assert audio.shape == (441_000, 2)
+    stager.load.assert_called_once_with()
+    stager.offload.assert_called_once_with()
+
+
+class _TinyMagiTransformer(nn.Module):
+    _layerwise_offload_blocks_attrs = ["block"]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.block = nn.ModuleList([nn.Linear(2, 2) for _ in range(3)])
+
+
+class _NoEagerMoveComponent(_Magi2StagedComponent):
+    def to(self, *args, **kwargs):
+        del args, kwargs
+        raise AssertionError("staged MAGI-2 component was moved eagerly")
+
+
+def _offload_pipeline(monkeypatch):
+    stagers = [Mock() for _ in range(4)]
+    stager_factory = Mock(side_effect=stagers)
+    monkeypatch.setattr(
+        "vllm_omni.platforms.current_omni_platform.Stream",
+        Mock(return_value=Mock()),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.magi2.pipeline_magi2.PinnedModuleStager",
+        stager_factory,
+    )
+    pipeline = object.__new__(Magi2Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.transformer = _TinyMagiTransformer()
+    components = [
+        _NoEagerMoveComponent(
+            nn.Linear(2, 2),
+            torch.device("cuda:0"),
+            pin_memory=False,
+        )
+        for _ in range(4)
+    ]
+    (
+        pipeline.text_encoder,
+        pipeline.image_vae,
+        pipeline.video_decoder,
+        pipeline.audio_decoder,
+    ) = components
+    return pipeline, stagers
+
+
+class _FakeLayerwiseHook:
+    def __init__(self) -> None:
+        self._prev_hook = None
+        self.current_slot = 0
+        self.prefetch_layer = Mock()
+        self.get_weights = Mock()
+
+
+def test_dlo_no_allgather_enable_preserves_staged_aux_and_streams_blocks(
+    monkeypatch,
+):
+    from vllm_omni.diffusion.offloader.base import (
+        OffloadConfig,
+        OffloadStrategy,
+    )
+    from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
+        DistributedLayerwiseOffloadBackend,
+    )
+
+    pipeline, stagers = _offload_pipeline(monkeypatch)
+    hooks: list[_FakeLayerwiseHook] = []
+
+    def fake_apply(*args, **kwargs):
+        del args, kwargs
+        hook = _FakeLayerwiseHook()
+        hooks.append(hook)
+        return hook
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.offloader.distributed_layerwise_backend.apply_distributed_block_hook",
+        fake_apply,
+    )
+    backend = DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=False,
+            dlo_use_allgather=False,
+        ),
+        torch.device("cpu"),
+    )
+    backend._allocate_shared_buffers = Mock(return_value=[{}, {}])
+    backend._cleanup_after_loading = Mock()
+    backend._release_mmap_handles = Mock()
+
+    backend.enable(pipeline)
+
+    assert backend.enabled
+    assert backend.dp_size == 1
+    assert len(hooks) == len(pipeline.transformer.block)
+    assert all(stager.offload.call_count == 1 for stager in stagers)
+    assert all(stager.load.call_count == 0 for stager in stagers)
+    hooks[-1].prefetch_layer.assert_called_once_with(
+        slot=hooks[0].current_slot,
+        non_blocking=False,
+    )

tests/diffusion/models/magi2/test_turbo_vae.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.magi2.turbo_vae import (
+    Magi2TurboVAEDecoder,
+    extract_turbo_decoder_state_dict,
+    turbo_unpatchify,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def test_turbo_unpatchify_matches_wan_channel_order():
+    # Four patch channels for one output channel. MAGI's order interleaves
+    # width before height when reconstructing each 2x2 patch.
+    patched = torch.arange(4, dtype=torch.float32).view(1, 4, 1, 1, 1)
+    output = turbo_unpatchify(patched, patch_size=2)
+    expected = torch.tensor([[[[[0.0, 2.0], [1.0, 3.0]]]]])
+    torch.testing.assert_close(output, expected)
+
+
+def test_extract_turbo_decoder_prefers_ema_and_drops_training_heads():
+    decoder_weight = torch.ones(1)
+    checkpoint = {
+        "state_dict": {"module.decoder.old.weight": torch.zeros(1)},
+        "ema_state_dict": {
+            "module.decoder.conv_in.conv.weight": decoder_weight,
+            "module.aligned_feature_projection_heads.0.weight": torch.empty(1),
+        },
+    }
+    assert extract_turbo_decoder_state_dict(checkpoint) == {"decoder.conv_in.conv.weight": decoder_weight}
+
+
+def test_turbo_temporal_tiles_preserve_chunk_overlap_contract():
+    class RepeatDecoder(nn.Module):
+        def forward(self, value: torch.Tensor, *, is_first_chunk: bool) -> torch.Tensor:
+            del is_first_chunk
+            return value.repeat_interleave(4, dim=2)
+
+    decoder = object.__new__(Magi2TurboVAEDecoder)
+    nn.Module.__init__(decoder)
+    decoder.z_dim = 1
+    decoder.first_chunk_size = 3
+    decoder.step_size = 2
+    decoder.temporal_compression_ratio = 4
+    decoder.decoder = RepeatDecoder()
+    decoder.use_tiling = False
+    decoder.register_buffer("latent_mean", torch.zeros(1), persistent=False)
+    decoder.register_buffer("latent_std", torch.ones(1), persistent=False)
+
+    latent = torch.arange(7, dtype=torch.float32).view(1, 1, 7, 1, 1)
+    actual = decoder.decode(latent)
+    expected = latent.repeat_interleave(4, dim=2)
+
+    torch.testing.assert_close(actual, expected)
+
+    decoder.latent_mean = torch.zeros(1, device="meta")
+    decoder.latent_std = torch.ones(1, device="meta")
+    prepared, _ = decoder._prepare_latent(latent)
+    assert prepared.device.type == "meta"