[Model] Add MiniMax H3 diffusion support (#5691)

tests/diffusion/models/minimax_h3/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -0,0 +1,375 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from multiprocessing.reduction import ForkingPickler
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_pipeline_import_registry_and_component_discovery():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.registry import (
+        _DIFFUSION_MODELS,
+        _DIFFUSION_POST_PROCESS_FUNCS,
+    )
+
+    assert _DIFFUSION_MODELS["MiniMaxH3Pipeline"] == (
+        "minimax_h3",
+        "pipeline_minimax_h3",
+        "MiniMaxH3Pipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] == "get_minimax_h3_post_process_func"
+    assert MiniMaxH3Pipeline._dit_modules == ["transformer"]
+    assert MiniMaxH3Pipeline._encoder_modules == ["text_encoder"]
+    assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
+
+
+def test_joint_postprocess_is_multiprocessing_picklable():
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        get_minimax_h3_post_process_func,
+    )
+
+    postprocess = get_minimax_h3_post_process_func(SimpleNamespace())
+    postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess))
+    video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
+    audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6)
+
+    result = postprocess((video, audio), output_type="np")
+
+    assert isinstance(result["video"], list)
+    assert result["video"][0].shape == (2, 4, 5, 3)
+    np.testing.assert_array_equal(result["audio"], audio.numpy())
+    assert result["audio_sample_rate"] == 32000
+    assert result["fps"] == 24
+
+
+def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    od_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(cfg_parallel_size=2),
+    )
+    with pytest.raises(ValueError, match="CFG-distilled"):
+        MiniMaxH3Pipeline(od_config=od_config)
+
+
+def test_shape_contract_matches_three_reference_tasks():
+    from vllm_omni.diffusion.models.minimax_h3.time_request import (
+        MINIMAX_H3_SHAPE_PLANNER,
+        minimax_h3_align_frame_count,
+    )
+
+    assert minimax_h3_align_frame_count(round(8.7 * 24)) == 209
+    assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(209) == 62
+    assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(8.7) == 348
+    assert minimax_h3_align_frame_count(round(5.0 * 24)) == 124
+    assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(124) == 37
+    assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(5.0) == 200
+    assert MINIMAX_H3_SHAPE_PLANNER.frame_count_from_video_latent_t(62) == 209
+
+
+def test_shifted_sigma_schedule_matches_reference_values():
+    from vllm_omni.diffusion.models.minimax_h3.time_request import (
+        minimax_h3_time_shift_sigmas,
+    )
+
+    sigmas = minimax_h3_time_shift_sigmas(num_steps=5, shift_scale=12.0)
+
+    assert sigmas == pytest.approx(
+        [1.0, 0.9729729891, 0.9230769277, 0.8000000119, 0.0],
+        abs=1e-7,
+    )
+
+
+def test_reference_image_resize_contract():
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _reference_image_shape,
+    )
+
+    assert _reference_image_shape(Image.new("RGB", (1080, 1440))) == (
+        2048,
+        2720,
+    )
+    with pytest.raises(ValueError, match="aspect ratio"):
+        _reference_image_shape(Image.new("RGB", (100, 501)))
+
+
+def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    expected = torch.ones(2, 3)
+    encoder.encode_ids = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    pixel_values = torch.ones(1, 4)
+    image_grid_thw = torch.tensor([[1, 1, 1]])
+
+    actual = encoder(
+        input_ids,
+        pixel_values=pixel_values,
+        image_grid_thw=image_grid_thw,
+    )
+
+    assert actual is expected
+    encoder.encode_ids.assert_called_once_with(
+        input_ids,
+        pixel_values=pixel_values,
+        image_grid_thw=image_grid_thw,
+    )
+
+
+def test_encoder_forward_forwards_video_inputs():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
+    torch.nn.Module.__init__(encoder)
+    expected = torch.ones(2, 3)
+    encoder.encode_ids = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    pixel_values_videos = torch.ones(1, 4)
+    video_grid_thw = torch.tensor([[2, 1, 1]])
+
+    actual = encoder(
+        input_ids,
+        pixel_values_videos=pixel_values_videos,
+        video_grid_thw=video_grid_thw,
+    )
+
+    assert actual is expected
+    encoder.encode_ids.assert_called_once_with(
+        input_ids,
+        pixel_values=None,
+        image_grid_thw=None,
+        pixel_values_videos=pixel_values_videos,
+        video_grid_thw=video_grid_thw,
+    )
+
+
+def test_reference_video_shape_uses_h3_adapt_shape_policy():
+    from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+        _reference_video_shape,
+    )
+
+    assert _reference_video_shape(1280, 720) == (1344, 768)
+    assert _reference_video_shape(3844, 2160) == (1344, 768)
+
+
+def test_text_encoder_stub_constructs_without_group_or_weights():
+    from vllm_omni.diffusion.models.minimax_h3.encoder import (
+        MiniMaxH3Qwen3VLEncoder,
+    )
+
+    encoder = MiniMaxH3Qwen3VLEncoder(
+        "/nonexistent/text_encoder",
+        device=torch.device("cpu"),
+        load_model=False,
+        encoder_group=None,
+    )
+    assert not encoder.is_loaded
+    assert encoder.tp_size == 1
+    # The stub has no parameters, so it never contributes to the runner's
+    # strict missing-parameter check on non-encoder ranks.
+    assert list(encoder.named_parameters()) == []
+
+
+def test_no_offload_keeps_text_encoder_resident():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+    input_ids = torch.tensor([1, 2])
+
+    actual = pipeline._encode_text_hidden(input_ids, {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_not_called()
+
+
+def test_model_offload_uses_hooked_text_encoder_call():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+    )
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder = Mock(return_value=expected)
+    input_ids = torch.tensor([1, 2])
+    vision_kwargs = {"pixel_values": torch.ones(1, 4)}
+
+    actual = pipeline._encode_text_hidden(input_ids, vision_kwargs)
+
+    assert actual is expected
+    pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs)
+    pipeline.text_encoder.load_to_device.assert_not_called()
+
+
+def test_layerwise_offload_releases_text_encoder():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=True,
+    )
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+
+
+def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+    class FakeRemote(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = torch.nn.Linear(1, 1).half()
+
+    monkeypatch.setattr(
+        vae_module,
+        "_load_component_config",
+        lambda _path: {
+            "latent_channels": 1,
+            "latents_mean": [0.0],
+            "latents_std": [1.0],
+        },
+    )
+    monkeypatch.setattr(
+        vae_module,
+        "_load_remote_component",
+        lambda _path, _config: FakeRemote(),
+    )
+
+    video_vae = vae_module.MiniMaxH3VideoVAE(
+        "unused",
+        device=torch.device("cpu"),
+    )
+
+    assert next(video_vae.parameters()).dtype == torch.float32
+
+
+def test_video_vae_encode_uses_configured_parallel_tiling():
+    from vllm_omni.diffusion.models.minimax_h3.vae import (
+        MiniMaxH3VideoVAE,
+    )
+
+    class FakeModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.ones(1))
+            self.parallel_tiling = True
+            self.encode_calls = []
+
+        def encode_videos(self, frames, *, use_fp16_latent):
+            assert self.parallel_tiling
+            self.encode_calls.append((frames, use_fp16_latent))
+            return [torch.ones(1, 1, 2, 2, 2)]
+
+    video_vae = object.__new__(MiniMaxH3VideoVAE)
+    torch.nn.Module.__init__(video_vae)
+    video_vae.model = FakeModel()
+    video_vae.config_dict = {
+        "latent_channels": 1,
+        "latents_mean": [0.0],
+        "latents_std": [1.0],
+    }
+
+    rows, shape = video_vae.encode_video("frames")
+
+    assert video_vae.model.parallel_tiling
+    assert video_vae.model.encode_calls == [("frames", True)]
+    assert rows.shape == (2, 4)
+    assert shape == (2, 2, 2)
+
+
+def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    prepared = [
+        {"prepared_path": "video-1.mp4"},
+        {"prepared_path": "video-2.mp4"},
+    ]
+
+    class FakeVideoVAE:
+        def __init__(self):
+            self.calls = []
+
+        def is_distributed_enabled(self):
+            return True
+
+        def encode_video(self, frames):
+            self.calls.append(frames)
+            index = len(self.calls)
+            return torch.full((1, 2), index, dtype=torch.float32), (index, 2, 3)
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.video_vae = FakeVideoVAE()
+
+    monkeypatch.setattr(
+        pipeline_module,
+        "_dit_rank_world",
+        lambda: ("dit-group", 1, 4),
+    )
+
+    def fake_broadcast_object_list(values, *, src, group, device):
+        assert values == [None]
+        assert (src, group, device) == (0, "dit-group", torch.device("cpu"))
+        values[0] = prepared
+
+    monkeypatch.setattr(
+        pipeline_module.dist,
+        "broadcast_object_list",
+        fake_broadcast_object_list,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "load_video_frames",
+        lambda path: f"frames:{path}",
+    )
+
+    rows, shapes = pipeline._encode_video_conditions(None, count=2)
+
+    assert pipeline.video_vae.calls == [
+        "frames:video-1.mp4",
+        "frames:video-2.mp4",
+    ]
+    torch.testing.assert_close(
+        rows,
+        torch.tensor([[1.0, 1.0], [2.0, 2.0]]),
+    )
+    assert shapes == [(1, 2, 3), (2, 2, 3)]

tests/diffusion/models/minimax_h3/test_minimax_h3_e2e.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import numpy as np
+import pytest
+
+pytestmark = [
+    pytest.mark.local_model,
+    pytest.mark.cuda,
+    pytest.mark.diffusion,
+    pytest.mark.parallel,
+]
+
+MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
+
+
+@pytest.mark.skipif(
+    not os.environ.get(MODEL_ENV),
+    reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
+)
+def test_minimax_h3_t2va_ulysses8_smoke():
+    import torch
+
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    if torch.accelerator.device_count() < 8:
+        pytest.skip("MiniMax H3 Ulysses smoke requires eight GPUs")
+
+    engine = Omni(
+        model=os.environ[MODEL_ENV],
+        parallel_config=DiffusionParallelConfig(ulysses_degree=8),
+        trust_remote_code=True,
+        enable_cpu_offload=True,
+        enforce_eager=True,
+    )
+    try:
+        outputs = engine.generate(
+            "A quiet cinematic night scene with matching ambient sound.",
+            OmniDiffusionSamplingParams(
+                height=256,
+                width=448,
+                num_frames=29,
+                fps=24,
+                num_inference_steps=2,
+                seed=42,
+                output_type="np",
+                extra_args={
+                    "task": "t2va",
+                    "flow_shift": 12.0,
+                    "audio_flow_shift": 3.0,
+                },
+            ),
+            use_tqdm=False,
+        )
+    finally:
+        engine.close()
+
+    assert len(outputs) == 1
+    frames = np.asarray(outputs[0].images[0])
+    assert frames.shape == (39, 256, 448, 3)
+    multimodal = outputs[0].multimodal_output
+    assert multimodal is not None
+    assert np.asarray(multimodal["audio"]).shape[1] == 2
+    assert multimodal["audio_sample_rate"] == 32000
+    assert multimodal["fps"] == 24

tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_video_patchify_round_trip_preserves_values():
+    from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
+        minimax_h3_patchify_video_latent,
+        minimax_h3_unpatchify_video_tokens,
+    )
+
+    latent = torch.arange(2 * 3 * 2 * 4 * 6).reshape(2, 3, 2, 4, 6)
+    rows = minimax_h3_patchify_video_latent(
+        latent,
+        patch_size=(1, 2, 2),
+    )
+    restored = minimax_h3_unpatchify_video_tokens(
+        rows,
+        latent_shape=(2, 2, 3, 3),
+        patch_size=(1, 2, 2),
+    )
+
+    torch.testing.assert_close(restored, latent)
+
+
+def test_audio_pack_round_trip_preserves_channel_major_order():
+    from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
+        minimax_h3_pack_audio_latent,
+        minimax_h3_unpack_audio_tokens,
+    )
+
+    latent = torch.arange(2 * 4 * 5).reshape(2, 4, 5)
+    rows = minimax_h3_pack_audio_latent(latent)
+    restored = minimax_h3_unpack_audio_tokens(
+        rows,
+        audio_t=10,
+        audio_channel=2,
+    )
+
+    torch.testing.assert_close(restored, latent)
+
+
+def test_t2va_and_fl2va_packing_keep_update_rows_separate():
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence,
+    )
+
+    common = dict(
+        text_len=4,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+    )
+    t2va = minimax_h3_packed_sequence(
+        **common,
+        include_keyframe_cond=False,
+    )
+    fl2va = minimax_h3_packed_sequence(
+        **common,
+        include_keyframe_cond=True,
+        keyframe_frame_indices=[0],
+        frame_count=5,
+    )
+
+    assert int(t2va["seq_len"]) == 64
+    assert t2va["img_pos"].numel() == 12
+    assert t2va["update_mask"].all()
+    assert fl2va["img_pos"].numel() == 18
+    assert fl2va["update_mask"].sum().item() == 12
+    assert (~fl2va["update_mask"]).sum().item() == 6
+    assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
+
+
+def test_ref2va_packing_tracks_video_and_audio_update_masks():
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence_ref2va_blocks,
+    )
+
+    packed = minimax_h3_packed_sequence_ref2va_blocks(
+        text_len=4,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+        ref_blocks=[
+            {"kind": "image", "latent_h": 4, "latent_w": 4},
+            {"kind": "audio", "ref_audio_t": 2},
+        ],
+    )
+
+    assert int(packed["seq_len"]) == 64
+    assert packed["img_pos"].numel() == 16
+    assert packed["update_mask"].sum().item() == 12
+    assert packed["audio_pos"].numel() == 10
+    assert packed["audio_update_mask"].sum().item() == 6
+    assert packed["cu_seqlens"].tolist() == [0, 30, 64]
+
+
+def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
+    from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
+        minimax_h3_audio_cond_noise_aug_rows,
+        minimax_h3_imgvid_cond_noise_aug_rows,
+    )
+
+    image_rows = torch.randn(4, 96)
+    audio_rows = torch.randn(6, 32)
+
+    torch.testing.assert_close(
+        minimax_h3_imgvid_cond_noise_aug_rows(
+            image_rows,
+            condition_shapes=[(1, 4, 4)],
+            target_latent_t=2,
+            imgvid_cond_num_frames=1,
+            seed=42,
+            noise_aug=1.0,
+        ),
+        image_rows,
+    )
+    torch.testing.assert_close(
+        minimax_h3_audio_cond_noise_aug_rows(
+            audio_rows,
+            condition_audio_t=[3],
+            seed=42,
+            noise_aug=1.0,
+        ),
+        audio_rows,
+    )
+
+    first = minimax_h3_audio_cond_noise_aug_rows(
+        audio_rows,
+        condition_audio_t=[3],
+        seed=42,
+        noise_aug=0.25,
+    )
+    second = minimax_h3_audio_cond_noise_aug_rows(
+        audio_rows,
+        condition_audio_t=[3],
+        seed=42,
+        noise_aug=0.25,
+    )
+    torch.testing.assert_close(first, second)

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_grouped_qkv_checkpoint_reorder():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        _reorder_grouped_qkv_to_qkv,
+    )
+
+    # Two groups with rows [q, k, v] become [q0, q1, k0, k1, v0, v1].
+    grouped = torch.arange(6, dtype=torch.float32).reshape(6, 1)
+    reordered = _reorder_grouped_qkv_to_qkv(
+        grouped,
+        num_query_groups=2,
+        heads_per_group=1,
+        head_dim=1,
+    )
+
+    assert reordered[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
+
+
+def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
+    from cache_dit import ForwardPattern
+
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTModel,
+    )
+
+    assert MiniMaxH3DiTModel._repeated_blocks == ["MiniMaxH3DiTBlock"]
+    assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
+    assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
+    assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
+    assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.token_refiner = nn.Module()
+    model.token_refiner.blocks = nn.ModuleList([nn.Linear(4, 4)])
+    model.final_layer = nn.Linear(4, 4)
+
+    matched = [
+        name
+        for name, module in model.named_modules()
+        if any(condition(name, module) for condition in MiniMaxH3DiTModel._hsdp_shard_conditions)
+    ]
+    assert matched == ["blocks.0", "blocks.1"]
+
+
+def test_packed_attention_is_a_regional_compile_boundary():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
+
+
+@pytest.mark.parametrize(
+    ("tp_size", "message"),
+    [
+        (3, "num_attention_heads"),
+        (5, "num_attention_heads"),
+    ],
+)
+def test_tp_rejects_non_divisible_head_counts(tp_size, message):
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTArchConfig,
+        MiniMaxH3DiTModel,
+    )
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    with pytest.raises(ValueError, match=message):
+        model._validate_tp_config(
+            arch=MiniMaxH3DiTArchConfig(),
+            tp_size=tp_size,
+        )
+
+
+def test_tp_accepts_checkpoint_supported_sizes():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTArchConfig,
+        MiniMaxH3DiTModel,
+    )
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    arch = MiniMaxH3DiTArchConfig()
+    for tp_size in (1, 2, 4, 7):
+        model._validate_tp_config(arch=arch, tp_size=tp_size)