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/ar_diffusion/test_capability_runner.py

@@ -323,12 +323,18 @@ def test_lingbot_like_sink_survives_sliding_window_eviction():
     kv = runner.kv_cache
     assert kv is not None
 
-    for _ in range(8):
+    state = commit_one_frame(runner, "s1", "main")
+    sink_block = kv.block_table(state.adapter("main"))[0]
+    for _ in range(7):
         state = commit_one_frame(runner, "s1", "main")
 
-    table = kv.block_table(state.adapter("main"))
-    assert table[0] != kv.null_block_id
-    assert table[1] == kv.null_block_id
+    adapter = state.adapter("main")
+    table = kv.block_table(adapter)
+    assert len(table) == 6  # one sink + five recent frames
+    assert table[0] == sink_block
+    assert kv.null_block_id not in table
+    assert adapter.absolute_num_computed_tokens == 8 * BLOCK
+    assert adapter.num_computed_tokens == 6 * BLOCK
 
     ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=False)[0].forward_ctx
     visible, _ = ctx.video_block_table(torch.device("cpu"))

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
 
 Covers the request adapter, the chunk-window spec/manager (registration + the
@@ -269,9 +270,14 @@ def test_compiled_paged_write_does_not_clone_the_pool():
     # drives the reinplace pass, so a drift here would silently test nothing.
     real_schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema
     real_mutated = [arg.name for arg in real_schema.arguments if arg.alias_info and arg.alias_info.is_write]
-    assert real_mutated == ["key_pool", "value_pool"], (
+    # The staging buffers are mutated too when reuse_history_staging is on, but they are separate
+    # allocations and not what this test is about; the pools are what must not be cloned.
+    assert real_mutated[:2] == ["key_pool", "value_pool"], (
         f"the real paged-write op now mutates {real_mutated}; this test still models two pools"
     )
+    assert set(real_mutated) - {"key_pool", "value_pool"} <= {"stage_key", "stage_value"}, (
+        f"unexpected mutated argument in the real paged-write op: {real_mutated}"
+    )
 
     num_blocks, heads, dim = 8, 4, 16
     cache_shape = (num_blocks, BLOCK, heads, dim)
@@ -357,6 +363,7 @@ def _make_tiny_capacity_kv(
     available_bytes: int,
     gpu_memory_fraction: float = 1.0,
     model_owned_state_bytes_per_session: int = 0,
+    reuse_history_staging: bool = False,
 ) -> ARDiffusionKVCache:
     return ARDiffusionKVCache(
         ARDiffusionKVConfig(
@@ -365,6 +372,7 @@ def _make_tiny_capacity_kv(
             window_chunks=3,
             sink_chunks=3,
             gpu_memory_fraction=gpu_memory_fraction,
+            reuse_history_staging=reuse_history_staging,
         ),
         num_layers=1,
         num_kv_heads=1,
@@ -382,6 +390,33 @@ def _make_tiny_capacity_kv(
     )
 
 
+def test_history_staging_is_reserved_in_the_kv_budget(monkeypatch):
+    """The per-layer staging pair is worker-wide memory allocated after admission; the budget must hold it."""
+    monkeypatch.setenv("VLLM_OMNI_AR_DIFFUSION_KV_GATHER", "1")
+    # Same geometry as the capacity-two test: 192 bytes fits two sessions exactly without staging.
+    plain = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
+    assert plain.session_capacity == 2 and plain.history_staging_reserved_bytes == 0
+
+    staged = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192, reuse_history_staging=True)
+    # 2 (K, V) x 1 layer x (sink 3 + window 3 + one action-capacity block) tokens x 1 head x 1 dim x 4 bytes.
+    assert staged.history_staging_reserved_bytes == 2 * 1 * 7 * 1 * 1 * 4
+    # It comes out of the same budget, so the second session no longer fits.
+    assert staged.session_capacity == 1
+    assert staged.num_blocks_total * 8 + staged.history_staging_reserved_bytes <= 192
+
+    # Boundary: one session fits at 128 bytes without staging; with staging the same budget is rejected
+    # up front instead of failing on the first forward.
+    _make_tiny_capacity_kv(requested_capacity=1, available_bytes=128)
+    with pytest.raises(ValueError, match="cannot fit one session"):
+        _make_tiny_capacity_kv(requested_capacity=1, available_bytes=128, reuse_history_staging=True)
+
+
+def test_history_staging_is_not_reserved_when_the_gather_path_is_off(monkeypatch):
+    monkeypatch.delenv("VLLM_OMNI_AR_DIFFUSION_KV_GATHER", raising=False)
+    staged = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192, reuse_history_staging=True)
+    assert staged.history_staging_reserved_bytes == 0 and staged.session_capacity == 2
+
+
 def test_capacity_two_retains_both_windows_and_allocates_next_block():
     # LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
     # 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
@@ -588,3 +623,70 @@ def test_non_contiguous_branch_indices_rejected():
             kv_branches=(ARDiffusionKVBranchSpec("main", 1),),
             session_capacity=1,
         )
+
+
+def test_staged_window_reuse_matches_a_full_gather():
+    """Reusing the staged history must be byte-identical to re-gathering the whole window.
+
+    This is the property the optimisation rests on: the leading blocks were staged by an earlier forward of
+    the same AR block, so skipping them can only be correct if what is already there equals what a fresh
+    gather would write. The test stages once, mutates only the current chunk's blocks in the pool (what a
+    later probe of the same block does), restages the tail alone, and compares against a full gather.
+    """
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import _stage_window
+
+    torch.manual_seed(0)
+    num_blocks, block_size, heads, dim = 6, 4, 2, 8
+    current_blocks = 2
+    key_cache = torch.randn(num_blocks, block_size, heads, dim)
+    value_cache = torch.randn(num_blocks, block_size, heads, dim)
+    block_ids = torch.tensor([5, 4, 3, 2, 1, 0])
+    stage_key = torch.zeros(num_blocks * block_size, heads, dim)
+    stage_value = torch.zeros(num_blocks * block_size, heads, dim)
+
+    _stage_window(stage_key, stage_value, key_cache, value_cache, block_ids, num_blocks, block_size, first_block=0)
+
+    # A later probe rewrites only the current chunk's blocks in the pool.
+    for slot in block_ids[num_blocks - current_blocks :]:
+        key_cache[slot] = torch.randn(block_size, heads, dim)
+        value_cache[slot] = torch.randn(block_size, heads, dim)
+
+    _stage_window(
+        stage_key,
+        stage_value,
+        key_cache,
+        value_cache,
+        block_ids,
+        num_blocks,
+        block_size,
+        first_block=num_blocks - current_blocks,
+    )
+
+    full_key = torch.zeros_like(stage_key)
+    full_value = torch.zeros_like(stage_value)
+    _stage_window(full_key, full_value, key_cache, value_cache, block_ids, num_blocks, block_size, first_block=0)
+
+    torch.testing.assert_close(stage_key, full_key, rtol=0, atol=0)
+    torch.testing.assert_close(stage_value, full_value, rtol=0, atol=0)
+
+
+def test_staged_window_tail_refresh_leaves_history_untouched():
+    """The tail restage must not write the history rows -- that is what makes it cheaper."""
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import _stage_window
+
+    torch.manual_seed(1)
+    num_blocks, block_size, heads, dim = 5, 4, 1, 4
+    key_cache = torch.randn(num_blocks, block_size, heads, dim)
+    value_cache = torch.randn(num_blocks, block_size, heads, dim)
+    block_ids = torch.arange(num_blocks)
+    stage_key = torch.full((num_blocks * block_size, heads, dim), -1.0)
+    stage_value = torch.full((num_blocks * block_size, heads, dim), -1.0)
+
+    _stage_window(
+        stage_key, stage_value, key_cache, value_cache, block_ids, num_blocks, block_size, first_block=num_blocks - 1
+    )
+
+    history_rows = (num_blocks - 1) * block_size
+    assert (stage_key[:history_rows] == -1.0).all(), "history rows were rewritten by a tail restage"
+    assert (stage_value[:history_rows] == -1.0).all()
+    torch.testing.assert_close(stage_key[history_rows:], key_cache[-1].reshape(block_size, heads, dim), rtol=0, atol=0)

tests/diffusion/ar_diffusion/test_kv_orchestration.py

@@ -6,6 +6,7 @@ import torch
 
 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
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 DIMS = dict(
@@ -108,3 +109,104 @@ def test_state_close_frees_both_branch_blocks():
     state.close()
     # All blocks returned — no leak across an evicted session.
     assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+@pytest.mark.parametrize(
+    "frames_per_call,block_size,max_model_len,num_calls",
+    [(1, 16, 113, 1000), (3, 16, 113, 1000), (3, 1560, 1 << 20, 241)],
+    ids=["single-frame", "three-frame", "lingbot-two-minutes"],
+)
+@pytest.mark.parametrize("sink_chunks,reset_at_boundary", [(0, False), (2, False), (0, True), (2, True)])
+def test_rollout_keeps_metadata_bounded(
+    frames_per_call, block_size, max_model_len, num_calls, sink_chunks, reset_at_boundary
+):
+    """Long rollouts retain exact global-position K/V with bounded storage indices."""
+    kv = ARDiffusionKVCache(
+        ARDiffusionKVConfig(
+            enable=True,
+            chunk_size=block_size,
+            window_chunks=3,
+            sink_chunks=sink_chunks,
+            reset_at_boundary=reset_at_boundary,
+        ),
+        num_layers=1,
+        num_kv_heads=1,
+        head_size=4,
+        dtype=torch.float32,
+        block_size=block_size,
+        max_model_len=max_model_len,
+        available_bytes=1 << 24,
+        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+        session_capacity=1,
+        frames_per_block=frames_per_call,
+        device=torch.device("cpu"),
+    )
+    adapter = kv.begin_request("long-rollout")
+    state = ARDiffusionKVState(kv, "long-rollout", {"main": adapter}, num_layers=1)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+    pool_shape = kv.key_cache(0).shape
+    retained_frames = sink_chunks + (0 if reset_at_boundary else 3)
+    for call_index in range(num_calls):
+        start = call_index * frames_per_call * block_size
+        stop = start + frames_per_call * block_size
+        if frames_per_call == 1:
+            kv.allocate_chunk(adapter)
+            slots = kv.chunk_write_slots(adapter)
+        else:
+            context = state.get_kv_caches("main", seq_len=stop - start, commit_current=True)[0].forward_ctx
+            context.ensure_video_slots(torch.device("cpu"))
+            slots = context.current_video_slot_mapping
+        assert len(kv.block_table(adapter)) <= retained_frames + frames_per_call
+        values = torch.arange(start, stop, dtype=torch.float32).view(-1, 1, 1).expand(-1, 1, 4)
+        kv.key_cache(0).flatten(0, 1)[slots] = values
+        kv.value_cache(0).flatten(0, 1)[slots] = -values
+        if frames_per_call == 1:
+            kv.commit_chunk(adapter)
+        else:
+            state.commit_paged_context("main")
+            assert state._committed["main"] == stop
+        sink_end = min(stop, sink_chunks * block_size)
+        tail_start = stop if reset_at_boundary else max(sink_end, stop - 3 * block_size)
+        expected = torch.cat((torch.arange(sink_end), torch.arange(tail_start, stop))).float()
+        expected = expected.view(-1, 1, 1).expand(-1, 1, 4)
+        resident = kv.window_block_ids(adapter)
+        torch.testing.assert_close(kv.key_cache(0)[resident].flatten(0, 1), expected)
+        torch.testing.assert_close(kv.value_cache(0)[resident].flatten(0, 1), -expected)
+        assert adapter.completed_chunks * block_size == stop
+        assert adapter.absolute_num_computed_tokens == stop
+        assert len(kv.block_table(adapter)) <= retained_frames
+        assert adapter.num_computed_tokens <= retained_frames * block_size
+        assert kv.key_cache(0).shape == pool_shape
+        assert free_total - kv.manager.block_pool.get_num_free_blocks() == min(stop // block_size, retained_frames)
+    state.reset()
+    fresh = state.adapter("main")
+    assert fresh.compacted_tokens == fresh.absolute_num_computed_tokens == fresh.num_computed_tokens == 0
+    assert kv.block_table(fresh) == []
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+    kv.allocate_chunk(fresh)
+    assert len(kv.chunk_write_slots(fresh)) == block_size
+    state.close()
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_compaction_offsets_are_request_local():
+    kv = make_cache()
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+    first, second = kv.begin_request("first"), kv.begin_request("second")
+    for index in range(8):
+        kv.allocate_chunk(first)
+        kv.commit_chunk(first)
+        if index in (0, 4):
+            kv.allocate_chunk(second)
+            kv.commit_chunk(second)
+    assert first.absolute_num_computed_tokens == 8 * 16
+    assert first.compacted_tokens == 6 * 16
+    assert second.absolute_num_computed_tokens == 2 * 16
+    assert second.compacted_tokens == 0
+    assert first.num_computed_tokens == second.num_computed_tokens == 2 * 16
+    assert set(kv.block_table(first)).isdisjoint(kv.block_table(second))
+    kv.end_request(first)
+    kv.end_request(second)
+    group = kv.manager.coordinator.single_type_managers[0]
+    assert not group.req_to_blocks and not group.num_cached_block
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for AR-Diffusion paged self-attention contexts."""
 
 from __future__ import annotations
@@ -20,6 +21,8 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
     ar_diffusion_paged_attention,
     paged_write_attn,
 )
+from vllm_omni.experimental.ar_diffusion.kv_cache import paged_attention as paged_attention_module
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -32,8 +35,12 @@ POS = "positive"
 NEG = "negative"
 
 
-def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
-    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+def make_state(
+    *, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False
+):
+    cfg = ARDiffusionKVConfig(
+        enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
+    )
     kv = ARDiffusionKVCache(
         cfg,
         num_layers=num_layers,
@@ -236,6 +243,86 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
     assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
 
 
+@pytest.mark.parametrize("history_chunks", [0, 1, 2, 3])
+def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks):
+    """A second probe of the same AR block restages its current K/V where the table actually holds it.
+
+    The padded table always ends in at least one action-capacity block, and while the window is still
+    growing in unused window capacity too, so "the last current_blocks entries" is padding. The refresh
+    has to start right after the visible history: empty (window growing), partial, full, and after a slide.
+    """
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device = torch.device("cpu")
+    dtype = torch.float32
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, reuse_history_staging=True)
+    if history_chunks:
+        _commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
+
+    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    key_cache, value_cache = kv.key_cache(0), kv.value_cache(0)
+
+    def probe(step: int):
+        # Each probe of the block writes new current K/V, then stages.
+        kv._k_pools[0][ctx.current_video_slot_mapping] = torch.full((BLOCK, N_HEADS, HEAD_DIM), float(step))
+        kv._v_pools[0][ctx.current_video_slot_mapping] = torch.full((BLOCK, N_HEADS, HEAD_DIM), -float(step))
+        # Store the metadata on the context the way prepare() does.
+        (ctx.block_table, ctx.query_start_loc, ctx.seq_lens, ctx.max_query_len, ctx.max_seq_len) = (
+            ctx.build_block_table(action_len=0, query_len=BLOCK, device=device)
+        )
+        block_table, max_seq_len = ctx.block_table, ctx.max_seq_len
+        ctx._prepare_history_staging(0)
+        n_blocks = max_seq_len // BLOCK
+        block_ids = block_table[0, :n_blocks].to(torch.long)
+        stage_k, stage_v = ctx.history_staging(0)
+        paged_attention_module._stage_window(
+            stage_k,
+            stage_v,
+            key_cache,
+            value_cache,
+            block_ids,
+            n_blocks,
+            BLOCK,
+            first_block=ctx.stage_first_block if ctx.reuse_history else 0,
+        )
+        full_k = key_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
+        full_v = value_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
+        return stage_k, stage_v, full_k, full_v, n_blocks
+
+    stage_k, stage_v, full_k, full_v, n_blocks = probe(1)
+    assert ctx.reuse_history is False
+    assert torch.equal(stage_k, full_k) and torch.equal(stage_v, full_v)
+
+    stage_k, stage_v, full_k, full_v, n_blocks = probe(2)
+    assert ctx.reuse_history is True
+    visible_history_blocks = min(history_chunks, 2 - 1)  # window of 2 blocks minus the current one
+    assert ctx.stage_first_block == visible_history_blocks
+    # The padded table is wider than the live window, so the end-of-table guess is not the offset.
+    assert n_blocks - 1 != ctx.stage_first_block
+    assert torch.equal(stage_k, full_k) and torch.equal(stage_v, full_v)
+
+
+@pytest.mark.parametrize("gather_enabled", [False, True])
+def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabled):
+    if gather_enabled:
+        monkeypatch.setenv(KV_GATHER_ENV, "1")
+    else:
+        monkeypatch.delenv(KV_GATHER_ENV, raising=False)
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, reuse_history_staging=True)
+    # The manager owns the pairs and allocates them only when their consumer is on.
+    assert bool(kv.history_staging) is gather_enabled
+    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    (ctx.block_table, ctx.query_start_loc, ctx.seq_lens, ctx.max_query_len, ctx.max_seq_len) = ctx.build_block_table(
+        action_len=0, query_len=BLOCK, device=device
+    )
+    ctx._prepare_history_staging(0)
+    assert ctx.staging_enabled is gather_enabled
+    stage_k, _ = ctx.history_staging(0)
+    assert (stage_k is not None) is gather_enabled
+
+
 @pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
 @pytest.mark.parametrize("history_chunks", [1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
@@ -376,13 +463,27 @@ def test_custom_op_registration_idempotent():
     assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
 
 
-def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
+def test_custom_op_mutable_arguments_cannot_be_elided_as_defaults():
+    # Older PyTorch ADInplaceOrView handlers index positional mutable inputs
+    # directly. Default-valued trailing inputs can disappear before that handler.
+    schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema
+    mutable = [arg for arg in schema.arguments if arg.alias_info is not None and arg.alias_info.is_write]
+    assert {arg.name for arg in mutable} == {"key_pool", "value_pool", "stage_key", "stage_value"}
+    assert all(not arg.has_default_value() for arg in mutable)
+
+
+@pytest.mark.parametrize("reuse_history_staging", [False, True])
+def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypatch, reuse_history_staging):
     """The op must trace as one opaque node: fullgraph OK, and changed tensor
     VALUES (new slots / block ids) must not trigger recompilation."""
     import torch._dynamo
 
     device = torch.device("cpu")
-    kv, st = make_state(num_layers=2, window_chunks=2)
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    kv, st = make_state(num_layers=2, window_chunks=2, reuse_history_staging=reuse_history_staging)
+    if reuse_history_staging:
+        # Hold the host-side staging offset fixed while table/slot values change.
+        _commit_video_span(kv, st, kv_branch=POS, n_chunks=2, dtype=torch.float32, device=device)
 
     def run_one_forward(commit):
         contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit)
@@ -391,9 +492,8 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
         q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
         k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
         v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
-        # Both layers through ONE compiled fn: layer_idx is a tensor, so a
-        # different layer must NOT recompile (all 40 DiT blocks share the
-        # block-forward code object in production).
+        # Both layers use one compiled function; the tensor-valued layer index
+        # must not specialize it.
         for layer_ctx in contexts:
             out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
         st.commit_paged_context(POS)
@@ -419,3 +519,52 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
         # Leave a clean dynamo state for later suites in the same pytest process
         # (e.g. model_executor transformers models).
         torch._dynamo.reset()
+
+
+@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@pytest.mark.parametrize("history_chunks", [0, 1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_chunks, action_len):
+    """VLLM_OMNI_AR_DIFFUSION_KV_GATHER=1 gathers the visible blocks and runs varlen FA3 without a block table.
+
+    Covers an empty, partial and full window plus a partially filled action block:
+    the gather reads the tail-padding null block, so its (zeroed) rows must be masked.
+    """
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+    if history_chunks:
+        _commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)
+
+    layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    action_k = action_v = None
+    if action_len:
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
+    inputs = layer_ctx.to_layer_inputs()
+
+    def run() -> torch.Tensor:
+        return paged_write_attn(
+            inputs,
+            query[0],
+            current_k[0],
+            current_v[0],
+            action_k[0] if action_len else None,
+            action_v[0] if action_len else None,
+            HEAD_DIM**-0.5,
+        ).unsqueeze(0)
+
+    monkeypatch.delenv(KV_GATHER_ENV, raising=False)
+    paged = run()
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    gathered = run()
+    assert torch.isfinite(gathered).all()
+    # Same kernel family on the same K/V: only accumulation order differs.
+    torch.testing.assert_close(gathered, paged, rtol=2e-3, atol=2e-3)

tests/diffusion/ar_diffusion/test_streaming_decode.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Streaming VAE decode against the real Wan autoencoder, on CPU.
+
+The autoencoder is built from a config rather than a checkpoint -- weights are
+random, only shapes and the temporal cache protocol matter -- so these run
+without a device or a download.
+
+The load-bearing test is equivalence: decoding a session chunk by chunk must
+produce exactly what decoding the whole clip produces. Everything else about
+streaming is worthless if that does not hold.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.streaming_decode import (
+    StreamingDecodeState,
+    SupportsStreamingDecode,
+    WanStreamingDecoder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Small enough to decode on CPU in a test, same architecture as the shipped
+# checkpoint: base_dim 96 / dim_mult [1,2,4,4] would be far too slow, so the
+# widths are scaled down while the causal structure is unchanged.
+VAE_CONFIG = {
+    "base_dim": 8,
+    "z_dim": 4,
+    "dim_mult": [1, 2],
+    "num_res_blocks": 1,
+    "temperal_downsample": [True],
+    "attn_scales": [],
+    "dropout": 0.0,
+}
+LATENT_H = LATENT_W = 4
+
+
+@pytest.fixture(scope="module")
+def vae():
+    diffusers = pytest.importorskip("diffusers")
+    torch.manual_seed(0)
+    model = diffusers.AutoencoderKLWan(**VAE_CONFIG).eval()
+    for parameter in model.parameters():
+        parameter.requires_grad_(False)
+    return model
+
+
+@pytest.fixture(scope="module")
+def decoder(vae):
+    return WanStreamingDecoder(vae)
+
+
+def _latent(num_frames: int, *, seed: int = 0) -> torch.Tensor:
+    generator = torch.Generator().manual_seed(seed)
+    return torch.randn(1, VAE_CONFIG["z_dim"], num_frames, LATENT_H, LATENT_W, generator=generator)
+
+
+@torch.no_grad()
+def _decode_whole(vae, latent: torch.Tensor) -> torch.Tensor:
+    """The non-streaming reference: one call, cache cleared at both ends."""
+    return vae._decode(latent, return_dict=False)[0]
+
+
+# --------------------------------------------------------------------------
+# Equivalence: the property that makes streaming worth doing at all
+# --------------------------------------------------------------------------
+
+
+def _stream(decoder, latent: torch.Tensor, chunk_sizes, *, session_id: str = "s") -> torch.Tensor:
+    state = decoder.new_decode_state(session_id)
+    pieces, offset = [], 0
+    for size in chunk_sizes:
+        pieces.append(decoder.decode_chunk(latent[:, :, offset : offset + size], state))
+        offset += size
+    return torch.cat(pieces, dim=2)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("chunk_sizes", [(3, 3, 3), (1, 1, 1, 1, 1), (2, 4, 3)])
+def test_chunking_is_neutral(vae, decoder, chunk_sizes) -> None:
+    """Splitting a session into chunks must change nothing about its output.
+
+    This is the property streaming actually introduces, so it is the one
+    asserted exactly. The control is the same session decoded in a single
+    call through the same code path -- not ``_decode`` -- because ``_decode``
+    differs in a way that has nothing to do with chunking; see
+    ``test_distance_from_the_one_shot_path_is_a_post_quant_conv_artifact``.
+    """
+    total = sum(chunk_sizes)
+    latent = _latent(total)
+    reference = _stream(decoder, latent, (total,), session_id="ref")
+    streamed = _stream(decoder, latent, chunk_sizes)
+
+    assert streamed.shape == reference.shape
+    torch.testing.assert_close(streamed, reference, rtol=0, atol=0)
+
+
+@torch.no_grad()
+def test_distance_from_the_one_shot_path_is_a_post_quant_conv_artifact(vae, decoder) -> None:
+    """Explain, and bound, the gap against diffusers' ``_decode``.
+
+    ``_decode`` applies ``post_quant_conv`` to the whole latent at once; the
+    tiled path in this repo, and this decoder, apply it per frame. It is a
+    1x1x1 convolution, so the two are mathematically identical -- but they can
+    select different kernels, and on bf16 CUDA that difference is real and is
+    amplified by the causal convolutions downstream.
+
+    Asserting equality with ``_decode`` would therefore pin a kernel-selection
+    coincidence rather than a property of streaming. Assert instead that the
+    gap is entirely explained by where ``post_quant_conv`` runs.
+    """
+    latent = _latent(6)
+    per_frame_pq = torch.cat([vae.post_quant_conv(latent[:, :, k : k + 1]) for k in range(latent.shape[2])], dim=2)
+    whole_pq = vae.post_quant_conv(latent)
+
+    streamed = _stream(decoder, latent, (3, 3))
+    one_shot = _decode_whole(vae, latent)
+
+    # Same inputs to the decoder => same output, whatever the backend does.
+    if torch.equal(whole_pq, per_frame_pq):
+        torch.testing.assert_close(streamed, one_shot, rtol=0, atol=0)
+    else:  # pragma: no cover - backend dependent
+        pytest.skip("backend batches post_quant_conv differently; the gap is that, not chunking")
+
+
+@torch.no_grad()
+def test_a_restarted_stream_does_not_match_a_continued_one(decoder) -> None:
+    """Negative control: the equivalence test above must be discriminating.
+
+    Clearing state between chunks -- what the non-streaming path does on every
+    call -- has to change the pixels, otherwise the temporal cache carries
+    nothing and the assertion above holds for a write-only cache too.
+
+    Restarting also changes the frame count, because a fresh state expands its
+    opening latent frame to one raw frame. Comparing totals would therefore
+    pass on the frame count alone and never look at a pixel, so this compares
+    the tail both runs agree on: the frames from the chunk's later latents,
+    which are the full temporal factor wide on either side.
+    """
+    latent = _latent(6)
+    state = decoder.new_decode_state("continued")
+    decoder.decode_chunk(latent[:, :, :3], state)
+    continued = decoder.decode_chunk(latent[:, :, 3:], state)
+    restarted = decoder.decode_chunk(latent[:, :, 3:], decoder.new_decode_state("restarted"))
+
+    tail = min(continued.shape[2], restarted.shape[2]) - 1
+    assert tail > 0, "the comparison needs at least one fully expanded frame on both sides"
+    continued_tail, restarted_tail = continued[:, :, -tail:], restarted[:, :, -tail:]
+    assert continued_tail.shape == restarted_tail.shape, "the tails must be comparable for this to mean anything"
+    assert not torch.equal(continued_tail, restarted_tail), (
+        "restarting the decoder produced the same pixels as continuing it, so the carried cache is not being read"
+    )
+
+
+# --------------------------------------------------------------------------
+# Cross-session isolation: the failure that timing metrics cannot see
+# --------------------------------------------------------------------------
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("seed_a, seed_b", [(1, 2), (11, 12)])
+def test_interleaved_sessions_keep_independent_temporal_context(vae, decoder, seed_a: int, seed_b: int) -> None:
+    """Two sessions ticking alternately must each match their solo decode.
+
+    A session that reads another's temporal cache produces a perfectly
+    plausible but wrong video, and every latency metric stays green, so this
+    has to be asserted against a recorded solo run rather than inferred. The
+    control is a solo stream rather than ``_decode``, so it differs from the
+    run under test only in the interleaving.
+    """
+    latent_a = _latent(6, seed=seed_a)
+    latent_b = _latent(6, seed=seed_b)
+    solo_a = _stream(decoder, latent_a, (6,), session_id="solo-a")
+    solo_b = _stream(decoder, latent_b, (6,), session_id="solo-b")
+    assert not torch.equal(solo_a, solo_b), "negative control: the two sessions must differ"
+
+    state_a = decoder.new_decode_state("a")
+    state_b = decoder.new_decode_state("b")
+    out_a, out_b = [], []
+    for start in (0, 3):
+        window = slice(start, start + 3)
+        out_a.append(decoder.decode_chunk(latent_a[:, :, window], state_a))
+        out_b.append(decoder.decode_chunk(latent_b[:, :, window], state_b))
+
+    torch.testing.assert_close(torch.cat(out_a, dim=2), solo_a, rtol=0, atol=0)
+    torch.testing.assert_close(torch.cat(out_b, dim=2), solo_b, rtol=0, atol=0)
+
+
+@torch.no_grad()
+def test_the_module_cache_is_untouched_so_sessions_can_share_one_decoder(vae, decoder) -> None:
+    vae.clear_cache()
+    before = [entry for entry in vae._feat_map]
+    decoder.decode_chunk(_latent(3), decoder.new_decode_state("s"))
+    assert vae._feat_map == before, "streaming decode must not write the module-owned cache"
+
+
+# --------------------------------------------------------------------------
+# Causal geometry and boundedness
+# --------------------------------------------------------------------------
+
+
+@torch.no_grad()
+def test_opening_chunk_is_shorter_and_later_chunks_are_full_length(vae, decoder) -> None:
+    """(n - 1) * factor + 1 for the opening chunk, n * factor after it."""
+    factor = 2 ** sum(VAE_CONFIG["temperal_downsample"])
+    state = decoder.new_decode_state("s")
+    first = decoder.decode_chunk(_latent(3), state)
+    second = decoder.decode_chunk(_latent(3, seed=5), state)
+
+    assert first.shape[2] == (3 - 1) * factor + 1
+    assert second.shape[2] == 3 * factor
+    assert state.chunks_decoded == 2
+    assert state.frames_decoded == 6
+
+
+@torch.no_grad()
+def test_resident_state_does_not_grow_with_session_length(vae, decoder) -> None:
+    """Each causal convolution retains at most CACHE_T temporal slices."""
+    state = decoder.new_decode_state("s")
+    decoder.decode_chunk(_latent(3), state)
+    baseline = state.nbytes()
+    assert baseline > 0
+
+    for step in range(8):
+        decoder.decode_chunk(_latent(3, seed=step + 10), state)
+        assert state.nbytes() == baseline
+
+    assert sum(state.nbytes_by_device().values()) == baseline
+
+
+@torch.no_grad()
+def test_release_returns_the_session_to_its_pre_stream_state(vae, decoder) -> None:
+    state = decoder.new_decode_state("s")
+    decoder.decode_chunk(_latent(3), state)
+    assert state.started and state.nbytes() > 0
+
+    decoder.release(state)
+    assert state.nbytes() == 0
+    assert not state.started
+    assert state.chunks_decoded == 0
+
+    # A released session restarts rather than resumes, so its first chunk is
+    # short again -- the same semantics reset() has for AR KV.
+    factor = 2 ** sum(VAE_CONFIG["temperal_downsample"])
+    assert decoder.decode_chunk(_latent(3), state).shape[2] == (3 - 1) * factor + 1
+
+
+# --------------------------------------------------------------------------
+# Contract
+# --------------------------------------------------------------------------
+
+
+def test_the_vaes_execution_context_is_entered_for_every_frame(vae) -> None:
+    """Decode must run inside the autoencoder's own execution context.
+
+    Every public compute entry point on ``OmniAutoencoderKLWan`` enters
+    ``_execution_context()``, which is where autocast is established for
+    fp16/bf16 parameters -- and on NPU the only place it is established at all.
+    Driving ``post_quant_conv`` and ``decoder`` directly bypasses it unless this
+    decoder enters it itself, which would compute streamed chunks under a
+    different numeric context from the whole-clip decode they are compared to.
+    """
+    entered: list[str] = []
+
+    class _ContextualVAE:
+        """The parts of the contract this decoder uses, plus the hook."""
+
+        config = vae.config
+        dtype = torch.float32
+
+        def __init__(self) -> None:
+            self.decoder = vae.decoder
+            self.post_quant_conv = vae.post_quant_conv
+            self._cached_conv_counts = vae._cached_conv_counts
+
+        @contextmanager
+        def _execution_context(self):
+            entered.append("enter")
+            try:
+                yield
+            finally:
+                entered.append("exit")
+
+    decoder = WanStreamingDecoder(_ContextualVAE())
+    with torch.no_grad():
+        decoder.decode_chunk(_latent(3), decoder.new_decode_state("s"))
+
+    # One context for the call, not one per frame, and it closes.
+    assert entered == ["enter", "exit"]
+
+
+def test_a_plain_diffusers_autoencoder_needs_no_execution_context(decoder) -> None:
+    """The hook is optional: a VAE without one still decodes."""
+    with torch.no_grad():
+        out = decoder.decode_chunk(_latent(1), decoder.new_decode_state("s"))
+    assert out.shape[2] == 1
+
+
+def test_wan_decoder_satisfies_the_protocol(decoder) -> None:
+    assert isinstance(decoder, SupportsStreamingDecode)
+
+
+def test_state_from_another_decoder_is_rejected(decoder) -> None:
+    foreign = StreamingDecodeState(session_id="s", feat_map=[None])
+    with pytest.raises(ValueError, match="does not belong to this decoder"):
+        decoder.decode_chunk(_latent(1), foreign)
+
+
+def test_a_non_wan_module_is_rejected() -> None:
+    with pytest.raises(TypeError, match="Wan-family autoencoder"):
+        WanStreamingDecoder(object())
+
+
+def test_empty_and_malformed_latents_are_rejected(decoder) -> None:
+    state = decoder.new_decode_state("s")
+    with pytest.raises(ValueError, match=r"\[B, C, T, H, W\]"):
+        decoder.decode_chunk(torch.zeros(1, 4, 4, 4), state)
+    with pytest.raises(ValueError, match="at least one frame"):
+        decoder.decode_chunk(torch.zeros(1, 4, 0, LATENT_H, LATENT_W), state)
+
+
+def test_a_batch_is_refused_rather_than_threaded_through_one_cache(decoder) -> None:
+    """One state is one session, so a batch has no correct interpretation here.
+
+    ``AutoencoderKLWan.decode`` handles a batch by decoding each sample with
+    its own cleared cache (``use_slicing``). This decoder carries exactly one
+    temporal context, so a batch would mix samples into it and still return
+    correctly shaped pixels.
+    """
+    state = decoder.new_decode_state("s")
+    batched = torch.zeros(2, VAE_CONFIG["z_dim"], 3, LATENT_H, LATENT_W)
+    with pytest.raises(ValueError, match="one sample per session state"):
+        decoder.decode_chunk(batched, state)
+    assert state.frames_decoded == 0, "a refused call must not advance the session"
+
+
+def test_session_id_must_be_meaningful(decoder) -> None:
+    with pytest.raises(ValueError, match="session_id"):
+        decoder.new_decode_state("  ")
+
+
+def test_declared_state_bytes_scales_with_area_and_dtype(vae) -> None:
+    # A local decoder, not the module-scoped fixture: the measured constant
+    # belongs to this test, and leaking it would change what a reordered run
+    # of the default-None case sees. Passing it here is also the constructor
+    # argument's intended use.
+    # Measured for the shipped checkpoint: 37832 KiB of fp32 cache at 64x64.
+    decoder = WanStreamingDecoder(vae, bytes_per_pixel_fp32=37832 * 1024 / (64 * 64))
+    fp32 = decoder.declared_state_bytes(height=480, width=832, dtype=torch.float32)
+    bf16 = decoder.declared_state_bytes(height=480, width=832, dtype=torch.bfloat16)
+    half_area = decoder.declared_state_bytes(height=240, width=832, dtype=torch.float32)
+    assert bf16 == fp32 // 2
+    assert half_area == fp32 // 2
+    assert fp32 / 2**20 == pytest.approx(3602.0, abs=5.0)

tests/diffusion/attention/test_attention_config.py

@@ -10,6 +10,7 @@ Tests cover:
 - AttentionMetadata.extra field
 """
 
+from dataclasses import replace
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -17,7 +18,12 @@ import pytest
 import torch
 
 import vllm_omni.diffusion.attention.layer as layer_mod
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+import vllm_omni.diffusion.attention.selector as selector_mod
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+    QueryRange,
+)
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.config import (
@@ -477,6 +483,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
     @pytest.mark.parametrize("mindiesd_available", [False, True])
     def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available):
         attention = Attention.__new__(Attention)
+        attention.allow_fp32_fallback = False
         attention._scheduler_paged_kv = True
         attention.paged_kv_cache_role = "primary"
         attention.backend_pref = "FLASH_ATTN"
@@ -519,6 +526,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
     @pytest.mark.parametrize("fa4_available", [False, True])
     def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch, fa4_available):
         attention = Attention.__new__(Attention)
+        attention.allow_fp32_fallback = False
         attention._scheduler_paged_kv = True
         attention.paged_kv_cache_role = "primary"
         attention.backend_pref = "FLASH_ATTN"
@@ -561,6 +569,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
 
     def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
         attention = Attention.__new__(Attention)
+        attention.allow_fp32_fallback = False
         attention._scheduler_paged_kv = False
         attention.paged_kv_cache_role = "primary"
         attention.backend_pref = "FLASH_ATTN"
@@ -900,6 +909,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
             return sentinel
 
         fake_attention = SimpleNamespace(
+            allow_fp32_fallback=False,
             attention=SimpleNamespace(
                 forward=_selected_forward,
             ),
@@ -938,6 +948,164 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
             Attention._run_ring_attention(fake_attention, query, query, query, metadata)
 
 
+class TestOptInFloat32Fallback:
+    @staticmethod
+    def _make_attention(monkeypatch, *, backend=None, allow=True, kv_cache_dtype=None):
+        events, state = [], {}
+        context = object()
+
+        class _SelectedImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+            @staticmethod
+            def supports_kv_cache_dtype(*args):
+                return True
+
+            def forward(self, query, key, value, metadata):
+                events.append("selected")
+                state["kernel_call"] = (query, key, value, metadata)
+                state["kernel_output"] = query + 4
+                return state["kernel_output"]
+
+        class _FallbackImpl(_SelectedImpl):
+            def forward(self, query, key, value, metadata):
+                output = super().forward(query, key, value, metadata)
+                events[-1] = "fallback"
+                return output
+
+        class _FlashBackend(FlashAttentionBackend):
+            @staticmethod
+            def get_impl_cls():
+                return _SelectedImpl
+
+        class _SDPABackend(layer_mod.SDPABackend):
+            @staticmethod
+            def get_impl_cls():
+                return _SelectedImpl
+
+        class _Strategy:
+            def pre_attention(self, query, key, value, metadata):
+                events.append("pre")
+                state["prepared"] = (
+                    query + 1,
+                    key + 2,
+                    value + 3,
+                    replace(metadata) if metadata is not None else None,
+                )
+                return (*state["prepared"], context)
+
+            def post_attention(self, output, ctx):
+                events.append("post")
+                assert output is state["kernel_output"]
+                assert ctx is context
+                return output + 7
+
+        # Keep real role/config resolution and construction; replace hardware
+        # lookup and kernels. These are CPU dispatch tests, not SP collectives.
+        monkeypatch.setattr(
+            selector_mod,
+            "_cached_get_backend_cls",
+            lambda name, *args, **kwargs: _SDPABackend if name == "TORCH_SDPA" else _FlashBackend,
+        )
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FallbackImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: _Strategy())
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+        config = OmniDiffusionConfig(
+            diffusion_attention_config={"default": backend} if backend else {},
+            diffusion_kv_cache_dtype=kv_cache_dtype,
+        )
+        with set_current_diffusion_config(config):
+            attention = Attention(
+                num_heads=4,
+                num_kv_heads=2,
+                head_size=8,
+                causal=False,
+                softmax_scale=0.25,
+                **({"allow_fp32_fallback": True} if allow else {}),
+            )
+        return attention, events, state
+
+    @pytest.mark.parametrize(
+        "backend,allow,dtype,is_cuda,expected",
+        [
+            pytest.param(None, True, torch.float32, True, "fallback", id="auto-fp32"),
+            pytest.param(None, False, torch.float32, True, "selected", id="default-opt-out"),
+            pytest.param("FLASH_ATTN", True, torch.float32, True, "selected", id="explicit-flash"),
+            pytest.param("TORCH_SDPA", True, torch.float32, True, "selected", id="explicit-sdpa"),
+            pytest.param(None, True, torch.bfloat16, True, "selected", id="bf16"),
+            pytest.param(None, True, torch.float32, False, "selected", id="cpu-fp32"),
+        ],
+    )
+    def test_public_forward_preserves_strategy_and_backend_choice(
+        self, monkeypatch, backend, allow, dtype, is_cuda, expected
+    ):
+        attention, events, state = self._make_attention(monkeypatch, backend=backend, allow=allow)
+        # Only the dispatch predicate is simulated; every tensor stays on CPU.
+        monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: is_cuda))
+        query = torch.ones(1, 2, 4, 8, dtype=dtype)
+        key = torch.full((1, 2, 2, 8), 2, dtype=dtype)
+        value = torch.full_like(key, 3)
+        metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False]]))
+
+        output = attention(query, key, value, metadata)
+
+        assert events == ["pre", expected, "post"]
+        assert all(got is prepared for got, prepared in zip(state["kernel_call"], state["prepared"]))
+        assert state["kernel_call"][3] is not metadata
+        assert state["kernel_call"][3].attn_mask is metadata.attn_mask
+        torch.testing.assert_close(output, query + 12)
+        assert attention.backend_explicit is (backend is not None)
+        impl = attention.sdpa_fallback if expected == "fallback" else attention.attention
+        assert impl.kwargs["num_kv_heads"] == 2
+        assert impl.kwargs["softmax_scale"] == 0.25
+
+    def test_explicit_backend_failure_is_not_retried_with_sdpa(self, monkeypatch):
+        attention, events, _ = self._make_attention(monkeypatch, backend="FLASH_ATTN")
+        monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
+        selected = Mock(side_effect=RuntimeError("selected backend rejects this dtype"))
+        monkeypatch.setattr(attention.attention, "forward", selected)
+        query = torch.ones(1, 2, 4, 8)
+
+        with pytest.raises(RuntimeError, match="selected backend rejects this dtype"):
+            attention(query, query, query)
+
+        selected.assert_called_once()
+        assert events == ["pre"]
+
+    @pytest.mark.parametrize("contract", ["packed", "packed-padding", "piecewise", "quantized-kv"])
+    def test_backend_metadata_is_not_silently_sent_to_dense_sdpa(self, monkeypatch, contract):
+        quant_dtype = "fp8" if contract == "quantized-kv" else None
+        attention, events, state = self._make_attention(monkeypatch, kv_cache_dtype=quant_dtype)
+        monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
+        lengths = torch.tensor([0, 1, 2], dtype=torch.int32)
+        metadata = AttentionMetadata()
+        if contract == "packed":
+            metadata.extra = {
+                "cu_seqlens_q": lengths,
+                "cu_seqlens_k": lengths,
+                "max_seqlen_q": 1,
+                "max_seqlen_k": 1,
+            }
+        elif contract == "packed-padding":
+            metadata.packed_padding = PackedPaddingMetadata(1, 1, lengths[:2], lengths[:2])
+        elif contract == "piecewise":
+            metadata.full_attn_spans = [[(0, 1)]]
+            metadata.query_ranges = (QueryRange(0, 2, 0),)
+        query = torch.ones(1, 2, 4, 8)
+
+        attention(query, query, query, metadata)
+
+        assert events == ["pre", "selected", "post"]
+        forwarded = state["kernel_call"][3]
+        if quant_dtype:
+            # Quantization is injected by the public forward from config.
+            assert forwarded.extra == {"kv_cache_dtype": "fp8"}
+            assert metadata.extra == {}
+        else:
+            assert forwarded is state["prepared"][3]
+
+
 class TestDiffusionKvCacheQuantization:
     @staticmethod
     def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True):

tests/diffusion/attention/test_block_sparse_ops.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import sys
+import types
+
+import pytest
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+from vllm_omni.diffusion.attention.ops.block_sparse import (
+    build_prefix_dense_block_map,
+    fastvideo_block_sparse_attn_bshd,
+    mean_pool_tiles,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("prefix_blocks,topk", [(0, 1), (2, 0), (2, 1), (2, 3), (2, 5)])
+def test_block_map_keeps_prefix_dense_and_selects_non_prefix_keys(prefix_blocks, topk):
+    sparse_blocks = 3
+    blocks = prefix_blocks + sparse_blocks
+    scores = torch.arange(2 * blocks * blocks, dtype=torch.float32).reshape(1, 2, blocks, blocks)
+    block_map = build_prefix_dense_block_map(scores, prefix_blocks, sparse_blocks, topk)
+    expected = torch.zeros_like(scores, dtype=torch.bool)
+    # Scores increase with key index, so top-k always selects the last keys.
+    for query_block in range(blocks):
+        for key_block in range(blocks):
+            if (
+                query_block < prefix_blocks
+                or key_block < prefix_blocks
+                or key_block >= blocks - min(topk, sparse_blocks)
+            ):
+                expected[..., query_block, key_block] = True
+    assert torch.equal(block_map, expected)
+
+
+@pytest.mark.parametrize("block_size", [4, 32, 64])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_pooling_uses_valid_tile_sizes_and_fp32_accumulation(block_size, dtype):
+    sizes = torch.tensor([block_size, block_size - 1, 0], dtype=torch.int32)
+    tiles = torch.arange(2 * 3 * block_size * 2 * 3).reshape(2, 3, block_size, 2, 3).to(dtype)
+    tiles[:, 1, -1] = 0
+    tiles[:, 2] = 0
+    expected = torch.stack(
+        [tiles[:, 0].float().mean(dim=1), tiles[:, 1, :-1].float().mean(dim=1), torch.zeros(2, 2, 3)], dim=2
+    )
+    result = mean_pool_tiles(tiles.flatten(1, 2), sizes, block_size)
+    assert result.dtype == torch.float32
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+def test_provider_wrapper_restores_transport_rows_and_bshd_layout(monkeypatch):
+    monkeypatch.delenv("FASTVIDEO_VSA_SM100A", raising=False)
+    calls = []
+
+    def block_sparse_attn(q, k, v, block_map, sizes):
+        calls.append((q.shape, block_map.shape, sizes.tolist()))
+        assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
+        assert sizes.dtype == torch.int32
+        return q + k + v, 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)
+    query = torch.randn(1, 4 * 64, 2, 8)
+    block_map = torch.ones(1, 2, 4, 4, dtype=torch.bool)
+    sizes = torch.tensor([64, 17, 64, 0])
+    output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
+    expected = torch.zeros_like(query)
+    expected[:, : 3 * 64] = query[:, : 3 * 64] * 3
+    torch.testing.assert_close(output, expected)
+    assert calls == [(torch.Size([1, 2, 3 * 64, 8]), torch.Size([1, 2, 3, 3]), [64, 17, 64])]
+
+
+def test_custom_op_fake_preserves_metadata_without_loading_provider(monkeypatch):
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", None)
+    with FakeTensorMode():
+        query = torch.empty(1, 4 * 64, 2, 8, dtype=torch.bfloat16)
+        block_map = torch.empty(1, 2, 4, 4, dtype=torch.bool)
+        sizes = torch.empty(4, dtype=torch.int32)
+        output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
+        assert output.shape == query.shape
+        assert output.dtype == query.dtype
+        assert output.device == query.device
+
+
+@pytest.mark.parametrize("error_type", [ImportError, RuntimeError, torch.AcceleratorError])
+def test_native_provider_failure_only_recovers_without_accelerator_fault(monkeypatch, mocker, error_type):
+    monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1")
+    error = error_type("native provider failure")
+    native = mocker.Mock()
+    native.is_supported.return_value = True
+    native.block_sparse_attn_sm100a.side_effect = error
+    provider = types.ModuleType("fastvideo_kernel")
+    setattr(provider, "block_sparse_attn_sm100a", native)
+    index = types.ModuleType("fastvideo_kernel.triton_kernels.index")
+    setattr(
+        index,
+        "map_to_index",
+        mocker.Mock(return_value=(torch.zeros(1, dtype=torch.int32), torch.ones(1, dtype=torch.int32))),
+    )
+    triton = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    fallback = mocker.Mock(side_effect=lambda q, *args: (q.clone(), None))
+    setattr(triton, "block_sparse_attn", fallback)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", provider)
+    monkeypatch.setitem(
+        sys.modules, "fastvideo_kernel.triton_kernels", types.ModuleType("fastvideo_kernel.triton_kernels")
+    )
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.triton_kernels.index", index)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", triton)
+    query = torch.arange(128 * 2 * 8, dtype=torch.float32).reshape(1, 128, 2, 8)
+    block_map = torch.ones(1, 2, 2, 2, dtype=torch.bool)
+    sizes = torch.tensor([64, 64], dtype=torch.int32)
+
+    if error_type is torch.AcceleratorError:
+        with pytest.raises(torch.AcceleratorError, match="native provider failure") as raised:
+            fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
+        assert raised.value is error
+        fallback.assert_not_called()
+    else:
+        output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
+        torch.testing.assert_close(output, query, rtol=0, atol=0)
+        fallback.assert_called_once()

tests/diffusion/attention/test_fastvideo_vsa.py

@@ -1,7 +1,6 @@
 # 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
@@ -12,14 +11,10 @@ import torch
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
     PackedPaddingMetadata,
-    VideoTokenLayout,
-    VideoTokenSpan,
 )
 from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
     FastVideoVSABackend,
     FastVideoVSAImpl,
-    _build_h3_block_map,
-    _get_h3_tile_metadata,
 )
 from vllm_omni.diffusion.attention.backends.registry import (
     DiffusionAttentionBackendEnum,
@@ -218,144 +213,8 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
     assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
 
 
-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 FastVideoVSAImpl(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_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
-    # The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
+    # The backend advertises supports_packed_mask_free, so a model may skip
     # building the padding mask. Every fallback out of the VSA path must honour
     # that contract itself; SDPA reads attn_mask and nothing else.
     impl = FastVideoVSAImpl(
@@ -374,6 +233,8 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
         cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
     )
 
+    monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
+
     baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
     perturbed_input = query.clone()
     perturbed_input[:, valid:] += 100.0
@@ -384,22 +245,3 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
     assert baseline.shape == query.shape
     torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
     assert torch.count_nonzero(baseline[:, 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())
-
-
-def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
-    scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
-    block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
-    assert block_map[:, :, :2].all()
-    assert block_map[..., :2].all()
-    assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()

tests/diffusion/attention/test_flash_attn.py

@@ -16,13 +16,14 @@ from unittest.mock import Mock
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
 from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
 HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
@@ -42,10 +43,12 @@ flash_attn_func = fa.flash_attn_func  # noqa: N813
         (9, "4", frozenset((2, 3, 4)), 4),
     ],
 )
+@pytest.mark.cpu
 def test_versioned_flash_attention_selection(device_major, requested, supported, expected):
     assert fa._choose_vllm_flash_attn_version(device_major, requested, supported) == expected
 
 
+@pytest.mark.cpu
 def test_versioned_flash_attention_selection_rejects_invalid_or_unavailable_version():
     with pytest.raises(ValueError, match="must be 2, 3, or 4"):
         fa._choose_vllm_flash_attn_version(9, "5", frozenset((2, 3, 4)))
@@ -92,7 +95,8 @@ def pad_tensor(tensor: torch.Tensor, target_seq_len: int, pad_value: float = 0.0
     return torch.cat([tensor, padding], dim=1)
 
 
-@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires an accelerator")
 def test_padding_equivalence():
     """
     Case 1: Test that padded and unpadded inputs produce similar outputs.
@@ -185,7 +189,8 @@ def test_padding_equivalence():
     print("✓ Case 1 PASSED: Padded and unpadded outputs are very close!")
 
 
-@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires an accelerator")
 def test_fa_vs_sdpa():
     """
     Case 2: Compare FlashAttention and SDPA backends with padding.
@@ -302,7 +307,8 @@ def test_fa_vs_sdpa():
     print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")
 
 
-@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires an accelerator")
 def test_flash_attn_func_preferred_over_varlen():
     """Test flash_attn_func availability and basic forward call."""
     if not HAS_FLASH_ATTN:
@@ -332,7 +338,8 @@ def test_flash_attn_func_preferred_over_varlen():
     print("✓ flash_attn_func forward works correctly!")
 
 
-@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires an accelerator")
 @pytest.mark.parametrize("k_len", [40, 256])
 def test_cross_attn_key_padding_vs_sdpa(k_len):
     """
@@ -390,6 +397,7 @@ def test_cross_attn_key_padding_vs_sdpa(k_len):
     print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
 
 
+@pytest.mark.cpu
 def test_varlen_masked_routing_by_role(monkeypatch):
     """The unpad route is picked by role, not Q/K length equality; runs without a GPU."""
     calls = []
@@ -424,6 +432,7 @@ def test_varlen_masked_routing_by_role(monkeypatch):
     assert cu_seqlens_k.tolist() == [0, 2, 5]
 
 
+@pytest.mark.cpu
 def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     calls = []
 
@@ -454,6 +463,7 @@ def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     assert calls[0]["softmax_scale"] == 0.5
 
 
+@pytest.mark.cpu
 def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch):
     calls = []
 
@@ -486,6 +496,7 @@ def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch):
     assert calls[0][3]["max_seqlen_q"] == 6
 
 
+@pytest.mark.cpu
 def test_packed_varlen_metadata_must_be_complete(monkeypatch):
     monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
     impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
@@ -503,11 +514,12 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
 #   - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
 #     sync, no MindIE-SD import);
 #   - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
+#   - native torch_npu dispatch for right-down causal attention;
 #   - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
 #
-# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
-# into ``sys.modules`` (same pattern as ``test_paged_attention`` /
-# ``benchmarks/conftest``). No production code is exercised for real kernels.
+# They run on CPU without a real NPU by injecting fake ``mindiesd`` and
+# ``torch_npu`` modules into ``sys.modules`` (same pattern as
+# ``test_paged_attention`` / ``benchmarks/conftest``). No real kernel runs.
 
 
 def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
@@ -529,9 +541,19 @@ def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_var
     return fake
 
 
+def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None):
+    """Install the minimal torch_npu surface used by causal NPU attention."""
+    fake = SimpleNamespace(
+        npu_fusion_attention=npu_fusion_attention or Mock(return_value=(torch.zeros(1), None, None, None, 0, 0, 0)),
+    )
+    monkeypatch.setitem(sys.modules, "torch_npu", fake)
+    return fake
+
+
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 
 
+@pytest.mark.cpu
 def test_resolve_packed_seq_accepts_real_plus_pad_contract():
     impl = _npu_impl()
     q = torch.randn(1, 8, 2, 4)
@@ -541,6 +563,7 @@ def test_resolve_packed_seq_accepts_real_plus_pad_contract():
     assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
 
 
+@pytest.mark.cpu
 def test_resolve_packed_seq_single_document_no_padding():
     impl = _npu_impl()
     q = torch.randn(1, 8, 2, 4)
@@ -564,6 +587,7 @@ def test_resolve_packed_seq_single_document_no_padding():
         "pad_longer_than_real",
     ],
 )
+@pytest.mark.cpu
 def test_resolve_packed_seq_rejects_malformed_metadata(case):
     """Any metadata outside the exact [real, pad] contract returns None so the
     caller falls back to the masked path."""
@@ -620,6 +644,7 @@ def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
     return impl, sentinel
 
 
+@pytest.mark.cpu
 def test_npu_env_unset_uses_varlen(monkeypatch):
     monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
     _fake_mindiesd(monkeypatch)
@@ -632,6 +657,7 @@ def test_npu_env_unset_uses_varlen(monkeypatch):
     impl._forward_prefix_kv_slice_npu.assert_not_called()
 
 
+@pytest.mark.cpu
 def test_npu_env_laser_uses_slice(monkeypatch):
     monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
     _fake_mindiesd(monkeypatch)
@@ -645,6 +671,7 @@ def test_npu_env_laser_uses_slice(monkeypatch):
 
 
 @pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
+@pytest.mark.cpu
 def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
     """Locks in the CURRENT dispatch behavior (PR #5891).
 
@@ -665,6 +692,7 @@ def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
     impl._forward_prefix_kv_slice_npu.assert_not_called()
 
 
+@pytest.mark.cpu
 def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
     """Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
     the existing masked attention_forward path; the new branches are skipped."""
@@ -685,9 +713,115 @@ def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
     assert out is fake_forward.return_value
 
 
-# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
+# --- Test group C: native NPU causal attention dispatch ----------------------
 
 
+@pytest.mark.parametrize(
+    ("query_length", "key_length", "expected_inner_precise"),
+    [(4, 4, 0), (2, 4, 0), (4, 2, 2)],
+)
+@pytest.mark.cpu
+def test_npu_causal_uses_native_right_down_mode(
+    monkeypatch,
+    query_length,
+    key_length,
+    expected_inner_precise,
+):
+    expected_out = torch.randn(1, query_length, 2, 4)
+    fusion_attention = Mock(return_value=(expected_out, None, None, None, 0, 0, 0))
+    fake_torch_npu = _fake_torch_npu(monkeypatch, npu_fusion_attention=fusion_attention)
+    impl = FlashAttentionImpl(num_heads=8, head_size=4, softmax_scale=0.5, causal=True)
+    query = torch.randn(1, query_length, 2, 4)
+    key = torch.randn(1, key_length, 2, 4)
+
+    out = impl.forward_fa_npu(query, key, key)
+
+    fake_torch_npu.npu_fusion_attention.assert_called_once()
+    args, kwargs = fake_torch_npu.npu_fusion_attention.call_args
+    assert torch.equal(args[0], query)
+    assert torch.equal(args[1], key)
+    assert torch.equal(args[2], key)
+    assert all(tensor.is_contiguous() for tensor in args[:3])
+    assert kwargs["head_num"] == query.shape[2]
+    assert kwargs["input_layout"] == "BSND"
+    assert kwargs["scale"] == 0.5
+    assert kwargs["keep_prob"] == 1.0
+    assert kwargs["sparse_mode"] == 3
+    assert kwargs["inner_precise"] == expected_inner_precise
+
+    block_mask = kwargs["atten_mask"]
+    assert block_mask.shape == (2048, 2048)
+    assert block_mask.dtype == torch.bool
+    assert block_mask.is_contiguous()
+    expected_corner = torch.tensor(
+        [
+            [False, True, True, True],
+            [False, False, True, True],
+            [False, False, False, True],
+            [False, False, False, False],
+        ]
+    )
+    assert torch.equal(block_mask[:4, :4], expected_corner)
+    assert out is expected_out
+
+
+@pytest.mark.cpu
+def test_npu_causal_composes_explicit_keep_mask(monkeypatch):
+    expected_out = torch.randn(1, 4, 2, 4)
+    fusion_attention = Mock(return_value=(expected_out, None, None, None, 0, 0, 0))
+    fake_torch_npu = _fake_torch_npu(monkeypatch, npu_fusion_attention=fusion_attention)
+    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=True)
+    query = torch.randn(1, 4, 2, 4)
+    key_keep_mask = torch.tensor([[True, False, True, True]])
+
+    out = impl.forward_fa_npu(
+        query,
+        query,
+        query,
+        AttentionMetadata(attn_mask=key_keep_mask),
+    )
+
+    fake_torch_npu.npu_fusion_attention.assert_called_once()
+    kwargs = fake_torch_npu.npu_fusion_attention.call_args.kwargs
+    assert kwargs["sparse_mode"] == 1
+    assert kwargs["inner_precise"] == 2
+    expected_block_mask = torch.tensor(
+        [
+            [
+                [
+                    [False, True, True, True],
+                    [False, True, True, True],
+                    [False, True, False, True],
+                    [False, True, False, False],
+                ]
+            ]
+        ]
+    )
+    assert torch.equal(kwargs["atten_mask"], expected_block_mask)
+    assert kwargs["atten_mask"].is_contiguous()
+    assert out is expected_out
+
+
+@pytest.mark.cpu
+def test_npu_noncausal_without_explicit_mask_stays_unmasked(monkeypatch):
+    captured: dict = {}
+
+    def fake_attention_forward(query, key, value, **kwargs):
+        captured.update(kwargs)
+        return torch.zeros_like(query)
+
+    _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
+    query = torch.randn(1, 4, 2, 4)
+
+    _npu_impl(causal=False).forward_fa_npu(query, query, query)
+
+    assert captured["attn_mask"] is None
+
+
+# --- Test group D: laser input pre-scaling in _forward_prefix_kv_slice_npu --
+
+
+@pytest.mark.cpu
 def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
     monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
     captured: dict = {}
@@ -726,6 +860,7 @@ def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
     torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
 
 
+@pytest.mark.cpu
 def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
     monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
     captured: dict = {}

tests/diffusion/attention/test_flash_attn_npu.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
+
+
+def _bottom_right_causal_reference(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    scale: float,
+    key_keep_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * scale
+    query_positions = torch.arange(query.shape[1]).unsqueeze(1)
+    key_positions = torch.arange(key.shape[1]).unsqueeze(0)
+    keep_mask = key_positions <= query_positions + (key.shape[1] - query.shape[1])
+    if key_keep_mask is not None:
+        assert query.shape[0] == 1
+        keep_mask = keep_mask & key_keep_mask[0].bool().unsqueeze(0)
+    valid_rows = keep_mask.any(dim=-1)
+
+    output = torch.zeros(
+        query.shape[0],
+        query.shape[1],
+        query.shape[2],
+        value.shape[-1],
+        dtype=torch.float32,
+    )
+    valid_scores = scores[:, :, valid_rows].masked_fill(~keep_mask[valid_rows], float("-inf"))
+    probabilities = torch.softmax(valid_scores, dim=-1)
+    output[:, valid_rows] = torch.einsum("bhqk,bkhd->bqhd", probabilities, value.float())
+    return output
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native causal attention requires Ascend NPU")
+def test_npu_bottom_right_causal_fully_masked_rows_are_zero():
+    torch.manual_seed(0)
+    batch_size, query_length, key_length = 1, 4, 2
+    num_heads, head_size = 2, 64
+    scale = head_size**-0.5
+
+    query_cpu = torch.randn(batch_size, query_length, num_heads, head_size).to(torch.bfloat16)
+    key_cpu = torch.randn(batch_size, key_length, num_heads, head_size).to(torch.bfloat16)
+    value_cpu = torch.randn(batch_size, key_length, num_heads, head_size).to(torch.bfloat16)
+    expected = _bottom_right_causal_reference(query_cpu, key_cpu, value_cpu, scale)
+
+    device = torch.device(current_omni_platform.device_type)
+    query = query_cpu.to(device)
+    key = key_cpu.to(device)
+    value = value_cpu.to(device)
+    impl = FlashAttentionImpl(
+        num_heads=num_heads,
+        head_size=head_size,
+        softmax_scale=scale,
+        causal=True,
+    )
+
+    output = impl.forward_fa_npu(query, key, value)
+    output_cpu = output.float().cpu()
+
+    assert torch.isfinite(output_cpu).all()
+    assert torch.equal(output_cpu[:, : query_length - key_length], torch.zeros_like(output_cpu[:, :2]))
+    torch.testing.assert_close(
+        output_cpu[:, query_length - key_length :],
+        expected[:, query_length - key_length :],
+        rtol=2e-2,
+        atol=2e-2,
+    )
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native causal attention requires Ascend NPU")
+def test_npu_bottom_right_causal_composes_explicit_keep_mask():
+    torch.manual_seed(1)
+    batch_size, sequence_length = 1, 4
+    num_heads, head_size = 2, 64
+    scale = head_size**-0.5
+    key_keep_mask_cpu = torch.tensor([[True, False, True, True]])
+
+    query_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16)
+    key_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16)
+    value_cpu = torch.randn(batch_size, sequence_length, num_heads, head_size).to(torch.bfloat16)
+    expected = _bottom_right_causal_reference(
+        query_cpu,
+        key_cpu,
+        value_cpu,
+        scale,
+        key_keep_mask_cpu,
+    )
+
+    device = torch.device(current_omni_platform.device_type)
+    query = query_cpu.to(device)
+    key = key_cpu.to(device)
+    value = value_cpu.to(device)
+    impl = FlashAttentionImpl(
+        num_heads=num_heads,
+        head_size=head_size,
+        softmax_scale=scale,
+        causal=True,
+    )
+
+    output = impl.forward_fa_npu(
+        query,
+        key,
+        value,
+        AttentionMetadata(attn_mask=key_keep_mask_cpu.to(device)),
+    )
+    output_cpu = output.float().cpu()
+
+    assert torch.isfinite(output_cpu).all()
+    torch.testing.assert_close(output_cpu, expected, rtol=2e-2, atol=2e-2)

tests/diffusion/attention/test_impl_overrides.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from vllm_omni.diffusion.attention import layer
+from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSABackend, FastVideoVSAImpl
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+from vllm_omni.diffusion.data import AttentionSpec
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _SpecializedImpl(FastVideoVSAImpl):
+    pass
+
+
+@pytest.fixture
+def local_attention(monkeypatch):
+    monkeypatch.setattr(layer, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention())
+
+
+@pytest.mark.parametrize("explicit", [False, True])
+@pytest.mark.parametrize("override", [False, True])
+def test_specialization_preserves_selected_backend_and_options(mocker, local_attention, explicit, override):
+    spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=7) if explicit else None
+    select = mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, spec))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        role="video.self",
+        role_category="self",
+        qkv_layout="BSND",
+        impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl} if override else None,
+    )
+
+    assert attention.attn_backend is FastVideoVSABackend
+    assert type(attention.attention) is (_SpecializedImpl if override else FastVideoVSAImpl)
+    assert attention.attention.topk == (7 if explicit else 64)
+    assert attention.attention.qkv_layout == "BSND"
+    assert attention.backend_pref == "FASTVIDEO_VSA"
+    assert attention.backend_explicit is explicit
+    assert attention.attn_spec is spec
+    assert isinstance(attention.parallel_strategy, NoParallelAttention)
+    select.assert_called_once_with(
+        role="video.self", head_size=8, attention_config=None, role_category="self", allow_trtllm_default=False
+    )
+
+
+def test_unselected_specialization_does_not_replace_dense_backend(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(SDPABackend, None))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+    )
+    assert attention.attn_backend is SDPABackend
+    assert type(attention.attention) is SDPABackend.get_impl_cls()
+
+
+def test_specialization_does_not_bypass_selection_errors(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", side_effect=ImportError("optional kernel unavailable"))
+    with pytest.raises(ImportError, match="optional kernel unavailable"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+        )
+
+
+def test_specialization_must_extend_selected_implementation(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, None))
+    with pytest.raises(TypeError, match="must subclass the selected implementation"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": SDPABackend.get_impl_cls()},
+        )
+
+
+def test_platform_backend_capabilities_survive_model_specialization(mocker, local_attention):
+    from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import MiniMaxH3VSAImpl
+
+    class PlatformBackend(FastVideoVSABackend):
+        @classmethod
+        def supports_packed_mask_free(cls):
+            return False
+
+        @staticmethod
+        def get_supported_head_sizes():
+            return [128]
+
+    # A registry/platform may refine capabilities without changing the
+    # implementation. The model must not replace those capability decisions.
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        causal=False,
+        impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl},
+    )
+    assert attention.attn_backend is PlatformBackend
+    assert not attention.attn_backend.supports_packed_mask_free()
+    assert attention.attn_backend.get_supported_head_sizes() == [128]
+    assert type(attention.attention) is MiniMaxH3VSAImpl
+
+
+def test_incompatible_platform_implementation_is_not_discarded(mocker, local_attention):
+    class PlatformImpl(FastVideoVSAImpl):
+        pass
+
+    class PlatformBackend(FastVideoVSABackend):
+        @staticmethod
+        def get_impl_cls():
+            return PlatformImpl
+
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
+    with pytest.raises(TypeError, match="selected implementation .*PlatformImpl"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+        )

tests/diffusion/attention/test_mindie_public_integration.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cross-repository CPU dispatch contracts, requiring the matching MindIE-SD source.
+
+Only native operators are replaced. Public adapters, rotations, layouts, padding,
+scale preparation, sequence metadata and output cropping execute their real code.
+Operator substitutes carry floating tensors; these are not numerical quantization tests.
+"""
+
+import importlib
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends import flash_attn, rainfusion_attn
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.platforms.npu.quant.kv_quant_npu import get_quant_attention_rotation
+
+mindiesd = pytest.importorskip("mindiesd")
+torch_npu = importlib.import_module("torch_npu")
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    "precision,batch,q_len,kv_len,q_heads",
+    [
+        ("fp8", 1, 130, 130, 2),
+        ("mxfp8", 1, 130, 130, 2),
+        ("mxfp4", 1, 130, 130, 2),
+        ("mxfp8", 2, 17, 8, 4),
+        ("mxfp8", 2, 8, 17, 4),
+    ],
+)
+@pytest.mark.parametrize("layout", ["BSND", "BNSD"])
+def test_omni_calls_real_dense_public_api(monkeypatch, precision, layout, batch, q_len, kv_len, q_heads):
+    assert callable(getattr(mindiesd, "quant_attention", None)), "MindIE-SD must export quant_attention"
+    monkeypatch.setattr(flash_attn, "current_omni_platform", SimpleNamespace(is_npu=lambda: True, device_name="npu"))
+    monkeypatch.setattr(flash_attn, "get_current_diffusion_config_or_none", lambda: None)
+    quantized_inputs = []
+
+    def dynamic_quant(tensor, **kwargs):
+        quantized_inputs.append((tensor.clone(), kwargs))
+        shape = list(tensor.shape)
+        axis = kwargs.get("axis", -1)
+        shape[axis] = max(1, shape[axis] // 32)
+        return tensor, torch.ones(shape, dtype=torch.uint8)
+
+    monkeypatch.setattr(torch_npu, "npu_dynamic_block_quant", dynamic_quant, raising=False)
+    monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", dynamic_quant, raising=False)
+    execute = Mock(side_effect=lambda query, *args, **kwargs: (query, None))
+    metadata = Mock(return_value=torch.zeros(1, dtype=torch.int32))
+    if precision == "fp8":
+        monkeypatch.setattr(torch.ops.mindiesd, "fused_infer_attention_score_v2", execute, raising=False)
+    elif precision == "mxfp8":
+        monkeypatch.setattr(torch_npu, "npu_fused_infer_attention_score_v2", execute, raising=False)
+    else:
+        monkeypatch.setattr(torch.ops.mindiesd, "quant_flash_attn", execute, raising=False)
+        monkeypatch.setattr(torch.ops.mindiesd, "quant_flash_attn_metadata", metadata, raising=False)
+
+    query = torch.randn(batch, q_len, q_heads, 64, dtype=torch.bfloat16)
+    key = torch.randn(batch, kv_len, 2, 64, dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    if layout == "BNSD":
+        query, key, value = (tensor.transpose(1, 2) for tensor in (query, key, value))
+    impl = flash_attn.FlashAttentionImpl(
+        num_heads=q_heads,
+        head_size=64,
+        softmax_scale=0.37,
+        qkv_layout=layout,
+    )
+    output = impl.forward_npu(query, key, value, AttentionMetadata(extra={"kv_cache_dtype": precision}))
+    rotation = get_quant_attention_rotation(query.device, query.dtype, 64)
+    expected = query @ rotation
+    torch.testing.assert_close(output, expected)
+    assert output.shape == query.shape and len(quantized_inputs) == 3
+    execute.assert_called_once()
+    kwargs = execute.call_args.kwargs
+    assert kwargs["softmax_scale"] == 0.37
+    if precision == "fp8":
+        assert kwargs["input_layout"] == "BNSD"
+        assert quantized_inputs[0][0].shape == (2, 130, 64)
+        assert quantized_inputs[0][1]["row_block_size"] == 128
+    elif precision == "mxfp8":
+        args = execute.call_args.args
+        assert args[0].shape == (batch * q_len, q_heads, 64)
+        assert args[1].shape == args[2].shape == (batch * kv_len, 2, 64)
+        assert kwargs["input_layout"] == "TND"
+        assert kwargs["num_query_heads"] == q_heads and kwargs["num_key_value_heads"] == 2
+        assert kwargs["actual_seq_qlen"] == [q_len * (i + 1) for i in range(batch)]
+        assert kwargs["actual_seq_kvlen"] == [kv_len * (i + 1) for i in range(batch)]
+        assert [entry[1]["axis"] for entry in quantized_inputs] == [-1, -1, 0]
+    else:
+        seq_axis = 1 if layout == "BSND" else 2
+        assert quantized_inputs[0][0].shape[seq_axis] == 512
+        assert kwargs["seqused_q"].tolist() == kwargs["seqused_kv"].tolist() == [512]
+        assert kwargs["layout_q"] == kwargs["layout_out"] == layout
+        assert kwargs["q_dtype"] == torch_npu.float4_e2m1fn_x2
+        assert kwargs["q_descale_dtype"] == torch_npu.float8_e8m0fnu
+        assert execute.call_args.args[3].shape[:2] == (1, 2)
+        assert [entry[1]["axis"] for entry in quantized_inputs] == [-1, -1, seq_axis]
+        assert kwargs["metadata"] is metadata.return_value
+        metadata.assert_called_once()
+
+
+@pytest.mark.parametrize("precision", ["fp8", "mxfp4"])
+@pytest.mark.parametrize("length", [65, 600])
+def test_omni_calls_real_sparse_public_api(monkeypatch, precision, length):
+    rf = importlib.import_module("mindiesd.layers.flash_attn.sparse_flash_attn_rf_v3")
+    monkeypatch.setattr(flash_attn, "current_omni_platform", SimpleNamespace(is_npu=lambda: True, device_name="npu"))
+    for module in (flash_attn, rainfusion_attn):
+        monkeypatch.setattr(module, "get_current_diffusion_config_or_none", lambda: None)
+    rainfusion_attn._mindiesd_supports_precision.cache_clear()
+    monkeypatch.setattr(rf, "do_tensor_rearrange_only", lambda q, k, v, *a, **kw: (q, k, v))
+    monkeypatch.setattr(rf, "avgpool", lambda *a, **kw: torch.ones(1))
+    monkeypatch.setattr(rf, "_generate_mask_direct", lambda *a, **kw: torch.ones(1, 2, 8, 8))
+    monkeypatch.setattr(rf, "_bsa_inv_rearrange", lambda out, *a: out)
+
+    def quantize(tensor, **kwargs):
+        shape = list(tensor.shape)
+        axis = kwargs.get("axis", -1)
+        shape[axis] = max(1, shape[axis] // 32)
+        return tensor, torch.ones(shape, dtype=torch.uint8)
+
+    monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", quantize, raising=False)
+    monkeypatch.setattr(torch_npu, "npu_dynamic_block_quant", quantize, raising=False)
+    execute = Mock(side_effect=lambda **kw: (torch.ones_like(kw["query"]), None))
+    monkeypatch.setattr(torch.ops.mindiesd, "block_sparse_attention", execute, raising=False)
+    query = torch.randn(1, length, 2, 64, dtype=torch.bfloat16)
+    impl = rainfusion_attn.RainFusionAttentionImpl(
+        num_heads=2,
+        head_size=64,
+        softmax_scale=0.37,
+        qkv_layout="BSND",
+        backend_kwargs={"sparsity": 0.8},
+    )
+    plan = rainfusion_attn.RainFusionPlan(used_len=length, prefix_len=0, latent_shape=(1, 1, length))
+    try:
+        output = impl._forward_sparse_npu(
+            query,
+            query,
+            query,
+            plan,
+            AttentionMetadata(extra={"kv_cache_dtype": precision}),
+        )
+    finally:
+        rainfusion_attn._mindiesd_supports_precision.cache_clear()
+    torch.testing.assert_close(output, torch.ones_like(query))
+    execute.assert_called_once()
+    kwargs = execute.call_args.kwargs
+    assert kwargs["actual_seq_lengths"] == kwargs["actual_seq_lengths_kv"] == [length]
+    assert kwargs["q_input_layout"] == "BNSD"
+    assert kwargs["block_sparse_mask"] is not None
+    if precision == "mxfp4":
+        assert kwargs["query"].shape[2] == (length + 63) // 64 * 64
+        assert kwargs["quant_mode"] == 2
+        assert kwargs["q_dtype"] == torch_npu.float4_e2m1fn_x2
+        assert kwargs["v_scale_dtype"] == torch_npu.float8_e8m0fnu

tests/diffusion/attention/test_mindie_runtime.py

@@ -0,0 +1,323 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU contract tests for MindIE adapters; no NPU kernels are mocked in production."""
+
+import sys
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention import layer as layer_mod
+from vllm_omni.diffusion.attention.backends import flash_attn, rainfusion_attn
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+_DEFAULT_SKIP_STEPS = frozenset({0, 1, 38, 39})
+
+
+@pytest.fixture
+def runtime(monkeypatch):
+    mod = ModuleType("mindiesd")
+
+    def execute(q, k, v, *, precision, layout, scale, q_rot=None, k_rot=None):
+        return q
+
+    setattr(mod, "quant_attention", Mock(side_effect=execute))
+    monkeypatch.setitem(sys.modules, mod.__name__, mod)
+    for module in (flash_attn, layer_mod):
+        monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(is_npu=lambda: True, device_name="npu"))
+    monkeypatch.setattr(flash_attn, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(rainfusion_attn, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(rainfusion_attn, "is_forward_context_available", lambda: False)
+    supports_precision = rainfusion_attn._mindiesd_supports_precision
+    supports_precision.cache_clear()
+    yield mod
+    supports_precision.cache_clear()
+
+
+def flash(layout="BSND"):
+    return flash_attn.FlashAttentionImpl(
+        num_heads=2,
+        head_size=64,
+        softmax_scale=0.37,
+        qkv_layout=layout,
+    )
+
+
+def sparse(**kwargs):
+    return rainfusion_attn.RainFusionAttentionImpl(
+        num_heads=2,
+        head_size=64,
+        softmax_scale=0.37,
+        qkv_layout="BSND",
+        backend_kwargs={"sparsity": 0.8, **kwargs},
+    )
+
+
+def video_metadata(**extra):
+    return AttentionMetadata(
+        video_layout=VideoTokenLayout(prefix_len=0, latent_grid=(4, 32, 32), used_len=4096),
+        extra={"max_seqlen_q": 4096, **extra},
+    )
+
+
+@pytest.mark.parametrize("method", ["fp8", "mxfp8", "mxfp4"])
+@pytest.mark.parametrize("layout", [None, "BSND", "BNSD"])
+def test_exact_runtime_layout_scale_and_rotation(runtime, method, layout):
+    q = torch.randn(1, 128, 2, 64, dtype=torch.bfloat16)
+    if layout == "BNSD":
+        q = q.transpose(1, 2)
+    impl = flash(layout=layout)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": method})
+    assert impl.forward_npu(q, q, q, metadata) is q
+    args, kwargs = runtime.quant_attention.call_args
+    assert args[0] is q  # no unconditional transpose or extra quantization
+    assert kwargs["layout"] == (layout or "BSND") and kwargs["scale"] == 0.37 and kwargs["precision"] == method
+    assert "attn_mask" not in kwargs
+    from vllm_omni.platforms.npu.quant.kv_quant_npu import get_quant_attention_rotation
+
+    assert kwargs["q_rot"] is kwargs["k_rot"]
+    torch.testing.assert_close(kwargs["q_rot"], get_quant_attention_rotation(q.device, q.dtype, 64))
+
+
+def test_missing_runtime_requires_config_change(runtime):
+    q = torch.randn(1, 128, 2, 64, dtype=torch.bfloat16)
+    del runtime.quant_attention
+    with pytest.raises(ImportError, match="diffusion_kv_cache_dtype"):
+        flash().forward_npu(q, q, q, AttentionMetadata(extra={"kv_cache_dtype": "mxfp4"}))
+
+
+def test_operator_failure_never_falls_back(runtime):
+    q = torch.randn(1, 128, 2, 64, dtype=torch.bfloat16)
+    runtime.quant_attention.side_effect = RuntimeError("operator failure")
+    with pytest.raises(RuntimeError, match="operator failure"):
+        flash().forward_npu(q, q, q, AttentionMetadata(extra={"kv_cache_dtype": "mxfp4"}))
+    runtime.quant_attention.assert_called_once()
+
+
+@pytest.mark.parametrize("method", ["fp8", "mxfp8", "mxfp4"])
+def test_unsupported_shape_requires_config_change(runtime, method):
+    q = torch.randn(1, 128, 2, 96, dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="diffusion_kv_cache_dtype"):
+        flash().forward_npu(q, q, q, AttentionMetadata(extra={"kv_cache_dtype": method}))
+    runtime.quant_attention.assert_not_called()
+
+
+def test_packed_metadata_never_reaches_quant_runtime(runtime):
+    q = torch.randn(1, 128, 2, 64, dtype=torch.bfloat16)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": "mxfp8", "cu_seqlens_q": torch.tensor([0, 64, 128])})
+    with pytest.raises(ValueError, match="packed/varlen"):
+        flash().forward_npu(q, q, q, metadata)
+    runtime.quant_attention.assert_not_called()
+
+
+def make_layer(runtime, monkeypatch, method="mxfp8", skip_steps=_DEFAULT_SKIP_STEPS):
+    layer = object.__new__(layer_mod.Attention)
+    torch.nn.Module.__init__(layer)
+    layer.attention = flash()
+    layer.attn_backend = flash_attn.FlashAttentionBackend
+    layer._disable_kv_quant = False
+    layer.layer_idx = 3
+    cfg = SimpleNamespace(
+        diffusion_kv_cache_dtype=method,
+        diffusion_kv_cache_skip_step_indices=skip_steps,
+        diffusion_kv_cache_skip_layer_indices=None,
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    layer._init_kv_cache_quantization(cfg)
+    return layer
+
+
+@pytest.mark.parametrize("method", ["fp8", "mxfp8", "mxfp4"])
+def test_global_quant_method_is_published_per_forward(runtime, monkeypatch, method):
+    layer = make_layer(runtime, monkeypatch, method=method, skip_steps=None)
+    assert layer._with_kv_cache_dtype(None).extra["kv_cache_dtype"] == method
+
+
+@pytest.mark.parametrize("method", [None, "auto", "float"])
+def test_optout_without_quant_policy_keeps_metadata_unchanged(runtime, monkeypatch, method):
+    layer = make_layer(runtime, monkeypatch)
+    layer._disable_kv_quant = True
+    layer._init_kv_cache_quantization(
+        SimpleNamespace(
+            diffusion_kv_cache_dtype=method,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+            parallel_config=SimpleNamespace(ring_degree=1),
+        )
+    )
+    source = video_metadata()
+    assert layer._with_kv_cache_dtype(source) is source
+
+
+def test_step_skip_is_request_local_and_cross_optout(runtime, monkeypatch):
+    layer = make_layer(runtime, monkeypatch)
+    ctx = SimpleNamespace(denoise_step_idx=0)
+    monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(layer_mod, "get_forward_context", lambda: ctx)
+    source = video_metadata()
+    for step in list(range(40)) * 2:
+        ctx.denoise_step_idx = step
+        resolved = layer._with_kv_cache_dtype(source)
+        assert resolved.extra["kv_cache_dtype"] == ("float" if step in (0, 1, 38, 39) else "mxfp8")
+        assert resolved.video_layout is source.video_layout
+    assert source.extra == {"max_seqlen_q": 4096}
+    layer._disable_kv_quant = True
+    assert layer._with_kv_cache_dtype(source).extra["kv_cache_dtype"] == "float"
+
+
+def test_ring_attention_is_rejected(runtime, monkeypatch):
+    layer = make_layer(runtime, monkeypatch)
+    cfg = SimpleNamespace(diffusion_kv_cache_dtype="mxfp8", parallel_config=SimpleNamespace(ring_degree=2))
+    with pytest.raises(ValueError, match="ring"):
+        layer._init_kv_cache_quantization(cfg)
+
+
+def test_sparse_tail_steps_and_dense_quant_dispatch(runtime, monkeypatch):
+    ctx = SimpleNamespace(denoise_step_idx=49, total_denoise_steps=50)
+    monkeypatch.setattr(rainfusion_attn, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(rainfusion_attn, "get_forward_context", lambda: ctx)
+    impl = sparse(start_step=2, end_step=2)
+    metadata = video_metadata(kv_cache_dtype="fp8")
+    assert impl._resolve_plan(metadata) is None
+    q = torch.randn(1, 4096, 2, 64, dtype=torch.bfloat16)
+    impl.forward_npu(q, q, q, metadata)
+    assert runtime.quant_attention.call_args.kwargs["precision"] == "fp8"
+
+
+@pytest.mark.parametrize("method", ["fp8", "mxfp4"])
+def test_sparse_quant_uses_public_runtime(runtime, monkeypatch, method):
+    runtime.sparse_attention = Mock(side_effect=lambda q, k, v, **kwargs: q)
+    monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: True)
+    q = torch.randn(1, 4224, 2, 64, dtype=torch.bfloat16)
+    out = sparse().forward_npu(q, q, q, video_metadata(kv_cache_dtype=method))
+    # The sparse path reuses the query tail as the output padding instead of
+    # zero-filling a fresh buffer; rows past the 4096-row valid prefix are
+    # never read back downstream (the model unpads by cu_seqlens), so pin the
+    # passthrough rather than requiring zeros.
+    assert out.shape == q.shape and torch.equal(out[:, 4096:], q[:, 4096:])
+    kwargs = runtime.sparse_attention.call_args.kwargs
+    assert (kwargs["precision"], kwargs["sparse_type"], kwargs["inner_precise"]) == (method, "rf_v3", 4)
+    runtime.quant_attention.assert_not_called()
+
+
+def test_sparse_custom_mask_stays_dense(runtime):
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones(1, 4096, dtype=torch.bool),
+        video_layout=VideoTokenLayout(prefix_len=0, latent_grid=(4, 32, 32)),
+        extra={"max_seqlen_q": 4096},
+    )
+    assert sparse()._resolve_plan(metadata) is None
+
+
+def test_sparse_model_padding_mask_is_equivalent_to_trimming(runtime):
+    metadata = video_metadata()
+    metadata.attn_mask = torch.arange(4224)[None] < 4096
+    assert sparse()._resolve_plan(metadata).used_len == 4096
+
+
+def test_bsa_operator_failure_never_retries(runtime, monkeypatch):
+    runtime.sparse_attention = Mock(side_effect=RuntimeError("BSA operator failure"))
+    monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: True)
+    q = torch.randn(1, 4096, 2, 64, dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError, match="BSA operator failure"):
+        sparse().forward_npu(q, q, q, video_metadata(kv_cache_dtype="mxfp4"))
+    runtime.sparse_attention.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "method,shape,supports_precision,error",
+    [
+        ("mxfp8", (1, 4096, 2, 64), True, "sparse mxfp8 is not supported"),
+        ("mxfp4", (1, 4096, 2, 64), False, "explicitly support precision"),
+        ("fp8", (2, 4096, 2, 64), True, "batch size 1"),
+        ("fp8", (1, 4096, 2, 32), True, "power-of-two head dimension"),
+        ("mxfp4", (1, 4096, 2, 96), True, "power-of-two head dimension"),
+    ],
+)
+def test_invalid_sparse_quant_requires_config_change(runtime, monkeypatch, method, shape, supports_precision, error):
+    runtime.sparse_attention = Mock(side_effect=lambda q, k, v, **kwargs: q)
+    monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: supports_precision)
+    q = torch.randn(shape, dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match=error):
+        sparse().forward_npu(q, q, q, video_metadata(kv_cache_dtype=method))
+    runtime.sparse_attention.assert_not_called()
+
+
+def test_layer_selector_requires_index_except_cross_optout(runtime, monkeypatch):
+    layer = make_layer(runtime, monkeypatch)
+    layer.layer_idx = None
+    config = SimpleNamespace(
+        diffusion_kv_cache_dtype="mxfp8",
+        diffusion_kv_cache_skip_step_indices=None,
+        diffusion_kv_cache_skip_layer_indices={0, 3},
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    with pytest.raises(ValueError, match="parseable transformer block index"):
+        layer._init_kv_cache_quantization(config)
+    layer._disable_kv_quant = True
+    layer._init_kv_cache_quantization(config)
+    assert layer._with_kv_cache_dtype(None).extra["kv_cache_dtype"] == "float"
+
+
+@pytest.mark.parametrize("prefix", ["transformer.blocks.3.attn1", "transformer_2.blocks.3.attn1"])
+@pytest.mark.parametrize("method", ["fp8", "mxfp4"])
+def test_expert_layer_and_step_skip_remain_sparse(runtime, monkeypatch, prefix, method):
+    def execute(q, k, v, *, precision="bf16", **kwargs):
+        return q
+
+    runtime.sparse_attention = Mock(wraps=execute)
+    monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: True)
+    layer = make_layer(runtime, monkeypatch)
+    layer.layer_idx = rainfusion_attn._try_extract_layer_index(prefix)
+    assert layer.layer_idx == 3
+    config = SimpleNamespace(
+        diffusion_kv_cache_dtype=method,
+        diffusion_kv_cache_skip_step_indices={0, 1, 38, 39},
+        diffusion_kv_cache_skip_layer_indices={3},
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    impl = sparse()
+    layer.attention = impl
+    layer.attn_backend = rainfusion_attn.RainFusionAttentionBackend
+    layer._init_kv_cache_quantization(config)
+    ctx = SimpleNamespace(denoise_step_idx=0)
+    monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(layer_mod, "get_forward_context", lambda: ctx)
+    source = video_metadata(kv_cache_dtype="mxfp8")
+    q = torch.randn(1, 4096, 2, 64, dtype=torch.bfloat16)
+    for layer_idx, step in ((3, 0), (3, 2), (4, 0), (4, 1), (4, 2), (4, 37), (4, 38), (4, 39), (4, 0)):
+        expected = "bf16" if layer_idx == 3 or step in (0, 1, 38, 39) else method
+        layer.layer_idx = layer_idx
+        ctx.denoise_step_idx = step
+        metadata = layer._with_kv_cache_dtype(source)
+        impl.forward_npu(q, q, q, metadata)
+        assert runtime.sparse_attention.call_args.kwargs["precision"] == expected
+    runtime.quant_attention.assert_not_called()
+    assert source.extra["kv_cache_dtype"] == "mxfp8"
+
+
+def test_custom_attention_initialization_keeps_upstream_optout(monkeypatch):
+    monkeypatch.setattr(layer_mod, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kw: object())
+    monkeypatch.setattr(layer_mod, "NoParallelAttention", lambda: object())
+    custom = torch.nn.Identity()
+    layer = layer_mod.Attention(
+        num_heads=2,
+        head_size=64,
+        softmax_scale=0.125,
+        causal=False,
+        custom_attention=custom,
+        skip_sequence_parallel=True,
+    )
+    assert layer.attention is custom and layer._kv_cache_dtype is None
+
+
+def test_legacy_global_fp8_uses_implicit_bsnd(runtime):
+    q = torch.randn(1, 17, 2, 64, dtype=torch.bfloat16)
+    impl = flash(layout=None)
+    assert impl.forward_npu(q, q, q, AttentionMetadata(extra={"kv_cache_dtype": "fp8"})) is q
+    assert runtime.quant_attention.call_args.kwargs["layout"] == "BSND"
+    assert runtime.quant_attention.call_args.kwargs["precision"] == "fp8"

tests/diffusion/attention/test_rainfusion_plan.py

@@ -133,18 +133,17 @@ def test_invalid_ref2va_video_spans_fall_back_to_dense():
     assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None
 
 
-def test_validate_available_rejects_legacy_mindiesd(monkeypatch):
+def test_validate_available_accepts_single_video_mindiesd(monkeypatch):
     mindiesd = types.ModuleType("mindiesd")
 
     def sparse_attention(query, key, value, **kwargs):
         return query
 
-    mindiesd.sparse_attention = sparse_attention
+    monkeypatch.setattr(mindiesd, "sparse_attention", sparse_attention, raising=False)
     monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
     monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
 
-    with pytest.raises(ValueError, match="video_spans"):
-        RainFusionAttentionBackend.validate_available()
+    RainFusionAttentionBackend.validate_available()
 
 
 def test_validate_available_accepts_new_mindiesd(monkeypatch):
@@ -153,13 +152,25 @@ def test_validate_available_accepts_new_mindiesd(monkeypatch):
     def sparse_attention(query, key, value, *, video_spans=None, **kwargs):
         return query
 
-    mindiesd.sparse_attention = sparse_attention
+    monkeypatch.setattr(mindiesd, "sparse_attention", sparse_attention, raising=False)
     monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
     monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
 
     RainFusionAttentionBackend.validate_available()
 
 
+def test_prefix_kv_slicing_contract_is_pinned():
+    # Regression pin for the #5543 packed-padding crash: reverting the flag
+    # would pass CI while reintroducing the ValueError for non-64-aligned
+    # requests. Both halves of the contract must hold together — the impl
+    # trims [real, pad] packed tensors to the valid prefix itself (so it may
+    # advertise prefix slicing) and it never reads attn_mask (so it must keep
+    # refusing one), which is what lets the model skip materializing the
+    # padding mask that _assert_metadata_compatible would reject.
+    assert RainFusionAttentionBackend.supports_prefix_kv_slicing is True
+    assert RainFusionAttentionBackend.supports_attention_mask() is False
+
+
 @pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
 def test_short_video_stays_dense(grid):
     assert make_impl()._resolve_plan(make_metadata(grid)) is None
@@ -300,22 +311,21 @@ def _fake_mindiesd_module():
     import types
 
     fake = types.ModuleType("mindiesd")
-    fake.sparse_attention = lambda *args, **kwargs: None
+    setattr(fake, "sparse_attention", lambda *args, **kwargs: None)
     return fake
 
 
 def test_precision_non_bf16_requires_mindiesd_support():
-    """precision != bf16 against a mindiesd lacking the kwarg must raise RuntimeError."""
+    """Unsupported sparse precision must fail before any operator executes."""
     import sys
 
     impl = make_impl(precision="mix")
-    sys.modules["mindiesd"] = _fake_mindiesd_module()
-    try:
+    with mock.patch.dict(sys.modules, {"mindiesd": _fake_mindiesd_module()}):
         with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=False):
-            with pytest.raises(RuntimeError, match="requires MindIE-SD"):
-                impl._forward_sparse_npu(None, None, None, None)
-    finally:
-        sys.modules.pop("mindiesd", None)
+            with pytest.raises(ValueError, match="explicitly support precision"):
+                impl._forward_sparse_npu(
+                    None, None, None, RainFusionPlan(prefix_len=0, used_len=8, latent_shape=[2, 2, 2])
+                )
 
 
 def test_precision_non_bf16_passes_gate_when_supported():
@@ -323,8 +333,7 @@ def test_precision_non_bf16_passes_gate_when_supported():
     import sys
 
     impl = make_impl(precision="mix")
-    sys.modules["mindiesd"] = _fake_mindiesd_module()
-    try:
+    with mock.patch.dict(sys.modules, {"mindiesd": _fake_mindiesd_module()}):
         with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=True):
             # q/k/v shapes: [B, S, N, D]; plan geometry must match S.
             q = torch.randn(1, 8, 4, 128)
@@ -332,5 +341,3 @@ def test_precision_non_bf16_passes_gate_when_supported():
             # The gate must pass; the fake mindiesd returns None so no crash.
             out = impl._forward_sparse_npu(q, q, q, plan)
             assert out is None
-    finally:
-        sys.modules.pop("mindiesd", None)

tests/diffusion/attention/test_ring_padding.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from dataclasses import dataclass
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends import ring_flash_attn, ring_pytorch_attn
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
+from vllm_omni.diffusion.attention.backends.ring.ring_utils import ring_kv_block_valid_length
+from vllm_omni.diffusion.attention.parallel import ring as ring_parallel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _SingleRankComm:
+    rank = 0
+    world_size = 1
+
+
+@dataclass(frozen=True)
+class _FakeSequenceParallelGroup:
+    ring_group: object
+
+
+def _fake_attention_result(q, _k, _v, **_kwargs):
+    return torch.zeros_like(q), torch.zeros(q.shape[0], q.shape[2], q.shape[1])
+
+
+@pytest.mark.parametrize(("block_rank", "expected"), [(0, 4), (1, 1), (2, 0)])
+def test_global_prefix_is_partitioned_across_ring_blocks(block_rank, expected):
+    assert ring_kv_block_valid_length(5, 4, block_rank, 3) == expected
+
+
+def test_absent_prefix_length_keeps_the_full_ring_block():
+    assert ring_kv_block_valid_length(None, 4, 2, 3) == 4
+
+
+def test_pytorch_ring_trims_padding_from_each_kv_block(monkeypatch):
+    captured: dict[str, Any] = {}
+
+    def fake_attention(q, k, v, **kwargs):
+        captured.update(k=k, v=v, **kwargs)
+        return _fake_attention_result(q, k, v, **kwargs)
+
+    monkeypatch.setattr(ring_pytorch_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", fake_attention)
+    qkv = torch.randn(1, 4, 2, 8)
+
+    output = ring_pytorch_attn.ring_pytorch_attn_func(
+        qkv,
+        qkv,
+        qkv,
+        group=object(),
+        valid_kv_length=3,
+    )
+
+    assert output.shape == qkv.shape
+    assert captured["k"].shape[1] == 3
+    assert captured["v"].shape[1] == 3
+
+
+def test_pytorch_ring_trims_circulated_blocks_by_source_rank(monkeypatch):
+    observed_lengths = []
+    # Each ring step sends K and V separately.
+    blocks = [torch.randn(1, 4, 2, 8) for _ in range(4)]
+
+    class FakeRingComm:
+        rank = 2
+        world_size = 3
+
+        def __init__(self, _group):
+            self._next_block = 0
+
+        def send_recv(self, _tensor):
+            block = blocks[self._next_block]
+            self._next_block += 1
+            return block
+
+        def commit(self):
+            pass
+
+        def wait(self):
+            pass
+
+    def fake_attention(q, k, v, **kwargs):
+        observed_lengths.append(k.shape[1])
+        return _fake_attention_result(q, k, v, **kwargs)
+
+    monkeypatch.setattr(ring_pytorch_attn, "RingComm", FakeRingComm)
+    monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", fake_attention)
+    qkv = torch.randn(1, 4, 2, 8)
+
+    ring_pytorch_attn.ring_pytorch_attn_func(
+        qkv,
+        qkv,
+        qkv,
+        group=object(),
+        valid_kv_length=5,
+    )
+
+    # Rank 2 starts with global block 2, then receives blocks 1 and 0.
+    # The first block is empty and therefore does not launch attention.
+    assert observed_lengths == [1, 4]
+
+
+def test_flash_ring_trims_padding_from_each_kv_block(monkeypatch):
+    captured: dict[str, Any] = {}
+
+    def fake_attention(q, k, v, **kwargs):
+        captured.update(k=k, v=v, **kwargs)
+        return _fake_attention_result(q, k, v, **kwargs)
+
+    monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(ring_flash_attn, "select_flash_attn_impl", lambda *_args, **_kwargs: fake_attention)
+    qkv = torch.randn(1, 4, 2, 8)
+
+    output, _ = ring_flash_attn.ring_flash_attn_forward(
+        object(),
+        qkv,
+        qkv,
+        qkv,
+        softmax_scale=0.25,
+        causal=False,
+        attn_type=AttnType.FA,
+        valid_kv_length=3,
+    )
+
+    assert output.shape == qkv.shape
+    assert captured["k"].shape[1] == 3
+    assert captured["v"].shape[1] == 3
+
+
+def test_ring_dispatch_forwards_global_valid_kv_length(monkeypatch):
+    captured: dict[str, Any] = {}
+
+    def fake_ring_attention(*_args, **kwargs):
+        captured.update(kwargs)
+        return "output"
+
+    monkeypatch.setattr(ring_pytorch_attn, "ring_pytorch_attn_func", fake_ring_attention)
+    strategy = ring_parallel.RingParallelAttention(
+        _FakeSequenceParallelGroup(ring_group=object()),
+        attn_backend_pref="sdpa",
+    )
+    metadata = AttentionMetadata(extra={"valid_kv_length": 3})
+    query = torch.randn(1, 4, 2, 8)
+
+    assert strategy.run_attention(query, query, query, metadata) == "output"
+    assert captured["valid_kv_length"] == 3
+
+
+def test_pytorch_ring_rejects_causal_with_a_valid_prefix(monkeypatch):
+    """Trimming a K/V block under causal=True would silently shift the diagonal.
+
+    The query block keeps the full padded length while the circulated K/V block
+    is trimmed, and a causal mask over a rectangular tile is bottom-right
+    aligned, so the combination must be refused rather than computed.
+    """
+    monkeypatch.setattr(ring_pytorch_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", _fake_attention_result)
+    qkv = torch.randn(1, 4, 2, 8)
+
+    with pytest.raises(ValueError, match="valid_kv_length is not supported with causal=True"):
+        ring_pytorch_attn.ring_pytorch_attn_func(
+            qkv,
+            qkv,
+            qkv,
+            causal=True,
+            group=object(),
+            valid_kv_length=3,
+        )
+
+
+def test_flash_ring_rejects_causal_with_a_valid_prefix(monkeypatch):
+    """Same refusal on the Flash-Attention ring path."""
+    monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(
+        ring_flash_attn,
+        "select_flash_attn_impl",
+        lambda *_args, **_kwargs: _fake_attention_result,
+    )
+    qkv = torch.randn(1, 4, 2, 8)
+
+    with pytest.raises(ValueError, match="valid_kv_length is not supported with causal=True"):
+        ring_flash_attn.ring_flash_attn_forward(
+            object(),
+            qkv,
+            qkv,
+            qkv,
+            softmax_scale=0.25,
+            causal=True,
+            attn_type=AttnType.FA,
+            valid_kv_length=3,
+        )
+
+
+def test_ring_paths_still_accept_causal_without_a_valid_prefix(monkeypatch):
+    """The guard is scoped to the padded case: plain causal Ring is unchanged."""
+    monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(
+        ring_flash_attn,
+        "select_flash_attn_impl",
+        lambda *_args, **_kwargs: _fake_attention_result,
+    )
+    qkv = torch.randn(1, 4, 2, 8)
+
+    output, _ = ring_flash_attn.ring_flash_attn_forward(
+        object(),
+        qkv,
+        qkv,
+        qkv,
+        softmax_scale=0.25,
+        causal=True,
+        attn_type=AttnType.FA,
+    )
+
+    assert output.shape == qkv.shape

tests/diffusion/cache/test_seacache.py

@@ -0,0 +1,635 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import datetime
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.cache.seacache import (
+    SeaCacheBackend,
+    SeaCacheConfig,
+    SeaCacheRootHook,
+    apply_sea_cache_hook,
+)
+from vllm_omni.diffusion.cache.seacache.sea_filter import (
+    apply_sea_filter,
+    extrapolate_residual,
+    indicator_distance,
+)
+from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.cache.teacache.extractors import EXTRACTOR_REGISTRY, CacheContext, get_extractor
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TinyCosmos3Transformer(torch.nn.Module):
+    """Small model that exposes a cache-neutral execution boundary."""
+
+    def _run_gen_layers(self, hidden_gen: torch.Tensor) -> torch.Tensor:
+        return hidden_gen * 2
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_ids: torch.Tensor | None = None,
+        text_mask: torch.Tensor | None = None,
+        video_shape: tuple[int, int, int] | None = None,
+        noisy_frame_mask: torch.Tensor | None = None,
+        control_latents: list[torch.Tensor] | torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del timestep, text_ids, text_mask, video_shape, noisy_frame_mask
+        if kwargs:
+            raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
+        controls = (
+            []
+            if control_latents is None
+            else [control_latents]
+            if isinstance(control_latents, torch.Tensor)
+            else list(control_latents)
+        )
+        inputs = [*controls, hidden_states]
+        gen_input = torch.cat(
+            [value.movedim(1, -1).flatten(1, 3) for value in inputs],
+            dim=1,
+        )
+        return self._run_gen_layers(gen_input)
+
+
+def _extract_tiny_cosmos3_context(
+    module: TinyCosmos3Transformer,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    text_ids: torch.Tensor | None = None,
+    text_mask: torch.Tensor | None = None,
+    video_shape: tuple[int, int, int] | None = None,
+    noisy_frame_mask: torch.Tensor | None = None,
+    control_latents: list[torch.Tensor] | torch.Tensor | None = None,
+    **kwargs,
+) -> CacheContext:
+    del timestep, text_ids, text_mask, video_shape
+    if kwargs:
+        raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
+    controls = (
+        []
+        if control_latents is None
+        else [control_latents]
+        if isinstance(control_latents, torch.Tensor)
+        else list(control_latents)
+    )
+    vision_items = [*controls, hidden_states]
+    gen_input = torch.cat(
+        [value.movedim(1, -1).flatten(1, 3) for value in vision_items],
+        dim=1,
+    )
+    return CacheContext(
+        modulated_input=gen_input,
+        hidden_states=gen_input,
+        encoder_hidden_states=None,
+        temb=torch.zeros_like(gen_input[:, 0]),
+        run_transformer_blocks=lambda: (module._run_gen_layers(gen_input),),
+        postprocess=lambda output: output,
+        extra_states={
+            "sea_cache_latents": vision_items,
+            "sea_cache_noisy_frame_mask": noisy_frame_mask,
+        },
+    )
+
+
+class Cosmos3OmniDiffusersPipeline:
+    def __init__(self) -> None:
+        self.transformer = TinyCosmos3Transformer()
+        self._current_step_index: int | None = None
+        self._current_sigma: float | None = None
+        self._num_timesteps: int | None = None
+
+    @property
+    def current_step_index(self) -> int | None:
+        return self._current_step_index
+
+    @property
+    def current_sigma(self) -> float | None:
+        return self._current_sigma
+
+    @property
+    def num_timesteps(self) -> int | None:
+        return self._num_timesteps
+
+
+def _latent(value: float) -> torch.Tensor:
+    return torch.full((1, 2, 2, 2, 2), value)
+
+
+def _run_step(
+    transformer: TinyCosmos3Transformer,
+    timestep: int,
+    value: float,
+    *,
+    hook: SeaCacheRootHook | None = None,
+    context: str = "cond",
+    control: float | None = None,
+    noisy_frame_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    controls = None if control is None else [_latent(control)]
+    cache_context = hook.cache_context(context) if hook is not None else nullcontext()
+    with torch.inference_mode(), cache_context:
+        return transformer(
+            hidden_states=_latent(value),
+            timestep=torch.tensor([timestep]),
+            noisy_frame_mask=noisy_frame_mask,
+            control_latents=controls,
+        )
+
+
+def _apply_test_hook(
+    transformer: TinyCosmos3Transformer,
+    metadata: SimpleNamespace,
+    config: SeaCacheConfig | None = None,
+) -> SeaCacheRootHook:
+    return apply_sea_cache_hook(
+        transformer,
+        config or SeaCacheConfig(threshold=100.0),
+        current_step_callback=lambda: metadata.step,
+        current_sigma_callback=lambda: metadata.sigma,
+        num_inference_steps_callback=lambda: metadata.num_steps,
+        extractor_fn=_extract_tiny_cosmos3_context,
+    )
+
+
+def test_config_validation() -> None:
+    assert SeaCacheConfig().threshold == 0.25
+    with pytest.raises(ValueError, match="residual_order"):
+        SeaCacheConfig(residual_order=-1)
+    with pytest.raises(ValueError, match="max_consecutive_cached"):
+        SeaCacheConfig(max_consecutive_cached=-1)
+
+
+def test_sea_filter_matches_reference_equation() -> None:
+    hidden = torch.randn(3, 4, 5, 2, dtype=torch.float32)
+    sigma = 0.4
+    power_exp = 3.0
+
+    spectrum = torch.fft.fftn(hidden, dim=(0, 1, 2))
+    gain = None
+    for axis in (0, 1, 2):
+        frequencies = torch.fft.fftfreq(hidden.shape[axis], dtype=torch.float32)
+        clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16)
+        axis_gain = (1.0 - sigma) * clean_power / ((1.0 - sigma) ** 2 * clean_power + sigma**2 + 1e-16)
+        shape = [1] * hidden.ndim
+        shape[axis] = hidden.shape[axis]
+        gain = axis_gain.reshape(shape) if gain is None else gain * axis_gain.reshape(shape)
+    assert gain is not None
+    gain = gain / gain.mean()
+    expected = torch.fft.ifftn(spectrum * gain, dim=(0, 1, 2)).real
+
+    torch.testing.assert_close(
+        apply_sea_filter(hidden, sigma=sigma, power_exp=power_exp),
+        expected,
+    )
+
+
+def test_indicator_distance_and_linear_extrapolation() -> None:
+    previous = [torch.ones(2, 2)]
+    current = [torch.full((2, 2), 1.5)]
+    assert indicator_distance(current, previous) == pytest.approx(0.5)
+    assert indicator_distance([torch.ones(3)], previous) == float("inf")
+
+    history = [
+        (0, torch.full((2, 2), 2.0)),
+        (2, torch.full((2, 2), 6.0)),
+    ]
+    torch.testing.assert_close(
+        extrapolate_residual(history, step=3, order=1),
+        torch.full((2, 2), 8.0),
+    )
+    torch.testing.assert_close(
+        extrapolate_residual(history, step=3, order=0),
+        torch.full((2, 2), 6.0),
+    )
+    quadratic_history = [
+        (0, torch.zeros(2, 2)),
+        (1, torch.ones(2, 2)),
+        (2, torch.full((2, 2), 4.0)),
+    ]
+    torch.testing.assert_close(
+        extrapolate_residual(quadratic_history, step=3, order=2),
+        torch.full((2, 2), 9.0),
+    )
+
+
+def test_hook_skips_middle_steps_and_forces_endpoints() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=4)
+    hook = _apply_test_hook(
+        transformer,
+        metadata,
+        SeaCacheConfig(threshold=100.0, max_consecutive_cached=2),
+    )
+    hook.refresh(transformer)
+
+    for step, timestep in enumerate((1000, 750, 500, 250)):
+        metadata.step = step
+        metadata.sigma = timestep / 1000
+        _run_step(transformer, timestep, 1.0 - step * 0.01, hook=hook)
+
+    assert hook.full_count == 2
+    assert hook.skip_count == 2
+    assert [step for step, _ in hook.state_manager._states["cond"].history] == [0, 3]
+
+
+def test_parameter_sharded_hook_skips_when_all_shard_ranks_agree(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook._parameter_sharded = True
+    fs_group = object()
+    reduced_groups: list[object] = []
+
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_fs_group",
+        lambda: SimpleNamespace(world_size=2, device_group=fs_group),
+    )
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
+
+    def all_reduce(decision, *, op, group):
+        assert op == torch.distributed.ReduceOp.MAX
+        reduced_groups.append(group)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
+
+    _run_step(transformer, 1000, 1.0, hook=hook)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.99, hook=hook)
+
+    assert hook.full_count == 1
+    assert hook.skip_count == 1
+    assert reduced_groups == [fs_group, fs_group]
+
+
+def test_parameter_sharded_peer_forces_full_compute(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook._parameter_sharded = True
+
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_fs_group",
+        lambda: SimpleNamespace(world_size=2, device_group=object()),
+    )
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
+    reduce_count = 0
+
+    def all_reduce(decision, *, op, group):
+        nonlocal reduce_count
+        del op, group
+        reduce_count += 1
+        if reduce_count == 2:
+            decision.fill_(1)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
+
+    _run_step(transformer, 1000, 1.0, hook=hook)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.99, hook=hook)
+
+    assert hook.full_count == 2
+    assert hook.skip_count == 0
+    assert hook.state_manager._states["cond"].accumulated_distance == 0.0
+
+
+def _uneven_cfg_sharded_worker(
+    rank: int,
+    init_method: str,
+    result_queue: torch.multiprocessing.Queue,
+) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    torch.distributed.init_process_group(
+        "gloo",
+        init_method=init_method,
+        rank=rank,
+        world_size=4,
+        timeout=datetime.timedelta(seconds=10),
+    )
+    try:
+        fs_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        sp_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        cfg_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 2], [1, 3])
+        ]
+        fs_group = next(group for ranks, group in fs_groups if rank in ranks)
+        sp_group = next(group for ranks, group in sp_groups if rank in ranks)
+        cfg_group = next(group for ranks, group in cfg_groups if rank in ranks)
+        parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
+        parallel_state._SP = SimpleNamespace(world_size=2, device_group=sp_group)
+
+        hook = SeaCacheRootHook(SeaCacheConfig())
+        hook._parameter_sharded = True
+        hook._synchronize_compute(False, torch.device("cpu"))
+        if rank in (0, 1):
+            hook._synchronize_compute(False, torch.device("cpu"))
+
+        gathered = [torch.zeros(1) for _ in range(2)]
+        torch.distributed.all_gather(
+            gathered,
+            torch.tensor([rank], dtype=torch.float32),
+            group=cfg_group,
+        )
+        result_queue.put((rank, [int(value.item()) for value in gathered]))
+    finally:
+        parallel_state._FS = None
+        parallel_state._SP = None
+        torch.distributed.destroy_process_group()
+
+
+def test_parameter_sharded_hook_handles_uneven_cfg_branch_dispatch() -> None:
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    result_queue = manager.Queue()
+    torch.multiprocessing.spawn(
+        _uneven_cfg_sharded_worker,
+        args=(get_distributed_init_method("seacache_uneven_cfg_"), result_queue),
+        nprocs=4,
+    )
+
+    results = sorted(result_queue.get(timeout=2) for _ in range(4))
+    assert results == [
+        (0, [0, 2]),
+        (1, [1, 3]),
+        (2, [0, 2]),
+        (3, [1, 3]),
+    ]
+
+
+def _hybrid_sp_worker(
+    rank: int,
+    init_method: str,
+    result_queue: torch.multiprocessing.Queue,
+) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    torch.distributed.init_process_group(
+        "gloo",
+        init_method=init_method,
+        rank=rank,
+        world_size=4,
+        timeout=datetime.timedelta(seconds=10),
+    )
+    try:
+        fs_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        ulysses_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        ring_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 2], [1, 3])
+        ]
+        sp_group = torch.distributed.new_group(
+            [0, 1, 2, 3],
+            timeout=datetime.timedelta(seconds=10),
+        )
+        fs_group = next(group for ranks, group in fs_groups if rank in ranks)
+        ulysses_group = next(group for ranks, group in ulysses_groups if rank in ranks)
+        ring_group = next(group for ranks, group in ring_groups if rank in ranks)
+        parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
+        parallel_state._SP = SimpleNamespace(
+            world_size=4,
+            device_group=sp_group,
+            ulysses_world_size=2,
+            ulysses_group=ulysses_group,
+            ring_world_size=2,
+            ring_group=ring_group,
+        )
+
+        compute = rank == 3
+        non_sharded_hook = SeaCacheRootHook(SeaCacheConfig())
+        non_sharded_result = non_sharded_hook._synchronize_compute(compute, torch.device("cpu"))
+
+        sharded_hook = SeaCacheRootHook(SeaCacheConfig())
+        sharded_hook._parameter_sharded = True
+        sharded_result = sharded_hook._synchronize_compute(compute, torch.device("cpu"))
+        result_queue.put((rank, non_sharded_result, sharded_result))
+    finally:
+        parallel_state._FS = None
+        parallel_state._SP = None
+        torch.distributed.destroy_process_group()
+
+
+def test_hook_synchronizes_full_hybrid_sequence_parallel_group() -> None:
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    result_queue = manager.Queue()
+    torch.multiprocessing.spawn(
+        _hybrid_sp_worker,
+        args=(get_distributed_init_method("seacache_hybrid_sp_"), result_queue),
+        nprocs=4,
+    )
+
+    results = sorted(result_queue.get(timeout=2) for _ in range(4))
+    assert results == [
+        (0, True, True),
+        (1, True, True),
+        (2, True, True),
+        (3, True, True),
+    ]
+
+
+def test_parameter_sharded_hook_fails_open_without_distributed_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    hook = SeaCacheRootHook(SeaCacheConfig())
+    hook._parameter_sharded = True
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+
+    assert hook._synchronize_compute(False, torch.device("cpu")) is True
+
+
+def test_hook_keeps_three_transfer_branches_separate() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+
+    contexts = ("cond", "cond_no_control", "uncond")
+    for step, (timestep, value) in enumerate(((1000, 1.0), (500, 0.9))):
+        metadata.step = step
+        metadata.sigma = timestep / 1000
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[0], control=0.5)
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[1])
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[2], control=0.5)
+
+    assert set(hook.state_manager._states) == set(contexts)
+    assert hook.full_count == 3
+    assert hook.skip_count == 3
+    assert len(hook.state_manager._states["cond"].previous_indicator) == 2
+    assert len(hook.state_manager._states["cond_no_control"].previous_indicator) == 1
+
+
+def test_hook_fails_open_without_noisy_vision() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    all_clean = torch.zeros(1, 1, 2, 1, 1)
+
+    _run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.9, hook=hook, noisy_frame_mask=all_clean)
+
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+
+
+def test_conditioning_only_model_failure_runs_forward_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    forward_calls = 0
+
+    def fail_forward(hidden_gen: torch.Tensor) -> torch.Tensor:
+        nonlocal forward_calls
+        forward_calls += 1
+        raise RuntimeError("model forward failed")
+
+    monkeypatch.setattr(transformer, "_run_gen_layers", fail_forward)
+    all_clean = torch.zeros(1, 1, 2, 1, 1)
+
+    with pytest.raises(RuntimeError, match="model forward failed"):
+        _run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
+
+    assert forward_calls == 1
+
+
+def test_hook_fails_open_without_explicit_context() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=2)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+
+    _run_step(transformer, 1000, 1.0)
+
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+    assert hook.state_manager._states == {}
+
+
+def test_hook_uses_exact_sigma_callback(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.cache.seacache import hook as hook_module
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=0.37, num_steps=2)
+    observed_sigmas: list[float] = []
+    original_filter = hook_module.apply_sea_filter
+
+    def recording_filter(hidden_states: torch.Tensor, sigma: float, power_exp: float) -> torch.Tensor:
+        observed_sigmas.append(sigma)
+        return original_filter(hidden_states, sigma, power_exp)
+
+    monkeypatch.setattr(hook_module, "apply_sea_filter", recording_filter)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    _run_step(transformer, timestep=999, value=1.0, hook=hook)
+
+    assert observed_sigmas
+    assert all(sigma == pytest.approx(0.37) for sigma in observed_sigmas)
+
+
+def test_backend_selector_and_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
+
+    class FSDPTinyCosmos3Transformer(TinyCosmos3Transformer):
+        pass
+
+    assert get_extractor(FSDPTinyCosmos3Transformer) is _extract_tiny_cosmos3_context
+    backend = get_cache_backend(
+        "sea_cache",
+        {
+            "sea_threshold": 0.4,
+            "sea_residual_order": 0,
+        },
+    )
+    assert isinstance(backend, SeaCacheBackend)
+    assert backend.config.sea_threshold == 0.4
+
+    pipeline = Cosmos3OmniDiffusersPipeline()
+    backend.enable(pipeline)
+    hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
+    assert isinstance(hook, SeaCacheRootHook)
+    assert callable(getattr(pipeline, "_cache_context_factory", None))
+    for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
+        assert not hasattr(pipeline.transformer, name)
+    pipeline._current_step_index = 0
+    pipeline._current_sigma = 1.0
+    pipeline._num_timesteps = 7
+    _run_step(pipeline.transformer, 1000, 1.0, hook=hook)
+    assert hook.full_count == 1
+
+    backend.refresh(pipeline, num_inference_steps=7)
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+    assert hook.state_manager._states == {}
+
+
+def test_backend_uses_resolved_pipeline_metadata_not_refresh_argument(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
+    pipeline = Cosmos3OmniDiffusersPipeline()
+    backend = SeaCacheBackend(DiffusionCacheConfig())
+    backend.enable(pipeline)
+    hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
+    assert isinstance(hook, SeaCacheRootHook)
+    backend.refresh(pipeline, num_inference_steps=35)
+    pipeline._num_timesteps = 50
+    pipeline._current_step_index = 0
+    pipeline._current_sigma = 1.0
+    _run_step(pipeline.transformer, 1000, 1.0, hook=hook)
+    assert hook.full_count == 1
+    assert hook.num_inference_steps_callback is not None
+    assert hook.num_inference_steps_callback() == 50
+
+    pipeline._current_step_index = 1
+    pipeline._current_sigma = 0.98
+    _run_step(pipeline.transformer, 980, 0.99, hook=hook)
+    assert hook.full_count == 1
+    assert hook.skip_count == 1
+
+
+def test_shared_config_defaults() -> None:
+    config = DiffusionCacheConfig()
+    assert config.sea_threshold == 0.25
+    assert config.sea_residual_order == 1
+    assert config.sea_max_consecutive_cached == 2
+    assert config.sea_power_exp == 3.0

tests/diffusion/cache/test_teacache_extractors.py

@@ -68,12 +68,12 @@ class BaseExtractorTest(ABC):
         pass
 
     @abstractmethod
-    def get_module(self):
+    def get_module(self, module, /):
         """Return model module instance."""
         pass
 
     @abstractmethod
-    def get_sample_inputs(self):
+    def get_sample_inputs(self, sample_inputs, /):
         """Return sample inputs for model."""
         pass
 
@@ -737,10 +737,10 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
         monkeypatch,
     ):
         """Strict SP must not mix local TeaCache state with gathered block rows."""
-        from vllm_omni.diffusion.attention.ops import minimax_h3_modulation
         from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
         from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
         from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.layers import indexed_modulation
 
         seq_len = sample_inputs["x"].shape[1]
         local_len = seq_len // 2
@@ -844,13 +844,13 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
         )
         monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
 
-        original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_
+        original_indexed_scale_shift = indexed_modulation.indexed_scale_shift_
 
         def capture_indexed_scale_shift(hidden, shift, scale, indices):
             captured_modulation_indices.append(indices.clone())
             return original_indexed_scale_shift(hidden, shift, scale, indices)
 
-        monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
+        monkeypatch.setattr(indexed_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
 
         local_img_pos = torch.tensor([0, 1])
         local_audio_pos = torch.tensor([2, 3])

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -325,7 +325,7 @@ def _runner(
         model_config=SimpleNamespace(max_model_len=max_len),
         scheduler_config=SimpleNamespace(
             max_num_batched_tokens=max_num_batched_tokens,
-            max_num_seqs=max_num_seqs,
+            max_num_seqs=max_num_seqs * max_rows_per_request,
         ),
         parallel_config=SimpleNamespace(
             decode_context_parallel_size=cp_size,
@@ -403,7 +403,7 @@ def test_initialize_failure_does_not_publish_partial_cache_state(
     successful_init_cache = model_runner_backend_module.init_kv_cache
     vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(max_model_len=16),
-        scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=1),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=2),
         parallel_config=SimpleNamespace(
             decode_context_parallel_size=1,
             cp_kv_cache_interleave_size=1,
@@ -445,7 +445,7 @@ def test_initialize_failure_does_not_publish_partial_cache_state(
     assert backend.block_tables is None
     assert backend.paged_attention_adapter is None
     assert backend._kv_cache_layer_adapters["layer-0"].kv_cache == "placeholder-0"
-    assert backend.vllm_config.scheduler_config.max_num_seqs == 1
+    assert backend.vllm_config.scheduler_config.max_num_seqs == 2
 
     monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", successful_init_cache)
     backend.initialize_kv_cache(config)

tests/diffusion/diffusion_kv/test_initialization.py

@@ -58,6 +58,7 @@ def _od_config(**overrides):
         gpu_memory_utilization=0.9,
         kv_cache_memory_bytes=None,
         max_num_seqs=1,
+        diffusion_kv_max_rows_per_request=1,
         max_num_batched_tokens=64,
         num_gpus=1,
     )

tests/diffusion/diffusion_kv/test_kv_connector.py

@@ -1,16 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from contextlib import nullcontext
+from types import SimpleNamespace
 from unittest import mock
 
 import pytest
 import torch
 from vllm.config import KVTransferConfig
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.outputs import KVConnectorOutput
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.kv_connector import (
+    KVTransferRegistrationError,
+    commit_kv_load,
     parse_kv_transfer_config,
     shutdown_kv_connector,
+    wait_for_kv_load,
 )
 from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
 from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
@@ -138,3 +145,205 @@ def test_shutdown_kv_connector_is_idempotent() -> None:
 
     assert ensure_shutdown.call_count == 2
     scheduler_connector.shutdown.assert_called_once()
+
+
+def test_commit_aligns_cfg_transport_pages_without_extending_computed_prefix() -> None:
+    requests = tuple(
+        SimpleNamespace(
+            request_id=f"req/{sequence_id}",
+            kv_transfer_params={"num_transfer_tokens": 9},
+            num_tokens=16,
+            num_computed_tokens=0,
+        )
+        for sequence_id in range(2)
+    )
+    manager = SimpleNamespace(
+        get_blocks=lambda _request_id: KVCacheBlocks(([0, 1, 2, 3],)),
+        kv_cache_config=SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=4))]),
+    )
+    connector = mock.Mock()
+
+    expected = commit_kv_load(connector, manager, requests, [7, 5])
+
+    assert expected == {"req/0", "req/1"}
+    calls = connector.update_state_after_alloc.call_args_list
+    assert [len(call.args[1].blocks[0]) for call in calls] == [3, 3]
+    assert [call.args[2] for call in calls] == [7, 5]
+    assert [request.num_computed_tokens for request in requests] == [7, 5]
+
+
+def test_wait_without_pending_load_uses_vllm_post_forward_once() -> None:
+    connector = mock.Mock()
+    output = mock.Mock(invalid_block_ids=set())
+    active_connector = mock.Mock(kv_connector=connector)
+    active_connector.post_forward.return_value = output
+    scheduler_output = SimpleNamespace(
+        kv_transfer_request_ids=set(),
+        finished_req_ids={"finished"},
+        kv_finished_request_ids={"finished/diffusion-kv/0", "finished/diffusion-kv/1"},
+    )
+
+    result = wait_for_kv_load(active_connector, scheduler_output, timeout=1.0)
+
+    assert result is output
+    active_connector.pre_forward.assert_called_once_with(scheduler_output)
+    active_connector.post_forward.assert_called_once_with({"finished/diffusion-kv/0", "finished/diffusion-kv/1"})
+    connector.get_finished.assert_not_called()
+
+
+def test_timeout_returns_partial_completion_without_releasing_pages():
+    active = mock.Mock()
+    active.kv_connector.get_finished.return_value = (set(), {"cfg0"})
+    active.kv_connector.get_block_ids_with_load_errors.return_value = set()
+    active.post_forward.return_value = KVConnectorOutput()
+    scheduled = SimpleNamespace(kv_transfer_request_ids={"cfg0", "cfg1"}, kv_finished_request_ids=set())
+    output = wait_for_kv_load(active, scheduled, timeout=0)
+    assert output.finished_recving == {"cfg0"}
+    active.kv_connector.request_finished.assert_not_called()
+    active.post_forward.assert_called_once_with(set())
+
+
+@pytest.fixture
+def cfg_registration():
+    requests = tuple(
+        SimpleNamespace(
+            request_id=f"cfg{i}",
+            status=None,
+            num_tokens=8,
+            kv_transfer_params={"num_transfer_tokens": 4, "transfer_id": "ticket", "do_remote_prefill": True},
+        )
+        for i in range(2)
+    )
+    manager = SimpleNamespace(
+        get_blocks=lambda _: KVCacheBlocks(([0, 1],)),
+        kv_cache_config=SimpleNamespace(kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=4))]),
+    )
+    return requests, manager
+
+
+def test_cfg_validation_is_atomic_before_connector_registration(cfg_registration):
+    requests, manager = cfg_registration
+    requests[1].kv_transfer_params["num_transfer_tokens"] = 9
+    connector = mock.Mock()
+    with pytest.raises(KVTransferRegistrationError, match="boundary"):
+        commit_kv_load(connector, manager, requests, [4, 4])
+    connector.update_state_after_alloc.assert_not_called()
+
+
+@pytest.mark.parametrize("fail_init", [False, True])
+def test_sp_connector_init_snapshots_topology_and_restores_model_tp(monkeypatch, fail_init):
+    import vllm.distributed.parallel_state as parallel_state
+
+    import vllm_omni.diffusion.diffusion_kv.kv_connector as module
+    import vllm_omni.diffusion.distributed.parallel_state as omni_parallel
+
+    tp = SimpleNamespace(world_size=1, rank_in_group=0)
+    sp = SimpleNamespace(world_size=2, rank_in_group=1)
+    monkeypatch.setattr(parallel_state, "_TP", tp)
+    monkeypatch.setattr(omni_parallel, "get_sp_group", lambda: sp)
+    snapshot: dict[str, int] = {}
+
+    def initialize(*args):
+        snapshot.update(
+            rank=parallel_state.get_tensor_model_parallel_rank(),
+            size=parallel_state.get_tensor_model_parallel_world_size(),
+        )
+        if fail_init:
+            raise RuntimeError("initialization failed")
+
+    monkeypatch.setattr(module, "ensure_kv_transfer_initialized", initialize)
+    config = SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(kv_connector="MooncakeConnector", engine_id="dit"),
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
+    )
+    with pytest.raises(RuntimeError, match="initialization failed") if fail_init else nullcontext():
+        module.init_worker_kv_connector(config, object())
+    assert parallel_state._TP is tp
+    assert snapshot == {"rank": 1, "size": 2}
+
+
+def test_partial_mooncake_registration_is_replaced_by_empty_notifications(cfg_registration):
+    from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
+        MooncakeConnector,
+        MooncakeConnectorScheduler,
+    )
+
+    scheduler = object.__new__(MooncakeConnectorScheduler)
+    scheduler.is_kv_producer = False
+    scheduler.is_kv_consumer = True
+    scheduler._reqs_need_recv = {}
+    connector = object.__new__(MooncakeConnector)
+    connector.connector_scheduler = scheduler
+    requests, manager = cfg_registration
+    for request in requests:
+        request.kv_transfer_params.update(remote_engine_id="ar", remote_bootstrap_addr="http://localhost:8998")
+
+    def register(request, blocks, count):
+        scheduler._reqs_need_recv[request.request_id] = (request, blocks)
+        request.kv_transfer_params["do_remote_prefill"] = False
+        if request.request_id == "cfg1":
+            raise RuntimeError("partial registration")
+
+    connector.update_state_after_alloc = register
+    with pytest.raises(KVTransferRegistrationError, match="partial registration"):
+        commit_kv_load(connector, manager, requests, [4, 4])
+    assert set(scheduler._reqs_need_recv) == {"cfg0", "cfg1"}
+    assert all(blocks == [] for _, blocks in scheduler._reqs_need_recv.values())
+    metadata = scheduler.build_connector_meta(None)
+    # Even partial registration failure sends every row in one per-rank batch,
+    # preserving the producer's complete fan-out count for the shared ticket.
+    assert set(metadata.reqs_to_recv["ar"]) == {"cfg0", "cfg1"}
+    assert all(
+        meta.transfer_id == "ticket" and meta.local_block_ids == [] for meta in metadata.reqs_to_recv["ar"].values()
+    )
+
+
+@pytest.mark.parametrize("timeout", [0, -1, float("inf"), "60"])
+def test_native_transfer_timeout_must_be_positive_finite(timeout):
+    with pytest.raises(ValueError, match="transfer_timeout"):
+        parse_kv_transfer_config(
+            {
+                **KV_TRANSFER_CONFIG,
+                "kv_connector_extra_config": {"transfer_timeout": timeout},
+            }
+        )
+
+
+def test_receive_completion_requires_all_ranks_across_polls():
+    from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+    from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
+
+    executor = SimpleNamespace(od_config=SimpleNamespace(num_gpus=2))
+    executor.collective_rpc = mock.Mock(
+        side_effect=[
+            [KVConnectorOutput(finished_recving={"cfg0", "cfg1"}), KVConnectorOutput(finished_recving={"cfg0"})],
+            [KVConnectorOutput(), KVConnectorOutput(finished_recving={"cfg1"})],
+        ]
+    )
+    scheduled = DiffusionSchedulerOutput(
+        step_id=0,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        finished_req_ids=set(),
+        num_running_reqs=0,
+        num_waiting_reqs=0,
+        kv_connector_metadata=object(),
+        kv_transfer_request_ids={"cfg0", "cfg1"},
+    )
+    result = DiffusionExecutor.prepare_kv_for_forward(executor, scheduled)
+    assert result.finished_recving == {"cfg0"}
+    scheduled.kv_transfer_request_ids = set()
+    result = DiffusionExecutor.prepare_kv_for_forward(executor, scheduled)
+    assert result.finished_recving == {"cfg1"}
+    assert not executor._kv_receive_completed_ranks
+
+
+def test_native_transfer_rejects_sleep_without_changing_legacy_sleep():
+    with pytest.raises(ValueError, match="registered pages must remain mapped"):
+        OmniDiffusionConfig.from_kwargs(
+            diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=2,
+            kv_transfer_config=dict(KV_TRANSFER_CONFIG),
+            enable_sleep_mode=True,
+        )
+    assert OmniDiffusionConfig.from_kwargs(enable_sleep_mode=True).enable_sleep_mode

tests/diffusion/diffusion_kv/test_metadata.py

@@ -61,13 +61,14 @@ class _MetadataScheduler(BaseScheduler):
 
 
 def test_scheduler_copies_opaque_params_to_sequence_requests() -> None:
-    params = {"transfer_id": "xfer-req-0", "do_remote_prefill": True}
+    params = {"transfer_id": "xfer-req-0", "do_remote_prefill": True, "num_transfer_tokens": 3}
     sequence = DiffusionKVRequest(
         "req-0/diffusion-kv/0",
         sequence_id=0,
         prefix_len=4,
         target_len=2,
         seq_len=6,
+        prompt_token_ids=[10, 11, 12, 13],
     )
     request = OmniDiffusionRequest(
         prompt="prompt",
@@ -83,6 +84,9 @@ def test_scheduler_copies_opaque_params_to_sequence_requests() -> None:
 
     assert state.diffusion_kv_requests[0].kv_transfer_params == params
     assert state.diffusion_kv_requests[0].kv_transfer_params is not params
+    assert sequence.prompt_token_ids == [10, 11, 12]
+    assert sequence.num_prompt_tokens == 3
+    assert sequence.seq_len == 6
 
 
 def test_diffusion_kv_metadata_uses_native_cache_group_block_ids() -> None:

tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise upstream Mooncake completion accounting without device/network I/O."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+
+import msgspec
+import pytest
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import mooncake_connector as mc
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.outputs import KVConnectorOutput
+
+from vllm_omni.diffusion.diffusion_kv.kv_connector import install_mooncake_cfg_fanout
+from vllm_omni.worker.mixins import OmniWorkerMixin
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture
+def producer(monkeypatch):
+    worker = object.__new__(mc.MooncakeConnectorWorker)
+    worker.shutdown = Mock()  # No hardware constructor, sockets or listener thread.
+    worker.is_kv_consumer = False
+    worker.tp_rank, worker.tp_size = 0, 1
+    worker.transfer_topo = SimpleNamespace(handshake_target_ranks=lambda size: list(range(size)))
+    worker.reqs_need_send, worker.finished_sending_reqs = {}, set()
+    for name in (
+        "kv_caches_base_addr",
+        "block_len_per_layer",
+        "kv_block_len_per_layer",
+        "registered_layer_names",
+        "registered_layer_indices",
+        "registered_group_indices",
+    ):
+        setattr(worker, name, [])
+    worker._get_transfer_regions = Mock(return_value=[])
+    worker._producer_cache_is_replicated = Mock(return_value=False)
+    worker._encoder = msgspec.msgpack.Encoder()
+    worker._sender_executor = None
+    worker._send_blocks = Mock(return_value=0)
+    worker.xfer_stats = Mock()
+    worker._build_transfer_params = AsyncMock(return_value=([1], [2], [16], [], None))
+    monkeypatch.setattr(mc, "_align_transfer_regions", lambda *_: ([], [], None))
+    monkeypatch.setattr(mc, "_validate_asymmetric_region_lengths", lambda **_: None)
+    monkeypatch.setattr(mc.envs, "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", 0.05)
+    connector = object.__new__(mc.MooncakeConnector)
+    connector.connector_worker = worker
+    return connector
+
+
+async def register(worker, ticket):
+    worker.sender_loop = asyncio.get_running_loop()
+    for blocks in ([], [[10]]):
+        meta = mc.MooncakeConnectorMetadata()
+        meta.add_new_req(ticket, blocks, {"transfer_id": ticket}, load_remote_cache=False)
+        await worker.record_send_reqs(meta)
+    return worker.reqs_need_send[ticket]
+
+
+def packet(rank, ranks, tickets, empty=False):
+    return mc.MooncakeXferMetadata(
+        remote_hostname="localhost",
+        remote_port=1234 + rank,
+        remote_tp_size=ranks,
+        remote_tp_rank=rank,
+        req_blocks={
+            f"{ticket}/{row}": (ticket, [] if empty else [[100 + rank * 10 + row]])
+            for ticket, rows in tickets.items()
+            for row in range(rows)
+        },
+        kv_caches_base_addr=[],
+        block_lens=[],
+        kv_block_lens=[],
+    )
+
+
+async def send(worker, rank, ranks, tickets, empty=False):
+    sock = SimpleNamespace(send_multipart=AsyncMock())
+    await worker.send_kv_to_decode(str(rank).encode(), sock, packet(rank, ranks, tickets, empty))
+    return [
+        msgspec.msgpack.decode(c.args[0][1], type=mc.MooncakeXferResponse) for c in sock.send_multipart.call_args_list
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("rows", [1, 2])
+@pytest.mark.parametrize("ranks", [2, 4])
+@pytest.mark.parametrize("in_flight", [False, True])
+async def test_source_pages_wait_for_every_rank_and_row(producer, rows, ranks, in_flight):
+    worker = producer.connector_worker
+    entered, release = asyncio.Event(), asyncio.Event()
+    original = worker._build_transfer_params
+
+    async def plan(ready, meta, *regions):
+        if in_flight and meta.remote_tp_rank == ranks - 1:
+            entered.set()
+            await release.wait()
+        return await original(ready, meta, *regions)
+
+    worker._build_transfer_params = plan
+    install_mooncake_cfg_fanout(producer)
+    adapted = worker._build_transfer_params
+    install_mooncake_cfg_fanout(producer)
+    assert worker._build_transfer_params is adapted  # Do not multiply twice.
+    ticket = await register(worker, "request")
+    pending = asyncio.create_task(send(worker, ranks - 1, ranks, {"request": rows})) if in_flight else None
+    try:
+        if pending is not None:
+            await asyncio.wait_for(entered.wait(), 2)
+        for rank in range(ranks - 1):
+            await send(worker, rank, ranks, {"request": rows})
+            assert ticket.need_send == ranks * rows
+            assert ticket.sent == (rank + 1) * rows
+            assert worker.reqs_need_send["request"] is ticket
+            finished = await worker.fetch_finished_sending_reqs()
+            scheduler = SimpleNamespace(connector=None, requests={"request": object()}, _free_blocks=Mock())
+            Scheduler._update_from_kv_xfer_finished(scheduler, KVConnectorOutput(finished_sending=finished))
+            scheduler._free_blocks.assert_not_called()
+        release.set()
+        responses = (
+            await asyncio.wait_for(pending, 2) if pending else await send(worker, ranks - 1, ranks, {"request": rows})
+        )
+        assert sum(len(response.ok_reqs or []) for response in responses) == rows
+        assert all(not response.err_reqs for response in responses)
+        assert "request" not in worker.reqs_need_send
+        finished = await worker.fetch_finished_sending_reqs()
+        assert finished == {"request"}
+        Scheduler._update_from_kv_xfer_finished(scheduler, KVConnectorOutput(finished_sending=finished))
+        scheduler._free_blocks.assert_called_once()
+        assert not await worker.fetch_finished_sending_reqs()
+    finally:
+        release.set()
+        if pending is not None:
+            await asyncio.gather(pending, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("empty", [False, True])
+@pytest.mark.parametrize("ranks", [1, 2])
+async def test_mixed_tickets_and_empty_block_notifications(producer, empty, ranks):
+    worker = producer.connector_worker
+    if empty:
+        worker._build_transfer_params.return_value = ([], [], [], [], None)
+    install_mooncake_cfg_fanout(producer)
+    single, cfg = await register(worker, "single"), await register(worker, "cfg")
+    await asyncio.gather(*(send(worker, rank, ranks, {"single": 1, "cfg": 2}, empty) for rank in range(ranks)))
+    assert (single.need_send, single.sent) == (ranks, ranks)
+    assert (cfg.need_send, cfg.sent) == (2 * ranks, 2 * ranks)
+    assert await worker.fetch_finished_sending_reqs() == {"single", "cfg"}
+    assert not worker.reqs_need_send
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failure", ["row_count", "transport"])
+async def test_failed_rank_does_not_report_source_complete(producer, failure):
+    worker = producer.connector_worker
+    install_mooncake_cfg_fanout(producer)
+    ticket = await register(worker, "request")
+    await send(worker, 0, 2, {"request": 2})
+    if failure == "transport":
+        worker._send_blocks.return_value = -1
+    responses = await send(worker, 1, 2, {"request": 1 if failure == "row_count" else 2})
+    assert any(response.err_reqs for response in responses)
+    assert ticket.sent == 2 and ticket.need_send == 4 and ticket.sending == 0
+    assert worker.reqs_need_send["request"] is ticket
+    assert not await worker.fetch_finished_sending_reqs()
+    # Preserve upstream timeout cleanup after all in-flight writes have stopped.
+    ticket.expire_time = 0
+    assert await worker.fetch_finished_sending_reqs() == {"request"}
+    assert not worker.reqs_need_send
+
+
+def test_worker_installs_adapter_after_native_initialization(producer, monkeypatch):
+    import vllm.distributed.kv_transfer as kv_transfer
+
+    worker = producer.connector_worker
+    original = worker._build_transfer_params
+
+    class UpstreamWorker:
+        def initialize_from_config(self, config):
+            assert config == "config"
+            assert worker._build_transfer_params is original
+
+    class OmniWorker(OmniWorkerMixin, UpstreamWorker):
+        pass
+
+    monkeypatch.setattr(kv_transfer, "has_kv_transfer_group", lambda: True)
+    monkeypatch.setattr(kv_transfer, "get_kv_transfer_group", lambda: producer)
+    object.__new__(OmniWorker).initialize_from_config("config")
+    assert worker._build_transfer_params is not original
+
+
+@pytest.mark.parametrize("role", ["other_connector", "scheduler", "consumer"])
+def test_adapter_leaves_other_roles_untouched(producer, role):
+    worker = producer.connector_worker
+    original = worker._build_transfer_params
+    if role == "scheduler":
+        producer.connector_worker = None
+    elif role == "consumer":
+        worker.is_kv_consumer = True
+    install_mooncake_cfg_fanout(object() if role == "other_connector" else producer)
+    assert worker._build_transfer_params is original

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -94,7 +94,7 @@ class _FakeAttentionGroup:
         return self.builder
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass(slots=True)
 class _FakeNativeMetadata:
     build_id: int
     causal: bool
@@ -102,6 +102,7 @@ class _FakeNativeMetadata:
     query_start_loc_cpu: torch.Tensor
     positions: torch.Tensor
     slot_mappings: torch.Tensor
+    max_num_splits: int = 0
 
 
 class _FakeLayer:
@@ -288,6 +289,8 @@ def test_omni_paged_backend_consumes_context_and_restores_diffusion_shape(
     assert output.shape == query.shape
     assert torch.equal(output, query)
     assert layer.calls[0][0].shape == (5, 2, 4)
+    assert layer.calls[0][1] is None
+    assert layer.calls[0][2] is None
     assert layer.native_events == ["update", "forward"]
     assert layer.calls[0][3] is events[0][2]
 
@@ -440,9 +443,7 @@ def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch:
 
 def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, layer, events = _make_adapter(monkeypatch)
-    # Exercise the Ascend-only homogeneous batch path. CUDA keeps its native
-    # output-buffer contract and is covered by the GPU adapter tests.
-    monkeypatch.setattr(adapter_module.current_omni_platform, "is_npu", lambda: True)
+    layer.impl.vllm_flash_attn_version = 2
     batch = adapter.prepare_batch(
         [
             DiffusionPagedAttentionRow(
@@ -499,10 +500,10 @@ def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.
     assert torch.equal(layer.updates[0][0], key.reshape(12, 2, 4))
     assert torch.equal(layer.updates[0][1], value.reshape(12, 2, 4))
     assert [call[0].shape[0] for call in layer.calls] == [4, 6, 2]
-    assert [call[1].shape[0] for call in layer.calls] == [4, 6, 2]
-    assert [call[2].shape[0] for call in layer.calls] == [4, 6, 2]
+    assert all(call[1] is None and call[2] is None for call in layer.calls)
     assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True))
     assert [metadata.build_id for metadata in context.piecewise_native_metadata] == [1, 2, 3]
+    assert [metadata.max_num_splits for metadata in context.piecewise_native_metadata] == [1, 1, 1]
 
     # The first metadata build is the normal whole-query path. Piecewise calls
     # then use causal [3, 5), full [5, 8), and causal [8, 9) segments.
@@ -539,6 +540,51 @@ def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.
     ]
 
 
+def test_homogeneous_piecewise_packs_strided_query_only_per_segment(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    adapter, _, _, _ = _make_adapter(monkeypatch)
+    batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=6,
+                seq_len=9,
+                kv_start_pos=3,
+            ),
+            DiffusionPagedAttentionRow(
+                request_id="req-1",
+                sequence_id=1,
+                query_len=6,
+                seq_len=9,
+                kv_start_pos=3,
+            ),
+        ]
+    )
+    packed_qkv = torch.randn(2, 6, 3, 2, 4)
+    query, key, value = packed_qkv.unbind(dim=2)
+    assert not query.is_contiguous()
+    metadata = SimpleNamespace(
+        attn_mask=None,
+        full_attn_spans=[[(5, 8)], [(5, 8)]],
+        query_ranges=None,
+        extra={},
+    )
+
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context(
+            "layer-0",
+            query,
+            key,
+            value,
+            omni_attn_metadata=metadata,
+        )
+
+    assert not context.query.is_contiguous()
+    assert context.query.data_ptr() == query.data_ptr()
+
+
 def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None:
     adapter, _, _, _ = _make_adapter(monkeypatch)
     shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32)
@@ -1232,7 +1278,7 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         ulysses_degree=2,
     )
 
-    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 2)]
+    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 1)]
     assert specialized_backends == [(native_backend, 2)]
     assert config.attention_config.backend is None
     assert config.attention_config.backend_per_kind is original_backend_per_kind

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

@@ -43,8 +43,7 @@ class _SmokeDiffusionAttention(nn.Module):
         self.softmax_scale = _HEAD_SIZE**-0.5
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
-def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
+def _native_adapter(num_rows: int = 1):
     device = torch.device("cuda", torch.accelerator.current_device_index())
     vllm_config = VllmConfig(
         cache_config=CacheConfig(
@@ -89,8 +88,8 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
 
     canonical_spec = native_layer.spec
     kv_cache_config = KVCacheConfig(
-        num_blocks=4,
-        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME])],
+        num_blocks=4 * num_rows,
+        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4 * num_rows, [_LAYER_NAME])],
         kv_cache_groups=[
             KVCacheGroupSpec(
                 layer_names=[_LAYER_NAME],
@@ -120,14 +119,15 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
         )
         block_tables = BlockTables(
             block_sizes=[_BLOCK_SIZE],
-            max_num_reqs=1,
-            max_num_batched_tokens=32,
+            max_num_reqs=num_rows,
+            max_num_batched_tokens=32 * num_rows,
             max_num_blocks_per_group=[2],
             device=device,
             kernel_block_sizes=kernel_block_sizes,
         )
 
-    block_tables.append_block_ids(0, ([3, 1],), overwrite=True)
+    for row in range(num_rows):
+        block_tables.append_block_ids(row, ([4 * row + 3, 4 * row + 1],), overwrite=True)
     block_tables.apply_staged_writes()
     adapter = DiffusionPagedAttentionAdapter(
         vllm_config=vllm_config,
@@ -136,11 +136,17 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
         block_tables=block_tables,
         attn_groups=attn_groups,
         layers={_LAYER_NAME: native_layer},
-        resolve_row=lambda _request_id, _sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
-            row_index=0,
+        resolve_row=lambda _request_id, sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
+            row_index=sequence_id,
             max_seq_len=19,
         ),
     )
+    return adapter, diffusion_layer, device
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
+def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
+    adapter, diffusion_layer, device = _native_adapter()
     prefix_batch = adapter.prepare_batch(
         [
             DiffusionPagedAttentionRow(
@@ -258,3 +264,48 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
     torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2)
     torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2)
     torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
+def test_hunyuan_homogeneous_cfg_paged_attention_graph_replay() -> None:
+    """Real paged kernels, strided QKV, mixed image spans and changed replay inputs."""
+    adapter, layer, device = _native_adapter(num_rows=2)
+    packed = torch.randn(2, 19, 3, _NUM_HEADS, _HEAD_SIZE, device=device, dtype=torch.float16)
+    query, key, value = packed.unbind(dim=2)
+    backend = FlashAttentionImpl(
+        num_heads=_NUM_HEADS, head_size=_HEAD_SIZE, softmax_scale=layer.softmax_scale, num_kv_heads=_NUM_HEADS
+    )
+    batch = adapter.prepare_batch(
+        [DiffusionPagedAttentionRow(request_id="cfg", sequence_id=i, query_len=19, seq_len=19) for i in range(2)]
+    )
+    metadata = AttentionMetadata(full_attn_spans=[[(5, 10)], [(5, 10)]])
+    with adapter.activate(batch):
+        context = adapter.prepare_layer_context(_LAYER_NAME, query, key, value, omni_attn_metadata=metadata)
+    assert context.piecewise_plan.homogeneous_batch_shape == (2, 19)
+    warmup = torch.cuda.Stream()
+    warmup.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(warmup):
+        for _ in range(3):
+            backend.forward_paged(context)
+    torch.cuda.current_stream().wait_stream(warmup)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = backend.forward_paged(context)
+    output_ptr = captured.data_ptr()
+    mask = torch.ones(19, 19, device=device, dtype=torch.bool).tril()
+    mask[5:10, 5:10] = True
+    for _ in range(2):
+        packed.copy_(torch.randn_like(packed))
+        graph.replay()
+        actual = captured.clone()
+        eager = backend.forward_paged(context)
+        reference = F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            attn_mask=mask,
+            scale=layer.softmax_scale,
+        ).transpose(1, 2)
+        torch.testing.assert_close(actual.reshape_as(reference), reference, rtol=2e-2, atol=2e-2)
+        torch.testing.assert_close(actual, eager, rtol=2e-2, atol=2e-2)
+        assert captured.data_ptr() == output_ptr

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -176,10 +176,11 @@ def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
     profile_request = object()
 
     assert worker.determine_available_kv_memory(profile_request) == [4096]
-    worker.model_runner.profile_run.assert_called_once_with(profile_request)
+    worker.model_runner.profile_run.assert_not_called()
 
 
 def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None:
+    monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", SimpleNamespace(empty_cache=Mock()))
     worker = object.__new__(DiffusionWorker)
     worker.rank = 0
     worker.init_snapshot = SimpleNamespace(free_memory=900)
@@ -209,6 +210,7 @@ def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch
 
 @pytest.mark.parametrize("non_kv_cache_memory", [750, 800])
 def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None:
+    monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", SimpleNamespace(empty_cache=Mock()))
     worker = object.__new__(DiffusionWorker)
     worker.rank = 0
     worker.init_snapshot = SimpleNamespace(free_memory=900)
@@ -245,6 +247,7 @@ def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cach
 
 
 def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None:
+    monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", SimpleNamespace(empty_cache=Mock()))
     worker = object.__new__(DiffusionWorker)
     worker.rank = 0
     worker.init_snapshot = SimpleNamespace(free_memory=1200)
@@ -310,3 +313,31 @@ def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None:
         )
 
     assert gathered == [(False, "ValueError: local failure")]
+
+
+def test_hunyuan_native_layer_identity_matches_ar_spec(monkeypatch):
+    import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer as hy3
+
+    monkeypatch.setattr(hy3, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(hy3, "get_sequence_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(hy3, "get_allgather_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(hy3, "get_sequence_parallel_rank", lambda: 0)
+    monkeypatch.setattr(hy3, "QKVParallelLinear", lambda **_: nn.Identity())
+    monkeypatch.setattr(hy3, "RowParallelLinear", lambda **_: nn.Identity())
+    monkeypatch.setattr(hy3, "get_rope", lambda **_: nn.Identity())
+
+    # Only backend construction is fake; execute both HY3 constructors and the
+    # real runner spec-discovery path, including the actual Attention prefix.
+    def attention(**kwargs):
+        return _attention(enabled=kwargs.get("paged_kv_cache_role") is not None, prefix=kwargs["prefix"])
+
+    monkeypatch.setattr(hy3, "Attention", attention)
+    model = hy3.HunYuanAttention(
+        config=SimpleNamespace(num_key_value_heads=2, attention_head_dim=8),
+        hidden_size=16,
+        num_heads=2,
+        num_kv_heads=2,
+        prefix="layers.3.self_attn",
+    )
+    runner = _runner(model.image_attn.attn)
+    assert set(runner.get_kv_cache_spec()) == {"model.layers.3.self_attn.attn"}

tests/diffusion/distributed/test_cfg_parallel.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
 """Tests for CFG (Classifier-Free Guidance) parallel functionality.
 
 CPU tests mock the CFG process group and validate mixin logic without GPUs.
@@ -32,8 +32,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.platforms import current_omni_platform
 
-_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
-_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
+_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2)
+_L4_THREE_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=3)
 
 
 def _set_random_seeds(seed: int) -> None:
@@ -1009,7 +1009,7 @@ def _run_multi_branch_gpu_parity(
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=2)
 @pytest.mark.parametrize("batch_size", [2])
 @pytest.mark.parametrize("cfg_normalize", [False, True])
 @pytest.mark.parametrize("extra_kwargs", [None, {"step_i": 1}])

tests/diffusion/distributed/test_comm.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
 """Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives.
 
 CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same
@@ -24,8 +24,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.platforms import current_omni_platform
 
-_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
-_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=4)
 
 DeviceKind = Literal["cpu", "cuda"]
 
@@ -461,3 +461,46 @@ def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
         device_kind="cuda",
         master_port=29501,
     )
+
+
+def _run_fused_qkv_matches_three_4d(
+    local_rank: int,
+    world_size: int,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    """LingBot's fused (B, S/N, 3, H, D) exchange must equal three per-tensor 4D exchanges."""
+    from vllm_omni.diffusion.distributed.comm import all_to_all_5D
+
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
+    _init_worker(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(ulysses_degree=world_size)
+    sp_group = get_sp_group().ulysses_group
+    try:
+        batch, seq_per_rank, heads, head_size = 1, 6, 4 * world_size, 8
+        # Distinguishable values: rank, q/k/v index, position, head, feature.
+        base = torch.arange(seq_per_rank * heads * head_size, dtype=torch.float32, device=device).reshape(
+            batch, seq_per_rank, heads, head_size
+        )
+        tensors = [base + 1000.0 * local_rank + 100000.0 * which for which in range(3)]
+        separate = [SeqAllToAll4D.apply(sp_group, t, 2, 1, False) for t in tensors]
+        fused = all_to_all_5D(torch.stack(tensors, dim=2), scatter_idx=3, gather_idx=1, group=sp_group)
+        assert fused.shape == (batch, seq_per_rank * world_size, 3, heads // world_size, head_size)
+        for which, expected in enumerate(separate):
+            assert torch.equal(fused.unbind(2)[which], expected), f"q/k/v index {which} differs on rank {local_rank}"
+    finally:
+        destroy_distributed_env()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
+    torch.multiprocessing.spawn(
+        _run_fused_qkv_matches_three_4d,
+        args=(world_size, "cpu", 29613),
+        nprocs=world_size,
+    )

tests/diffusion/distributed/test_expert_parallel.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
 
 """
 System test for Expert Parallel (EP).

tests/diffusion/distributed/test_flux2_vae_patch_parallel.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
 
 """FLUX.2-dev VAE patch-parallel correctness tests.
 

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for SP subgroup construction in parallel_state.py."""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for SP subgroup construction for diffusion parallel state."""
 
 from __future__ import annotations
 
@@ -8,12 +8,55 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+import vllm.distributed.parallel_state as vllm_parallel_state
 
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.distributed import parallel_state as omni_parallel_state
 from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.cpu]
 
 
+def ensure_parallel_state_initialized():
+    """Ensure torch distributed is initialized and that all world / count vars are set."""
+    assert torch.distributed.is_initialized()
+    assert omni_parallel_state._WORLD is not None
+    assert vllm_parallel_state._WORLD is not None
+    assert vllm_parallel_state._NODE_COUNT is not None
+
+
+def ensure_parallel_state_not_initialized():
+    """Ensure torch distributed is not initialized and that all world / count vars are unset."""
+    assert not torch.distributed.is_initialized()
+    assert omni_parallel_state._WORLD is None
+    assert vllm_parallel_state._WORLD is None
+    assert vllm_parallel_state._NODE_COUNT is None
+
+
+def test_omni_manages_vllm_distributed_state(monkeypatch):
+    """Verify Omni initializes and tears down vLLM's distributed state properly.
+
+    This is a regression test for ensuring vLLM Omni also sets up vLLM's vars correctly
+    to avoid potential misalignment in cases where Omni uses vLLM's native coordinator,
+    e.g., MoE + diffusion.
+    """
+    ensure_parallel_state_not_initialized()
+    monkeypatch.setattr(omni_parallel_state.current_omni_platform, "get_device_count", lambda: 1)
+    monkeypatch.setattr(omni_parallel_state.current_omni_platform, "set_device", lambda _device: None)
+    omni_parallel_state.init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method=get_distributed_init_method(),
+        backend="gloo",
+    )
+    try:
+        ensure_parallel_state_initialized()
+    finally:
+        omni_parallel_state.destroy_distributed_env()
+    ensure_parallel_state_not_initialized()
+
+
 def _fake_new_group_factory(created_groups: list[SimpleNamespace]):
     def _fake_new_group(ranks, *args, **kwargs):
         group = SimpleNamespace(ranks=list(ranks))

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -31,9 +31,9 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.platforms import current_omni_platform
 
-_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
-_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
-_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2)
+_L4_THREE_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=3)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=4)
 
 DeviceKind = Literal["cpu", "cuda"]
 _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

tests/diffusion/distributed/test_sana_video_sp_numeric.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
 """Real multi-process numeric equivalence gate for SANA-Video sequence parallel.
 
 Runs the native transformer directly (no pipeline, no topology validation) with
@@ -273,7 +273,7 @@ def _baseline_worker(
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=2)
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 @pytest.mark.parametrize("num_frames", [21, 5])
 def test_sp2_matches_sp1_fp32(task: str, num_frames: int) -> None:
@@ -290,7 +290,7 @@ def test_sp2_matches_sp1_fp32(task: str, num_frames: int) -> None:
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=4)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4)
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 @pytest.mark.parametrize("num_frames", [21, 5])
 def test_sp4_matches_sp1_fp32(task: str, num_frames: int) -> None:
@@ -307,7 +307,7 @@ def test_sp4_matches_sp1_fp32(task: str, num_frames: int) -> None:
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=4)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4)
 def test_sp2_cfg2_matches_sp1_fp32() -> None:
     """Four processes: CFG branches with distinct inputs each run their own SP
     group; catches CFG/SP group crossover and deadlocks."""
@@ -324,7 +324,7 @@ def test_sp2_cfg2_matches_sp1_fp32() -> None:
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=4)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4)
 def test_tp2_sp2_matches_sp1_fp32() -> None:
     """Four processes: TP shards heads and channels while SP shards frames; the
     combined mesh must still reproduce the serial output."""
@@ -342,7 +342,7 @@ def test_tp2_sp2_matches_sp1_fp32() -> None:
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=2)
 def test_sp2_bf16_drift_informational() -> None:
     """Records the bf16 reduction-reorder drift magnitude; the load-bearing
     numeric gate is the fp32 suite above."""

tests/diffusion/distributed/test_sequence_parallel.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
 
 """
 System test for Sequence Parallel (SP) backends: Ulysses and Ring attention.

tests/diffusion/distributed/test_vae_decode_parallelism.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
 
 """VAE patch-parallel decode parity tests for diffusion models.
 
@@ -145,7 +145,7 @@ def _run_generate(
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 4})
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards={"cuda": 4})
 @pytest.mark.parametrize("model_case", MODEL_CASES)
 def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
     if current_omni_platform.is_npu():

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import os
 from types import SimpleNamespace
 
@@ -353,6 +356,94 @@ def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.Mon
         wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
 
 
+def _reference_conv_input(conv, x, cache_x):
+    """The conv input the previous implementation built: cat, F.pad, then halo_exchange."""
+    padding = list(conv._padding)
+    if cache_x is not None and padding[4] > 0:
+        x = torch.cat([cache_x, x], dim=2)
+        padding[4] -= cache_x.shape[2]
+    x = torch.nn.functional.pad(x, padding)
+    x_padded, _, _ = wan_spatial_shard.halo_exchange(
+        x, group=object(), halo_size=conv.halo_size, split_dim=conv.split_dim
+    )
+    return x_padded, x.shape[conv.split_tensor_dim]
+
+
+def _dist_conv(split_dim, kernel=(3, 3, 3), padding=(1, 1, 1)):
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d
+
+    torch.manual_seed(0)
+    source = WanCausalConv3d(3, 4, kernel, padding=padding)
+    return source, wan_spatial_shard.WanDistCausalConv3d(source, object(), split_dim=split_dim)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("split_dim", ["width", "height"])
+@pytest.mark.parametrize("with_cache", [False, True])
+def test_single_pass_conv_input_matches_cat_pad_halo_on_one_rank(monkeypatch, split_dim, with_cache):
+    """One rank, halo-bearing kernel: same bytes as the previous cat -> F.pad -> halo_exchange path."""
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1))
+    _, conv = _dist_conv(split_dim)
+    x = torch.randn(1, 3, 2, 6, 8)
+    cache_x = torch.randn(1, 3, 2, 6, 8) if with_cache else None
+
+    expected, expected_extent = _reference_conv_input(conv, x, cache_x)
+    got, got_extent = conv._assemble_input(x, cache_x)
+
+    assert got.shape == expected.shape and got_extent == expected_extent
+    assert torch.equal(got, expected)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_single_pass_conv_rejects_cache_longer_than_causal_padding():
+    _, conv = _dist_conv("width")
+    with pytest.raises(AssertionError, match="temporal cache exceeds"):
+        conv._assemble_input(torch.randn(1, 3, 2, 6, 8), torch.randn(1, 3, 3, 6, 8))
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("split_dim", ["width", "height"])
+def test_single_pass_conv_without_a_halo_matches_the_stock_causal_conv(monkeypatch, split_dim):
+    """A kernel with no spatial extent along the split (the temporal conv) needs no halo and equals the stock conv."""
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1))
+    source, conv = _dist_conv(split_dim, kernel=(3, 1, 1), padding=(1, 0, 0))
+    x = torch.randn(1, 3, 2, 6, 8)
+    cache_x = torch.randn(1, 3, 2, 6, 8)
+
+    assert torch.equal(conv(x, cache_x), source(x, cache_x))
+    assert torch.equal(conv(x, None), source(x, None))
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("split_dim", ["width", "height"])
+@pytest.mark.parametrize("rank", [0, 1])
+@pytest.mark.parametrize("kernel,padding", [((3, 3, 3), (1, 1, 1)), ((3, 1, 1), (0, 0, 0))])
+def test_single_pass_conv_input_matches_cat_pad_halo_on_two_ranks(monkeypatch, split_dim, rank, kernel, padding):
+    """Same bytes as cat -> F.pad -> halo_exchange, halos included, on either rank of a pair."""
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2))
+
+    def fake_p2p(*, rank, world_size, group, top_row_ref, bottom_row_ref, recv_top_buf, recv_bottom_buf):
+        # A neighbour's halo is whatever it sent; stand in with a marker derived from our own edge so
+        # both the reference and the single-pass path see the same exchange. Global edges stay zero.
+        recv_top_buf.copy_(top_row_ref * 0 + 7.0) if rank > 0 else recv_top_buf.zero_()
+        recv_bottom_buf.copy_(bottom_row_ref * 0 + 9.0) if rank < world_size - 1 else recv_bottom_buf.zero_()
+
+    monkeypatch.setattr(wan_spatial_shard, "_halo_exchange_p2p", fake_p2p)
+    _, conv = _dist_conv(split_dim, kernel=kernel, padding=padding)
+    x = torch.randn(1, 3, 1, 6, 8)
+    cache_x = torch.randn(1, 3, 2, 6, 8)
+
+    expected, expected_extent = _reference_conv_input(conv, x, cache_x)
+    got, got_extent = conv._assemble_input(x, cache_x)
+
+    assert got.shape == expected.shape and got_extent == expected_extent
+    assert torch.equal(got, expected)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
@@ -538,7 +629,7 @@ def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "H100"}, num_cards=_SPATIAL_SHARD_WORLD_SIZE)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=_SPATIAL_SHARD_WORLD_SIZE)
 @pytest.mark.parametrize("split_dim", ["height", "width"])
 def test_spatial_shard_decode_matches_reference(split_dim: str):
     manager = mp.get_context("spawn").Manager()

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -41,13 +41,13 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     initialize_model_parallel,
     model_parallel_is_initialized,
 )
-from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import (
+from vllm_omni.diffusion.models.bagel.mot.mot_qkv_parallel_linear import (
     MoTQKVParallelLinear,
 )
-from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import (
+from vllm_omni.diffusion.models.bagel.mot.mot_row_parallel_linear import (
     MoTRowParallelLinear,
 )
-from vllm_omni.diffusion.layers.mot.ops.mot_gemm import (
+from vllm_omni.diffusion.models.bagel.mot.ops.mot_gemm import (
     get_best_mot_config,
 )
 

tests/diffusion/kernels/mot/test_mot_norm.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803, E741
 """Layer-level correctness & performance test for MoTRMSNorm.
 
@@ -20,7 +23,7 @@ import torch
 from vllm import ir
 from vllm.config import VllmConfig, set_current_vllm_config
 
-from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 

tests/diffusion/kernels/mot/test_mot_norm_npu.py

@@ -1,8 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 import torch
 
 import vllm_omni.diffusion.layers.custom_op as custom_op
-from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 

tests/diffusion/layers/test_custom_op.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 the CustomOp platform-dispatch base class."""
 
 import pytest
@@ -39,7 +39,7 @@ def test_forward_cuda_still_requires_an_implementation():
 
 def test_mot_rmsnorm_dispatches_on_xpu_without_an_override():
     """MoTRMSNorm relies on the inherited fallback (regression test)."""
-    from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+    from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
     assert "forward_xpu" not in MoTRMSNorm.__dict__
 

tests/diffusion/layers/test_fused_qk_norm_rope.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
@@ -80,3 +80,113 @@ def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
 
     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)
+
+
+# ---------------------------------------------------------------------------
+# General geometry (any even head_dim, here Boogu-Image's 120) in both
+# pairing modes, against the module's own eager reference at the same
+# tolerance as the MiniMax-H3 test above.
+# ---------------------------------------------------------------------------
+
+_BOOGU_HEAD_DIM = 120
+
+
+def _boogu_inputs(strided: bool):
+    torch.manual_seed(11)
+    if strided:
+        # Slices of a wider head axis: the op must honour q/k strides
+        # (merged-QKV projections hand the op such views).
+        q = torch.randn(4139, 35, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)[:, :28]
+        k = torch.randn(4139, 35, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)[:, :7]
+    else:
+        q = torch.randn(4139, 28, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+        k = torch.randn(4139, 7, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    q_weight = torch.randn(_BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(_BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    freqs = torch.randn(4139, _BOOGU_HEAD_DIM // 2, device="cuda", dtype=torch.float32)
+    rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1)
+    return q, k, q_weight, k_weight, rope_table
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize("strided", [False, True])
+def test_fused_qk_norm_rope_interleaved(strided):
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+        _eager_qk_norm_rope,
+        _launch_fused_qk_norm_rope,
+        fused_qk_norm_rope,
+    )
+
+    q, k, q_weight, k_weight, rope_table = _boogu_inputs(strided)
+    expected = _eager_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, _BOOGU_HEAD_DIM, _BOOGU_HEAD_DIM, True)
+    # Check the public op AND the launcher directly: the latter cannot fall
+    # back to eager, so a silent dispatch regression cannot go green here.
+    for actual in (
+        fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=True),
+        _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=True),
+    ):
+        torch.testing.assert_close(actual[0], expected[0], atol=0.0625, rtol=0.02)
+        torch.testing.assert_close(actual[1], expected[1], atol=0.0625, rtol=0.02)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_fused_qk_norm_rope_half_split_general_dim():
+    """head_dim=120 off the MiniMax-H3 128 pin, half-split pairing.
+
+    Exercises the generalised combined kernel's half-split mode via the
+    launcher. It is not routed in production (half-split traffic keeps the
+    untouched pre-existing per-tensor kernel and its 128/96 contract, so the
+    public op falls back to eager at this geometry), pending the
+    maintainers' call.
+    """
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+        _eager_qk_norm_rope,
+        _launch_fused_qk_norm_rope,
+    )
+
+    q, k, q_weight, k_weight, rope_table = _boogu_inputs(strided=False)
+    expected = _eager_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, _BOOGU_HEAD_DIM, _BOOGU_HEAD_DIM)
+    actual = _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=False)
+    torch.testing.assert_close(actual[0], expected[0], atol=0.0625, rtol=0.02)
+    torch.testing.assert_close(actual[1], expected[1], atol=0.0625, rtol=0.02)
+
+
+def test_fused_qk_norm_rope_min_tokens_resolution(monkeypatch):
+    """The op-level token-gate override: unset or blank -> caller's default; a
+    non-negative integer string overrides it (``0`` = always fuse); anything
+    else is rejected naming the variable."""
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope_min_tokens
+
+    env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+    monkeypatch.delenv(env, raising=False)
+    assert fused_qk_norm_rope_min_tokens(2048) == 2048
+    for blank in ("", "  "):
+        monkeypatch.setenv(env, blank)
+        assert fused_qk_norm_rope_min_tokens(2048) == 2048
+    monkeypatch.setenv(env, "0")
+    assert fused_qk_norm_rope_min_tokens(2048) == 0
+    monkeypatch.setenv(env, "4096")
+    assert fused_qk_norm_rope_min_tokens(2048) == 4096
+    for bad in ("-1", "abc"):
+        monkeypatch.setenv(env, bad)
+        with pytest.raises(ValueError, match=env):
+            fused_qk_norm_rope_min_tokens(2048)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or not HAS_TRITON, reason="CUDA and Triton required")
+@pytest.mark.parametrize("interleaved", [False, True])
+def test_fused_qk_norm_rope_large_storage_offsets(interleaved):
+    """Only three tokens, but the last row lies beyond signed 32-bit indexing."""
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
+
+    # 4 GiB of backing storage; initialize only the three small visible rows.
+    q = torch.empty_strided((3, 1, 128), (2**30, 128, 1), device="cuda", dtype=torch.bfloat16)
+    q.fill_(1)
+    weight = torch.ones(128, device="cuda", dtype=torch.bfloat16)
+    table = torch.zeros(3, 96, device="cuda", dtype=torch.bfloat16)
+    table[:, :48] = 1
+    actual = fused_qk_norm_rope(q, q, weight, weight, table, _EPS, interleaved=interleaved)
+    for value in actual:
+        torch.testing.assert_close(value, torch.ones_like(value))

tests/diffusion/layers/test_fused_qk_rope.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+from diffusers.models.embeddings import apply_rotary_emb
+from vllm.triton_utils import HAS_TRITON
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+def test_fused_qk_rope_fake_and_cpu_eligibility():
+    from vllm_omni.diffusion.layers.fused_qk_rope import (
+        _fused_qk_rope_fake,
+        fused_qk_rope,
+        fused_qk_rope_supported,
+    )
+
+    q = torch.randn(2, 3, 4, 8, dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    cos = torch.randn(3, 8)
+    sin = torch.randn_like(cos)
+
+    fake_q, fake_k = _fused_qk_rope_fake(q, k, cos, sin)
+    assert fake_q.shape == q.shape and fake_q.dtype == q.dtype
+    assert fake_k.shape == k.shape and fake_k.dtype == k.dtype
+    assert not fused_qk_rope_supported(q, k, cos, sin)
+    with pytest.raises(ValueError, match="contiguous CUDA BF16"):
+        fused_qk_rope(q, k, cos, sin)
+
+
+def _cuda_inputs(batch: int, sequence: int, heads: int = 24, head_dim: int = 128):
+    q = torch.randn(batch, sequence, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    # Deliberately do not repeat adjacent table columns.  Diffusers consumes
+    # full-width tables, so even and odd outputs may use different values.
+    cos = torch.randn(sequence, head_dim, device="cuda", dtype=torch.float32)
+    sin = torch.randn_like(cos)
+    return q, k, cos, sin
+
+
+def _reference(q, k, cos, sin):
+    rotary_emb = (cos, sin)
+    return (
+        apply_rotary_emb(q, rotary_emb, sequence_dim=1),
+        apply_rotary_emb(k, rotary_emb, sequence_dim=1),
+    )
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize("sequence", [512, 4096, 4608])
+def test_fused_qk_rope_is_bit_exact_at_longcat_production_shapes(sequence):
+    from vllm_omni.diffusion.layers.fused_qk_rope import _launch_fused_qk_rope, fused_qk_rope
+
+    torch.manual_seed(sequence)
+    with torch.inference_mode():
+        q, k, cos, sin = _cuda_inputs(1, sequence)
+        expected = _reference(q, k, cos, sin)
+        for actual in (
+            _launch_fused_qk_rope(q, k, cos, sin),
+            fused_qk_rope(q, k, cos, sin),
+        ):
+            assert torch.equal(actual[0], expected[0])
+            assert torch.equal(actual[1], expected[1])
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_fused_qk_rope_is_bit_exact_for_batch_and_full_width_odd_even_tables():
+    from vllm_omni.diffusion.layers.fused_qk_rope import fused_qk_rope
+
+    torch.manual_seed(11)
+    with torch.inference_mode():
+        q, k, cos, sin = _cuda_inputs(2, 17, heads=3)
+        assert not torch.equal(cos[:, ::2], cos[:, 1::2])
+        assert not torch.equal(sin[:, ::2], sin[:, 1::2])
+        expected = _reference(q, k, cos, sin)
+        actual = fused_qk_rope(q, k, cos, sin)
+
+    assert torch.equal(actual[0], expected[0])
+    assert torch.equal(actual[1], expected[1])
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_fused_qk_rope_rejects_noncontiguous_and_non_fp32_tables():
+    from vllm_omni.diffusion.layers.fused_qk_rope import fused_qk_rope_supported
+
+    q, k, cos, sin = _cuda_inputs(1, 8, heads=2)
+    q_noncontiguous = torch.randn(1, 8, 2, 256, device="cuda", dtype=torch.bfloat16)[..., ::2]
+    k_noncontiguous = torch.randn_like(q_noncontiguous)
+    cos_noncontiguous = torch.randn(8, 256, device="cuda")[:, ::2]
+    sin_noncontiguous = torch.randn(8, 256, device="cuda")[:, ::2]
+    with torch.inference_mode():
+        assert fused_qk_rope_supported(q, k, cos, sin)
+        assert not fused_qk_rope_supported(q_noncontiguous, k_noncontiguous, cos, sin)
+        assert not fused_qk_rope_supported(q, k, cos_noncontiguous, sin_noncontiguous)
+        assert not fused_qk_rope_supported(q, k, cos.to(torch.bfloat16), sin.to(torch.bfloat16))
+
+        q_wide = torch.randn(1, 8, 2, 256, device="cuda", dtype=torch.bfloat16)
+        k_wide = torch.randn_like(q_wide)
+        cos_wide = torch.randn(8, 256, device="cuda")
+        sin_wide = torch.randn_like(cos_wide)
+        assert not fused_qk_rope_supported(q_wide, k_wide, cos_wide, sin_wide)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_fused_qk_rope_custom_op_has_fullgraph_fake():
+    from vllm_omni.diffusion.layers.fused_qk_rope import fused_qk_rope
+
+    torch.manual_seed(29)
+    with torch.inference_mode():
+        q, k, cos, sin = _cuda_inputs(1, 19, heads=3)
+        expected = _reference(q, k, cos, sin)
+        compiled = torch.compile(fused_qk_rope, dynamic=True, fullgraph=True)
+        actual = compiled(q, k, cos, sin)
+
+    assert torch.equal(actual[0], expected[0])
+    assert torch.equal(actual[1], expected[1])

tests/diffusion/layers/test_qwen_select01_modulation.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.layers import qwen_select01_modulation
+from vllm_omni.diffusion.layers.qwen_select01_modulation import (
+    can_use_qwen_select01_triton,
+    fused_layernorm_select01,
+    fused_residual_layernorm_select01,
+    select01_modulation_native,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("platform_is_cuda", "tensor_is_cuda", "has_triton", "expected"),
+    [
+        (True, True, True, True),
+        (False, True, True, False),
+        (True, False, True, False),
+        (True, True, False, False),
+    ],
+)
+def test_qwen_select01_triton_requires_cuda_platform(
+    monkeypatch,
+    platform_is_cuda: bool,
+    tensor_is_cuda: bool,
+    has_triton: bool,
+    expected: bool,
+):
+    class FakePlatform:
+        @staticmethod
+        def is_cuda() -> bool:
+            return platform_is_cuda
+
+    class FakeTensor:
+        is_cuda = tensor_is_cuda
+
+    monkeypatch.setattr(qwen_select01_modulation, "current_platform", FakePlatform())
+    monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", has_triton)
+
+    assert can_use_qwen_select01_triton(FakeTensor()) is expected
+
+
+def _make_inputs(
+    batch_size: int,
+    seq_len: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, ...]:
+    generator = torch.Generator(device=device)
+    generator.manual_seed(202817)
+    x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator)
+    residual = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator)
+    residual_gate = torch.randn(batch_size, 1, hidden_size, dtype=dtype, device=device, generator=generator)
+    mod_params = torch.randn(batch_size * 2, hidden_size * 3, dtype=dtype, device=device, generator=generator)
+    index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int64, device=device, generator=generator)
+    return x, residual, residual_gate, mod_params, index
+
+
+def _reference_layernorm_select01(
+    x: torch.Tensor,
+    mod_params: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    scale, shift, gate = select01_modulation_native(mod_params, index)
+    out = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).to(x.dtype)
+    return out * (1 + scale) + shift, gate
+
+
+def _reference_residual_layernorm_select01(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    residual_gate: torch.Tensor,
+    mod_params: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    scale, shift, gate = select01_modulation_native(mod_params, index)
+    residual_out = residual + residual_gate * x
+    out = F.layer_norm(residual_out.float(), (residual_out.shape[-1],), eps=eps).to(residual_out.dtype)
+    return out * (1 + scale) + shift, residual_out, gate
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("broadcast_index", [False, True])
+def test_qwen_select01_native_fallback_matches_reference(dtype: torch.dtype, broadcast_index: bool):
+    device = torch.device("cpu")
+    x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 64, dtype, device)
+    if broadcast_index:
+        index = index[:1]
+    eps = 1e-6
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate,
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate,
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    ("dtype", "atol", "rtol"),
+    [
+        (torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, 5e-2, 5e-2),
+    ],
+)
+@pytest.mark.parametrize("seq_len", [1, 7, 128, 1024])
+def test_qwen_select01_triton_matches_reference(dtype: torch.dtype, atol: float, rtol: float, seq_len: int):
+    device = torch.device("cuda:0")
+    x, residual, residual_gate, mod_params, index = _make_inputs(2, seq_len, 256, dtype, device)
+    eps = 1e-6
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    ("dtype", "atol", "rtol"),
+    [
+        (torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, 5e-2, 5e-2),
+    ],
+)
+def test_qwen_select01_triton_matches_reference_for_structured_multi_batch_index(
+    dtype: torch.dtype,
+    atol: float,
+    rtol: float,
+):
+    device = torch.device("cuda:0")
+    x, residual, residual_gate, mod_params, _ = _make_inputs(3, 16, 256, dtype, device)
+    index = torch.tensor(
+        [
+            [0] * 16,
+            [1] * 16,
+            [0, 1] * 8,
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+    eps = 1e-6
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    ("dtype", "atol", "rtol"),
+    [
+        (torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, 5e-2, 5e-2),
+    ],
+)
+def test_qwen_select01_triton_broadcasts_single_index_batch(dtype: torch.dtype, atol: float, rtol: float):
+    device = torch.device("cuda:0")
+    x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 256, dtype, device)
+    index = index[:1]
+    eps = 1e-6
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_qwen_select01_cuda_falls_back_without_triton(monkeypatch):
+    device = torch.device("cuda:0")
+    x, residual, residual_gate, mod_params, index = _make_inputs(2, 7, 256, torch.bfloat16, device)
+    residual_gate = residual_gate.expand_as(residual)
+    eps = 1e-6
+
+    def fail_if_launched(*args, **kwargs):
+        pytest.fail("Triton launcher must not run when HAS_TRITON is false")
+
+    monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", False)
+    monkeypatch.setattr(qwen_select01_modulation, "_launch_layernorm_select01", fail_if_launched)
+    monkeypatch.setattr(qwen_select01_modulation, "_launch_residual_layernorm_select01", fail_if_launched)
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate,
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate,
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=0, rtol=0)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)

tests/diffusion/lora/test_loader.py

@@ -18,6 +18,7 @@ from vllm_omni.diffusion.lora.loader import (
     WanLoraLoaderMixin,
     _prepare_lora_delta,
     _remap_state_dict_keys,
+    _resolve_lora_compute_device,
 )
 
 HEAD_DIM = 24
@@ -509,3 +510,93 @@ def test_lora_loader_mixin_lora_is_fused_property(mocker: MockerFixture):
     assert pipeline2.lora_is_fused is False
     pipeline2.lora_is_fused = True
     assert pipeline2.lora_is_fused is True
+
+
+def test_prepare_lora_delta_with_compute_device():
+    base_key = "blocks.0.attn.to_q"
+    state_dict = {
+        f"{base_key}.lora_A.weight": torch.ones(RANK, HEAD_DIM),
+        f"{base_key}.lora_B.weight": torch.ones(HEAD_DIM, RANK),
+    }
+    delta, used_keys = _prepare_lora_delta(state_dict, base_key, compute_device=torch.device("cpu"))
+    assert delta is not None
+    assert delta.shape == (HEAD_DIM, HEAD_DIM)
+    assert len(used_keys) == 2
+
+
+def test_resolve_lora_compute_device_probes_unusable_accelerator(mocker: MockerFixture):
+    """cuda:N/npu:N may be reported without a physical device; the probe must reject it."""
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader.get_local_device",
+        return_value=torch.device("cuda:999"),
+    )
+    assert _resolve_lora_compute_device() is None
+
+
+def test_load_lora_into_module_cpu_weights_match_cpu_reference(mocker: MockerFixture):
+    """Regression test: params on CPU used to force compute_device=None and skip the
+    #7005 accelerator path. With an unusable accelerator, fusion must fall back to
+    CPU and still produce base + lora_B @ lora_A."""
+    pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
+    original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+
+    lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
+    lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
+
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader.get_local_device",
+        return_value=torch.device("cuda:999"),
+    )
+
+    used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
+    assert len(used_keys) > 0
+    assert not (set(lora_state_dict.keys()) - used_keys)
+
+    for name, param in pipeline.transformer.named_parameters():
+        if not name.endswith(".weight"):
+            continue
+        base_key = name[: -len(".weight")]
+        key_a = f"{base_key}.lora_A.weight"
+        key_b = f"{base_key}.lora_B.weight"
+        if key_a in lora_state_dict and key_b in lora_state_dict:
+            expected = original_params[name] + (
+                lora_state_dict[key_b].to(param.dtype) @ lora_state_dict[key_a].to(param.dtype)
+            )
+            assert_close(param, expected)
+
+
+def test_load_lora_into_module_forwards_usable_accel_device(mocker: MockerFixture):
+    """A usable accelerator must be forwarded as compute_device for CPU params.
+    A spy forces the actual math back onto CPU so the test runs without a device."""
+    fake_dev = torch.device("cuda:0")
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader._resolve_lora_compute_device",
+        return_value=fake_dev,
+    )
+
+    seen_devices = []
+    original = _prepare_lora_delta
+
+    def spy(state_dict, base_key, *args, **kwargs):
+        seen_devices.append(kwargs.get("compute_device"))
+        kwargs["compute_device"] = None
+        return original(state_dict, base_key, *args, **kwargs)
+
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader._prepare_lora_delta",
+        side_effect=spy,
+    )
+
+    pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
+    original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+
+    lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
+    lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
+
+    used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
+
+    # Every CPU param with a LoRA delta must have received the accelerator device.
+    assert seen_devices
+    assert all(d == fake_dev for d in seen_devices)
+    assert not (set(lora_state_dict.keys()) - used_keys)
+    assert_params_not_equal(pipeline.transformer, original_params)

tests/diffusion/lora/test_lora_manager.py

@@ -411,6 +411,50 @@ def test_lora_manager_rolls_back_all_layers_when_activation_fails():
     assert len(first.set_calls) == first_calls_after_success + 2
 
 
+def test_lora_manager_rejects_adapter_that_binds_no_layer():
+    """An adapter whose target modules match nothing must fail, not silently no-op."""
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+    layer = _DummyLoRALayer(n_slices=1, output_slices=(2,))
+    manager._lora_modules = {"transformer.blocks.0.attn.to_q": layer}
+
+    # Adapter-side name the engine does not expose, e.g. a diffusers-style
+    # checkpoint against a differently named engine layout.
+    mismatched = LoRALayerWeights(
+        module_name="unet.down_blocks.0.attn.to_q",
+        rank=2,
+        lora_alpha=2,
+        lora_a=torch.ones((2, 2)),
+        lora_b=torch.ones((2, 2)),
+    )
+    manager._registered_adapters = {
+        3: type(
+            "LM",
+            (),
+            {
+                "id": 3,
+                "loras": {"unet.down_blocks.0.attn.to_q": mismatched},
+                "get_lora": lambda self, key: self.loras.get(key),
+            },
+        )()
+    }
+
+    with pytest.raises(ValueError, match="applies to no layer") as excinfo:
+        manager._activate_adapter(3, scale=1.0)
+
+    # The message must name what was received so the mismatch is diagnosable.
+    assert "unet.down_blocks.0.attn.to_q" in str(excinfo.value)
+
+    # Nothing was bound and the adapter must not be left marked active.
+    assert manager._active_adapter_id is None
+    assert layer.set_calls == []
+    assert layer.reset_calls >= 1
+
+
 def _dummy_lora_request(adapter_id: int) -> LoRARequest:
     return LoRARequest(
         lora_name=f"adapter_{adapter_id}",
@@ -685,8 +729,10 @@ def test_lora_manager_discovers_bagel_component(monkeypatch):
         lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
     )
 
-    # Pipeline with a 'bagel' component (no 'transformer')
+    # Pipeline with a 'bagel' component (no 'transformer') declared via
+    # _lora_components, matching how BagelPipeline opts in.
     pipeline = torch.nn.Module()
+    pipeline._lora_components = ["bagel"]
     pipeline.bagel = torch.nn.Module()
     pipeline.bagel.language_model = torch.nn.Module()
     pipeline.bagel.language_model.qkv_proj = _FakeLinearBase()

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -6,9 +6,11 @@ Tests for the DiffusersPipelineLoader.
 """
 
 import json
+import os
+import tempfile
 from pathlib import Path
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -260,17 +262,16 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
     assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
 
 
-def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present():
+def test_maybe_fuse_distilled_lora_skips_when_lora_path_unset():
     cfg = SimpleNamespace(
         lora_backend="distill",
-        lora_path="/path/to/lora.safetensors",
+        lora_path=None,
         lora_scale=1.0,
         dtype=torch.bfloat16,
         quantization_config=None,
         parallel_config=SimpleNamespace(use_hsdp=False),
     )
     loader = DiffusersPipelineLoader(LoadConfig(), cfg)
-    loader._hwr_state = {"warm_snapshot": {"transformer.weight": (1234, "hash")}}
 
     model = nn.Module()
     model.load_lora_weights = MagicMock()
@@ -278,7 +279,7 @@ def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present():
     loader._maybe_fuse_distilled_lora(model)
 
     model.load_lora_weights.assert_not_called()
-    assert getattr(model, "lora_is_fused", False) is True
+    assert getattr(model, "lora_is_fused", False) is False
 
 
 def test_maybe_fuse_distilled_lora_fuses_when_not_warm_snapshot():
@@ -1053,6 +1054,34 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
     assert loader.take_host_weight_plan() is None
 
 
+def test_dlo_mmap_plan_with_distilled_lora_falls_back_to_ordinary_loader(monkeypatch):
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        lora_backend=LoRABackend.DISTILL,
+        lora_path="/fake/lora.safetensors",
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.load_lora_weights = MagicMock()
+    calls: list[str] = []
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: calls.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    assert loader.load_model(load_device="cpu") is model
+    assert calls == ["load", "process"]
+    assert model.load_lora_weights.call_count == 1
+    assert loader.take_host_weight_plan() is None
+
+
 def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
     from vllm.model_executor.layers.quantization.online.fp8 import (
         Fp8PerTensorOnlineLinearMethod,
@@ -1212,6 +1241,169 @@ def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
         loader.load_model(load_device="cpu")
 
 
+def test_dlo_allgather_allows_unquantized_host_fallback():
+    """Host-loaded unquantized fallback layers are plain contiguous bf16 —
+    the same runtime layout DLO already shards on the ordinary path — so the
+    allowlist must not reject them."""
+    from vllm_omni.quantization.int8_config import UnquantizedHostLinearMethod
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object.__new__(UnquantizedHostLinearMethod)
+
+    assert DiffusersPipelineLoader._unsupported_dlo_allgather_online_quant_methods(model) == ()
+
+
+def test_dlo_load_model_keeps_host_fallback_on_cpu_through_post_load_sweep(monkeypatch, mocker):
+    """Through load_model(): DLO + online quant must build the model inside
+    load_unquantizable_fallback_on_cpu(), and the over-wide fallback weight
+    must survive the post-load sweep (the _process_weights_after_loading pass
+    plus model.to("cpu")) as the same host tensor it was loaded into — never
+    bounced to the accelerator."""
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.quantization import int8_config
+    from vllm_omni.quantization.int8_config import (
+        NPU_QUANT_MATMUL_MAX_OUT_FEATURES,
+        DiffusionInt8Config,
+        NPUInt8OnlineLinearMethod,
+        UnquantizedHostLinearMethod,
+    )
+
+    # Same stand-in as TestHostFallbackLoading: the eager fallback delegates to
+    # UnquantizedLinearMethod, which reads the TP group.
+    mock_group = mocker.Mock()
+    mock_group.rank_in_group = 0
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    out_features = NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1
+    built_with_ctx: list[bool] = []
+
+    def copy_loader(param, loaded_weight, *args, **kwargs):
+        param.data.copy_(loaded_weight)
+
+    def build_model(*_args, **_kwargs):
+        built_with_ctx.append(int8_config._LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get())
+        quant_config = DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic")
+        method = NPUInt8OnlineLinearMethod(quant_config)
+        layer = nn.Module()
+        layer.quant_method = method
+        method.create_weights(
+            layer,
+            input_size_per_partition=8,
+            output_partition_sizes=[out_features],
+            input_size=8,
+            output_size=out_features,
+            params_dtype=torch.bfloat16,
+            weight_loader=copy_loader,
+        )
+        model = nn.Module()
+        model.transformer = layer
+        return model
+
+    def load_weights(_model, **_kwargs):
+        layer = _model.transformer
+        loaded = torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16)
+        layer.weight.weight_loader(layer.weight, loaded)
+
+    loader._init_from_load_format = build_model  # type: ignore[method-assign]
+    loader.load_weights = load_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+    # Diffusion DiT models have no vLLM Attention layers, so the upstream
+    # layerwise finalize is a no-op here; stub it to keep the real
+    # _process_weights_after_loading sweep CPU-only.
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.reload.layerwise.finalize_layerwise_processing",
+        lambda *_args, **_kwargs: None,
+    )
+
+    model = loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+    # Construction ran inside load_unquantizable_fallback_on_cpu(), so the
+    # over-wide layer took the host-loading fallback...
+    assert built_with_ctx == [True]
+    assert type(model.transformer.quant_method) is UnquantizedHostLinearMethod
+    # ...and the real post-load sweep skipped it via the fully-loaded flag and
+    # model.to("cpu") left the loaded host tensor untouched.
+    assert model.transformer._already_called_process_weights_after_loading
+    assert model.transformer.weight.device.type == "cpu"
+    expected = torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16)
+    assert torch.equal(model.transformer.weight, expected)
+
+
+@pytest.mark.parametrize(("offload_after_quant", "ctx_expected"), [(True, True), (False, False)])
+def test_hsdp_enters_host_fallback_context_only_when_offloading_after_quant(mocker, offload_after_quant, ctx_expected):
+    """The HSDP path must apply the same host-fallback bound as the ordinary
+    path: with a quant config it initializes on the accelerator, so the
+    load_unquantizable_fallback_on_cpu() context is what keeps over-wide
+    fallback weights off the device before sharding."""
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+    from vllm_omni.quantization import int8_config
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader.quant_config = object()
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    ctx_seen: list[bool] = []
+
+    def build_model(*_args, **_kwargs):
+        ctx_seen.append(int8_config._LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get())
+        return model
+
+    loader._init_from_load_format = build_model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
+        side_effect=lambda _model: None,
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: None,
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"), offload_after_quant=offload_after_quant)
+
+    assert ctx_seen == [ctx_expected]
+
+
 def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
     mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
     NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod
@@ -1487,3 +1679,413 @@ def test_load_model(prefetch_helios_model, mock_tp_group):
     )
     model = loader.load_model(load_device="cpu")
     assert isinstance(model, HeliosPipeline)
+
+
+def test_hsdp_broadcast_weight_load_rank0(mocker):
+    """Ensure that on rank 0 with enable_broadcast_weight_load=True, weights are loaded, fused, and broadcast."""
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=4)
+    mocker.patch("torch.distributed.get_rank", return_value=0)
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=4,
+        ),
+        enable_broadcast_weight_load=True,
+        lora_backend=LoRABackend.DISTILL,
+        lora_path="/fake/lora.safetensors",
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._maybe_fuse_distilled_lora = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model: events.append("fuse")
+    )
+    loader._broadcast_model_weights = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, **_kwargs: events.append("broadcast")
+    )
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, _device: events.append("process")
+    )
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: events.append("shard"),
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"))
+
+    assert events == ["load", "fuse", "broadcast", "process", "shard"]
+
+
+def test_hsdp_broadcast_weight_load_rank_nonzero(mocker):
+    """Ensure that on rank > 0 with enable_broadcast_weight_load=True, loading is skipped and broadcast is received."""
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=4)
+    mocker.patch("torch.distributed.get_rank", return_value=1)
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=4,
+        ),
+        enable_broadcast_weight_load=True,
+        lora_backend=LoRABackend.DISTILL,
+        lora_path="/fake/lora.safetensors",
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.load_lora_weights = mocker.Mock()  # type: ignore[assignment]
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._maybe_fuse_distilled_lora = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model: events.append("fuse")
+    )
+    loader._broadcast_model_weights = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, **_kwargs: events.append("broadcast")
+    )
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, _device: events.append("process")
+    )
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: events.append("shard"),
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"))
+
+    # Rank 1 should NOT call load_weights or _maybe_fuse_distilled_lora from disk
+    assert "load" not in events
+    assert "fuse" not in events
+    assert events == ["broadcast", "process", "shard"]
+    assert getattr(model, "lora_is_fused", False) is True
+
+
+def test_broadcast_model_weights_invokes_dist_broadcast(mocker):
+    broadcast_calls = []
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=2)
+    mocker.patch("torch.distributed.get_rank", return_value=0)
+    mocker.patch("torch.distributed.get_backend", return_value="gloo")
+    mocker.patch("torch.distributed.broadcast", side_effect=lambda tensor, src: broadcast_calls.append((tensor, src)))
+    mocker.patch("torch.distributed.barrier")
+
+    model = nn.Module()
+    model.linear = nn.Linear(2, 2)
+    model.register_buffer("buf", torch.tensor([1.0, 2.0]))
+    model.register_buffer("int_buf", torch.tensor([1, 2, 3], dtype=torch.long))
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=True),
+        enable_broadcast_weight_load=True,
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader._broadcast_model_weights(model, target_device=torch.device("cpu"), src_rank=0)
+
+    # float32 tensors (linear.weight, linear.bias, buf) are coalesced into 1 bucket,
+    # and int64 tensor (int_buf) is in a separate bucket -> 2 coalesced broadcast calls
+    assert len(broadcast_calls) == 2
+    assert all(src == 0 for _, src in broadcast_calls)
+
+
+def test_broadcast_model_weights_receiver_copies_data(mocker):
+    """Ensure receiver rank copies broadcasted flat tensor data into its model parameters/buffers."""
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=2)
+    mocker.patch("torch.distributed.get_rank", return_value=1)
+    mocker.patch("torch.distributed.get_backend", return_value="gloo")
+    mocker.patch("torch.distributed.barrier")
+
+    # Simulate broadcast by filling the flat tensor with source values
+    def mock_broadcast(tensor, src):
+        if tensor.dtype == torch.float32:
+            tensor.fill_(3.14)
+        elif tensor.dtype == torch.long:
+            tensor.fill_(42)
+
+    mocker.patch("torch.distributed.broadcast", side_effect=mock_broadcast)
+
+    model = nn.Module()
+    model.linear = nn.Linear(2, 2)
+    model.register_buffer("buf", torch.zeros(2))
+    model.register_buffer("int_buf", torch.zeros(3, dtype=torch.long))
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=True),
+        enable_broadcast_weight_load=True,
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader._broadcast_model_weights(model, target_device=torch.device("cpu"), src_rank=0)
+
+    assert torch.allclose(model.linear.weight, torch.tensor(3.14))
+    assert torch.allclose(model.linear.bias, torch.tensor(3.14))
+    assert torch.allclose(model.buf, torch.tensor(3.14))
+    assert torch.equal(model.int_buf, torch.tensor([42, 42, 42], dtype=torch.long))
+
+
+def test_broadcast_model_weights_mixed_cpu_and_accelerator_devices(mocker):
+    """Ensure that when some pipeline modules are on CPU (transformer) and others on accelerator (VAE/encoder),
+    broadcasting normalizes tensors to CPU before concatenation and succeeds without device mismatch.
+    """
+    dev = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=2)
+    mocker.patch("torch.distributed.get_rank", return_value=0)
+    mocker.patch("torch.distributed.get_backend", return_value="nccl" if dev.type == "cuda" else "gloo")
+    mocker.patch("torch.distributed.broadcast")
+    mocker.patch("torch.distributed.barrier")
+
+    model = nn.Module()
+    # Transformer parameters on CPU
+    model.transformer = nn.Linear(4, 4)
+    # VAE / Encoder parameters on accelerator
+    model.vae = nn.Linear(4, 4).to(dev)
+    model.text_encoder = nn.Linear(4, 4).to(dev)
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=True),
+        enable_broadcast_weight_load=True,
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    # Should not raise RuntimeError: Tensors must have same device
+    loader._broadcast_model_weights(model, target_device=dev, src_rank=0)
+
+
+def test_broadcast_model_weights_mixed_devices_receiver(mocker):
+    """Ensure receiver rank with mixed CPU/accelerator tensors correctly receives broadcast data."""
+    dev = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=2)
+    mocker.patch("torch.distributed.get_rank", return_value=1)
+    mocker.patch("torch.distributed.get_backend", return_value="nccl" if dev.type == "cuda" else "gloo")
+    mocker.patch("torch.distributed.barrier")
+
+    def mock_broadcast(tensor, src):
+        tensor.fill_(7.0)
+
+    mocker.patch("torch.distributed.broadcast", side_effect=mock_broadcast)
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2)  # CPU
+    model.vae = nn.Linear(2, 2).to(dev)  # Accelerator
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=True),
+        enable_broadcast_weight_load=True,
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader._broadcast_model_weights(model, target_device=dev, src_rank=0)
+
+    assert torch.allclose(model.transformer.weight, torch.tensor(7.0))
+    assert torch.allclose(model.vae.weight, torch.tensor(7.0, device=dev))
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_hsdp_broadcast_weight_load_falls_back_when_online_quantization_enabled(mocker, rank):
+    """Ensure that when online quantization is enabled, HSDP weight loading falls back to
+    ordinary per-rank loading on all ranks rather than rank-0 broadcast.
+    """
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    mocker.patch("torch.distributed.is_initialized", return_value=True)
+    mocker.patch("torch.distributed.get_world_size", return_value=4)
+    mocker.patch("torch.distributed.get_rank", return_value=rank)
+
+    mock_quant_config = SimpleNamespace(
+        is_checkpoint_quantized=False,
+    )
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=4,
+        ),
+        enable_broadcast_weight_load=True,
+        lora_backend=None,
+        lora_path=None,
+        quantization_config=mock_quant_config,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._broadcast_model_weights = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, **_kwargs: events.append("broadcast")
+    )
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, _device: events.append("process")
+    )
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: events.append("shard"),
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"))
+
+    # Both rank 0 and non-zero ranks MUST execute ordinary load_weights, not broadcast
+    assert "load" in events
+    assert "broadcast" not in events
+    assert events == ["load", "process", "shard"]
+
+
+def _mp_worker_online_quant(rank: int, world_size: int, rendezvous: str, temp_dir: str) -> None:
+    torch.set_num_threads(1)
+    import torch.distributed as dist
+
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=world_size)
+    try:
+
+        class _MockOnlineQuantModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = nn.Linear(4, 4)
+                self.load_called = False
+
+            def load_weights(self, weights):
+                self.load_called = True
+                return {"transformer.weight", "transformer.bias"}
+
+        mock_quant_config = SimpleNamespace(is_checkpoint_quantized=False)
+        od_config = SimpleNamespace(
+            dtype=torch.float32,
+            parallel_config=SimpleNamespace(
+                use_hsdp=True,
+                hsdp_replicate_size=1,
+                hsdp_shard_size=world_size,
+            ),
+            enable_broadcast_weight_load=True,
+            lora_backend=None,
+            lora_path=None,
+            quantization_config=mock_quant_config,
+        )
+        loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+        model = _MockOnlineQuantModel()
+        loader._init_from_load_format = lambda *args, **kwargs: model  # type: ignore[method-assign]
+        loader._process_weights_after_loading = lambda *args, **kwargs: None  # type: ignore[method-assign]
+
+        with (
+            patch.object(
+                loader_module.ModuleDiscovery,
+                "discover",
+                return_value=PipelineModules(
+                    dits=[model.transformer],
+                    dit_names=["transformer"],
+                    vaes=[],
+                    encoders=[],
+                    encoder_names=[],
+                    resident_modules=[],
+                    resident_names=[],
+                ),
+            ),
+            patch.object(loader_module, "apply_hsdp_to_model", return_value=None),
+        ):
+            loader._load_model_with_hsdp(torch.device("cpu"))
+
+        if model.load_called:
+            with open(os.path.join(temp_dir, f"rank_{rank}_success.flag"), "w") as f:
+                f.write("ok")
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(
+    not torch.distributed.is_available() or not torch.distributed.is_gloo_available(),
+    reason="requires torch.distributed gloo",
+)
+def test_hsdp_broadcast_weight_load_online_quant_multiprocess():
+    """Multi-process regression test verifying that with online quantization and broadcast enabled,
+    all ranks fall back to independent load_weights in parallel processes without hanging.
+    """
+    import torch.multiprocessing as mp
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
+        world_size = 2
+        with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
+            mp.spawn(
+                _mp_worker_online_quant,
+                args=(world_size, rendezvous, temp_dir),
+                nprocs=world_size,
+                join=True,
+            )
+        assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag"))
+        assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag"))

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)