Merge remote-tracking branch 'origin/main' into qwen-image-2.1

# Conflicts:
#	vllm_omni/diffusion/executor/multiproc_executor.py
#	vllm_omni/diffusion/offloader/sequential_backend.py

tests/diffusion/models/anima/__init__.py

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

tests/diffusion/models/anima/test_anima.py

@@ -0,0 +1,454 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import json
+from pathlib import Path
+from shutil import rmtree
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from tests.model_tests.diffusion.anima_builder import CHECKPOINT_FILENAME, tiny_anima_builder
+from vllm_omni.diffusion.attention import selector
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
+from vllm_omni.diffusion.models.anima import pipeline_anima
+from vllm_omni.diffusion.models.anima.anima_text_conditioner import AnimaTextConditioner
+from vllm_omni.diffusion.models.anima.anima_transformer import AnimaTransformer3DModel
+from vllm_omni.diffusion.models.anima.pipeline_anima import AnimaPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "offload_args",
+    [
+        {"enable_cpu_offload": True},
+        {"enable_layerwise_offload": True},
+        {"enable_distributed_layerwise_offload": True},
+        {"diffusion_offload_config": {"mode": "module", "components": ["dit"]}},
+        {
+            "diffusion_offload_config": {
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "rank-local"}},
+            }
+        },
+    ],
+)
+def test_native_anima_rejects_offload(offload_args):
+    config = OmniDiffusionConfig(model="anima.safetensors", **offload_args)
+    with pytest.raises(NotImplementedError, match="offload"):
+        AnimaPipeline(od_config=config, device=torch.device("cpu"))
+
+
+def test_native_anima_initializes_without_offload():
+    config = OmniDiffusionConfig(model="anima.safetensors")
+    AnimaPipeline(od_config=config, device=torch.device("cpu"))
+
+
+def test_anima_registration() -> None:
+    """The native Anima pipeline is registered and importable."""
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    assert DiffusionModelRegistry._try_load_model_cls("AnimaPipeline") is not None
+
+
+@pytest.mark.parametrize("model_class_name", ["AnimaPipeline", "AnimaModularPipeline"])
+def test_enrich_config_native_anima_checkpoint(tmp_path: Path, model_class_name: str) -> None:
+    """Native and reference Anima class names resolve to AnimaPipeline."""
+    dummy_checkpoint = tmp_path / "model.safetensors"
+    dummy_checkpoint.write_text("dummy")
+
+    config = OmniDiffusionConfig(
+        model=str(dummy_checkpoint),
+        model_class_name=model_class_name,
+    )
+    config.enrich_config()
+
+    assert config.diffusion_load_format == "default"
+    assert config.model_class_name == "AnimaPipeline"
+    assert config.diffusers_pipeline_cls is None
+
+
+def test_native_anima_component_paths_use_custom_pipeline_args() -> None:
+    """Anima resolves shared and component-specific paths from native args."""
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.od_config = SimpleNamespace(
+        custom_pipeline_args={
+            "components_path": "/tmp/anima-components",
+            "vae_path": "/tmp/anima-vae",
+            "text_encoder_model": "/tmp/anima-text-encoder",
+        }
+    )
+
+    assert pipeline._component_path("vae", "/tmp/default") == "/tmp/anima-vae"
+    assert pipeline._component_path("text_encoder", "/tmp/default") == "/tmp/anima-text-encoder"
+    assert pipeline._component_path("tokenizer", "/tmp/default") == "/tmp/default"
+
+
+def test_native_anima_converts_original_cosmos_transformer_keys() -> None:
+    """Original Cosmos checkpoint names map to native module names."""
+    converted = AnimaPipeline._convert_original_transformer_state_dict(
+        {
+            "net.x_embedder.proj.1.weight": "patch",
+            "net.blocks.0.self_attn.q_proj.weight": "q",
+            "net.blocks.0.self_attn.q_norm.weight": "q_norm",
+            "net.blocks.0.mlp.layer1.weight": "mlp",
+            "net.final_layer.linear.weight": "out",
+            "net.accum_iteration": "drop",
+        }
+    )
+
+    assert converted == {
+        "patch_embed.proj.weight": "patch",
+        "transformer_blocks.0.attn1.to_q.weight": "q",
+        "transformer_blocks.0.attn1.norm_q.weight": "q_norm",
+        "transformer_blocks.0.ff.net.0.proj.weight": "mlp",
+        "proj_out.weight": "out",
+    }
+
+
+def test_native_anima_resolves_vae_scale_factor_from_loaded_vae() -> None:
+    """The loaded VAE determines the spatial scale factor."""
+    from vllm_omni.diffusion.models.anima.pipeline_anima import _anima_vae_scale_factor_from_vae
+
+    vae = SimpleNamespace(config=SimpleNamespace(spatial_compression_ratio=16))
+
+    assert _anima_vae_scale_factor_from_vae(vae) == 16
+
+
+@pytest.mark.parametrize("use_config_file", [False, True])
+def test_native_anima_loads_synthetic_checkpoint(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, use_config_file: bool
+) -> None:
+    """Native modules load exact weights from an Anima safetensors checkpoint."""
+    tiny_transformer_config = {
+        "in_channels": 1,
+        "out_channels": 1,
+        "num_attention_heads": 1,
+        "attention_head_dim": 12,
+        "num_layers": 1,
+        "mlp_ratio": 1.0,
+        "text_embed_dim": 4,
+        "adaln_lora_dim": 3,
+        "max_size": (1, 2, 2),
+        "patch_size": (1, 1, 1),
+        "rope_scale": (1.0, 1.0, 1.0),
+        "concat_padding_mask": True,
+        "extra_pos_embed_type": None,
+    }
+    tiny_text_conditioner_config = {
+        "source_dim": 4,
+        "target_dim": 4,
+        "model_dim": 4,
+        "num_layers": 1,
+        "num_attention_heads": 1,
+        "target_vocab_size": 8,
+        "min_sequence_length": 4,
+    }
+    if use_config_file:
+        (tmp_path / "anima.json").write_text(
+            json.dumps({"transformer": tiny_transformer_config, "text_conditioner": tiny_text_conditioner_config})
+        )
+    else:
+        monkeypatch.setattr(pipeline_anima, "ANIMA_TRANSFORMER_CONFIG", tiny_transformer_config)
+        monkeypatch.setattr(pipeline_anima, "ANIMA_TEXT_CONDITIONER_CONFIG", tiny_text_conditioner_config)
+
+    monkeypatch.setattr(selector, "_cached_get_backend_cls", lambda *_args, **_kwargs: SDPABackend)
+    transformer = AnimaTransformer3DModel(**tiny_transformer_config)
+    text_conditioner = AnimaTextConditioner(**tiny_text_conditioner_config)
+    transformer_state = {name: tensor.detach().clone() for name, tensor in transformer.state_dict().items()}
+    text_conditioner_state = {name: tensor.detach().clone() for name, tensor in text_conditioner.state_dict().items()}
+    checkpoint_state = {
+        **{f"transformer.{name}": tensor for name, tensor in transformer_state.items()},
+        **{f"text_conditioner.{name}": tensor for name, tensor in text_conditioner_state.items()},
+    }
+
+    checkpoint_path = tmp_path / "anima.safetensors"
+    save_file(checkpoint_state, str(checkpoint_path))
+
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.od_config = SimpleNamespace(model=str(checkpoint_path), dtype=torch.float32)
+    pipeline.device = torch.device("cpu")
+
+    def assert_loaded(loaded_transformer, loaded_text_conditioner):
+        for name, tensor in transformer_state.items():
+            assert torch.equal(loaded_transformer.state_dict()[name], tensor)
+        for name, tensor in text_conditioner_state.items():
+            assert torch.equal(loaded_text_conditioner.state_dict()[name], tensor)
+
+    loaded_transformer, loaded_text_conditioner = pipeline._load_native_denoiser_components(dict(checkpoint_state))
+    assert_loaded(loaded_transformer, loaded_text_conditioner)
+
+    loaded_transformer, loaded_text_conditioner = pipeline._load_native_denoiser_components()
+    assert_loaded(loaded_transformer, loaded_text_conditioner)
+
+    if use_config_file:
+        tiny_transformer_config["num_layers"] = 2
+        (tmp_path / "anima.json").write_text(
+            json.dumps({"transformer": tiny_transformer_config, "text_conditioner": tiny_text_conditioner_config})
+        )
+        with pytest.raises(RuntimeError, match="Missing key"):
+            pipeline._load_native_denoiser_components()
+
+
+def test_tiny_anima_generates_images(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Load every saved component and exercise CFG, decoding, and multiple outputs on CPU."""
+    monkeypatch.setattr(selector, "_cached_get_backend_cls", lambda *_args, **_kwargs: SDPABackend)
+    monkeypatch.setattr(current_omni_platform, "is_available", lambda: False)
+    monkeypatch.setattr(
+        DistributedAutoencoderKLQwenImage,
+        "init_distributed",
+        lambda self: setattr(self, "distributed_executor", SimpleNamespace(parallel_size=1)),
+    )
+    model_dir = tiny_anima_builder()
+    try:
+        config = OmniDiffusionConfig(
+            model=str(Path(model_dir) / CHECKPOINT_FILENAME), model_class_name="AnimaPipeline", dtype=torch.float32
+        )
+        pipeline = AnimaPipeline(od_config=config, device=torch.device("cpu"))
+        pipeline.load_weights()
+        pipeline.eval()
+        outputs = []
+        with torch.inference_mode():
+            for _ in range(2):
+                request = OmniDiffusionRequest(
+                    prompt="Dummy prompt",
+                    request_id="tiny-anima",
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=32,
+                        width=32,
+                        num_inference_steps=2,
+                        guidance_scale=4.0,
+                        num_outputs_per_prompt=2,
+                        generator=torch.Generator(device="cpu").manual_seed(42),
+                    ),
+                )
+                outputs.append(pipeline(DiffusionRequestBatch([request]))[0].output)
+        assert outputs[0].shape == (2, 3, 32, 32)
+        assert torch.isfinite(outputs[0]).all()
+        assert torch.equal(outputs[0], outputs[1])
+        assert not torch.equal(outputs[0][0], outputs[0][1])
+    finally:
+        rmtree(model_dir)
+
+
+def test_native_anima_profiler_setup_is_deferred_and_idempotent(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Profiler targets are wrapped once, after deferred components are loaded."""
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.od_config = SimpleNamespace(enable_diffusion_pipeline_profiler=True)
+    pipeline._profiler_initialized = False
+    pipeline.vae = SimpleNamespace(decode=lambda: "decoded")
+    pipeline.text_encoder = SimpleNamespace(forward=lambda: "encoded")
+
+    monkeypatch.setattr(current_omni_platform, "is_available", lambda: False)
+    pipeline._setup_profiler()
+    wrapped_decode = pipeline.vae.decode
+    wrapped_text_encoder = pipeline.text_encoder.forward
+    pipeline._setup_profiler()
+
+    assert pipeline.vae.decode is wrapped_decode
+    assert pipeline.text_encoder.forward is wrapped_text_encoder
+    assert pipeline.vae.decode() == "decoded"
+    assert pipeline.text_encoder.forward() == "encoded"
+    assert "AnimaPipeline.vae.decode" in pipeline.stage_durations
+    assert "AnimaPipeline.text_encoder.forward" in pipeline.stage_durations
+
+
+def _make_anima_forward_probe():
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.vae_scale_factor = 8
+    pipeline.transformer = SimpleNamespace(dtype=torch.float32)
+    pipeline.text_encoder = SimpleNamespace(dtype=torch.float32)
+    pipeline._current_timestep = None
+    pipeline._num_timesteps = 0
+    pipeline._guidance_scale = 0.0
+    captured = {}
+
+    def encode_prompt(**_kwargs):
+        return {
+            "qwen_prompt_embeds": torch.zeros(1, 2, 4),
+            "qwen_attention_mask": torch.ones(1, 2),
+            "t5_input_ids": torch.ones(1, 2, dtype=torch.long),
+            "t5_attention_mask": torch.ones(1, 2),
+            "negative_qwen_prompt_embeds": torch.zeros(1, 2, 4),
+            "negative_qwen_attention_mask": torch.ones(1, 2),
+            "negative_t5_input_ids": torch.ones(1, 2, dtype=torch.long),
+            "negative_t5_attention_mask": torch.ones(1, 2),
+        }
+
+    def condition_prompt_embeds(**_kwargs):
+        return torch.zeros(1, 2, 4)
+
+    def prepare_latents(**kwargs):
+        captured["prepare_latents"] = kwargs
+        return torch.zeros(1, 16, 1, kwargs["height"] // 8, kwargs["width"] // 8)
+
+    def diffuse(**kwargs):
+        captured["diffuse"] = kwargs
+        return kwargs["latents"]
+
+    pipeline.encode_prompt = encode_prompt
+    pipeline.condition_prompt_embeds = condition_prompt_embeds
+    pipeline.prepare_latents = prepare_latents
+    pipeline.prepare_timesteps = lambda **_kwargs: (torch.ones(1), 1)
+    pipeline.diffuse = diffuse
+
+    def decode_latents(latents, output_type="pil"):
+        captured["output_type"] = output_type
+        return DiffusionOutput(output=latents)
+
+    pipeline.decode_latents = decode_latents
+    return pipeline, captured
+
+
+def test_native_anima_forward_uses_official_default_resolution() -> None:
+    """Forward uses Anima's reference resolution when none is requested."""
+    pipeline, captured = _make_anima_forward_probe()
+    req = OmniDiffusionRequest(
+        prompt="a red cube",
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_id="anima-defaults",
+    )
+
+    pipeline.forward(DiffusionRequestBatch([req]))
+
+    assert captured["prepare_latents"]["height"] == 1024
+    assert captured["prepare_latents"]["width"] == 1024
+
+
+def test_native_anima_forward_honors_sampling_output_type() -> None:
+    """Request-level output type controls Anima decoding."""
+    pipeline, captured = _make_anima_forward_probe()
+    req = OmniDiffusionRequest(
+        prompt="a red cube",
+        sampling_params=OmniDiffusionSamplingParams(output_type="latent"),
+        request_id="anima-output-type",
+    )
+
+    outputs = pipeline.forward(DiffusionRequestBatch([req]))
+
+    assert len(outputs) == 1
+    assert isinstance(outputs[0], DiffusionOutput)
+    assert outputs[0].output.shape == (1, 16, 1, 128, 128)
+    assert captured["output_type"] == "latent"
+
+
+@pytest.mark.parametrize("request_count", [0, 2])
+def test_native_anima_forward_rejects_non_single_request_batches(request_count: int) -> None:
+    pipeline, captured = _make_anima_forward_probe()
+    requests = [
+        OmniDiffusionRequest(
+            prompt="a red cube",
+            sampling_params=OmniDiffusionSamplingParams(),
+            request_id=f"anima-{index}",
+        )
+        for index in range(request_count)
+    ]
+
+    with pytest.raises(ValueError, match="supports one request per forward"):
+        pipeline.forward(DiffusionRequestBatch(requests))
+
+    assert captured == {}
+
+
+def test_native_anima_explicit_guidance_scale_drives_cfg_multiplier() -> None:
+    """An explicit guidance scale controls native classifier-free guidance."""
+    pipeline, captured = _make_anima_forward_probe()
+    req = OmniDiffusionRequest(
+        prompt="a red cube",
+        sampling_params=OmniDiffusionSamplingParams(guidance_scale=5.0),
+        request_id="anima-guidance",
+    )
+
+    pipeline.forward(DiffusionRequestBatch([req]))
+
+    assert captured["diffuse"]["do_true_cfg"] is True
+    assert captured["diffuse"]["true_cfg_scale"] == 5.0
+
+
+def test_native_anima_true_cfg_scale_overrides_guidance_multiplier() -> None:
+    """An explicit true-CFG scale overrides the guidance-scale multiplier."""
+    pipeline, captured = _make_anima_forward_probe()
+    req = OmniDiffusionRequest(
+        prompt="a red cube",
+        sampling_params=OmniDiffusionSamplingParams(guidance_scale=5.0, true_cfg_scale=3.0),
+        request_id="anima-true-cfg",
+    )
+
+    pipeline.forward(DiffusionRequestBatch([req]))
+
+    assert captured["diffuse"]["do_true_cfg"] is True
+    assert captured["diffuse"]["true_cfg_scale"] == 3.0
+
+
+def test_native_anima_cfg_equation() -> None:
+    """The denoising loop applies the standard true-CFG equation."""
+    pipeline = AnimaPipeline.__new__(AnimaPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline._interrupt = False
+
+    class MockProgressBar:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            pass
+
+        def update(self):
+            pass
+
+    pipeline.progress_bar = lambda **_kwargs: MockProgressBar()
+
+    class MockTransformer:
+        dtype = torch.float32
+
+        def __call__(self, hidden_states, timestep, encoder_hidden_states, padding_mask, return_dict=False):
+            # Return double the encoder_hidden_states to simulate transformer output
+            return (encoder_hidden_states * 2.0,)
+
+    pipeline.transformer = MockTransformer()
+
+    class MockScheduler:
+        def set_begin_index(self, index):
+            pass
+
+        def step(self, noise_pred, t, latents, return_dict=False):
+            # Return the noise_pred itself to inspect it
+            return (noise_pred,)
+
+    pipeline.scheduler = MockScheduler()
+
+    prompt_embeds = torch.tensor([[[2.0, 3.0]]])
+    negative_prompt_embeds = torch.tensor([[[1.0, 2.0]]])
+    latents = torch.zeros(1, 16, 1, 16, 16)
+    padding_mask = torch.zeros(1, 1, 16, 16)
+    timesteps = torch.tensor([500.0])
+
+    # Cond = prompt_embeds * 2.0 = [4.0, 6.0]
+    # Uncond = negative_prompt_embeds * 2.0 = [2.0, 4.0]
+    # Expected noise_pred = uncond + true_cfg_scale * (cond - uncond)
+    #                     = [2.0, 4.0] + 4.0 * ([4.0, 6.0] - [2.0, 4.0])
+    #                     = [2.0, 4.0] + 4.0 * [2.0, 2.0]
+    #                     = [2.0, 4.0] + [8.0, 8.0] = [10.0, 12.0]
+    out_latents = pipeline.diffuse(
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+        latents=latents,
+        padding_mask=padding_mask,
+        timesteps=timesteps,
+        do_true_cfg=True,
+        true_cfg_scale=4.0,
+    )
+
+    expected = torch.tensor([[[10.0, 12.0]]])
+    assert torch.allclose(out_latents, expected)

tests/diffusion/models/auk/test_auk_transformer_parity.py

@@ -0,0 +1,439 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK transformer parity: the correctness oracle for the in-tree port.
+
+Checks that :class:`AuKTransformer` reproduces the upstream ``auk`` reference backbone
+on the released ``auk_base`` weights. Both models are built from the same
+state dict and fed the same inputs, under the reference's own inference
+configuration: fp32 weights with ``torch.autocast(bfloat16)``. That
+configuration matters. The reference keeps its rotary frequencies in a buffer,
+so casting its weights to bf16 would round both the frequencies and the
+integer positions and move the numbers for reasons unrelated to the port.
+
+Cases:
+
+* ``single`` - one conditional forward
+* ``cfg`` - the doubled ``cfg_infer`` forward, both branches
+* ``masked`` - batch of 2 with padded targets, references and text
+* ``rope`` - the local rotary implementation against ``x_transformers``, and
+  its survival of a cast to bf16
+* ``sample`` - an 8-step Euler trajectory against the reference backbone
+  driven by a transcription of upstream's ODE loop
+* ``generator`` - the sampler's two noise paths agreeing
+* ``cache`` - the text projection cache and ``clear_cache()``
+* ``mask flag`` - ``attn_mask_enabled=False`` reproducing the reference with
+  its own flag off
+
+This directory deliberately has no ``__init__.py``, unlike its siblings.
+pytest's prepend import mode puts the nearest non-package directory on
+``sys.path``, and an ``auk`` package here would shadow the upstream ``auk``
+distribution this test compares against.
+
+Needs the released checkpoint and a GPU, so it runs locally rather than in
+CI. Point ``AUK_CKPT_DIR`` at a directory holding ``auk_base.safetensors`` and
+``config.yaml``; the test skips when that is unset or when the upstream
+``auk`` package is not importable::
+
+    AUK_CKPT_DIR=/path/to/ckpts/AuK CUDA_VISIBLE_DEVICES=6 \
+        python -m pytest tests/diffusion/models/auk/test_auk_transformer_parity.py -v -s
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import os
+from pathlib import Path
+
+import pytest
+import torch
+import yaml
+
+from vllm_omni.diffusion.models.auk.auk_transformer import (
+    AuKTransformer,
+    Rotary,
+    _apply_rope,
+    dit_state_dict,
+    sample_latents,
+)
+
+CKPT_DIR = os.environ.get("AUK_CKPT_DIR")
+
+# Tolerances. The port evaluates rotary embeddings in fp32 and skips the
+# reference's redundant uncond text projection, so the two graphs are not
+# bit-identical under autocast; they agree well inside bf16 noise.
+FORWARD_ATOL = 3e-2
+FORWARD_COS = 0.999
+SAMPLE_ATOL = 5e-2
+SAMPLE_COS = 0.999
+ROPE_ATOL = 1e-5
+
+TARGET_FRAMES = 300
+REF_FRAMES = 150
+TEXT_TOKENS = 40
+SHORT_LENS = (220, 90, 25)  # target, reference, text for the padded batch row
+TIME = 0.37
+SAMPLE_STEPS = 8
+SWAY_COEF = -1.0
+SEED = 1234
+DEVICE = "cuda:0"
+AUTOCAST_DTYPE = torch.bfloat16
+
+
+def _skip_reason() -> str | None:
+    if CKPT_DIR is None:
+        return "AUK_CKPT_DIR is not set"
+    if not Path(CKPT_DIR, "auk_base.safetensors").is_file():
+        return f"no auk_base.safetensors under {CKPT_DIR}"
+    try:
+        # Resolving the submodule, not just "auk", also catches the shadow an
+        # __init__.py in this directory would create.
+        upstream = importlib.util.find_spec("auk.model")
+    except ModuleNotFoundError:
+        upstream = None
+    if upstream is None:
+        return "the upstream 'auk' package is not importable or is shadowed by this directory"
+    if not torch.cuda.is_available():
+        return "no CUDA device"
+    return None
+
+
+SKIP_REASON = _skip_reason()
+
+pytestmark = [
+    pytest.mark.local_model,
+    pytest.mark.diffusion,
+    pytest.mark.gpu,
+    pytest.mark.cuda,
+    pytest.mark.skipif(SKIP_REASON is not None, reason=str(SKIP_REASON)),
+]
+
+
+def _arch() -> dict:
+    """Read the backbone geometry out of the checkpoint's config."""
+    config = yaml.safe_load(Path(CKPT_DIR, "config.yaml").read_text())["model"]
+    return {
+        "dim": config["arch"]["dim"],
+        "heads": config["arch"]["heads"],
+        "dim_head": config["arch"].get("dim_head", 64),
+        "ff_mult": config["arch"]["ff_mult"],
+        "text_hidden_dim": config["arch"]["text_hidden_dim"],
+        "num_layers": config["arch"]["num_layers"],
+        "num_single_layers": config["arch"]["num_single_layers"],
+        "latent_dim": config["vae"]["latent_dim"],
+    }
+
+
+def _stats(port: torch.Tensor, ref: torch.Tensor) -> tuple[float, float, float]:
+    """Return max abs diff, cosine similarity and the reference's RMS."""
+    a = port.flatten().float()
+    b = ref.flatten().float()
+    return (
+        (a - b).abs().max().item(),
+        torch.nn.functional.cosine_similarity(a, b, dim=0).item(),
+        b.pow(2).mean().sqrt().item(),
+    )
+
+
+def _assert_close(name: str, port: torch.Tensor, ref: torch.Tensor, atol: float, min_cos: float) -> None:
+    max_abs, cos, rms = _stats(port, ref)
+    print(f"[{name}] max|Δ|={max_abs:.4e} cos={cos:.8f} ref_rms={rms:.4f} rel={max_abs / rms:.3e}")
+    assert max_abs < atol, f"{name}: max|Δ|={max_abs:.4e} >= {atol:.1e} (reference rms {rms:.4f})"
+    assert cos > min_cos, f"{name}: cos={cos:.8f} <= {min_cos}"
+
+
+def _autocast():
+    return torch.autocast("cuda", dtype=AUTOCAST_DTYPE)
+
+
+@pytest.fixture(scope="module")
+def models() -> tuple[torch.nn.Module, AuKTransformer]:
+    """The reference and the port, same weights, fp32 on the accelerator."""
+    from auk.model.flux2_edit import Flux2Edit
+    from safetensors.torch import load_file
+
+    arch = _arch()
+    device = torch.device(DEVICE)
+    state = dit_state_dict(load_file(str(Path(CKPT_DIR, "auk_base.safetensors")), device="cpu"))
+
+    reference = Flux2Edit(dropout=0.0, attn_backend="torch", attn_mask_enabled=True, **arch)
+    reference.load_state_dict(state, strict=True)
+    reference = reference.to(device).eval().requires_grad_(False)
+
+    port = AuKTransformer(**arch)
+    # strict=True is the point of the port: every released tensor lands.
+    port.load_state_dict(state, strict=True)
+    port = port.to(device).eval().requires_grad_(False)
+
+    del state
+    return reference, port
+
+
+@pytest.fixture(scope="module")
+def inputs() -> dict[str, torch.Tensor]:
+    """Fixed random inputs, as a batch of 2 whose second row is padded."""
+    device = torch.device(DEVICE)
+    generator = torch.Generator(device="cpu").manual_seed(SEED)
+    arch = _arch()
+
+    def randn(*shape: int) -> torch.Tensor:
+        return torch.randn(*shape, generator=generator, dtype=torch.float32).to(device)
+
+    def pad_mask(width: int, short: int) -> torch.Tensor:
+        mask = torch.zeros(2, width, dtype=torch.bool, device=device)
+        mask[0, :] = True
+        mask[1, :short] = True
+        return mask
+
+    mask2 = pad_mask(TARGET_FRAMES, SHORT_LENS[0])
+    ref_mask2 = pad_mask(REF_FRAMES, SHORT_LENS[1])
+    c_mask2 = pad_mask(TEXT_TOKENS, SHORT_LENS[2])
+
+    # Padding rows are all-zero in the text encoder's output, so keep the
+    # inputs consistent with the masks.
+    x2 = randn(2, TARGET_FRAMES, arch["latent_dim"]) * mask2.unsqueeze(-1)
+    ref2 = randn(2, REF_FRAMES, arch["latent_dim"]) * ref_mask2.unsqueeze(-1)
+    text2 = randn(2, TEXT_TOKENS, arch["text_hidden_dim"]) * c_mask2.unsqueeze(-1)
+
+    return {
+        "x": x2[:1].contiguous(),
+        "ref": ref2[:1].contiguous(),
+        "text": text2[:1].contiguous(),
+        "c_mask": c_mask2[:1].contiguous(),
+        "x2": x2,
+        "ref2": ref2,
+        "text2": text2,
+        "mask2": mask2,
+        "ref_mask2": ref_mask2,
+        "c_mask2": c_mask2,
+        "time": torch.tensor(TIME, dtype=torch.float32, device=device),
+    }
+
+
+def test_single_forward(models, inputs) -> None:
+    """One conditional forward: reference against port."""
+    reference, port = models
+    kwargs = dict(text=inputs["text"], time=inputs["time"], c_mask=inputs["c_mask"], ref=inputs["ref"])
+    with torch.no_grad(), _autocast():
+        want = reference(x=inputs["x"], **kwargs)
+        got = port(inputs["x"], **kwargs)
+
+    assert got.shape == want.shape == (1, TARGET_FRAMES, port.latent_dim)
+    _assert_close("single", got, want, FORWARD_ATOL, FORWARD_COS)
+
+
+def test_cfg_forward(models, inputs) -> None:
+    """The doubled cfg_infer forward, cond branch first."""
+    reference, port = models
+    kwargs = dict(
+        text=inputs["text"],
+        time=inputs["time"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        cfg_infer=True,
+    )
+    with torch.no_grad(), _autocast():
+        want = reference(x=inputs["x"], **kwargs)
+        got = port(inputs["x"], **kwargs)
+
+    assert got.shape == want.shape == (2, TARGET_FRAMES, port.latent_dim)
+    _assert_close("cfg-cond", got[:1], want[:1], FORWARD_ATOL, FORWARD_COS)
+    _assert_close("cfg-uncond", got[1:], want[1:], FORWARD_ATOL, FORWARD_COS)
+
+
+def test_masked_batch_forward(models, inputs) -> None:
+    """Batch of 2 with padded targets, references and text.
+
+    Only valid target frames are compared. Padded frames are unconstrained:
+    the reference leaves them to whatever masked attention emits there.
+    """
+    reference, port = models
+    kwargs = dict(
+        text=inputs["text2"],
+        time=inputs["time"],
+        mask=inputs["mask2"],
+        c_mask=inputs["c_mask2"],
+        ref=inputs["ref2"],
+        ref_mask=inputs["ref_mask2"],
+    )
+    with torch.no_grad(), _autocast():
+        want = reference(x=inputs["x2"], **kwargs)
+        got = port(inputs["x2"], **kwargs)
+
+    assert got.shape == want.shape == (2, TARGET_FRAMES, port.latent_dim)
+    valid = inputs["mask2"].unsqueeze(-1).expand_as(got)
+    _assert_close("masked", got[valid], want[valid], FORWARD_ATOL, FORWARD_COS)
+
+
+def _set_mask_flag(reference, port, enabled: bool) -> None:
+    """Flip attn_mask_enabled on both models in place, so no second copy is loaded."""
+    for module in port.modules():
+        if hasattr(module, "attn_mask_enabled"):
+            module.attn_mask_enabled = enabled
+    for blocks in (reference.transformer_blocks, reference.single_transformer_blocks):
+        for block in blocks:
+            block.attn.processor.attn_mask_enabled = enabled
+
+
+def test_mask_flag_disabled(models, inputs) -> None:
+    """``attn_mask_enabled=False`` must track the reference with its own flag off.
+
+    With the flag off both models run attention unmasked and still zero the
+    padded outputs, so the padded batch is compared in full.
+    """
+    reference, port = models
+    kwargs = dict(
+        text=inputs["text2"],
+        time=inputs["time"],
+        mask=inputs["mask2"],
+        c_mask=inputs["c_mask2"],
+        ref=inputs["ref2"],
+        ref_mask=inputs["ref_mask2"],
+    )
+    with torch.no_grad(), _autocast():
+        masked = port(inputs["x2"], **kwargs)
+        _set_mask_flag(reference, port, False)
+        try:
+            want = reference(x=inputs["x2"], **kwargs)
+            got = port(inputs["x2"], **kwargs)
+        finally:
+            _set_mask_flag(reference, port, True)
+
+    _assert_close("mask-flag-off", got, want, FORWARD_ATOL, FORWARD_COS)
+    # Guard against a vacuous pass: the flag has to actually change something.
+    assert (got.float() - masked.float()).abs().max().item() > 1e-3
+
+
+def test_sampler_generator_matches_seed(models, inputs) -> None:
+    """A generator seeded with 0 must draw exactly what ``seed=0`` draws."""
+    _, port = models
+    common = dict(
+        gen_frames=TARGET_FRAMES,
+        nfe=2,
+        cfg_strength=0.0,
+        sway_sampling_coef=SWAY_COEF,
+        text=inputs["text"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        ref_mask=None,
+    )
+    with torch.no_grad(), _autocast():
+        by_seed = sample_latents(port, seed=0, **common)
+        by_generator = sample_latents(port, generator=torch.Generator(device=DEVICE).manual_seed(0), **common)
+
+    print(f"[generator] max|Δ|={(by_generator.float() - by_seed.float()).abs().max().item():.4e}")
+    torch.testing.assert_close(by_generator, by_seed, rtol=0, atol=0)
+
+
+def test_cache_needs_clearing(models, inputs) -> None:
+    """The cached text projection is reused until ``clear_cache()``.
+
+    ``sample_latents`` clears on the way out, so a pipeline that reuses the
+    module across requests is safe without doing anything. A caller driving
+    ``forward(cache=True)`` itself has to clear when the text changes.
+    """
+    _, port = models
+    other_text = inputs["text"] * 0.5
+    kwargs = dict(time=inputs["time"], c_mask=inputs["c_mask"], ref=inputs["ref"], cfg_infer=True)
+
+    port.clear_cache()
+    with torch.no_grad(), _autocast():
+        want = port(inputs["x"], inputs["text"], **kwargs)
+        port(inputs["x"], other_text, cache=True, **kwargs)
+        stale = port(inputs["x"], inputs["text"], cache=True, **kwargs)
+        port.clear_cache()
+        fresh = port(inputs["x"], inputs["text"], cache=True, **kwargs)
+        port.clear_cache()
+
+    # The cache is real: the second call reused the first call's projection.
+    assert (stale.float() - want.float()).abs().max().item() > 1e-3
+    torch.testing.assert_close(fresh, want, rtol=0, atol=0)
+    assert port.text_cond is None and port.text_uncond is None
+
+
+def test_rope_matches_x_transformers() -> None:
+    """The local rotary implementation against x_transformers, in fp32."""
+    from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb
+
+    device = torch.device(DEVICE)
+    dim_head, seq_len = 64, TEXT_TOKENS + REF_FRAMES + TARGET_FRAMES
+    q = torch.randn(2, 24, seq_len, dim_head, dtype=torch.float32, device=device)
+
+    freqs = RotaryEmbedding(dim_head).to(device).forward_from_seq_len(seq_len)[0]
+    want = apply_rotary_pos_emb(q, freqs)
+    got = _apply_rope(q, Rotary(dim_head).to(device)(seq_len))
+
+    max_abs, cos, _ = _stats(got, want)
+    print(f"[rope] max|Δ|={max_abs:.4e} cos={cos:.10f}")
+    assert max_abs < ROPE_ATOL, f"rope: max|Δ|={max_abs:.4e} >= {ROPE_ATOL:.1e}"
+
+
+def test_rope_survives_half_precision_cast() -> None:
+    """Casting the module to bf16 must not round the rotary frequencies.
+
+    The reference stores them in a buffer, so ``.to(bfloat16)`` rounds both the
+    frequencies and the integer positions and silently changes every rotation.
+    The port recomputes them instead.
+    """
+    dim_head, seq_len = 64, TEXT_TOKENS + REF_FRAMES + TARGET_FRAMES
+    exact = Rotary(dim_head)(seq_len)
+    got = Rotary(dim_head).to(torch.bfloat16)(seq_len)
+
+    assert got.dtype == torch.float32
+    torch.testing.assert_close(got, exact, rtol=0, atol=0)
+
+
+def _reference_euler(reference, *, x, text, c_mask, ref, steps, cfg_strength, sway) -> torch.Tensor:
+    """Upstream's ODE loop, transcribed.
+
+    ``CFMEdit.sample`` hands the backbone to ``torchdiffeq.odeint(method=
+    "euler")``, which on a fixed grid is ``y += (t1 - t0) * f(t0, y)``.
+    Constructing a real ``CFMEdit`` would drag in the Qwen-Omni text encoder,
+    so only the backbone here comes from upstream.
+    """
+    t = torch.linspace(0, 1, steps + 1, device=x.device, dtype=torch.float32)
+    if sway is not None:
+        t = t + sway * (torch.cos(math.pi / 2 * t) - 1 + t)
+
+    shared = dict(text=text, c_mask=c_mask, ref=ref, mask=None, ref_mask=None, cache=True)
+    for i in range(steps):
+        if cfg_strength < 1e-5:
+            v = reference(x=x, time=t[i], **shared)
+        else:
+            pred = reference(x=x, time=t[i], cfg_infer=True, **shared)
+            v_cond, v_uncond = torch.chunk(pred, 2, dim=0)
+            v = v_cond + (v_cond - v_uncond) * cfg_strength
+        x = x + (t[i + 1] - t[i]) * v
+    reference.clear_cache()
+    return x
+
+
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_euler_sample(models, inputs, cfg_strength: float) -> None:
+    """An 8-step Euler trajectory, with and without guidance."""
+    reference, port = models
+    shared = dict(text=inputs["text"], c_mask=inputs["c_mask"], ref=inputs["ref"])
+
+    with torch.no_grad(), _autocast():
+        got = sample_latents(
+            port,
+            gen_frames=TARGET_FRAMES,
+            ref_mask=None,
+            nfe=SAMPLE_STEPS,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=SWAY_COEF,
+            seed=SEED,
+            **shared,
+        )
+        # Re-seeding reproduces the noise sample_latents drew for itself.
+        torch.manual_seed(SEED)
+        noise = torch.randn(
+            TARGET_FRAMES, port.latent_dim, device=inputs["ref"].device, dtype=inputs["ref"].dtype
+        ).unsqueeze(0)
+        want = _reference_euler(
+            reference, x=noise, steps=SAMPLE_STEPS, cfg_strength=cfg_strength, sway=SWAY_COEF, **shared
+        )
+
+    assert got.shape == want.shape == (1, TARGET_FRAMES, port.latent_dim)
+    _assert_close(f"sample-cfg{cfg_strength:g}", got, want, SAMPLE_ATOL, SAMPLE_COS)

tests/diffusion/models/auk/test_auk_vae_parity.py

@@ -0,0 +1,254 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK VAE parity: the numeric oracle for the vendored BigVGAN-flow codec.
+
+Checks ``AuKVAE`` against the AuK reference implementation on real speech, with
+both models loaded from the same released ``vae.safetensors``:
+
+  * encoder statistics (mean / log-std) agree,
+  * ``encode(sample=False)`` reproduces the reference mean path,
+  * ``encode(sample=True)`` reproduces ``encoding_and_normalization`` for a
+    matching RNG seed,
+  * ``decode`` reproduces ``denormalize`` + ``inference_from_latents``, in fp32
+    and (reported, not gated tightly) under bf16 autocast,
+  * the waveform survives a mean-path round trip.
+
+Runs on GPU in fp32 and needs both the released checkpoint and the reference
+``auk`` package importable, so it is a local_model test rather than CI. Point
+``AUK_CKPT_DIR`` at the directory holding ``vae.safetensors`` and ``AUK_REF_WAV``
+at a speech clip (default ``$AUK_CKPT_DIR/ref.wav``)::
+
+    AUK_CKPT_DIR=... AUK_REF_WAV=... python -m pytest \
+        tests/diffusion/models/auk/test_auk_vae_parity.py -v -s
+
+Unlike its sibling directories this one carries no ``__init__.py`` on purpose: it would put
+``tests/diffusion/models`` on ``sys.path`` ahead of site-packages, and this directory's own name
+would then shadow the reference ``auk`` package the parity check imports.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+from pathlib import Path
+
+import pytest
+import soundfile as sf
+import torch
+import torchaudio.functional as AF
+
+from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE
+
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, pytest.mark.gpu, pytest.mark.cuda]
+
+# The released model.vae.model_init_kwargs, shared by both models under test.
+VAE_CONFIG = {
+    "upsample_rates": [5, 4, 3, 2, 2, 2],
+    "upsample_kernel_sizes": [10, 8, 6, 4, 4, 4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3, 7, 11],
+    "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+    "downsample_rates": [2, 2, 2, 3, 4, 5],
+    "downsample_channels": [12, 24, 48, 96, 192, 384, 768],
+    "snake_logscale": True,
+    "latent_dim": 64,
+    "use_vae": True,
+    "causal": True,
+    "flow_hidden_channels": 256,
+    "act_causal": True,
+}
+
+SAMPLE_RATE = 24000
+HOP_SIZE = 480
+LATENT_DIM = 64
+SEED = 1234
+# Keep the clip short: the parity run shares its GPU with other jobs.
+MAX_SECONDS = 4.0
+TOLERANCE = 1e-5
+DEVICE = "cuda"
+
+_CKPT_DIR = os.environ.get("AUK_CKPT_DIR")
+_VAE_PATH = Path(_CKPT_DIR) / "vae.safetensors" if _CKPT_DIR else None
+_REF_WAV = Path(os.environ.get("AUK_REF_WAV") or (Path(_CKPT_DIR) / "ref.wav" if _CKPT_DIR else "unset"))
+
+
+def _reference_importable() -> bool:
+    """Resolve ``auk.model``, not ``auk``.
+
+    An ``__init__.py`` in this directory would make its own name resolve as ``auk`` and satisfy a
+    top-level probe, leaving the fixtures to fail at setup instead of skipping. The submodule
+    probe fails in that case, which is the honest answer.
+    """
+    try:
+        return importlib.util.find_spec("auk.model") is not None
+    except (ImportError, ValueError):
+        return False
+
+
+def _skip_reason() -> str | None:
+    if not _CKPT_DIR:
+        return "set AUK_CKPT_DIR to the directory holding vae.safetensors"
+    if _VAE_PATH is None or not _VAE_PATH.is_file():
+        return f"no vae.safetensors under {_CKPT_DIR}"
+    if not _REF_WAV.is_file():
+        return f"no reference clip at {_REF_WAV} (set AUK_REF_WAV)"
+    if not _reference_importable():
+        return "cannot import the reference auk.model package (an __init__.py here would shadow it)"
+    if not torch.cuda.is_available():
+        return "parity runs in fp32 on GPU"
+    return None
+
+
+pytestmark.append(pytest.mark.skipif(_skip_reason() is not None, reason=_skip_reason() or ""))
+
+
+def _max_abs(lhs: torch.Tensor, rhs: torch.Tensor) -> float:
+    return (lhs.float() - rhs.float()).abs().max().item()
+
+
+def _snr_db(signal: torch.Tensor, estimate: torch.Tensor) -> float:
+    noise = (signal.float() - estimate.float()).pow(2).sum()
+    return (10.0 * torch.log10(signal.float().pow(2).sum() / noise)).item()
+
+
+@pytest.fixture(scope="module")
+def clip() -> torch.Tensor:
+    """Mono 24 kHz speech, trimmed to a whole number of latent frames."""
+    audio, sample_rate = sf.read(str(_REF_WAV), dtype="float32", always_2d=True)  # (T, C)
+    wav = torch.from_numpy(audio).mean(dim=1, keepdim=True).transpose(0, 1)
+    if sample_rate != SAMPLE_RATE:
+        wav = AF.resample(wav, sample_rate, SAMPLE_RATE)
+    samples = min(int(MAX_SECONDS * SAMPLE_RATE), wav.size(-1)) // HOP_SIZE * HOP_SIZE
+    wav = wav[:, :samples].contiguous().to(DEVICE)
+    print(f"\n[clip] {_REF_WAV.name}: {samples} samples, {samples / SAMPLE_RATE:.2f} s, {samples // HOP_SIZE} frames")
+    return wav
+
+
+@pytest.fixture(scope="module")
+def reference(clip: torch.Tensor) -> dict[str, torch.Tensor]:
+    """Reference tensors from the upstream model, which is then released so one model fits at a time."""
+    from auk.model.vae import load_vae_model
+    from auk.model.vae.bigvgan_flow_vae import BigVGANFlowVAEConfig
+
+    model = load_vae_model(
+        vae_name="BigVGANFlowVAE",
+        vae_cfg=BigVGANFlowVAEConfig.from_dict(VAE_CONFIG),
+        vae_ckpt=str(_VAE_PATH),
+        map_location="cpu",
+    )
+    model = model.to(DEVICE).eval()
+    model.requires_grad_(False)
+
+    with torch.no_grad():
+        mean, log_std = model.audio_encoder(clip.unsqueeze(1)).chunk(2, 1)
+        scale = torch.sqrt(model.global_log_std.float())
+        mean_latents = (mean.transpose(1, 2).float() - model.global_mean.float()) / scale
+        torch.manual_seed(SEED)
+        sampled_latents, _ = model.encoding_and_normalization(clip.unsqueeze(1))
+        decoded = model.inference_from_latents(model.denormalize(mean_latents).permute(0, 2, 1))
+
+    tensors = {
+        "mean": mean.clone(),
+        "log_std": log_std.clone(),
+        "mean_latents": mean_latents.clone(),
+        "sampled_latents": sampled_latents.clone(),
+        "decoded": decoded.squeeze(1).clone(),
+    }
+    del model
+    torch.accelerator.empty_cache()
+    return tensors
+
+
+@pytest.fixture(scope="module")
+def port(reference: dict[str, torch.Tensor]) -> tuple[AuKVAE, list[str], list[str]]:
+    """The vendored module, loaded from the same checkpoint. Depends on ``reference`` for load ordering.
+
+    Placed on the accelerator before loading so weight norm is folded there; see ``load_weights``.
+    """
+    model = AuKVAE.from_config(VAE_CONFIG).to(DEVICE)
+    missing, unexpected = model.load_weights(_VAE_PATH)
+    model = model.eval()
+    model.requires_grad_(False)
+    return model, missing, unexpected
+
+
+def test_checkpoint_covers_every_inference_tensor(port: tuple[AuKVAE, list[str], list[str]]) -> None:
+    model, missing, unexpected = port
+    print(f"[keys] missing={len(missing)} unexpected={len(unexpected)}")
+    print(f"[keys] unexpected prefixes={sorted({key.split('.')[0] for key in unexpected})}")
+    assert missing == []
+    assert unexpected and all(key.startswith("flow.") for key in unexpected)
+    assert (model.hop_size, model.sample_rate, model.latent_dim) == (HOP_SIZE, SAMPLE_RATE, LATENT_DIM)
+
+
+def test_encoder_statistics_match(
+    clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]
+) -> None:
+    model = port[0]
+    with torch.no_grad():
+        mean, log_std = model.audio_encoder(clip.unsqueeze(1)).chunk(2, 1)
+    mean_diff = _max_abs(mean, reference["mean"])
+    log_std_diff = _max_abs(log_std, reference["log_std"])
+    print(f"[encoder] max|dmean|={mean_diff:.3e} max|dlog_std|={log_std_diff:.3e}")
+    assert mean_diff < TOLERANCE
+    assert log_std_diff < TOLERANCE
+    # Folded and hook-computed weights are bit-identical on one device, so the whole encoder is.
+    assert (mean_diff, log_std_diff) == (0.0, 0.0)
+
+
+def test_encode_mean_path_matches(
+    clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]
+) -> None:
+    model = port[0]
+    with torch.no_grad():
+        latents = model.encode(clip)
+    assert latents.shape == (1, clip.size(-1) // HOP_SIZE, LATENT_DIM)
+    diff = _max_abs(latents, reference["mean_latents"])
+    print(f"[encode mean] shape={tuple(latents.shape)} max|d|={diff:.3e}")
+    assert diff == 0.0
+
+
+def test_encode_sampled_matches_seeded_reference(
+    clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]
+) -> None:
+    model = port[0]
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(SEED)
+    with torch.no_grad():
+        latents = model.encode(clip, sample=True, generator=generator)
+    diff = _max_abs(latents, reference["sampled_latents"])
+    spread = _max_abs(latents, reference["mean_latents"])
+    print(f"[encode sample] max|d vs seeded reference|={diff:.3e} max|d vs mean path|={spread:.3e}")
+    assert diff < TOLERANCE
+    assert diff == 0.0  # a fresh device Generator replays the default one after torch.manual_seed
+    assert spread > TOLERANCE  # the draw actually moved the latents
+
+
+def test_decode_matches(reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]) -> None:
+    model = port[0]
+    latents = reference["mean_latents"]
+    with torch.no_grad():
+        wav = model.decode(latents)
+    assert wav.shape == (1, latents.size(1) * HOP_SIZE)
+    diff = _max_abs(wav, reference["decoded"])
+    print(f"[decode fp32] shape={tuple(wav.shape)} max|d|={diff:.3e}")
+    assert diff < TOLERANCE
+    assert diff == 0.0
+
+    with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        wav_bf16 = model.decode(latents)
+    bf16_diff = _max_abs(wav_bf16, reference["decoded"])
+    print(f"[decode bf16 autocast] max|d| vs fp32 reference={bf16_diff:.3e} snr={_snr_db(wav, wav_bf16):.2f} dB")
+    assert torch.isfinite(wav_bf16).all()
+    assert bf16_diff < 0.5  # sanity only: bf16 is reported, not held to fp32 parity
+
+
+def test_round_trip_snr(clip: torch.Tensor, port: tuple[AuKVAE, list[str], list[str]]) -> None:
+    model = port[0]
+    with torch.no_grad():
+        wav = model.decode(model.encode(clip))
+    assert wav.shape == clip.shape
+    snr = _snr_db(clip, wav)
+    print(f"[round trip] snr={snr:.2f} dB peak={wav.abs().max().item():.4f}")
+    assert snr > 10.0

tests/diffusion/models/auk/test_pipeline_auk.py

@@ -0,0 +1,514 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK stage-1 pipeline: request parsing on CPU, generation parity on a GPU.
+
+The CPU tests stub the transformer, the codec and the ODE so they exercise only
+what the pipeline itself owns: how a request turns into a target length, a
+sampling schedule and a noise generator.
+
+The parity test replays the saved upstream reference for the zero-shot TTS case
+and reports latent, per-frame and waveform distances rather than asserting bit
+equality, because two differences live in the conditioning rather than in the
+port: the saved text condition is an fp32 thinker forward while the reference
+generated under bf16 autocast, and the reference's stochastic VAE draw came from
+the process-global RNG seeded at 1234.
+
+It replays with ``vae_sample=False``. That is the deliberate choice: the
+reference re-seeds the global RNG immediately before drawing the ODE noise, so a
+per-request generator reproduces that noise exactly (verified bit-identical) as
+long as nothing else has drawn from it. Sampling the VAE posterior consumes the
+generator first, which moves the noise and produces a different, equally valid
+realization; comparing latents then measures the noise, not the port. The
+posterior mean costs a reference latent that is 0.045 MSE (0.969 cosine) away
+from the draw the reference used, which is the smaller of the two errors.
+Generate the reference artifacts once with the upstream package installed
+(``tools/auk_parity_reference.py --auk-repo ... --ckpt-dir ... --qwen-dir ...
+--out /path/to/auk-parity``; it writes ``parity_ref/<variant>`` and
+``fusion_ref``), then run::
+
+    AUK_OMNI_CKPT_DIR=/path/to/auk-omni-base \
+    AUK_PARITY_REF=/path/to/auk-parity/parity_ref/base \
+    python -m pytest -s tests/diffusion/models/auk/test_pipeline_auk.py
+
+``AUK_OMNI_CKPT_DIR`` is the assembled directory this pipeline loads, which is
+not the same thing as the released AuK snapshot that the transformer and codec
+parity tests in this directory read from ``AUK_CKPT_DIR``. ``AUK_CKPT_DIR`` is
+accepted as a fallback, and the test skips rather than fails when the directory
+it resolves to is not an assembled one.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+import torchaudio
+from torch import nn
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.models.auk import pipeline_auk
+from vllm_omni.diffusion.models.auk.pipeline_auk import AuKPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+LATENT_DIM = 64
+HOP = 480
+SAMPLE_RATE = 24000
+TEXT_HIDDEN_DIM = 2048
+
+# The distilled student's grid, exact values from the released checkpoint.
+FLASH_T_GRID = [0.0, 0.07612049579620361, 0.2928932309150696, 0.6173166036605835, 1.0]
+
+CKPT_DIR = os.environ.get("AUK_OMNI_CKPT_DIR") or os.environ.get("AUK_CKPT_DIR")
+PARITY_REF = os.environ.get("AUK_PARITY_REF")
+# An assembled directory is the only thing this pipeline can load.
+IS_ASSEMBLED = bool(CKPT_DIR) and (Path(CKPT_DIR).expanduser() / "config.json").is_file()
+
+# Measured on the released base checkpoint: mean frame cosine 0.961 to 0.968,
+# relative latent MSE 0.051 to 0.063, log-mel distance 0.34. That spread is
+# bf16 against fp32 DiT weights plus where the codec folds its weight norm,
+# neither of which is worth gating on. A port that has lost the conditioning
+# sits at 0.110 cosine, 1.82 relative MSE and 2.88 log-mel, so these gates
+# still separate the two by an order of magnitude.
+PARITY_MIN_FRAME_COSINE = 0.93
+PARITY_MAX_RELATIVE_MSE = 0.15
+PARITY_MAX_MEL_DISTANCE = 0.6
+
+
+class _StubTransformer(nn.Module):
+    """Records its construction arguments; owns one weight so strict loading runs."""
+
+    def __init__(
+        self,
+        *,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        ff_mult: float,
+        latent_dim: int,
+        text_hidden_dim: int,
+        num_layers: int,
+        num_single_layers: int,
+        attn_mask_enabled: bool = True,
+    ) -> None:
+        super().__init__()
+        self.init_kwargs = {
+            "dim": dim,
+            "heads": heads,
+            "dim_head": dim_head,
+            "ff_mult": ff_mult,
+            "latent_dim": latent_dim,
+            "text_hidden_dim": text_hidden_dim,
+            "num_layers": num_layers,
+            "num_single_layers": num_single_layers,
+            "attn_mask_enabled": attn_mask_enabled,
+        }
+        self.latent_dim = latent_dim
+        self.proj = nn.Linear(latent_dim, latent_dim, bias=False)
+
+
+class _StubVAE(nn.Module):
+    """Deterministic stand-in for the codec, one latent frame per hop."""
+
+    hop_size = HOP
+    sample_rate = SAMPLE_RATE
+    latent_dim = LATENT_DIM
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.encode_calls: list[dict[str, Any]] = []
+        self.decode_calls = 0
+        self.weights_path: str | None = None
+        self.config: dict[str, Any] = {}
+
+    @classmethod
+    def from_config(cls, cfg: dict[str, Any]) -> _StubVAE:
+        vae = cls()
+        vae.config = dict(cfg)
+        return vae
+
+    def load_weights(self, path: str, **_: Any) -> tuple[list[str], list[str]]:
+        self.weights_path = str(path)
+        return [], []
+
+    def encode(
+        self,
+        wav: torch.Tensor,
+        *,
+        sample: bool = False,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        frames = wav.shape[-1] // HOP
+        self.encode_calls.append({"samples": int(wav.shape[-1]), "sample": sample, "generator": generator})
+        if sample:
+            return torch.randn(1, frames, LATENT_DIM, generator=generator, device=wav.device)
+        return torch.zeros(1, frames, LATENT_DIM, device=wav.device)
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        self.decode_calls += 1
+        return torch.zeros(1, latents.shape[1] * HOP, device=latents.device)
+
+
+def _stub_sampler(calls: list[dict[str, Any]]):
+    """Record the ODE arguments and return noise drawn from the request generator."""
+
+    def sample_latents(dit: nn.Module, **kwargs: Any) -> torch.Tensor:
+        calls.append(kwargs)
+        shape = (1, kwargs["gen_frames"], kwargs["latent_dim"])
+        generator = kwargs.get("generator")
+        if generator is None:
+            return torch.zeros(*shape, device=kwargs["device"], dtype=kwargs["dtype"])
+        return torch.randn(*shape, generator=generator, device=kwargs["device"], dtype=kwargs["dtype"])
+
+    return sample_latents
+
+
+def _write_checkpoint(root: Path, variant: str) -> Path:
+    """Write a tiny assembled checkpoint the stubs can load."""
+
+    from safetensors.torch import save_file
+
+    model_dir = root / f"auk-{variant}"
+    model_dir.mkdir(parents=True, exist_ok=True)
+    config = {
+        "model_type": "auk",
+        "architectures": ["AuKForConditionalGeneration"],
+        "dit": {
+            "dim": 16,
+            "heads": 2,
+            "dim_head": 8,
+            "ff_mult": 2,
+            "text_hidden_dim": TEXT_HIDDEN_DIM,
+            "num_layers": 1,
+            "num_single_layers": 1,
+            "attn_mask_enabled": True,
+        },
+        "vae": {
+            "latent_dim": LATENT_DIM,
+            "downsample_rate": HOP,
+            "target_sample_rate": SAMPLE_RATE,
+            "model_init_kwargs": {"latent_dim": LATENT_DIM},
+        },
+        "variant": variant,
+        "flash_t_grid": FLASH_T_GRID,
+        "defaults": {"nfe": 32, "cfg": 2.0, "sway": -1.0},
+    }
+    (model_dir / "config.json").write_text(json.dumps(config))
+    save_file(
+        {
+            "transformer.proj.weight": torch.zeros(LATENT_DIM, LATENT_DIM),
+            "layer_weights": torch.zeros(36),
+            "layer_scale": torch.ones(1),
+        },
+        str(model_dir / "auk.safetensors"),
+    )
+    save_file({"global_mean": torch.zeros(LATENT_DIM)}, str(model_dir / "vae.safetensors"))
+    return model_dir
+
+
+@pytest.fixture
+def build_pipeline(tmp_path, monkeypatch):
+    """Build an AuKPipeline whose transformer, codec and ODE are stubs."""
+
+    monkeypatch.setattr(pipeline_auk, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(pipeline_auk, "AuKTransformer", _StubTransformer)
+    monkeypatch.setattr(pipeline_auk, "AuKVAE", _StubVAE)
+
+    def _build(variant: str = "base") -> tuple[AuKPipeline, list[dict[str, Any]]]:
+        calls: list[dict[str, Any]] = []
+        monkeypatch.setattr(pipeline_auk, "sample_latents", _stub_sampler(calls))
+        od_config = OmniDiffusionConfig(
+            model=str(_write_checkpoint(tmp_path, variant)),
+            dtype=torch.float32,
+            model_class_name="AuKPipeline",
+        )
+        return AuKPipeline(od_config=od_config), calls
+
+    return _build
+
+
+def _prompt(
+    *,
+    tokens: int = 8,
+    audio: Any = None,
+    knobs: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Build the diffusion prompt the stage input processor would emit."""
+
+    prompt: dict[str, Any] = {
+        "prompt": "",
+        "prompt_embeds": torch.zeros(tokens, TEXT_HIDDEN_DIM),
+        "additional_information": {"auk": dict(knobs or {})},
+    }
+    if audio is not None:
+        prompt["multi_modal_data"] = {"audio": audio}
+    return prompt
+
+
+def _batch(prompt: dict[str, Any], **sampling: Any) -> DiffusionRequestBatch:
+    request = OmniDiffusionRequest(
+        prompt=prompt,
+        sampling_params=OmniDiffusionSamplingParams(**sampling),
+        request_id="auk-test-0",
+    )
+    return DiffusionRequestBatch(requests=[request])
+
+
+def _silence(seconds: float, sample_rate: int = SAMPLE_RATE) -> tuple[np.ndarray, int]:
+    return np.zeros(int(seconds * sample_rate), dtype=np.float32), sample_rate
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+class TestRequestParsing:
+    def test_pipeline_declares_audio_output(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        assert pipeline.support_audio_output is True
+        assert pipeline.audio_sample_rate == SAMPLE_RATE
+        assert pipeline.supports_request_batch is False
+        # The warmup request cannot carry an encoder-stage text condition.
+        assert pipeline.dummy_run_num_frames == 0
+
+    def test_config_reaches_the_transformer(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        assert pipeline.dit.init_kwargs["latent_dim"] == LATENT_DIM
+        assert pipeline.dit.init_kwargs["text_hidden_dim"] == TEXT_HIDDEN_DIM
+        assert pipeline.dit.init_kwargs["attn_mask_enabled"] is True
+
+    def test_gen_seconds_sets_the_target_length(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        outputs = pipeline.forward(_batch(_prompt(audio=_silence(2.0), knobs={"gen_seconds": 6.0}), seed=1))
+
+        assert calls[0]["gen_frames"] == math.ceil(6.0 * SAMPLE_RATE / HOP)
+        assert calls[0]["ref"].shape == (1, 100, LATENT_DIM)
+        assert calls[0]["ref_mask"].shape == (1, 100)
+        assert calls[0]["c_mask"].shape == (1, 8)
+        assert outputs[0].output.shape == (300 * HOP,)
+        assert outputs[0].output.dtype is torch.float32
+
+    def test_target_length_falls_back_to_the_source_clip(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(audio=_silence(2.0), knobs={"gen_seconds": None}), seed=1))
+
+        assert calls[0]["gen_frames"] == 100
+
+    def test_text_only_request_without_a_duration_is_rejected(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        with pytest.raises(ValueError, match="gen_seconds"):
+            pipeline.forward(_batch(_prompt(knobs={"gen_seconds": None}), seed=1))
+        assert calls == []
+
+    def test_text_only_request_gets_an_empty_reference(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(knobs={"gen_seconds": 3.5}), seed=1))
+
+        assert calls[0]["ref"].shape == (1, 0, LATENT_DIM)
+        assert calls[0]["gen_frames"] == 175
+
+    def test_source_clip_is_resampled_to_the_codec_rate(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(audio=_silence(1.0, 16000), knobs={"gen_seconds": 1.0}), seed=1))
+
+        assert pipeline.vae.encode_calls[0]["samples"] == SAMPLE_RATE
+
+    def test_stereo_source_is_mixed_to_mono(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        stereo = np.zeros((2, SAMPLE_RATE), dtype=np.float32)
+
+        pipeline.forward(_batch(_prompt(audio=(stereo, SAMPLE_RATE), knobs={"gen_seconds": 1.0}), seed=1))
+
+        assert pipeline.vae.encode_calls[0]["samples"] == SAMPLE_RATE
+
+    def test_missing_text_condition_is_rejected(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        prompt = _prompt(knobs={"gen_seconds": 1.0})
+        del prompt["prompt_embeds"]
+
+        with pytest.raises(ValueError, match="prompt_embeds"):
+            pipeline.forward(_batch(prompt, seed=1))
+
+    def test_base_variant_uses_the_requested_schedule(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        pipeline.forward(
+            _batch(
+                _prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": -1.0}),
+                seed=1,
+                num_inference_steps=16,
+                guidance_scale=3.0,
+            )
+        )
+
+        assert calls[0]["nfe"] == 16
+        assert calls[0]["cfg_strength"] == 3.0
+        assert calls[0]["sway_sampling_coef"] == -1.0
+        assert calls[0]["t_grid"] is None
+
+    def test_schedule_defaults_come_from_the_checkpoint(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": None}), seed=1))
+
+        assert calls[0]["nfe"] == 32
+        assert calls[0]["cfg_strength"] == 2.0
+        assert calls[0]["sway_sampling_coef"] == -1.0
+
+    def test_flash_variant_locks_the_distilled_recipe(self, build_pipeline):
+        pipeline, calls = build_pipeline("flash")
+
+        pipeline.forward(
+            _batch(
+                _prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": -1.0}),
+                seed=1,
+                num_inference_steps=32,
+                guidance_scale=2.0,
+            )
+        )
+
+        assert calls[0]["nfe"] == 4
+        assert calls[0]["cfg_strength"] == 0.0
+        assert calls[0]["sway_sampling_coef"] is None
+        assert calls[0]["t_grid"] == FLASH_T_GRID
+
+    def test_noise_comes_from_a_seeded_request_generator(self, build_pipeline):
+        pipeline, calls = build_pipeline()
+
+        outputs = pipeline.forward(
+            _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=0, output_type="latent"),
+        )
+
+        expected = torch.randn(1, 50, LATENT_DIM, generator=torch.Generator().manual_seed(0))
+        assert torch.equal(outputs[0].output, expected)
+        # The global RNG is never seeded on the pipeline's behalf.
+        assert calls[0]["seed"] is None
+        assert isinstance(calls[0]["generator"], torch.Generator)
+
+    def test_latent_output_type_skips_the_decoder(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        outputs = pipeline.forward(_batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1, output_type="latent"))
+
+        assert outputs[0].output.shape == (1, 50, LATENT_DIM)
+        assert pipeline.vae.decode_calls == 0
+
+    def test_vae_sample_draws_from_the_request_generator(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+
+        pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1))
+
+        call = pipeline.vae.encode_calls[0]
+        assert call["sample"] is True
+        assert isinstance(call["generator"], torch.Generator)
+
+    def test_one_request_per_forward(self, build_pipeline):
+        pipeline, _ = build_pipeline()
+        batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1)
+        batch.requests.append(batch.requests[0])
+
+        with pytest.raises(AssertionError, match="one request per forward"):
+            pipeline.forward(batch)
+
+
+def _reference_audio_path(messages: list[dict[str, Any]]) -> str:
+    for message in messages:
+        for item in message.get("content") or ():
+            if isinstance(item, dict) and item.get("type") == "audio":
+                path = item.get("audio") or item.get("audio_url")
+                if path:
+                    return str(path)
+    raise AssertionError("the saved reference case carries no source clip")
+
+
+def _log_mel_distance(left: torch.Tensor, right: torch.Tensor, sample_rate: int) -> float:
+    """Mean absolute log-mel difference between two mono waveforms."""
+
+    mel = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate, n_fft=1024, hop_length=256, n_mels=80)
+    length = min(left.shape[-1], right.shape[-1])
+    spectra = [torch.log(mel(wav[..., :length].float()) + 1e-5) for wav in (left, right)]
+    return float((spectra[0] - spectra[1]).abs().mean())
+
+
+@pytest.mark.local_model
+@pytest.mark.diffusion
+@pytest.mark.skipif(
+    not (IS_ASSEMBLED and PARITY_REF),
+    reason="needs AUK_PARITY_REF and AUK_OMNI_CKPT_DIR pointing at an assembled checkpoint directory",
+)
+def test_zero_shot_tts_tracks_the_upstream_reference():
+    """Replay the saved zero-shot TTS case and report the distances."""
+
+    parity_dir = Path(PARITY_REF)
+    fusion_dir = Path(os.environ.get("AUK_FUSION_REF") or parity_dir.parent.parent / "fusion_ref")
+    reference = torch.load(parity_dir / "zs_tts_en.pt", map_location="cpu", weights_only=False)
+    fusion = torch.load(fusion_dir / "zs_tts_audio.pt", map_location="cpu", weights_only=False)
+
+    # soundfile rather than torchaudio.load: the repo's test convention, and
+    # torchaudio's decoder needs an ffmpeg the port environment does not have.
+    samples, source_rate = sf.read(_reference_audio_path(reference["messages"]), dtype="float32", always_2d=True)
+    source = torch.from_numpy(samples).transpose(0, 1)
+    knobs = {"gen_seconds": reference["gen_seconds"], "sway": -1.0, "t_grid": None, "vae_sample": False}
+    prompt = {
+        "prompt": "",
+        # The encoder stage emits bf16; the saved fusion is an fp32 forward.
+        "prompt_embeds": fusion["fused"].to(torch.bfloat16),
+        "multi_modal_data": {"audio": (source, int(source_rate))},
+        "additional_information": {"auk": knobs},
+    }
+    od_config = OmniDiffusionConfig(model=CKPT_DIR, dtype=torch.bfloat16, model_class_name="AuKPipeline")
+    pipeline = AuKPipeline(od_config=od_config)
+
+    def run(output_type: str) -> torch.Tensor:
+        batch = _batch(
+            dict(prompt),
+            seed=reference["seed"],
+            num_inference_steps=32,
+            guidance_scale=2.0,
+            output_type=output_type,
+        )
+        return pipeline.forward(batch)[0].output
+
+    # Two deterministic runs: the ODE is seeded, so the latents belong to the
+    # waveform even though the decode happens in the second call.
+    latents = run("latent")
+    waveform = run("pt")
+
+    ref_frames = int(reference["ref_latent_lens"][0])
+    expected = reference["generated"][:, ref_frames:, :].float()
+    assert latents.shape == expected.shape
+    assert torch.isfinite(latents).all()
+    assert torch.isfinite(waveform).all()
+    assert waveform.shape[-1] == expected.shape[1] * HOP
+
+    cosine = F.cosine_similarity(latents[0], expected[0], dim=-1)
+    metrics = {
+        "gen_frames": int(expected.shape[1]),
+        "ref_frames": ref_frames,
+        "latent_mse": float(torch.mean((latents - expected) ** 2)),
+        "latent_relative_mse": float(torch.mean((latents - expected) ** 2) / expected.var()),
+        "frame_cosine_mean": float(cosine.mean()),
+        "frame_cosine_min": float(cosine.min()),
+        "log_mel_distance": _log_mel_distance(waveform, reference["audio"].reshape(-1), int(reference["sr"])),
+        "rms_port": float(waveform.pow(2).mean().sqrt()),
+        "rms_reference": float(reference["audio"].reshape(-1).pow(2).mean().sqrt()),
+    }
+    print("\nAuK zero-shot TTS parity:", json.dumps(metrics, indent=1))
+
+    assert metrics["frame_cosine_mean"] > PARITY_MIN_FRAME_COSINE
+    assert metrics["latent_relative_mse"] < PARITY_MAX_RELATIVE_MSE
+    assert metrics["log_mel_distance"] < PARITY_MAX_MEL_DISTANCE

tests/diffusion/models/bagel/test_bagel_lora.py

@@ -122,8 +122,10 @@ class TestStage1DiTLoRA:
             lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
         )
 
-        # Build pipeline with bagel component
+        # Build pipeline with bagel component declared via _lora_components,
+        # matching how BagelPipeline opts in.
         pipeline = torch.nn.Module()
+        pipeline._lora_components = ["bagel"]
         pipeline.bagel = torch.nn.Module()
         lm = torch.nn.Module()
         lm.stacked_params_mapping = [
@@ -194,8 +196,11 @@ class TestBagelLoRARoundTrip:
 
         monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA)
 
-        # Build pipeline with bagel.language_model.foo (simple non-packed layer)
+        # Build pipeline with bagel.language_model.foo (simple non-packed layer);
+        # the bagel component is declared via _lora_components, matching how
+        # BagelPipeline opts in.
         pipeline = torch.nn.Module()
+        pipeline._lora_components = ["bagel"]
         pipeline.bagel = torch.nn.Module()
         lm = torch.nn.Module()
         lm.foo = _FakeLinearBase()

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -26,7 +26,9 @@ NORM_EPS = 1e-5
 def _init_distributed(request):
     """Initialize the minimal single-rank distributed environment required by
     the vLLM parallel linear layers (tensor-parallel group must exist)."""
-    if request.node.name.startswith("test_real_rotary_"):
+    if request.node.name.startswith(
+        ("test_real_rotary_", "test_packed_rope_table", "test_fused_qk_norm_rope", "test_fallback_with_no")
+    ):
         yield
         return
 
@@ -438,29 +440,86 @@ def test_transformer_rejects_prompt_tuning():
         BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True}))
 
 
-def _native_to_checkpoint_name(name: str) -> str:
-    """Inverse of ``load_weights`` remapping: native param name -> diffusers name.
+_QKV_FANOUT = {
+    "to_qkv": ("to_q", "to_k", "to_v"),
+    "img_to_qkv": ("img_to_q", "img_to_k", "img_to_v"),
+    "instruct_to_qkv": ("instruct_to_q", "instruct_to_k", "instruct_to_v"),
+}
+
+
+def _native_to_checkpoint_weights(name: str, param: torch.Tensor) -> list[tuple[str, torch.Tensor]]:
+    """Inverse of ``load_weights`` remapping: native param -> diffusers weights.
+
+    The diffusers checkpoint stores fused projections as separate matrices, so a
+    single merged native param (QKV or FFN gate/up) fans out into several
+    checkpoint weights. Returns ``(diffusers_name, value)`` pairs whose values,
+    when fed through ``load_weights``, reassemble into the original ``param``.
 
     - ``.to_out.<suffix>`` -> ``.to_out.0.<suffix>`` (diffusers ModuleList wrap).
     - promoted joint-attention projections move back under ``.processor.``.
     """
+    q_size = NUM_HEADS * HEAD_DIM
+    kv_size = NUM_KV_HEADS * HEAD_DIM
+
+    # Fused QKV projections split back into per-matrix q/k/v weights.
+    for token, (q_name, k_name, v_name) in _QKV_FANOUT.items():
+        marker = f".{token}."
+        if marker in name:
+            q, k, v = param[:q_size], param[q_size : q_size + kv_size], param[q_size + kv_size :]
+            results = []
+            for sub_name, value in ((q_name, q), (k_name, k), (v_name, v)):
+                ckpt_name = name.replace(marker, f".{sub_name}.")
+                if token != "to_qkv":
+                    # Joint-attention projections live under `.processor.` upstream.
+                    ckpt_name = ckpt_name.replace(".img_instruct_attn.", ".img_instruct_attn.processor.")
+                results.append((ckpt_name, value))
+            return results
+
+    # Fused FFN gate/up splits into linear_1 (gate) / linear_3 (input).
+    if ".gate_up_proj." in name:
+        inner = param.shape[0] // 2
+        return [
+            (name.replace(".gate_up_proj.", ".linear_1."), param[:inner]),
+            (name.replace(".gate_up_proj.", ".linear_3."), param[inner:]),
+        ]
+
+    # Default: 1:1 with the existing diffusers name promotions.
     if ".to_out." in name:
         name = name.replace(".to_out.", ".to_out.0.")
-    for proj in (
-        "img_to_q",
-        "img_to_k",
-        "img_to_v",
-        "instruct_to_q",
-        "instruct_to_k",
-        "instruct_to_v",
-        "instruct_out",
-        "img_out",
-    ):
+    for proj in ("instruct_out", "img_out"):
         token = f".img_instruct_attn.{proj}."
         if token in name:
             name = name.replace(token, f".img_instruct_attn.processor.{proj}.")
             break
-    return name
+    return [(name, param)]
+
+
+def test_transformer_exposes_stacked_params_mapping():
+    """The packed -> sub-layer mapping must be discoverable from the module tree.
+
+    ``diffusion/lora/loader.py`` and the quantized weight loaders read
+    ``stacked_params_mapping`` off the model, so it has to exist before (and
+    independently of) ``load_weights``.
+    """
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _BOOGU_STACKED_PARAMS_MAPPING,
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+
+    assert tuple(model.stacked_params_mapping) == _BOOGU_STACKED_PARAMS_MAPPING
+    # A per-instance copy, so consumers cannot mutate the module constant.
+    assert model.stacked_params_mapping is not _BOOGU_STACKED_PARAMS_MAPPING
+
+    # Every entry is a (param, shard, shard_id) triple.
+    for param_name, shard_name, shard_id in model.stacked_params_mapping:
+        assert param_name.startswith(".") and shard_name.startswith(".")
+        assert shard_id in {"q", "k", "v", 0, 1}
+
+    # The mapping targets exactly the fused projections this port creates.
+    mapped_leaves = {param.strip(".").split(".")[-1] for param, _, _ in model.stacked_params_mapping}
+    assert mapped_leaves == {"to_qkv", "img_to_qkv", "instruct_to_qkv", "gate_up_proj"}
 
 
 def test_transformer_load_weights_round_trip():
@@ -471,24 +530,59 @@ def test_transformer_load_weights_round_trip():
     model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
     native_params = dict(model.named_parameters())
 
-    # Build synthetic diffusers-named weights (one per native parameter).
-    checkpoint_weights = {}
+    # Build synthetic diffusers-named weights. Fused native projections fan out
+    # into the separate matrices the diffusers checkpoint stores.
+    expected: dict[str, torch.Tensor] = {}
+    checkpoint_weights: dict[str, torch.Tensor] = {}
     for native_name, param in native_params.items():
-        checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param)
+        full = torch.randn_like(param)
+        expected[native_name] = full
+        for ckpt_name, value in _native_to_checkpoint_weights(native_name, full):
+            checkpoint_weights[ckpt_name] = value
 
-    # The remapping must be a bijection over the parameter set.
-    assert len(checkpoint_weights) == len(native_params)
+    # Every checkpoint name is unique; merged params fan out to >=1 weight.
+    assert len(checkpoint_weights) >= len(native_params)
 
     loaded = model.load_weights(list(checkpoint_weights.items()))
 
     # No missing / unexpected parameters.
     assert loaded == set(native_params.keys())
 
-    # Values landed on the right parameters (TP=1: weight_loader copies verbatim).
+    # Values landed on the right parameters (TP=1: weight_loader copies verbatim;
+    # fused shards are placed at their q/k/v / gate/up offsets).
     reloaded = dict(model.named_parameters())
     for native_name in native_params:
-        expected = checkpoint_weights[_native_to_checkpoint_name(native_name)]
-        assert torch.allclose(reloaded[native_name], expected)
+        assert torch.allclose(reloaded[native_name], expected[native_name])
+
+
+def test_transformer_load_weights_rejects_partial_fused_params():
+    """A fused parameter must not count as loaded until all its shards arrive.
+
+    The loader compares parameter *names*, so reporting ``to_qkv`` complete
+    after a single ``to_q`` would let a checkpoint carrying only ``to_q`` start
+    up with the ``k``/``v`` slices left uninitialized.
+    """
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+    native_params = dict(model.named_parameters())
+
+    complete_name = "noise_refiner.0.attn.to_qkv.weight"
+    partial_name = "noise_refiner.0.feed_forward.gate_up_proj.weight"
+    assert complete_name in native_params
+    assert partial_name in native_params
+
+    weights = list(_native_to_checkpoint_weights(complete_name, native_params[complete_name]))
+    # Only the gate half of the FFN, so its fused parameter stays incomplete.
+    gate_ckpt, gate_value = _native_to_checkpoint_weights(partial_name, native_params[partial_name])[0]
+    weights.append((gate_ckpt, gate_value))
+
+    loaded = model.load_weights(weights)
+
+    assert complete_name in loaded
+    assert partial_name not in loaded
 
 
 def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
@@ -511,12 +605,15 @@ def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
     try:
         model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
         native_params = dict(model.named_parameters())
-        loaded_name = next(iter(native_params))
-        checkpoint_name = _native_to_checkpoint_name(loaded_name)
+        # A non-fused parameter, so one checkpoint entry loads it completely.
+        loaded_name = "x_embedder.weight"
+        assert loaded_name in native_params
+        (checkpoint_name, checkpoint_value), *_ = _native_to_checkpoint_weights(loaded_name, native_params[loaded_name])
+        checkpoint_value = torch.randn_like(checkpoint_value)
 
         loaded = model.load_weights(
             [
-                (checkpoint_name, torch.randn_like(native_params[loaded_name])),
+                (checkpoint_name, checkpoint_value),
                 ("unexpected.weight", torch.ones(1)),
             ]
         )
@@ -600,3 +697,152 @@ def test_transformer_forward_ti2i_shape():
 
     assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
     assert torch.isfinite(out).all()
+
+
+# ---------------------------------------------------------------------------
+# The three tests below are for the fused qk-norm+RoPE kernel's wiring into
+# this transformer: the packed-table layout conversion, the fused path against
+# the previous eager chain, and the fallback when no packed table is present.
+# They need CUDA (the module-level marks are CPU; these add the cuda mark and
+# skip on CPU-only runners) and real Boogu geometry rather than the mock dims
+# above.
+# ---------------------------------------------------------------------------
+
+_FUSED_HEAD_DIM = 120
+_FUSED_Q_HEADS = 28
+_FUSED_KV_HEADS = 7
+_FUSED_EPS = 1e-5
+
+
+def _fused_rotary_pair(tokens: int):
+    theta = torch.randn(1, tokens, _FUSED_HEAD_DIM // 2, device="cuda", dtype=torch.float32)
+    return (
+        torch.cos(theta).repeat_interleave(2, dim=-1),
+        torch.sin(theta).repeat_interleave(2, dim=-1),
+    )
+
+
+def _fused_norms():
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.layernorm import RMSNorm
+
+    with set_current_vllm_config(VllmConfig()):
+        norm_q = RMSNorm(_FUSED_HEAD_DIM, eps=_FUSED_EPS).cuda().to(torch.bfloat16)
+        norm_k = RMSNorm(_FUSED_HEAD_DIM, eps=_FUSED_EPS).cuda().to(torch.bfloat16)
+    with torch.no_grad():
+        norm_q.weight.copy_(torch.randn(_FUSED_HEAD_DIM))
+        norm_k.weight.copy_(torch.randn(_FUSED_HEAD_DIM))
+    return norm_q, norm_k
+
+
+def _operand_ulp_bound(x, weight, cos, sin):
+    """One bf16 rounding of a normalized operand, propagated through RoPE."""
+    import torch.nn.functional as F
+
+    n = F.rms_norm(x, (_FUSED_HEAD_DIM,), weight, _FUSED_EPS).float().abs()
+    c = cos[..., ::2].float().abs().unsqueeze(2)
+    s = sin[..., ::2].float().abs().unsqueeze(2)
+    even_mag = n[..., ::2] * c + n[..., 1::2] * s
+    odd_mag = n[..., ::2] * s + n[..., 1::2] * c
+    return 2.0**-6 * torch.stack((even_mag, odd_mag), dim=-1).flatten(-2) + 1e-6
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_packed_rope_table_layout(monkeypatch):
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _with_packed_rope_table,
+    )
+
+    monkeypatch.delenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", raising=False)
+    torch.manual_seed(0)
+    tokens = 2311  # above _FUSED_MIN_TOKENS
+    cos, sin = _fused_rotary_pair(tokens)
+    out = _with_packed_rope_table((cos, sin))
+    assert len(out) == 3
+    assert out[0] is cos and out[1] is sin
+    packed = out[2]
+    assert packed.shape == (tokens, _FUSED_HEAD_DIM)
+    assert packed.dtype == torch.float32 and packed.is_contiguous()
+    half = _FUSED_HEAD_DIM // 2
+    assert torch.equal(packed[:, :half], cos[0, :, ::2])
+    assert torch.equal(packed[:, half:], sin[0, :, ::2])
+    # Below _FUSED_MIN_TOKENS the tuple passes through unchanged (host-bound
+    # regime: the fused path would cost more than it saves).
+    short = _with_packed_rope_table(_fused_rotary_pair(17))
+    assert len(short) == 2
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_fused_qk_norm_rope_matches_eager_chain():
+    """The fused path vs the previous eager chain at real Boogu shapes.
+
+    Every element must sit within one bf16 ulp of a normalized operand
+    propagated through the rotation (the two paths legitimately differ by
+    single- vs double-rounding of the norm's weight multiply).
+    """
+    from vllm.triton_utils import HAS_TRITON
+
+    if not HAS_TRITON:
+        pytest.skip("Triton required")
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _qk_norm_rope,
+        _with_packed_rope_table,
+    )
+
+    torch.manual_seed(0)
+    norm_q, norm_k = _fused_norms()
+    q = torch.randn(1, 4139, _FUSED_Q_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 4139, _FUSED_KV_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    cos, sin = _fused_rotary_pair(4139)
+    with_table = _with_packed_rope_table((cos, sin))
+
+    fused = _qk_norm_rope(q, k, norm_q, norm_k, with_table, torch.bfloat16)
+    eager = _qk_norm_rope(q, k, norm_q, norm_k, (cos, sin), torch.bfloat16)
+    for fused_t, eager_t, x, w in zip(fused, eager, (q, k), (norm_q.weight, norm_k.weight)):
+        bound = _operand_ulp_bound(x, w, cos, sin)
+        diff = (fused_t.float() - eager_t.float()).abs()
+        assert (diff <= bound).all(), "fused path beyond one operand ulp of the eager chain"
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_fallback_with_no_packed_table():
+    """Without a packed table the helper is bit-exact to the old chain."""
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _qk_norm_rope,
+        apply_rotary_emb,
+    )
+
+    torch.manual_seed(0)
+    norm_q, norm_k = _fused_norms()
+    q = torch.randn(1, 64, _FUSED_Q_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 64, _FUSED_KV_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    cos, sin = _fused_rotary_pair(64)
+
+    out_q, out_k = _qk_norm_rope(q, k, norm_q, norm_k, (cos, sin), torch.bfloat16)
+    ref_q = apply_rotary_emb(norm_q(q), (cos, sin)).to(torch.bfloat16)
+    ref_k = apply_rotary_emb(norm_k(k), (cos, sin)).to(torch.bfloat16)
+    assert torch.equal(out_q, ref_q)
+    assert torch.equal(out_k, ref_k)
+
+
+def test_packed_rope_table_env_override(monkeypatch):
+    """VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS steers the packing gate: 0 packs
+    even a 43-token table, a huge value leaves a 4139-token table untouched,
+    unset keeps the 2048 default routing."""
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _with_packed_rope_table,
+    )
+
+    env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+    short, long = _identity_rotary_emb(1, 43), _identity_rotary_emb(1, 4139)
+
+    monkeypatch.setenv(env, "0")
+    assert len(_with_packed_rope_table(short)) == 3
+    monkeypatch.setenv(env, str(10**9))
+    assert _with_packed_rope_table(long) is long
+    monkeypatch.delenv(env)
+    assert len(_with_packed_rope_table(long)) == 3
+    assert _with_packed_rope_table(short) is short

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -13,11 +13,13 @@ Two groups:
    CFG handling, and reshape logic can be verified numerically on CPU.
 """
 
+import os
 from types import SimpleNamespace
 
 import pytest
 import torch
 from torch import nn
+from transformers import Qwen3VLConfig
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig, TransformerConfig
 
@@ -56,9 +58,17 @@ def mock_dependencies(mocker, monkeypatch):
         f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
         lambda *a, **k: mock_scheduler,
     )
-    monkeypatch.setattr(
+    mllm_loader = mocker.patch(
         f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
-        lambda *a, **k: mllm_wrapper,
+        return_value=mllm_wrapper,
+    )
+    mllm_config_loader = mocker.patch(
+        f"{_MODULE}.Qwen3VLConfig.from_pretrained",
+        return_value=Qwen3VLConfig(),
+    )
+    mllm_builder = mocker.patch(
+        f"{_MODULE}.create_transformers_model_with_vllm_linears",
+        return_value=nn.Linear(1, 1),
     )
     monkeypatch.setattr(
         f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
@@ -74,12 +84,17 @@ def mock_dependencies(mocker, monkeypatch):
     mock_transformer_cls.return_value = mock_transformer_instance
     monkeypatch.setattr(f"{_MODULE}.BooguImageTransformer2DModel", mock_transformer_cls)
 
-    # Treat the dummy model id as a local path: skips hub prefetch.
-    mocker.patch("os.path.exists", return_value=True)
+    # Treat only the dummy model id as local. Other filesystem checks (for
+    # example lazy imports in the quantization registry) must remain real.
+    path_exists = os.path.exists
+    mocker.patch("os.path.exists", side_effect=lambda path: str(path).startswith("dummy-boogu") or path_exists(path))
 
     return {
         "inner_encoder": inner_encoder,
         "mllm_wrapper": mllm_wrapper,
+        "mllm_loader": mllm_loader,
+        "mllm_config_loader": mllm_config_loader,
+        "mllm_builder": mllm_builder,
         "processor": mock_processor,
         "vae": mock_vae,
         "scheduler": mock_scheduler,
@@ -128,6 +143,9 @@ def test_constructor_wires_components(boogu_pipeline, mock_dependencies):
     assert boogu_pipeline.vae_scale_factor == 8
     assert boogu_pipeline.default_sample_size == 128
     assert hasattr(boogu_pipeline, "load_weights")
+    assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
+    mock_dependencies["mllm_config_loader"].assert_not_called()
+    mock_dependencies["mllm_builder"].assert_not_called()
 
 
 def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
@@ -188,7 +206,7 @@ def test_constructor_weights_sources(boogu_pipeline):
     assert source.fall_back_to_pt is True
 
 
-def test_constructor_routes_only_transformer_config(mock_dependencies, mocker):
+def test_constructor_rejects_unsupported_mllm_quantization(mock_dependencies, mocker):
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
     from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
@@ -204,10 +222,143 @@ def test_constructor_routes_only_transformer_config(mock_dependencies, mocker):
             {"transformer": transformer_config, "mllm": encoder_config, "vae": None}
         ),
     )
+    with pytest.raises(ValueError, match="Boogu MLLM only supports FP8 quantization"):
+        BooguImagePipeline(od_config=od_config)
+    mock_dependencies["mllm_loader"].assert_not_called()
+    mock_dependencies["mllm_config_loader"].assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("quantization_config", "quantize_mllm"),
+    [
+        pytest.param("fp8", True, id="global-fp8"),
+        pytest.param({"mllm": {"method": "fp8"}, "transformer": {"method": "fp8"}}, True, id="mllm-and-dit-fp8"),
+        pytest.param({"mllm": None, "transformer": "fp8"}, False, id="dit-only-fp8"),
+    ],
+)
+def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm):
+    from transformers import AutoModel
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        revision="test-revision",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.bfloat16,
+        quantization_config=quantization_config,
+    )
+    pipeline = BooguImagePipeline(od_config=od_config)
+
+    if quantize_mllm:
+        mock_dependencies["mllm_loader"].assert_not_called()
+        mock_dependencies["mllm_config_loader"].assert_called_once_with(
+            "dummy-boogu", subfolder="mllm", local_files_only=True, revision="test-revision"
+        )
+        builder = mock_dependencies["mllm_builder"]
+        builder.assert_called_once()
+        auto_cls, hf_config, online_quant_config = builder.call_args.args
+        assert auto_cls is AutoModel
+        assert hf_config is mock_dependencies["mllm_config_loader"].return_value
+        assert isinstance(online_quant_config, OnlineQuantizationConfig)
+        assert online_quant_config.args.linear.weight == kFp8Static128BlockSym
+        assert builder.call_args.kwargs == {
+            "dtype": torch.bfloat16,
+            "device": pipeline._execution_device,
+            "prefix": "mllm",
+            "skip_modules": ("mllm.visual",),
+        }
+        assert pipeline.mllm is builder.return_value
+        assert not pipeline.mllm.training
+        assert all(not parameter.requires_grad for parameter in pipeline.mllm.parameters())
+        transformer_source, mllm_source = pipeline.weights_sources
+        assert transformer_source.subfolder == "transformer"
+        assert mllm_source.model_or_path == "dummy-boogu"
+        assert mllm_source.subfolder == "mllm"
+        assert mllm_source.prefix == "mllm."
+        assert mllm_source.revision == "test-revision"
+    else:
+        assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
+        assert pipeline.mllm is mock_dependencies["inner_encoder"]
+        mock_dependencies["mllm_config_loader"].assert_not_called()
+        mock_dependencies["mllm_builder"].assert_not_called()
+        assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
+    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config)
+
+
+def test_constructor_preserves_mllm_ignored_layers(mock_dependencies):
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    quant_config = Fp8Config(
+        ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.bfloat16,
+        quantization_config=quant_config,
+    )
     BooguImagePipeline(od_config=od_config)
-    kwargs = mock_dependencies["transformer_cls"].call_args.kwargs
-    assert kwargs["quant_config"] is transformer_config
-    assert kwargs["prefix"] == "transformer"
+
+    online_quant_config = mock_dependencies["mllm_builder"].call_args.args[2]
+    assert online_quant_config.ignored_layers == quant_config.ignored_layers
+
+
+def test_constructor_preserves_serialized_mllm_quantization(mock_dependencies):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    mock_dependencies["mllm_config_loader"].return_value = Qwen3VLConfig(
+        quantization_config={"quant_method": "fp8", "modules_to_not_convert": ["model.visual"]}
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu-fp8",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.bfloat16,
+        quantization_config="fp8",
+    )
+    pipeline = BooguImagePipeline(od_config=od_config)
+
+    # No override: HF reads the serialized checkpoint's scales and skip list.
+    assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
+    assert pipeline.mllm is mock_dependencies["inner_encoder"]
+    mock_dependencies["mllm_builder"].assert_not_called()
+    assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
+
+
+def test_load_weights_maps_mllm_and_preserves_transformer():
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    pipeline = object.__new__(BooguImagePipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.mllm = nn.ModuleDict(
+        {
+            "language_model": nn.Linear(2, 2, bias=False),
+            "visual": nn.Linear(2, 2, bias=False),
+        }
+    )
+    pipeline.transformer = nn.Linear(2, 2, bias=False)
+    language_weight = torch.full((2, 2), 1.0)
+    visual_weight = torch.full((2, 2), 2.0)
+    transformer_weight = torch.full((2, 2), 3.0)
+
+    loaded = pipeline.load_weights(
+        [
+            ("mllm.lm_head.weight", torch.full((2, 2), 9.0)),
+            ("mllm.model.language_model.weight", language_weight),
+            ("mllm.model.visual.weight", visual_weight),
+            ("transformer.weight", transformer_weight),
+        ]
+    )
+
+    assert loaded == {"mllm.language_model.weight", "mllm.visual.weight", "transformer.weight"}
+    torch.testing.assert_close(pipeline.mllm["language_model"].weight, language_weight)
+    torch.testing.assert_close(pipeline.mllm["visual"].weight, visual_weight)
+    torch.testing.assert_close(pipeline.transformer.weight, transformer_weight)
 
 
 @pytest.mark.parametrize(

tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from importlib.util import find_spec
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_FRAMEWORK_AVAILABLE = find_spec("cosmos_framework") is not None
+requires_cosmos_framework = pytest.mark.skipif(
+    not _FRAMEWORK_AVAILABLE,
+    reason="cosmos_framework source is not present on PYTHONPATH",
+)
+
+
+def test_framework_preflight_checks_every_policy_symbol(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils
+
+    imported = []
+
+    def record_import(module_name: str, symbol_name: str, error_message: str):
+        imported.append((module_name, symbol_name, error_message))
+        return object()
+
+    monkeypatch.setattr(cosmos3_utils, "lazy_import", record_import)
+
+    cosmos3_utils.preflight_cosmos3_action_framework_imports()
+
+    expected_symbols = {
+        "ActionTransformPipeline",
+        "FlowUniPCMultistepScheduler",
+        "build_abs_pose_from_components",
+        "convert_rotation",
+        "get_domain_id",
+        "pose_abs_to_rel",
+        "pose_rel_to_abs",
+    }
+    assert {symbol_name for _, symbol_name, _ in imported} == expected_symbols
+    assert [(module_name, symbol_name) for module_name, symbol_name, _ in imported] == list(
+        cosmos3_utils._ROBOLAB_FRAMEWORK_IMPORTS.values()
+    )
+
+
+def test_lazy_import_reports_framework_symbol_skew(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils
+
+    monkeypatch.setattr(cosmos3_utils, "import_module", lambda _module_name: object())
+
+    with pytest.raises(ImportError, match="MissingSymbol") as exc_info:
+        cosmos3_utils.lazy_import(
+            "cosmos_framework.fake_module",
+            "MissingSymbol",
+            "Cosmos Framework compatibility error.",
+        )
+
+    assert isinstance(exc_info.value.__cause__, AttributeError)
+
+
+@requires_cosmos_framework
+def test_action_domain_table_matches_current_framework_and_legacy_aliases() -> None:
+    from cosmos_framework.data.generator.action.utils.domain_utils import (
+        EMBODIMENT_TO_DOMAIN_ID as FRAMEWORK_DOMAIN_IDS,
+    )
+
+    from vllm_omni.diffusion.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID as VLLM_DOMAIN_IDS
+    from vllm_omni.diffusion.models.cosmos3.utils import _LEGACY_EMBODIMENT_ALIASES
+
+    canonical_vllm_domain_ids = {
+        name: domain_id for name, domain_id in VLLM_DOMAIN_IDS.items() if name not in _LEGACY_EMBODIMENT_ALIASES
+    }
+
+    assert canonical_vllm_domain_ids == FRAMEWORK_DOMAIN_IDS
+    assert {alias: VLLM_DOMAIN_IDS[alias] for alias in _LEGACY_EMBODIMENT_ALIASES} == {
+        alias: FRAMEWORK_DOMAIN_IDS[canonical_name] for alias, canonical_name in _LEGACY_EMBODIMENT_ALIASES.items()
+    }
+
+
+@requires_cosmos_framework
+def test_current_cosmos_framework_policy_modules_import_and_instantiate() -> None:
+    from cosmos_framework.data.generator.action.utils.domain_utils import get_domain_id
+    from cosmos_framework.data.generator.action.utils.pose_utils import (
+        build_abs_pose_from_components,
+        convert_rotation,
+        pose_abs_to_rel,
+        pose_rel_to_abs,
+    )
+    from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline
+    from cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc import (
+        FlowUniPCMultistepScheduler,
+    )
+
+    transform = ActionTransformPipeline(
+        max_action_dim=64,
+        cfg_dropout_rate=0.0,
+        format_prompt_as_json=True,
+    )
+    scheduler = FlowUniPCMultistepScheduler(
+        num_train_timesteps=1000,
+        shift=1.0,
+        use_dynamic_shifting=False,
+    )
+    scheduler.set_timesteps(4, device=torch.device("cpu"), shift=5.0)
+
+    assert transform.prompt_json_formatter is not None
+    assert len(scheduler.timesteps) == 4
+    assert get_domain_id("droid_lerobot") == 8
+    assert convert_rotation([0.0, 0.0, 0.0, 1.0], "quat_xyzw", "rot6d").shape == (6,)
+    assert callable(pose_abs_to_rel)
+    assert callable(pose_rel_to_abs)
+    assert callable(build_abs_pose_from_components)
+
+
+@requires_cosmos_framework
+def test_real_action_transform_pipeline_matches_droid_policy_contract() -> None:
+    from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline
+
+    transform = ActionTransformPipeline(
+        max_action_dim=64,
+        cfg_dropout_rate=0.0,
+        format_prompt_as_json=True,
+    )
+    sample = {
+        "ai_caption": "Pick up the cube.",
+        "video": torch.zeros((3, 3, 192, 320), dtype=torch.uint8),
+        "action": torch.zeros((3, 8), dtype=torch.float32),
+        "conditioning_fps": torch.tensor(15, dtype=torch.long),
+        "mode": "wam",
+        "domain_id": torch.tensor(8, dtype=torch.long),
+        "viewpoint": "concat_view",
+        "additional_view_description": "Wrist view above two exterior views.",
+    }
+
+    transformed = transform(sample, "256")
+
+    assert isinstance(transformed["ai_caption"], dict)
+    assert transformed["action"].shape == (3, 64)
+    assert int(transformed["raw_action_dim"].item()) == 8
+    assert transformed["sequence_plan"].has_action is True
+    assert transformed["mode"] == "wam"
+    assert int(transformed["domain_id"].item()) == 8

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -1,13 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
+import json
 import sys
 import types
+from contextlib import contextmanager
 from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
+from unittest.mock import Mock
 
 import numpy as np
 import pytest
@@ -21,6 +24,20 @@ from vllm_omni.experimental.world_models.session_state import SessionStateManage
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.mark.parametrize(
+    ("alias", "canonical_name"),
+    [
+        ("galbot", "embodiment_b"),
+        ("agibot_gear_gripper", "embodiment_c_gripper"),
+        ("agibot_gear_gripper_ext", "embodiment_c_gripper_ext"),
+    ],
+)
+def test_action_domain_table_preserves_legacy_aliases(alias: str, canonical_name: str) -> None:
+    from vllm_omni.diffusion.models.cosmos3.action import resolve_domain_id
+
+    assert resolve_domain_id(domain_name=alias) == resolve_domain_id(domain_name=canonical_name)
+
+
 def test_pipeline_declares_layerwise_offload_components() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
 
@@ -69,6 +86,7 @@ class StubScheduler:
         flow_shift: float = 1.0,
     ) -> None:
         self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
+        self.sigmas = self.timesteps.float() / 10
         self.config = SimpleNamespace(
             num_train_timesteps=1000,
             flow_shift=flow_shift,
@@ -97,6 +115,7 @@ class StubScheduler:
         else:
             assert num_inference_steps is not None
             self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device)
+        self.sigmas = self.timesteps.float()
 
     def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs):
         del kwargs
@@ -234,7 +253,7 @@ def passthrough_progress_bar(iterable):
 
 @pytest.fixture(autouse=True)
 def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
-    module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
+    module: Any = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
     module.is_guardrails_enabled = lambda od_config, sampling_params=None: False
     module.ensure_initialized = lambda od_config: None
     module.check_text_safety = lambda text: None
@@ -267,6 +286,8 @@ def make_cosmos3_pipeline():
         pipeline.is_edge_model = False
         pipeline._guidance_scale = None
         pipeline._num_timesteps = None
+        pipeline._current_step_index = None
+        pipeline._current_sigma = None
         pipeline._cosmos3_branch_caches = None
         pipeline._cache_dit_requires_paired_cfg = False
         pipeline._sound_tokenizer = None
@@ -284,7 +305,7 @@ def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 def make_sampling_params(**overrides: Any) -> SimpleNamespace:
-    values = {
+    values: dict[str, Any] = {
         "height": None,
         "width": None,
         "num_frames": None,
@@ -361,23 +382,31 @@ def _capture_tokenize_calls(pipeline: Any) -> list[dict[str, Any]]:
 
 
 @pytest.mark.parametrize(
-    ("provided", "value", "default", "is_distilled", "expected"),
+    ("provided", "value", "default", "is_distilled", "expected", "expected_warning"),
     [
-        (False, 1.0, 7.0, False, 7.0),
-        (True, 1.0, 7.0, False, 1.0),
-        (True, 4.5, 7.0, False, 4.5),
-        (False, 1.0, 7.0, True, 1.0),
-        (True, 4.5, 7.0, True, 1.0),
+        (False, 1.0, 7.0, False, 7.0, False),
+        (True, 1.0, 7.0, False, 1.0, False),
+        (True, 4.5, 7.0, False, 4.5, False),
+        (False, 1.0, 7.0, True, 1.0, False),
+        (True, 1.0, 7.0, True, 1.0, False),
+        (True, 4.5, 7.0, True, 1.0, True),
+        (True, 0.0, 7.0, True, 1.0, True),
     ],
 )
 def test_resolve_guidance_scale(
     make_cosmos3_pipeline,
+    monkeypatch: pytest.MonkeyPatch,
     provided: bool,
     value: float,
     default: float,
     is_distilled: bool,
     expected: float,
+    expected_warning: bool,
 ) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+    warning_once = Mock()
+    monkeypatch.setattr(pipeline_cosmos3.logger, "warning_once", warning_once)
     pipeline = make_cosmos3_pipeline()
     pipeline.is_distilled_model = is_distilled
     sp = make_sampling_params(
@@ -386,6 +415,12 @@ def test_resolve_guidance_scale(
     )
 
     assert pipeline._resolve_guidance_scale(sp, default) == expected
+    if expected_warning:
+        warning_once.assert_called_once()
+        assert "overridden to 1.0" in warning_once.call_args.args[0]
+        assert "negative_prompt does not affect generation" in warning_once.call_args.args[0]
+    else:
+        warning_once.assert_not_called()
 
 
 def test_distilled_generation_accepts_t2i_and_i2v(make_cosmos3_pipeline) -> None:
@@ -479,6 +514,69 @@ def test_forward_threads_request_id_to_robolab(make_cosmos3_pipeline) -> None:
     assert captured["session_id"] == "robolab-request-7"
 
 
+@pytest.mark.parametrize("format_prompt_as_json", [False, True])
+def test_robolab_input_builder_threads_prompt_format_and_uses_wam(
+    make_cosmos3_pipeline,
+    monkeypatch: pytest.MonkeyPatch,
+    format_prompt_as_json: bool,
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+    pipeline = make_cosmos3_pipeline()
+    pipeline.transformer = StubCosmos3Transformer(action_gen=True, action_dim=64)
+    captured: dict[str, Any] = {}
+
+    def fake_transform(sample, resolution):
+        captured["sample_mode"] = sample["mode"]
+        captured["resolution"] = resolution
+        sample["sequence_plan"] = SimpleNamespace(
+            condition_frame_indexes_action=[0],
+            action_start_frame_offset=1,
+        )
+        sample["raw_action_dim"] = torch.tensor(8)
+        sample["image_size"] = torch.tensor([16, 16, 16, 16])
+        if format_prompt_as_json:
+            sample["ai_caption"] = {"actions": {"instruction": sample["ai_caption"]}}
+        return sample
+
+    def fake_get_transform(*, format_prompt_as_json: bool):
+        captured["format_prompt_as_json"] = format_prompt_as_json
+        return fake_transform
+
+    pipeline._get_robolab_transform = fake_get_transform
+    monkeypatch.setattr(pipeline_cosmos3, "get_robolab_domain_id", lambda name: 8)
+    obs = {
+        "prompt": "Pick up the cube.",
+        "observation/image": np.zeros((16, 16, 3), dtype=np.uint8),
+        "observation/joint_position": np.zeros(7, dtype=np.float32),
+        "observation/gripper_position": np.zeros(1, dtype=np.float32),
+    }
+    sampling_params = make_sampling_params(
+        extra_args={
+            "robot_obs": obs,
+            "action_chunk_size": 2,
+            "image_height": 16,
+            "image_width": 16,
+            "format_prompt_as_json": format_prompt_as_json,
+        }
+    )
+
+    inputs = pipeline._build_robolab_policy_inputs(sampling_params, request_id="request-1")
+
+    assert inputs is not None
+    assert captured == {
+        "sample_mode": "wam",
+        "resolution": "480",
+        "format_prompt_as_json": format_prompt_as_json,
+    }
+    assert inputs.domain_id == 8
+    assert inputs.raw_action_dim == 8
+    if format_prompt_as_json:
+        assert json.loads(inputs.prompt) == {"actions": {"instruction": "Pick up the cube."}}
+    else:
+        assert inputs.prompt == "Pick up the cube."
+
+
 @pytest.mark.parametrize(
     ("prompt", "sampling_params", "message"),
     [
@@ -657,6 +755,7 @@ def _make_od_config(
         custom_pipeline_args={},
         model_config=model_config or {},
         tf_model_config=tf_model_config,
+        parallel_config=SimpleNamespace(cfg_parallel_size=1, ulysses_degree=1),
     )
 
 
@@ -692,6 +791,7 @@ def test_pipeline_init_uses_flow_unipc_with_cosmos3_defaults(stub_real_pipeline_
     assert pipeline._engine_init_flow_shift == 2.5
 
 
+@pytest.mark.parametrize("cfg_parallel_size,ulysses_degree", [(1, 1), (1, 2), (2, 1), (2, 2)])
 @pytest.mark.parametrize(
     ("scheduler_class_name", "expected_distilled"),
     [
@@ -705,6 +805,8 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file(
     monkeypatch: pytest.MonkeyPatch,
     scheduler_class_name: str,
     expected_distilled: bool,
+    cfg_parallel_size: int,
+    ulysses_degree: int,
 ) -> None:
     import json
 
@@ -717,7 +819,7 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file(
     t_list = [1.0, 0.75, 0.5, 0.25]
     scheduler_dir = tmp_path / "scheduler"
     scheduler_dir.mkdir()
-    scheduler_config = {"_class_name": scheduler_class_name}
+    scheduler_config: dict[str, Any] = {"_class_name": scheduler_class_name}
     if expected_distilled:
         scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list}
     (scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config))
@@ -753,6 +855,20 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file(
 
     od_config = _make_od_config(sound_gen=False)
     od_config.model = str(tmp_path)
+    od_config.parallel_config.cfg_parallel_size = cfg_parallel_size
+    od_config.parallel_config.ulysses_degree = ulysses_degree
+    if expected_distilled and cfg_parallel_size > 1:
+        monkeypatch.setattr(
+            pipeline_cosmos3.AutoTokenizer,
+            "from_pretrained",
+            lambda *args, **kwargs: pytest.fail("component loading must not start for distilled CFG parallelism"),
+        )
+        with pytest.raises(ValueError, match="Set --cfg-parallel-size 1 and use --ulysses-degree"):
+            Cosmos3OmniDiffusersPipeline(od_config=od_config)
+        assert StubFlowMatchScheduler.from_config_calls == []
+        assert StubFlowUniPCScheduler.from_config_calls == []
+        return
+
     pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config)
 
     assert pipeline.is_distilled_model is expected_distilled
@@ -1420,7 +1536,7 @@ def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.Monke
         fused_add_rms_norm: list[str]
         custom_op: list[str]
 
-    fake_kernel = types.ModuleType("vllm.config.kernel")
+    fake_kernel: Any = types.ModuleType("vllm.config.kernel")
     fake_kernel.IrOpPriorityConfig = FakeIrOpPriorityConfig
     monkeypatch.setitem(sys.modules, fake_kernel.__name__, fake_kernel)
 
@@ -2043,6 +2159,47 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
     torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
 
 
+def test_diffuse_publishes_exact_seacache_metadata_and_cfg_contexts(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.scheduler.sigmas = torch.tensor([0.91, 0.17])
+    observations: list[tuple[str, int | None, float | None, int | None]] = []
+
+    class RecordingHook:
+        @contextmanager
+        def cache_context(self, name: str):
+            sigma = pipeline.current_sigma
+            observations.append(
+                (
+                    name,
+                    pipeline.current_step_index,
+                    None if sigma is None else float(sigma),
+                    pipeline.num_timesteps,
+                )
+            )
+            yield
+
+    pipeline._cache_context_factory = RecordingHook().cache_context
+    pipeline.diffuse(
+        latents=torch.zeros(1, 2, 1, 1, 1),
+        timesteps=torch.tensor([900, 100]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
+    )
+
+    assert observations == [
+        ("cond", 0, pytest.approx(0.91), 2),
+        ("uncond", 0, pytest.approx(0.91), 2),
+        ("cond", 1, pytest.approx(0.17), 2),
+        ("uncond", 1, pytest.approx(0.17), 2),
+    ]
+    assert pipeline.current_step_index is None
+    assert pipeline.current_sigma is None
+
+
 def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     pipeline._use_session_state = True
@@ -2102,6 +2259,53 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
     torch.testing.assert_close(result, torch.full_like(latents, 254.0))
 
 
+def test_diffuse_transfer_uses_named_seacache_contexts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.scheduler.sigmas = torch.tensor([0.42])
+    contexts: list[tuple[str, int | None, float | None, int | None]] = []
+
+    class RecordingHook:
+        @contextmanager
+        def cache_context(self, name: str):
+            sigma = pipeline.current_sigma
+            contexts.append(
+                (
+                    name,
+                    pipeline.current_step_index,
+                    None if sigma is None else float(sigma),
+                    pipeline.num_timesteps,
+                )
+            )
+            yield
+
+    pipeline._cache_context_factory = RecordingHook().cache_context
+    latents = torch.zeros(1, 2, 1, 1, 1)
+    velocity_mask = torch.ones(1, 1, 1, 1, 1)
+    pipeline.diffuse_transfer(
+        latents=latents,
+        timesteps=torch.tensor([7]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        control_guidance=1.5,
+        control_guidance_interval=None,
+        control_latents=[torch.zeros_like(latents)],
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
+        velocity_mask=velocity_mask,
+        condition_latents=torch.zeros_like(latents),
+    )
+
+    assert contexts == [
+        ("cond", 0, pytest.approx(0.42), 1),
+        ("cond_no_control", 0, pytest.approx(0.42), 1),
+        ("uncond", 0, pytest.approx(0.42), 1),
+    ]
+    assert pipeline.current_step_index is None
+    assert pipeline.current_sigma is None
+
+
 def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     pipeline._use_session_state = True
@@ -2472,7 +2676,7 @@ def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -
 
 class TestForwardRouting:
     def _install_forward_stubs(self, pipeline):
-        captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []}
+        captured: dict[str, Any] = {"diffuse_calls": [], "prepare_calls": []}
 
         def fake_format(
             prompt,

tests/diffusion/models/cosmos3/test_cosmos3_transformer.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
 
 from __future__ import annotations
 
@@ -513,6 +513,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
+def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.cache.teacache.extractors import extract_cosmos3_context
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    class TrackingNorm(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.calls = 0
+
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            self.calls += 1
+            return hidden_states + 5.0
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    norm = TrackingNorm()
+    model.norm_moe_gen = norm
+    captured: dict[str, torch.Tensor] = {}
+
+    def run_gen_layers(hidden_gen: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        captured["input"] = hidden_gen.detach().clone()
+        return hidden_gen + 2.0
+
+    monkeypatch.setattr(model, "_run_gen_layers", run_gen_layers)
+    forward_kwargs = {
+        "hidden_states": torch.zeros(1, 2, 1, 2, 2),
+        "timestep": torch.tensor([1.0]),
+        "text_ids": torch.tensor([[1, 2]], dtype=torch.long),
+        "text_mask": torch.ones(1, 2, dtype=torch.long),
+        "video_shape": (1, 2, 2),
+        "fps": 24.0,
+    }
+
+    full_output = model(**forward_kwargs)
+
+    assert norm.calls == 1
+    norm.calls = 0
+
+    ctx = extract_cosmos3_context(model, **forward_kwargs)
+    execution_input = ctx.hidden_states.detach().clone()
+    execution_output = ctx.run_transformer_blocks()[0]
+    residual = execution_output - execution_input
+    extracted_output = ctx.postprocess(execution_output)
+
+    assert norm.calls == 1
+    expected_residual = ((captured["input"] + 2.0) + 5.0) - captured["input"]
+    assert torch.equal(residual, expected_residual)
+    assert torch.equal(extracted_output, full_output)
+
+    norm.calls = 0
+
+    def fail_if_gen_layers_run(*args, **kwargs):
+        del args, kwargs
+        raise AssertionError("SeaCache hit unexpectedly executed GEN layers")
+
+    monkeypatch.setattr(model, "_run_gen_layers", fail_if_gen_layers_run)
+
+    cached_ctx = extract_cosmos3_context(model, **forward_kwargs)
+    cached_output = cached_ctx.postprocess(cached_ctx.hidden_states + residual)
+
+    assert norm.calls == 0
+    assert torch.equal(cached_output, full_output)
+    for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
+        assert not hasattr(model, name)
+
+
+def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    class CountingNorm(nn.Module):
+        def __init__(self, original: nn.Module) -> None:
+            super().__init__()
+            self.original = original
+            self.calls = 0
+
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            self.calls += 1
+            return self.original(hidden_states)
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    norm = CountingNorm(model.norm_moe_gen)
+    model.norm_moe_gen = norm
+
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+    )
+
+    assert norm.calls == 1
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+
+
 def test_qwen_gen_mlp_remains_gated() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
 

tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+from vllm.triton_utils import HAS_TRITON
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import _apply_rotary_emb
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.fixture(autouse=True)
+def reset_fused_rope_state():
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    fused_rope._FAILED_KEYS.clear()
+    yield
+    fused_rope._FAILED_KEYS.clear()
+
+
+def _inputs(shape: tuple[int, int, int, int]):
+    batch, sequence, _, head_dim = shape
+    torch.manual_seed(17)
+    query = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn_like(query)
+    freqs_cos = torch.randn(
+        batch,
+        sequence,
+        head_dim // 2,
+        device="cuda",
+        dtype=torch.float32,
+    )
+    freqs_sin = torch.randn_like(freqs_cos)
+    return query, key, freqs_cos, freqs_sin
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (1, 1, 1, 128),
+        (2, 257, 7, 128),
+        (2, 4608, 32, 128),
+    ],
+)
+def test_fused_qk_rope_is_bit_exact(shape):
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    query, key, freqs_cos, freqs_sin = _inputs(shape)
+    with torch.inference_mode():
+        expected_query = _apply_rotary_emb(query, freqs_cos, freqs_sin)
+        expected_key = _apply_rotary_emb(key, freqs_cos, freqs_sin)
+        actual = fused_rope.try_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin)
+
+    assert actual is not None
+    actual_query, actual_key = actual
+    assert torch.equal(actual_query, expected_query)
+    assert torch.equal(actual_key, expected_key)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_launch_failure_disables_runtime_key(monkeypatch):
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    query, key, freqs_cos, freqs_sin = _inputs((1, 17, 3, 128))
+    launch_count = 0
+
+    def failing_launch(*args):
+        nonlocal launch_count
+        launch_count += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(
+        fused_rope,
+        "_launch_fused_qk_rotary_emb",
+        failing_launch,
+    )
+    with torch.inference_mode():
+        for _ in range(2):
+            assert fused_rope.try_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin) is None
+    assert launch_count == 1
+    assert len(fused_rope._FAILED_KEYS) == 1
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+def test_failed_runtime_key_cache_is_bounded():
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    for index in range(fused_rope._FAILED_KEYS_MAX_SIZE + 1):
+        fused_rope._record_failed_runtime_key((index,))
+
+    assert len(fused_rope._FAILED_KEYS) == fused_rope._FAILED_KEYS_MAX_SIZE
+    assert (0,) not in fused_rope._FAILED_KEYS
+    assert (fused_rope._FAILED_KEYS_MAX_SIZE,) in fused_rope._FAILED_KEYS
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_compile_path_does_not_launch_fused_kernel(monkeypatch):
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    query, key, freqs_cos, freqs_sin = _inputs((1, 17, 3, 128))
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+
+    def unexpected_launch(*args):
+        raise AssertionError("compile path must remain native")
+
+    monkeypatch.setattr(
+        fused_rope,
+        "_launch_fused_qk_rotary_emb",
+        unexpected_launch,
+    )
+    with torch.inference_mode():
+        assert fused_rope.try_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin) is None
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize("requires_grad_input", ["query", "key", "cos", "sin"])
+def test_gradient_inputs_do_not_launch_fused_kernel(monkeypatch, requires_grad_input):
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    values = list(_inputs((1, 17, 3, 128)))
+    input_index = {"query": 0, "key": 1, "cos": 2, "sin": 3}[requires_grad_input]
+    values[input_index].requires_grad_(True)
+
+    def unexpected_launch(*args):
+        raise AssertionError("autograd inputs must remain on the native path")
+
+    monkeypatch.setattr(
+        fused_rope,
+        "_launch_fused_qk_rotary_emb",
+        unexpected_launch,
+    )
+    assert (
+        fused_rope.try_fused_qk_rotary_emb(
+            *values,
+        )
+        is None
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="NVIDIA CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_mismatched_key_dtype_or_device_does_not_launch(monkeypatch):
+    from vllm_omni.diffusion.models.ernie_image import fused_rope
+
+    query, key, freqs_cos, freqs_sin = _inputs((1, 17, 3, 128))
+
+    def unexpected_launch(*args):
+        raise AssertionError("mismatched Q/K inputs must remain native")
+
+    monkeypatch.setattr(
+        fused_rope,
+        "_launch_fused_qk_rotary_emb",
+        unexpected_launch,
+    )
+    for unsupported_key in (key.float(), key.cpu()):
+        assert (
+            fused_rope.try_fused_qk_rotary_emb(
+                query,
+                unsupported_key,
+                freqs_cos,
+                freqs_sin,
+            )
+            is None
+        )

tests/diffusion/models/ernie_image/test_ernie_image_rope.py

@@ -1,8 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 import torch
 
 from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import (
     ErnieImageEmbedND3,
+    _apply_qk_rotary_emb,
     _apply_rotary_emb,
     rope,
 )
@@ -185,3 +189,30 @@ class TestErnieImageRopePositionEmbedding:
         out_diffusers = x * cos_ + rotated * sin_
 
         assert torch.allclose(out_vllm, out_diffusers, atol=1e-6), "vllm-omni and diffusers outputs should match"
+
+    def test_apply_qk_rotary_emb_uses_native_fallback(self, monkeypatch):
+        """Unsupported routes preserve both existing eager calls exactly."""
+        from vllm_omni.diffusion.models.ernie_image import (
+            ernie_image_transformer,
+        )
+
+        B, S, H, D = 2, 17, 3, 128
+        query = torch.randn(B, S, H, D)
+        key = torch.randn(B, S, H, D)
+        freqs_cos = torch.randn(B, S, D // 2)
+        freqs_sin = torch.randn(B, S, D // 2)
+        monkeypatch.setattr(
+            ernie_image_transformer,
+            "try_fused_qk_rotary_emb",
+            lambda *args: None,
+        )
+
+        actual_query, actual_key = _apply_qk_rotary_emb(query, key, freqs_cos, freqs_sin)
+        assert torch.equal(
+            actual_query,
+            _apply_rotary_emb(query, freqs_cos, freqs_sin),
+        )
+        assert torch.equal(
+            actual_key,
+            _apply_rotary_emb(key, freqs_cos, freqs_sin),
+        )

tests/diffusion/models/flux2/test_flux2_transformer_tp.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 import torch
 from pytest_mock import MockerFixture
@@ -265,7 +268,7 @@ class TestFlux2TransformerLayerwiseOffload:
         ]
 
     def test_get_blocks_from_dit(self):
-        from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend
+        from vllm_omni.diffusion.offloader.block_discovery import get_blocks_from_dit
 
         # Block discovery only needs the ModuleList structure, not production dims.
         model = Flux2Transformer2DModel(
@@ -275,6 +278,6 @@ class TestFlux2TransformerLayerwiseOffload:
             attention_head_dim=4,
             joint_attention_dim=16,
         )
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == ["transformer_blocks", "single_transformer_blocks"]
         assert len(blocks) == 4

tests/diffusion/models/glm_image/test_glm_image_quantization.py

@@ -399,15 +399,14 @@ class TestGlmImagePrepareModule:
         prepare = GlmImagePrepare(
             image_projector=projector,
             rope=rope,
-            patch_size=2,
         )
 
         # Test forward pass
         hidden_states = torch.randn(1, 16, 64, 64)  # [B, C, H, W]
         result = prepare(hidden_states)
 
-        assert len(result) == 5
-        hidden_out, rope_cos, rope_sin, height, width = result
+        assert len(result) == 3
+        hidden_out, rope_cos, rope_sin = result
         assert hidden_out.shape[0] == 1  # batch size
         assert hidden_out.shape[1] == 1024  # seq_len (32 * 32)
         assert hidden_out.shape[2] == 2560  # hidden_dim

tests/diffusion/models/helios/test_helios_usp_2rank.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Two-process USP smoke test for the Helios transformer.
+
+Pins the multi-rank mechanism that ``test_helios_usp_shard.py`` (unit tests
+of ``_sp_split_seq`` alone) cannot see:
+
+  * the manual ``_sp_shard_depth`` bump in ``forward()`` must make attn1 run
+    the Ulysses all-to-all (not ``NoParallelAttention``);
+  * attn2 cross-attention must stay rank-local (``skip_sequence_parallel=
+    True``): if replicated text K/V were fed through Ulysses they would be
+    ws-fold duplicated and outputs would diverge from the single-rank run;
+  * end-to-end correctness: ws=2 output equals the ws=1 baseline for the same
+    weights/inputs (divisible shapes, so the padding path is not involved).
+
+Runs ``HeliosTransformer3DModel.forward`` with tiny random weights via
+``torch.multiprocessing.spawn``: first a baseline (world_size=1), then an
+SP run (world_size=2, ulysses_degree=2); rank-0 outputs are compared.
+
+Requires >= 2 devices; skipped otherwise.
+"""
+
+import os
+import tempfile
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelConfig,
+    get_sp_plan_from_model,
+)
+from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
+from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
+from vllm_omni.platforms import current_omni_platform
+
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, *_L4_TWO_GPU]
+
+RANDOM_SEED = 42
+DTYPE = torch.float32  # keep both runs on the SDPA fallback for tight tolerance
+
+# Tiny but architecturally faithful config:
+#   inner_dim = 4 * 32 = 128, rope_dim sums to head_dim (32),
+#   current chunk 2x4x6 latent -> (1,2,2) patch -> 2*2*3 = 12 tokens,
+#   short history 1x4x6 latent -> (1,2,2) patch -> 6 tokens.
+# All component lengths are even, so ws=2 needs no padding replicas.
+TINY_MODEL_KWARGS = dict(
+    patch_size=(1, 2, 2),
+    num_attention_heads=4,
+    attention_head_dim=32,
+    in_channels=4,
+    out_channels=4,
+    text_dim=64,
+    freq_dim=32,
+    ffn_dim=128,
+    num_layers=2,
+    rope_dim=(12, 10, 10),
+    rope_theta=10000.0,
+    guidance_cross_attn=True,
+    zero_history_timestep=True,
+    has_multi_term_memory_patch=True,
+    is_amplify_history=False,
+)
+
+BATCH = 1
+CUR_FRAMES, LAT_H, LAT_W = 2, 4, 6  # -> 12 current tokens
+HIST_FRAMES = 1  # short history -> 6 tokens
+TEXT_LEN = 5
+
+
+def _update_env(envs_dict: dict[str, str]) -> None:
+    for k, v in envs_dict.items():
+        os.environ[k] = v
+
+
+def _make_inputs(device: torch.device) -> dict[str, torch.Tensor]:
+    torch.manual_seed(RANDOM_SEED + 1)
+    return {
+        "hidden_states": torch.randn(BATCH, 4, CUR_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE),
+        "timestep": torch.full((BATCH,), 500.0, device=device),
+        "encoder_hidden_states": torch.randn(
+            BATCH, TEXT_LEN, TINY_MODEL_KWARGS["text_dim"], device=device, dtype=DTYPE
+        ),
+        "indices_latents_history_short": torch.tensor([[0]], device=device),
+        "latents_history_short": torch.randn(BATCH, 4, HIST_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE),
+    }
+
+
+def _spy_strategies(model: HeliosTransformer3DModel) -> dict[str, list[str]]:
+    """Record the parallel strategy each Attention resolves to during forward."""
+    seen: dict[str, list[str]] = {"attn1": [], "attn2": []}
+    for name, attn in (("attn1", model.blocks[0].attn1.attn), ("attn2", model.blocks[0].attn2.attn)):
+        original = attn._get_active_parallel_strategy
+
+        def spy(original=original, name=name):
+            strategy = original()
+            seen[name].append(getattr(strategy, "name", type(strategy).__name__))
+            return strategy
+
+        attn._get_active_parallel_strategy = spy
+    return seen
+
+
+def helios_forward_on_model(
+    local_rank: int,
+    world_size: int,
+    ulysses_degree: int,
+    output_file: str,
+    model_state_file: str,
+    is_baseline: bool,
+) -> None:
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    torch.set_default_device(device)
+    torch.set_default_dtype(DTYPE)
+
+    _update_env(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": "12361",
+        }
+    )
+    init_distributed_environment()
+
+    parallel_config = DiffusionParallelConfig(
+        pipeline_parallel_size=1,
+        data_parallel_size=1,
+        tensor_parallel_size=1,
+        sequence_parallel_size=ulysses_degree,
+        ulysses_degree=ulysses_degree,
+        ring_degree=1,
+        allgather_degree=1,
+        cfg_parallel_size=1,
+    )
+    od_config = OmniDiffusionConfig.from_kwargs(
+        model="test_model",
+        dtype=DTYPE,
+        parallel_config=parallel_config,
+        diffusion_attention_backend="TORCH_SDPA",
+    )
+    initialize_model_parallel(
+        data_parallel_size=1,
+        cfg_parallel_size=1,
+        sequence_parallel_size=ulysses_degree,
+        ulysses_degree=ulysses_degree,
+        ring_degree=1,
+        allgather_degree=1,
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+    )
+
+    with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
+        torch.manual_seed(RANDOM_SEED)
+        model = HeliosTransformer3DModel(**TINY_MODEL_KWARGS).to(device).to(DTYPE)
+
+        if is_baseline and local_rank == 0:
+            with open(model_state_file, "wb") as f:
+                torch.save({k: v.cpu() for k, v in model.state_dict().items()}, f)
+
+        if world_size > 1:
+            torch.distributed.barrier()
+
+        with open(model_state_file, "rb") as f:
+            model_state = torch.load(f)
+        model.load_state_dict({k: v.to(device).to(DTYPE) for k, v in model_state.items()})
+
+        if not is_baseline:
+            # Mirror registry._apply_sequence_parallel_if_enabled: install the
+            # proj_out gather hook from the model's own _sp_plan and flip
+            # sp_plan_hooks_applied so sp_active follows _sp_shard_depth
+            # (i.e. the manual bump in Helios forward is what enables Ulysses).
+            plan = get_sp_plan_from_model(model)
+            assert plan is not None, "HeliosTransformer3DModel must define _sp_plan"
+            apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=ulysses_degree), plan)
+            get_forward_context().sp_plan_hooks_applied = True
+
+        seen = _spy_strategies(model)
+        inputs = _make_inputs(device)
+        output = model(**inputs).sample
+
+        assert output.shape == (BATCH, 4, CUR_FRAMES, LAT_H, LAT_W), f"unexpected output shape {output.shape}"
+
+        if is_baseline:
+            assert seen["attn1"] and all(n == "none" for n in seen["attn1"]), (
+                f"baseline attn1 must use NoParallelAttention ('none'), saw {seen['attn1']}"
+            )
+        else:
+            # attn1: SP communication must be active inside the block loop
+            # (regression here = the manual _sp_shard_depth bump was lost).
+            assert seen["attn1"] and all(n == "ulysses" for n in seen["attn1"]), (
+                f"attn1 must run Ulysses under USP (manual _sp_shard_depth bump), saw {seen['attn1']}"
+            )
+            # attn2: replicated text K/V must stay rank-local
+            # (regression here = the ws-fold duplicated-keys bug).
+            assert seen["attn2"] and all(n == "none" for n in seen["attn2"]), (
+                f"attn2 must skip sequence parallel (replicated text K/V), saw {seen['attn2']}"
+            )
+
+        if local_rank == 0:
+            with open(output_file, "wb") as f:
+                torch.save(output.detach().cpu().float(), f)
+
+        destroy_distributed_env()
+
+
+def test_helios_usp2_matches_single_rank():
+    """ws=2 (Ulysses) must reproduce the ws=1 output for identical weights/inputs."""
+    available_devices = current_omni_platform.get_device_count()
+    if available_devices < 2:
+        pytest.skip(f"Test requires 2 devices but only {available_devices} available")
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
+        baseline_output_file = f.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
+        sp_output_file = f.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
+        model_state_file = f.name
+
+    try:
+        torch.multiprocessing.spawn(
+            helios_forward_on_model,
+            args=(1, 1, baseline_output_file, model_state_file, True),
+            nprocs=1,
+        )
+        torch.multiprocessing.spawn(
+            helios_forward_on_model,
+            args=(2, 2, sp_output_file, model_state_file, False),
+            nprocs=2,
+        )
+
+        with open(baseline_output_file, "rb") as f:
+            baseline = torch.load(f)
+        with open(sp_output_file, "rb") as f:
+            sp = torch.load(f)
+
+        assert baseline.shape == sp.shape
+        max_abs_diff = (baseline - sp).abs().max().item()
+        assert torch.allclose(baseline, sp, atol=1e-3, rtol=1e-3), (
+            f"USP=2 output diverges from single-rank baseline: max_abs_diff={max_abs_diff:.6e}"
+        )
+    finally:
+        for path in [baseline_output_file, sp_output_file, model_state_file]:
+            if os.path.exists(path):
+                os.remove(path)

tests/diffusion/models/helios/test_helios_usp_shard.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for Helios per-component sequence splitting.
+
+These tests verify the manual sharding logic that replaces the broken
+rope/blocks.0 hooks in the Helios ``_sp_plan``.  They exercise the *real*
+``HeliosTransformer3DModel._sp_split_seq`` (the parallel-state getters it
+imports lazily inside the method body are monkeypatched), so the model-side
+helper cannot drift silently from what is tested.  For the multi-rank
+mechanism around the helper (manual ``_sp_shard_depth`` bump, attn1/attn2
+parallel strategy, ws=2 vs ws=1 output equivalence) see
+``test_helios_usp_2rank.py``.
+
+Coverage:
+  1. USP disabled (ws == 1) returns the tensor unchanged.
+  2. Divisible split gives each rank the expected slice.
+  3. Non-divisible split pads with replicated tokens (no crash).
+  4. Per-component split: history + current tokens are distributed so
+     that *both* ranks receive some current tokens (no mosaic regression).
+  5. Edge cases: 1-D tensor, empty seq, 4-D tensor.
+"""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.distributed.parallel_state as sp_parallel_state
+from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@contextmanager
+def _fake_sp_state(ws: int, rank: int) -> Iterator[None]:
+    """Patch the parallel-state getters used by the real helper.
+
+    ``_sp_split_seq`` imports these names from the module *inside* its
+    function body, so replacing the module attributes takes effect.
+    """
+    prev_ws = sp_parallel_state.get_sequence_parallel_world_size
+    prev_rank = sp_parallel_state.get_sequence_parallel_rank
+    sp_parallel_state.get_sequence_parallel_world_size = lambda: ws
+    sp_parallel_state.get_sequence_parallel_rank = lambda: rank
+    try:
+        yield
+    finally:
+        sp_parallel_state.get_sequence_parallel_world_size = prev_ws
+        sp_parallel_state.get_sequence_parallel_rank = prev_rank
+
+
+def _sp_split_seq(x: torch.Tensor, ws: int, rank: int) -> torch.Tensor:
+    """Call the real model-side helper without constructing a model."""
+    with _fake_sp_state(ws, rank):
+        return HeliosTransformer3DModel._sp_split_seq(None, x)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestSpSplitSeqDisabled:
+    """When USP is off (ws == 1) the helper must be a no-op."""
+
+    def test_ws1_returns_unchanged(self):
+        x = torch.randn(2, 100, 64)
+        out = _sp_split_seq(x, ws=1, rank=0)
+        assert torch.equal(out, x)
+        assert out is x
+
+
+class TestSpSplitSeqDivisible:
+    """Divisible sequences: each rank gets exactly seq_len / ws tokens."""
+
+    def test_each_rank_gets_correct_slice(self):
+        for ws in [2, 4, 8]:
+            seq_len = ws * 10
+            x = torch.arange(seq_len * 4).float().reshape(1, seq_len, 4)
+            for rank in range(ws):
+                out = _sp_split_seq(x, ws=ws, rank=rank)
+                n = seq_len // ws
+                expected = x[:, rank * n : (rank + 1) * n, :]
+                assert out.shape == (1, n, 4)
+                assert torch.equal(out, expected)
+
+    def test_output_is_contiguous(self):
+        x = torch.randn(1, 20, 8)
+        out = _sp_split_seq(x, ws=2, rank=0)
+        assert out.is_contiguous()
+
+
+class TestSpSplitSeqNonDivisible:
+    """Non-divisible sequences are padded with replicated tokens."""
+
+    def test_non_divisible_pads_instead_of_crash(self):
+        """400x272 → 9*25*17 = 3825 tokens, USP=2 → not divisible.
+        Should pad to 3826 and split, not crash."""
+        x = torch.randn(1, 3825, 8)
+        for rank in range(2):
+            out = _sp_split_seq(x, ws=2, rank=rank)
+            # 3825 → pad 1 → 3826 → /2 = 1913
+            assert out.shape[1] == 1913
+
+    def test_non_divisible_pads_with_replicated_last_token(self):
+        """Verify the padding is the replicated last token, not zeros."""
+        x = torch.randn(1, 7, 4)  # 7 % 2 = 1, pad 1
+        last_token = x[:, -1:, :].clone()
+        out = _sp_split_seq(x, ws=2, rank=0)
+        # rank 0 gets tokens [0:4] (first 4 of the padded 8)
+        assert out.shape == (1, 4, 4)
+        # rank 1 gets tokens [4:8] — token 7 (the padded one) = last token
+        out1 = _sp_split_seq(x, ws=2, rank=1)
+        assert out1.shape == (1, 4, 4)
+        assert torch.equal(out1[:, -1:, :], last_token)
+
+    def test_non_divisible_7_4(self):
+        """7 tokens, ws=4 → pad 1 → 8 → /4 = 2 per rank."""
+        x = torch.randn(1, 7, 8)
+        for rank in range(4):
+            out = _sp_split_seq(x, ws=4, rank=rank)
+            assert out.shape[1] == 2
+
+    def test_non_divisible_15_4(self):
+        """15 tokens, ws=4 → pad 1 → 16 → /4 = 4 per rank."""
+        x = torch.randn(1, 15, 8)
+        for rank in range(4):
+            out = _sp_split_seq(x, ws=4, rank=rank)
+            assert out.shape[1] == 4
+
+
+class TestSpSplitSeqPerComponent:
+    """Simulate Helios forward: split current + history components so that
+    *both* ranks receive some current tokens (no mosaic regression)."""
+
+    def test_both_ranks_get_current_tokens(self):
+        """history=2400, current=540, ws=2.
+        After per-component split, each rank must have > 0 current tokens."""
+        current = torch.randn(1, 540, 64)
+        history = torch.randn(1, 2400, 64)
+
+        for rank in range(2):
+            cur_local = _sp_split_seq(current, ws=2, rank=rank)
+            hist_local = _sp_split_seq(history, ws=2, rank=rank)
+
+            assert cur_local.shape[1] > 0, f"Rank {rank} received 0 current tokens"
+            assert cur_local.shape[1] == 270  # 540 / 2
+            assert hist_local.shape[1] == 1200  # 2400 / 2
+
+            original_context_length = cur_local.shape[1]
+            assert original_context_length > 0  # no mosaic regression
+
+    def test_gather_restores_full_current(self):
+        """After splitting current into 2 ranks, simulating proj_out
+        gather must restore the full current sequence length so that
+        the original unpatchify reshape succeeds."""
+        current = torch.randn(1, 540, 64)
+        local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)]
+        gathered = torch.cat(local_chunks, dim=1)
+        assert gathered.shape[1] == 540  # full current restored
+
+    def test_gather_with_padding_truncates_correctly(self):
+        """Non-divisible current (3825) → pad to 3826 → split →
+        gather → 3826. Slice to 3825 for unpatchify."""
+        current = torch.randn(1, 3825, 64)
+        local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)]
+        gathered = torch.cat(local_chunks, dim=1)
+        assert gathered.shape[1] == 3826  # padded size
+        # Simulate the slice before unpatchify
+        expected_seq = 3825
+        if gathered.shape[1] > expected_seq:
+            gathered = gathered[:, :expected_seq, :]
+        assert gathered.shape[1] == 3825  # correct after slice
+
+    def test_rotary_emb_aligned_with_hidden_states(self):
+        """rotary_emb and hidden_states must be split at the same
+        positions so that tokens align inside the attention layer."""
+        hidden = torch.randn(1, 540, 5120)
+        rotary = torch.randn(1, 540, 256)
+        for rank in range(2):
+            h_local = _sp_split_seq(hidden, ws=2, rank=rank)
+            r_local = _sp_split_seq(rotary, ws=2, rank=rank)
+            assert h_local.shape[1] == r_local.shape[1]
+
+    def test_whole_split_causes_mosaic_but_per_component_does_not(self):
+        """Demonstrates why per-component split is needed:
+        with whole split (simulated), rank 0 gets 0 current tokens;
+        with per-component split, both ranks get current tokens."""
+        history = torch.randn(1, 2400, 64)
+        current = torch.randn(1, 540, 64)
+        full = torch.cat([history, current], dim=1)  # 2940
+
+        # Whole split (the OLD broken behavior)
+        whole_n = full.shape[1] // 2  # 1470
+        rank0_whole = full[:, :whole_n, :]
+        rank0_current = rank0_whole.shape[1] - 2400  # negative → 0 current
+        assert rank0_current <= 0, "Whole split should give rank 0 no current"
+
+        # Per-component split (the NEW correct behavior)
+        for rank in range(2):
+            cur_local = _sp_split_seq(current, ws=2, rank=rank)
+            assert cur_local.shape[1] == 270  # both ranks have current
+
+
+class TestSpSplitSeqEdgeCases:
+    """Edge cases that should not crash or silently misbehave."""
+
+    def test_1d_tensor_skipped(self):
+        """1-D tensors must be returned unchanged (dim < 2 guard)."""
+        x = torch.randn(100)
+        out = _sp_split_seq(x, ws=2, rank=0)
+        assert torch.equal(out, x)
+
+    def test_empty_seq_skipped(self):
+        """Empty sequence dimension must be returned unchanged."""
+        x = torch.randn(1, 0, 64)
+        out = _sp_split_seq(x, ws=2, rank=0)
+        assert out.shape == x.shape
+
+    def test_4d_tensor_split_on_dim1(self):
+        """4-D tensor [B, S, H, D] should split along dim=1."""
+        x = torch.randn(1, 20, 8, 32)
+        out = _sp_split_seq(x, ws=2, rank=1)
+        assert out.shape == (1, 10, 8, 32)

tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.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
 
 from types import SimpleNamespace
 
@@ -362,6 +362,20 @@ def test_prepared_model_inputs_match_local_tokenization(
     local_inputs = pipeline.prepare_model_inputs(**common_kwargs)
     prepared_inputs = pipeline.prepare_model_inputs(**common_kwargs, prepared_layout=prepared_layout)
 
+    if with_reference_image:
+        # Preparation runs without a forward context. Coverage must come from
+        # the scheduler metadata copied onto this request, including every CFG row.
+        reused = pipeline.prepare_model_inputs(
+            **common_kwargs, prepared_layout=prepared_layout, kv_computed_tokens=(5,) * len(prefix_lens)
+        )
+        assert reused["cond_vae_images"] is None and reused["cond_vit_images"] is None
+        partial = pipeline.prepare_model_inputs(
+            **common_kwargs,
+            prepared_layout=prepared_layout,
+            kv_computed_tokens=(5,) * (len(prefix_lens) - 1) + (4,),
+        )
+        assert partial["cond_vae_images"] is not None and partial["cond_vit_images"] is not None
+
     comparable_fields = (
         "input_ids",
         "position_ids",
@@ -474,3 +488,33 @@ def test_paged_preprocess_attaches_layout_and_scheduler_kv_requests(monkeypatch)
         (12, 17, 32)
     ]
     assert len(tokenizer.calls) == 1
+
+
+@pytest.mark.parametrize("transfer_tokens, expected", [(9, [7, 5]), (4, [4, 4])])
+def test_cfg_rows_preserve_reusable_token_ids_and_native_transfer_boundary(transfer_tokens, expected):
+    from vllm_omni.diffusion.diffusion_kv.kv_connector import prepare_kv_requests
+
+    tokenizer, image_processor = _components([12, 14])
+    request = _request(guidance_scale=5.0)
+    layout = _prepare(request, tokenizer, image_processor)
+    layout.tokenizer_output.think_recaption_end_pos = [[7], [5]]
+    layout.tokenizer_output.uncond_cfg_start_pos = [[None], [5]]
+    rows = build_hunyuan_diffusion_kv_requests(request, layout)
+    assert [row.prompt_token_ids for row in rows] == [
+        layout.tokenizer_output.tokens[0, :7].tolist(),
+        layout.tokenizer_output.tokens[1, :5].tolist(),
+    ]
+    prepare_kv_requests(rows, {"num_transfer_tokens": transfer_tokens})
+    assert [row.num_prompt_tokens for row in rows] == expected
+    assert all(row.kv_transfer_params["num_transfer_tokens"] == transfer_tokens for row in rows)
+    layout.tokenizer_output.uncond_cfg_start_pos = [[5]]
+    with pytest.raises(ValueError, match="one boundary per CFG row"):
+        build_hunyuan_diffusion_kv_requests(request, layout)
+
+
+@pytest.mark.parametrize("computed, expected", [((8, 8), True), ((8, 5), False), ((0, 0), False), ((), False)])
+def test_native_image_reuse_requires_coverage_in_every_cfg_row(computed, expected):
+    from vllm_omni.diffusion.models.hunyuan_image3.request_layout import native_kv_covers_cond_images
+
+    output = TokenizerEncodeOutput(joint_image_slices=[[slice(2, 8)], [slice(2, 6)]])
+    assert native_kv_covers_cond_images(output, computed) is expected

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.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 HunyuanImage3 AR sampler logic (stage transitions,
 ratio restriction, comprehension blocking)."""
 
@@ -21,6 +21,10 @@ RATIO_END = 210
 RATIO_OTHER_START = 220
 RATIO_OTHER_END = 223
 TIMESTEP = 224
+CFG = 110
+IMG_TOK = 111
+JOINT_SEP = 112
+EOI = 113
 
 
 class FakeSamplerModel:
@@ -47,10 +51,13 @@ class FakeSamplerModel:
             self._stage_transitions[END_OF_THINK] = [RECAPTION]
             self._stage_transitions[END_OF_RECAPTION] = [ANSWER, BOI, SIZE_TOKEN]
 
-        self._blocked_token_ids: set[int] = set()
-        if is_comprehension:
-            self._blocked_token_ids.update([BOI, SIZE_TOKEN])
-            self._blocked_token_ids.update(self._all_ratio_ids)
+        # Mirrors the model: the blocklist is always built so generation-mode
+        # engines can apply it to per-request comprehension rows (#6088).
+        self._cfg_token_id = CFG
+        self._blocked_token_ids: set[int] = {BOI, EOI, SIZE_TOKEN, CFG, IMG_TOK, TIMESTEP, JOINT_SEP}
+        self._blocked_token_ids.update(self._all_ratio_ids)
+        self._blocked_token_ids_tensor = None
+        self._sampler = None
 
     # Bind the real methods from the model class.
     from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import (
@@ -59,6 +66,8 @@ class FakeSamplerModel:
 
     _get_forced_token = _Real._get_forced_token
     _apply_ratio_restriction = _Real._apply_ratio_restriction
+    _row_is_comprehension = staticmethod(_Real._row_is_comprehension)
+    sample = _Real.sample
 
 
 class FakeTokenEmbedding:
@@ -250,3 +259,114 @@ class TestForceEosAfterRatio:
         assert logits[0, EOS].item() == 0
         non_eos_max = logits[0, :EOS].max().item()
         assert non_eos_max == min_score
+
+
+class _RecordingSampler:
+    """Stub sampler: returns the (mutated) logits so tests can assert on them."""
+
+    def __call__(self, *, logits, sampling_metadata):
+        return logits
+
+
+def _metadata(output_token_ids):
+    from types import SimpleNamespace
+
+    return SimpleNamespace(output_token_ids=output_token_ids)
+
+
+class TestPerRequestTaskMode:
+    """Generation-mode engines must honor per-request comprehension rows (#6088).
+
+    A text chat (I2T/T2T) served by an AR+DiT deployment is marked with
+    extra_args={"ar_task_mode": "comprehension"}; its rows get the stateless
+    image-token blocking instead of stage-transition forcing, while untagged
+    rows keep the generation behavior.
+    """
+
+    VOCAB = 300
+
+    def _run(self, model, histories, per_req_extra_args):
+        logits = torch.zeros(len(histories), self.VOCAB)
+        logits[:, 50] = 7.0  # ordinary text token, must survive comprehension
+        model._sampler = _RecordingSampler()
+        out = model.sample(logits, _metadata(histories), per_req_extra_args=per_req_extra_args)
+        return out
+
+    def test_comprehension_row_skips_forcing_and_blocks_image_tokens(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        min_score = torch.finfo(torch.float32).min
+
+        logits = self._run(
+            model,
+            histories=[[END_OF_THINK], [END_OF_THINK]],
+            per_req_extra_args=[{"ar_task_mode": "comprehension"}, None],
+        )
+
+        # Row 0 (comprehension): no </think> -> <recaption> forcing (ordinary
+        # text token survives, which a force-fill would have erased) and every
+        # DiT scaffold token is masked.
+        assert logits[0, 50].item() == 7.0
+        for tid in (BOI, EOI, SIZE_TOKEN, CFG, IMG_TOK, TIMESTEP, JOINT_SEP, RATIO_START, RATIO_OTHER_START):
+            assert logits[0, tid].item() == min_score, f"token {tid} not blocked"
+
+        # Row 1 (generation): forced to <recaption> exactly as before.
+        assert logits[1, RECAPTION].item() == 0
+        assert logits[1, 50].item() == min_score
+
+    def test_no_extra_args_keeps_generation_behavior(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        min_score = torch.finfo(torch.float32).min
+
+        logits = self._run(model, histories=[[END_OF_THINK]], per_req_extra_args=None)
+
+        assert logits[0, RECAPTION].item() == 0
+        assert logits[0, 50].item() == min_score
+
+    def test_empty_or_foreign_extra_args_keeps_generation_behavior(self):
+        model = FakeSamplerModel(is_comprehension=False)
+
+        logits = self._run(
+            model,
+            histories=[[END_OF_THINK], [END_OF_THINK]],
+            per_req_extra_args=[{}, {"other_key": 1}],
+        )
+
+        assert logits[0, RECAPTION].item() == 0
+        assert logits[1, RECAPTION].item() == 0
+
+    def test_comprehension_row_without_trigger_history(self):
+        model = FakeSamplerModel(is_comprehension=False)
+        min_score = torch.finfo(torch.float32).min
+
+        logits = self._run(
+            model,
+            histories=[[1, 2, 3]],
+            per_req_extra_args=[{"ar_task_mode": "comprehension"}],
+        )
+
+        assert logits[0, 50].item() == 7.0
+        assert logits[0, CFG].item() == min_score
+
+    def test_static_comprehension_engine_blocks_cfg(self):
+        model = FakeSamplerModel(is_comprehension=True)
+        min_score = torch.finfo(torch.float32).min
+
+        logits = self._run(model, histories=[[1, 2]], per_req_extra_args=None)
+
+        for tid in (BOI, EOI, SIZE_TOKEN, CFG, IMG_TOK, TIMESTEP, JOINT_SEP):
+            assert logits[0, tid].item() == min_score
+        assert logits[0, 50].item() == 7.0
+
+    def test_extra_args_shorter_than_batch_defaults_to_generation(self):
+        model = FakeSamplerModel(is_comprehension=False)
+
+        logits = self._run(
+            model,
+            histories=[[END_OF_THINK], [END_OF_THINK]],
+            per_req_extra_args=[{"ar_task_mode": "comprehension"}],
+        )
+        min_score = torch.finfo(torch.float32).min
+
+        assert logits[0, 50].item() == 7.0  # tagged row: comprehension
+        assert logits[1, RECAPTION].item() == 0  # missing entry: generation
+        assert logits[1, 50].item() == min_score

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import math
 from contextlib import ExitStack, contextmanager
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
@@ -143,6 +144,104 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
     assert mgr.state_dict() == {}
 
 
+def test_cache_manager_initializes_with_quantization_skip_layers(monkeypatch):
+    import vllm_omni.diffusion.attention.layer as layer
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ImageKVCacheManager
+
+    config = SimpleNamespace(
+        diffusion_attention_config=None,
+        parallel_config=SimpleNamespace(ring_degree=1),
+        diffusion_kv_cache_dtype="auto",
+        diffusion_kv_cache_skip_layer_indices={3},
+    )
+    monkeypatch.setattr(layer, "get_attn_backend_for_role", lambda **_: (layer.SDPABackend, None))
+    monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **_: object())
+    with patched_mgr_env(), set_current_diffusion_config(config):
+        # Keep the real Attention constructor and its skip-layer validation.
+        with patch(f"{_TRANSFORMER_MODULE}.Attention", layer.Attention):
+            mgr = ImageKVCacheManager(
+                num_heads=NUM_HEADS,
+                num_kv_heads=NUM_KV_HEADS,
+                head_dim=HEAD_DIM,
+                scaling=SCALING,
+                prefix="model.layers.3.self_attn",
+            )
+    assert mgr.attn.layer_idx == 3
+    assert not mgr.attn._should_apply_kv_cache_quant()
+
+
+@pytest.mark.parametrize("first_step", [True, False])
+def test_native_connector_keeps_imported_prefix_in_paged_attention(first_step) -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr()
+    mgr.attn.paged_kv_active = True
+    query_len, prefix_len = IMAGE_TOKEN_LEN, 19
+    key, value = _make_known_kv(2 * query_len)
+    spans = [[(prefix_len, prefix_len + query_len)]] * 2
+    with (
+        set_forward_context(),
+        override_paged_kv_adapter(object()),
+        patch.object(mgr, "_forward_dense_legacy", side_effect=AssertionError("dense KV path")),
+        patch(f"{_TRANSFORMER_MODULE}.repeat_kv", side_effect=AssertionError("KV head expansion")),
+        patch(f"{_TRANSFORMER_MODULE}.get_paged_kv_computed_tokens", return_value=(prefix_len, prefix_len)),
+    ):
+        _call_mgr(
+            mgr,
+            2,
+            query_len,
+            prefix_len + query_len,
+            key,
+            value,
+            first_step=first_step,
+            full_attn_spans=spans,
+        )
+    assert len(mgr.attn.paged_calls) == 1
+    assert mgr.attn.paged_calls[0][1].shape == (2, query_len, NUM_KV_HEADS, HEAD_DIM)
+    assert mgr.attn.paged_metadata[0].attn_mask is None
+    assert mgr.attn.paged_metadata[0].full_attn_spans == spans
+    assert mgr.image_kv_cache_map is None
+    assert not hasattr(mgr, "_materialize_paged_prefix")
+
+
+def test_legacy_tensor_prefix_still_uses_dense_cache_and_reuse() -> None:
+    mgr = _make_cache_mgr()
+    prefix_len, prompt_len = 3, 2
+    imported = _make_known_kv(prefix_len, base=10.0)
+    mgr._injected_ar_kv = [imported]
+    first_len = prompt_len + IMAGE_TOKEN_LEN
+    key, value = _make_known_kv(first_len, base=20.0)
+    with (
+        patch.object(mgr, "_forward_paged", side_effect=AssertionError("legacy entered paged path")),
+    ):
+        _call_mgr(
+            mgr,
+            1,
+            first_len,
+            prefix_len + first_len,
+            key,
+            value,
+            first_step=True,
+            gen_timestep_scatter_index=_gen_timestep_index(1, prompt_len),
+        )
+        assert mgr.image_kv_cache_map is not None
+        assert mgr.image_kv_cache_map[0].shape[1] == prefix_len + prompt_len
+        next_key, next_value = _make_known_kv(IMAGE_TOKEN_LEN, base=30.0)
+        _call_mgr(
+            mgr,
+            1,
+            IMAGE_TOKEN_LEN,
+            prefix_len + first_len,
+            next_key,
+            next_value,
+            position_ids=torch.arange(prefix_len + prompt_len, prefix_len + first_len).unsqueeze(0),
+        )
+    assert len(mgr.attn.calls) == 2
+    assert not mgr.attn.paged_calls
+    assert mgr.attn.calls[-1][1].shape == (1, prefix_len + first_len, NUM_HEADS, HEAD_DIM)
+
+
 def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
     from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
 

tests/diffusion/models/hunyuan_image3/test_prompt_utils.py

@@ -13,9 +13,11 @@ from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
     HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS,
     available_bot_tasks,
     available_tasks,
+    build_ar_prompt_inputs,
     build_prompt,
     build_prompt_tokens,
     resolve_stop_token_ids,
+    validate_special_token_ids,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -296,19 +298,42 @@ def _repo_root() -> pathlib.Path:
     return pathlib.Path(__file__).resolve().parents[4]
 
 
-def test_end2end_routes_through_shared_prompt_utils():
-    end2end_path = _repo_root() / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
-    tree = ast.parse(end2end_path.read_text(encoding="utf-8"))
+def test_example_routes_through_shared_seam():
+    """Neither the shared task examples nor the ``model_extras`` seam itself
+    should duplicate AR prompt tokenization.
 
-    local_func_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)}
-    assert not (local_func_names & {"build_prompt", "build_prompt_tokens"})
+    After the model_extras migration, t2i / image-editing run through the
+    shared ``text_to_image.py`` / ``image_edit.py`` examples (both call
+    ``build_ar_stage_inputs``), and t2t / i2t run through the shared
+    ``x_to_text.py`` example (which calls ``build_x_to_text_prompt``). This
+    test pins that both ``model_extras.hunyuan_image3`` seam functions
+    delegate to ``prompt_utils.build_ar_prompt_inputs`` / ``build_ar_stage_inputs``
+    rather than re-implementing tokenization inline.
+    """
+    seam_path = _repo_root() / "vllm_omni" / "model_extras" / "hunyuan_image3.py"
+    seam_tree = ast.parse(seam_path.read_text(encoding="utf-8"))
 
-    imported_from_prompt_utils: set[str] = set()
-    for node in ast.walk(tree):
+    local_func_names = {n.name for n in ast.walk(seam_tree) if isinstance(n, ast.FunctionDef)}
+    assert not (local_func_names & {"build_prompt", "build_prompt_tokens", "resolve_stop_token_ids"})
+
+    # build_x_to_text_prompt must route through build_ar_stage_inputs, not
+    # re-derive the AR prefill / stop-token logic itself.
+    build_x_to_text_prompt_node = next(
+        n for n in ast.walk(seam_tree) if isinstance(n, ast.FunctionDef) and n.name == "build_x_to_text_prompt"
+    )
+    called_names = {
+        node.func.id
+        for node in ast.walk(build_x_to_text_prompt_node)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
+    }
+    assert "build_ar_stage_inputs" in called_names
+
+    # The model_extras seam itself must delegate to the shared prompt_utils.
+    routed_from_prompt_utils: set[str] = set()
+    for node in ast.walk(seam_tree):
         if isinstance(node, ast.ImportFrom) and node.module and node.module.endswith("hunyuan_image3.prompt_utils"):
-            imported_from_prompt_utils.update(alias.name for alias in node.names)
-    expected_imports = {"build_prompt_tokens", "resolve_stop_token_ids", "resolve_sys_type"}
-    assert expected_imports <= imported_from_prompt_utils
+            routed_from_prompt_utils.update(alias.name for alias in node.names)
+    assert "build_ar_prompt_inputs" in routed_from_prompt_utils
 
 
 _HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct"
@@ -331,3 +356,87 @@ def test_segment_tokenize_diverges_from_full_string_encode():
     full_ids = tok.encode(build_prompt(user_prompt, task="i2t", bot_task=None), add_special_tokens=False)
     assert seg_ids != full_ids
     assert len(seg_ids) >= len(full_ids)
+
+
+def test_build_ar_prompt_inputs_tokenizer_path_matches_build_prompt_tokens():
+    """With a tokenizer, the helper delegates to build_prompt_tokens (parity)."""
+    tok = FakeTokenizer()
+    ar = build_ar_prompt_inputs("a red panda", task="t2i", bot_task="think", tokenizer=tok)
+    expected = build_prompt_tokens("a red panda", FakeTokenizer(), task="t2i", bot_task="think")
+    assert ar.prompt is None
+    assert ar.prompt_token_ids == expected.token_ids
+    assert ar.system_prompt_type == expected.system_prompt_type
+    # think + no explicit size -> AR predicts ratio, so stop range is non-trivial.
+    assert len(ar.stop_token_ids) > 1
+
+
+def test_build_ar_prompt_inputs_string_fallback_without_tokenizer():
+    """Without a tokenizer, the helper returns the build_prompt string form."""
+    ar = build_ar_prompt_inputs("a red panda", task="t2i", bot_task="recaption")
+    assert ar.prompt_token_ids is None
+    assert ar.prompt == build_prompt("a red panda", task="t2i", bot_task="recaption")
+    assert ar.system_prompt_type
+    assert ar.stop_token_ids
+
+
+def test_build_ar_prompt_inputs_explicit_size_stops_at_terminator():
+    """Explicit image_size makes 'recaption' stop at </recaption> only."""
+    ar = build_ar_prompt_inputs("x", task="t2i", bot_task="recaption", image_size="1024x1024")
+    assert ar.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]]
+
+
+def test_build_ar_prompt_inputs_comprehension_stops_on_answer():
+    """Text-output tasks (t2t / i2t) stop on <answer>."""
+    ar = build_ar_prompt_inputs("hi", task="t2t")
+    assert ar.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]]
+
+
+class _MatchingTokenizer:
+    """Reports exactly HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS -- a "real" tokenizer
+    that hasn't drifted from the hardcoded table."""
+
+    def convert_tokens_to_ids(self, tok: str) -> int:
+        return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok]
+
+
+class _DriftedTokenizer:
+    """Reports a different id for one token -- simulates a model/tokenizer
+    revision bump shifting special-token ids out from under the hardcoded
+    table."""
+
+    def convert_tokens_to_ids(self, tok: str) -> int:
+        if tok == "<think>":
+            return 999999
+        return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok]
+
+
+class _BaseCheckpointTokenizer:
+    """Simulates tencent/HunyuanImage-3.0 (base): reports the same
+    model_class_name as -Instruct, but its vocab doesn't include the
+    Instruct-only <img_ratio_33>-<img_ratio_36> tokens -- convert_tokens_to_ids
+    falls back to unk_token_id for those, like a real HF tokenizer would."""
+
+    unk_token_id = 0
+    _INSTRUCT_ONLY = {"<img_ratio_33>", "<img_ratio_36>"}
+
+    def convert_tokens_to_ids(self, tok: str) -> int:
+        if tok in self._INSTRUCT_ONLY:
+            return self.unk_token_id
+        return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok]
+
+
+def test_validate_special_token_ids_passes_for_matching_tokenizer():
+    validate_special_token_ids(_MatchingTokenizer())  # must not raise
+
+
+def test_validate_special_token_ids_raises_on_drift():
+    with pytest.raises(ValueError, match="<think>.*expected 128023.*got 999999"):
+        validate_special_token_ids(_DriftedTokenizer())
+
+
+def test_validate_special_token_ids_tolerates_base_checkpoint_missing_tokens():
+    """The base checkpoint's tokenizer not recognizing Instruct-only ratio
+    tokens (unk_token_id, not a real drifted id) must not be flagged as a
+    mismatch -- both checkpoints share model_class_name, so this hook fires
+    for either one."""
+    validate_special_token_ids(_BaseCheckpointTokenizer())  # must not raise

tests/diffusion/models/lance/test_pipeline_lance.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from vllm_omni.diffusion.models.lance.lance_transformer import LanceBagel
+from vllm_omni.diffusion.models.lance.pipeline_lance import LancePipeline
+from vllm_omni.model_executor.models.utils import normalize_decoded_video_frames
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FramesWithMetadata(list[Image.Image]):
+    def __init__(self, frames: list[Image.Image], **metadata: object) -> None:
+        super().__init__(frames)
+        for name, value in metadata.items():
+            setattr(self, name, value)
+
+
+class StopAfterVideoPreprocessError(Exception):
+    pass
+
+
+def test_normalize_decoded_video_frames_preserves_rgb_pixels_and_fps():
+    red = Image.new("RGB", (3, 2), color=(255, 0, 0))
+    rgba = Image.new("RGBA", (3, 2), color=(0, 255, 0, 64))
+    rgba_bytes = rgba.tobytes()
+
+    video, fps = normalize_decoded_video_frames(FramesWithMetadata([red, rgba], fps=8), default_fps=24.0)
+
+    assert video.shape == (2, 2, 3, 3)
+    assert video.dtype == np.uint8
+    assert video.flags.c_contiguous
+    np.testing.assert_array_equal(video[0, 0, 0], [255, 0, 0])
+    np.testing.assert_array_equal(video[1, 0, 0], [0, 255, 0])
+    assert fps == 8.0
+    assert rgba.mode == "RGBA"
+    assert rgba.tobytes() == rgba_bytes
+
+
+def test_normalize_decoded_video_frames_uses_frame_rate_when_fps_is_missing():
+    frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, frame_rate=12)
+
+    _, fps = normalize_decoded_video_frames(frames, default_fps=24.0)
+
+    assert fps == 12.0
+
+
+def test_normalize_decoded_video_frames_prefers_valid_fps_over_frame_rate():
+    frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps=8, frame_rate=12)
+
+    _, fps = normalize_decoded_video_frames(frames, default_fps=24.0)
+
+    assert fps == 8.0
+
+
+def test_normalize_decoded_video_frames_uses_default_fps_for_plain_list_and_tuple():
+    frames = [Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2))]
+
+    assert normalize_decoded_video_frames(frames, default_fps=24.0)[1] == 24.0
+    assert normalize_decoded_video_frames(tuple(frames), default_fps=24.0)[1] == 24.0
+
+
+def test_normalize_decoded_video_frames_accepts_numeric_string_fps():
+    frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps="12.5")
+
+    _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0)
+
+    assert frame_rate == 12.5
+
+
+@pytest.mark.parametrize(
+    ("fps", "frame_rate", "expected"),
+    [
+        (float("nan"), 12, 12.0),
+        (0, 15, 15.0),
+        (-1, 8, 8.0),
+        (float("inf"), 10, 10.0),
+        ("invalid", "6.5", 6.5),
+    ],
+)
+def test_normalize_decoded_video_frames_uses_valid_frame_rate_after_invalid_fps(fps, frame_rate, expected):
+    frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps=fps, frame_rate=frame_rate)
+
+    _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0)
+
+    assert frame_rate == expected
+
+
+def test_normalize_decoded_video_frames_uses_default_when_both_metadata_values_are_invalid():
+    frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps="invalid", frame_rate=0)
+
+    _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0)
+
+    assert frame_rate == 24.0
+
+
+@pytest.mark.parametrize(
+    ("frames", "match"),
+    [
+        ([], "empty decoded video frame sequence"),
+        ([Image.new("RGB", (2, 2)), object()], "must be a PIL.Image.Image"),
+        (
+            [Image.new("RGB", (2, 2)), Image.new("RGB", (3, 2))],
+            "must have identical dimensions",
+        ),
+    ],
+)
+def test_normalize_decoded_video_frames_rejects_invalid_sequences(frames, match):
+    with pytest.raises(ValueError, match=match):
+        normalize_decoded_video_frames(frames, default_fps=24.0)
+
+
+@pytest.mark.parametrize(
+    ("frames", "match"),
+    [
+        ([Image.new("RGB", (0, 16))], r"index 0.*0x16"),
+        ([Image.new("RGB", (16, 0))], r"index 0.*16x0"),
+        ([Image.new("RGB", (0, 0))], r"index 0.*0x0"),
+        (
+            [Image.new("RGB", (16, 16)), Image.new("RGB", (0, 16))],
+            r"index 1.*0x16",
+        ),
+    ],
+)
+def test_normalize_decoded_video_frames_rejects_zero_dimensions(frames, match):
+    with pytest.raises(ValueError, match=match):
+        normalize_decoded_video_frames(frames, default_fps=24.0)
+
+
+@pytest.mark.parametrize(
+    ("video_input", "expected_fps"),
+    [
+        (FramesWithMetadata([Image.new("RGB", (3, 2))] * 2, fps=8), 8.0),
+        (tuple([Image.new("RGBA", (3, 2))] * 2), 12.0),
+    ],
+)
+def test_lance_video_edit_normalizes_decoded_frame_sequences(monkeypatch, video_input, expected_fps):
+    captured = {}
+
+    def capture_preprocessed_input(video, fps):
+        captured["video"] = video
+        captured["fps"] = fps
+        raise StopAfterVideoPreprocessError
+
+    monkeypatch.setattr(LanceBagel, "_lance_video_preprocess", staticmethod(capture_preprocessed_input))
+    request = SimpleNamespace(
+        prompts=[
+            {
+                "modalities": ["video"],
+                "user_text": "edit the video",
+                "multi_modal_data": {"video": video_input},
+                "extra_args": {"origin_fps": 12.0},
+            }
+        ],
+        sampling_params=SimpleNamespace(extra_args={}, num_inference_steps=1),
+    )
+
+    with pytest.raises(StopAfterVideoPreprocessError):
+        object.__new__(LancePipeline)._forward_video_edit(request)
+
+    video = captured["video"]
+    assert isinstance(video, np.ndarray)
+    assert video.shape == (2, 2, 3, 3)
+    assert video.dtype == np.uint8
+    assert video.flags.c_contiguous
+    assert captured["fps"] == expected_fps

tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py

@@ -12,6 +12,7 @@ from vllm_omni.diffusion.models.lingbot_world.actions import (
     LINGBOT_CAMERA_ACTION_SCHEMA,
     LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
     LingBotCameraControlReducer,
+    camera_trajectory_from_absolute_pose,
     integrate_lingbot_camera_actions,
     parse_lingbot_camera_action_frames,
     parse_lingbot_camera_action_script,
@@ -212,6 +213,11 @@ def test_action_integrator_matches_lingbot_motion_and_calibration() -> None:
         torch.tensor([500.0, 500.0, 208.0, 120.0]),
     )
 
+    with pytest.raises(ValueError, match="resolution must be positive"):
+        camera_trajectory_from_absolute_pose(trajectory.poses[:1], width=0, height=480)
+    with pytest.raises(ValueError, match=r"\[frames, 4, 4\]"):
+        camera_trajectory_from_absolute_pose(torch.zeros(4, 4), width=832, height=480)
+
 
 @pytest.mark.parametrize(
     ("data", "message"),

tests/diffusion/models/lingbot_world/test_lingbot_world_attention.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
 
 from __future__ import annotations
 
@@ -7,7 +7,7 @@ import importlib.util
 import math
 import sys
 from pathlib import Path
-from types import ModuleType
+from types import ModuleType, SimpleNamespace
 
 import pytest
 import torch
@@ -33,6 +33,10 @@ _STUBBED_MODULE_NAMES = (
     "vllm_omni.diffusion",
     "vllm_omni.diffusion.attention",
     "vllm_omni.diffusion.attention.layer",
+    "vllm_omni.diffusion.distributed",
+    "vllm_omni.diffusion.distributed.comm",
+    "vllm_omni.diffusion.distributed.parallel_state",
+    "vllm_omni.diffusion.distributed.sp_plan",
     "vllm_omni.diffusion.layers",
     "vllm_omni.diffusion.layers.norm",
     "vllm_omni.diffusion.layers.rope",
@@ -58,6 +62,58 @@ def _install_vllm_stubs() -> None:
     distributed.get_tensor_model_parallel_world_size = lambda: 1
     distributed.tensor_model_parallel_all_reduce = lambda value: value
 
+    class _SeqAllToAll4D:
+        @staticmethod
+        def apply(group, value, scatter_idx, gather_idx, use_sync=False):
+            del group, scatter_idx, gather_idx, use_sync
+            return value
+
+    setattr(
+        sys.modules["vllm_omni.diffusion.distributed.comm"],
+        "SeqAllToAll4D",
+        _SeqAllToAll4D,
+    )
+
+    def _all_to_all_5D(value, scatter_idx=3, gather_idx=1, group=None, use_sync=False):
+        # Single-rank stub: the fused q/k/v exchange is the identity when N == 1.
+        del scatter_idx, gather_idx, group, use_sync
+        return value
+
+    setattr(
+        sys.modules["vllm_omni.diffusion.distributed.comm"],
+        "all_to_all_5D",
+        _all_to_all_5D,
+    )
+
+    def get_sp_group():
+        return SimpleNamespace(
+            ulysses_world_size=1,
+            ulysses_rank=0,
+            ulysses_group=None,
+        )
+
+    setattr(
+        sys.modules["vllm_omni.diffusion.distributed.parallel_state"],
+        "get_sp_group",
+        get_sp_group,
+    )
+
+    class _SequenceParallelInput:
+        def __init__(self, split_dim, expected_dims=None, split_output=False, auto_pad=False):
+            self.split_dim = split_dim
+            self.expected_dims = expected_dims
+            self.split_output = split_output
+            self.auto_pad = auto_pad
+
+    class _SequenceParallelOutput:
+        def __init__(self, gather_dim, expected_dims=None):
+            self.gather_dim = gather_dim
+            self.expected_dims = expected_dims
+
+    sp_plan = sys.modules["vllm_omni.diffusion.distributed.sp_plan"]
+    setattr(sp_plan, "SequenceParallelInput", _SequenceParallelInput)
+    setattr(sp_plan, "SequenceParallelOutput", _SequenceParallelOutput)
+
     def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None:
         for name, value in attrs.items():
             setattr(weight, name, value)
@@ -541,3 +597,51 @@ def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyP
     norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0]))
 
     torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
+
+
+@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention"])
+def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class):
+    module = _load_module()
+    monkeypatch.setattr(
+        module,
+        "get_sp_group",
+        lambda: SimpleNamespace(
+            ulysses_world_size=3,
+            ulysses_rank=0,
+            ulysses_group=None,
+        ),
+    )
+    with pytest.raises(ValueError, match="heads must be divisible"):
+        getattr(module, attention_class)(dim=8, num_heads=4)
+
+
+def test_supplied_camera_modulation_matches_computing_it_inline() -> None:
+    """Passing a prebuilt modulation must be arithmetically identical to letting the block build its own."""
+    module = _load_module()
+    torch.manual_seed(0)
+    block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+    camera = torch.randn(1, 3, 4)
+
+    scale, shift = block.camera_modulation(camera)
+    inline_scale, inline_shift = block.camera_modulation(camera)
+
+    torch.testing.assert_close(scale, inline_scale, rtol=0, atol=0)
+    torch.testing.assert_close(shift, inline_shift, rtol=0, atol=0)
+
+
+def test_the_camera_modulation_reads_only_the_camera_tokens() -> None:
+    """The reuse window is only sound because the modulation depends on nothing else that varies per forward."""
+    module = _load_module()
+    torch.manual_seed(0)
+    block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+    camera = torch.randn(1, 3, 4)
+
+    first = block.camera_modulation(camera)
+    # Same camera, different everything else a forward would carry: the modulation cannot move.
+    second = block.camera_modulation(camera.clone())
+
+    torch.testing.assert_close(first[0], second[0], rtol=0, atol=0)
+    torch.testing.assert_close(first[1], second[1], rtol=0, atol=0)
+    # A different camera must move it, or the cache key would be hiding a real dependency.
+    other = block.camera_modulation(camera + 1)
+    assert not torch.equal(first[0], other[0])

tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py

@@ -8,6 +8,7 @@ import math
 import subprocess
 import sys
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pytest
@@ -20,7 +21,10 @@ _CAMERA_MODULE = importlib.util.module_from_spec(_SPEC)
 sys.modules[_SPEC.name] = _CAMERA_MODULE
 _SPEC.loader.exec_module(_CAMERA_MODULE)
 
-CameraTrajectory = _CAMERA_MODULE.CameraTrajectory
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory
+else:
+    CameraTrajectory = _CAMERA_MODULE.CameraTrajectory
 build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding
 interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory
 load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory
@@ -138,21 +142,20 @@ def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_pa
         load_camera_trajectory(_trusted(tmp_path))
 
 
-def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None:
-    _write_trajectory(
-        tmp_path,
-        _identity_poses(269),
-        np.arange(269 * 4, dtype=np.float32).reshape(269, 4),
-    )
+@pytest.mark.parametrize("num_frames", [129, 269, 4096])
+def test_load_camera_trajectory_preserves_all_bounded_source_frames(tmp_path: Path, num_frames: int) -> None:
+    poses = _identity_poses(num_frames)
+    poses[:, 0, 3] = np.arange(num_frames)
+    intrinsics = np.arange(num_frames * 4, dtype=np.float32).reshape(num_frames, 4)
+    _write_trajectory(tmp_path, poses, intrinsics)
 
     trajectory = load_camera_trajectory(_trusted(tmp_path))
 
-    assert trajectory.poses.shape == (117, 4, 4)
-    assert trajectory.intrinsics.shape == (117, 4)
-    np.testing.assert_array_equal(
-        trajectory.intrinsics.numpy(),
-        np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117],
-    )
+    assert trajectory.poses.shape == (num_frames, 4, 4)
+    assert trajectory.intrinsics.shape == (num_frames, 4)
+    assert trajectory.poses.dtype == trajectory.intrinsics.dtype == torch.float32
+    np.testing.assert_array_equal(trajectory.poses.numpy(), poses.astype(np.float32))
+    np.testing.assert_array_equal(trajectory.intrinsics.numpy(), intrinsics)
 
 
 def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None:
@@ -419,7 +422,17 @@ def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_transla
     torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction)
 
 
-def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None:
+@pytest.mark.parametrize(
+    ("translation_scale", "expected_origins"),
+    [
+        (None, [[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]]),
+        (0.05, [[0.0, 0.0, 0.0], [60.0, 80.0, 0.0], [120.0, 160.0, 0.0]]),
+    ],
+)
+def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm(
+    translation_scale: float | None,
+    expected_origins: list[list[float]],
+) -> None:
     poses = torch.eye(4).repeat(3, 1, 1)
     poses[:, :3, 3] = torch.tensor(
         [
@@ -437,10 +450,37 @@ def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm()
         target_width=1,
         device=torch.device("cpu"),
         dtype=torch.float32,
+        translation_scale=translation_scale,
     )
 
-    expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]])
-    torch.testing.assert_close(result[:, :3, 0, 0], expected_origins)
+    torch.testing.assert_close(result[:, :3, 0, 0], torch.tensor(expected_origins))
+
+    from vllm_omni.diffusion.models.lingbot_world.actions import LINGBOT_CONTROLLER_TRANSLATION_UNIT
+
+    if translation_scale != LINGBOT_CONTROLLER_TRANSLATION_UNIT:
+        return
+    # Fixed controller unit keeps relative speed; max-norm would collapse both to 1.
+    slow = torch.eye(4).repeat(2, 1, 1)
+    slow[1, 2, 3] = 0.005
+    fast = torch.eye(4).repeat(2, 1, 1)
+    fast[1, 2, 3] = 0.5
+    kwargs = dict(
+        height=1,
+        width=1,
+        target_height=1,
+        target_width=1,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+        translation_scale=LINGBOT_CONTROLLER_TRANSLATION_UNIT,
+    )
+    slow_origin = build_plucker_embedding(_trajectory(slow), **kwargs)[1, :3, 0, 0]
+    fast_origin = build_plucker_embedding(_trajectory(fast), **kwargs)[1, :3, 0, 0]
+    torch.testing.assert_close(slow_origin, torch.tensor([0.0, 0.0, 0.1]))
+    torch.testing.assert_close(fast_origin, torch.tensor([0.0, 0.0, 10.0]))
+    wasd = torch.eye(4).repeat(2, 1, 1)
+    wasd[1, 2, 3] = LINGBOT_CONTROLLER_TRANSLATION_UNIT
+    wasd_origin = build_plucker_embedding(_trajectory(wasd), **kwargs)[1, :3, 0, 0]
+    torch.testing.assert_close(wasd_origin, torch.tensor([0.0, 0.0, 1.0]))
 
 
 def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None:

tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.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
 
 from __future__ import annotations
 
@@ -7,14 +7,17 @@ import hashlib
 import inspect
 import json
 import math
+from datetime import timedelta
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 import torch
 
 from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests
+from tests.helpers.mark import hardware_test
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 _FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture"
 _SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture"
@@ -32,11 +35,13 @@ def _tiny_model(
     num_layers: int = 2,
     num_frames_per_block: int = 1,
     sliding_window_num_frames: int = 3,
+    attention_head_dim: int = 2,
+    rope_max_seq_len: int = 16,
 ):
     return module.CausalLingBotWorldTransformer3DModel(
         patch_size=(1, 2, 2),
         num_attention_heads=2,
-        attention_head_dim=2,
+        attention_head_dim=attention_head_dim,
         in_channels=36,
         out_channels=2,
         text_dim=6,
@@ -45,7 +50,7 @@ def _tiny_model(
         num_layers=num_layers,
         cross_attn_norm=True,
         eps=1e-6,
-        rope_max_seq_len=16,
+        rope_max_seq_len=rope_max_seq_len,
         sink_size=1,
         num_frames_per_block=num_frames_per_block,
         sliding_window_num_frames=sliding_window_num_frames,
@@ -117,7 +122,9 @@ def _assert_self_cache_unchanged(cache, snapshot) -> None:
         attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot)
 
 
-def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None:
+@pytest.mark.cpu
+@pytest.mark.parametrize("num_chunks", [4, 18])
+def test_tiny_transformer_runs_contiguous_chunks_with_explicit_cache_commit_and_camera_path(num_chunks) -> None:
     torch.manual_seed(7)
     module = attention_tests._load_module()
     model = _tiny_model(module).eval()
@@ -131,7 +138,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera
     assert model.patch_embedding.in_channels == 36
     assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2
 
-    for start_frame in range(4):
+    for start_frame in range(num_chunks):
         hidden_states = torch.randn(1, 36, 1, 4, 4)
         timestep = torch.tensor([float(start_frame + 1)])
         snapshot = _self_cache_snapshot(cache)
@@ -165,7 +172,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera
         expected_token_end = (start_frame + 1) * 2 * 2
         assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention)
 
-    assert len(outputs) == 4
+    assert len(outputs) == num_chunks
     assert all(len(outputs) == 1 for outputs in cross_key_outputs)
     assert len(text_projection_outputs) == 1
 
@@ -191,6 +198,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera
     assert not torch.equal(camera_output, alternate_output)
 
 
+@pytest.mark.cpu
 def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(module, num_layers=2).eval()
@@ -229,6 +237,7 @@ def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -
 
 
 @pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks"))
+@pytest.mark.cpu
 def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None:
     module = attention_tests._load_module()
     model = _tiny_model(
@@ -257,6 +266,7 @@ def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames:
     assert cache.cross_attention == [None]
 
 
+@pytest.mark.cpu
 def test_forward_accepts_exactly_one_configured_frame_block() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(
@@ -279,6 +289,7 @@ def test_forward_accepts_exactly_one_configured_frame_block() -> None:
     assert output.shape == (1, 2, 3, 4, 4)
 
 
+@pytest.mark.cpu
 def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(
@@ -303,6 +314,7 @@ def test_transformer_allocates_request_cache_from_its_configured_geometry() -> N
     assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
 
 
+@pytest.mark.cpu
 def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
@@ -323,6 +335,7 @@ def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None:
     torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0]))
 
 
+@pytest.mark.cpu
 def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
@@ -363,6 +376,7 @@ def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None:
     torch.testing.assert_close(output, expected)
 
 
+@pytest.mark.cpu
 def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
@@ -381,6 +395,7 @@ def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None:
     torch.testing.assert_close(output, expected)
 
 
+@pytest.mark.cpu
 def test_constructor_defaults_match_official_checkpoint_config() -> None:
     module = attention_tests._load_module()
     parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters
@@ -406,6 +421,7 @@ def test_constructor_defaults_match_official_checkpoint_config() -> None:
     assert {name: parameters[name].default for name in expected} == expected
 
 
+@pytest.mark.cpu
 def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(module)
@@ -413,12 +429,14 @@ def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None:
     assert model.dtype == next(model.parameters()).dtype
 
 
+@pytest.mark.cpu
 def test_transformer_declares_regional_compile_block() -> None:
     module = attention_tests._load_module()
 
     assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"]
 
 
+@pytest.mark.cpu
 def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None:
     module = attention_tests._load_module()
     reduced_values: list[torch.Tensor] = []
@@ -440,6 +458,7 @@ def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None:
     torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]]))
 
 
+@pytest.mark.cpu
 def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None:
     module = attention_tests._load_module()
     monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
@@ -448,6 +467,7 @@ def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None:
         module.LingBotSelfAttention(dim=4, num_heads=2)
 
 
+@pytest.mark.cpu
 def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None:
     module = attention_tests._load_module()
 
@@ -456,6 +476,7 @@ def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None:
 
 
 @pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"])
+@pytest.mark.cpu
 def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None:
     module = attention_tests._load_module()
 
@@ -463,6 +484,7 @@ def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None
         module.CausalLingBotWorldTransformer3DModel(**{field: 4})
 
 
+@pytest.mark.cpu
 def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
     module = attention_tests._load_module()
 
@@ -470,6 +492,7 @@ def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> No
         module.CausalLingBotWorldTransformer3DModel(quant_config=object())
 
 
+@pytest.mark.cpu
 def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None:
     module = attention_tests._load_module()
 
@@ -522,6 +545,7 @@ def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> N
         ("sink_size", 0),
     ],
 )
+@pytest.mark.cpu
 def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None:
     module = attention_tests._load_module()
     config = {
@@ -554,6 +578,7 @@ def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object
         module.CausalLingBotWorldTransformer3DModel.from_config(config)
 
 
+@pytest.mark.cpu
 def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None:
     module = attention_tests._load_module()
     config = {
@@ -569,6 +594,7 @@ def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None:
     assert model.config.num_layers == 40
 
 
+@pytest.mark.cpu
 def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(module, num_layers=1)
@@ -606,6 +632,7 @@ def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() ->
         model.load_weights([("unexpected_model.weight", torch.ones(1))])
 
 
+@pytest.mark.cpu
 def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None:
     module = attention_tests._load_module()
     model = _tiny_model(module, num_layers=1)
@@ -619,6 +646,7 @@ def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None:
     assert model.load_weights(weights()) == {first_name}
 
 
+@pytest.mark.cpu
 def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None:
     utils = pytest.importorskip("vllm.model_executor.models.utils")
     module = attention_tests._load_module()
@@ -638,6 +666,7 @@ def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> No
     assert set(dict(pipeline.named_parameters())) <= loaded
 
 
+@pytest.mark.cpu
 def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None:
     module = attention_tests._load_module()
     fixture = json.loads(_FIXTURE_PATH.read_text())
@@ -667,6 +696,7 @@ def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None:
     assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values())
 
 
+@pytest.mark.cpu
 def test_official_default_shapes_match_public_safetensors_header_fixture() -> None:
     module = attention_tests._load_module()
     fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text())
@@ -683,3 +713,377 @@ def test_official_default_shapes_match_public_safetensors_header_fixture() -> No
         shape, dtype = parameter_specs[name]
         assert tuple(shape) == tuple(expected_shape), name
         assert dtype == expected_dtype
+
+
+@pytest.mark.cpu
+def test_timestep_expansion_only_runs_with_ulysses(monkeypatch):
+    dtype = torch.bfloat16
+    module = attention_tests._load_module()
+    hidden = torch.randn(2, 12, 4).to(dtype)
+    camera = torch.randn_like(hidden)
+    table = torch.randn(2, 3, 6, 4).to(dtype)
+    rotary = (torch.randn(12, 2), torch.randn(12, 2))
+    prepare = module._LingBotSPPrepare()
+    assert prepare(hidden, camera, table, rotary)[2] is table
+    monkeypatch.setattr(module, "get_sp_group", lambda: SimpleNamespace(ulysses_world_size=2))
+    expanded = prepare(hidden, camera, table, rotary)[2]
+    assert expanded.shape == (2, 12, 6, 4)
+    torch.testing.assert_close(expanded, table.repeat_interleave(4, dim=1), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("unsharded", [None, 3], ids=["missing-hook", "unsharded-rope"])
+@pytest.mark.cpu
+def test_missing_or_partial_sp_split_fails_before_attention(monkeypatch, unsharded):
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_frames_per_block=3, sliding_window_num_frames=6)
+    cache = _cache(module, model)
+    for block in model.blocks:
+        block.self_attn.ulysses_world_size = 2
+    monkeypatch.setattr(module, "get_sp_group", lambda: SimpleNamespace(ulysses_world_size=2))
+    original = model.sp_prepare.forward
+
+    def partial_split(*args):
+        values = original(*args)
+        if unsharded is None:
+            return values
+        return tuple(
+            value if i == unsharded else value.chunk(2, dim=1 if i < 3 else 0)[0] for i, value in enumerate(values)
+        )
+
+    monkeypatch.setattr(model.sp_prepare, "forward", partial_split)
+
+    def forbidden(*args, **kwargs):
+        raise AssertionError("attention/collectives must not run before shard validation")
+
+    monkeypatch.setattr(model.blocks[0], "forward", forbidden)
+    with pytest.raises(RuntimeError, match="SP input hooks"):
+        model(
+            torch.randn(1, 36, 3, 4, 4),
+            torch.tensor([1.0]),
+            torch.randn(1, 3, 6),
+            torch.randn(1, 384, 3, 4, 4),
+            cache=cache,
+            start_frame=0,
+            update_cache=False,
+        )
+
+
+# Real multi-rank regression: synthetic weights, native SP/TP and attention.
+_HEADS, _HEAD_DIM, _LAYERS = 8, 32, 2
+_FRAMES, _SIDE, _TOKENS_PER_FRAME = 3, 32, 256
+
+
+def _model(dtype):
+    from vllm_omni.diffusion.models.lingbot_world.transformer import CausalLingBotWorldTransformer3DModel
+
+    model = CausalLingBotWorldTransformer3DModel(
+        num_attention_heads=_HEADS,
+        attention_head_dim=_HEAD_DIM,
+        num_layers=_LAYERS,
+        in_channels=4,
+        out_channels=4,
+        text_dim=8,
+        freq_dim=16,
+        ffn_dim=512,
+        patch_size=(1, 2, 2),
+        sink_size=1,
+        num_frames_per_block=_FRAMES,
+        sliding_window_num_frames=6,
+        rope_max_seq_len=32,
+    ).eval()
+    return model.to(device=torch.device("cuda", torch.accelerator.current_device_index()), dtype=dtype)
+
+
+def _checkpoint(model):
+    weights = {}
+    generator = torch.Generator().manual_seed(17)
+    for name, param in sorted(model.named_parameters()):
+        data = torch.randn(param.shape, generator=generator) * 0.02
+        if "norm" in name and name.endswith("weight"):
+            data.fill_(1)
+        if name.endswith("modulation") and data.shape[-2] == 6:
+            data[..., 2, :] = 1
+            data[..., 5, :] = 1
+        data = data.to(param.dtype)
+        param.copy_(data)
+        if ".self_attn.qkv." in name:
+            for component, chunk in zip(("q", "k", "v"), data.chunk(3), strict=True):
+                weights[name.replace(".self_attn.qkv.", f".self_attn.{component}.")] = chunk
+        else:
+            weights[name] = data
+    return weights
+
+
+def _rollout(model, mode, dtype, batch):
+    from vllm_omni.diffusion.models.lingbot_world.transformer import LingBotTransformerCache
+    from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+    from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+    device = next(model.parameters()).device
+    state = None
+    if mode == "paged":
+        kv = ARDiffusionKVCache(
+            ARDiffusionKVConfig(enable=True, chunk_size=_TOKENS_PER_FRAME, window_chunks=5, sink_chunks=1),
+            num_layers=_LAYERS,
+            num_kv_heads=model.blocks[0].self_attn.num_sp_heads,
+            head_size=_HEAD_DIM,
+            dtype=dtype,
+            block_size=_TOKENS_PER_FRAME,
+            max_model_len=4096,
+            available_bytes=1 << 27,
+            kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+            session_capacity=1,
+            frames_per_block=_FRAMES,
+            max_scratch_tokens_per_branch=_FRAMES * _TOKENS_PER_FRAME,
+            cross_attention_lengths={"text": 5},
+            # Cross-attention keeps every local head on every SP rank, so its pool is wider than the
+            # self-attention share above.
+            cross_attention_kv_heads={"text": model.blocks[0].cross_attn.num_local_heads},
+            device=device,
+        )
+        state = ARDiffusionKVState(kv, "numeric", {"main": kv.begin_request("numeric")}, num_layers=_LAYERS)
+        cache = LingBotTransformerCache(self_attention=[], cross_attention=[None] * _LAYERS)
+    else:
+        cache = model.allocate_cache(
+            batch_size=batch, latent_height=_SIDE, latent_width=_SIDE, device=device, dtype=dtype
+        )
+    generator = torch.Generator().manual_seed(123)
+    text = torch.randn(batch, 5, 8, generator=generator).to(device, dtype)
+    outputs = []
+    try:
+        for start in range(0, 12, _FRAMES):
+            latent = torch.randn(batch, 4, _FRAMES, _SIDE, _SIDE, generator=generator).to(device, dtype)
+            camera = torch.randn(batch, 384, _FRAMES, _SIDE, _SIDE, generator=generator).to(device, dtype)
+            for commit in (False, False, True):
+                if state is not None:
+                    cache.self_attention = state.get_kv_caches(
+                        "main", seq_len=_FRAMES * _TOKENS_PER_FRAME, commit_current=commit
+                    )
+                output = model(
+                    latent,
+                    torch.tensor([100.0, 400.0, 900.0], device=device).expand(batch, -1),
+                    text,
+                    camera,
+                    cache=cache,
+                    start_frame=start,
+                    update_cache=commit,
+                )
+                outputs.append(output.float().cpu())
+                if state is not None:
+                    if not state.is_cross_attention_populated("main", "text"):
+                        state.populate_cross_attention(
+                            "main", "text", [(layer.key, layer.value) for layer in cache.cross_attention]
+                        )
+                        for layer, pool in zip(
+                            cache.cross_attention, state.get_cross_attention_kv("main", "text"), strict=True
+                        ):
+                            layer.key, layer.value = pool["k"], pool["v"]
+                    state.commit_paged_context("main")
+        cross = [
+            (layer.key.detach().float().cpu(), layer.value.detach().float().cpu()) for layer in cache.cross_attention
+        ]
+        return torch.stack(outputs), cross
+    finally:
+        if state is not None:
+            state.close()
+
+
+def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous):
+    from vllm.config import VllmConfig
+    from vllm.config.vllm import set_current_vllm_config
+    from vllm.distributed import get_tensor_model_parallel_rank
+
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        destroy_distributed_env,
+        destroy_model_parallel,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    from vllm_omni.diffusion.forward_context import set_forward_context
+    from vllm_omni.diffusion.registry import _apply_sequence_parallel_if_enabled
+    from vllm_omni.platforms import current_omni_platform
+
+    current_omni_platform.set_device(torch.device("cuda", rank))
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    torch.distributed.init_process_group(
+        "nccl", init_method=f"file://{rendezvous}", world_size=world_size, rank=rank, timeout=timedelta(seconds=90)
+    )
+    init_distributed_environment(world_size=world_size, rank=rank, local_rank=rank, backend="nccl")
+    try:
+        with torch.inference_mode(), set_current_vllm_config(VllmConfig()):
+            for baseline in (True, False):
+                sp, tp = (1, 1) if baseline else (sp_size, tp_size)
+                initialize_model_parallel(
+                    data_parallel_size=world_size // (sp * tp),
+                    sequence_parallel_size=sp,
+                    ulysses_degree=sp,
+                    tensor_parallel_size=tp,
+                )
+                config = OmniDiffusionConfig(
+                    model=str(Path(rendezvous).parent),
+                    dtype=dtype,
+                    enforce_eager=True,
+                    parallel_config=DiffusionParallelConfig(
+                        data_parallel_size=world_size // (sp * tp),
+                        sequence_parallel_size=sp,
+                        ulysses_degree=sp,
+                        tensor_parallel_size=tp,
+                    ),
+                    diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")),
+                )
+                with set_current_diffusion_config(config), set_forward_context(omni_diffusion_config=config):
+                    model = _model(dtype)
+                    if baseline:
+                        weights = _checkpoint(model)
+                    else:
+                        model.load_weights(iter(weights.items()))
+                    _apply_sequence_parallel_if_enabled(SimpleNamespace(transformer=model), config)
+                    result, cross = _rollout(model, mode, dtype, batch)
+                    if baseline:
+                        expected, expected_cross = result, cross
+                    else:
+                        assert expected.abs().max() > 1e-3
+                        bound = 1e-5 if dtype == torch.float32 else 1e-2
+                        error = (result - expected).double().norm() / expected.double().norm()
+                        assert error <= bound, f"rank={rank}, TP{tp} SP{sp}: relative L2 {error.item():.3g} > {bound}"
+                        # Cross-attention K/V is replicated across the Ulysses group, not split over it, so
+                        # every SP rank holds its whole TP shard of heads. Only TP slices the cache now.
+                        local_heads = _HEADS // tp
+                        first = get_tensor_model_parallel_rank() * local_heads
+                        for (k, v), (ek, ev) in zip(cross, expected_cross, strict=True):
+                            tolerance = 1e-5 if dtype == torch.float32 else 2e-2
+                            torch.testing.assert_close(
+                                k, ek[:, :, first : first + local_heads], rtol=tolerance, atol=tolerance
+                            )
+                            torch.testing.assert_close(
+                                v, ev[:, :, first : first + local_heads], rtol=tolerance, atol=tolerance
+                            )
+                        if rank == 0:
+                            print(f"{mode} {dtype} TP{tp} SP{sp}: relative L2={error.item():.3g}", flush=True)
+                    del model
+                torch.distributed.barrier()
+                destroy_model_parallel()
+    finally:
+        destroy_distributed_env()
+
+
+def _run(tmp_path: Path, sp, tp, mode, dtype, batch):
+    if torch.accelerator.device_count() < sp * tp:
+        pytest.skip(f"requires {sp * tp} GPUs")
+    torch.multiprocessing.spawn(
+        _worker, args=(sp * tp, sp, tp, mode, dtype, batch, str(tmp_path / "rendezvous")), nprocs=sp * tp
+    )
+
+
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_sp2_direct_matches_sp1_fp32(tmp_path):
+    _run(tmp_path, 2, 1, "direct", torch.float32, 2)
+
+
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=4)
+def test_sp4_direct_matches_sp1_bf16(tmp_path):
+    _run(tmp_path, 4, 1, "direct", torch.bfloat16, 1)
+
+
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=4)
+def test_tp2_sp2_paged_matches_sp1_bf16(tmp_path):
+    _run(tmp_path, 2, 2, "paged", torch.bfloat16, 1)
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("start_frame", [0, 1021, 1023, 1024, 21600])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_temporal_rope_matches_absolute_positions_without_growing_cache(start_frame, dtype):
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1, attention_head_dim=128, rope_max_seq_len=1024).to(dtype)
+    buffers = {name: (value.data_ptr(), value.shape) for name, value in model.named_buffers()}
+    cosine, sine = model._rotary_embedding(
+        frames=3, height=2, width=3, start_frame=start_frame, dtype=dtype, device=torch.device("cpu")
+    )
+    # Independent complex-valued reference uses absolute (time, height, width)
+    # coordinates, including blocks crossing the original precomputed table.
+    grids = torch.meshgrid(
+        torch.arange(start_frame, start_frame + 3, dtype=torch.float64),
+        torch.arange(2, dtype=torch.float64),
+        torch.arange(3, dtype=torch.float64),
+        indexing="ij",
+    )
+    reference = []
+    for positions, dim in zip(grids, (44, 42, 42), strict=True):
+        frequencies = 1.0 / 10000 ** (torch.arange(0, dim, 2, dtype=torch.float64) / dim)
+        phase = positions.flatten()[:, None] * frequencies[None, :]
+        reference.append(torch.polar(torch.ones_like(phase), phase))
+    expected = torch.cat(reference, dim=-1)
+    torch.testing.assert_close(cosine, expected.real.float().to(dtype), rtol=0, atol=0)
+    torch.testing.assert_close(sine, expected.imag.float().to(dtype), rtol=0, atol=0)
+    assert {name: (value.data_ptr(), value.shape) for name, value in model.named_buffers()} == buffers
+
+
+@pytest.mark.cpu
+def test_temporal_rope_extension_keeps_spatial_limit():
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1)
+    with pytest.raises(ValueError, match="Spatial RoPE"):
+        model._rotary_embedding(
+            frames=3, height=17, width=1, start_frame=16, dtype=torch.float32, device=torch.device("cpu")
+        )
+
+
+@pytest.mark.cpu
+def test_the_camera_injector_runs_once_per_held_cache_and_gives_the_same_output() -> None:
+    """Two forwards sharing one caller-held cache: the injector runs once per block, outputs unchanged."""
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=2).eval()
+    hidden_states = torch.randn(1, 36, 1, 4, 4)
+    timestep = torch.tensor([1.0])
+    encoder_hidden_states = torch.randn(1, 3, 6)
+    camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+
+    def run(cache, camera_cache=None):
+        return model(
+            hidden_states,
+            timestep,
+            encoder_hidden_states,
+            camera_hidden_states,
+            cache=cache,
+            start_frame=0,
+            update_cache=False,
+            camera_modulation_cache=camera_cache,
+        )
+
+    # Baseline: no cache, so every forward rebuilds the injector in every block.
+    uncached_first = run(_cache(module, model))
+    uncached_second = run(_cache(module, model))
+    injector_calls = [attention_tests._record_outputs(block.cam_injector_layer1) for block in model.blocks]
+    held = module.CameraModulationCache()
+    cached_first = run(_cache(module, model), held)
+    cached_second = run(_cache(module, model), held)
+
+    # One build per block for both forwards, instead of one per block per forward.
+    assert [len(calls) for calls in injector_calls] == [1, 1]
+    torch.testing.assert_close(cached_first, uncached_first, rtol=0, atol=0)
+    torch.testing.assert_close(cached_second, uncached_second, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_timestep_projection_is_staged_into_one_static_buffer_per_shape() -> None:
+    """Two forwards' projections land in the same buffer; the blocks never see the fresh per-forward tensor."""
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1).eval()
+    first = torch.randn(1, 8, 6, model.dim)
+    second = torch.randn(1, 8, 6, model.dim)
+
+    staged_first = model._stage_timestep_projection(first)
+    assert staged_first is not first and torch.equal(staged_first, first)
+    staged_second = model._stage_timestep_projection(second)
+    assert staged_second is staged_first and torch.equal(staged_second, second)
+    # A different shape gets its own buffer; the first one is kept.
+    other = model._stage_timestep_projection(torch.randn(1, 4, 6, model.dim))
+    assert other is not staged_first and len(model._timestep_projection_buffers) == 2

tests/diffusion/models/longcat_image/test_longcat_image_transformer.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+from diffusers.models.embeddings import apply_rotary_emb
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(autouse=True)
+def _clear_qk_rope_signature_caches():
+    from vllm_omni.diffusion.models.longcat_image import longcat_image_transformer as longcat
+
+    longcat._FAILED_QK_ROPE_SIGNATURES.clear()
+    yield
+    longcat._FAILED_QK_ROPE_SIGNATURES.clear()
+
+
+def _inputs(sequence: int = 5, head_dim: int = 8):
+    q = torch.randn(2, sequence, 3, head_dim, dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    cos = torch.randn(sequence, head_dim)
+    sin = torch.randn_like(cos)
+    return q, k, (cos, sin)
+
+
+def _reference(q, k, rotary_emb):
+    return (
+        apply_rotary_emb(q, rotary_emb, sequence_dim=1),
+        apply_rotary_emb(k, rotary_emb, sequence_dim=1),
+    )
+
+
+def test_prepare_rotary_emb_resolves_threshold_once_only_when_enabled(monkeypatch):
+    from vllm_omni.diffusion.models.longcat_image import longcat_image_transformer as longcat
+
+    calls = 0
+
+    def resolve(default):
+        nonlocal calls
+        calls += 1
+        assert default == 512
+        return 37
+
+    monkeypatch.setattr(longcat, "fused_qk_norm_rope_min_tokens", resolve)
+    txt_cos, txt_sin = torch.randn(2, 8), torch.randn(2, 8)
+    img_cos, img_sin = torch.randn(3, 8), torch.randn(3, 8)
+    native = longcat._prepare_rotary_emb(txt_cos, txt_sin, img_cos, img_sin, enable_fusion=False)
+    fused = longcat._prepare_rotary_emb(txt_cos, txt_sin, img_cos, img_sin, enable_fusion=True)
+
+    assert calls == 1
+    assert len(native) == 2
+    assert len(fused) == 3 and fused[2] == 37
+    assert torch.equal(native[0], fused[0])
+    assert torch.equal(native[1], fused[1])
+
+
+@pytest.mark.parametrize(
+    ("mode", "sp_size"),
+    [
+        ("unsupported", 1),
+        ("compile", 1),
+        ("grad", 1),
+        ("capture", 1),
+        ("sequence_parallel", 2),
+        ("below_threshold", 1),
+    ],
+)
+def test_qk_rope_non_eager_modes_use_exact_original_path(monkeypatch, mode, sp_size):
+    from vllm_omni.diffusion.models.longcat_image import longcat_image_transformer as longcat
+
+    torch.manual_seed(1)
+    q, k, rotary_emb = _inputs()
+    expected = _reference(q, k, rotary_emb)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: mode == "compile")
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: mode == "capture")
+    monkeypatch.setattr(longcat, "fused_qk_rope_supported", lambda *args: mode != "unsupported")
+    monkeypatch.setattr(
+        longcat,
+        "fused_qk_rope",
+        lambda *args: (_ for _ in ()).throw(AssertionError("fallback called fused op")),
+    )
+
+    context = nullcontext() if mode == "grad" else torch.no_grad()
+    min_tokens = q.shape[0] * q.shape[1] + 1 if mode == "below_threshold" else 0
+    with context:
+        actual = longcat._apply_qk_rope(q, k, rotary_emb, sp_size, min_tokens)
+
+    assert torch.equal(actual[0], expected[0])
+    assert torch.equal(actual[1], expected[1])
+
+
+def test_qk_rope_eligible_path_skips_runtime_parity_check(monkeypatch):
+    from vllm_omni.diffusion.models.longcat_image import longcat_image_transformer as longcat
+
+    torch.manual_seed(2)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    monkeypatch.setattr(longcat, "fused_qk_rope_supported", lambda *args: True)
+    fused_calls = 0
+
+    def fused(query, key, cos, sin):
+        nonlocal fused_calls
+        fused_calls += 1
+        return query + 1, key + 2
+
+    monkeypatch.setattr(
+        longcat,
+        "_apply_qk_rope_reference",
+        lambda *args: (_ for _ in ()).throw(AssertionError("eligible path ran the native parity check")),
+    )
+    monkeypatch.setattr(longcat, "fused_qk_rope", fused)
+
+    with torch.no_grad():
+        q, k, rotary_emb = _inputs()
+        actual = longcat._apply_qk_rope(q, k, rotary_emb, 1, 0)
+
+    assert fused_calls == 1
+    assert torch.equal(actual[0], q + 1)
+    assert torch.equal(actual[1], k + 2)
+
+
+def test_qk_rope_launch_failure_permanently_falls_back(monkeypatch):
+    from vllm_omni.diffusion.models.longcat_image import longcat_image_transformer as longcat
+
+    torch.manual_seed(3)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    monkeypatch.setattr(longcat, "fused_qk_rope_supported", lambda *args: True)
+    fused_calls = 0
+
+    def broken_fused(query, key, cos, sin):
+        nonlocal fused_calls
+        fused_calls += 1
+        raise RuntimeError("kernel failed")
+
+    monkeypatch.setattr(longcat, "fused_qk_rope", broken_fused)
+    q, k, rotary_emb = _inputs()
+    expected = _reference(q, k, rotary_emb)
+    with torch.no_grad():
+        first = longcat._apply_qk_rope(q, k, rotary_emb, 1, 0)
+        second = longcat._apply_qk_rope(q, k, rotary_emb, 1, 0)
+
+    assert fused_calls == 1
+    assert len(longcat._FAILED_QK_ROPE_SIGNATURES) == 1
+    for actual in (first, second):
+        assert torch.equal(actual[0], expected[0])
+        assert torch.equal(actual[1], expected[1])
+
+
+def test_failed_qk_rope_signature_cache_is_bounded():
+    from vllm_omni.diffusion.models.longcat_image import longcat_image_transformer as longcat
+
+    for index in range(longcat._FAILED_QK_ROPE_SIGNATURES_MAX_SIZE + 1):
+        longcat._record_failed_qk_rope_signature((index,))
+
+    assert len(longcat._FAILED_QK_ROPE_SIGNATURES) == longcat._FAILED_QK_ROPE_SIGNATURES_MAX_SIZE
+    assert (0,) not in longcat._FAILED_QK_ROPE_SIGNATURES
+    assert (longcat._FAILED_QK_ROPE_SIGNATURES_MAX_SIZE,) in longcat._FAILED_QK_ROPE_SIGNATURES
+
+
+class _FakeQKV(torch.nn.Module):
+    def __init__(self, heads: int, head_dim: int) -> None:
+        super().__init__()
+        self.num_heads = heads
+        self.num_kv_heads = heads
+        self.width = heads * head_dim
+
+    def forward(self, x):
+        return torch.cat((x, x + 1, x + 2), dim=-1), None
+
+
+class _AddNorm(torch.nn.Module):
+    def __init__(self, value: float) -> None:
+        super().__init__()
+        self.value = value
+
+    def forward(self, x):
+        return x + self.value
+
+
+class _CaptureAttention(torch.nn.Module):
+    def forward(self, query, key, value, metadata=None):
+        del key, value, metadata
+        return query
+
+
+class _PassthroughLinear(torch.nn.Module):
+    def forward(self, x):
+        return x, None
+
+
+def test_dual_attention_keeps_native_norm_and_text_image_concat_order(monkeypatch):
+    from vllm_omni.diffusion.models.longcat_image import longcat_image_transformer as longcat
+
+    attention = longcat.LongCatImageAttention.__new__(longcat.LongCatImageAttention)
+    torch.nn.Module.__init__(attention)
+    attention.parallel_config = SimpleNamespace(sequence_parallel_size=1)
+    attention.head_dim = 4
+    attention.added_kv_proj_dim = 8
+    attention.to_qkv = _FakeQKV(2, 4)
+    attention.add_kv_proj = _FakeQKV(2, 4)
+    attention.norm_q = _AddNorm(10)
+    attention.norm_k = _AddNorm(20)
+    attention.norm_added_q = _AddNorm(100)
+    attention.norm_added_k = _AddNorm(200)
+    attention.attn = _CaptureAttention()
+    attention.to_out = _PassthroughLinear()
+    attention.to_add_out = _PassthroughLinear()
+
+    captured = []
+
+    def paired(query, key, rotary_emb, sp_size, min_tokens):
+        captured.append((query.clone(), key.clone(), rotary_emb, sp_size, min_tokens))
+        return query, key
+
+    monkeypatch.setattr(longcat, "_apply_qk_rope", paired)
+    image = torch.randn(1, 3, 8)
+    text = torch.randn(1, 2, 8)
+    rotary_pair = (torch.randn(5, 4), torch.randn(5, 4))
+    rotary_emb = (*rotary_pair, 0)
+    image_out, text_out = attention(image, text, rotary_emb)
+
+    image_q = image.unflatten(-1, (2, 4))
+    text_q = text.unflatten(-1, (2, 4))
+    expected_q = torch.cat((text_q + 100, image_q + 10), dim=1)
+    expected_k = torch.cat((text_q + 1 + 200, image_q + 1 + 20), dim=1)
+    assert len(captured) == 1
+    assert torch.equal(captured[0][0], expected_q)
+    assert torch.equal(captured[0][1], expected_k)
+    assert captured[0][2][0] is rotary_pair[0]
+    assert captured[0][2][1] is rotary_pair[1]
+    assert captured[0][3:] == (1, 0)
+    assert image_out.shape == image.shape
+    assert text_out.shape == text.shape

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -233,6 +233,7 @@ def _fake_ltx_parallel_attention():
 
     layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
     nn.Module.__init__(layer)
+    layer.allow_fp32_fallback = False
     layer.attn_backend = Backend
     layer.attention = SimpleNamespace(forward=run("native"))
     layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))

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,78 @@
+# 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
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [
+    pytest.mark.diffusion,
+    pytest.mark.core_model,
+    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+    pytest.mark.skipif(not current_omni_platform.is_cuda(), 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_dlo_allgather.py

@@ -36,6 +36,7 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadBackend,
     DistributedLayerwiseOffloadHook,
 )
+from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent
 
 
 @dataclass(frozen=True)
@@ -46,12 +47,6 @@ class _WeightSourceStub:
     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]
 
@@ -189,7 +184,7 @@ def _load_mmap_transform_and_reconstruct(
         backend._mmap_transforms_by_tensor_id = {}
         backend._load_weights_via_mmap(
             pipeline,
-            _PipelineModulesStub(dits=[target], dit_names=["transformer"]),
+            (ResolvedComponent(path="transformer", module=target, selected=True),),
             plan_result.plan,
         )
 

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)

tests/diffusion/models/magi2/test_native_preview.py

@@ -90,6 +90,21 @@ def test_hsdp_policy_shards_individual_preview_layers():
     assert not condition("pre_adapter", model.pre_adapter)
 
 
+def test_layerwise_offload_contract() -> None:
+    from vllm_omni.diffusion.offloader import get_blocks_from_dit
+
+    model = Magi2PreviewTransformer(_tiny_config(num_layers=2))
+
+    assert model._layerwise_offload_blocks_attrs == ["block"]
+    assert "block" in dict(model.named_children())
+    assert list(model.block) == list(model.block.layers)
+    assert model.layers is model.block.layers
+
+    attr_names, blocks = get_blocks_from_dit(model)
+    assert attr_names == ["block"]
+    assert blocks == list(model.block.layers)
+
+
 def test_data_proxy_keeps_output_layout_request_scoped() -> None:
     proxy = Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=0))
 

tests/diffusion/models/mammoth_moda2/test_dit_attention.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MammothModa2 DiT self-attention runs through the shared Omni attention layer.
+
+``_reference_attention`` is the arithmetic the processor used before it was
+wired to the shared layer: SDPA over K/V replicated to the query head count,
+padded rows zeroed. Both sides read the same module weights, so any drift in
+the projections, the QK RMSNorm, the interleaved real RoPE, the softmax scale
+or the padding mask shows up as a mismatch here.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import (
+    TransformerBlock,
+    apply_real_rotary_emb,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Small stand-in for the real 2520 / 21 / 7 (head_dim 120) layout.
+DIM, HEADS = 48, 6
+SEQ, BATCH = 37, 2
+
+# The shared layer resolves its backend from the diffusion config at construction;
+# without one it takes the platform default, which on a CUDA host is FLASH_ATTN
+# and cannot run CPU tensors. Pin SDPA the way a deployment would.
+_SDPA_CONFIG = OmniDiffusionConfig(diffusion_attention_config={"default": {"backend": "TORCH_SDPA"}})
+
+
+def _block(kv_heads: int, modulation: bool = True) -> TransformerBlock:
+    torch.manual_seed(0)
+    with set_current_diffusion_config(_SDPA_CONFIG):
+        block = TransformerBlock(
+            DIM,
+            HEADS,
+            kv_heads,
+            multiple_of=8,
+            ffn_dim_multiplier=1.0,
+            norm_eps=1e-5,
+            modulation=modulation,
+        )
+    return block.eval()
+
+
+def _inputs(head_dim: int):
+    torch.manual_seed(1)
+    hidden = torch.randn(BATCH, SEQ, DIM)
+    mask = torch.ones(BATCH, SEQ, dtype=torch.bool)
+    mask[0, SEQ - 9 :] = False
+    mask[1, SEQ - 2 :] = False
+    angles = torch.rand(1, SEQ, head_dim) * 6.283
+    rotary = (angles.cos(), angles.sin())
+    temb = torch.randn(BATCH, min(DIM, 1024))
+    return hidden, mask, rotary, temb
+
+
+def _reference_attention(attn, hidden, mask, rotary):
+    """The pre-change processor, kept verbatim as the oracle."""
+    batch, seq, _ = hidden.shape
+    query, key, value = attn.to_q(hidden), attn.to_k(hidden), attn.to_v(hidden)
+    head_dim = query.shape[-1] // attn.heads
+    kv_heads = key.shape[-1] // head_dim
+    query = attn.norm_q(query.view(batch, seq, attn.heads, head_dim))
+    key = attn.norm_k(key.view(batch, seq, kv_heads, head_dim))
+    value = value.view(batch, seq, kv_heads, head_dim)
+    query = apply_real_rotary_emb(query, rotary[0], rotary[1])
+    key = apply_real_rotary_emb(key, rotary[0], rotary[1])
+    query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
+    if kv_heads < attn.heads:
+        key = key.repeat_interleave(attn.heads // kv_heads, dim=1)
+        value = value.repeat_interleave(attn.heads // kv_heads, dim=1)
+    out = F.scaled_dot_product_attention(query, key, value, attn_mask=mask.view(batch, 1, 1, seq), scale=attn.scale)
+    out = (out * mask[:, None, :, None]).transpose(1, 2).reshape(batch, seq, attn.heads * head_dim)
+    return attn.to_out[1](attn.to_out[0](out))
+
+
+@pytest.mark.parametrize("kv_heads", [2, HEADS], ids=["gqa_3to1", "mha"])
+def test_processor_matches_previous_arithmetic(kv_heads):
+    block = _block(kv_heads)
+    hidden, mask, rotary, _ = _inputs(block.head_dim)
+    with torch.no_grad():
+        got = block.attn(
+            hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary
+        )
+        want = _reference_attention(block.attn, hidden, mask, rotary)
+    torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5)
+    assert torch.count_nonzero(got[~mask]) == 0, "padded rows must stay zero"
+
+
+def test_no_mask_equals_all_valid_mask():
+    block = _block(2)
+    hidden, _, rotary, _ = _inputs(block.head_dim)
+    full = torch.ones(BATCH, SEQ, dtype=torch.bool)
+    with torch.no_grad():
+        unmasked = block.attn(
+            hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=None, image_rotary_emb=rotary
+        )
+        masked = block.attn(
+            hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=full, image_rotary_emb=rotary
+        )
+    torch.testing.assert_close(unmasked, masked, atol=1e-6, rtol=1e-6)
+
+
+@pytest.mark.parametrize("modulation", [True, False], ids=["noise_block", "context_refiner"])
+def test_block_forward_hands_native_kv_heads_to_shared_layer(monkeypatch, modulation):
+    """Through the block's real forward, on both of its branches (the modulated
+    noise/refiner blocks and the unmodulated context refiner): the shared layer
+    is what runs, and it receives K/V with their native head count -- no
+    replication before the call."""
+    block = _block(2, modulation=modulation)
+    hidden, mask, rotary, temb = _inputs(block.head_dim)
+    if not modulation:
+        temb = None
+    seen = []
+    original = block.attn.omni_attn.forward
+
+    def spy(query, key, value, attn_metadata=None):
+        seen.append((tuple(query.shape), tuple(key.shape), tuple(value.shape), attn_metadata))
+        return original(query, key, value, attn_metadata)
+
+    monkeypatch.setattr(block.attn.omni_attn, "forward", spy)
+    with torch.no_grad():
+        out = block(hidden, mask, rotary, temb)
+    assert out.shape == hidden.shape and torch.isfinite(out).all()
+    assert len(seen) == 1
+    q_shape, k_shape, v_shape, metadata = seen[0]
+    assert q_shape == (BATCH, SEQ, HEADS, block.head_dim)
+    assert k_shape == v_shape == (BATCH, SEQ, 2, block.head_dim)
+    assert metadata is not None and metadata.attn_mask.dtype == torch.bool
+    assert torch.equal(metadata.attn_mask, mask)
+
+
+def test_backend_selection_reaches_the_dit_layer():
+    """The diffusion config selects the backend when the DiT layer is constructed."""
+    assert _block(2).attn.omni_attn.attn_backend is DiffusionAttentionBackendEnum.TORCH_SDPA.get_class()
+    with set_current_diffusion_config(
+        OmniDiffusionConfig(diffusion_attention_config={"default": {"backend": "FLASH_ATTN"}})
+    ):
+        block = TransformerBlock(DIM, HEADS, 2, multiple_of=8, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+    assert block.attn.omni_attn.attn_backend is DiffusionAttentionBackendEnum.FLASH_ATTN.get_class()
+
+
+def test_shared_layer_owns_no_parameters_and_keeps_checkpoint_keys():
+    block = _block(2)
+    assert sum(p.numel() for p in block.attn.omni_attn.parameters()) == 0
+    assert not list(block.attn.omni_attn.buffers())
+    keys = set(block.state_dict())
+    assert not any("omni_attn" in key for key in keys)
+    # The projections and QK norms the checkpoint provides are still where they were.
+    for expected in (
+        "attn.to_q.weight",
+        "attn.to_k.weight",
+        "attn.to_v.weight",
+        "attn.to_out.0.weight",
+        "attn.norm_q.weight",
+        "attn.norm_k.weight",
+    ):
+        assert expected in keys, expected
+
+
+@pytest.mark.parametrize("with_mask", [True, False], ids=["empty_mask", "no_mask"])
+def test_empty_text_stream_skips_the_kernel(monkeypatch, with_mask):
+    """CFG's unconditional branch carries zero text tokens (the pipeline's default
+    negative_prompt_embeds has no rows), so the context refiner attends over an
+    empty sequence. The flash-attention varlen fallback cannot take that
+    (arange step 0), so the processor must not reach the shared layer at all."""
+    block = _block(kv_heads=2, modulation=False)
+    attn = block.attn
+
+    def _never(*args, **kwargs):
+        raise AssertionError("shared layer called with an empty sequence")
+
+    monkeypatch.setattr(attn.omni_attn, "forward", _never)
+    hidden = torch.randn(BATCH, 0, DIM)
+    mask = torch.ones(BATCH, 0, dtype=torch.bool) if with_mask else None
+    angles = torch.rand(1, 0, block.head_dim)
+    out = attn(
+        hidden_states=hidden,
+        encoder_hidden_states=hidden,
+        attention_mask=mask,
+        image_rotary_emb=(angles.cos(), angles.sin()),
+    )
+    assert out.shape == (BATCH, 0, DIM)

tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Real-shape check of the MammothModa2 DiT attention on CUDA: the shared
+backend (FLASH_ATTN by default, head_dim 120) against the pre-change arithmetic
+in bf16, with the padding mask the joint text+image sequence carries."""
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock
+
+from .test_dit_attention import _reference_attention
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.advanced_model,
+    pytest.mark.cuda,
+    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+]
+
+# MammothModa2-Preview gen_dit_config: hidden 2520, 21 heads, 7 KV heads, head_dim 120.
+DIM, HEADS, KV_HEADS = 2520, 21, 7
+
+
+@pytest.mark.parametrize("seq", [77 + 4096, 512], ids=["t2i_1024", "short"])
+def test_real_shape_matches_previous_arithmetic_bf16(seq):
+    torch.manual_seed(0)
+    block = (
+        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+        .cuda()
+        .to(torch.bfloat16)
+        .eval()
+    )
+    hidden = torch.randn(2, seq, DIM, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(2, seq, dtype=torch.bool, device="cuda")
+    mask[0, seq - 300 :] = False
+    mask[1, seq - 17 :] = False
+    angles = torch.rand(1, seq, block.head_dim, device="cuda") * 6.283
+    rotary = (angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16))
+    with torch.no_grad():
+        got = block.attn(
+            hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary
+        )
+        want = _reference_attention(block.attn, hidden, mask, rotary)
+    assert torch.isfinite(got).all()
+    assert torch.count_nonzero(got[~mask]) == 0
+    diff = (got.float() - want.float()).abs()[mask]
+    # bf16 kernels differ in accumulation order; measured on A800: max 1.0e-3, mean 5e-5.
+    assert diff.max().item() < 2e-2, diff.max().item()
+    assert diff.mean().item() < 1e-3, diff.mean().item()
+
+
+def test_empty_text_stream_on_the_default_backend():
+    """The recipe's text-to-image request with text_guidance_scale > 1 runs the
+    context refiner on a zero-token unconditional prompt. Before the guard this
+    raised ``RuntimeError: step must be nonzero`` from the FA varlen fallback."""
+    torch.manual_seed(0)
+    block = (
+        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5, modulation=False)
+        .cuda()
+        .to(torch.bfloat16)
+        .eval()
+    )
+    hidden = torch.randn(1, 0, DIM, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(1, 0, dtype=torch.bool, device="cuda")
+    angles = torch.rand(1, 0, block.head_dim, device="cuda")
+    with torch.no_grad():
+        out = block.attn(
+            hidden_states=hidden,
+            encoder_hidden_states=hidden,
+            attention_mask=mask,
+            image_rotary_emb=(angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)),
+        )
+    assert out.shape == (1, 0, DIM)
+
+
+@pytest.mark.parametrize(
+    ("seq", "force_gqa_fallback"),
+    [(512, False), (77 + 4096, False), (32, True)],
+    ids=["short", "t2i_1024", "unsupported_native_gqa"],
+)
+def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_gqa_fallback):
+    """FP32 keeps its dtype and reaches shared SDPA's GQA compatibility check.
+
+    Cover the long masked shape that otherwise falls back to quadratic math
+    attention, and force the unsupported-GQA branch independently of hardware.
+    """
+    torch.manual_seed(0)
+    dt = torch.float32
+    block = (
+        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+        .cuda()
+        .to(dt)
+        .eval()
+    )
+    hidden = torch.randn(1, seq, DIM, device="cuda", dtype=dt)
+    mask = torch.ones(1, seq, dtype=torch.bool, device="cuda")
+    mask[:, -13:] = False
+    angles = torch.rand(1, seq, block.head_dim, device="cuda")
+    rotary = (angles.cos().to(dt), angles.sin().to(dt))
+    capabilities, calls, phases = [], [], []
+    can_use_fused_gqa = sdpa_backend.can_sdpa_use_fused_gqa
+    sdpa = torch.nn.functional.scaled_dot_product_attention
+    strategy = block.attn.omni_attn._get_active_parallel_strategy()
+    pre_attention, post_attention = strategy.pre_attention, strategy.post_attention
+
+    def prepare(*args):
+        phases.append("pre")
+        return pre_attention(*args)
+
+    def restore(*args):
+        phases.append("post")
+        return post_attention(*args)
+
+    def check_gqa(*args):
+        supported = False if force_gqa_fallback else can_use_fused_gqa(*args)
+        capabilities.append(supported)
+        return supported
+
+    def record_sdpa(query, key, value, **kwargs):
+        phases.append("sdpa")
+        calls.append((query.shape, key.shape, value.shape, query.dtype, kwargs))
+        return sdpa(query, key, value, **kwargs)
+
+    with torch.no_grad(), monkeypatch.context() as patch:
+        patch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", check_gqa)
+        patch.setattr(torch.nn.functional, "scaled_dot_product_attention", record_sdpa)
+        patch.setattr(strategy, "pre_attention", prepare)
+        patch.setattr(strategy, "post_attention", restore)
+        got = block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary)
+
+    assert phases == ["pre", "sdpa", "post"], "FP32 must preserve the shared parallel dispatch"
+    assert len(capabilities) == len(calls) == 1, "FP32 must use shared SDPA's runtime GQA check"
+    q_shape, k_shape, v_shape, dtype, kwargs = calls[0]
+    expected_kv_heads = KV_HEADS if capabilities[0] else HEADS
+    assert q_shape == (1, HEADS, seq, block.head_dim)
+    assert k_shape == v_shape == (1, expected_kv_heads, seq, block.head_dim)
+    assert dtype == got.dtype == torch.float32
+    assert kwargs["enable_gqa"] == capabilities[0]
+    assert kwargs["scale"] == block.attn.scale
+    assert kwargs["is_causal"] is False
+    assert torch.equal(kwargs["attn_mask"], mask[:, None, None, :])
+
+    with torch.no_grad():
+        want = _reference_attention(block.attn, hidden, mask, rotary)
+    assert torch.isfinite(got).all()
+    assert torch.count_nonzero(got[~mask]) == 0
+    diff = (got.float() - want.float()).abs()[mask]
+    assert diff.max().item() < 1e-3, diff.max().item()

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

@@ -0,0 +1,392 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, TransformerConfig
+from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit
+from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
+    MammothModa2DiTPipeline,
+    _build_mammoth_config,
+    _root_weight_source,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _raw_config() -> dict:
+    return {
+        "model_type": "mammothmoda2",
+        "llm_config": {
+            "model_type": "mammothmoda2_qwen2_5_vl",
+            "text_config": {
+                "model_type": "mammothmoda2_qwen2_5_vl_text",
+                "hidden_size": 8,
+                "gen_vocab_start_index": 100,
+            },
+        },
+        "gen_vae_config": {"block_out_channels": [8, 8]},
+        "gen_dit_config": {"hidden_size": 8, "in_channels": 4},
+    }
+
+
+def _od_config() -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model="/models/MammothModa2-Preview",
+        model_class_name="MammothModa2DiTPipeline",
+        tf_model_config=TransformerConfig.from_dict(_raw_config()),
+    )
+
+
+def test_build_mammoth_config_uses_shared_transformer_projection() -> None:
+    config = _build_mammoth_config(_od_config())
+    assert config.model_type == "mammothmoda2"
+    assert config.get_text_config().hidden_size == 8
+    assert config.gen_dit_config == {"hidden_size": 8, "in_channels": 4}
+
+
+def test_build_mammoth_config_rejects_empty_shared_projection() -> None:
+    config = _od_config()
+    config.tf_model_config = TransformerConfig()
+    with pytest.raises(ValueError, match="root checkpoint config"):
+        _build_mammoth_config(config)
+
+
+def test_root_weight_source_loads_combined_checkpoint_once() -> None:
+    source = _root_weight_source(_od_config())
+    assert source.model_or_path == "/models/MammothModa2-Preview"
+    assert source.subfolder is None
+    assert source.prefix == ""
+    assert source.fall_back_to_pt is True
+
+
+def test_root_weight_source_forwards_revision() -> None:
+    config = _od_config()
+    config.revision = "rev-7"
+    assert _root_weight_source(config).revision == "rev-7"
+
+
+def test_pipeline_declares_native_components_and_single_request_mode_only() -> None:
+    assert MammothModa2DiTPipeline._dit_modules == ["gen_transformer"]
+    assert MammothModa2DiTPipeline._encoder_modules == ["gen_image_condition_refiner"]
+    assert MammothModa2DiTPipeline._vae_modules == ["gen_vae"]
+    assert MammothModa2DiTPipeline.supports_request_batch is False
+    assert MammothModa2DiTPipeline.supports_step_execution is False
+
+
+def test_mammoth_postprocess_denormalizes_nonnegative_raw_vae_output() -> None:
+    factory = getattr(pipeline_mammothmoda2_dit, "get_mammoth_moda2_post_process_func", None)
+    assert factory is not None
+
+    images = factory(_od_config())(torch.zeros(1, 3, 2, 2))
+
+    assert len(images) == 1
+    assert images[0].getpixel((0, 0)) == (128, 128, 128)
+
+
+def test_mammoth_postprocess_is_registered() -> None:
+    from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS
+
+    assert _DIFFUSION_POST_PROCESS_FUNCS["MammothModa2DiTPipeline"] == "get_mammoth_moda2_post_process_func"
+
+
+def test_root_weight_source_rejects_missing_model_path() -> None:
+    config = _od_config()
+    config.model = None
+    with pytest.raises(ValueError, match="model path"):
+        _root_weight_source(config)
+
+
+def _pipeline_shell() -> MammothModa2DiTPipeline:
+    pipeline = object.__new__(MammothModa2DiTPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.config = _build_mammoth_config(_od_config())
+    pipeline._llm_hidden_size = 8
+    return pipeline
+
+
+def _batch(
+    *,
+    request_id: str = "req-7",
+    prompt: object | None = None,
+    sampling: OmniDiffusionSamplingParams | None = None,
+) -> DiffusionRequestBatch:
+    if prompt is None:
+        prompt = {
+            "prompt": "",
+            "height": 32,
+            "width": 48,
+            "additional_information": {
+                "full_hidden_states": torch.arange(32, dtype=torch.float32).reshape(4, 8),
+                "full_token_ids": [10, 11, 100, 101],
+                "answer_start_index": 2,
+            },
+        }
+    if sampling is None:
+        sampling = OmniDiffusionSamplingParams(
+            height=32,
+            width=48,
+            seed=42,
+            guidance_scale=4.0,
+            num_inference_steps=7,
+            extra_args={"cfg_range": [0.2, 0.8]},
+        )
+    return DiffusionRequestBatch([OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id=request_id)])
+
+
+def test_parse_request_uses_standard_sampling_fields() -> None:
+    parsed = _pipeline_shell()._parse_request(_batch())
+    assert parsed.request_id == "req-7"
+    assert (parsed.height, parsed.width) == (32, 48)
+    assert parsed.num_inference_steps == 7
+    assert parsed.text_guidance_scale == 4.0
+    assert parsed.cfg_range == (0.2, 0.8)
+    assert parsed.seed == 42
+    assert parsed.answer_start_index == 2
+
+
+def test_parse_request_prefers_legacy_sampling_overrides() -> None:
+    sampling = OmniDiffusionSamplingParams(
+        guidance_scale=3.0,
+        num_inference_steps=5,
+        seed=9,
+        extra_args={"text_guidance_scale": 6.0, "num_inference_steps": 11, "cfg_range": [0.0, 0.5]},
+    )
+    parsed = _pipeline_shell()._parse_request(_batch(sampling=sampling))
+    assert parsed.text_guidance_scale == 6.0
+    assert parsed.num_inference_steps == 11
+    assert parsed.cfg_range == (0.0, 0.5)
+
+
+def test_parse_request_falls_back_to_request_level_sampling_values() -> None:
+    prompt = _batch().prompts[0]
+    prompt["additional_information"].update(
+        text_guidance_scale=[1.5],
+        num_inference_steps=[3],
+        cfg_range=[0.25, 0.75],
+    )
+
+    parsed = _pipeline_shell()._parse_request(_batch(prompt=prompt, sampling=OmniDiffusionSamplingParams()))
+
+    assert parsed.text_guidance_scale == 1.5
+    assert parsed.num_inference_steps == 3
+    assert parsed.cfg_range == (0.25, 0.75)
+
+
+def test_parse_request_standard_fields_precede_request_level_fallbacks() -> None:
+    prompt = _batch().prompts[0]
+    prompt["additional_information"].update(
+        text_guidance_scale=[1.5],
+        num_inference_steps=[3],
+    )
+    sampling = OmniDiffusionSamplingParams(guidance_scale=4.0, num_inference_steps=7)
+
+    parsed = _pipeline_shell()._parse_request(_batch(prompt=prompt, sampling=sampling))
+
+    assert parsed.text_guidance_scale == 4.0
+    assert parsed.num_inference_steps == 7
+
+
+def test_parse_request_rejects_multiple_requests_and_outputs() -> None:
+    pipeline = _pipeline_shell()
+    batch = _batch()
+    batch.requests.append(_batch(request_id="req-8").requests[0])
+    with pytest.raises(ValueError, match="exactly one request"):
+        pipeline._parse_request(batch)
+    with pytest.raises(ValueError, match="num_outputs_per_prompt=1"):
+        pipeline._parse_request(_batch(sampling=OmniDiffusionSamplingParams(num_outputs_per_prompt=2)))
+
+
+@pytest.mark.parametrize(
+    ("cfg_range", "message"),
+    [
+        ([0.5], "two values"),
+        ([-0.1, 0.5], "0 <= start <= end <= 1"),
+        ([0.7, 0.2], "0 <= start <= end <= 1"),
+        ([0.2, 1.1], "0 <= start <= end <= 1"),
+    ],
+)
+def test_parse_request_rejects_invalid_cfg_range(cfg_range, message) -> None:
+    batch = _batch(sampling=OmniDiffusionSamplingParams(extra_args={"cfg_range": cfg_range}))
+    with pytest.raises(ValueError, match=message):
+        _pipeline_shell()._parse_request(batch)
+
+
+def test_parse_request_requires_real_ar_conditions() -> None:
+    with pytest.raises(ValueError, match="req-missing"):
+        _pipeline_shell()._parse_request(_batch(request_id="req-missing", prompt={"prompt": "draw a cat"}))
+
+
+def test_parse_request_rejects_hidden_state_token_count_mismatch() -> None:
+    batch = _batch()
+    batch.prompts[0]["additional_information"]["full_hidden_states"] = torch.zeros(3, 8)
+    with pytest.raises(ValueError, match="hidden-state/token-count mismatch"):
+        _pipeline_shell()._parse_request(batch)
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "message"),
+    [
+        (0, 32, "Invalid image size"),
+        (32, -1, "Invalid image size"),
+        (30, 32, "multiples of 16"),
+        (32, 31, "multiples of 16"),
+    ],
+)
+def test_parse_request_rejects_invalid_dimensions(height, width, message) -> None:
+    batch = _batch()
+    batch.prompts[0].update(height=height, width=width)
+    with pytest.raises(ValueError, match=message):
+        _pipeline_shell()._parse_request(batch)
+
+
+def test_parse_request_rejects_explicit_zero_steps() -> None:
+    with pytest.raises(ValueError, match="num_inference_steps must be positive"):
+        _pipeline_shell()._parse_request(_batch(sampling=OmniDiffusionSamplingParams(num_inference_steps=0)))
+
+
+@pytest.mark.parametrize(("standard_steps", "expected_steps"), [(7, 7), (None, 50)])
+def test_parse_request_falls_through_null_legacy_step_count(standard_steps, expected_steps) -> None:
+    sampling = OmniDiffusionSamplingParams(num_inference_steps=standard_steps, extra_args={"num_inference_steps": None})
+    parsed = _pipeline_shell()._parse_request(_batch(sampling=sampling))
+    assert parsed.num_inference_steps == expected_steps
+
+
+@pytest.mark.parametrize(("standard_guidance", "expected_guidance"), [(4.0, 4.0), (None, 9.0)])
+def test_parse_request_falls_through_null_legacy_guidance(standard_guidance, expected_guidance) -> None:
+    sampling = OmniDiffusionSamplingParams(guidance_scale=standard_guidance, extra_args={"text_guidance_scale": None})
+    parsed = _pipeline_shell()._parse_request(_batch(sampling=sampling))
+    assert parsed.text_guidance_scale == expected_guidance
+
+
+def test_parse_request_defaults_null_cfg_range() -> None:
+    sampling = OmniDiffusionSamplingParams(extra_args={"cfg_range": None})
+    parsed = _pipeline_shell()._parse_request(_batch(sampling=sampling))
+    assert parsed.cfg_range == (0.0, 1.0)
+
+
+@pytest.mark.parametrize("generators", [[], [torch.Generator(), torch.Generator()]])
+def test_parse_request_rejects_invalid_generator_list_cardinality(generators) -> None:
+    sampling = OmniDiffusionSamplingParams(generator=generators)
+    with pytest.raises(ValueError, match="exactly one generator.*req-generators"):
+        _pipeline_shell()._parse_request(_batch(request_id="req-generators", sampling=sampling))
+
+
+def test_parse_request_synthesizes_dummy_ar_conditions() -> None:
+    batch = _batch(
+        request_id="dummy_req_id",
+        prompt={"prompt": "dummy run"},
+        sampling=OmniDiffusionSamplingParams(height=512, width=512, seed=1, guidance_scale=0.0, num_inference_steps=2),
+    )
+    parsed = _pipeline_shell()._parse_request(batch)
+    assert parsed.full_hidden_states.shape == (2, 8)
+    assert parsed.full_token_ids == [0, 100]
+    assert parsed.answer_start_index == 1
+
+
+@dataclass
+class _FakeTransformerConfig:
+    in_channels: int = 4
+
+
+@dataclass
+class _FakeTimeCaptionEmbed:
+    image_embedder: nn.Module | None = None
+
+
+class _FakeTransformer(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.anchor = nn.Parameter(torch.zeros(1))
+        self.config = _FakeTransformerConfig()
+        self.time_caption_embed = _FakeTimeCaptionEmbed()
+        self.calls = 0
+
+    def forward(self, *, hidden_states, **kwargs):
+        self.calls += 1
+        return torch.zeros_like(hidden_states)
+
+
+@dataclass
+class _FakeVaeConfig:
+    scaling_factor: float | None = None
+    shift_factor: float | None = None
+
+
+class _FakeVae(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.anchor = nn.Parameter(torch.zeros(1))
+        self.config = _FakeVaeConfig()
+
+    def decode(self, latents, return_dict=False):
+        assert return_dict is False
+        return (torch.zeros(1, 3, 32, 48, dtype=latents.dtype),)
+
+
+class _FakeScheduler:
+    def __init__(self) -> None:
+        self.timesteps = torch.tensor([])
+        self.requested_steps = None
+
+    def set_timesteps(self, *, num_inference_steps, device, num_tokens):
+        self.requested_steps = num_inference_steps
+        self.timesteps = torch.arange(num_inference_steps, device=device, dtype=torch.float32)
+
+    def step(self, model_pred, timestep, latents, return_dict=False):
+        assert return_dict is False
+        return (latents - model_pred,)
+
+
+def test_forward_returns_diffusion_output_with_request_sampling(mocker) -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+    captured = {}
+
+    def fake_randn_tensor(shape, *, generator, device, dtype):
+        captured["shape"] = shape
+        captured["seed"] = generator.initial_seed()
+        return torch.zeros(shape, device=device, dtype=dtype)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    mocker.patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler)
+    mocker.patch(f"{module}.randn_tensor", side_effect=fake_randn_tensor)
+    result = pipeline.forward(
+        _batch(
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=42, guidance_scale=1.0, num_inference_steps=2
+            )
+        )
+    )
+
+    assert isinstance(result, DiffusionOutput)
+    assert result.output.shape == (1, 3, 32, 48)
+    assert captured["shape"] == (1, 4, 4, 6)
+    assert captured["seed"] == 42
+    assert scheduler.requested_steps == 2
+    assert pipeline.gen_transformer.calls == 2
+
+
+def test_forward_rejects_missing_visual_tokens_before_model_access() -> None:
+    prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(3, 8),
+            "full_token_ids": [10, 11, 12],
+            "answer_start_index": 2,
+        },
+    }
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-empty"):
+        _pipeline_shell().forward(_batch(request_id="req-empty", prompt=prompt))

tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py

@@ -0,0 +1,679 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable, Iterable
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint import (
+    FASTH3_BASE_MODEL,
+    FASTH3_V2_BASE_SCHEDULE,
+    FASTH3_V2_MODEL_ID,
+    FastH3CheckpointSpec,
+)
+from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+    MiniMaxH3DiTArchConfig,
+    MiniMaxH3DiTModel,
+)
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+    resolve_minimax_h3_diffusion_model_path,
+)
+from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
+from vllm_omni.errors import OmniClientError
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
+    is_minimax_h3_modular,
+    resolve_minimax_h3_model_root,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _fastvideo_metadata() -> dict[str, object]:
+    """Return the release-owned metadata read from fastvideo_inference.json."""
+    return {
+        "model_id": FASTH3_V2_MODEL_ID,
+        "schema_version": "fasth3-inference-contract-v1",
+        "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
+        "transformer_forwards": 8,
+        "video_scheduler_shift": 10.0,
+        "audio_scheduler_shift": 3.0,
+        "guidance_scale": 1.0,
+        "attention_backend": "VIDEO_SPARSE_ATTN_H3",
+        "vsa_sparsity": 0.8,
+        "vsa_tile_size": 64,
+        "task": "t2av",
+    }
+
+
+def _sampling(
+    *,
+    num_inference_steps: int | str | None = 8,
+    guidance_scale: float | None = 1.0,
+    extra_args: dict[str, object] | None = None,
+    lora_request: LoRARequest | None = None,
+    timesteps: torch.Tensor | None = None,
+    sigmas: list[float] | None = None,
+) -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        # Keep the string form to exercise malformed client input validation.
+        num_inference_steps=num_inference_steps,  # type: ignore[arg-type]
+        guidance_scale=guidance_scale,
+        extra_args=extra_args or {},
+        lora_request=lora_request,
+        timesteps=timesteps,
+        sigmas=sigmas,
+    )
+
+
+def _od_config(
+    *,
+    backend: str = "FASTVIDEO_VSA",
+    topk: int | None = None,
+    per_role: dict[str, AttentionSpec] | None = None,
+    lora_path: str | None = None,
+    ring_degree: int = 1,
+    allgather_degree: int = 1,
+) -> OmniDiffusionConfig:
+    attention_config = AttentionConfig(
+        default=AttentionSpec(backend=backend, fastvideo_vsa_topk=topk),
+        per_role=per_role or {},
+    )
+    return OmniDiffusionConfig(
+        diffusion_attention_config=attention_config,
+        lora_path=lora_path,
+        parallel_config=DiffusionParallelConfig(ring_degree=ring_degree, allgather_degree=allgather_degree),
+    )
+
+
+def test_v2_release_has_nine_sigma_nodes_and_eight_intervals():
+    assert FASTH3_V2_BASE_SCHEDULE.base_schedule == (
+        0.999,
+        0.874,
+        0.749,
+        0.624,
+        0.5,
+        0.375,
+        0.25,
+        0.125,
+        0.0,
+    )
+    assert len(FASTH3_V2_BASE_SCHEDULE.base_schedule) == 9
+    assert FASTH3_V2_BASE_SCHEDULE.num_inference_steps == 8
+    assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(10.0) == pytest.approx(
+        [
+            0.9998999099189271,
+            0.985788405143244,
+            0.9675752486758817,
+            0.94316807738815,
+            0.9090909090909091,
+            0.8571428571428571,
+            0.7692307692307693,
+            0.5882352941176471,
+            0.0,
+        ],
+        abs=1e-12,
+    )
+    assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(3.0) == pytest.approx(
+        [
+            0.9996664442961973,
+            0.9541484716157204,
+            0.8995196156925539,
+            0.8327402135231315,
+            0.75,
+            0.6428571428571429,
+            0.5,
+            0.3,
+            0.0,
+        ],
+        abs=1e-12,
+    )
+
+
+def test_fastvideo_metadata_builds_release_sampling_metadata():
+    spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+
+    assert spec.vsa_sparsity == 0.8
+    assert spec.release_metadata() == {
+        "partition": "fl2va",
+        "tasks": ["t2va"],
+        "sigma_shift_scales": {"video": 10.0, "audio": 3.0},
+        "base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
+    }
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [
+        ("model_id", "FastVideo/FastH3-other"),
+        ("schema_version", "fasth3-inference-contract-v0"),
+        ("dmd_denoising_steps", [999, 749, 500, 250]),
+        ("transformer_forwards", 4),
+        ("video_scheduler_shift", 12.0),
+        ("audio_scheduler_shift", 4.0),
+        ("guidance_scale", 2.0),
+        ("attention_backend", "FLASH_ATTN"),
+        ("vsa_sparsity", 0.9),
+        ("vsa_tile_size", 32),
+        ("task", "t2va"),
+    ],
+)
+def test_fastvideo_metadata_rejects_contract_drift(field, value):
+    metadata = _fastvideo_metadata()
+    metadata[field] = value
+
+    with pytest.raises(ValueError, match=field):
+        FastH3CheckpointSpec.from_metadata(metadata)
+
+
+def test_fastvideo_metadata_rejects_missing_contract_field():
+    metadata = _fastvideo_metadata()
+    metadata.pop("task")
+
+    with pytest.raises(ValueError, match="task"):
+        FastH3CheckpointSpec.from_metadata(metadata)
+
+
+def test_legacy_h3_metadata_is_not_treated_as_fastvideo_release():
+    with pytest.raises(ValueError, match="model_id"):
+        FastH3CheckpointSpec.from_metadata(
+            {
+                "schema_version": 1,
+                "partition": "fl2va",
+                "tasks": ["t2va", "fl2va"],
+                "sigma_shift_scales": {"video": 12.0, "audio": 3.0},
+            }
+        )
+
+
+def test_v2_serving_contract_accepts_backend_selected_for_h3_self_role():
+    config = _od_config(
+        backend="TORCH_SDPA",
+        per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA")},
+    )
+    FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).check_serving_contract(
+        partition="fl2va", od_config=config
+    )
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "partition", "match"),
+    [
+        ({}, "ref2va", "task-type fl2va"),
+        ({"lora_path": "/tmp/adapter"}, "fl2va", "LoRA"),
+        ({"backend": "TORCH_SDPA"}, "fl2va", "FASTVIDEO_VSA"),
+        ({"topk": 2}, "fl2va", "fixed fastvideo_vsa_topk"),
+        ({"ring_degree": 2}, "fl2va", "local attention"),
+        ({"allgather_degree": 2}, "fl2va", "local attention"),
+    ],
+)
+def test_v2_serving_contract_rejects_incompatible_runtime(kwargs, partition, match):
+    spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+    with pytest.raises(ValueError, match=match):
+        spec.check_serving_contract(partition=partition, od_config=_od_config(**kwargs))
+
+
+def test_v2_serving_contract_reads_a_per_role_fixed_topk_override():
+    spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+    config = _od_config(
+        per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=1)},
+    )
+    with pytest.raises(ValueError, match="fixed fastvideo_vsa_topk"):
+        spec.check_serving_contract(partition="fl2va", od_config=config)
+
+
+def test_v2_request_contract_accepts_omitted_steps_and_default_shifts():
+    spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+    spec.check_request(_sampling(num_inference_steps=None))
+    spec.check_request(_sampling(extra_args={"flow_shift": 10, "audio_flow_shift": 3}, guidance_scale=1))
+
+
+@pytest.mark.parametrize(
+    "sampling",
+    [
+        _sampling(num_inference_steps=None),
+        _sampling(timesteps=torch.tensor([0.9, 0.5, 0.0])),
+        _sampling(sigmas=[1.0, 0.5, 0.0]),
+    ],
+)
+def test_v2_step_execution_requires_pinned_count_and_schedule(sampling):
+    spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+    with pytest.raises(OmniClientError, match="step execution"):
+        spec.check_request(sampling, step_execution=True)
+
+
+def test_v2_step_execution_accepts_only_explicit_pinned_count():
+    spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+    spec.check_request(_sampling(num_inference_steps=8), step_execution=True)
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "match"),
+    [
+        ({"num_inference_steps": 7}, "num_inference_steps=8"),
+        ({"num_inference_steps": 9}, "num_inference_steps=8"),
+        ({"num_inference_steps": "8"}, "num_inference_steps=8"),
+        ({"extra_args": {"flow_shift": 9}}, "flow_shift=10"),
+        ({"extra_args": {"audio_flow_shift": 4}}, "audio_flow_shift=3"),
+        ({"extra_args": {"flow_shift": "bad"}}, "flow_shift=10"),
+        ({"guidance_scale": 2.0}, "guidance_scale=1"),
+        (
+            {"lora_request": LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/test.safetensors")},
+            "per-request LoRA",
+        ),
+    ],
+)
+def test_v2_request_contract_rejects_sampling_drift(kwargs, match):
+    spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+    with pytest.raises(OmniClientError, match=match):
+        spec.check_request(_sampling(**kwargs))
+
+
+def test_resolve_native_vaes_uses_pinned_revision_and_only_vae_patterns(tmp_path, monkeypatch):
+    model_root = tmp_path / "FastH3"
+    model_root.mkdir()
+    revision = "base-revision-123"
+    (model_root / "provenance.json").write_text(
+        json.dumps({"base_model": f"hf://{FASTH3_BASE_MODEL}@{revision}"}),
+        encoding="utf-8",
+    )
+    snapshot = tmp_path / "hf-snapshot"
+    calls: dict[str, object] = {}
+
+    def fake_download(**kwargs: object) -> str:
+        calls.update(kwargs)
+        snapshot.mkdir()
+        return str(snapshot)
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint.download_weights_from_hf_specific",
+        fake_download,
+    )
+
+    result = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
+
+    assert result == snapshot / "FL2VA"
+    assert calls == {
+        "model_name_or_path": FASTH3_BASE_MODEL,
+        "cache_dir": None,
+        "allow_patterns": ["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
+        "revision": revision,
+        "require_all": True,
+    }
+    assert sorted(path.name for path in model_root.iterdir()) == ["provenance.json"]
+    assert not (model_root / "model_index.json").exists()
+    assert not (model_root / "conversion.json").exists()
+
+
+@pytest.mark.parametrize(
+    "base_model",
+    [
+        "MiniMaxAI/MiniMax-H3@main",
+        "hf://MiniMaxAI/MiniMax-H3@",
+        "hf://Other/MiniMax-H3@revision",
+        "",
+        None,
+    ],
+)
+def test_resolve_native_vaes_rejects_unpinned_provenance(tmp_path, base_model):
+    model_root = tmp_path / "FastH3"
+    model_root.mkdir()
+    (model_root / "provenance.json").write_text(json.dumps({"base_model": base_model}), encoding="utf-8")
+
+    with pytest.raises(ValueError, match="provenance"):
+        FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
+
+
+def test_modular_local_model_root_uses_text_encoder_without_partition_files(tmp_path):
+    root = tmp_path / "FastH3"
+    (root / "text_encoder").mkdir(parents=True)
+    (root / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
+    (root / "modular_model_index.json").write_text(
+        json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
+        encoding="utf-8",
+    )
+
+    assert is_minimax_h3_modular(str(root)) is True
+    assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "text_encoder")
+    assert not (root / "FL2VA").exists()
+
+
+def test_modular_hub_model_root_downloads_only_text_encoder(monkeypatch, tmp_path):
+    from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
+
+    snapshot = tmp_path / "modular-snapshot"
+    calls: dict[str, object] = {}
+
+    def fake_index(model: str, *, revision: str | None = None) -> dict[str, str]:
+        assert model == "FastVideo/FastVideo-FastH3-8-Step-V2"
+        assert revision == "release-revision"
+        return {"_class_name": "MiniMaxH3ModularPipeline"}
+
+    def fake_download(**kwargs: object) -> str:
+        calls.update(kwargs)
+        snapshot.mkdir()
+        return str(snapshot)
+
+    monkeypatch.setattr(checkpoint_module, "get_diffusion_model_index", fake_index)
+    monkeypatch.setattr(checkpoint_module, "download_weights_from_hf_specific", fake_download)
+
+    result = resolve_minimax_h3_model_root(
+        "FastVideo/FastVideo-FastH3-8-Step-V2",
+        "release-revision",
+        "fl2va",
+    )
+
+    assert result == str(snapshot / "text_encoder")
+    assert calls == {
+        "model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
+        "cache_dir": None,
+        "allow_patterns": ["text_encoder/**"],
+        "revision": "release-revision",
+        "require_all": True,
+    }
+
+
+def test_native_hub_model_with_modular_class_keeps_partition_layout(monkeypatch):
+    from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
+
+    monkeypatch.setattr(
+        checkpoint_module,
+        "get_diffusion_model_index",
+        lambda _model, *, revision=None: {"_class_name": "MiniMaxH3ModularPipeline"},
+    )
+
+    assert checkpoint_module.is_minimax_h3_modular("MiniMaxAI/MiniMax-H3") is False
+
+
+def test_native_local_model_root_keeps_partition_layout(tmp_path):
+    root = tmp_path / "MiniMax-H3"
+    (root / "FL2VA" / "text_encoder").mkdir(parents=True)
+    (root / "FL2VA" / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
+
+    assert is_minimax_h3_modular(str(root)) is False
+    assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
+
+
+def test_diffusion_resolver_returns_local_modular_root_without_partition_files(tmp_path):
+    root = tmp_path / "FastH3"
+    root.mkdir()
+    (root / "modular_model_index.json").write_text(
+        json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
+        encoding="utf-8",
+    )
+
+    assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(root)
+    assert not (root / "FL2VA").exists()
+    assert not (root / "Ref2VA").exists()
+
+
+def test_diffusion_resolver_hub_modular_downloads_transformer_contract_only(monkeypatch, tmp_path):
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+    snapshot = tmp_path / "modular-snapshot"
+    calls: dict[str, object] = {}
+
+    def fake_modular(_model: str, _revision: str | None = None) -> bool:
+        return True
+
+    def fake_download(**kwargs: object) -> str:
+        calls.update(kwargs)
+        snapshot.mkdir()
+        return str(snapshot)
+
+    monkeypatch.setattr(pipeline_module, "is_minimax_h3_modular", fake_modular)
+    monkeypatch.setattr(pipeline_module, "download_weights_from_hf_specific", fake_download)
+
+    result = resolve_minimax_h3_diffusion_model_path(
+        "FastVideo/FastVideo-FastH3-8-Step-V2", "release-revision", "fl2va"
+    )
+
+    assert result == str(snapshot)
+    assert calls == {
+        "model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
+        "cache_dir": None,
+        "allow_patterns": [
+            "modular_model_index.json",
+            "fastvideo_inference.json",
+            "provenance.json",
+            "transformer/**",
+        ],
+        "revision": "release-revision",
+        "require_all": True,
+    }
+    patterns = calls["allow_patterns"]
+    assert isinstance(patterns, list)
+    assert all("video_vae" not in pattern and "audio_vae" not in pattern for pattern in patterns)
+
+
+def test_diffusers_arch_config_normalizes_all_aliases():
+    arch = MiniMaxH3DiTArchConfig.from_mapping(
+        {
+            "num_refiner_layers": 3,
+            "ffn_dim": 12,
+            "in_channels": 4,
+            "audio_in_channels": 6,
+            "freq_dim": 8,
+            "time_embed_hidden_dim": 10,
+            "rope_freq_dim": 5,
+            "patch_size": [1, 2, 2],
+        }
+    )
+
+    assert arch.token_refiner_num_layers == 3
+    assert arch.ffn_hidden_size == 12
+    assert arch.latents_dim == 4
+    assert arch.audio_latents_dim == 6
+    assert arch.timestep_input_dim == 8
+    assert arch.time_embed_hidden_size == 10
+    assert arch.rope_inv_freq_len == 5
+    assert arch.patch_size == (1, 2, 2)
+
+
+class _WeightTarget(nn.Module):
+    def __init__(self, shape: tuple[int, ...], loader: Callable[..., None]) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(shape))
+        self.weight.weight_loader = loader
+
+
+class _PipelineLoadTransformer(nn.Module):
+    def __init__(self, loaded: set[str]) -> None:
+        super().__init__()
+        self.arch = MiniMaxH3DiTArchConfig(num_layers=1)
+        self._loaded = loaded
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        list(weights)
+        return self._loaded
+
+    def post_load_weights(self) -> None:
+        pass
+
+
+def test_pipeline_load_weights_rejects_missing_fast_h3_compression_gate():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.transformer = _PipelineLoadTransformer({"blocks.0.attn.qkv_proj.weight"})
+    pipeline.video_vae = None
+    pipeline.audio_vae = None
+    pipeline._fasth3 = None
+    pipeline._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+
+    with pytest.raises(ValueError, match="missing compression gates"):
+        pipeline.load_weights([("transformer.blocks.0.attn.qkv_proj.weight", torch.zeros(1))])
+
+
+def _load_fixture(
+    *, diffusers_weights: bool
+) -> tuple[MiniMaxH3DiTModel, list[tuple[str | None, torch.Tensor]], list[tuple[int, torch.Tensor]]]:
+    model = object.__new__(MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.arch = MiniMaxH3DiTArchConfig(
+        num_layers=1,
+        hidden_size=3,
+        num_attention_heads=2,
+        attention_head_dim=2,
+        ffn_hidden_size=6,
+        rope_inv_freq_len=2,
+    )
+    model._diffusers_weights = diffusers_weights
+    model._rope_theta = 10000.0
+    qkv_calls: list[tuple[str | None, torch.Tensor]] = []
+    fc1_calls: list[tuple[int, torch.Tensor]] = []
+
+    def qkv_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str | None = None) -> None:
+        if shard_id is None:
+            assert loaded_weight.shape == (12, 3)
+            param.data.copy_(loaded_weight)
+            qkv_calls.append((None, loaded_weight.clone()))
+            return
+        assert loaded_weight.shape == (4, 3)
+        assert shard_id in {"q", "k", "v"}
+        offset = {"q": 0, "k": 4, "v": 8}[shard_id]
+        param.data[offset : offset + 4].copy_(loaded_weight)
+        qkv_calls.append((shard_id, loaded_weight.clone()))
+
+    def fc1_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: int) -> None:
+        assert loaded_weight.shape == (3, 3)
+        param.data[shard_id * 3 : (shard_id + 1) * 3].copy_(loaded_weight)
+        fc1_calls.append((shard_id, loaded_weight.clone()))
+
+    model.blocks = nn.ModuleList([nn.Module()])
+    model.blocks[0].attn = nn.Module()
+    model.blocks[0].attn.qkv_proj = _WeightTarget((12, 3), qkv_loader)
+    model.blocks[0].mlp = nn.Module()
+    model.blocks[0].mlp.fc1 = _WeightTarget((6, 3), fc1_loader)
+    model.rope = nn.Module()
+    model.rope.register_buffer("inv_freq", torch.zeros(2))
+    return model, qkv_calls, fc1_calls
+
+
+def test_diffusers_load_weights_uses_shape_aware_qkv_and_mlp_slices():
+    model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=True)
+    q = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 10
+    k = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100
+    v = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 200
+    mlp = torch.arange(18, dtype=torch.float32).reshape(6, 3) + 300
+
+    loaded = model.load_weights(
+        [
+            ("transformer_blocks.0.attn.to_q.weight", q),
+            ("transformer_blocks.0.attn.to_k.weight", k),
+            ("transformer_blocks.0.attn.to_v.weight", v),
+            ("transformer_blocks.0.ff.net.0.proj.weight", mlp),
+        ]
+    )
+
+    assert loaded == {
+        "blocks.0.attn.qkv_proj.weight",
+        "blocks.0.mlp.fc1.weight",
+        "rope.inv_freq",
+    }
+    assert [shard_id for shard_id, _ in qkv_calls] == ["q", "k", "v"]
+    assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
+    torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, torch.cat((q, k, v)))
+    torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, torch.cat((mlp[3:], mlp[:3])))
+    torch.testing.assert_close(
+        model.rope.inv_freq,
+        1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)),
+        rtol=0,
+        atol=0,
+    )
+
+
+def test_native_load_weights_keeps_grouped_qkv_and_gate_first_layouts():
+    model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=False)
+    grouped_qkv = torch.arange(36, dtype=torch.float32).reshape(12, 3)
+    fc1 = torch.arange(18, dtype=torch.float32).reshape(6, 3)
+
+    loaded = model.load_weights(
+        [
+            ("blocks.0.attn.qkv_proj.weight", grouped_qkv),
+            ("blocks.0.mlp.fc1.weight", fc1),
+        ]
+    )
+
+    expected_qkv = torch.cat(
+        (
+            grouped_qkv[:2],
+            grouped_qkv[6:8],
+            grouped_qkv[2:4],
+            grouped_qkv[8:10],
+            grouped_qkv[4:6],
+            grouped_qkv[10:],
+        )
+    )
+    assert loaded == {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight"}
+    assert [shard_id for shard_id, _ in qkv_calls] == [None]
+    assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
+    torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, expected_qkv)
+    torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, fc1)
+    assert torch.equal(model.rope.inv_freq, torch.zeros(2))
+
+
+def test_diffusers_load_weights_rejects_unknown_name():
+    model, _, _ = _load_fixture(diffusers_weights=True)
+
+    with pytest.raises(ValueError, match="unsupported Diffusers H3 weight"):
+        model.load_weights([("transformer_blocks.0.attn.not_a_projection.weight", torch.zeros(4, 3))])
+
+
+def test_diffusers_load_weights_rejects_duplicate_source_name():
+    model, _, _ = _load_fixture(diffusers_weights=True)
+    weight = torch.zeros(4, 3)
+
+    with pytest.raises(ValueError, match="duplicate Diffusers H3 weight"):
+        model.load_weights(
+            [
+                ("transformer_blocks.0.attn.to_q.weight", weight),
+                ("transformer_blocks.0.attn.to_q.weight", weight),
+            ]
+        )
+
+
+def test_diffusers_load_weights_rejects_incomplete_qkv_group():
+    model, _, _ = _load_fixture(diffusers_weights=True)
+
+    with pytest.raises(ValueError, match="incomplete Diffusers H3 QKV"):
+        model.load_weights(
+            [
+                ("transformer_blocks.0.attn.to_q.weight", torch.zeros(4, 3)),
+                ("transformer_blocks.0.attn.to_k.weight", torch.zeros(4, 3)),
+            ]
+        )
+
+
+def test_diffusers_load_weights_rejects_missing_compression_gate_parameter():
+    model, _, _ = _load_fixture(diffusers_weights=True)
+
+    with pytest.raises(ValueError, match="has no model parameter"):
+        model.load_weights([("transformer_blocks.0.attn.to_gate_compress.weight", torch.zeros(4, 3))])
+
+
+def test_resolve_sigma_positions_uses_release_metadata_schedule():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline._fasth3 = None
+    pipeline._lora_sigma_schedules = {}
+    spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
+    pipeline._base_schedule_by_partition = {
+        "fl2va": DMD2SigmaSchedule.from_metadata(spec.release_metadata()),
+    }
+
+    positions, steps = pipeline._resolve_sigma_positions("t2va", _sampling())
+    assert positions == FASTH3_V2_BASE_SCHEDULE.base_schedule
+    assert steps == 8 == len(positions) - 1
+    with pytest.raises(OmniClientError, match="must be 8"):
+        pipeline._resolve_sigma_positions("t2va", _sampling(num_inference_steps=9))

tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+import time
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+def _collective_worker(rank, rendezvous, backend):
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+    from vllm_omni.errors import OmniClientError
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
+    from vllm_omni.platforms import current_omni_platform
+
+    torch.set_num_threads(1)
+    device = torch.device("cuda", rank) if backend == "nccl" else torch.device("cpu")
+    if backend == "nccl":
+        current_omni_platform.set_device(device)
+    dist.init_process_group(
+        backend, init_method=f"file://{rendezvous}", rank=rank, world_size=2, timeout=timedelta(seconds=45)
+    )
+    try:
+        module._dit_rank_world = lambda: (dist.group.WORLD, rank, 2)
+
+        class VisualCodec:
+            fail = False
+
+            def is_distributed_enabled(self):
+                return True
+
+            def collective(self):
+                value = torch.ones(1, device=device)
+                dist.all_reduce(value)
+                assert value.item() == 2
+
+            def encode_image(self, image):
+                self.collective()
+                assert image.size == (32, 32)
+                return torch.ones(1, 96)
+
+            def encode_video(self, frames):
+                self.collective()
+                assert frames.shape == (4, 32, 32, 3)
+                if self.fail and rank == 1:
+                    raise ValueError("codec failed on rank 1")
+                return torch.full((1, 96), 2.0), (1, 2, 2)
+
+        class AudioCodec:
+            def encode_waveform(self, waveform, sample_rate):
+                assert rank == 0
+                assert sample_rate == 32000
+                assert waveform.shape == (64000,)
+                return torch.full((160, 32), 3.0), 80
+
+        model = object.__new__(module.MiniMaxH3Pipeline)
+        torch.nn.Module.__init__(model)
+        model.device = device
+        model.od_config = SimpleNamespace(enable_layerwise_offload=False)
+        model.video_vae = VisualCodec()
+        model.audio_vae = AudioCodec()
+        model.partition = "combined"
+        model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+        model.text_encoder_tp_size = 2
+        model.text_encoder_group = SimpleNamespace(
+            world_size=2, ranks=[0, 1], rank_in_group=rank, device_group=dist.group.WORLD
+        )
+        model.processor = SimpleNamespace(image_processor=SimpleNamespace(merge_size=2))
+        model.tokenizer = lambda text, **kwargs: {"input_ids": [1, 2, 3]}
+
+        def encode_text(ids, kwargs):
+            assert ids.tolist() == [1, 2, 3]
+            assert not kwargs
+            hidden = torch.ones(3, 5120, device=device, dtype=torch.bfloat16)
+            dist.all_reduce(hidden)
+            return hidden
+
+        model._encode_text_hidden = encode_text
+        sampling = OmniDiffusionSamplingParams(
+            height=32, width=32, num_frames=96, extra_args={"task": "t2va", "aspect_ratio": "1:1"}
+        )
+        text_conditioning, window_text = model._prepare_local_conditioning({"prompt": "test"}, sampling)
+        assert window_text is None
+        torch.testing.assert_close(text_conditioning.hidden_states, torch.full_like(text_conditioning.hidden_states, 2))
+        assert text_conditioning.visual_condition is None
+
+        media = (
+            MiniMaxH3EncoderMediaInput(
+                task="ref2va",
+                height=32,
+                width=32,
+                num_frames=107,
+                latent_t=32,
+                audio_t=178,
+                images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
+                videos=(torch.zeros(4, 32, 32, 3, dtype=torch.uint8),),
+                video_audios=(None,),
+                audios=((torch.zeros(64000), 32000),),
+            )
+            if rank == 0
+            else None
+        )
+        conditioning = model._encode_local_media(media)
+        assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video", "audio"]
+        torch.testing.assert_close(
+            conditioning.visual_condition.cpu(), torch.cat([torch.ones(1, 96), torch.full((1, 96), 2.0)])
+        )
+        torch.testing.assert_close(conditioning.audio_condition.cpu(), torch.full((160, 32), 3.0))
+        assert conditioning.audio_condition_lengths == (80,)
+
+        model.video_vae.fail = True
+        with pytest.raises(OmniClientError, match="codec failed on rank 1"):
+            model._encode_local_media(media)
+        with pytest.raises(OmniClientError, match="non-empty prompt"):
+            model._prepare_local_conditioning({"prompt": ""}, sampling)
+        model.video_vae.fail = False
+        assert model._encode_local_media(media).visual_condition_shapes == ((1, 2, 2), (1, 2, 2))
+    finally:
+        dist.destroy_process_group()
+
+
+def _run_collectives(tmp_path, backend):
+    context = mp.spawn(_collective_worker, args=(str(tmp_path / "rendezvous"), backend), nprocs=2, join=False)
+    deadline = time.monotonic() + 150
+    try:
+        while not context.join(timeout=1):
+            if time.monotonic() > deadline:
+                pytest.fail("MiniMax H3 collectives did not finish within 150 seconds")
+    finally:
+        for process in context.processes:
+            if process.is_alive():
+                process.terminate()
+            process.join(timeout=5)
+
+
+@pytest.mark.cpu
+def test_two_rank_conditioning_and_errors_gloo(tmp_path):
+    _run_collectives(tmp_path, "gloo")
+
+
+@pytest.mark.cuda
+@pytest.mark.parallel
+@pytest.mark.skipif(
+    os.environ.get("VLLM_TEST_MINIMAX_H3_COLLECTIVES") != "1",
+    reason="set VLLM_TEST_MINIMAX_H3_COLLECTIVES=1 with two reserved GPUs",
+)
+def test_two_rank_conditioning_and_errors_nccl(tmp_path):
+    if torch.accelerator.device_count() < 2:
+        pytest.skip("requires two GPUs")
+    _run_collectives(tmp_path, "nccl")

tests/diffusion/models/minimax_h3/test_minimax_h3_continuation.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.minimax_h3.continuation import (
+    diffuse_continuation,
+    plan_continuation_windows,
+    resolve_continuation,
+)
+from vllm_omni.diffusion.models.minimax_h3.packed_sequence import minimax_h3_packed_sequence_ref2va_blocks
+from vllm_omni.errors import OmniClientError
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("vary_prompts", [False, True])
+@pytest.mark.parametrize("total,window", [(1807, 277), (1807, 243), (124, 277)])
+def test_continuation_retains_prefix_and_global_stereo_timeline(total, window, vary_prompts):
+    audio_t = round(total * 40 / 24)
+    source = torch.arange(2 * audio_t * 32).reshape(2, audio_t, 32).float()
+    calls = []
+    plan = plan_continuation_windows(total, window, 22)
+    texts = [(torch.full((3 + i, 8), float(i)), torch.ones(3 + i, dtype=torch.long)) for i in range(len(plan))]
+
+    def sample(**kw):
+        calls.append(kw)
+        t = kw["latent_t"]
+        if vary_prompts:
+            assert kw["text_embeddings"] is texts[len(calls) - 1][0]
+            assert kw["text_tags"] is texts[len(calls) - 1][1]
+            assert kw["media_time_origin"] == len(texts[-1][0])
+        # Different constants expose replacement of an old prefix or an
+        # accidental append of the hidden overlap.
+        video = torch.full((1, 24, t, 2, 2), float(len(calls)))
+        audio = kw["locked_audio_rows"].reshape(2, kw["audio_t"], 32).permute(0, 2, 1)
+        if len(calls) > 1:
+            assert torch.all(kw["visual_condition"] == len(calls) - 1)
+            assert kw["ref_blocks"][-1]["kind"] == "latent_guide"
+        return video, audio
+
+    kwargs = dict(
+        num_frames=total,
+        latent_t=(total - 5) // 17 * 5 + 2,
+        latent_h=2,
+        latent_w=2,
+        audio_t=audio_t,
+        locked_audio_rows=source.reshape(-1, 32),
+    )
+    video, audio = diffuse_continuation(
+        sample, kwargs, window_frames=window, overlap_frames=22, text_conditioning=texts if vary_prompts else None
+    )
+    torch.testing.assert_close(audio, source.permute(0, 2, 1))
+    plan = plan_continuation_windows(total, window, 22)
+    assert [call["temporal_offset"] for call in calls] == pytest.approx([part.start * 40 / 24 for part in plan])
+    offset = 0
+    for i, part in enumerate(plan):
+        new_frames = part.end - part.start - part.overlap
+        count = (new_frames - 5) // 17 * 5 + 2 if i == 0 else new_frames // 17 * 5
+        assert torch.all(video[:, :, offset : offset + count] == i + 1)
+        offset += count
+    assert offset == video.shape[2] == kwargs["latent_t"]
+
+
+@pytest.mark.parametrize("temporal_offset", [0.0, 425.0, 221 * 40 / 24])
+def test_latent_guide_shares_target_clock_and_remains_condition_only(temporal_offset):
+    packed = minimax_h3_packed_sequence_ref2va_blocks(
+        temporal_offset=temporal_offset,
+        text_len=3,
+        latent_t=12,
+        latent_h=4,
+        latent_w=4,
+        audio_t=68,
+        ref_blocks=[
+            {"kind": "image", "latent_h": 4, "latent_w": 4},
+            {"kind": "latent_guide", "latent_t": 7, "latent_h": 4, "latent_w": 4, "ref_audio_t": 37},
+        ],
+    )
+    pos = packed["img_position_ids"]
+    visual = packed["img_pos"]
+    target = visual[packed["update_mask"]]
+    guide = visual[~packed["update_mask"]][4:]
+    torch.testing.assert_close(pos[guide], pos[target[:28]])
+    audio = packed["audio_pos"]
+    guide_audio = audio[~packed["audio_update_mask"]].reshape(2, 37)
+    target_audio = audio[packed["audio_update_mask"]].reshape(2, 68)
+    torch.testing.assert_close(pos[guide_audio], pos[target_audio[:, :37]])
+    assert pos[target[0], 0] == pytest.approx(4 + temporal_offset)
+    torch.testing.assert_close(pos[:3, 0], torch.arange(3, dtype=torch.float64))
+    assert torch.all(pos[3:7, 0] == 3)  # static reference image stays fixed
+
+
+@pytest.mark.parametrize(
+    "extra,step",
+    [
+        ({"long_video_mode": "bad"}, False),
+        ({"long_video_mode": "continuation"}, True),
+        ({"long_video_mode": "continuation", "continuation_overlap_frames": 23}, False),
+        ({"long_video_mode": "continuation", "continuation_window_frames": True}, False),
+        ({"long_video_mode": "continuation", "continuation_overlap_frames": 277}, False),
+    ],
+)
+def test_invalid_continuation_options(extra, step):
+    with pytest.raises(OmniClientError):
+        resolve_continuation(extra, task="ref2va", step_execution=step)
+
+
+def test_global_offset_moves_temporal_references_but_not_images_or_padding():
+    kwargs = dict(
+        text_len=3,
+        latent_t=12,
+        latent_h=4,
+        latent_w=4,
+        audio_t=68,
+        ref_blocks=[
+            {"kind": "image", "latent_h": 4, "latent_w": 4},
+            {"kind": "audio", "ref_audio_t": 3},
+            {"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 4},
+            {"kind": "image", "latent_h": 4, "latent_w": 4},
+            {"kind": "latent_guide", "ref_audio_t": 37, "latent_t": 7, "latent_h": 4, "latent_w": 4},
+        ],
+    )
+    base = minimax_h3_packed_sequence_ref2va_blocks(**kwargs)
+    shifted = minimax_h3_packed_sequence_ref2va_blocks(**kwargs, temporal_offset=425.0)
+    moving = base["image_mask"] | base["audio_mask"]
+    moving[3:7] = False  # first static image
+    moving[25:29] = False  # static image after temporal references
+    delta = shifted["img_position_ids"] - base["img_position_ids"]
+    torch.testing.assert_close(delta[:, 0], moving.double() * 425)
+    assert torch.count_nonzero(delta[:, 1:]) == 0
+    for key, value in base.items():
+        if isinstance(value, torch.Tensor) and key != "img_position_ids":
+            torch.testing.assert_close(shifted[key], value)
+
+
+def test_long_ref2va_defaults_to_continuation_with_explicit_full_opt_out():
+    assert resolve_continuation({"long_video": True}, task="ref2va", step_execution=False) == (277, 22)
+    assert resolve_continuation({}, task="ref2va", step_execution=False) is None
+    assert (
+        resolve_continuation({"long_video": True, "long_video_mode": "full"}, task="ref2va", step_execution=False)
+        is None
+    )
+
+
+def test_different_chunk_prompt_lengths_keep_the_same_media_clock():
+    kwargs = dict(
+        latent_t=7,
+        latent_h=4,
+        latent_w=4,
+        audio_t=37,
+        ref_blocks=[{"kind": "image", "latent_h": 4, "latent_w": 4}],
+        media_time_origin=20,
+    )
+    first = minimax_h3_packed_sequence_ref2va_blocks(text_len=3, **kwargs)
+    second = minimax_h3_packed_sequence_ref2va_blocks(text_len=15, temporal_offset=425, **kwargs)
+    for key, mask in (("img_pos", "update_mask"), ("audio_pos", "audio_update_mask")):
+        a = first["img_position_ids"][first[key][first[mask]]]
+        b = second["img_position_ids"][second[key][second[mask]]]
+        torch.testing.assert_close(b[:, 0] - a[:, 0], torch.full_like(a[:, 0], 425))
+        torch.testing.assert_close(b[:, 1:], a[:, 1:])
+
+
+@pytest.mark.parametrize("prompt_count", [None, 7, 6])
+def test_request_encodes_only_window_prompts_and_prepares_references_once(monkeypatch, prompt_count):
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaConditioning
+
+    pipeline = object.__new__(module.MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.load_text_encoder = True
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config = None
+    pipeline._active_turbo_spec = lambda sampling: None
+    pipeline._has_active_native_lora = lambda sampling: False
+    pipeline._resolve_task = lambda *args, **kwargs: "ref2va"
+    encoded, prepared_inputs, encoded_media = [], [], []
+    prepare = module.prepare_encoder_inputs
+
+    def prepare_once(*args, **kwargs):
+        result = prepare(*args, **kwargs)
+        prepared_inputs.append(result)
+        return result
+
+    def encode_text(prepared):
+        encoded.append(prepared)
+        return torch.full((3, 5120), len(encoded), dtype=torch.bfloat16), torch.zeros(3, dtype=torch.long)
+
+    def encode_media(media):
+        encoded_media.append(media)
+        return MiniMaxH3EncoderMediaConditioning(
+            task=media.task,
+            height=media.height,
+            width=media.width,
+            num_frames=media.num_frames,
+            latent_t=media.latent_t,
+            audio_t=media.audio_t,
+        )
+
+    monkeypatch.setattr(module, "prepare_encoder_inputs", prepare_once)
+    pipeline.encode_prompt = encode_text
+    pipeline._encode_local_media = encode_media
+    pipeline._prepare_encoder_conditioning_inputs = lambda conditioning, sampling: {
+        "text_embeddings": conditioning.hidden_states,
+        "continuation": (277, 22),
+        "num_frames": conditioning.num_frames,
+        "task": conditioning.task,
+    }
+    extra = {"task": "ref2va", "long_video": True, "duration": 75}
+    if prompt_count is not None:
+        extra["continuation_prompts"] = [f"window {i}" for i in range(prompt_count)]
+    sampling = OmniDiffusionSamplingParams(width=64, height=64, fps=24, extra_args=extra)
+    request = {"prompt": "main prompt", "multi_modal_data": {"image": Image.new("RGB", (256, 256))}}
+    if prompt_count == 6:
+        with pytest.raises(OmniClientError, match="exactly 7"):
+            pipeline._prepare_request_inputs(request, sampling)
+        assert not encoded and not encoded_media
+        return
+    context = pipeline._prepare_request_inputs(request, sampling)
+    assert [item.prompt for item in encoded] == extra.get("continuation_prompts", ["main prompt"])
+    assert len(prepared_inputs) == len(encoded_media) == 1
+    assert all(item.media is prepared_inputs[0].media for item in encoded)
+    if prompt_count is not None:
+        texts = context["continuation_text_conditioning"]
+        assert len(texts) == 7
+        assert texts[0][0] is context["text_embeddings"]
+        assert [int(hidden[0, 0]) for hidden, _ in texts] == list(range(1, 8))
+    else:
+        assert "continuation_text_conditioning" not in context

tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import Mock
+
+import pytest
+import torch
+from PIL import Image
+
+from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture
+def pipeline():
+    model = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(model)
+    model.load_text_encoder = True
+    model.load_vae_encoder = True
+    model.device = torch.device("cpu")
+    model.partition = "combined"
+    model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    model.default_video_shift = 12.0
+    model.default_audio_shift = 3.0
+    model.od_config = SimpleNamespace(enable_layerwise_offload=False)
+    model._quality_policy = SimpleNamespace(resolve=Mock(return_value=SimpleNamespace(cache_dit=None)))
+    model._cache_dit_runtime = SimpleNamespace(prepare=Mock())
+    model.encode_prompt = Mock(
+        return_value=(torch.ones(2, 5120, dtype=torch.bfloat16), torch.ones(2, dtype=torch.long))
+    )
+    model.video_vae = Mock()
+    model.video_vae.is_distributed_enabled.return_value = False
+    model.video_vae.encode_image.side_effect = lambda image: torch.ones((image.height // 32) * (image.width // 32), 96)
+    model.audio_vae = Mock()
+    model.transformer = torch.nn.Identity()
+    model.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+    model.decode = Mock(return_value=(torch.zeros(1, 3, 1, 32, 32), torch.zeros(1)))
+    return model
+
+
+def _request(task):
+    prompt: dict[str, Any] = {"prompt": "A person waves."}
+    if task != "t2va":
+        prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))}
+    sampling = OmniDiffusionSamplingParams(
+        height=32, width=32, num_frames=96, num_inference_steps=2, extra_args={"task": task, "aspect_ratio": "1:1"}
+    )
+    return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
+
+
+@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"])
+def test_single_stage_encodes_without_stage_zero(pipeline, task):
+    output = pipeline.forward(_request(task))
+    pipeline.encode_prompt.assert_called_once()
+    assert pipeline.video_vae.encode_image.call_count == int(task != "t2va")
+    pipeline.diffuse.assert_called_once()
+    assert pipeline.diffuse.call_args.kwargs["task"] == task
+    video, audio = output.output
+    torch.testing.assert_close(video, torch.zeros(1, 1, 32, 32, 3, dtype=torch.uint8))
+    torch.testing.assert_close(audio, pipeline.decode.return_value[1])
+
+
+def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch):
+    import numpy as np
+
+    from vllm_omni.data_entry_keys import flatten_payload
+    from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import (
+        MiniMaxH3EncoderConditioning,
+        MiniMaxH3TextConditioning,
+    )
+    from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
+    from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion, prepare_encoder_prompt
+
+    frames = np.zeros((48, 32, 32, 3), dtype=np.uint8)
+    decode_frames = Mock(return_value=frames)
+    sampler = Mock(wraps=processing.sample_reference_video_frames)
+    monkeypatch.setattr(processing, "load_video_frames", decode_frames)
+    monkeypatch.setattr(processing, "sample_reference_video_frames", sampler)
+    monkeypatch.setattr(
+        processing,
+        "prepare_reference_videos",
+        Mock(
+            return_value=[{"prepared_path": "prepared.mp4", "original_path": "original.mp4", "input_has_audio": True}]
+        ),
+    )
+    monkeypatch.setattr(processing, "load_video_audio", Mock(return_value=(torch.full((64000,), 3.0), 32000)))
+    pipeline.video_vae.encode_video.return_value = (torch.full((1, 96), 2.0), (1, 2, 2))
+
+    def encode_audio(waveform, sample_rate):
+        length = round(waveform.shape[-1] / sample_rate * 40)
+        return torch.full((length * 2, 32), float(waveform[0])), length
+
+    pipeline.audio_vae.encode_waveform.side_effect = encode_audio
+    request = _request("ref2va")
+    request.prompts[0]["multi_modal_data"].update(video="original.mp4", audio=(torch.full((192000,), 4.0), 32000))
+    pipeline.forward(request)
+    local_kwargs = pipeline.diffuse.call_args.kwargs
+
+    transformed = prepare_encoder_prompt(request.prompts[0], [request.sampling_params])
+    stage_zero = object.__new__(MiniMaxH3Encoder)
+    torch.nn.Module.__init__(stage_zero)
+    stage_zero._component_leader = True
+    stage_zero.video_vae = pipeline.video_vae
+    stage_zero.audio_vae = pipeline.audio_vae
+    media = stage_zero._encode_media(stage_zero._media_input(transformed["additional_information"]))
+    text = MiniMaxH3TextConditioning(*pipeline.encode_prompt.return_value)
+    payload = MiniMaxH3EncoderConditioning.from_components(text, media).to_omni_payload()
+    source = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=flatten_payload(payload))])
+    bridged = encoder2diffusion([source], request.prompts[0])
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
+    pipeline.encode_prompt.reset_mock()
+    pipeline.forward(SimpleNamespace(prompts=[bridged], sampling_params=request.sampling_params))
+    external_kwargs = pipeline.diffuse.call_args.kwargs
+    for key, value in local_kwargs.items():
+        if isinstance(value, torch.Tensor):
+            torch.testing.assert_close(value, external_kwargs[key])
+        else:
+            assert value == external_kwargs[key]
+    pipeline.encode_prompt.assert_not_called()
+    assert decode_frames.call_count == sampler.call_count == 2
+    assert all(call.kwargs["decoded_frames"] is frames for call in sampler.call_args_list)
+    assert [block["kind"] for block in local_kwargs["ref_blocks"]] == ["image", "video_audio", "audio"]
+    assert local_kwargs["audio_condition_lengths"] == [80, round(107 / 24 * 40)]
+    assert all(
+        call.args[0].shape[-1] == round(107 / 24 * 32000)
+        for call in pipeline.audio_vae.encode_waveform.call_args_list[1::2]
+    )

tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py

@@ -262,7 +262,10 @@ def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
 
 def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path):
     path = tmp_path / "fasth3" / "adapter_model.safetensors"
-    _write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))})
+    _write_adapter(
+        path,
+        tensors={"transformer_blocks.0.attn.not_a_projection.lora_A.weight": torch.ones((_RANK, _HIDDEN))},
+    )
     # Dropping it silently would load a model that is not the distilled student.
     with pytest.raises(FastH3AdapterError, match="name no known"):
         _load(path.parent)

tests/diffusion/models/minimax_h3/test_minimax_h3_latent_edit_pipeline.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Pipeline-level MiniMax H3 latent initialization tests."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_prepare_edit_video_clones_the_last_frame_with_lossless_rgb(monkeypatch, tmp_path):
+    from vllm_omni.model_executor.models.minimax_h3 import reference_video
+
+    source = tmp_path / "source.webm"
+    commands: list[list[str]] = []
+    monkeypatch.setattr(
+        reference_video,
+        "_probe_video",
+        lambda _path: {"frame_count": 12, "audio_codecs": ("opus",)},
+    )
+    monkeypatch.setattr(
+        reference_video.subprocess,
+        "run",
+        lambda command, **_kwargs: commands.append(command),
+    )
+
+    prepared = reference_video.prepare_edit_video(
+        source,
+        target_width=448,
+        target_height=256,
+        target_frame_count=73,
+        workdir=str(tmp_path / "prepared"),
+    )
+
+    assert prepared["input_has_audio"] is True
+    assert prepared["frame_count"] == 73
+    command = commands[0]
+    assert command[command.index("-vf") + 1].endswith("tpad=stop_mode=clone:stop=-1")
+    assert command[command.index("-c:v") + 1] == "libx264rgb"
+    assert command[command.index("-pix_fmt") + 1] == "rgb24"
+
+
+@pytest.mark.parametrize("source_audio", [None, "override.wav"])
+def test_encoder_preparation_routes_edit_sources_without_qwen_inputs(
+    monkeypatch,
+    tmp_path,
+    source_audio,
+):
+    from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing
+
+    latent_t, audio_t = 7, 37
+    video_clean = torch.full((42, 96), 3.0)
+    audio_source = torch.zeros(4, 32)
+    audio_source[:, 0] = torch.tensor([1.0, 2.0, 11.0, 12.0])
+    waveform = torch.zeros(3_200)
+    frames = torch.zeros(22, 64, 96, 3, dtype=torch.uint8).numpy()
+    calls: dict[str, Any] = {}
+
+    def prepare_video(value, **kwargs):
+        calls["prepare"] = (value, kwargs)
+        return {
+            "original_path": "source.mov",
+            "prepared_path": str(tmp_path / "prepared.mp4"),
+            "input_has_audio": True,
+        }
+
+    def load_video_audio(value, **kwargs):
+        calls["video_audio"] = (value, kwargs)
+        return waveform, 32_000
+
+    def load_audio_file(value):
+        calls["audio_file"] = value
+        return waveform, 32_000
+
+    class VideoVAE:
+        def encode_video(self, value):
+            calls["video_encode_shape"] = value.shape
+            return video_clean.clone(), (latent_t, 4, 6)
+
+    class AudioVAE:
+        def encode_waveform(self, value, sample_rate):
+            calls["audio_encode"] = (value.shape, sample_rate)
+            return audio_source.clone(), 2
+
+    monkeypatch.setattr(processing, "resolve_minimax_h3_shape", lambda *_args: (64, 96, 22, latent_t, audio_t))
+    monkeypatch.setattr(processing, "prepare_edit_video", prepare_video)
+    monkeypatch.setattr(processing, "load_video_frames", lambda _path: frames)
+    monkeypatch.setattr(processing, "load_video_audio", load_video_audio)
+    monkeypatch.setattr(processing, "load_audio_file", load_audio_file)
+
+    prepared = processing.prepare_encoder_inputs(
+        {
+            "prompt": "edit the clip",
+            "multi_modal_data": {
+                "source_video": "source.mov",
+                "source_audio": source_audio,
+                "video_noise_mask": torch.zeros(1, 1, latent_t, 4, 6),
+                "audio_noise_mask": torch.tensor([0.0, 0.5, *([0.0] * (audio_t - 2))]).reshape(1, 1, -1),
+            },
+        },
+        SimpleNamespace(extra_args={"task": "t2va"}),
+    )
+    media = type(prepared.media).from_mm_tensors(
+        prepared.media.to_mm_tensors(),
+        prepared.media.to_metadata(),
+    )
+    conditioning = processing.encode_media(
+        media,
+        video_vae=VideoVAE(),
+        audio_vae=AudioVAE(),
+        emit_conditioning=True,
+    )
+
+    assert prepared.images == []
+    assert prepared.qwen_videos == []
+    assert prepared.condition_labels == []
+    assert prepared.media.images == ()
+    assert prepared.media.videos == ()
+    assert calls["prepare"][0] == "source.mov"
+    assert calls["prepare"][1]["target_width"] == 96
+    assert calls["prepare"][1]["target_height"] == 64
+    assert calls["prepare"][1]["target_frame_count"] == 22
+    if source_audio is None:
+        assert calls["video_audio"] == (
+            "source.mov",
+            {"duration_seconds": pytest.approx(22 / 24)},
+        )
+        assert "audio_file" not in calls
+    else:
+        assert calls["audio_file"] == "override.wav"
+        assert "video_audio" not in calls
+
+    assert conditioning is not None
+    torch.testing.assert_close(conditioning.video_edit_clean_rows, video_clean)
+    assert conditioning.video_edit_mask.shape == (latent_t, 4, 6)
+    torch.testing.assert_close(conditioning.video_edit_mask, torch.zeros(latent_t, 4, 6))
+    assert conditioning.audio_edit_source_t == 2
+    assert conditioning.audio_edit_mask.shape == (2, audio_t)
+    torch.testing.assert_close(
+        conditioning.audio_edit_clean_rows[:, 0],
+        torch.tensor([1.0, 2.0, *([0.0] * (audio_t - 2)), 11.0, 12.0, *([0.0] * (audio_t - 2))]),
+    )
+    torch.testing.assert_close(
+        conditioning.audio_edit_mask,
+        torch.tensor([0.0, 0.5, *([0.0] * (audio_t - 2))]).repeat(2, 1),
+    )
+
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import (
+        MiniMaxH3EncoderConditioning,
+        MiniMaxH3TextConditioning,
+    )
+
+    combined = MiniMaxH3EncoderConditioning.from_components(
+        MiniMaxH3TextConditioning(
+            torch.zeros(1, 5120, dtype=torch.bfloat16),
+            torch.ones(1, dtype=torch.int64),
+        ),
+        conditioning,
+    )
+    round_trip = MiniMaxH3EncoderConditioning.from_omni_payload(combined.to_omni_payload())
+    torch.testing.assert_close(round_trip.video_edit_clean_rows, video_clean)
+    torch.testing.assert_close(round_trip.video_edit_mask, torch.zeros(latent_t, 4, 6))
+    torch.testing.assert_close(round_trip.audio_edit_clean_rows, conditioning.audio_edit_clean_rows)
+    torch.testing.assert_close(
+        round_trip.audio_edit_mask,
+        torch.tensor([0.0, 0.5, *([0.0] * (audio_t - 2))]).repeat(2, 1),
+    )
+    assert round_trip.audio_edit_source_t == 2
+
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.default_video_shift = 12.0
+    pipeline.default_audio_shift = 3.0
+    pipeline.od_config = SimpleNamespace(step_execution=False)
+    pipeline._fasth3 = None
+    pipeline._active_turbo_spec = lambda _sampling: None
+    pipeline._has_active_native_lora = lambda _sampling: False
+    pipeline._resolve_task = lambda task, **_kwargs: task
+    pipeline._resolve_sigma_positions = lambda _task, _sampling: (None, 3)
+    pipeline._quality_policy = SimpleNamespace(resolve=lambda **_kwargs: SimpleNamespace(cache_dit=None))
+    pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda _plan: None)
+    context = pipeline._prepare_encoder_conditioning_inputs(
+        round_trip,
+        SimpleNamespace(extra_args={}, quality=None, seed=17, num_outputs_per_prompt=1),
+    )
+    torch.testing.assert_close(context["video_edit_clean_rows"], video_clean)
+    torch.testing.assert_close(context["video_edit_mask_rows"], torch.zeros(42))
+    assert context["video_edit_restore_mask_rows"].shape == (42, 96)
+    torch.testing.assert_close(
+        context["audio_edit_mask_rows"],
+        torch.tensor([0.0, 0.5, *([1.0] * (audio_t - 2))]).repeat(2),
+    )
+    torch.testing.assert_close(context["audio_edit_mask_rows"], context["audio_edit_restore_mask_rows"])
+
+
+def _build_inputs(*, video_mask: torch.Tensor, audio_mask: torch.Tensor):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    latent_t, latent_h, latent_w, audio_t = 2, 4, 6, 3
+    video_clean = torch.linspace(-1.0, 1.0, latent_t * (latent_h // 2) * (latent_w // 2) * 96).reshape(-1, 96)
+    audio_clean = torch.linspace(-0.5, 0.5, audio_t * 2 * 32).reshape(-1, 32)
+    inputs = pipeline._build_denoise_inputs(
+        task="t2va",
+        text_embeddings=torch.zeros(2, 4),
+        text_tags=torch.ones(2, dtype=torch.long),
+        seed=17,
+        latent_t=latent_t,
+        latent_h=latent_h,
+        latent_w=latent_w,
+        audio_t=audio_t,
+        num_frames=22,
+        num_steps=3,
+        video_shift=12.0,
+        audio_shift=3.0,
+        base_schedule=None,
+        visual_condition=None,
+        visual_condition_shape=None,
+        audio_condition=None,
+        ref_audio_t=None,
+        video_edit_clean_rows=video_clean,
+        video_edit_mask_rows=video_mask,
+        video_edit_restore_mask_rows=video_mask,
+        audio_edit_clean_rows=audio_clean,
+        audio_edit_mask_rows=audio_mask,
+        audio_edit_restore_mask_rows=audio_mask,
+    )
+    return pipeline, inputs, video_clean, audio_clean
+
+
+def test_build_denoise_inputs_uses_h3_source_anchors():
+    pipeline, inputs, video_clean, audio_clean = _build_inputs(
+        video_mask=torch.zeros(12),
+        audio_mask=torch.zeros(6),
+    )
+    video_edit = inputs["video_edit"]
+    audio_edit = inputs["audio_edit"]
+    assert video_edit is not None and audio_edit is not None
+    initial_video, _ = pipeline._initial_noise(
+        seed=17,
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+    )
+    torch.testing.assert_close(video_edit.clean_rows, video_clean)
+    torch.testing.assert_close(video_edit.anchor_rows, 0.999 * video_clean + 0.001 * initial_video)
+    torch.testing.assert_close(audio_edit.clean_rows, audio_clean)
+    torch.testing.assert_close(audio_edit.anchor_rows, audio_clean)

tests/diffusion/models/minimax_h3/test_minimax_h3_latent_mask.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Focused CPU tests for MiniMax H3 latent-mask editing."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_video_full_grid_quantization_and_cell_level_restore():
+    from vllm_omni.diffusion.models.minimax_h3.latent_mask import minimax_h3_video_edit_masks
+
+    shape = dict(latent_t=2, latent_h=4, latent_w=6)
+    full = torch.arange(48, dtype=torch.float32).reshape(2, 4, 6) / 48
+    cells = full.reshape(2, 2, 2, 3, 2).permute(0, 1, 3, 2, 4)
+    parsed = minimax_h3_video_edit_masks(full, **shape)
+    expected_model = torch.ceil(cells.amax(dim=(3, 4)).reshape(-1) * 256) / 256
+    expected_restore = cells.unsqueeze(3).expand(2, 2, 3, 24, 2, 2).reshape(12, 96)
+    torch.testing.assert_close(parsed.model_mask_rows, expected_model)
+    torch.testing.assert_close(parsed.restore_mask_rows, expected_restore)
+
+
+def test_audio_grid_uses_channel_major_order_and_preserves_raw_values():
+    from vllm_omni.diffusion.models.minimax_h3.latent_mask import minimax_h3_audio_edit_masks
+
+    grid = torch.tensor([[0.0, 0.1, 1.0], [0.25, 0.5, 0.75]])
+    expected = grid.flatten()
+    parsed = minimax_h3_audio_edit_masks(grid, audio_t=3)
+    torch.testing.assert_close(parsed.model_mask_rows, torch.ceil(expected * 256) / 256)
+    torch.testing.assert_close(parsed.restore_mask_rows, expected)
+
+
+def test_zero_fractional_one_semantics_and_cell_level_restore():
+    from vllm_omni.diffusion.models.minimax_h3.latent_mask import (
+        MiniMaxH3LatentEdit,
+        minimax_h3_video_edit_masks,
+    )
+
+    clean = torch.full((3, 2), 10.0)
+    anchor = torch.tensor([[8.0, 8.0], [6.0, 6.0], [4.0, 4.0]])
+    state = torch.tensor([[2.0, 2.0], [4.0, 4.0], [6.0, 6.0]])
+    velocity = torch.full((3, 2), 2.0)
+    mask = torch.tensor([0.0, 0.5, 1.0])
+    edit = MiniMaxH3LatentEdit.from_rows(clean, anchor, mask, mask)
+    assert edit is not None
+
+    model_rows, target_timesteps = edit.prepare(state, 0.25, 0.999, sigma=0.75)
+    torch.testing.assert_close(model_rows, torch.lerp(anchor, state, mask[:, None]))
+    torch.testing.assert_close(
+        target_timesteps,
+        torch.tensor([0.999, 0.625, 0.25]),
+    )
+    predicted = model_rows + 0.75 * velocity * mask[:, None]
+    torch.testing.assert_close(edit.x0(model_rows, velocity, 0.25), torch.lerp(clean, predicted, mask[:, None]))
+
+    # A single active cell drives the full model token, while raw restoration
+    # preserves the other three cells independently in every packed channel.
+    parsed = minimax_h3_video_edit_masks(
+        torch.tensor([[[1.0, 0.0], [0.0, 0.0]]]),
+        latent_t=1,
+        latent_h=2,
+        latent_w=2,
+    )
+    restore = torch.tensor([1.0, 0.0, 0.0, 0.0]).repeat(24).reshape(1, 96)
+    torch.testing.assert_close(parsed.model_mask_rows, torch.ones(1))
+    torch.testing.assert_close(parsed.restore_mask_rows, restore)
+
+    cell_edit = MiniMaxH3LatentEdit.from_rows(
+        torch.full((1, 96), 10.0),
+        torch.full((1, 96), 8.0),
+        parsed.model_mask_rows,
+        parsed.restore_mask_rows,
+    )
+    assert cell_edit is not None
+    cell_state = torch.full((1, 96), 2.0)
+    cell_velocity = torch.full((1, 96), 4.0)
+    expected = torch.where(restore.bool(), cell_state + 0.75 * cell_velocity, cell_edit.clean_rows)
+    torch.testing.assert_close(cell_edit.x0(cell_state, cell_velocity, 0.25), expected)

tests/diffusion/models/minimax_h3/test_minimax_h3_long_video_limits.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Duration limits apply before encoding and to transported encoder payloads."""
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from PIL import Image
+
+from vllm_omni.errors import OmniClientError
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
+    prepare_encoder_inputs,
+    resolve_minimax_h3_shape,
+)
+from vllm_omni.model_executor.models.minimax_h3.long_video import validate_encoded_frame_limit
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("mode,limit", [("full", 30), ("continuation", 300), (None, 300)])
+@pytest.mark.parametrize("input_kind", ["duration", "duration_seconds", "target", "num_frames"])
+def test_long_video_limits_cover_all_duration_inputs(mode, limit, input_kind):
+    extra: dict[str, Any] = {"long_video": True}
+    if mode is not None:
+        extra["long_video_mode"] = mode
+    sampling = OmniDiffusionSamplingParams(width=64, height=64, fps=24, extra_args=extra)
+
+    def set_duration(seconds):
+        if input_kind == "target":
+            extra["target"] = {"duration_seconds": seconds}
+        elif input_kind == "num_frames":
+            sampling.num_frames = seconds * 24
+        else:
+            extra[input_kind] = seconds
+
+    set_duration(limit)
+    _, _, frames, _, _ = resolve_minimax_h3_shape("ref2va", sampling, None)
+    assert frames >= limit * 24
+    validate_encoded_frame_limit(extra, "ref2va", frames)
+    with pytest.raises(OmniClientError, match="request limit"):
+        validate_encoded_frame_limit(extra, "ref2va", frames + 17)
+    set_duration(limit + 1)
+    with pytest.raises(OmniClientError, match="duration"):
+        resolve_minimax_h3_shape("ref2va", sampling, None)
+
+
+@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"])
+def test_full_mode_rejects_300_seconds_and_accepts_comparison(task):
+    sampling = OmniDiffusionSamplingParams(
+        width=64,
+        height=64,
+        fps=24,
+        extra_args={"long_video": True, "long_video_mode": "full", "duration": 300, "aspect_ratio": "1:1"},
+    )
+    image = Image.new("RGB", (64, 64)) if task == "fl2va" else None
+    with pytest.raises(OmniClientError, match="30"):
+        resolve_minimax_h3_shape(task, sampling, image)
+    sampling.extra_args["duration"] = 20
+    assert resolve_minimax_h3_shape(task, sampling, image)[2] == 481
+
+
+@pytest.mark.parametrize("mode", ["full", "continuation"])
+def test_mode_alone_does_not_bypass_long_video_opt_in(mode):
+    sampling = OmniDiffusionSamplingParams(
+        width=64,
+        height=64,
+        extra_args={"long_video_mode": mode, "duration": 20},
+    )
+    with pytest.raises(OmniClientError, match="15"):
+        resolve_minimax_h3_shape("ref2va", sampling, None)
+
+
+@pytest.mark.parametrize("kind", ["video", "audio"])
+def test_continuation_edit_rejected_before_loading_source(kind):
+    sampling = OmniDiffusionSamplingParams(
+        width=64,
+        height=64,
+        extra_args={"long_video": True, "task": "ref2va", "duration": 20},
+    )
+    with pytest.raises(OmniClientError, match="continuation does not support latent-mask editing"):
+        prepare_encoder_inputs(
+            {
+                "prompt": "test",
+                "multi_modal_data": {
+                    "image": Image.new("RGB", (64, 64)),
+                    f"source_{kind}": "must-not-be-loaded",
+                    f"{kind}_noise_mask": 0.5,
+                },
+            },
+            sampling,
+        )
+
+
+def test_locked_audio_rejects_audio_edit_before_loading_source():
+    sampling = OmniDiffusionSamplingParams(
+        width=64,
+        height=64,
+        extra_args={"task": "t2va", "aspect_ratio": "1:1", "audio_mode": "lock_source"},
+    )
+    with pytest.raises(OmniClientError, match="cannot be combined"):
+        prepare_encoder_inputs(
+            {
+                "prompt": "test",
+                "multi_modal_data": {
+                    "audio": (torch.zeros(8000), 8000),
+                    "source_audio": "must-not-be-loaded",
+                    "audio_noise_mask": 0.5,
+                },
+            },
+            sampling,
+        )
+
+
+@pytest.mark.parametrize("kind", ["video", "audio"])
+def test_external_encoder_cannot_bypass_continuation_edit_validation(kind):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(step_execution=False)
+    conditioning = SimpleNamespace(
+        task="ref2va",
+        num_frames=481,
+        video_edit_clean_rows=torch.zeros(1) if kind == "video" else None,
+        audio_edit_clean_rows=torch.zeros(1) if kind == "audio" else None,
+    )
+    sampling = OmniDiffusionSamplingParams(extra_args={"long_video": True})
+    with pytest.raises(OmniClientError, match="continuation does not support latent-mask editing"):
+        pipeline._prepare_encoder_conditioning_inputs(conditioning, sampling)
+
+
+def test_external_encoder_cannot_bypass_full_duration_limit():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(step_execution=False)
+    sampling = OmniDiffusionSamplingParams(extra_args={"long_video": True, "long_video_mode": "full"})
+    with pytest.raises(OmniClientError, match="30-second request limit"):
+        pipeline._prepare_encoder_conditioning_inputs(SimpleNamespace(task="ref2va", num_frames=7213), sampling)

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -6,10 +6,11 @@ import torch
 from vllm.triton_utils import HAS_TRITON
 
 from tests.helpers.mark import hardware_marks
-from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
+from vllm_omni.diffusion.layers.indexed_modulation import (
     _MAX_1D_GRID_SIZE,
     _iter_row_chunks,
     _launch_row_chunks,
+    _use_hopper_bf16_affine_semantics,
     indexed_gate,
     indexed_gate_rms_norm_scale_shift,
     indexed_scale_shift_,
@@ -190,3 +191,55 @@ def test_fused_modulation_preserves_bf16_residual_boundary() -> None:
     )
 
     assert torch.equal(modulated_out, expected)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(
+    torch.cuda.is_available() and torch.cuda.get_device_capability()[0] != 9,
+    reason="SM90-specific precision path",
+)
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_hopper_fused_modulation_matches_pytorch_reference() -> None:
+    torch.manual_seed(42)
+    rows, hidden_size, conditions = 128, 3072, 4
+    residual = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
+    gate = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    branch = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    indices = torch.arange(rows, device="cuda") % conditions
+    eps = 1e-6
+
+    reference_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16)
+    normalized = reference_residual.float()
+    variance = normalized.pow(2).mean(-1, keepdim=True)
+    normalized = (weight.float() * (normalized * torch.rsqrt(variance + eps))).to(torch.bfloat16)
+    reference_modulated = (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(
+        torch.bfloat16
+    )
+
+    actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift(
+        residual,
+        gate,
+        branch,
+        weight,
+        shift,
+        scale,
+        indices,
+        eps,
+    )
+
+    torch.testing.assert_close(actual_residual, reference_residual, atol=5e-2, rtol=5e-2)
+    torch.testing.assert_close(actual_modulated, reference_modulated, atol=5e-2, rtol=5e-2)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_hopper_precision_dispatch_is_architecture_scoped(monkeypatch) -> None:
+    tensor = torch.empty(0, device="cuda")
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0))
+    assert _use_hopper_bf16_affine_semantics(tensor)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
+    assert not _use_hopper_bf16_affine_semantics(tensor)

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -1,82 +1,282 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-from contextlib import contextmanager
+import weakref
+from contextlib import contextmanager, nullcontext
 from types import SimpleNamespace
 from typing import Any
 from unittest.mock import Mock
 
 import pytest
 import torch
-from PIL import Image
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-def test_h3_profiler_targets_cover_reference_encoders_and_vae_decoders():
+def test_h3_full_pipeline_profiler_includes_local_encoding():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     targets = MiniMaxH3Pipeline._PROFILER_TARGETS
-    assert "_encode_visual_conditions" in targets
-    assert "_encode_reference_audio_conditions" in targets
-    assert "video_vae.decode_latent" in targets
-    assert "audio_vae.decode_latent" in targets
-    assert "_encode_video_conditions" not in targets
-    assert "_encode_video_audio_conditions" not in targets
-    assert "_encode_audio_conditions" not in targets
+    assert targets == [
+        "encode_prompt",
+        "_encode_local_media",
+        "diffuse",
+        "decode",
+        "video_vae.decode_latent",
+        "audio_vae.decode_latent",
+        "prepare_encode",
+        "denoise_step",
+        "post_decode",
+    ]
 
 
-@pytest.mark.parametrize("distributed", [False, True])
-def test_manual_component_offload_wins_over_requested_model_offload(distributed):
+def test_decoded_output_is_unchanged_without_active_model_offload():
     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=not distributed,
-        enable_distributed_layerwise_offload=distributed,
-    )
     pipeline._model_cpu_offload_modules = []
-    pipeline.text_encoder = Mock()
-    expected = torch.ones(2, 3)
-    pipeline.text_encoder.encode_ids.return_value = expected
-    component = Mock()
+    output = Mock()
 
-    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
-    with pipeline._component_on_device(component):
-        component.load_to_device.assert_called_once_with()
-        component.offload_to_cpu.assert_not_called()
-
-    assert actual is expected
-    pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
-    component.offload_to_cpu.assert_called_once_with()
+    assert pipeline._offload_model_cpu_stage_output(output) is output
+    output.cpu.assert_not_called()
 
 
-def test_requested_model_offload_without_backend_uses_resident_components():
+def test_reply_rank_copies_decoded_output_to_host(monkeypatch):
     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,
-        enable_distributed_layerwise_offload=False,
+    pipeline._model_cpu_offload_modules = [Mock()]
+    monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: True)
+    output = Mock()
+    host_output = Mock()
+    output.cpu.return_value = host_output
+
+    assert pipeline._offload_model_cpu_stage_output(output) is host_output
+    output.cpu.assert_called_once_with()
+
+
+def test_non_reply_rank_drops_decoded_output_storage(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline._model_cpu_offload_modules = [Mock()]
+    monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: False)
+    output = torch.ones((2, 3), dtype=torch.float16)
+
+    placeholder = pipeline._offload_model_cpu_stage_output(output)
+
+    assert placeholder.device.type == "meta"
+    assert placeholder.shape == output.shape
+    assert placeholder.dtype == output.dtype
+
+
+def test_decode_releases_only_audio_before_return(monkeypatch):
+    """Video is released by the callers of ``decode`` -- after quantization."""
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    video = torch.ones((1, 1, 4, 4))
+    audio = torch.ones((1, 8))
+    pipeline = SimpleNamespace(
+        device=torch.device("cpu"),
+        video_vae=Mock(decode_latent=Mock(return_value=video)),
+        audio_vae=Mock(decode_latent=Mock(return_value=audio)),
+        _component_on_device=lambda component: nullcontext(),
+        _offload_model_cpu_stage_output=Mock(side_effect=lambda value: value),
+        _release_stage_cache=Mock(),
+    )
+    monkeypatch.setattr(
+        module.current_omni_platform,
+        "create_autocast_context",
+        lambda **kwargs: nullcontext(),
     )
-    pipeline._model_cpu_offload_modules = []
-    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]), {})
+    actual_video, actual_audio = MiniMaxH3Pipeline.decode(
+        pipeline,
+        torch.zeros(1),
+        torch.zeros(1),
+        height=2,
+        width=3,
+    )
 
-    assert actual is expected
-    pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.encode_ids.assert_called_once()
+    assert actual_video.shape[-2:] == (2, 3)
+    assert actual_audio is audio
+    calls = pipeline._offload_model_cpu_stage_output.call_args_list
+    assert len(calls) == 1
+    assert calls[0].args[0] is audio
 
 
-def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
+def _recording_release(events, releases):
+    """Stand-in for ``_offload_model_cpu_stage_output`` that records what it got."""
+
+    def release(value):
+        events.append("release")
+        releases.append(value)
+        return value
+
+    return release
+
+
+def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    # MiniMaxH3VideoVAE.decode_latent ends in frames.float(), so this is what
+    # the release hook would see if it ran before the quantizer.
+    decoded_dtype = torch.float32
+    events: list[str] = []
+    releases: list[torch.Tensor] = []
+
+    def diffuse(**kwargs):
+        events.append(f"diffuse:{kwargs['seed']}")
+        return torch.zeros(1), torch.zeros(1)
+
+    def decode(video_latent, audio_latent, *, height, width):
+        events.append("decode")
+        return (
+            torch.zeros((1, 3, 2, height, width), dtype=decoded_dtype),
+            torch.zeros((1, 8)),
+        )
+
+    # ``forward`` indexes ``preencode_mp4`` (#7018): this test drives the
+    # tensor-decode branch, which is the one the release hook wraps.
+    context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4, "preencode_mp4": False}
+    pipeline = SimpleNamespace(
+        od_config=SimpleNamespace(),
+        _extract_prompt=lambda raw: ("a prompt", {}),
+        _extract_text_conditioning=lambda raw: None,
+        _extract_prepared_reference_videos=lambda raw: None,
+        _prepare_request_inputs=lambda *args, **kwargs: context,
+        _denoise_kwargs=lambda ctx: {},
+        diffuse=diffuse,
+        decode=decode,
+        _offload_model_cpu_stage_output=_recording_release(events, releases),
+        _release_stage_cache=lambda: None,
+    )
+    request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace())
+
+    output = MiniMaxH3Pipeline.forward(pipeline, request)
+
+    # The release must see the quantized frames, not the decoded floats.
+    assert len(releases) == 2
+    assert [tensor.dtype for tensor in releases] == [torch.uint8, torch.uint8]
+    assert decoded_dtype is not torch.uint8
+    # ... and it must still happen before the next seed reloads the DiT.
+    assert events == ["diffuse:41", "decode", "release", "diffuse:42", "decode", "release"]
+    assert output.output[0].dtype is torch.uint8
+
+
+def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypatch):
+    """Releasing the copy is not enough if a local still pins the original.
+
+    ``_offload_model_cpu_stage_output`` hands back a *new* tensor, so the
+    decoded frames stay alive for as long as anything still references them.
+    The whole point of this branch is that they are gone before the next seed
+    reloads the DiT, and ``diffuse`` is where that reload happens.
+    """
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    decoded_refs: list[weakref.ref] = []
+    alive_at_diffuse: list[int] = []
+
+    def diffuse(**kwargs):
+        alive_at_diffuse.append(sum(1 for ref in decoded_refs if ref() is not None))
+        return torch.zeros(1), torch.zeros(1)
+
+    def decode(video_latent, audio_latent, *, height, width):
+        video = torch.zeros((1, 3, 2, height, width), dtype=torch.float32)
+        decoded_refs.append(weakref.ref(video))
+        return video, torch.zeros((1, 8))
+
+    # ``forward`` indexes ``preencode_mp4`` (#7018): this test drives the
+    # tensor-decode branch, which is the one the release hook wraps.
+    context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4, "preencode_mp4": False}
+    pipeline = SimpleNamespace(
+        od_config=SimpleNamespace(),
+        _extract_prompt=lambda raw: ("a prompt", {}),
+        _extract_text_conditioning=lambda raw: None,
+        _extract_prepared_reference_videos=lambda raw: None,
+        _prepare_request_inputs=lambda *args, **kwargs: context,
+        _denoise_kwargs=lambda ctx: {},
+        diffuse=diffuse,
+        decode=decode,
+        _offload_model_cpu_stage_output=lambda value: value.clone(),
+        _release_stage_cache=lambda: None,
+    )
+    request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace())
+
+    MiniMaxH3Pipeline.forward(pipeline, request)
+
+    assert len(decoded_refs) == 2
+    # Seed 1's diffuse sees nothing decoded yet; seed 2's must not see seed 1's
+    # decoded frames still resident.
+    assert alive_at_diffuse == [0, 0]
+
+
+def test_post_decode_releases_video_after_quantization(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    events: list[str] = []
+    releases: list[torch.Tensor] = []
+
+    def decode(video_latent, audio_latent, *, height, width):
+        events.append("decode")
+        return (
+            torch.zeros((1, 3, 2, height, width), dtype=torch.float16),
+            torch.zeros((1, 8)),
+        )
+
+    pipeline = SimpleNamespace(
+        od_config=SimpleNamespace(),
+        _unpack_denoised_rows=lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)),
+        decode=decode,
+        _offload_model_cpu_stage_output=_recording_release(events, releases),
+        _release_stage_cache=lambda: None,
+    )
+    state = SimpleNamespace(
+        latents=torch.zeros(1),
+        extra={
+            module._STEP_BRANCH: object(),
+            module._STEP_AUDIO_ROWS: torch.zeros(1),
+            module._STEP_SHAPE: {
+                "latent_t": 1,
+                "latent_h": 1,
+                "latent_w": 1,
+                "audio_t": 1,
+                "height": 4,
+                "width": 4,
+            },
+        },
+    )
+
+    output = MiniMaxH3Pipeline.post_decode(pipeline, state)
+
+    assert events == ["decode", "release"]
+    assert len(releases) == 1
+    assert releases[0].dtype is torch.uint8
+    assert output.output[0] is releases[0]
+
+
+@pytest.mark.parametrize(
+    ("available", "initialized", "rank", "expected"),
+    [(False, False, 7, True), (True, False, 7, True), (True, True, 0, True), (True, True, 1, False)],
+)
+def test_output_owner_rank_matches_executor_contract(monkeypatch, available, initialized, rank, expected):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    monkeypatch.setattr(module.dist, "is_available", lambda: available)
+    monkeypatch.setattr(module.dist, "is_initialized", lambda: initialized)
+    monkeypatch.setattr(module.dist, "get_rank", lambda: rank)
+
+    assert MiniMaxH3Pipeline._is_output_owner_rank() is expected
+
+
+@pytest.mark.parametrize("load_text_encoder", [True, False])
+def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch, load_text_encoder):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
 
@@ -84,9 +284,11 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     torch.nn.Module.__init__(pipeline)
     pipeline.transformer = torch.nn.Linear(2, 2)
     pipeline.transformers_ref = torch.nn.Linear(2, 2)
-    pipeline.text_encoder = torch.nn.Linear(2, 2)
     pipeline.video_vae = torch.nn.Linear(2, 2)
     pipeline.audio_vae = torch.nn.Linear(2, 2)
+    pipeline._encoder_modules = ["text_encoder"] if load_text_encoder else []
+    if load_text_encoder:
+        pipeline.text_encoder = torch.nn.Linear(2, 2)
     apply_offload = Mock()
     remove_offload = Mock()
     monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
@@ -99,7 +301,7 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     )
 
     dits = [pipeline.transformer, pipeline.transformers_ref]
-    stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
+    stages = [*([pipeline.text_encoder] if load_text_encoder else []), pipeline.video_vae, pipeline.audio_vae]
     apply_offload.assert_called_once_with(
         dit_modules=dits,
         encoder_modules=stages,
@@ -114,6 +316,46 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     remove_offload.assert_called_once_with([*dits, *stages])
 
 
+@pytest.mark.parametrize("decode_fails", [False, True])
+def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    component = torch.nn.Linear(2, 2)
+    events: list[tuple[Any, ...]] = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    pipeline._model_cpu_offload_modules = [component]
+
+    def decode():
+        with pipeline._component_on_device(component):
+            events.append(("decode", component))
+            if decode_fails:
+                raise RuntimeError("decode failed")
+
+    if decode_fails:
+        with pytest.raises(RuntimeError, match="decode failed"):
+            decode()
+    else:
+        decode()
+
+    assert events == [
+        ("activate", component),
+        ("decode", component),
+        ("offload", component),
+    ]
+
+
 def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
@@ -174,44 +416,20 @@ def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
     apply_offload.assert_not_called()
 
 
-@pytest.mark.parametrize("decode_fails", [False, True])
-def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+def _encode_media(pipeline, **kwargs):
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
+    from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media
 
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    component = torch.nn.Linear(2, 2)
-    events: list[tuple[Any, ...]] = []
-
-    @contextmanager
-    def record_component(value):
-        events.append(("activate", value))
-        try:
-            yield
-        finally:
-            events.append(("offload", value))
-
-    monkeypatch.setattr(module, "sequential_offload_component", record_component)
-    pipeline._model_cpu_offload_modules = [component]
-
-    def decode():
-        with pipeline._component_on_device(component):
-            events.append(("decode", component))
-            if decode_fails:
-                raise RuntimeError("decode failed")
-
-    if decode_fails:
-        with pytest.raises(RuntimeError, match="decode failed"):
-            decode()
-    else:
-        decode()
-
-    assert events == [
-        ("activate", component),
-        ("decode", component),
-        ("offload", component),
-    ]
+    media = MiniMaxH3EncoderMediaInput(
+        task="ref2va", height=32, width=32, num_frames=124, latent_t=31, audio_t=208, **kwargs
+    )
+    return encode_media(
+        media,
+        video_vae=getattr(pipeline, "video_vae", None),
+        audio_vae=getattr(pipeline, "audio_vae", None),
+        emit_conditioning=True,
+        component_scope=pipeline._component_on_device,
+    )
 
 
 @pytest.mark.parametrize("encode_fails", [False, True])
@@ -221,7 +439,6 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
     pipeline.video_vae = Mock()
     pipeline._model_cpu_offload_modules = [pipeline.video_vae]
     events = []
@@ -242,17 +459,15 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode
 
     pipeline.video_vae.encode_image.side_effect = encode_image
     monkeypatch.setattr(module, "sequential_offload_component", record_component)
-    monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
-    monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
-    images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
+    images = (torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8))
 
     if encode_fails:
         with pytest.raises(RuntimeError, match="encode failed"):
-            pipeline._encode_visual_conditions(images, None, video_count=0)
+            _encode_media(pipeline, images=images)
     else:
-        rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0)
-        assert rows.shape == (3, 4)
-        assert shapes == [(1, 1, 1), (1, 1, 2)]
+        result = _encode_media(pipeline, images=images)
+        assert result.visual_condition.shape == (3, 4)
+        assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2))
 
     assert events == [
         ("activate", pipeline.video_vae),
@@ -268,9 +483,7 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
     pipeline.video_vae = Mock()
-    pipeline.video_vae.is_distributed_enabled.return_value = False
     pipeline._model_cpu_offload_modules = [pipeline.video_vae]
     events = []
 
@@ -287,31 +500,26 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
         return torch.ones(1, 4)
 
     def encode_video(frames):
-        events.append(("video", frames))
+        events.append(("video", frames.shape))
         return torch.ones(3, 4), (3, 2, 2)
 
     pipeline.video_vae.encode_image.side_effect = encode_image
     pipeline.video_vae.encode_video.side_effect = encode_video
     monkeypatch.setattr(module, "sequential_offload_component", record_component)
-    monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
-    monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
-    monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}")
-    images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
-    prepared_videos = [{"prepared_path": "reference.mp4"}]
-
-    rows, shapes = pipeline._encode_visual_conditions(
-        images,
-        prepared_videos,
-        video_count=1,
+    result = _encode_media(
+        pipeline,
+        images=(torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8)),
+        videos=(torch.zeros(12, 32, 32, 3, dtype=torch.uint8),),
+        video_audios=(None,),
     )
 
-    assert rows.shape == (5, 4)
-    assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)]
+    assert result.visual_condition.shape == (5, 4)
+    assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2), (3, 2, 2))
     assert events == [
         ("activate", pipeline.video_vae),
         ("image", (16, 16)),
         ("image", (32, 16)),
-        ("video", "frames:reference.mp4"),
+        ("video", (12, 32, 32, 3)),
         ("offload", pipeline.video_vae),
     ]
 
@@ -325,7 +533,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
     torch.nn.Module.__init__(pipeline)
     pipeline.audio_vae = Mock()
     pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
-    events: list[tuple[Any, ...]] = []
+    events: list[tuple] = []
 
     @contextmanager
     def record_component(value):
@@ -335,26 +543,20 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
         finally:
             events.append(("offload", value))
 
-    def encode_embedded(*args, **kwargs):
-        events.append(("embedded",))
+    def encode_waveform(waveform, sample_rate):
+        events.append(("audio", int(waveform[0]), sample_rate))
         if encode_fails:
             raise RuntimeError("audio encode failed")
-        return torch.ones(1, 2), [1]
+        return torch.ones(80, 32), 80
 
-    def encode_standalone(*args, **kwargs):
-        events.append(("standalone",))
-        return torch.ones(1, 2), [1]
-
-    pipeline._encode_video_audio_conditions_resident = encode_embedded
-    pipeline._encode_audio_conditions_resident = encode_standalone
+    pipeline.audio_vae.encode_waveform.side_effect = encode_waveform
     monkeypatch.setattr(module, "sequential_offload_component", record_component)
 
     def call():
-        return pipeline._encode_reference_audio_conditions(
-            [{"input_has_audio": True}],
-            has_audio=[True],
-            standalone_audios=[(torch.ones(1), 16_000)],
-            max_duration_seconds=1.0,
+        return _encode_media(
+            pipeline,
+            video_audios=(),
+            audios=((torch.ones(32_000), 16_000),),
         )
 
     if encode_fails:
@@ -362,10 +564,61 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
             call()
     else:
         result = call()
-        assert result[0].shape == (1, 2)
-        assert result[2].shape == (1, 2)
+        assert result.audio_condition.shape == (80, 32)
+        assert result.audio_condition_lengths == (80,)
+        assert events[1:-1] == [("audio", 1, 16_000)]
 
     assert events[0] == ("activate", pipeline.audio_vae)
     assert events[-1] == ("offload", pipeline.audio_vae)
     assert events.count(("activate", pipeline.audio_vae)) == 1
     assert events.count(("offload", pipeline.audio_vae)) == 1
+
+
+@pytest.mark.parametrize("distributed", [False, True])
+def test_manual_component_offload_wins_over_requested_model_offload(distributed):
+    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=not distributed,
+        enable_distributed_layerwise_offload=distributed,
+    )
+    pipeline._model_cpu_offload_modules = []
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+    component = Mock()
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+    with pipeline._component_on_device(component):
+        component.load_to_device.assert_called_once_with()
+        component.offload_to_cpu.assert_not_called()
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+    component.offload_to_cpu.assert_called_once_with()
+
+
+def test_requested_model_offload_without_backend_uses_resident_components():
+    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,
+        enable_distributed_layerwise_offload=False,
+    )
+    pipeline._model_cpu_offload_modules = []
+    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.encode_ids.assert_called_once()

tests/diffusion/models/minimax_h3/test_minimax_h3_packing.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
 
 import pytest
 import torch
@@ -190,3 +190,90 @@ def test_condition_noise_accepts_a_reference_video_longer_than_the_target():
         noise_aug=0.5,
     )
     assert result.shape == rows.shape
+
+
+def test_explicit_seq_len_pins_one_shape_across_prompt_lengths():
+    """Prompts of different token counts must land on the same packed length.
+
+    Without the pin they fall into different 64-row buckets, and each bucket is
+    a packed shape the compiled transformer has not seen before.
+    """
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence,
+    )
+
+    common = dict(
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+        include_keyframe_cond=False,
+    )
+    # used = text_len + 18 rows: 22 -> bucket 64; 118 -> bucket 128.
+    short = minimax_h3_packed_sequence(text_len=4, **common)
+    long = minimax_h3_packed_sequence(text_len=100, **common)
+    assert int(short["seq_len"]) == 64
+    assert int(long["seq_len"]) == 128
+
+    short_pinned = minimax_h3_packed_sequence(text_len=4, seq_len=192, **common)
+    long_pinned = minimax_h3_packed_sequence(text_len=100, seq_len=192, **common)
+    assert int(short_pinned["seq_len"]) == int(long_pinned["seq_len"]) == 192
+    # The pin only adds padding rows; the used prefix is untouched.
+    assert int(short_pinned["cu_seqlens"][1]) == int(short["cu_seqlens"][1]) == 22
+    assert int(long_pinned["cu_seqlens"][1]) == int(long["cu_seqlens"][1]) == 118
+    assert short_pinned["img_pos"].tolist() == short["img_pos"].tolist()
+
+
+def test_the_ref2va_packer_pins_one_shape_across_prompt_lengths_too():
+    """Both packers share one resolver, so both honour the pin."""
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence_ref2va_blocks,
+    )
+
+    common = dict(
+        latent_t=2,
+        latent_h=4,
+        latent_w=6,
+        audio_t=3,
+        ref_blocks=[{"kind": "audio", "ref_audio_t": 1}],
+    )
+    short = minimax_h3_packed_sequence_ref2va_blocks(text_len=4, **common)
+    long = minimax_h3_packed_sequence_ref2va_blocks(text_len=100, **common)
+    assert int(short["seq_len"]) != int(long["seq_len"])
+
+    short_pinned = minimax_h3_packed_sequence_ref2va_blocks(text_len=4, seq_len=192, **common)
+    long_pinned = minimax_h3_packed_sequence_ref2va_blocks(text_len=100, seq_len=192, **common)
+    assert int(short_pinned["seq_len"]) == int(long_pinned["seq_len"]) == 192
+    assert int(short_pinned["cu_seqlens"][1]) == int(short["cu_seqlens"][1])
+    assert int(long_pinned["cu_seqlens"][1]) == int(long["cu_seqlens"][1])
+
+
+def test_explicit_seq_len_below_used_rows_is_a_client_error():
+    """Both packers reject an under-sized pin as a 4xx, not a server error."""
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+        minimax_h3_packed_sequence,
+        minimax_h3_packed_sequence_ref2va_blocks,
+    )
+    from vllm_omni.errors import OmniClientError
+
+    with pytest.raises(OmniClientError, match="used rows"):
+        minimax_h3_packed_sequence(
+            text_len=4,
+            latent_t=2,
+            latent_h=4,
+            latent_w=6,
+            audio_t=3,
+            include_keyframe_cond=False,
+            seq_len=8,
+        )
+
+    with pytest.raises(OmniClientError, match="used rows"):
+        minimax_h3_packed_sequence_ref2va_blocks(
+            text_len=4,
+            latent_t=2,
+            latent_h=4,
+            latent_w=6,
+            audio_t=3,
+            ref_blocks=[{"kind": "audio", "ref_audio_t": 1}],
+            seq_len=8,
+        )

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -41,7 +42,7 @@ class _FakeAttention(nn.Module):
         super().__init__()
 
 
-def _small_od_config():
+def _small_od_config(*, class_name=None):
     arch = {
         "num_layers": 1,
         "token_refiner_num_layers": 1,
@@ -60,6 +61,8 @@ def _small_od_config():
         "final_adaln_out_features": 2 * 8,
         "rope_inv_freq_len": 2,
     }
+    if class_name is not None:
+        arch["_class_name"] = class_name
     return SimpleNamespace(
         tf_model_config=arch,
         parallel_config=SimpleNamespace(ulysses_degree=1),
@@ -133,6 +136,24 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
     )
 
 
+def test_explicit_native_weight_format_overrides_diffusers_class_name(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "Attention", _FakeAttention)
+    monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
+
+    model = h3.MiniMaxH3DiTModel(
+        _small_od_config(class_name="MiniMaxH3Transformer3DModel"),
+        diffusers_weights=False,
+    )
+
+    assert model._diffusers_weights is False
+
+
 class _WeightTarget(nn.Module):
     def __init__(self, loader):
         super().__init__()

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py

@@ -175,7 +175,7 @@ def _generate_joint_output(omni, config: QualityTestConfig):
         pytest.param(
             _QUALITY_CONFIG,
             id=_QUALITY_CONFIG.id,
-            marks=hardware_marks(res={"cuda": _QUALITY_CONFIG.gpu}, num_cards=2),
+            marks=hardware_marks(res={"cuda": ["H100", "B200"]}, num_cards=2),
         )
     ],
 )
@@ -186,10 +186,15 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
     common_kwargs = {
         "model": model_ref,
         "enforce_eager": True,
-        "tensor_parallel_size": 2,
         # The fused BF16 baseline only fits on H100-80GB with encoder TP.
+        # This test loads the direct FL2VA diffusion pipeline (not the
+        # disaggregated two-stage MiniMax-H3 pipeline), so configure both the
+        # DiT world and the Qwen3-VL encoder TP explicitly.
+        "tensor_parallel_size": 2,
         "text_encoder_tp_size": 2,
         "vae_use_tiling": True,
+        # MiniMax H3's VAE is implemented by code shipped with the checkpoint.
+        "trust_remote_code": True,
     }
 
     omni_bl = Omni(**common_kwargs)

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -27,6 +27,9 @@ class _SegmentMeanModel:
     assignment shows up as a numeric difference instead of passing silently.
     """
 
+    def __init__(self):
+        self.calls: list[dict[str, torch.Tensor]] = []
+
     def __call__(self, **kwargs):
         x = kwargs["x"][0]
         audio_x = kwargs["audio_x"][0]
@@ -45,6 +48,14 @@ class _SegmentMeanModel:
             pooled_audio[start:stop] = audio_x[start:stop].mean(dim=0, keepdim=True)
 
         row_timesteps = kwargs["unique_timesteps"][kwargs["inverse_indices"]].unsqueeze(-1)
+        self.calls.append(
+            {
+                "video_rows": x[img_pos].clone(),
+                "audio_rows": audio_x[audio_pos].clone(),
+                "video_timesteps": row_timesteps[img_pos, 0].clone(),
+                "audio_timesteps": row_timesteps[audio_pos, 0].clone(),
+            }
+        )
         video = (pooled_video + 0.5 * x + row_timesteps)[img_pos]
         audio = (pooled_audio + 0.5 * audio_x + row_timesteps)[audio_pos]
         if not kwargs.get("skip_mask_out_condition", False):
@@ -110,6 +121,8 @@ def _step_pipeline(model, *, packed_batch_supported: bool = True):
     from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
     pipeline.transformer = model
     pipeline.device = torch.device("cpu")
     pipeline._transformer_for_task = lambda task: model
@@ -154,9 +167,77 @@ def test_step_execution_matches_request_mode_denoise_loop():
     torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio)
 
 
-def test_batched_step_execution_matches_independent_requests():
+def test_step_execution_matches_request_mode_with_latent_edits():
+    """Masked model rows, row timesteps, and scheduler math match both paths."""
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+    from vllm_omni.diffusion.models.minimax_h3.latent_mask import MiniMaxH3LatentEdit
+
+    model = _SegmentMeanModel()
+    branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=51)
+    sigmas_video = _sigmas(5, 12.0)
+    sigmas_audio = _sigmas(5, 3.0)
+    video_clean = torch.linspace(-1.0, 1.0, video_rows.numel()).reshape_as(video_rows)
+    audio_clean = torch.linspace(1.0, -1.0, audio_rows.numel()).reshape_as(audio_rows)
+    video_mask = torch.linspace(0.0, 1.0, video_rows.shape[0])
+    video_edit = MiniMaxH3LatentEdit.from_rows(
+        video_clean,
+        0.999 * video_clean + 0.001 * video_rows,
+        video_mask,
+        video_mask,
+    )
+    audio_mask = torch.tensor([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
+    audio_edit = MiniMaxH3LatentEdit.from_rows(
+        audio_clean,
+        audio_clean,
+        audio_mask,
+        audio_mask,
+    )
+
+    reference_video, reference_audio = minimax_h3_denoise_loop(
+        model=model,
+        positive=branch,
+        initial_video_rows=video_rows,
+        initial_audio_rows=audio_rows,
+        keyframe_cond_rows=None,
+        video_edit=video_edit,
+        audio_edit=audio_edit,
+        sigmas_video=sigmas_video,
+        sigmas_audio=sigmas_audio,
+        device=torch.device("cpu"),
+    )
+    request_first_call = model.calls[0]
+    torch.testing.assert_close(request_first_call["video_rows"], video_edit.model_rows(video_rows))
+    torch.testing.assert_close(request_first_call["audio_rows"], audio_edit.model_rows(audio_rows))
+    torch.testing.assert_close(
+        request_first_call["video_timesteps"],
+        video_edit.target_timesteps(1.0 - sigmas_video[0], 0.999, sigma=sigmas_video[0]),
+    )
+    torch.testing.assert_close(
+        request_first_call["audio_timesteps"],
+        audio_edit.target_timesteps(1.0 - sigmas_audio[0], 1.0, sigma=sigmas_audio[0]),
+    )
+    model.calls.clear()
+
+    state = _make_state("req-edit", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+    state.extra[mod._STEP_VIDEO_EDIT] = video_edit
+    state.extra[mod._STEP_AUDIO_EDIT] = audio_edit
+    pipeline = _step_pipeline(model)
+    while not state.denoise_completed:
+        velocity = pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state])
+        pipeline.step_scheduler(state, velocity)
+
+    for name, expected in request_first_call.items():
+        torch.testing.assert_close(model.calls[0][name], expected)
+    torch.testing.assert_close(state.latents, reference_video)
+    torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio)
+
+
+@pytest.mark.parametrize("lock_audio", [False, True])
+def test_batched_step_execution_matches_independent_requests(lock_audio):
     """Two co-batched requests must land where they would have landed alone."""
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+    from vllm_omni.diffusion.models.minimax_h3.latent_mask import MiniMaxH3LatentEdit
 
     model = _SegmentMeanModel()
     specs = [
@@ -169,20 +250,40 @@ def test_batched_step_execution_matches_independent_requests():
 
     pipeline = _step_pipeline(model)
 
+    def make_state(request_id: str, index: int):
+        branch, video_rows, audio_rows = _make_branch(**specs[index])
+        sigmas_video, sigmas_audio = schedules[index]
+        state = _make_state(request_id, model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+        video_mask = torch.linspace(0.1 * index, 0.8 + 0.1 * index, video_rows.shape[0])
+        video_edit = MiniMaxH3LatentEdit.from_rows(
+            torch.full_like(video_rows, 10.0 + index),
+            torch.full_like(video_rows, 20.0 + index),
+            video_mask,
+            video_mask,
+        )
+        audio_mask = torch.linspace(0.2 - 0.1 * index, 1.0 - 0.1 * index, audio_rows.shape[0])
+        audio_edit = MiniMaxH3LatentEdit.from_rows(
+            torch.full_like(audio_rows, 30.0 + index),
+            torch.full_like(audio_rows, 40.0 + index),
+            audio_mask,
+            audio_mask,
+        )
+        state.extra[mod._STEP_VIDEO_EDIT] = video_edit
+        if lock_audio and index == 0:
+            branch.locked_audio_rows = audio_rows.clone()
+        else:
+            state.extra[mod._STEP_AUDIO_EDIT] = audio_edit
+        return state
+
     alone: list[tuple[torch.Tensor, torch.Tensor]] = []
-    for spec, (sigmas_video, sigmas_audio) in zip(specs, schedules):
-        branch, video_rows, audio_rows = _make_branch(**spec)
-        state = _make_state("solo", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+    for index in range(len(specs)):
+        state = make_state("solo", index)
         while not state.denoise_completed:
             pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state]))
         alone.append((state.latents, state.extra[mod._STEP_AUDIO_ROWS]))
 
-    states = []
-    for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)):
-        branch, video_rows, audio_rows = _make_branch(**spec)
-        if index == 0:
-            assert branch.used_len == branch.seq_len
-        states.append(_make_state(f"req-{index}", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio))
+    states = [make_state(f"req-{index}", index) for index in range(len(specs))]
+    assert states[0].extra[mod._STEP_BRANCH].used_len == states[0].extra[mod._STEP_BRANCH].seq_len
 
     active = list(states)
     while active:
@@ -284,15 +385,20 @@ def test_mixed_step_batch_leaves_gated_attention_dense():
     assert recorder.denoise_timestep is None
 
 
-def test_prepare_encode_seeds_runner_visible_state(monkeypatch):
+@pytest.mark.parametrize("batch_frames", [1, 33])
+def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames):
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
 
     branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8)
+    video_edit = object()
+    audio_edit = object()
     sigmas_video = _sigmas(6, 12.0)
     sigmas_audio = _sigmas(6, 3.0)
     context = {
         "height": 96,
         "width": 64,
+        "preencode_mp4": True,
+        "preencode_batch_frames": batch_frames,
         "latent_t": 2,
         "latent_h": 4,
         "latent_w": 6,
@@ -301,8 +407,17 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch):
     }
 
     pipeline = _step_pipeline(_SegmentMeanModel())
-    monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {})))
-    monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context)
+    conditioning = object()
+    monkeypatch.setattr(
+        mod.MiniMaxH3Pipeline,
+        "_extract_encoder_conditioning",
+        staticmethod(lambda _: conditioning),
+    )
+    monkeypatch.setattr(
+        mod.MiniMaxH3Pipeline,
+        "_prepare_encoder_conditioning_inputs",
+        lambda self, value, sampling: context,
+    )
     monkeypatch.setattr(
         mod.MiniMaxH3Pipeline,
         "_build_denoise_inputs",
@@ -312,6 +427,8 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch):
             "audio_rows": audio_rows,
             "cond_anchor": None,
             "audio_anchor": None,
+            "video_edit": video_edit,
+            "audio_edit": audio_edit,
             "sigmas_video": sigmas_video,
             "sigmas_audio": sigmas_audio,
         },
@@ -333,8 +450,22 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch):
     assert state.do_true_cfg is False
     torch.testing.assert_close(state.current_timestep, torch.tensor(1.0 - sigmas_video[0]))
     assert state.extra[mod._STEP_BRANCH] is branch
+    assert state.extra[mod._STEP_VIDEO_EDIT] is video_edit
+    assert state.extra[mod._STEP_AUDIO_EDIT] is audio_edit
     assert state.extra[mod._STEP_SHAPE]["height"] == 96
 
+    pipeline.od_config = SimpleNamespace()
+    monkeypatch.setattr(pipeline, "_unpack_denoised_rows", lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)))
+    calls = []
+
+    def decode_to_mp4(*args, **kwargs):
+        calls.append(kwargs)
+        return b"mp4"
+
+    monkeypatch.setattr(pipeline, "decode_to_mp4", decode_to_mp4)
+    assert pipeline.post_decode(state).output == (b"mp4", None)
+    assert calls[0]["batch_frames"] == batch_frames
+
 
 def test_prepare_encode_rejects_request_mode_only_features():
     """Multi-output and DLO have no representation in the step contract."""
@@ -426,48 +557,31 @@ def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention):
     assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False
 
 
-def test_broadcast_rank0_exception_single_rank_reraises():
-    """Single-rank execution has no group; the helper just reraises."""
-    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
-        _broadcast_rank0_exception,
+def test_locked_driving_audio_is_clean_and_unchanged_during_denoising():
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+
+    branch, video, audio = _make_branch(text_len=3, latent_t=2, latent_h=2, latent_w=2, audio_t=3, seed=7)
+    branch.locked_audio_rows = audio.clone()
+    seen = []
+
+    def model(**kwargs):
+        positions = kwargs["audio_pos_info"]["position_ids"]
+        times = kwargs["unique_timesteps"][kwargs["inverse_indices"]]
+        torch.testing.assert_close(times[positions], torch.ones_like(times[positions]))
+        torch.testing.assert_close(kwargs["audio_x"][0, positions], audio)
+        seen.append(True)
+        return torch.ones_like(video), torch.ones_like(audio)
+
+    result_video, result_audio = minimax_h3_denoise_loop(
+        model=model,
+        positive=branch,
+        initial_video_rows=video,
+        initial_audio_rows=audio,
+        keyframe_cond_rows=None,
+        sigmas_video=[1.0, 0.5, 0.0],
+        sigmas_audio=[1.0, 0.25, 0.0],
+        device=torch.device("cpu"),
     )
-    from vllm_omni.errors import OmniClientError
-
-    _broadcast_rank0_exception(None)
-    with pytest.raises(OmniClientError, match="bad ref"):
-        _broadcast_rank0_exception(OmniClientError("bad ref"))
-
-
-def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch):
-    """A rank-0 error becomes a matching client error on every other DiT rank."""
-    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
-    from vllm_omni.errors import OmniClientError
-
-    def fake_rank_world(rank):
-        return lambda: (object(), rank, 4)
-
-    def make_broadcast(rank0_payload):
-        def fake_broadcast(payload_list, *, src, group):
-            payload_list[0] = rank0_payload
-
-        return fake_broadcast
-
-    monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0))
-    err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError")
-    rank0_payload = {
-        "type": type(err).__name__,
-        "message": str(err),
-        "status_code": err.status_code,
-        "error_type": err.error_type,
-    }
-    monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload))
-    with pytest.raises(OmniClientError) as rank0_info:
-        mod._broadcast_rank0_exception(err)
-    assert rank0_info.value is err
-
-    monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2))
-    with pytest.raises(OmniClientError) as rank2_info:
-        mod._broadcast_rank0_exception(None)
-    assert rank2_info.value.status_code == 422
-    assert rank2_info.value.error_type == "UnprocessableEntityError"
-    assert "invalid reference-video file" in str(rank2_info.value)
+    assert len(seen) == 2
+    torch.testing.assert_close(result_audio, audio)
+    assert not torch.equal(result_video, video)

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_decode_revert.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for in-place decode revert and stage-boundary cache release.
+
+``_revert_decoded_inplace`` must be numerically identical to the checkpoint's
+``revert_tensor`` (denormalize → clamp → rearrange) with zero whole-video
+copies, and ``_release_stage_cache`` must force the allocator-cache release
+(or fall back to a plain empty_cache when no component cache exists).
+"""
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 as pipeline_module
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+DENORM_MEAN = (-2.11628412, -2.03571429, -1.80444444)  # -imagenet_mean/std
+DENORM_STD = (4.36681223, 4.46428571, 4.44444444)  # 1/imagenet_std
+
+
+class _FakeDenormalize:
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+
+class _FakeProcessor:
+    def __init__(self):
+        self.transform_rev = _FakeDenormalize(DENORM_MEAN, DENORM_STD)
+        self.use_3d_conv = True
+        self.revert_calls = 0
+
+    def revert_tensor(self, tensor):
+        self.revert_calls += 1
+        return _legacy_revert(tensor)
+
+
+class _FakeCheckpointModel:
+    def __init__(self):
+        self.processor = _FakeProcessor()
+
+
+def _vae(model=None):
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.model = model if model is not None else _FakeCheckpointModel()
+    return vae
+
+
+def _legacy_revert(tensor: torch.Tensor) -> torch.Tensor:
+    """The checkpoint's revert_tensor, op for op, on the CPU (B == 1)."""
+    if tensor.ndim == 4:
+        tensor = tensor.unsqueeze(2)
+    assert tensor.ndim == 5
+    # Normalize runs on the (b t) c h w rearrangement; at B == 1 that is a
+    # view of (B, C, T, H, W), so the per-channel broadcast is equivalent.
+    mean = torch.tensor(DENORM_MEAN).view(1, 3, 1, 1, 1)
+    std = torch.tensor(DENORM_STD).view(1, 3, 1, 1, 1)
+    tensor = (tensor - mean) / std
+    return tensor.clamp(0, 1)
+
+
+def test_revert_inplace_matches_legacy_bitwise():
+    decoded = torch.randn(1, 3, 7, 6, 8).clamp_(-3, 3)
+    expected = _legacy_revert(decoded.clone())
+    out = _vae()._revert_decoded_inplace(decoded.clone())
+    assert out.shape == expected.shape == (1, 3, 7, 6, 8)
+    assert torch.equal(out, expected)
+
+
+def test_revert_inplace_mutates_without_extra_copies():
+    decoded = torch.randn(1, 3, 5, 4, 4)
+    out = _vae()._revert_decoded_inplace(decoded)
+    assert out is decoded  # in-place: no whole-video copy materializes
+    assert out.is_contiguous()
+
+
+def test_revert_inplace_unsqueezes_4d_input():
+    # 4-D decoded tensors follow the (B, C, H, W) image contract: unsqueeze(2)
+    # lands T=1, matching the checkpoint's revert_tensor branch.
+    decoded = torch.randn(1, 3, 6, 8)
+    expected = _legacy_revert(decoded.clone())
+    out = _vae()._revert_decoded_inplace(decoded.clone())
+    assert torch.equal(out, expected)
+
+
+def test_revert_inplace_falls_back_without_contract():
+    vae = _vae()
+    vae.model.processor.transform_rev = type("Empty", (), {})()
+    decoded = torch.randn(1, 3, 5, 4, 4)
+    out = vae._revert_decoded_inplace(decoded)
+    assert vae.model.processor.revert_calls == 1
+    assert torch.equal(out, _legacy_revert(decoded))
+
+
+def test_revert_inplace_falls_back_on_bad_arity():
+    vae = _vae()
+    vae.model.processor.transform_rev = _FakeDenormalize((0.5, 0.5), (1.0, 1.0, 1.0))
+    vae._revert_decoded_inplace(torch.randn(1, 3, 5, 4, 4))
+    assert vae.model.processor.revert_calls == 1
+
+
+class _RecordingCache:
+    def __init__(self):
+        self.calls = []
+
+    def release_if_needed(self, *, force=False):
+        self.calls.append(force)
+        return True
+
+
+def _pipeline():
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline._dlo_component_cache = None
+    return pipeline
+
+
+def test_stage_cache_release_forces_component_cache():
+    pipeline = _pipeline()
+    cache = _RecordingCache()
+    pipeline._dlo_component_cache = cache
+    pipeline._release_stage_cache()
+    assert cache.calls == [True]
+
+
+def test_stage_cache_release_falls_back_to_empty_cache(monkeypatch):
+    released = []
+    monkeypatch.setattr(
+        pipeline_module.current_omni_platform,
+        "empty_cache",
+        lambda: released.append(True),
+        raising=True,
+    )
+    pipeline = _pipeline()
+    pipeline._release_stage_cache()
+    assert released == [True]

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_encode_prep.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for the streaming uint8 preparation in the MiniMax-H3 video VAE.
+
+``_stream_prepare_video_tensor`` must be numerically identical to the
+checkpoint's whole-video numpy path (``convert_numpy_to_tensor`` →
+``transform_tensor`` → ``transpose``) while padding on the host, so the
+device-side copies never materialize.
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+CLIP_LENGTH = 4
+
+
+class _FakeNormalize:
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+
+class _FakeProcessor:
+    # Read by the stream prep via getattr(..., False); annotation only so
+    # tests can attach it per case.
+    isolated_last_frame: bool
+
+    def __init__(self):
+        self.transform = _FakeNormalize(IMAGENET_MEAN, IMAGENET_STD)
+
+
+class _FakeCheckpointModel:
+    def __init__(self, **overrides):
+        self.clip_length = CLIP_LENGTH
+        self.isolated_first_frame = False
+        self.frame_pre_padding = 0
+        self.processor = _FakeProcessor()
+        for key, value in overrides.items():
+            setattr(self, key, value)
+
+
+def _vae(model=None):
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.model = model if model is not None else _FakeCheckpointModel()
+    return vae
+
+
+def _legacy_reference(frames: np.ndarray) -> torch.Tensor:
+    """The checkpoint's numpy path, op for op, on the CPU.
+
+    convert_numpy_to_tensor: astype(fp32) → permute → ÷255
+    transform_tensor: Normalize = (x - mean) / std
+    encode_videos: transpose(0, 1)
+    """
+    tensor = torch.from_numpy(frames.astype(np.float32)).permute(0, 3, 1, 2)
+    tensor = tensor / 255.0
+    mean = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1)
+    std = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1)
+    tensor = (tensor - mean) / std
+    return tensor.transpose(0, 1).contiguous()
+
+
+def _video(num_frames, height=6, width=8, seed=0):
+    rng = np.random.default_rng(seed)
+    return rng.integers(0, 256, size=(num_frames, height, width, 3), dtype=np.uint8)
+
+
+def _pad_frames(frames: np.ndarray, pad: int) -> np.ndarray:
+    if not pad:
+        return frames
+    return np.concatenate([frames, np.repeat(frames[-1:], pad, axis=0)])
+
+
+def test_stream_prep_matches_legacy_path_bitwise():
+    frames = _video(10)
+    out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out is not None
+    assert out.shape == (3, 12, 6, 8)  # pad = (0 - 10) % 4 == 2
+    assert out.dtype == torch.float32
+    assert torch.equal(out, _legacy_reference(_pad_frames(frames, 2)))
+
+
+def test_stream_prep_isolated_first_frame_without_last_frame_skips_pad():
+    # Asymmetric checkpoint (isolated first frame, no isolated last frame):
+    # trim-stable lengths are k*clip + 0, which never satisfy the
+    # encode_temporal offset of 1, so the remote re-pads through its own
+    # whole-video cat regardless. Padding here would only upload frames the
+    # remote trim discards, so the legacy frame count is kept.
+    model = _FakeCheckpointModel(isolated_first_frame=True)
+    frames = _video(10)
+    out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out is not None
+    assert out.shape[1] == 10
+    assert torch.equal(out, _legacy_reference(frames))
+
+
+def test_stream_prep_symmetric_isolated_frames_pad_to_trim_stable():
+    # isolated first + last frames: offset == tail % clip == 1, so padding
+    # to the next trim-stable length makes the remote trim a no-op and keeps
+    # encode_temporal from re-padding: T=10, clip=4, tail=1 → 3*4+1 == 13.
+    model = _FakeCheckpointModel(isolated_first_frame=True)
+    model.processor.isolated_last_frame = True
+    frames = _video(10)
+    out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out is not None
+    assert out.shape[1] == 13
+    assert out.shape[1] % 4 == 1
+    assert torch.equal(out, _legacy_reference(_pad_frames(frames, 3)))
+
+
+def test_stream_prep_uses_processor_align_when_available():
+    calls = []
+
+    class _AlignedProcessor(_FakeProcessor):
+        def align_video_length(self, video_length, mode="pad", granularity="chunk"):
+            calls.append((video_length, mode, granularity))
+            return 5
+
+    model = _FakeCheckpointModel()
+    model.processor = _AlignedProcessor()
+    frames = _video(10)
+    out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out is not None
+    assert calls == [(10, "pad", "chunk")]
+    assert out.shape[1] == 15
+    assert torch.equal(out, _legacy_reference(_pad_frames(frames, 5)))
+
+
+def test_stream_prep_pads_with_last_frame():
+    frames = _video(2)  # pad = (0 - 2) % 4 == 2
+    out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    padded = np.concatenate([frames, np.repeat(frames[-1:], 2, axis=0)])
+    assert torch.equal(out, _legacy_reference(padded))
+
+
+def test_stream_prep_aligned_input_needs_no_pad():
+    frames = _video(8)  # 8 % 4 == 0
+    out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out.shape[1] == 8
+    assert torch.equal(out, _legacy_reference(frames))
+
+
+@pytest.mark.parametrize(
+    "frames",
+    [
+        _video(5).astype(np.float32),  # not uint8
+        _video(5)[:, :, :, :1],  # single channel
+        np.zeros((5, 6, 8, 4), dtype=np.uint8),  # four channels
+        np.zeros((5, 6, 8), dtype=np.uint8),  # missing channel axis
+        np.zeros((0, 6, 8, 3), dtype=np.uint8),  # empty
+    ],
+)
+def test_stream_prep_rejects_unsupported_inputs(frames):
+    assert _vae()._stream_prepare_video_tensor(frames, torch.device("cpu")) is None
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        _FakeCheckpointModel(clip_length=None),
+        _FakeCheckpointModel(clip_length=0),
+        type("NoProcessor", (), {"clip_length": CLIP_LENGTH})(),
+        type(
+            "NoTransform",
+            (),
+            {"clip_length": CLIP_LENGTH, "processor": type("P", (), {})()},
+        )(),
+    ],
+)
+def test_stream_prep_rejects_missing_checkpoint_contract(model):
+    frames = _video(5)
+    assert _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu")) is None
+
+
+def test_stream_prep_rejects_bad_norm_arity():
+    model = _FakeCheckpointModel()
+    model.processor.transform = _FakeNormalize((0.5, 0.5), IMAGENET_STD)
+    assert _vae(model)._stream_prepare_video_tensor(_video(5), torch.device("cpu")) is None
+
+
+class _RecordingModel(_FakeCheckpointModel):
+    """Feeds encode_video a well-shaped latent and records its input."""
+
+    def __init__(self):
+        super().__init__()
+        self.received = None
+        self.latent = torch.randn(1, 4, 3, 4, 4)  # (B, C, T', H', W'), even H/W
+
+    def encode_videos(self, videos, use_fp16_latent=False, **kwargs):
+        self.received = videos
+        assert use_fp16_latent is True
+        return [self.latent]
+
+
+def _vae_for_encode_video():
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.model = _RecordingModel()
+    vae.device_module = torch
+    vae.config_dict = {
+        "latent_channels": 4,
+        "latents_mean": [0.0] * 4,
+        "latents_std": [1.0] * 4,
+    }
+    # A plain tensor (not nn.Parameter): nn.Module.__setattr__ refuses
+    # Parameters on instances built via object.__new__.
+    vae._dummy_parameter = torch.zeros(1, dtype=torch.float32)
+    vae.parameters = lambda: iter([vae._dummy_parameter])
+    return vae
+
+
+def test_encode_video_streams_uint8_numpy(monkeypatch):
+    monkeypatch.delenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", raising=False)
+    vae = _vae_for_encode_video()
+    rows, shape = vae.encode_video(_video(10))
+    assert isinstance(vae.model.received, list)
+    assert isinstance(vae.model.received[0], torch.Tensor)
+    assert vae.model.received[0].shape == (3, 12, 6, 8)
+    assert rows.shape[0] == shape[0] * shape[1] * shape[2] // 4  # patchify rows
+
+
+def test_encode_video_legacy_env_escapes_to_checkpoint_path(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", "1")
+    vae = _vae_for_encode_video()
+    frames = _video(10)
+    vae.encode_video(frames)
+    assert vae.model.received is frames  # untouched numpy, checkpoint converts
+
+
+def test_encode_video_falls_back_when_preparation_unsupported(monkeypatch):
+    monkeypatch.delenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", raising=False)
+    vae = _vae_for_encode_video()
+    vae.model.clip_length = None  # contract not discoverable
+    frames = _video(10)
+    vae.encode_video(frames)
+    assert vae.model.received is frames

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py

@@ -0,0 +1,473 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for the pinned tiled-VAE gather buffer in the MiniMax-H3 video VAE.
+
+The checkpoint's ``_all_gather_tiled_results`` allocates its gather output on
+every call, so with a per-request ``empty_cache()`` the collective lands on a
+new device address every request and XCCL's per-address registrations pile up.
+On XPU the adapter replaces that method with an equal-shape gather into a
+buffer it keeps alive. These tests pin the properties that make it a fix
+rather than a rewrite: the address is stable, the tiles come back
+bit-identical and in the same order, the caller owns what it gets back, and no
+other device is touched.
+"""
+
+import logging
+import sys
+import types
+from typing import Any
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+TILE = (1, 3, 2, 4, 4)
+
+
+class _CheckpointGatherSentinel:
+    """Stands in for the checkpoint's own bound method.
+
+    Identity is the point: a test asserts this exact object is still installed
+    on the devices the adapter must leave alone.
+    """
+
+    def __call__(self, tasks, num_tiles):  # pragma: no cover - never invoked
+        raise AssertionError("the checkpoint gather must not run in these tests")
+
+
+class _FakeCheckpointModel:
+    def __init__(self):
+        self._all_gather_tiled_results = _CheckpointGatherSentinel()
+
+
+def _local_tiles(all_tiles, num_tiles, rank, sp_size):
+    """One rank's round-robin share, as ``klvae.py:275`` computes it."""
+    return [all_tiles[index] for index in range(rank, num_tiles, sp_size)]
+
+
+def _reference_gather(all_tiles, num_tiles, sp_size):
+    """The checkpoint's variable-shape list gather, reimplemented verbatim.
+
+    Written from the control flow of ``_all_gather_tiled_results`` in the
+    checkpoint's ``klvae.py`` (lines 252-273) -- stack each rank's share, gather
+    the list, walk it back into global tile order -- so that it is an
+    independent oracle rather than a restatement of the code under test.
+    """
+    gathered = [torch.stack(_local_tiles(all_tiles, num_tiles, rank, sp_size), dim=0) for rank in range(sp_size)]
+    results = [None] * num_tiles
+    for rank, rank_tensors in enumerate(gathered):
+        for k in range(rank_tensors.shape[0]):
+            global_index = k * sp_size + rank
+            if global_index >= num_tiles:
+                break
+            results[global_index] = rank_tensors[k]
+    return results
+
+
+def _peer_stacks(all_tiles, num_tiles, sp_size):
+    """What every rank hands to the collective, padded the way the adapter pads."""
+    max_tasks = -(-num_tiles // sp_size)
+    stacks = []
+    for rank in range(sp_size):
+        share = _local_tiles(all_tiles, num_tiles, rank, sp_size)
+        stacked = share[0].new_zeros((max_tasks, *share[0].shape))
+        stacked[: len(share)] = torch.stack(share, dim=0)
+        stacks.append(stacked)
+    return stacks
+
+
+def _tiles(num_tiles, seed=0, dtype=torch.float32):
+    generator = torch.Generator().manual_seed(seed)
+    return [torch.rand(TILE, generator=generator).to(dtype) for _ in range(num_tiles)]
+
+
+def _workspace_bytes(num_tiles, sp_size, dtype=torch.float32):
+    """Bytes one gather needs: ``sp_size`` copies of the padded local stack."""
+    max_tasks = -(-num_tiles // sp_size)
+    element_size = torch.empty(0, dtype=dtype).element_size()
+    numel = 1
+    for dim in TILE:
+        numel *= dim
+    return sp_size * max_tasks * numel * element_size
+
+
+@pytest.fixture
+def fake_collective(monkeypatch):
+    """Replace the collective with a single-process stand-in.
+
+    ``all_gather_into_tensor`` is the only distributed call on this path, so
+    filling the output buffer from a caller-supplied peer list exercises the
+    real padding, keying, addressing and unpacking code without a process
+    group.
+    """
+    calls: list[dict[str, Any]] = []
+
+    def fake_all_gather_into_tensor(output, input_tensor, group=None, async_op=False):
+        peers = calls[-1]["peers"]
+        assert output.shape == (len(peers), *input_tensor.shape)
+        calls[-1]["sent"] = input_tensor.clone()
+        calls[-1]["buf_ptr"] = output.data_ptr()
+        for rank, peer in enumerate(peers):
+            output[rank].copy_(peer)
+
+    monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather_into_tensor)
+    return calls
+
+
+class _Records(logging.Handler):
+    """vllm's logger sets ``propagate=False``, so caplog never sees these."""
+
+    def __init__(self):
+        super().__init__(level=logging.DEBUG)
+        self.messages = []
+        self.records = []
+
+    def emit(self, record):
+        self.messages.append(record.getMessage())
+        self.records.append((record.levelno, record.getMessage()))
+
+
+@pytest.fixture
+def gather_log():
+    handler = _Records()
+    logger = vae_module.logger
+    previous = logger.level
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+    try:
+        yield handler
+    finally:
+        logger.removeHandler(handler)
+        logger.setLevel(previous)
+
+
+def _vae(sp_size, sp_rank=0):
+    """An adapter instance with the override installed, no checkpoint needed."""
+    state = {
+        "sp_size": sp_size,
+        "sp_rank": sp_rank,
+        "sp_enabled": sp_size > 1,
+        "sp_process_group": object(),
+    }
+    module = types.ModuleType("fake_gather_ckpt.parallel")
+    module.get_parallel_state = lambda: state  # type: ignore[attr-defined]
+    sys.modules.setdefault("fake_gather_ckpt", types.ModuleType("fake_gather_ckpt"))
+    sys.modules["fake_gather_ckpt.parallel"] = module
+
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.remote = type("Remote", (), {"__module__": "fake_gather_ckpt.klvae"})()
+    vae.model = _FakeCheckpointModel()
+    vae._device_target = torch.device("xpu")
+    vae._tile_gather_workspace = None
+    vae._tile_gather_stats = {"hits": 0, "allocs": 0, "workspace_bytes": 0}
+    vae._checkpoint_tile_gather = None
+    vae._install_persistent_tile_gather()
+    return vae
+
+
+def _run(vae, calls, all_tiles, num_tiles, sp_size, sp_rank):
+    calls.append({"peers": _peer_stacks(all_tiles, num_tiles, sp_size)})
+    return vae.model._all_gather_tiled_results(_local_tiles(all_tiles, num_tiles, sp_rank, sp_size), num_tiles)
+
+
+# --------------------------------------------------------------------------
+# The address is stable across calls. This is what the accumulation is about:
+# drop the workspace and every call registers a new receive address. And the
+# workspace is one grow-only block, so varying request geometries cannot keep
+# adding resident buffers the way a per-geometry cache would.
+# --------------------------------------------------------------------------
+
+
+def test_the_landing_buffer_address_is_stable_across_calls(fake_collective):
+    vae = _vae(4)
+    pointers = []
+    for request in range(3):
+        _run(vae, fake_collective, _tiles(8, seed=request), 8, 4, 0)
+        pointers.append(fake_collective[-1]["buf_ptr"])
+
+    assert len(set(pointers)) == 1, f"gather buffer moved across calls: {pointers}"
+    assert vae._tile_gather_stats == {
+        "hits": 2,
+        "allocs": 1,
+        "workspace_bytes": _workspace_bytes(8, 4),
+    }
+    assert vae._tile_gather_workspace.numel() == _workspace_bytes(8, 4)
+    assert vae._tile_gather_workspace.dtype is torch.uint8
+
+
+def test_distinct_shapes_share_the_one_workspace(fake_collective):
+    """One decode gathers several temporal clips, and the head/tail clips differ."""
+    vae = _vae(4)
+    _run(vae, fake_collective, _tiles(8), 8, 4, 0)
+    first = fake_collective[-1]["buf_ptr"]
+    _run(vae, fake_collective, _tiles(12), 12, 4, 0)
+    second = fake_collective[-1]["buf_ptr"]
+    _run(vae, fake_collective, _tiles(8, seed=5), 8, 4, 0)
+    third = fake_collective[-1]["buf_ptr"]
+
+    # The larger geometry did not fit, so the block was replaced -- not added to.
+    assert second == third
+    assert vae._tile_gather_stats["allocs"] == 2
+    assert vae._tile_gather_stats["hits"] == 1
+    assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(12, 4)
+    # ``first`` is deliberately not compared: the replaced block is freed, so
+    # the allocator is free to hand its address back for the larger one.
+    assert isinstance(first, int)
+
+
+def test_workspace_is_bounded_across_geometries(fake_collective):
+    """Alternating geometries must not each retain a buffer of their own."""
+    sp_size = 4
+    large, small, largest = 12, 8, 20
+    vae = _vae(sp_size)
+
+    pointers = []
+    for round_index in range(3):
+        for num_tiles in (large, small):
+            _run(vae, fake_collective, _tiles(num_tiles, seed=round_index), num_tiles, sp_size, 0)
+            pointers.append(fake_collective[-1]["buf_ptr"])
+
+    assert len(pointers) == 6
+    assert len(set(pointers)) == 1, f"workspace moved across geometries: {pointers}"
+    assert vae._tile_gather_stats["allocs"] == 1
+    assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(large, sp_size)
+    assert vae._tile_gather_workspace.numel() == _workspace_bytes(large, sp_size)
+
+    # A geometry that does not fit grows the single block; it does not add one.
+    _run(vae, fake_collective, _tiles(largest), largest, sp_size, 0)
+    grown = fake_collective[-1]["buf_ptr"]
+    assert vae._tile_gather_stats["allocs"] == 2
+    assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(largest, sp_size)
+    assert vae._tile_gather_workspace.numel() == _workspace_bytes(largest, sp_size)
+
+    # ...and the geometries seen before it now fit, so nothing reallocates again.
+    for num_tiles in (large, small):
+        _run(vae, fake_collective, _tiles(num_tiles, seed=7), num_tiles, sp_size, 0)
+        assert fake_collective[-1]["buf_ptr"] == grown
+    assert vae._tile_gather_stats["allocs"] == 2
+    assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(largest, sp_size)
+
+
+def test_workspace_grows_only_upward(fake_collective):
+    """Resident bytes track the largest geometry seen, not the sum of them."""
+    sp_size = 4
+    small, large = 8, 12
+    vae = _vae(sp_size)
+
+    _run(vae, fake_collective, _tiles(small), small, sp_size, 0)
+    assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(small, sp_size)
+
+    _run(vae, fake_collective, _tiles(large), large, sp_size, 0)
+
+    assert vae._tile_gather_stats["allocs"] == 2
+    assert vae._tile_gather_stats["hits"] == 0
+    assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(large, sp_size)
+    assert vae._tile_gather_stats["workspace_bytes"] < _workspace_bytes(small, sp_size) + _workspace_bytes(
+        large, sp_size
+    )
+    assert vae._tile_gather_workspace.numel() == _workspace_bytes(large, sp_size)
+
+
+def test_mixed_dtype_shares_one_workspace(fake_collective):
+    """The workspace is bytes, so a narrower dtype reuses it instead of adding one."""
+    sp_size, num_tiles = 4, 8
+    vae = _vae(sp_size)
+    fp32_bytes = _workspace_bytes(num_tiles, sp_size, dtype=torch.float32)
+
+    pointers = []
+    for round_index in range(2):
+        for dtype in (torch.float32, torch.bfloat16):
+            tiles = _tiles(num_tiles, seed=round_index, dtype=dtype)
+            produced = _run(vae, fake_collective, tiles, num_tiles, sp_size, 0)
+            pointers.append(fake_collective[-1]["buf_ptr"])
+            assert all(tile.dtype is dtype for tile in produced)
+            for index, tile in enumerate(produced):
+                assert torch.equal(tile, tiles[index])
+
+    assert len(set(pointers)) == 1, f"workspace moved across dtypes: {pointers}"
+    assert vae._tile_gather_stats["allocs"] == 1
+    assert vae._tile_gather_stats["hits"] == 3
+    assert vae._tile_gather_stats["workspace_bytes"] == fp32_bytes
+    assert vae._tile_gather_workspace.numel() == fp32_bytes
+
+
+# --------------------------------------------------------------------------
+# Equivalence with the checkpoint's own gather, padded branch included. A
+# reordered tile grid is silent at runtime, so order is asserted as well.
+# --------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("num_tiles", [5, 8, 9, 28])
+@pytest.mark.parametrize("sp_rank", [0, 3])
+def test_gathered_tiles_match_the_checkpoint_gather(fake_collective, num_tiles, sp_rank):
+    sp_size = 4
+    all_tiles = _tiles(num_tiles, seed=num_tiles)
+    vae = _vae(sp_size, sp_rank=sp_rank)
+
+    produced = _run(vae, fake_collective, all_tiles, num_tiles, sp_size, sp_rank)
+    expected = _reference_gather(all_tiles, num_tiles, sp_size)
+
+    assert len(produced) == num_tiles
+    assert all(tile is not None for tile in produced)
+    for index, (got, want) in enumerate(zip(produced, expected)):
+        assert torch.equal(got, want), f"tile {index} differs from the checkpoint gather"
+    # ...and it is the tile that belongs at that index, not a permutation that
+    # happens to agree with an equally permuted oracle.
+    for index, tile in enumerate(produced):
+        assert torch.equal(tile, all_tiles[index])
+
+
+def test_the_padded_branch_is_actually_exercised(fake_collective):
+    """9 tiles over 4 ranks: ranks 1-3 are one short, so the padding runs."""
+    sp_size, num_tiles = 4, 9
+    all_tiles = _tiles(num_tiles, seed=9)
+    assert len(_local_tiles(all_tiles, num_tiles, 1, sp_size)) < -(-num_tiles // sp_size)
+
+    vae = _vae(sp_size, sp_rank=1)
+    produced = _run(vae, fake_collective, all_tiles, num_tiles, sp_size, 1)
+
+    sent = fake_collective[-1]["sent"]
+    assert sent.shape[0] == 3
+    assert torch.equal(sent[2], torch.zeros_like(sent[2]))
+    for index, tile in enumerate(produced):
+        assert torch.equal(tile, all_tiles[index])
+
+
+def test_an_empty_local_share_still_raises(fake_collective):
+    """A rank with no tiles fails loudly, exactly as the checkpoint does."""
+    vae = _vae(4)
+    with pytest.raises(ValueError, match="empty tasks on sp rank"):
+        vae.model._all_gather_tiled_results([], 8)
+
+
+def test_more_tiles_than_round_robin_allows_is_refused(fake_collective):
+    """The unpacking arithmetic is only valid for the checkpoint's ownership."""
+    vae = _vae(4)
+    with pytest.raises(ValueError, match="round-robin"):
+        vae.model._all_gather_tiled_results(_tiles(5), 8)
+
+
+# --------------------------------------------------------------------------
+# Ownership: the buffer is overwritten by the next gather, so returned tiles
+# must not alias it.
+# --------------------------------------------------------------------------
+
+
+def test_returned_tiles_do_not_alias_the_pinned_buffer(fake_collective):
+    sp_size, num_tiles = 4, 8
+    first_tiles = _tiles(num_tiles, seed=1)
+    vae = _vae(sp_size)
+    produced = _run(vae, fake_collective, first_tiles, num_tiles, sp_size, 0)
+
+    buffer_storage = vae._tile_gather_workspace.untyped_storage().data_ptr()
+    for tile in produced:
+        assert tile.untyped_storage().data_ptr() != buffer_storage
+
+    # The next gather reuses that memory; the first call's tiles must survive it.
+    _run(vae, fake_collective, _tiles(num_tiles, seed=2), num_tiles, sp_size, 0)
+    for index, tile in enumerate(produced):
+        assert torch.equal(tile, first_tiles[index])
+
+
+# --------------------------------------------------------------------------
+# Non-XPU behaviour is unchanged: the checkpoint's own method stays installed.
+# --------------------------------------------------------------------------
+
+
+def _construct(monkeypatch, device, model):
+    """Run the real ``__init__`` against a stub checkpoint component."""
+
+    class _Remote(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.zeros(1))
+            self.model = model
+
+    monkeypatch.setattr(
+        vae_module,
+        "_load_component_config",
+        lambda path: {"latent_channels": 1, "latents_mean": [0.0], "latents_std": [1.0]},
+    )
+    # **kwargs: the loader gained a trust_remote_code argument at some point,
+    # and this stub does not care either way.
+    monkeypatch.setattr(vae_module, "_load_remote_component", lambda path, config, **kwargs: _Remote())
+    # The adapter only stores this handle; resolving it for real needs an
+    # accelerator, which a CPU test runner does not have.
+    monkeypatch.setattr(torch, "get_device_module", lambda *args, **kwargs: types.SimpleNamespace())
+    return MiniMaxH3VideoVAE("unused", device=torch.device(device))
+
+
+@pytest.mark.parametrize("device", ["cpu", "meta"])
+def test_non_xpu_devices_keep_the_checkpoint_gather(monkeypatch, device):
+    model = _FakeCheckpointModel()
+    original = model._all_gather_tiled_results
+
+    vae = _construct(monkeypatch, device, model)
+
+    assert vae._tile_gather_reuse_enabled() is False
+    assert vae.model._all_gather_tiled_results is original
+    assert vae._checkpoint_tile_gather is None
+
+
+@pytest.mark.parametrize(
+    ("device", "enabled"),
+    [("xpu", True), ("cpu", False), ("cuda", False), ("meta", False)],
+)
+def test_only_xpu_asks_for_the_pinned_buffer(device, enabled):
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae._device_target = torch.device(device)
+    assert vae._tile_gather_reuse_enabled() is enabled
+
+
+def test_construction_installs_the_override_when_the_device_asks_for_it(monkeypatch):
+    """Close the loop: ``__init__`` must consult the predicate, not ignore it."""
+    monkeypatch.setattr(MiniMaxH3VideoVAE, "_tile_gather_reuse_enabled", lambda self: True)
+    model = _FakeCheckpointModel()
+    original = model._all_gather_tiled_results
+
+    vae = _construct(monkeypatch, "cpu", model)
+
+    assert vae.model._all_gather_tiled_results is not original
+    assert vae._checkpoint_tile_gather is original
+
+
+# --------------------------------------------------------------------------
+# The decision-point receipt carries quantities, not just presence: a line
+# that only says "installed" cannot tell reuse from reallocation.
+# --------------------------------------------------------------------------
+
+
+def test_the_receipt_line_reports_hits_allocs_and_the_address(fake_collective, gather_log):
+    vae = _vae(4)
+    install = [
+        (level, line) for level, line in gather_log.records if "persistent tile gather installed device=xpu" in line
+    ]
+    # The one-shot install line is what an operator needs at info level.
+    assert install and all(level == logging.INFO for level, _ in install)
+    gather_log.messages.clear()
+    gather_log.records.clear()
+
+    _run(vae, fake_collective, _tiles(8), 8, 4, 0)
+    first_ptr = fake_collective[-1]["buf_ptr"]
+    _run(vae, fake_collective, _tiles(8, seed=3), 8, 4, 0)
+
+    records = [(level, line) for level, line in gather_log.records if "[H3_VAE_GATHER] reuse=" in line]
+    # The per-gather line fires once per decoder tile batch, so it is debug.
+    assert all(level == logging.DEBUG for level, _ in records)
+    lines = [line for _, line in records]
+    assert len(lines) == 2
+    assert "reuse=alloc" in lines[0] and "hits=0 allocs=1" in lines[0]
+    assert "reuse=hit" in lines[1] and "hits=1 allocs=1" in lines[1]
+    # Resident bytes are reported, not just the bytes this one gather needed.
+    needed = _workspace_bytes(8, 4) / (1024.0 * 1024.0)
+    for line in lines:
+        assert f"need_mib={needed:.2f}" in line
+        assert f"workspace_mib={needed:.2f}" in line
+    assert f"buf_ptr=0x{first_ptr:x}" in lines[0]
+    assert f"buf_ptr=0x{first_ptr:x}" in lines[1]

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py

@@ -198,7 +198,7 @@ def test_video_vae_installs_exact_optimizations(monkeypatch):
     monkeypatch.setattr(
         vae_module,
         "_load_remote_component",
-        lambda _path, _config: FakeRemote(),
+        lambda _path, _config, **_kwargs: FakeRemote(),
     )
     install = Mock(return_value=True)
     monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install)

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_split_residency.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for the split encoder/decoder residency of the MiniMax-H3 VAE.
+
+The encode half (CNN encoder + quant_conv) and the decode half
+(post_quant_conv + ViT decoder) stage independently: loading one half must
+rebind only its own parameters while the other half stays on its pinned
+CPU master.
+"""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.minimax_h3.vae import (
+    MiniMaxH3VideoVAE,
+    _VideoVAEPartProxy,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeRemoteModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Linear(8, 8)
+        self.quant_conv = nn.Linear(4, 4)
+        self.post_quant_conv = nn.Linear(4, 4)
+        self.decoder = nn.Linear(16, 16)
+
+
+class _FakeRemote(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = _FakeRemoteModel()
+
+
+def _vae() -> MiniMaxH3VideoVAE:
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    nn.Module.__init__(vae)  # bare instance: enable Module attribute assignment
+    vae._device_target = torch.device("cpu")
+    vae.remote = _FakeRemote()
+    vae.model = vae.remote.model
+    vae._stager = None
+    vae._encoder_stager = None
+    vae._decoder_stager = None
+    vae._build_residency_stagers(torch.device("cpu"))
+    vae.encoder_component = _VideoVAEPartProxy(vae, "encoder")
+    vae.decoder_component = _VideoVAEPartProxy(vae, "decoder")
+    return vae
+
+
+def _ptr(param: torch.nn.Parameter) -> int:
+    return param.untyped_storage().data_ptr()
+
+
+def test_split_stagers_snapshot_disjoint_halves():
+    vae = _vae()
+    enc_master = _ptr(vae.model.encoder.weight)
+    dec_master = _ptr(vae.model.decoder.weight)
+    pqc_master = _ptr(vae.model.post_quant_conv.weight)
+    assert enc_master != dec_master
+
+    vae._load_part_to_device("encoder")
+    # The encoder half is rebound to freshly allocated staging storage...
+    assert _ptr(vae.model.encoder.weight) != enc_master
+    assert _ptr(vae.model.quant_conv.weight) != enc_master
+    # ...while every decode-half tensor stays on its own CPU master.
+    assert _ptr(vae.model.decoder.weight) == dec_master
+    assert _ptr(vae.model.post_quant_conv.weight) == pqc_master
+
+    vae._offload_part_to_cpu("encoder")
+    assert _ptr(vae.model.encoder.weight) == enc_master
+
+
+def test_part_load_is_idempotent_and_independent():
+    vae = _vae()
+    vae._load_part_to_device("decoder")
+    staged_ptr = _ptr(vae.model.decoder.weight)
+    vae._load_part_to_device("decoder")
+    assert _ptr(vae.model.decoder.weight) == staged_ptr
+    vae._offload_part_to_cpu("decoder")
+
+
+def test_whole_component_load_and_offload_cover_both_halves():
+    vae = _vae()
+    enc_master = _ptr(vae.model.encoder.weight)
+    dec_master = _ptr(vae.model.decoder.weight)
+    vae.load_to_device()
+    assert _ptr(vae.model.encoder.weight) != enc_master
+    assert _ptr(vae.model.decoder.weight) != dec_master
+    vae.offload_to_cpu()
+    assert _ptr(vae.model.encoder.weight) == enc_master
+    assert _ptr(vae.model.decoder.weight) == dec_master
+
+
+def test_part_values_survive_a_load_offload_cycle():
+    vae = _vae()
+    expected = vae.model.encoder.weight.detach().clone()
+    vae._load_part_to_device("encoder")
+    assert torch.equal(vae.model.encoder.weight.detach(), expected)
+    vae._offload_part_to_cpu("encoder")
+    assert torch.equal(vae.model.encoder.weight.detach(), expected)
+
+
+def test_proxy_routes_to_the_matching_half():
+    vae = _vae()
+    enc_master = _ptr(vae.model.encoder.weight)
+    dec_master = _ptr(vae.model.decoder.weight)
+    vae.encoder_component.load_to_device()
+    assert _ptr(vae.model.encoder.weight) != enc_master
+    assert _ptr(vae.model.decoder.weight) == dec_master
+    vae.encoder_component.offload_to_cpu()
+    assert _ptr(vae.model.encoder.weight) == enc_master
+
+
+def test_proxy_does_not_register_the_adapter_as_submodule():
+    vae = _vae()
+    proxy = vae.encoder_component
+    assert proxy._vae is vae
+    # The back-reference must not recurse through module traversal.
+    assert all(child is not vae for child in proxy.children())
+    assert len(list(proxy.parameters())) == 0
+
+
+def test_component_on_device_unwraps_proxy_for_sequential_offload(monkeypatch):
+    """Model-level CPU offload hooks the real video_vae, never the part
+    proxies: _component_on_device must hand the hooked module to the
+    sequential context, or _get_sequential_offload_hook raises on the
+    hook-less proxy on the first encode/decode of every request."""
+    from contextlib import contextmanager
+
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    vae = _vae()
+    seen: list[nn.Module] = []
+
+    @contextmanager
+    def fake_sequential_offload_component(component):
+        seen.append(component)
+        yield
+
+    monkeypatch.setattr(pipeline_module, "sequential_offload_component", fake_sequential_offload_component)
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline._model_cpu_offload_modules = [vae]
+
+    with pipeline._component_on_device(vae.encoder_component):
+        pass
+    with pipeline._component_on_device(vae.decoder_component):
+        pass
+    # A plain component passes through unchanged.
+    with pipeline._component_on_device(vae):
+        pass
+    # Mocks auto-create any attribute, so the unwrap must be a type check:
+    # a Mock component must reach the sequential context as itself.
+    mock_vae = Mock()
+    pipeline._model_cpu_offload_modules = [mock_vae]
+    with pipeline._component_on_device(mock_vae):
+        pass
+
+    assert seen == [vae, vae, vae, mock_vae]
+
+
+def test_sequential_offload_target_is_the_hooked_module():
+    from vllm_omni.diffusion.offloader.sequential_backend import _get_sequential_offload_hook
+
+    vae = _vae()
+    proxy = vae.encoder_component
+    # The pre-fix failure mode: the proxy carries no hook registry.
+    with pytest.raises(RuntimeError, match="no sequential offload hook"):
+        _get_sequential_offload_hook(proxy)
+    assert proxy.sequential_offload_target is vae
+
+
+def test_cache_retention_forwarded_to_both_stagers():
+    vae = _vae()
+    cache = object()
+    vae.set_omni_component_cache(cache)
+    assert vae._encoder_stager.cache_retention is cache
+    assert vae._decoder_stager.cache_retention is cache
+
+
+def test_unknown_part_raises():
+    vae = _vae()
+    with pytest.raises(ValueError, match="unknown video VAE part"):
+        vae._part_stager("middle")
+
+
+def test_missing_checkpoint_structure_falls_back_to_whole_stager():
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    nn.Module.__init__(vae)
+    vae._device_target = torch.device("cpu")
+
+    class BareRemote(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = nn.Linear(4, 4)  # no encoder/decoder submodules
+
+    vae.remote = BareRemote()
+    vae.model = vae.remote.model
+    vae._stager = None
+    vae._encoder_stager = None
+    vae._decoder_stager = None
+    vae._build_residency_stagers(torch.device("cpu"))
+    assert vae._stager is not None
+    assert vae._encoder_stager is None
+    assert vae._decoder_stager is None
+    # Part loads fall back to the whole-module stager.
+    vae._load_part_to_device("decoder")
+    assert vae._stager.loaded
+    vae._offload_part_to_cpu("decoder")
+    assert not vae._stager.loaded

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_temporal_patches.py

@@ -0,0 +1,335 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for the temporal VAE forks and the uint8 output transfer.
+
+The encode fork must feed bit-identical frames to the checkpoint's chunk
+loop (padding moves from a whole-video cat into the tail chunk), the decode
+fork must produce the same video as the legacy fp32-buffer chain
+(revert -> clamp -> *255 -> round -> uint8), and the output transfer must
+match the previous ``.to(uint8, contiguous)`` conversion.
+"""
+
+from typing import Any, cast
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+    _prepare_minimax_h3_video_output,
+)
+from vllm_omni.diffusion.models.minimax_h3.vae_temporal import (
+    _decode_temporal_streaming_uint8,
+    _encode_temporal_tail_pad,
+    install_temporal_stream_patches,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+CLIP = 4
+
+
+# ---------------------------------------------------------------------------
+# encode_temporal fork
+# ---------------------------------------------------------------------------
+
+
+class _EncodeModel:
+    """Minimal stand-in implementing the checkpoint's encode_temporal."""
+
+    # Attached per test case as capability probes for the install checks;
+    # annotations only, the fork must keep working when they are absent.
+    encode_temporal: Any
+    _decode_temporal_streaming: Any
+    processor: Any
+
+    def __init__(self, *, isolated_first=False, isolated_last=False, key_frame=False, frame_drop=0):
+        self.clip_length = CLIP
+        self.isolated_first_frame = isolated_first
+        self.frame_pre_padding = 0
+        self.isolated_key_frame = key_frame
+        self.token_drop = 0
+        self.isolated_last_frame = isolated_last
+        self.frame_drop = frame_drop
+        self.encoded = []
+
+    def _adaptive_encode(self, clip):
+        self.encoded.append(clip.clone())
+        return clip  # identity: z_list entries are the clips themselves
+
+    def reference_encode_temporal(self, x):
+        """The checkpoint's original: whole-video pad, then chunk."""
+        offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0
+        if x.shape[2] % self.clip_length != offset_frame:
+            pad_size = (offset_frame - x.shape[2]) % self.clip_length
+            pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)
+            x = torch.cat([x, pad_frames], dim=2)
+        num_chunks = (x.shape[2] - offset_frame) // self.clip_length
+        z_list = []
+        for i in range(num_chunks):
+            start_idx = i * self.clip_length + offset_frame
+            end_idx = (i + 1) * self.clip_length + offset_frame
+            clip_x = x[:, :, start_idx:end_idx, :, :]
+            if self.isolated_key_frame:
+                key_frame = clip_x[:, :, :1, :, :]
+                z_key = self._adaptive_encode(key_frame)
+                if clip_x.shape[2] > 1:
+                    z_video = self._adaptive_encode(clip_x[:, :, 1:, :, :])
+                    z = torch.cat([z_key, z_video], dim=2)
+                else:
+                    z = z_key
+            else:
+                z = self._adaptive_encode(clip_x)
+            z_list.append(z)
+        z = torch.cat(z_list, dim=2)
+        if self.token_drop > 0:
+            z = z[:, :, : -self.token_drop]
+        if self.isolated_first_frame:
+            z_first = self._adaptive_encode(x[:, :, :1, :, :])
+            z = torch.cat([z_first, z], dim=2)
+        if self.isolated_last_frame:
+            frame_num = x.shape[2]
+            last_idx = frame_num - self.frame_drop + offset_frame
+            z_last = self._adaptive_encode(x[:, :, last_idx : last_idx + 1, :, :])
+            z = torch.cat([z, z_last], dim=2)
+        return z
+
+
+def _video_x(t, c=3, h=4, w=4):
+    return torch.arange(t * c * h * w, dtype=torch.float32).reshape(1, c, t, h, w) / 255.0
+
+
+@pytest.mark.parametrize("num_frames", [8, 10, 11, 13])
+@pytest.mark.parametrize(
+    "flags",
+    [
+        {},
+        {"isolated_first": True},
+        {"isolated_last": True},
+        {"isolated_first": True, "isolated_last": True},
+        {"key_frame": True},
+    ],
+)
+def test_encode_fork_matches_whole_video_pad_bitwise(num_frames, flags):
+    x = _video_x(num_frames)
+    ref_model = _EncodeModel(**flags)
+    expected = ref_model.reference_encode_temporal(x.clone())
+    # Empty-slice encodes (out-of-range isolated-last indexes are no-ops in
+    # the checkpoint) carry no content; compare only real calls.
+    ref_calls = [c.clone() for c in ref_model.encoded if c.shape[2] > 0]
+
+    fork_model = _EncodeModel(**flags)
+    out = _encode_temporal_tail_pad(fork_model, x.clone())
+
+    assert out.shape == expected.shape
+    assert torch.equal(out, expected)
+    assert len(fork_model.encoded) == len(ref_calls)
+    for got, want in zip(fork_model.encoded, ref_calls):
+        assert torch.equal(got, want)
+
+
+def test_encode_fork_avoids_whole_video_pad_for_misaligned_input():
+    # 375 frames with an asymmetric config (isolated last frame only):
+    # the checkpoint always pads (375 % 17 != 0); the fork must encode the
+    # tail chunk with an in-chunk pad instead of concatenating the video.
+    model = _EncodeModel(isolated_last=True, frame_drop=16)
+    model.clip_length = 17
+    x = _video_x(375, c=1, h=2, w=2)
+    _encode_temporal_tail_pad(model, x)
+    # 375 % 17 == 1 → pad 16 inside the tail chunk; the isolated last frame
+    # resolves to the final real frame (pad frames replicate it).
+    assert len(model.encoded) == (375 // 17) + 1 + 1
+    assert torch.equal(model.encoded[-1].flatten(), x[:, :, -1:, :, :].flatten())
+
+
+# ---------------------------------------------------------------------------
+# _decode_temporal_streaming uint8 fork
+# ---------------------------------------------------------------------------
+
+DENORM_MEAN = (-2.11628412, -2.03571429, -1.80444444)
+DENORM_STD = (4.36681223, 4.46428571, 4.44444444)
+
+
+class _FakeDenorm:
+    def __init__(self):
+        self.mean = DENORM_MEAN
+        self.std = DENORM_STD
+
+
+class _FakeProcessor:
+    def __init__(self):
+        self.transform_rev = _FakeDenorm()
+
+
+class _DecodeModel:
+    """Minimal stand-in for the checkpoint's streaming decode contract."""
+
+    def __init__(self, *, token_drop=0, frame_overlap=0):
+        self.tokens_chunk_size = 2
+        self.vae_ratio_t = 1
+        self.token_overlap = 1 if token_drop else 0
+        self.token_drop = token_drop
+        self.frame_pre_padding = 0
+        self.frame_overlap = frame_overlap
+        self.processor = _FakeProcessor()
+
+    def _adaptive_decode(self, clip_z):
+        # Deterministic pseudo-decode: map latent tokens to frames 1:1.
+        return torch.sin(clip_z * 3.0)
+
+    def blend(self, a, b, overlap, dim=-3):
+        return (a + b) / 2.0
+
+    def _decode_temporal_pad_frames(self, z, pad_tokens):
+        return int(pad_tokens)
+
+    def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens):
+        total = num_chunks * self.tokens_chunk_size * self.vae_ratio_t
+        if z_head is not None:
+            total += 1
+        if z_tail is not None:
+            total += 1
+        if self.token_drop > 0:
+            total += self.frame_overlap
+        return int(total), int(pad_tokens), int(total - pad_tokens)
+
+    def legacy_streaming(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype):
+        """The checkpoint's original: accumulate in the cat dtype, quantize later."""
+        total_frames, _, output_frames = self._decode_temporal_output_frame_plan(
+            z, z_head, z_tail, num_chunks, pad_tokens
+        )
+        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+        split_count = int(self.token_drop > 0) + 1
+        dec = None
+        dec_overlap = None
+        write_pos = 0
+
+        def write_part(part):
+            nonlocal dec, write_pos
+            if dec is None:
+                out_shape = list(part.shape)
+                out_shape[2] = output_frames
+                dec = torch.empty(out_shape, dtype=part.dtype)
+            n = min(part.shape[2], dec.shape[2] - write_pos)
+            if n > 0:
+                dec[:, :, write_pos : write_pos + n].copy_(part[:, :, :n])
+                write_pos += n
+
+        for i in range(num_chunks):
+            t_start = i * self.tokens_chunk_size
+            t_end = t_start + self.tokens_chunk_size + self.token_overlap
+            clip_z = z[:, :, t_start:t_end]
+            clip_dec = self._adaptive_decode(clip_z)
+            if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
+                clip_dec = clip_dec.to(temporal_cat_dtype)
+            for j in range(split_count):
+                chunk = clip_dec[:, :, j * chunk_dec : min(j * chunk_dec + chunk_dec, clip_dec.shape[2])]
+                if j == 0:
+                    if dec_overlap is not None:
+                        chunk = self.blend(dec_overlap, chunk, self.frame_overlap, dim=-3)
+                        dec_overlap = None
+                    write_part(chunk)
+                else:
+                    dec_overlap = chunk.contiguous()
+            if i == num_chunks - 1 and dec_overlap is not None:
+                write_part(dec_overlap)
+                dec_overlap = None
+        # Legacy post-decode chain: revert in place in the decoded dtype, then
+        # the adapter's ``frames.float()`` before the pipeline's quantizer.
+        assert dec is not None
+        mean = torch.tensor(DENORM_MEAN, dtype=dec.dtype).view(1, 3, 1, 1, 1)
+        std = torch.tensor(DENORM_STD, dtype=dec.dtype).view(1, 3, 1, 1, 1)
+        dec.sub_(mean).div_(std).clamp_(0.0, 1.0)
+        dec = dec.float()
+        dec.mul_(255.0).round_()
+        return dec.to(torch.uint8)
+
+
+@pytest.mark.parametrize("token_drop,frame_overlap", [(0, 0), (1, 1)])
+@pytest.mark.parametrize("temporal_cat_dtype", [None, torch.float16, torch.bfloat16])
+def test_decode_fork_matches_legacy_chain_bitwise(token_drop, frame_overlap, temporal_cat_dtype):
+    # The quantizer must see FP32 no matter the cat dtype: the legacy adapter
+    # returned frames.float() before the pipeline multiplied and rounded, so
+    # an fp16/bf16 cat dtype only narrows denormalization, never the rounding.
+    torch.manual_seed(0)
+    tokens = 2 * 2 + (1 if token_drop else 0)  # num_chunks * tokens_chunk_size
+    z = torch.randn(1, 3, tokens, 4, 4)
+    legacy = _DecodeModel(token_drop=token_drop, frame_overlap=frame_overlap)
+    expected = legacy.legacy_streaming(z.clone(), None, None, 2, 0, temporal_cat_dtype)
+
+    fork = _DecodeModel(token_drop=token_drop, frame_overlap=frame_overlap)
+    out = _decode_temporal_streaming_uint8(fork, z.clone(), None, None, 2, 0, temporal_cat_dtype)
+    assert out.dtype == torch.uint8
+    assert out.shape == expected.shape
+    assert torch.equal(out, expected)
+
+
+def test_decode_fork_rejects_empty_plan():
+    model = _DecodeModel()
+    model.tokens_chunk_size = 0
+    with pytest.raises(ValueError, match="non-positive output_frames"):
+        _decode_temporal_streaming_uint8(model, torch.randn(1, 3, 0, 4, 4), None, None, 0, 0, None)
+
+
+# ---------------------------------------------------------------------------
+# install + escape hatch
+# ---------------------------------------------------------------------------
+
+
+def test_install_binds_both_forks():
+    model = _EncodeModel()
+    # cast(Any, ...) keeps mypy from narrowing the probes to the lambda type,
+    # which has no __func__ for the bound-method assertions below.
+    model.encode_temporal = cast(Any, lambda x: x)  # placeholder for the capability check
+    model._decode_temporal_streaming = cast(Any, lambda *a: None)
+    model.processor = _FakeProcessor()
+    install_temporal_stream_patches(model)
+    assert model.encode_temporal.__func__ is _encode_temporal_tail_pad
+    assert model._decode_temporal_streaming.__func__ is _decode_temporal_streaming_uint8
+
+
+def test_install_skips_decode_fork_without_denorm_contract():
+    model = _EncodeModel()
+    model.encode_temporal = cast(Any, lambda x: x)
+    model._decode_temporal_streaming = cast(Any, lambda *a: None)
+    install_temporal_stream_patches(model)  # no processor -> keep original
+    assert model.encode_temporal.__func__ is _encode_temporal_tail_pad
+    assert not hasattr(model._decode_temporal_streaming, "__func__")
+
+
+def test_install_respects_legacy_env(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_VAE_LEGACY_TEMPORAL", "1")
+    model = _EncodeModel()
+    install_temporal_stream_patches(model)
+    assert not hasattr(getattr(model, "encode_temporal", None), "__func__")
+
+
+# ---------------------------------------------------------------------------
+# _prepare_minimax_h3_video_output
+# ---------------------------------------------------------------------------
+
+
+def test_prepare_output_matches_previous_conversion():
+    video = torch.randn(1, 3, 4, 6, 8).clamp_(-1, 2)
+    legacy = (
+        video.detach()
+        .float()
+        .clamp(0, 1)
+        .mul(255)
+        .round()
+        .permute(0, 2, 3, 4, 1)
+        .to(dtype=torch.uint8, memory_format=torch.contiguous_format)
+    )
+    out = _prepare_minimax_h3_video_output(video.clone())
+    assert out.dtype == torch.uint8
+    assert out.shape == (1, 4, 6, 8, 3)
+    assert out.is_contiguous()
+    assert torch.equal(out, legacy)
+
+
+def test_prepare_output_uint8_passthrough():
+    video = torch.randint(0, 256, (1, 3, 4, 6, 8), dtype=torch.uint8)
+    out = _prepare_minimax_h3_video_output(video)
+    assert out.dtype == torch.uint8
+    assert out.is_contiguous()
+    assert torch.equal(out, video.permute(0, 2, 3, 4, 1).contiguous())

tests/diffusion/models/minimax_h3/test_vsa_layout.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import math
+import sys
+import types
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+    VideoTokenLayout,
+    VideoTokenSpan,
+)
+from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import (
+    MiniMaxH3VSAImpl,
+    _get_h3_tile_metadata,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
+    extra: dict[str, Any] = {
+        "vsa_h3_prefix_segments": prefix_segments,
+    }
+    if gate is not None:
+        extra["gate_compress"] = gate
+    return AttentionMetadata(
+        packed_padding=packed_padding,
+        video_layout=VideoTokenLayout(
+            used_len=sum(prefix_segments) + math.prod(video_shape),
+            video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
+        ),
+        extra=extra,
+    )
+
+
+def _fake_block_sparse_kernel(monkeypatch, calls):
+    """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
+
+    def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
+        calls["block_map"] = block_map.clone()
+        calls["variable_block_sizes"] = variable_block_sizes.tolist()
+        calls["q_shape"] = tuple(q.shape)
+        blocks = block_map.shape[-1]
+        assert blocks * 64 == q.shape[2]
+        # Expand the per-block map to a per-token mask so the stand-in honours
+        # the sparsity contract the real kernel implements, and drop the rows
+        # each partial block pads with, which the real kernel skips through
+        # variable_block_sizes.
+        token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
+        within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
+        token_mask = token_mask & within.reshape(1, 1, 1, -1)
+        out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
+        return out.to(q.dtype), None
+
+    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    setattr(module, "block_sparse_attn", block_sparse_attn)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
+
+
+def _h3_impl(**backend_kwargs):
+    kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
+    kwargs.update(backend_kwargs)
+    return MiniMaxH3VSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
+
+
+def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5, 70), (2, 4, 4)
+    seq_len = sum(prefix_segments) + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    assert output.shape == query.shape
+    # topk covers every video block, so the block map is dense and the tiled
+    # kernel must reproduce plain attention over the same rows.
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
+    # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
+    assert calls["variable_block_sizes"] == [5, 64, 6, 32]
+
+
+def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=1)
+    prefix_segments, video_shape = (5,), (8, 4, 4)
+    seq_len = 5 + 8 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    # 1 prefix block + 2 video blocks. The odd count is padded to an even one
+    # for the paired-CTA kernel contract, and the Triton route slices that
+    # transport-only partner off again before launching.
+    block_map = calls["block_map"]
+    assert block_map.shape[-1] == 3
+    assert calls["variable_block_sizes"] == [5, 64, 64]
+    assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
+    assert block_map[..., :1].all(), "prefix keys are exempt from selection"
+    assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
+
+    # The same odd geometry must still untile back onto the original rows.
+    dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
+
+
+def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    seq_len = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    zero_gate = torch.zeros_like(query)
+    with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
+    live_gate = torch.full_like(query, 0.5)
+    with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
+
+    torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
+    assert not torch.allclose(with_live_gate.float(), without_gate.float())
+
+
+def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    valid = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
+    padding = PackedPaddingMetadata(
+        q_length=valid,
+        kv_length=valid,
+        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
+        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
+    )
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
+
+    assert output.shape == query.shape
+    assert torch.count_nonzero(output[:, valid:]) == 0
+
+
+def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
+    partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
+        (5, 70, 9), (5, 6, 6), torch.device("cpu")
+    )
+    assert sizes[:4].tolist() == [5, 64, 6, 9]
+    assert prefix_blocks == 4
+    assert video_blocks == 8
+    assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
+    assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
+
+
+@pytest.mark.parametrize("fallback_on_error", [False, True])
+def test_h3_kernel_failure_preserves_fallback_policy(monkeypatch, fallback_on_error):
+    impl = _h3_impl(fallback_on_error=fallback_on_error)
+    query = torch.randn(1, 37, 2, 8)
+    metadata = _h3_metadata((5,), (2, 4, 4))
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("kernel failure")
+
+    monkeypatch.setattr(impl, "_forward_h3", fail)
+    monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
+    if fallback_on_error:
+        output = impl.forward_cuda(query, query, query, metadata)
+        expected = impl.sdpa_fallback.forward_cuda(query, query, query, metadata)
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+    else:
+        with pytest.raises(RuntimeError, match="kernel failure"):
+            impl.forward_cuda(query, query, query, metadata)
+
+
+def test_h3_accelerator_failure_is_not_swallowed(monkeypatch):
+    impl = _h3_impl(fallback_on_error=True)
+    query = torch.randn(1, 37, 2, 8)
+
+    def fail(*args, **kwargs):
+        raise torch.AcceleratorError("device failure")
+
+    monkeypatch.setattr(impl, "_forward_h3", fail)
+    with pytest.raises(torch.AcceleratorError, match="device failure"):
+        impl.forward_cuda(query, query, query, _h3_metadata((5,), (2, 4, 4)))
+
+
+def test_h3_without_layout_delegates_to_shared_backend(monkeypatch):
+    from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSAImpl
+
+    impl = _h3_impl()
+    query = torch.randn(1, 37, 2, 8)
+    metadata = AttentionMetadata()
+    expected = torch.empty_like(query)
+    calls = []
+
+    def forward(self, q, k, v, attn_metadata):
+        calls.append((self, q, k, v, attn_metadata))
+        return expected
+
+    monkeypatch.setattr(FastVideoVSAImpl, "forward_cuda", forward)
+    assert impl.forward_cuda(query, query, query, metadata) is expected
+    assert len(calls) == 1
+    assert all(actual is wanted for actual, wanted in zip(calls[0], (impl, query, query, query, metadata)))

tests/diffusion/models/pi05/__init__.py

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

tests/diffusion/models/pi05/test_pi05_parity.py

@@ -0,0 +1,500 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""π0.5 parity against LeRobot.
+
+Reference: **lerobot 0.6.1**, transformers 5.5.4, torch 2.11.0+cpu. The oracle
+is LeRobot's own code, so a later version changing its π0.5 implementation
+fails this without anything here regressing — diff ``policies/pi05/`` first.
+
+Compares final action chunks under fixed noise: float32 exactly, bfloat16
+against LeRobot's own bfloat16 layout within ``BF16_ATOL``.
+
+Needs a separate ``lerobot[pi]`` venv (its dependencies conflict with
+vllm-omni's) with the vllm-omni ``pi05`` package importable::
+
+    python -m pytest tests/diffusion/models/pi05/test_pi05_parity.py -v -s
+
+Skips itself when LeRobot is absent. ``PI05_PARITY_MODEL_PATH`` points at a
+local checkpoint to skip the HF download.
+"""
+
+from __future__ import annotations
+
+import copy
+import importlib.util
+import os
+
+import pytest
+import torch
+
+# local_model: needs real weights + a lerobot venv, so it runs locally rather
+# than in the ready-CI. Additionally gated on lerobot being importable.
+_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None
+
+# cpu: the comparison runs on CPU in float32 so both implementations see
+# identical numerics; local_model: it needs a lerobot venv that CI does not have.
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+# ─── Config (fixed; matches LeRobot defaults for a π0.5 checkpoint) ────
+DEVICE = "cpu"
+DTYPE_STR = "float32"
+ATOL = 1e-4
+BF16_ATOL = 5e-2  # measured 2.95e-2, reproducible; margin is for version drift
+NUM_STEPS = 10
+BATCH_SIZE = 2
+ACTION_DIM = 32
+STATE_DIM = 32
+ACTION_HORIZON = 50
+MAX_TOKEN_LEN = 200  # π0 uses 48
+NUM_STATE_BINS = 256
+
+MODEL_PATH = os.environ.get("PI05_PARITY_MODEL_PATH", "lerobot/pi05_base")
+
+CAMERAS = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+
+
+def _resolve_checkpoint_dir() -> str:
+    if os.path.isdir(MODEL_PATH):
+        return MODEL_PATH
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(repo_id=MODEL_PATH, repo_type="model")
+
+
+# ─── Dummy dataset stats ──────────────────────────────────────────────
+def _dummy_dataset_stats() -> dict:
+    """Identity-ish stats. π0.5 normalizes with QUANTILES by default, so q01/q99
+    carry the load here where π0's mean/std would.
+
+    Using dummy stats means this test needs no checkpoint with real
+    ``norm_stats`` — it verifies numerical equivalence, not calibration.
+    """
+    return {
+        "observation.state": {
+            "mean": torch.zeros(STATE_DIM),
+            "std": torch.ones(STATE_DIM),
+            "q01": -torch.ones(STATE_DIM),
+            "q99": torch.ones(STATE_DIM),
+        },
+        "action": {
+            "mean": torch.zeros(ACTION_DIM),
+            "std": torch.ones(ACTION_DIM),
+            "q01": -torch.ones(ACTION_DIM),
+            "q99": torch.ones(ACTION_DIM),
+        },
+        "images": {
+            cam: {
+                "mean": torch.zeros(3, 224, 224),
+                "std": torch.ones(3, 224, 224),
+                "q01": torch.zeros(3, 224, 224),
+                "q99": torch.ones(3, 224, 224),
+            }
+            for cam in CAMERAS
+        },
+    }
+
+
+def _create_dummy_batch(batch_size: int = BATCH_SIZE, num_views: int = 3, device: str = DEVICE) -> dict:
+    """Reproducible dummy inputs — identical across both implementations."""
+    g = torch.Generator(device="cpu").manual_seed(0)
+    prompt = "Pick up the red block and place it in the bin"
+    batch = {
+        "observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device),
+        "action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device),
+        "task": [prompt for _ in range(batch_size)],
+    }
+    for cam in CAMERAS[:num_views]:
+        batch[f"observation.images.{cam}"] = torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
+            device
+        )
+    return batch
+
+
+# ─── Instantiation ────────────────────────────────────────────────────
+def _instantiate_lerobot(use_relative_actions: bool = False, action_feature_names=None, dtype: str = DTYPE_STR):
+    from lerobot.policies.pi05 import PI05Policy
+    from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors
+
+    policy = PI05Policy.from_pretrained(MODEL_PATH, strict=True, dtype=dtype)
+    policy.to(DEVICE)
+    policy.config.device = DEVICE
+    policy.config.use_relative_actions = use_relative_actions
+    if action_feature_names is not None:
+        # Must land on the config *before* the processors are built:
+        # make_pi05_pre_post_processors captures action_feature_names into
+        # RelativeActionsProcessorStep at construction time, and its _build_mask
+        # treats ``action_names is None`` as "shift every dim" — including the
+        # gripper that relative_exclude_joints is supposed to hold back. Setting
+        # it afterwards leaves the step with the all-True mask and silently
+        # compares two different transforms.
+        policy.config.action_feature_names = action_feature_names
+    policy.eval()
+
+    pre, post = make_pi05_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats())
+    return policy, pre, post
+
+
+def _instantiate_vllm_omni(
+    use_relative_actions: bool = False,
+    action_feature_names=None,
+    dtype: str = DTYPE_STR,
+):
+    """Build the vllm-omni π0.5 model in isolation (no pipeline, no engine)."""
+    from vllm_omni.diffusion.models.pi05 import Pi05Config, Pi05ForActionPrediction
+    from vllm_omni.diffusion.models.pi05.pipeline_pi05 import _set_inference_dtype
+
+    cfg = Pi05Config(
+        max_action_dim=ACTION_DIM,
+        max_state_dim=STATE_DIM,
+        chunk_size=ACTION_HORIZON,
+        num_inference_steps=NUM_STEPS,
+        tokenizer_max_length=MAX_TOKEN_LEN,
+        state_num_bins=NUM_STATE_BINS,
+        dtype=dtype,
+        use_relative_actions=use_relative_actions,
+        relative_exclude_joints=["gripper"] if action_feature_names else [],
+        action_feature_names=action_feature_names,
+    )
+    model = Pi05ForActionPrediction(cfg)
+    _set_inference_dtype(model, getattr(torch, dtype))
+    model.to(device=DEVICE)
+    _load_lerobot_weights(model)
+    model.eval()
+    return model, cfg
+
+
+def _load_lerobot_weights(model):
+    import safetensors.torch
+
+    path = os.path.join(_resolve_checkpoint_dir(), "model.safetensors")
+    state = safetensors.torch.load_file(path)
+    filled = model.load_weights(list(state.items()))
+    assert filled, "no weights were loaded — the remap rules are broken"
+
+
+def _extract_lerobot_model_inputs(lerobot_policy, processed_batch):
+    """Mimic what ``PI05Policy.predict_action_chunk`` feeds into
+    ``self.model.sample_actions``. Using these *exact* tensors for vllm-omni
+    means any divergence must come from the core model, not preprocessing.
+
+    Note π0.5 has no ``prepare_state`` step: the state is already inside
+    ``lang_tokens``.
+    """
+    from lerobot.utils.constants import (
+        OBS_LANGUAGE_ATTENTION_MASK,
+        OBS_LANGUAGE_TOKENS,
+    )
+
+    images, img_masks = lerobot_policy._preprocess_images(processed_batch)
+    return images, img_masks, processed_batch[OBS_LANGUAGE_TOKENS], processed_batch[OBS_LANGUAGE_ATTENTION_MASK]
+
+
+def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor:
+    g = torch.Generator(device="cpu").manual_seed(42)
+    return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device)
+
+
+# ─── Main parity test ─────────────────────────────────────────────────
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+@pytest.mark.parametrize("num_views", [1, 2, 3])
+def test_pi05_vllm_omni_vs_lerobot(num_views):
+    lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
+    omni_model, _ = _instantiate_vllm_omni()
+
+    raw_batch = _create_dummy_batch(num_views=num_views)
+    processed = lerobot_pre(copy.deepcopy(raw_batch))
+    images, img_masks, lang_tokens, lang_masks = _extract_lerobot_model_inputs(lerobot_policy, processed)
+    noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)
+
+    # The input contract: 256 tokens per view + a constant 200 text tokens.
+    assert lang_tokens.shape[1] == MAX_TOKEN_LEN, "π0.5 must pad text to exactly 200 tokens"
+
+    with torch.no_grad():
+        lerobot_actions = lerobot_policy.model.sample_actions(
+            images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS
+        )
+        omni_actions = omni_model.sample_actions(
+            images=images,
+            image_masks=img_masks,
+            lang_tokens=lang_tokens,
+            lang_masks=lang_masks,
+            noise=noise,
+            num_steps=NUM_STEPS,
+        )
+
+    diff = (lerobot_actions.float() - omni_actions.float()).abs()
+    print(f"[parity] views={num_views} |Δ| max={diff.max().item():.2e} mean={diff.mean().item():.2e}")
+    if not torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL):
+        _diagnose_divergence(lerobot_policy.model, omni_model, images, img_masks, lang_tokens, lang_masks, noise)
+    assert torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL), (
+        f"actions differ beyond atol={ATOL}; max_diff={diff.max().item():.2e}"
+    )
+
+
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+def test_pi05_bfloat16_matches_lerobot_bfloat16():
+    """Both sides in bfloat16, so this compares implementations rather than
+    measuring bfloat16's drift from float32."""
+    # LeRobot applies the layout in PaliGemmaWithExpertModel.__init__ from
+    # config.dtype, so the policy must be built as bfloat16 rather than cast after.
+    lerobot_policy, lerobot_pre, _ = _instantiate_lerobot(dtype="bfloat16")
+
+    omni_model, _ = _instantiate_vllm_omni(dtype="bfloat16")
+
+    raw_batch = _create_dummy_batch(num_views=3)
+    processed = lerobot_pre(copy.deepcopy(raw_batch))
+    images, img_masks, lang_tokens, lang_masks = _extract_lerobot_model_inputs(lerobot_policy, processed)
+    noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)
+
+    with torch.no_grad():
+        lerobot_actions = lerobot_policy.model.sample_actions(
+            images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS
+        )
+        omni_actions = omni_model.sample_actions(
+            images=images,
+            image_masks=img_masks,
+            lang_tokens=lang_tokens,
+            lang_masks=lang_masks,
+            noise=noise,
+            num_steps=NUM_STEPS,
+        )
+
+    diff = (lerobot_actions.float() - omni_actions.float()).abs()
+    print(f"[parity] bfloat16 |Δ| max={diff.max().item():.2e} mean={diff.mean().item():.2e}")
+    assert torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=BF16_ATOL), (
+        f"bfloat16 actions differ beyond atol={BF16_ATOL}; max_diff={diff.max().item():.2e}"
+    )
+
+
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+def test_pi05_prompt_parity():
+    """π0.5's state is carried as prompt tokens, so the tokenized prompt is
+    itself part of the contract. A bin-index mismatch here (e.g. from running
+    the discretizer before the normalizer) would otherwise only surface as a
+    small, plausible-looking action difference.
+    """
+    from transformers import AutoTokenizer
+
+    from vllm_omni.diffusion.models.pi05.processor_pi05 import (
+        apply_norm,
+        build_norm_stats,
+        build_pi05_prompt,
+        tokenize_prompt,
+    )
+
+    lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
+    raw_batch = _create_dummy_batch(batch_size=1)
+    processed = lerobot_pre(copy.deepcopy(raw_batch))
+
+    from lerobot.utils.constants import OBS_LANGUAGE_TOKENS
+
+    lerobot_tokens = processed[OBS_LANGUAGE_TOKENS][0]
+
+    stats = _dummy_dataset_stats()["observation.state"]
+    norm = build_norm_stats({"state": {"q01": stats["q01"], "q99": stats["q99"]}}, "state")
+    normalized = apply_norm(raw_batch["observation.state"][0].float(), norm)
+    prompt = build_pi05_prompt(
+        task=raw_batch["task"][0],
+        normalized_state=normalized.numpy(),
+        state_num_bins=NUM_STATE_BINS,
+    )
+    tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224", padding_side="right")
+    ids, _ = tokenize_prompt(tokenizer, prompt, MAX_TOKEN_LEN)
+
+    assert len(ids) == MAX_TOKEN_LEN
+    assert torch.equal(torch.tensor(ids), lerobot_tokens.cpu()), "π0.5 prompt tokens diverge from LeRobot"
+
+
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+@pytest.mark.parametrize("state_dim", [7, STATE_DIM])
+def test_pi05_prompt_parity_across_state_widths(state_dim):
+    """``pi05_base`` declares a 32-dim state, which is exactly ``max_state_dim``,
+    so the full-policy parity above cannot tell padding apart from the real
+    width. This drives LeRobot's tokenizer step directly, at both widths.
+    """
+    from lerobot.lerobot_types import TransitionKey
+    from lerobot.policies.pi05.processor_pi05 import Pi05PrepareStateTokenizerProcessorStep
+    from lerobot.utils.constants import OBS_STATE
+
+    from vllm_omni.diffusion.models.pi05.processor_pi05 import build_pi05_prompt
+
+    g = torch.Generator().manual_seed(state_dim)
+    state = torch.randn(state_dim, generator=g, dtype=torch.float32)
+    task = "pick up the red block"
+
+    step = Pi05PrepareStateTokenizerProcessorStep(max_state_dim=STATE_DIM)
+    transition = step(
+        {
+            TransitionKey.OBSERVATION: {OBS_STATE: state.unsqueeze(0)},
+            TransitionKey.COMPLEMENTARY_DATA: {"task": [task]},
+        }
+    )
+    lerobot_prompt = transition[TransitionKey.COMPLEMENTARY_DATA]["task"][0]
+
+    prompt = build_pi05_prompt(task=task, normalized_state=state.numpy())
+
+    assert prompt == lerobot_prompt
+    assert len(prompt.split("State: ")[1].split(";")[0].split()) == state_dim
+
+
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+@pytest.mark.parametrize(
+    "lerobot_mode,stats,ours,state,state_dim",
+    [
+        ("MEAN_STD", {"mean": [0.0], "std": [1e-7]}, {"mode": "mean_std", "mean": [0.0], "std": [1e-7]}, [1e-7], 1),
+        ("QUANTILES", {"q01": [0.0], "q99": [1e-7]}, {"mode": "quantile", "q01": [0.0], "q99": [1e-7]}, [1e-7], 1),
+        ("QUANTILES", {"q01": [2.0], "q99": [2.0]}, {"mode": "quantile", "q01": [2.0], "q99": [2.0]}, [2.0], 1),
+        (
+            "QUANTILES",
+            {"q01": [0.0] * 4, "q99": [10.0] * 4},
+            {"mode": "quantile", "q01": [0.0] * 4, "q99": [10.0] * 4},
+            [-10.0, 0.0, 5.0, 20.0],
+            4,
+        ),
+        (
+            "MEAN_STD",
+            {"mean": [1.0, 2.0], "std": [0.5, 4.0]},
+            {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 4.0]},
+            [3.0, -6.0],
+            2,
+        ),
+        ("MIN_MAX", {"min": [-1.0], "max": [3.0]}, {"mode": "min_max", "min": [-1.0], "max": [3.0]}, [2.0], 1),
+    ],
+)
+def test_pi05_state_normalization_parity(lerobot_mode, stats, ours, state, state_dim):
+    """State normalization against LeRobot's ``NormalizerProcessorStep``.
+
+    The full-policy parity test feeds both sides stats with a wide range, so it
+    cannot see the denominator rule. A wrongly normalized state still discretizes
+    into a well-formed prompt: for ``q01=0, q99=1e-7`` a threshold substitution
+    moves the state token from bin 255 to bin 0.
+    """
+    from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+    from lerobot.processor.normalize_processor import NormalizerProcessorStep
+    from lerobot.utils.constants import OBS_STATE
+
+    from vllm_omni.diffusion.models.pi05.processor_pi05 import apply_norm, build_norm_stats
+
+    step = NormalizerProcessorStep(
+        features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,))},
+        norm_map={FeatureType.STATE: getattr(NormalizationMode, lerobot_mode)},
+        stats={OBS_STATE: {k: torch.tensor(v, dtype=torch.float32) for k, v in stats.items()}},
+    )
+    expected = step._apply_transform(
+        torch.tensor(state, dtype=torch.float32), OBS_STATE, FeatureType.STATE, inverse=False
+    )
+
+    actual = apply_norm(torch.tensor(state, dtype=torch.float32), build_norm_stats({"state": ours}, "state"))
+
+    assert torch.allclose(actual, expected, atol=1e-6)
+
+
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+def test_pi05_relative_actions_parity():
+    """With ``use_relative_actions=True`` the checkpoint's norm_stats live in
+    relative space, so the post-pipeline must add the raw state back. Compares
+    the full pre → model → post path against LeRobot's."""
+    from lerobot.utils.constants import (
+        OBS_LANGUAGE_ATTENTION_MASK,
+        OBS_LANGUAGE_TOKENS,
+    )
+
+    from vllm_omni.diffusion.models.pi05.processor_pi05 import (
+        Pi05RelativeActions,
+        apply_norm,
+        build_norm_stats,
+    )
+
+    action_names = [f"joint_{i}" for i in range(ACTION_DIM - 1)] + ["gripper"]
+    lerobot_policy, lerobot_pre, lerobot_post = _instantiate_lerobot(
+        use_relative_actions=True, action_feature_names=action_names
+    )
+    omni_model, cfg = _instantiate_vllm_omni(use_relative_actions=True, action_feature_names=action_names)
+
+    raw_batch = _create_dummy_batch()
+    processed = lerobot_pre(copy.deepcopy(raw_batch))
+    images, img_masks = lerobot_policy._preprocess_images(processed)
+    lang_tokens = processed[OBS_LANGUAGE_TOKENS]
+    lang_masks = processed[OBS_LANGUAGE_ATTENTION_MASK]
+    noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)
+
+    with torch.no_grad():
+        lerobot_actions = lerobot_post(
+            lerobot_policy.model.sample_actions(
+                images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS
+            )
+        )
+        omni_raw = omni_model.sample_actions(
+            images=images,
+            image_masks=img_masks,
+            lang_tokens=lang_tokens,
+            lang_masks=lang_masks,
+            noise=noise,
+            num_steps=NUM_STEPS,
+        )
+
+    rel = Pi05RelativeActions(
+        enabled=True,
+        exclude_joints=cfg.relative_exclude_joints,
+        action_names=action_names,
+        max_action_dim=ACTION_DIM,
+    )
+    omni_actions = apply_norm(omni_raw, build_norm_stats(cfg.norm_stats, "action"), inverse=True)
+    # Per-sample reference state, matching LeRobot: its
+    # RelativeActionsProcessorStep caches the whole (B, D) state on the input
+    # side and shifts each sample by its own row. Passing a single row here
+    # would shift every sample by sample 0's state, which only agrees for B=1.
+    omni_actions = rel.to_absolute(omni_actions, raw_batch["observation.state"])
+
+    lerobot_actions = torch.as_tensor(lerobot_actions).float()
+    assert torch.allclose(lerobot_actions, omni_actions.float(), atol=ATOL), (
+        "relative-action path diverges from LeRobot; "
+        f"max_diff={(lerobot_actions - omni_actions.float()).abs().max().item():.2e}"
+    )
+
+
+# ─── Per-stage divergence diagnostics ─────────────────────────────────
+@torch.no_grad()
+def _diagnose_divergence(lerobot_flow_model, omni_model, images, img_masks, lang_tokens, lang_masks, noise):
+    """Localize a numerical mismatch to a specific stage:
+    1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang.
+    2. Prefix KV cache layer 0 — PaliGemma LM attention.
+    3. The AdaRMS timestep conditioning vector.
+    4. A single denoise_step velocity at t=1.0 — action expert forward.
+    """
+    from vllm_omni.diffusion.models.pi05.modeling_pi05 import (
+        make_att_2d_masks,
+        prepare_attention_masks_4d,
+    )
+
+    lr_embs, lr_pad, lr_att = lerobot_flow_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
+    sg_embs, sg_pad, sg_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
+    print(f"[diag] prefix_embs max|Δ| = {(lr_embs.float() - sg_embs.float()).abs().max().item():.2e}")
+    print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_pad, sg_pad)}")
+
+    img_len = 256 * len(images)
+    print(f"[diag]   image slice max|Δ| = {(lr_embs[:, :img_len] - sg_embs[:, :img_len]).abs().max().item():.2e}")
+    print(f"[diag]   lang  slice max|Δ| = {(lr_embs[:, img_len:] - sg_embs[:, img_len:]).abs().max().item():.2e}")
+
+    prefix_att_4d = prepare_attention_masks_4d(make_att_2d_masks(sg_pad, sg_att))
+    prefix_pos = torch.cumsum(sg_pad, dim=1) - 1
+    _, sg_kv = omni_model.paligemma_with_expert.forward(
+        attention_mask=prefix_att_4d,
+        position_ids=prefix_pos,
+        past_key_values=None,
+        inputs_embeds=[sg_embs, None],
+        use_cache=True,
+    )
+
+    t = torch.ones(noise.shape[0], dtype=torch.float32)
+    sg_cond = omni_model.embed_timestep(t)
+    print(f"[diag] adarms cond: shape={tuple(sg_cond.shape)} mean={sg_cond.mean().item():+.4f}")
+
+    sg_vt = omni_model.denoise_step(sg_pad, sg_kv, noise, t)
+    print(f"[diag] denoise_step(t=1) v_t mean={sg_vt.mean().item():+.6f}")
+
+
+if __name__ == "__main__":
+    test_pi05_vllm_omni_vs_lerobot(3)

tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
+    _qwen_image_qk_norm_rope,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+BATCH = 2
+Q_HEADS = 24
+K_HEADS = 24
+HEAD_DIM = 128
+EPS = 1e-6
+SEED = 2026
+
+
+@dataclass(frozen=True)
+class QwenImageQKInput:
+    q: torch.Tensor
+    k: torch.Tensor
+    norm_q: nn.RMSNorm
+    norm_k: nn.RMSNorm
+    freqs: torch.Tensor
+
+
+def _make_input(
+    *,
+    seq_len: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    packed_qkv_view: bool,
+) -> QwenImageQKInput:
+    gen = torch.Generator(device=device)
+    gen.manual_seed(SEED + seq_len)
+
+    if packed_qkv_view:
+        qkv_dim = (Q_HEADS + K_HEADS + K_HEADS) * HEAD_DIM
+        qkv = torch.randn(BATCH, seq_len, qkv_dim, device=device, dtype=dtype, generator=gen)
+        q, k, _ = qkv.split(
+            [
+                Q_HEADS * HEAD_DIM,
+                K_HEADS * HEAD_DIM,
+                K_HEADS * HEAD_DIM,
+            ],
+            dim=-1,
+        )
+        q = q.unflatten(-1, (Q_HEADS, HEAD_DIM))
+        k = k.unflatten(-1, (K_HEADS, HEAD_DIM))
+    else:
+        q = torch.randn(BATCH, seq_len, Q_HEADS, HEAD_DIM, device=device, dtype=dtype, generator=gen)
+        k = torch.randn(BATCH, seq_len, K_HEADS, HEAD_DIM, device=device, dtype=dtype, generator=gen)
+
+    norm_q = nn.RMSNorm(HEAD_DIM, eps=EPS, device=device, dtype=dtype)
+    norm_k = nn.RMSNorm(HEAD_DIM, eps=EPS, device=device, dtype=dtype)
+    norm_q.requires_grad_(False)
+    norm_k.requires_grad_(False)
+    with torch.no_grad():
+        norm_q.weight.copy_(torch.randn(HEAD_DIM, device=device, dtype=dtype, generator=gen))
+        norm_k.weight.copy_(torch.randn(HEAD_DIM, device=device, dtype=dtype, generator=gen))
+    angles = torch.randn(seq_len, HEAD_DIM // 2, device=device, dtype=torch.float32, generator=gen)
+    freqs = torch.polar(torch.ones_like(angles), angles)
+    return QwenImageQKInput(q=q, k=k, norm_q=norm_q, norm_k=norm_k, freqs=freqs)
+
+
+def _eager_rotary(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
+    q = data.norm_q(data.q)
+    k = data.norm_k(data.k)
+    rope = RotaryEmbedding(is_neox_style=False)
+    cos = data.freqs.real.to(data.q.dtype)
+    sin = data.freqs.imag.to(data.q.dtype)
+    return rope(q, cos, sin), rope(k, cos, sin)
+
+
+def _fp32_complex_rotary(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    """Local copy of the old CUDA helper; fused kernel still tracks this math."""
+    paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype)
+
+
+def _fused_kernel_reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]:
+    q = data.norm_q(data.q)
+    k = data.norm_k(data.k)
+    return _fp32_complex_rotary(q, data.freqs), _fp32_complex_rotary(k, data.freqs)
+
+
+def _run(data: QwenImageQKInput, *, use_fused: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
+    return _qwen_image_qk_norm_rope(
+        data.q,
+        data.k,
+        data.norm_q,
+        data.norm_k,
+        data.freqs,
+        RotaryEmbedding(is_neox_style=False),
+        EPS,
+        use_fused=use_fused,
+    )
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_qwen_image_qk_norm_rope_cuda_fp32_fallback_matches_reference():
+    data = _make_input(
+        seq_len=7,
+        dtype=torch.float32,
+        device=torch.device("cuda:0"),
+        packed_qkv_view=True,
+    )
+
+    actual_q, actual_k = _run(data, use_fused=False)
+    expected_q, expected_k = _eager_rotary(data)
+
+    torch.testing.assert_close(actual_q, expected_q, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_k, expected_k, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_qwen_image_qk_norm_rope_cuda_fp16_fallback_matches_reference():
+    data = _make_input(
+        seq_len=7,
+        dtype=torch.float16,
+        device=torch.device("cuda:0"),
+        packed_qkv_view=True,
+    )
+
+    actual_q, actual_k = _run(data, use_fused=False)
+    expected_q, expected_k = _eager_rotary(data)
+
+    torch.testing.assert_close(actual_q, expected_q, atol=1e-3, rtol=1e-3)
+    torch.testing.assert_close(actual_k, expected_k, atol=1e-3, rtol=1e-3)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("seq_len", [1, 7, 257, 1024])
+@pytest.mark.parametrize("packed_qkv_view", [False, True])
+def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference(
+    seq_len: int,
+    packed_qkv_view: bool,
+):
+    data = _make_input(
+        seq_len=seq_len,
+        dtype=torch.bfloat16,
+        device=torch.device("cuda:0"),
+        packed_qkv_view=packed_qkv_view,
+    )
+
+    actual_q, actual_k = _run(data)
+    expected_q, expected_k = _fused_kernel_reference(data)
+
+    torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
+    torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_qwen_image_fused_qk_norm_rope_torch_compile_fullgraph_capture():
+    data = _make_input(
+        seq_len=257,
+        dtype=torch.bfloat16,
+        device=torch.device("cuda:0"),
+        packed_qkv_view=True,
+    )
+    rope = RotaryEmbedding(is_neox_style=False)
+
+    def fn(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        freqs: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return _qwen_image_qk_norm_rope(
+            q,
+            k,
+            data.norm_q,
+            data.norm_k,
+            freqs,
+            rope,
+            EPS,
+        )
+
+    compiled_fn = torch.compile(fn, dynamic=True, fullgraph=True)
+    expected_q, expected_k = fn(data.q, data.k, data.freqs)
+    actual_q, actual_k = compiled_fn(data.q, data.k, data.freqs)
+
+    torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
+    torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)

tests/diffusion/models/qwen_image/test_qwen_image_pipeline_device.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Encoder/VAE placement under cpu_offload; DiT follows the loader device context."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_DEVICE_STACK: list[str] = ["cpu"]
+_REAL_TORCH_DEVICE = torch.device
+
+
+class _FakeDevice:
+    def __init__(self, type_: str):
+        self.type = type_
+        self.index = 0
+
+    def __enter__(self):
+        _DEVICE_STACK.append(self.type)
+        return self
+
+    def __exit__(self, *args):
+        _DEVICE_STACK.pop()
+        return False
+
+
+def _fake_torch_device(spec):
+    if isinstance(spec, _FakeDevice):
+        return spec
+    if isinstance(spec, _REAL_TORCH_DEVICE):
+        return _FakeDevice(spec.type)
+    return _FakeDevice(str(spec).split(":")[0])
+
+
+class _RecordingTransformer(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.probe_device_type = _DEVICE_STACK[-1]
+
+
+class _MoveableModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.placed_device = _FakeDevice("cpu")
+        self.weight = nn.Parameter(torch.zeros(1))
+
+    def to(self, device, *args, **kwargs):
+        if isinstance(device, _FakeDevice):
+            self.placed_device = device
+        elif isinstance(device, _REAL_TORCH_DEVICE):
+            self.placed_device = _FakeDevice(device.type)
+        else:
+            self.placed_device = _FakeDevice(str(device).split(":")[0])
+        return self
+
+    @property
+    def temperal_downsample(self):
+        return [True, True, True]
+
+
+def _run_pipeline_init(monkeypatch, *, enable_cpu_offload, loader_device):
+    import vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image as pipe_mod
+    from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline
+
+    enc = _MoveableModule()
+    vae = _MoveableModule()
+    _DEVICE_STACK.clear()
+    _DEVICE_STACK.append(loader_device)
+
+    monkeypatch.setattr(pipe_mod, "get_local_device", lambda: _REAL_TORCH_DEVICE("cuda"))
+    monkeypatch.setattr(pipe_mod.torch, "device", _fake_torch_device)
+    monkeypatch.setattr(pipe_mod, "prefetch_subfolders", lambda *a, **k: None)
+    monkeypatch.setattr(pipe_mod, "get_transformer_config_kwargs", lambda *a, **k: {})
+    monkeypatch.setattr(pipe_mod, "QwenImageTransformer2DModel", _RecordingTransformer)
+    monkeypatch.setattr(
+        pipe_mod.FlowMatchEulerDiscreteScheduler,
+        "from_pretrained",
+        classmethod(lambda cls, *a, **k: MagicMock()),
+    )
+    monkeypatch.setattr(
+        pipe_mod.Qwen2Tokenizer,
+        "from_pretrained",
+        classmethod(lambda cls, *a, **k: MagicMock()),
+    )
+
+    def _prefetch(factory, *args, **kwargs):
+        subfolder = kwargs.get("subfolder")
+        if subfolder == "text_encoder":
+            return enc
+        if subfolder == "vae":
+            return vae
+        raise AssertionError(f"unexpected prefetch subfolder={subfolder!r}")
+
+    monkeypatch.setattr(pipe_mod, "from_pretrained_with_prefetch", _prefetch)
+    monkeypatch.setattr(
+        QwenImagePipeline,
+        "setup_diffusion_pipeline_profiler",
+        lambda self, **kwargs: None,
+    )
+
+    od_config = SimpleNamespace(
+        model="/tmp/unused-qwen-image",
+        parallel_config=SimpleNamespace(),
+        tf_model_config={},
+        quantization_config=None,
+        enable_cpu_offload=enable_cpu_offload,
+        enable_diffusion_pipeline_profiler=False,
+    )
+    return QwenImagePipeline(od_config=od_config)
+
+
+def test_pipeline_cpu_offload_parks_encoder_vae_on_cpu(monkeypatch):
+    """#7555: cpu_offload keeps BF16 encoder/VAE off the GPU during DiT init."""
+    pipe = _run_pipeline_init(monkeypatch, enable_cpu_offload=True, loader_device="cuda")
+    assert pipe.text_encoder.placed_device.type == "cpu"
+    assert pipe.vae.placed_device.type == "cpu"
+    assert pipe.transformer.probe_device_type == "cuda"
+
+
+def test_pipeline_dit_follows_loader_cpu_without_override(monkeypatch):
+    """Layerwise / HSDP: DiT must not be forced onto CUDA by the pipeline."""
+    pipe = _run_pipeline_init(monkeypatch, enable_cpu_offload=False, loader_device="cpu")
+    assert pipe.transformer.probe_device_type == "cpu"
+    assert pipe.text_encoder.placed_device.type == "cuda"
+    assert pipe.vae.placed_device.type == "cuda"

tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py

@@ -3,37 +3,12 @@ import torch
 
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenEmbedRope,
-    _apply_qwen_image_rotary_emb,
 )
 from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("noncontiguous", [False, True])
-def test_qwen_rotary_preserves_frequency_precision(dtype, noncontiguous):
-    from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
-
-    generator = torch.Generator().manual_seed(42)
-    x = torch.randn(2, 17, 3, 128, generator=generator).to(dtype)
-    if noncontiguous:
-        x = x.transpose(1, 2).contiguous().transpose(1, 2)
-    angles = torch.randn(17, 64, generator=generator)
-    freqs = torch.polar(torch.ones_like(angles), angles)
-    expected = apply_rotary_emb_qwen(x, freqs, use_real=False)
-    actual = _apply_qwen_image_rotary_emb(x, freqs)
-    assert actual.dtype == dtype
-    assert actual.shape == x.shape
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-
-    if dtype == torch.bfloat16:
-        # A regression to early BF16 frequency rounding must be observable.
-        rounded_freqs = torch.complex(freqs.real.to(dtype).float(), freqs.imag.to(dtype).float())
-        rounded = _apply_qwen_image_rotary_emb(x, rounded_freqs)
-        assert not torch.equal(actual, rounded)
-
-
 def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count():
     prompt_embeds = torch.zeros(2, 64, 8)
     prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool)

tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.qwen_image import qwen_image_transformer
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenImageTransformerBlock
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _StaticOutput(nn.Module):
+    def __init__(self, output: torch.Tensor):
+        super().__init__()
+        self.output = output
+
+    def forward(self, _: torch.Tensor) -> torch.Tensor:
+        return self.output
+
+
+class _RecordingNorm(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+    def forward(self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor:
+        self.calls.append((x, scale, shift))
+        return x
+
+
+class _ZeroAttention(nn.Module):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        **_: object,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.zeros_like(hidden_states), torch.zeros_like(encoder_hidden_states)
+
+
+class _ZeroMLP(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.zeros_like(x)
+
+
+def test_non_cuda_select01_uses_adalayernorm_dispatch(monkeypatch):
+    hidden_size = 4
+    block = QwenImageTransformerBlock.__new__(QwenImageTransformerBlock)
+    nn.Module.__init__(block)
+    block.zero_cond_t = True
+    block.img_mod = _StaticOutput(torch.zeros(2, 6 * hidden_size))
+    block.txt_mod = _StaticOutput(torch.zeros(1, 6 * hidden_size))
+    block.img_norm1 = _RecordingNorm()
+    block.img_norm2 = _RecordingNorm()
+    block.txt_norm1 = _RecordingNorm()
+    block.txt_norm2 = _RecordingNorm()
+    block.attn = _ZeroAttention()
+    block.img_mlp = _ZeroMLP()
+    block.txt_mlp = _ZeroMLP()
+
+    monkeypatch.setattr(qwen_image_transformer, "can_use_qwen_select01_triton", lambda _: False)
+
+    def fail_if_called(*args, **kwargs):
+        pytest.fail("non-CUDA select01 must preserve the AdaLayerNorm dispatch")
+
+    monkeypatch.setattr(qwen_image_transformer, "fused_layernorm_select01", fail_if_called)
+    monkeypatch.setattr(qwen_image_transformer, "fused_residual_layernorm_select01", fail_if_called)
+
+    hidden_states = torch.randn(1, 3, hidden_size)
+    encoder_hidden_states = torch.randn(1, 2, hidden_size)
+    modulate_index = torch.tensor([[0, 1, 0]])
+    block(
+        hidden_states,
+        encoder_hidden_states,
+        torch.ones(1, 2),
+        torch.zeros(2, hidden_size),
+        (torch.empty(0), torch.empty(0)),
+        modulate_index=modulate_index,
+    )
+
+    assert len(block.img_norm1.calls) == 1
+    assert len(block.img_norm2.calls) == 1
+    for _, scale, shift in (*block.img_norm1.calls, *block.img_norm2.calls):
+        assert scale.shape == hidden_states.shape
+        assert shift.shape == hidden_states.shape

tests/diffusion/models/wan2_2/conftest.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from contextlib import contextmanager
@@ -33,7 +36,7 @@ class StubScheduler:
         self.config = SimpleNamespace(num_train_timesteps=1000)
         self.set_timesteps_calls: list[tuple[int, torch.device]] = []
 
-    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
+    def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
         self.set_timesteps_calls.append((num_steps, device))
 
 

tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json

@@ -0,0 +1,97 @@
+{
+  "source_url": "https://github.com/Wan-Video/Wan2.2/blob/42bf4cfaa384bc21833865abc2f9e6c0e67233dc/wan/utils/fm_solvers_unipc.py",
+  "source_sha256": "0dec8c7ed17f6f2049275c6848113314da6ccec1c8db5bdc89df43c05c6038d9",
+  "git_head": "42bf4cfaa384bc21833865abc2f9e6c0e67233dc",
+  "num_train_timesteps": 1000,
+  "use_dynamic_shifting": false,
+  "reproduction": "Load FlowUniPCMultistepScheduler from source_url; construct with num_train_timesteps=1000, shift=1.0, use_dynamic_shifting=False; call set_timesteps(40, device='cpu', shift=12.0); record timesteps.tolist() and sigmas.tolist().",
+  "constructor_shift": 1.0,
+  "runtime_shift": 12.0,
+  "steps": 40,
+  "timesteps": [
+    999,
+    997,
+    995,
+    993,
+    990,
+    988,
+    985,
+    982,
+    979,
+    976,
+    972,
+    969,
+    965,
+    961,
+    956,
+    952,
+    947,
+    941,
+    936,
+    929,
+    922,
+    915,
+    907,
+    898,
+    888,
+    877,
+    865,
+    852,
+    837,
+    819,
+    799,
+    776,
+    749,
+    717,
+    678,
+    631,
+    571,
+    492,
+    386,
+    235
+  ],
+  "sigmas": [
+    0.9999166131019592,
+    0.9977826476097107,
+    0.9955461621284485,
+    0.9931996464729309,
+    0.9907346963882446,
+    0.988142192363739,
+    0.9854119420051575,
+    0.9825326800346375,
+    0.9794917702674866,
+    0.9762753844261169,
+    0.9728676676750183,
+    0.9692512154579163,
+    0.9654061794281006,
+    0.961310088634491,
+    0.9569376707077026,
+    0.95225989818573,
+    0.9472436308860779,
+    0.9418508410453796,
+    0.9360373020172119,
+    0.9297519326210022,
+    0.9229347705841064,
+    0.9155154228210449,
+    0.9074103832244873,
+    0.8985199332237244,
+    0.888724148273468,
+    0.8778772950172424,
+    0.865800678730011,
+    0.8522725701332092,
+    0.8370144367218018,
+    0.8196719288825989,
+    0.7997865080833435,
+    0.7767546772956848,
+    0.7497654557228088,
+    0.717703104019165,
+    0.6789888143539429,
+    0.6313128471374512,
+    0.5711563229560852,
+    0.49288034439086914,
+    0.38684695959091187,
+    0.23510953783988953,
+    0.0
+  ],
+  "implementation": "Wan-Video/Wan2.2 wan/utils/fm_solvers_unipc.py"
+}

tests/diffusion/models/wan2_2/test_create_transformer_quant.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
 """Regression tests for transformer quant-config auto-detection.
 
 The loader path at pipeline_wan2_2.py carries two quantization contracts:
@@ -27,6 +27,7 @@ from types import SimpleNamespace
 import pytest
 
 import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_
 
     # od_config must remain unchanged — _create_transformer does not modify it.
     assert od_config.quantization_config is None
+
+
+def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch):
+    import json
+
+    import torch
+
+    import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module
+    import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module
+
+    monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
+    for component in ("transformer", "transformer_2"):
+        directory = tmp_path / component
+        directory.mkdir()
+        config = dict(_MIN_CFG)
+        expert_quant_config = (
+            disk_quant_config[0 if component == "transformer" else 1]
+            if isinstance(disk_quant_config, tuple)
+            else disk_quant_config
+        )
+        if expert_quant_config is not None:
+            config["quantization_config"] = expert_quant_config
+        (directory / "config.json").write_text(json.dumps(config))
+    (tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]}))
+
+    class FakeComponent(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
+
+    class FakeTransformer(FakeComponent):
+        def __init__(self, quant_config=None, **kwargs):
+            super().__init__()
+            self.received_quant_config = quant_config
+
+    for module in (wan22_module, i2v_module):
+        monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
+        monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
+        monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent())
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+    monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer)
+    pipeline_cls = {
+        "t2v": wan22_module.Wan22Pipeline,
+        "i2v": i2v_module.Wan22I2VPipeline,
+        "vace": vace_module.Wan22VACEPipeline,
+    }[pipeline_kind]
+    monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None)
+    od_config = OmniDiffusionConfig(
+        model=str(tmp_path),
+        model_class_name=pipeline_cls.__name__,
+        dtype=torch.bfloat16,
+        quantization_config=active,
+        boundary_ratio=0.875,
+        flow_shift=12.0,
+        enable_diffusion_pipeline_profiler=False,
+    )
+    od_config.enrich_config()
+    pipeline = pipeline_cls(od_config=od_config)
+    return pipeline
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
+@pytest.mark.parametrize("offline", [False, True])
+def test_pipeline_constructor_routes_each_expert_quantization(
+    pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker
+):
+    """Run the actual pipeline constructors, expert factories and quant-method dispatch."""
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    from vllm_omni.quantization import build_quant_config
+
+    active = build_quant_config(
+        {
+            "transformer": {
+                "method": method,
+                "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
+                "w4a8_fallback_steps": [0, 1],
+            },
+            "transformer_2": {
+                "method": method,
+                "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
+                "w4a8_fallback_steps": [],
+            },
+        }
+    )
+    disk_quant_config = None
+    if offline:
+        disk_quant_config = {
+            "quant_method": method,
+            flag: True,
+            "ignored_layers": ["blocks.12.ffn.net_2"],
+            "w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"],
+            "w4a8_fallback_steps": [9],
+        }
+    pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch)
+    for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)):
+        received = transformer.received_quant_config
+        assert received.get_name() == method
+        assert getattr(received, flag) == offline
+        assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else [])
+        selected = f"blocks.{10 + index}.attn1.to_qkv"
+        other = f"blocks.{11 - index}.attn1.to_qkv"
+        assert received.w4a8_fallback_layers == [selected]
+        assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer
+        assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer
+        if offline:
+            assert isinstance(
+                received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod
+            )
+    assert pipeline.od_config.quantization_config is active
+
+
+@pytest.mark.parametrize("explicit_none", [False, True])
+def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch):
+    from vllm_omni.quantization import build_quant_config
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+    components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None}
+    active = build_quant_config(components)
+    pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
+    pipeline._create_transformer(_MIN_CFG, component="transformer_2")
+    assert captured[-1].get("quant_config") is None
+    with pytest.raises(ValueError, match="disabled for component 'transformer_2'"):
+        pipeline._create_transformer(
+            {**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}},
+            component="transformer_2",
+        )
+
+
+def test_pipeline_component_default_reaches_expert_factory(monkeypatch):
+    from vllm_omni.quantization import build_quant_config
+
+    FakeTransformer, captured = _make_fake_transformer()
+    monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
+    active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}})
+    pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
+    pipeline._create_transformer(_MIN_CFG, component="transformer_2")
+    assert captured[-1]["quant_config"] is active.default_config
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize(
+    "configured_expert, with_default",
+    [
+        ("transformer", False),
+        ("transformer", True),
+        ("transformer_2", False),
+        ("transformer_2", True),
+    ],
+)
+def test_pipeline_partial_expert_config_does_not_match_sibling(
+    pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch
+):
+    from vllm_omni.quantization import build_quant_config
+
+    component_spec = {
+        configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
+    }
+    if with_default:
+        component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]}
+    active = build_quant_config(component_spec)
+    pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch)
+    sibling = "transformer_2" if configured_expert == "transformer" else "transformer"
+    selected = getattr(pipeline, configured_expert).received_quant_config
+    other = getattr(pipeline, sibling).received_quant_config
+    assert selected is active.component_configs[configured_expert]
+    assert other is active.default_config
+    if with_default:
+        assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"]
+    else:
+        assert other is None
+
+
+@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
+@pytest.mark.parametrize("different_ignored_layers", [False, True])
+def test_enriched_auto_quantization_preserves_each_expert_policy(
+    pipeline_kind, different_ignored_layers, tmp_path, monkeypatch
+):
+    """Startup auto-detection must not turn the first expert into a user override."""
+    high = {
+        "quant_method": "mxfp4",
+        "is_checkpoint_mxfp4_serialized": True,
+        "ignored_layers": ["blocks.0.ffn.net_2"],
+        "w4a8_fallback_steps": [],
+        "w4a8_fallback_layers": [],
+        "mxfp4_scale_alg": 0,
+        "require_smooth_scale": False,
+    }
+    low = {
+        **high,
+        "ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"],
+        "w4a8_fallback_steps": [37],
+        "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
+        "mxfp4_scale_alg": 2,
+        "require_smooth_scale": True,
+    }
+    pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch)
+    assert pipeline.od_config.quantization_config is not None  # enrich_config really detected the first expert.
+    for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]:
+        config = transformer.received_quant_config
+        for policy in (
+            "ignored_layers",
+            "w4a8_fallback_steps",
+            "w4a8_fallback_layers",
+            "mxfp4_scale_alg",
+            "require_smooth_scale",
+        ):
+            assert getattr(config, policy) == expected[policy], policy

tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import io
+from types import SimpleNamespace
+
+import av
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import (
+    WAN_DEFAULT_OUTPUT_FPS,
+    resolve_wan_output_fps,
+    resolve_wan_preencode_batch_frames,
+    resolve_wan_preencode_mp4,
+    resolve_wan_video_codec_options,
+    wan_preencoded_mp4_payload,
+)
+from vllm_omni.diffusion.utils.chunked_video import decode_to_mp4
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeWanVAE:
+    """Publish temporal chunks the way the Wan VAE seam does."""
+
+    chunk_value_range = (-1.0, 1.0)
+
+    def __init__(self, *, batch: int, chunks: int, frames_per_chunk: int, height: int, width: int):
+        self.batch = batch
+        self.chunks = chunks
+        self.frames_per_chunk = frames_per_chunk
+        self.height = height
+        self.width = width
+        self.chunk_calls = 0
+
+    def decode_with_chunks(self, latents, *, on_chunk):
+        del latents
+        assert on_chunk is not None
+        for index in range(self.chunks):
+            self.chunk_calls += 1
+            value = -1.0 + 2.0 * index / max(self.chunks - 1, 1)
+            on_chunk(torch.full((self.batch, 3, self.frames_per_chunk, self.height, self.width), value))
+
+
+def _decoded_frame_count(data: bytes) -> int:
+    with av.open(io.BytesIO(data)) as container:
+        return sum(1 for _ in container.decode(video=0))
+
+
+def test_decode_wan_latents_to_mp4_emits_one_playable_mp4_per_batch_entry():
+    vae = _FakeWanVAE(batch=2, chunks=6, frames_per_chunk=2, height=16, width=32)
+
+    videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=4)
+
+    assert len(videos) == 2
+    for data in videos:
+        assert data[:4] and _decoded_frame_count(data) == 12
+
+
+def test_decode_wan_latents_to_mp4_coalesces_transfers_without_losing_frames():
+    """batch_frames only changes how many transfers carry the same frames."""
+    unbatched = _FakeWanVAE(batch=1, chunks=8, frames_per_chunk=1, height=16, width=32)
+    batched = _FakeWanVAE(batch=1, chunks=8, frames_per_chunk=1, height=16, width=32)
+
+    per_chunk = decode_to_mp4(unbatched, torch.zeros(1), fps=24, batch_frames=1)
+    coalesced = decode_to_mp4(batched, torch.zeros(1), fps=24, batch_frames=4)
+
+    assert _decoded_frame_count(per_chunk[0]) == _decoded_frame_count(coalesced[0]) == 8
+
+
+def test_decode_wan_latents_to_mp4_returns_nothing_for_a_rank_without_output():
+    class SilentVAE:
+        chunk_value_range = (-1.0, 1.0)
+
+        def decode_with_chunks(self, latents, *, on_chunk):
+            del latents, on_chunk
+
+    assert decode_to_mp4(SilentVAE(), torch.zeros(1), fps=24) == []
+
+
+def test_decode_wan_latents_to_mp4_rejects_a_non_positive_batch():
+    with pytest.raises(ValueError, match="batch_frames"):
+        decode_to_mp4(
+            _FakeWanVAE(batch=1, chunks=1, frames_per_chunk=1, height=16, width=16), None, fps=24, batch_frames=0
+        )
+
+
+def test_decode_wan_latents_to_mp4_rejects_a_vae_without_the_capability():
+    class PlainVAE:
+        chunk_value_range = (-1.0, 1.0)
+
+        def decode(self, z, return_dict=True):
+            del z, return_dict
+
+    with pytest.raises(TypeError, match="chunked VAE decode capability"):
+        decode_to_mp4(PlainVAE(), torch.zeros(1), fps=24)
+
+
+def test_preencode_flag_is_off_unless_the_request_asks_for_it():
+    assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args=None), output_type="np") is False
+    assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args={}), output_type="np") is False
+    assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args={"preencode_mp4": True}), output_type="np") is True
+
+
+def test_preencode_rejects_combinations_it_cannot_serve():
+    interpolating = SimpleNamespace(extra_args={"preencode_mp4": True}, enable_frame_interpolation=True)
+    with pytest.raises(ValueError, match="enable_frame_interpolation"):
+        resolve_wan_preencode_mp4(interpolating, output_type="np")
+
+    with pytest.raises(ValueError, match="output_type"):
+        resolve_wan_preencode_mp4(SimpleNamespace(extra_args={"preencode_mp4": True}), output_type="pil")
+
+
+def test_output_fps_mirrors_the_serving_fallback():
+    assert resolve_wan_output_fps(SimpleNamespace(fps=16)) == 16
+    assert resolve_wan_output_fps(SimpleNamespace(fps=[16])) == 16
+    assert resolve_wan_output_fps(SimpleNamespace(fps=None)) == WAN_DEFAULT_OUTPUT_FPS
+    assert resolve_wan_output_fps(SimpleNamespace()) == WAN_DEFAULT_OUTPUT_FPS
+
+
+def test_codec_options_reach_the_encoder_as_strings():
+    params = SimpleNamespace(extra_args={"video_codec_options": {"preset": "ultrafast", "threads": 0}})
+    assert resolve_wan_video_codec_options(params) == {"preset": "ultrafast", "threads": "0"}
+    assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={})) == {"preset": "ultrafast", "threads": "0"}
+    assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={"video_codec_options": None})) is None
+    assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={"video_codec_options": {}})) == {}
+
+
+def test_preencoded_payload_passes_bytes_through_and_ignores_tensors():
+    assert wan_preencoded_mp4_payload(torch.zeros(1, 3, 2, 4, 4)) is None
+    assert wan_preencoded_mp4_payload([b"one", b"two"]) == {"payload": {"video": [b"one", b"two"]}, "metadata": {}}
+    assert wan_preencoded_mp4_payload(b"one") == {"payload": {"video": [b"one"]}, "metadata": {}}
+
+
+@pytest.mark.parametrize("value", [1, 5, 33])
+def test_request_batch_threshold_controls_transfers_without_losing_frames(monkeypatch, value):
+    from vllm_omni.diffusion.utils import chunked_video
+
+    transfers = []
+    quantize = chunked_video.quantize_chunk
+
+    def capture(chunk, value_range):
+        transfers.append(chunk.shape[2])
+        return quantize(chunk, value_range)
+
+    monkeypatch.setattr(chunked_video, "quantize_chunk", capture)
+    params = SimpleNamespace(extra_args={"preencode_mp4": True, "preencode_batch_frames": value})
+    vae = _FakeWanVAE(batch=1, chunks=4, frames_per_chunk=3, height=16, width=32)
+    videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=resolve_wan_preencode_batch_frames(params))
+
+    assert transfers == {1: [3, 3, 3, 3], 5: [6, 6], 33: [12]}[value]
+    assert _decoded_frame_count(videos[0]) == 12
+
+
+@pytest.mark.parametrize("value", [0, -1, True, 1.5, "17", None])
+def test_request_batch_threshold_rejects_invalid_values(value):
+    with pytest.raises(ValueError, match="preencode_batch_frames"):
+        resolve_wan_preencode_batch_frames(SimpleNamespace(extra_args={"preencode_batch_frames": value}))
+
+
+def test_request_batch_threshold_preserves_model_defaults():
+    params = SimpleNamespace(extra_args=None)
+    assert resolve_wan_preencode_batch_frames(params) == 17
+    assert resolve_wan_preencode_batch_frames(params, default=1) == 1

tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py

@@ -1,15 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 from PIL import Image
 from torch import nn
 
 from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN, build_wan_scheduler
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
     get_wan22_i2v_post_process_func,
@@ -50,7 +51,7 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
 
 
 def _make_i2v_sampling(**overrides):
-    values = {
+    values: dict[str, object] = {
         "height": 16,
         "width": 16,
         "num_frames": 5,
@@ -275,16 +276,18 @@ def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
     assert first_frame_mask.shape == (2, 1, 2, 2, 2)
 
 
-def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
+@pytest.mark.parametrize("solver", ["unipc", "euler"])
+@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
+def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch, solver: str, shift: float) -> None:
     monkeypatch.setattr(
         "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
         SimpleNamespace(is_available=lambda: False),
     )
     pipeline = _make_i2v_pipeline(expand_timesteps=True)
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 5.0
+    pipeline.scheduler = build_wan_scheduler(solver, shift)
+    pipeline.od_config = SimpleNamespace(flow_shift=shift)
+    pipeline._sample_solver = solver
+    pipeline._flow_shift = shift
     pipeline.boundary_ratio = 0.875
     pipeline.has_image_encoder = False
     pipeline._guidance_scale = None
@@ -327,6 +330,8 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -
                     generator=gen_a,
                     latents=latents_a,
                     num_outputs_per_prompt=2,
+                    num_inference_steps=50,
+                    extra_args={"sample_solver": solver, "flow_shift": shift},
                 ),
             ),
             SimpleNamespace(
@@ -336,12 +341,25 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -
                     generator=gen_b,
                     latents=latents_b,
                     num_outputs_per_prompt=2,
+                    num_inference_steps=50,
+                    extra_args={"sample_solver": solver, "flow_shift": shift},
                 ),
             ),
         ]
     )
 
     outputs = pipeline.forward(batch)
+    if solver == "unipc":
+        sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
+        sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
+        torch.testing.assert_close(
+            pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
+        )
+        assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
+    else:
+        reference = build_wan_scheduler("euler", shift)
+        reference.set_timesteps(50, device="cpu")
+        torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
 
     assert prepare_call["batch_size"] == 4
     assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]

tests/diffusion/models/wan2_2/test_wan22_mindie_metadata.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Wan integration contracts required by quantized dense/sparse attention."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformerBlock
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_block_publishes_unpadded_grid_without_partial_packed_metadata():
+    block = object.__new__(WanTransformerBlock)
+    torch.nn.Module.__init__(block)
+    block.scale_shift_table = torch.zeros(1, 6, 4)
+    block.norm1 = lambda x, *a: x
+    block.norm2 = lambda x: x
+    block.norm3 = lambda x, *a: x
+    block.attn1 = Mock(side_effect=lambda x, *a: torch.zeros_like(x))
+    block.attn2 = Mock(side_effect=lambda x, *a: torch.zeros_like(x))
+    block.ffn = lambda x: torch.zeros_like(x)
+    # Local SP tokens include alignment padding; grid describes the full valid video.
+    x = torch.zeros(1, 1152, 4)
+    block(
+        x,
+        x,
+        torch.zeros(1, 6, 4),
+        None,
+        hidden_states_mask=torch.ones(1, 1152, dtype=torch.bool),
+        vsa_dit_seq_shape=(5, 16, 27),
+    )
+    metadata = block.attn1.call_args.args[2]
+    assert metadata.video_layout.latent_grid == (5, 16, 27)
+    assert metadata.video_layout.used_len == 2160
+    assert metadata.video_layout.prefix_len == 0
+    assert "max_seqlen_q" not in metadata.extra  # CUDA must not infer incomplete packed varlen.
+    assert block.attn2.call_args.args[2] is None

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -9,7 +9,8 @@ import pytest
 import torch
 from torch import nn
 
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
+from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, build_wan_scheduler
 from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -18,6 +19,14 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.fixture(autouse=True)
+def _cpu_pipeline_runtime(monkeypatch):
+    # These lightweight pipelines have no accelerator or distributed groups.
+    module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
+    monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
+    monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1))
+
+
 class _StubTransformer(nn.Module):
     @property
     def dtype(self) -> torch.dtype:
@@ -30,14 +39,30 @@ class _StubTextEncoder(nn.Module):
         return torch.float32
 
 
+class _StubVaeConfig:
+    latents_mean = [0.0, 0.0, 0.0, 0.0]
+    latents_std = [1.0, 1.0, 1.0, 1.0]
+    z_dim = 4
+
+
+class _StubVae(nn.Module):
+    dtype = torch.float32
+    config = _StubVaeConfig()
+
+    def decode(self, latents, return_dict=False):
+        del return_dict
+        batch, _, frames, height, width = latents.shape
+        return (torch.zeros(batch, 3, frames, height, width),)
+
+
 class _StubScheduler:
     def __init__(self, timesteps: list[int]) -> None:
         self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
         self.config = SimpleNamespace(num_train_timesteps=1000)
-        self.set_timesteps_calls: list[tuple[int, torch.device]] = []
+        self.set_timesteps_calls: list[tuple[int, torch.device, float | None]] = []
 
-    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
-        self.set_timesteps_calls.append((num_steps, device))
+    def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
+        self.set_timesteps_calls.append((num_steps, device, shift))
 
 
 @contextmanager
@@ -75,6 +100,7 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline.transformer = _StubTransformer()
     pipeline.transformer_2 = None
     pipeline.text_encoder = _StubTextEncoder()
+    pipeline.vae = _StubVae()
     pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4)
     pipeline.scheduler = _StubScheduler([9, 5])
     pipeline.od_config = SimpleNamespace(flow_shift=5.0)
@@ -166,7 +192,24 @@ def test_forward_delegates_denoising_to_diffuse(
     assert captured["boundary_timestep"] == pytest.approx(875.0)
     assert captured["latent_condition"] is None
     assert captured["first_frame_mask"] is None
-    assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
+    assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"), 5.0)]
+
+
+@pytest.mark.parametrize("solver", ["unipc", "euler"])
+def test_forward_passes_request_shift_without_mutating_scheduler_config(solver: str) -> None:
+    pipeline = _make_pipeline()
+    pipeline.diffuse = lambda **kwargs: kwargs["latents"]
+    for shift in (3.0, 12.0, 5.0):
+        sampling = _make_sampling(num_inference_steps=5, extra_args={"sample_solver": solver, "flow_shift": shift})
+        request = OmniDiffusionRequest(prompt="prompt", request_id="schedule", sampling_params=sampling)
+        pipeline.forward(DiffusionRequestBatch(requests=[request]))
+        reference = build_wan_scheduler(solver, shift)
+        if solver == "unipc":
+            reference.set_timesteps(5, device="cpu", shift=shift)
+            assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
+        else:
+            reference.set_timesteps(5, device="cpu")
+        torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
 
 
 def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
@@ -227,6 +270,99 @@ def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
     torch.testing.assert_close(outputs[1].output, latents_b)
 
 
+def test_forward_emits_request_local_typed_media_after_vae_decode() -> None:
+    pipeline = _make_pipeline()
+
+    def _fake_diffuse(
+        *,
+        latents,
+        timesteps,
+        prompt_embeds,
+        negative_prompt_embeds,
+        guidance_low,
+        guidance_high,
+        boundary_timestep,
+        dtype,
+        attention_kwargs,
+        latent_condition,
+        first_frame_mask,
+        generator,
+    ):
+        del (
+            timesteps,
+            prompt_embeds,
+            negative_prompt_embeds,
+            guidance_low,
+            guidance_high,
+            boundary_timestep,
+            dtype,
+            attention_kwargs,
+            latent_condition,
+            first_frame_mask,
+            generator,
+        )
+        return torch.zeros_like(latents)
+
+    pipeline.diffuse = _fake_diffuse  # type: ignore[method-assign]
+    batch = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt="prompt",
+                request_id="request-0",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_frames=1,
+                    num_inference_steps=2,
+                    max_sequence_length=32,
+                    output_type="np",
+                ),
+            )
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert len(outputs) == 1
+    assert outputs[0].output is None
+    assert outputs[0].media is not None
+    assert outputs[0].media.prepared_for_transport is False
+    assert outputs[0].media.video.tensor.shape == (1, 3, 1, 8, 8)
+    assert outputs[0].media.video.spec.layout is VideoTensorLayout.BCTHW
+    assert outputs[0].media.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+    assert outputs[0].media.video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE
+
+
+def test_forward_keeps_legacy_output_on_non_owner_vae_rank() -> None:
+    # Distributed VAE decode uses broadcast_result=False, so non-owner ranks get
+    # an empty placeholder instead of the full video. Wrapping that as typed media
+    # would fail split_diffusion_output_by_request's batch check on every non-owner
+    # rank, so the pipeline must keep the placeholder on the legacy output field.
+    pipeline = _make_pipeline()
+    pipeline.vae.decode = lambda latents, return_dict=False: (torch.empty(0),)  # type: ignore[assignment]
+    pipeline.diffuse = lambda **kwargs: torch.zeros_like(kwargs["latents"])  # type: ignore[method-assign]
+
+    batch = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt="prompt",
+                request_id="request-0",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_frames=1,
+                    num_inference_steps=2,
+                    max_sequence_length=32,
+                    output_type="np",
+                ),
+            )
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert len(outputs) == 1
+    assert outputs[0].media is None
+    assert outputs[0].output is not None
+    assert outputs[0].output.numel() == 0
+
+
 def test_forward_batches_precomputed_prompt_embeddings() -> None:
     pipeline = _make_pipeline()
     diffuse_call = {}
@@ -343,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
     assert torch.equal(result, torch.full_like(latents, 10.0))
 
 
+@pytest.mark.parametrize("fail_first_request", [False, True])
+def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.forward_context import (
+        ForwardContext,
+        get_forward_context,
+        is_forward_context_available,
+        override_forward_context,
+    )
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
+
+    contexts: list[ForwardContext] = []
+    events = []
+
+    class InjectedPredictionError(Exception):
+        pass
+
+    class RecordingTransformer(_StubTransformer):
+        def __init__(self, expert):
+            super().__init__()
+            self.expert = expert
+
+        def forward(self, hidden_states, encoder_hidden_states, **kwargs):
+            context = get_forward_context()
+            assert context is contexts[-1]
+            request_index = len(contexts) - 1
+            step = context.denoise_step_idx
+            positive = bool(encoder_hidden_states[0, 0, 0] > 0)
+            events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2])))
+            if fail_first_request and request_index == 0 and step == 1 and not positive:
+                raise InjectedPredictionError
+            return (torch.ones_like(hidden_states) * (1 if positive else -1),)
+
+    def encode_prompt(**kwargs):
+        context = get_forward_context()
+        assert context.denoise_step_idx is None
+        contexts.append(context)
+        embeds = torch.ones(1, 8, 8)
+        return embeds, -embeds
+
+    pipeline = _make_pipeline()
+    pipeline.transformer = RecordingTransformer("high")
+    pipeline.transformer_2 = RecordingTransformer("low")
+    pipeline.scheduler = _StubScheduler([900, 500, 100])
+    monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt)
+    monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents)
+
+    # Use the real runner context manager, Wan forward/diffuse, and CFG
+    # dispatch. Only model computation, scheduler math and unrelated KV I/O
+    # are replaced; no test code publishes a denoise step.
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = pipeline
+    runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32)
+    runner.vllm_config = None
+    runner.cache_backend = None
+    monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None)
+    with override_forward_context(None):
+        for request_index in range(2):
+            request = OmniDiffusionRequest(
+                prompt="a fox walks",
+                request_id=f"request-{request_index}",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent"
+                ),
+            )
+
+            def execute():
+                return runner._execute_request_list(
+                    [request],
+                    od_config=runner.od_config,
+                    allow_single_output=True,
+                    require_request_batch_support=False,
+                    record_name="test_wan_step_context",
+                    record_output_peak_memory=False,
+                )
+
+            if fail_first_request and request_index == 0:
+                with pytest.raises(InjectedPredictionError):
+                    execute()
+            else:
+                execute()
+            assert not is_forward_context_available()
+
+    assert len(contexts) == 2 and contexts[0] is not contexts[1]
+    expected = [
+        (request, step, "high" if step == 0 else "low", positive, step in (0, 2))
+        for request in range(2)
+        for step in range(2 if fail_first_request and request == 0 else 3)
+        for positive in (True, False)
+    ]
+    assert events == expected
+
+
 class _StubDMDScheduler:
     def __init__(self) -> None:
         self.predict_clean_calls: list[tuple[float, float, float]] = []

tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.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 Wan2.2 quant_config propagation through transformer creation.
 
 Tests cover:
@@ -196,7 +196,7 @@ class TestPatchWanRmsNorm:
     """Test that patch_wan_rms_norm doesn't raise on concurrent module registration."""
 
     def test_patches_modules_with_wan_rms_norm(self):
-        from vllm_omni.diffusion.layers.norm import RMSNormVAE
+        from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
         from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm
 
         # Create a fake module that has WanRMS_norm

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -1,4 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
+from typing import Any
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -7,6 +11,7 @@ import pytest
 import torch
 import torch.nn as nn
 
+from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
     Wan22S2VPipeline,
     _make_clip_generators,
@@ -14,11 +19,32 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
 from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("configured_shift", [None, 12.0])
+def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float | None) -> None:
+    module = "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v"
+    config = SimpleNamespace(model="unused", flow_shift=configured_shift, enable_diffusion_pipeline_profiler=False)
+
+    def init_components(pipeline, *args):
+        pipeline.vae = SimpleNamespace(config=SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8))
+
+    with (
+        patch(f"{module}.get_local_device", return_value=torch.device("cpu")),
+        patch(f"{module}._resolve_model_path", return_value="unused"),
+        patch(f"{module}._is_diffusers_format", return_value=True),
+        patch.object(Wan22S2VPipeline, "_init_diffusers_format", init_components),
+        patch.object(Wan22S2VPipeline, "setup_diffusion_pipeline_profiler"),
+    ):
+        pipeline = Wan22S2VPipeline(od_config=config)
+    assert pipeline._flow_shift == (3.0 if configured_shift is None else configured_shift)
+    assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
+    assert pipeline.scheduler.sigma_max == float(np.float32(0.999))
 
 
 def _make_s2v_sampling(**overrides):
-    values = {
+    values: dict[str, Any] = {
         "height": 16,
         "width": 16,
         "num_frames": 8,
@@ -211,7 +237,8 @@ def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
         pipeline.forward(batch)
 
 
-def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
+@pytest.mark.parametrize("shift", [3.0, 12.0])
+def test_s2v_forward_batches_request_local_inputs_and_splits_outputs(shift: float) -> None:
     pipeline = object.__new__(Wan22S2VPipeline)
     nn.Module.__init__(pipeline)
     pipeline.device = torch.device("cpu")
@@ -227,7 +254,8 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
         enable_cpu_offload=False,
         parallel_config=SimpleNamespace(use_hsdp=False),
     )
-    pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
+    pipeline._flow_shift = shift
+    pipeline.scheduler = FlowUniPCMultistepScheduler(shift=1.0)
     pipeline.vae_scale_factor_spatial = 8
     pipeline.resolution_divisor = 16
     pipeline.motion_frames = 7
@@ -262,6 +290,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
                     "multi_modal_data": {"image": image, "audio": audio_a},
                 },
                 sampling_params=_make_s2v_sampling(
+                    num_inference_steps=5,
                     num_outputs_per_prompt=2,
                     generator=generators_a,
                     latents=latents_a,
@@ -275,6 +304,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
                     "multi_modal_data": {"image": image, "audio": audio_b},
                 },
                 sampling_params=_make_s2v_sampling(
+                    num_inference_steps=5,
                     num_outputs_per_prompt=2,
                     generator=generators_b,
                     latents=latents_b,
@@ -287,6 +317,12 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
         platform.is_available.return_value = False
         outputs = pipeline.forward(batch)
 
+    expected = np.linspace(float(np.float32(0.999)), 0.0, 6)[:-1]
+    expected = shift * expected / (1.0 + (shift - 1.0) * expected)
+    torch.testing.assert_close(
+        pipeline.scheduler.sigmas, torch.tensor(np.append(expected, 0.0), dtype=torch.float32), rtol=0, atol=0
+    )
+    assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
     assert len(outputs) == 2
     assert outputs[0].output[0].shape[0] == 2
     assert outputs[1].output[0].shape[0] == 2
@@ -519,6 +555,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
     mock_scheduler.timesteps = torch.linspace(999, 0, 5)
     mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),))
     pipeline.scheduler = mock_scheduler
+    pipeline._flow_shift = 3.0
 
     # -- Bind methods from the real class --
     pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline)
@@ -611,3 +648,245 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
     assert video.shape[1] == 3  # channels
     assert audio_waveform is not None
     assert audio_sr == 16000
+
+
+def _make_s2v_preencode_pipeline() -> Wan22S2VPipeline:
+    """Build the same stub pipeline the batching test drives, for preencode runs."""
+    pipeline = object.__new__(Wan22S2VPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = MagicMock()
+    pipeline.transformer.dtype = torch.float32
+    # A fresh iterator per call: the clip loop reads this once per clip.
+    pipeline.transformer.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float32)])
+    pipeline.transformer.casual_audio_encoder = None
+    pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
+    pipeline.vae = MagicMock()
+    pipeline.vae.dtype = torch.float32
+    # The shared consumer reads the published pixel range off the VAE.
+    pipeline.vae.chunk_value_range = (-1.0, 1.0)
+    # One decoded clip: [B, C, T, H, W] for a batch of four (2 requests x 2 outputs).
+    pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
+    pipeline._flow_shift = 3.0
+    pipeline.vae_scale_factor_spatial = 8
+    pipeline.resolution_divisor = 16
+    pipeline.motion_frames = 7
+    pipeline.drop_first_motion = True
+    pipeline._DEFAULT_INFER_FRAMES = 8
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline.check_inputs = lambda *args, **kwargs: None
+    pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
+    pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
+    pipeline.prepare_motion_latents = MagicMock(
+        side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
+    )
+    pipeline._denormalize_latents = lambda latents: latents
+    pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])
+    return pipeline
+
+
+def _make_s2v_preencode_batch(audio_a, audio_b, **sampling_overrides) -> DiffusionRequestBatch:
+    image = PIL.Image.new("RGB", (16, 16))
+    return DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first", "multi_modal_data": {"image": image, "audio": audio_a}},
+                sampling_params=_make_s2v_sampling(
+                    num_outputs_per_prompt=2,
+                    latents=torch.zeros(2, 16, 2, 2, 2),
+                    output_type="np",
+                    **sampling_overrides,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "multi_modal_data": {"image": image, "audio": audio_b}},
+                sampling_params=_make_s2v_sampling(
+                    num_outputs_per_prompt=2,
+                    latents=torch.ones(2, 16, 2, 2, 2),
+                    output_type="np",
+                    **sampling_overrides,
+                ),
+            ),
+        ]
+    )
+
+
+def _decode_mp4(data: bytes):
+    """Return (video frame count, first audio samples) for one encoded container."""
+    import io
+
+    import av
+
+    with av.open(io.BytesIO(data)) as container:
+        frames = sum(1 for _ in container.decode(video=0))
+    with av.open(io.BytesIO(data)) as container:
+        assert container.streams.audio, "the worker must mux the waveform into the container"
+        samples = np.concatenate(
+            [frame.to_ndarray().reshape(-1) for frame in container.decode(audio=0)][:4],
+        )
+    return frames, samples
+
+
+@pytest.mark.parametrize("batch_frames", [1, 1000])
+def test_s2v_preencode_returns_playable_mp4_bytes_per_request(monkeypatch, batch_frames) -> None:
+    """The clip loop hands finished clips to the encoder instead of concatenating."""
+    pipeline = _make_s2v_preencode_pipeline()
+    # Two clips, so the autoregressive motion feedback runs between pushes.
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
+    audio_a = np.zeros(16000, dtype=np.float32)
+    audio_b = np.full(16000, 0.5, dtype=np.float32)
+    batch = _make_s2v_preencode_batch(
+        audio_a, audio_b, extra_args={"preencode_mp4": True, "preencode_batch_frames": batch_frames}
+    )
+    from vllm_omni.diffusion.utils import chunked_video
+
+    transfers = []
+    quantize = chunked_video.quantize_chunk
+
+    def capture(chunk, value_range):
+        transfers.append(chunk.shape[2])
+        return quantize(chunk, value_range)
+
+    monkeypatch.setattr(chunked_video, "quantize_chunk", capture)
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        outputs = pipeline.forward(batch)
+
+    assert len(outputs) == 2
+    assert pipeline.vae.decode.call_count == 2, "both clips must reach the encoder"
+    assert len(transfers) == (2 if batch_frames == 1 else 1)
+
+    decoded = []
+    for request_output in outputs:
+        # Two outputs per prompt, each already a complete container.
+        assert isinstance(request_output.output, list)
+        assert len(request_output.output) == 2
+        for data in request_output.output:
+            assert isinstance(data, bytes)
+            frames, samples = _decode_mp4(data)
+            assert frames > 0
+            decoded.append(samples)
+
+    # Four entries: request a's waveform for its two outputs, then request b's.
+    # A silent and a non-silent waveform must not land on the same request.
+    assert np.abs(decoded[0]).max() == pytest.approx(0.0, abs=1e-3)
+    assert np.abs(decoded[1]).max() == pytest.approx(0.0, abs=1e-3)
+    assert np.abs(decoded[2]).max() > 0.1
+    assert np.abs(decoded[3]).max() > 0.1
+
+
+def test_s2v_preencode_aborts_encoders_when_a_later_clip_fails(monkeypatch) -> None:
+    """A failure after the first push must not strand encoder worker threads."""
+    pipeline = _make_s2v_preencode_pipeline()
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
+    batch = _make_s2v_preencode_batch(
+        np.zeros(16000, dtype=np.float32),
+        np.ones(16000, dtype=np.float32),
+        extra_args={"preencode_mp4": True, "preencode_batch_frames": 1},
+    )
+    sessions = []
+
+    from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
+
+    def capture_session(**kwargs):
+        session = ChunkedVideoMP4Session(**kwargs)
+        sessions.append(session)
+        return session
+
+    def fail_on_second_clip(**kwargs):
+        if pipeline.vae.decode.call_count:
+            raise RuntimeError("second clip failed")
+        return kwargs["latents"]
+
+    pipeline.diffuse = fail_on_second_clip
+    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.ChunkedVideoMP4Session", capture_session)
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        with pytest.raises(RuntimeError, match="second clip failed"):
+            pipeline.forward(batch)
+
+    assert len(sessions) == 1
+    assert len(sessions[0]._encoders) == 4
+    assert all(not encoder._thread.is_alive() for encoder in sessions[0]._encoders)
+
+
+def test_s2v_preencode_skips_encoding_on_a_vae_patch_parallel_peer(monkeypatch) -> None:
+    """A peer still receives every clip for the motion loop but encodes none of them."""
+    pipeline = _make_s2v_preencode_pipeline()
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
+    # A peer's patch-parallel decode returns an empty placeholder.
+    pipeline.vae.decode.return_value = (torch.empty(0),)
+    pipeline.vae._vae_pp_group = object()
+    broadcasts = []
+    monkeypatch.setattr(
+        "torch.distributed.broadcast",
+        lambda tensor, src, group: broadcasts.append(tensor.zero_()),
+    )
+    batch = _make_s2v_preencode_batch(
+        np.zeros(16000, dtype=np.float32),
+        np.ones(16000, dtype=np.float32),
+        extra_args={"preencode_mp4": True, "preencode_batch_frames": 1},
+    )
+    sessions = []
+
+    from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
+
+    def capture_session(**kwargs):
+        session = ChunkedVideoMP4Session(**kwargs)
+        sessions.append(session)
+        return session
+
+    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.ChunkedVideoMP4Session", capture_session)
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        outputs = pipeline.forward(batch)
+
+    assert len(broadcasts) == 2, "the peer must still join the broadcast for both clips"
+    # Once for the initial motion latents, once more from the peer's broadcast clip.
+    assert pipeline.prepare_motion_latents.call_count == 2, "the motion loop still consumes the peer's frames"
+    assert sessions[0]._encoders == []
+    assert [request_output.output for request_output in outputs] == [[], []]
+
+
+@pytest.mark.parametrize("output_type", ["pil", "pt", "latent"])
+def test_s2v_preencode_rejects_request_output_types_it_cannot_serve(output_type) -> None:
+    """Pre-encoding returns MP4 bytes, so a request asking for frames must fail."""
+    pipeline = _make_s2v_preencode_pipeline()
+    batch = _make_s2v_preencode_batch(
+        np.zeros(16000, dtype=np.float32),
+        np.ones(16000, dtype=np.float32),
+        extra_args={"preencode_mp4": True},
+    )
+    for request in batch.requests:
+        request.sampling_params.output_type = output_type
+
+    with pytest.raises(ValueError, match="output_type"):
+        pipeline.forward(batch)
+
+
+def test_s2v_preencode_keeps_the_full_decode_path_untouched() -> None:
+    """Without the flag the loop still returns the (video, audio, rate) tuple."""
+    pipeline = _make_s2v_preencode_pipeline()
+    audio_a = np.zeros(16000, dtype=np.float32)
+    audio_b = np.ones(16000, dtype=np.float32)
+    batch = _make_s2v_preencode_batch(audio_a, audio_b)
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        outputs = pipeline.forward(batch)
+
+    assert outputs[0].output[0].shape[0] == 2
+    np.testing.assert_array_equal(outputs[0].output[1], audio_a)
+    np.testing.assert_array_equal(outputs[1].output[1], audio_b)

tests/diffusion/models/wan2_2/test_wan22_scheduler.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare actual Wan schedules to native Wan outputs and independent sigma arithmetic."""
+
+import json
+from pathlib import Path
+
+import numpy as np
+import numpy.typing as npt
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    FASTWAN_DMD_SCHEDULER_SHIFT,
+    build_wan_scheduler,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_wan_unipc_matches_captured_native_full_schedule() -> None:
+    # The fixture pins the official Wan source and records how to regenerate it.
+    golden = json.loads((Path(__file__).parent / "fixtures/native_unipc_shift12.json").read_text())
+    scheduler = build_wan_scheduler("unipc", golden["runtime_shift"])
+    for steps in (golden["steps"], 20, golden["steps"]):
+        scheduler.set_timesteps(steps, device="cpu", shift=golden["runtime_shift"])
+        if steps == golden["steps"]:
+            assert scheduler.timesteps.tolist() == golden["timesteps"]
+            torch.testing.assert_close(scheduler.sigmas, torch.tensor(golden["sigmas"]), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("shift", [1.0, 3.0, 5.0, 12.0])
+@pytest.mark.parametrize("steps", [20, 40])
+def test_wan_unipc_shifts_unmodified_training_endpoints(shift: float, steps: int) -> None:
+    # Native Wan: float32 training sigma_max=(1-1/1000), sigma_min=0;
+    # interpolate in float64, apply the runtime shift, then cast sigma to FP32.
+    unshifted = np.linspace(float(np.float32(0.999)), 0.0, steps + 1)[:-1]
+    shifted = shift * unshifted / (1.0 + (shift - 1.0) * unshifted)
+    expected_timesteps = (shifted * 1000).astype(np.int64)
+    expected_sigmas: npt.NDArray[np.float32] = np.append(shifted, 0.0).astype(np.float32)
+    scheduler = build_wan_scheduler("unipc", shift)
+    scheduler.set_timesteps(steps, device="cpu", shift=shift)
+    assert scheduler.config.shift == scheduler.config["shift"] == 1.0
+    torch.testing.assert_close(scheduler.timesteps, torch.from_numpy(expected_timesteps), rtol=0, atol=0)
+    torch.testing.assert_close(scheduler.sigmas, torch.from_numpy(expected_sigmas), rtol=0, atol=0)
+
+
+def test_wan_euler_keeps_dmd_shift_semantics() -> None:
+    scheduler = build_wan_scheduler("euler", FASTWAN_DMD_SCHEDULER_SHIFT)
+    scheduler.set_timesteps(3, device="cpu")
+    # The unchanged Euler branch starts at sigma=1, with shift=8 for DMD.
+    expected = torch.tensor([1.0, 16.0 / 17.0, 0.8, 0.0])
+    torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=1e-7)
+    torch.testing.assert_close(scheduler.timesteps, expected[:-1] * 1000, rtol=0, atol=1e-4)

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -1,15 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from dataclasses import dataclass, field
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 from PIL import Image
 from torch import nn
 
 from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import build_wan_scheduler
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
     Wan22VACEPipeline,
     create_vace_transformer_from_config,
@@ -33,6 +35,8 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
         )
         enable_layerwise_offload: bool = True
         post_init_calls: int = 0
+        _resolved_diffusion_offload: object = None
+        _diffusion_offload_flags_materialized: bool = False
 
         def __post_init__(self) -> None:
             self.post_init_calls += 1
@@ -41,7 +45,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
     resolved_offload = object()
     config._resolved_diffusion_offload = resolved_offload
     config._diffusion_offload_flags_materialized = True
-    captured = {}
+    captured: dict[str, object] = {}
 
     def capture_base_init(self, *, od_config, prefix="") -> None:
         captured.update(od_config=od_config, prefix=prefix)
@@ -51,6 +55,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
     Wan22VACEPipeline(od_config=config, prefix="stage")
 
     forwarded = captured["od_config"]
+    assert isinstance(forwarded, RuntimeConfig)
     assert forwarded is not config
     assert config.flow_shift is None
     assert forwarded.flow_shift == 3.0
@@ -75,7 +80,7 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
 
 
 def _make_vace_sampling(**overrides):
-    values = {
+    values: dict[str, object] = {
         "height": 16,
         "width": 16,
         "num_frames": 5,
@@ -123,7 +128,7 @@ def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
 
 
 def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None:
-    captured = {}
+    captured: dict[str, object] = {}
 
     class FakeVACETransformer:
         def __init__(self, **kwargs) -> None:
@@ -203,7 +208,9 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
     torch.testing.assert_close(result, torch.ones_like(latents))
 
 
-def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
+@pytest.mark.parametrize("solver", ["unipc", "euler"])
+@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
+def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch, solver: str, shift: float) -> None:
     monkeypatch.setattr(
         "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
         SimpleNamespace(is_available=lambda: False),
@@ -211,10 +218,10 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
     pipeline = _make_vace_pipeline()
     pipeline.transformer.vace_patch_embedding = None
     pipeline.transformer_2 = None
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 3.0
+    pipeline.scheduler = build_wan_scheduler(solver, shift)
+    pipeline.od_config = SimpleNamespace(flow_shift=shift)
+    pipeline._sample_solver = solver
+    pipeline._flow_shift = shift
     pipeline.boundary_ratio = None
     pipeline._guidance_scale = None
     pipeline._num_timesteps = None
@@ -247,6 +254,8 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
                     generator=gen_a,
                     latents=latents_a,
                     num_outputs_per_prompt=2,
+                    num_inference_steps=50,
+                    extra_args={"sample_solver": solver, "flow_shift": shift},
                 ),
             ),
             SimpleNamespace(
@@ -256,12 +265,25 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
                     generator=gen_b,
                     latents=latents_b,
                     num_outputs_per_prompt=2,
+                    num_inference_steps=50,
+                    extra_args={"sample_solver": solver, "flow_shift": shift},
                 ),
             ),
         ]
     )
 
     outputs = pipeline.forward(batch)
+    if solver == "unipc":
+        sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
+        sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
+        torch.testing.assert_close(
+            pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
+        )
+        assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
+    else:
+        reference = build_wan_scheduler("euler", shift)
+        reference.set_timesteps(50, device="cpu")
+        torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
 
     assert prepare_call["batch_size"] == 4
     assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]