Merge branch 'main' into codex-diffusion-shutdown-escalation

tests/diffusion/ar_diffusion/test_capability_runner.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """CPU contracts for capability-driven AR-Diffusion sessions."""
 
 from __future__ import annotations
@@ -10,6 +11,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.experimental.ar_diffusion.capability import (
     ARDiffusionCrossAttentionKVSpec,
@@ -27,6 +29,22 @@ NEG = "negative"
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def platform_synchronize_calls(monkeypatch) -> list[None]:
+    """Stub the platform barrier the runner takes at every chunk boundary.
+
+    These tests run on a CPU host, where the resolved platform is
+    ``UnspecifiedOmniPlatform`` and ``synchronize()`` raises. Tests that need
+    a failing barrier override this with their own fake.
+    """
+    calls: list[None] = []
+    monkeypatch.setattr(
+        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
+        SimpleNamespace(synchronize=lambda: calls.append(None)),
+    )
+    return calls
+
+
 def lingbot_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
     """Single-kv_branch causal DMD: 3 latent frames/block + sink/window + text KV."""
     return ARDiffusionKVCacheSpec(
@@ -114,6 +132,43 @@ class BatchCapablePipeline(CapablePipeline):
     supports_request_batch = True
 
 
+class StepCapablePipeline(CapablePipeline):
+    supports_step_execution = True
+
+    def prepare_encode(self, state, **kwargs):
+        del kwargs
+        return state
+
+    def denoise_step(self, input_batch, *, states=None, **kwargs):
+        del input_batch, states, kwargs
+        return None
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del state, noise_pred, kwargs
+
+    def post_decode(self, state, **kwargs):
+        del state, kwargs
+        from vllm_omni.diffusion.data import DiffusionOutput
+
+        return DiffusionOutput()
+
+
+def scheduler_output(
+    request_ids: tuple[str, ...] = ("req-1",),
+    *,
+    finished: tuple[str, ...] = (),
+) -> DiffusionSchedulerOutput:
+    """Build the real scheduler payload; stepwise requests arrive as cached rows."""
+    return DiffusionSchedulerOutput(
+        step_id=0,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData(request_ids=list(request_ids)),
+        finished_req_ids=set(finished),
+        num_running_reqs=len(request_ids),
+        num_waiting_reqs=0,
+    )
+
+
 def make_runner(
     pipeline: object,
     *,
@@ -140,6 +195,7 @@ def make_runner(
     runner._sessions = OrderedDict()
     runner._session_capacity = 0
     runner._perf_e2e_times = []
+    runner._stepwise_chunk_started = {}
     runner._preallocate_kv_cache(available_bytes=available_bytes)
     return runner
 
@@ -395,12 +451,16 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
         ctx.ensure_video_slots(torch.device("cpu"))
         return object()
 
-    def synchronize_boom(device):
+    def synchronize_boom():
         raise RuntimeError("asynchronous kernel failed")
 
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
-    monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
-    runner.device = torch.device("cuda")
+    # The runner defers the device barrier to the platform, so the failure is
+    # injected there rather than on a device-specific torch entry point.
+    monkeypatch.setattr(
+        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
+        SimpleNamespace(synchronize=synchronize_boom),
+    )
     request = SimpleNamespace(
         request_id="broken-request",
         sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
@@ -424,6 +484,11 @@ def test_ar_runner_rejects_step_and_request_batch_modes():
         make_runner(BatchCapablePipeline(lingbot_like_spec()))
 
 
+def test_ar_runner_allows_step_execution_when_pipeline_implements_the_contract():
+    runner = make_runner(StepCapablePipeline(lingbot_like_spec()), step_execution=True)
+    assert runner.kv_cache is not None
+
+
 def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
     runner = object.__new__(ARDiffusionModelRunner)
     with pytest.raises(RuntimeError, match="request-batch execution"):
@@ -432,6 +497,144 @@ def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
         runner.execute_stepwise(None)
 
 
+def test_execute_stepwise_binds_request_id_session_and_unbinds_after_call(monkeypatch):
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    bound_during = []
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        bound_during.append(pipeline.bound_state is not None)
+        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False)])
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    output = runner.execute_stepwise(scheduler_output())
+
+    assert bound_during == [True]
+    assert pipeline.bound_state is None
+    assert pipeline.binds == ["req-1"]
+    assert "req-1" in runner._sessions
+    assert output.get_request_output("req-1") is not None
+    assert not pipeline.closes
+
+
+def test_execute_stepwise_closes_session_when_request_finishes(monkeypatch):
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=True)])
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    runner.execute_stepwise(scheduler_output())
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["req-1"]
+    assert "req-1" not in runner._sessions
+
+
+def test_execute_stepwise_exception_releases_kv_and_does_not_resume_pages(monkeypatch):
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    first = runner._get_or_create_session("req-1")
+
+    def boom(self, scheduler_output):
+        del self, scheduler_output
+        raise RuntimeError("step failed")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", boom)
+    with pytest.raises(RuntimeError, match="step failed"):
+        runner.execute_stepwise(scheduler_output())
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["req-1"]
+    assert "req-1" not in runner._sessions
+    second = runner._get_or_create_session("req-1")
+    assert second is not first
+
+
+def test_execute_stepwise_closes_session_for_scheduler_aborted_request(monkeypatch):
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    runner._get_or_create_session("req-1")
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        return BatchRunnerOutput.from_list([])
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    # An aborted request is never scheduled again and never reports finished,
+    # so the retire path is the only thing that can free its KV.
+    runner.execute_stepwise(scheduler_output(request_ids=(), finished=("req-1",)))
+
+    assert "req-1" not in runner._sessions
+    assert pipeline.closes == ["req-1"]
+
+
+def test_execute_stepwise_times_once_per_chunk_not_once_per_step(monkeypatch, platform_synchronize_calls):
+    from vllm_omni.diffusion.data import DiffusionOutput
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    results = [None, None, None, DiffusionOutput()]
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=results.pop(0))])
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    for _ in range(4):
+        runner.execute_stepwise(scheduler_output())
+
+    # Four denoise steps, one emitted AR block: one timing sample, matching the
+    # one-sample-per-block meaning request mode already has.
+    assert len(runner._perf_e2e_times) == 1
+    assert len(platform_synchronize_calls) == 1
+    assert "req-1" not in runner._stepwise_chunk_started
+
+
+def test_execute_stepwise_synchronize_exception_fails_closed(monkeypatch):
+    """The chunk barrier is where an asynchronous device error surfaces. It must
+    clean up like an in-step failure instead of leaving the session, the cached
+    step state and the timing entry to the scheduler's next cycle."""
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    runner.state_cache = {"req-1": object()}
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=object())])
+
+    def synchronize_boom():
+        raise RuntimeError("asynchronous kernel failed")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    monkeypatch.setattr(
+        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
+        SimpleNamespace(synchronize=synchronize_boom),
+    )
+
+    with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
+        runner.execute_stepwise(scheduler_output())
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["req-1"]
+    assert "req-1" not in runner._sessions
+    assert "req-1" not in runner.state_cache
+    assert "req-1" not in runner._stepwise_chunk_started
+    assert not runner._perf_e2e_times
+
+
 def test_model_specific_warmup_provider_is_consumed(monkeypatch):
     requests = [object(), object()]
     pipeline = WarmupPipeline(lingbot_like_spec(), requests)

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -25,12 +25,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
 BLOCK = 16
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
@@ -155,14 +152,155 @@ def test_paged_pool_layout_exposes_flat_slot_views():
         dtype=torch.float32,
         device=torch.device("cpu"),
     )
-    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
+    key_cache, value_cache = kv_pools[0]
+    assert key_cache.shape == value_cache.shape == (4, BLOCK, 4, 64)
     assert k_pools[0].shape == (4 * BLOCK, 4, 64)
     assert v_pools[0].shape == (4 * BLOCK, 4, 64)
 
+    # The flat views must alias their block-shaped cache, so a slot write is
+    # visible through the layout the attention kernel reads.
     k_pools[0][BLOCK + 3].fill_(7)
     v_pools[0][2 * BLOCK + 5].fill_(11)
-    assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
-    assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
+    assert torch.equal(key_cache[1, 3], k_pools[0][BLOCK + 3])
+    assert torch.equal(value_cache[2, 5], v_pools[0][2 * BLOCK + 5])
+
+
+def test_key_and_value_caches_do_not_share_storage():
+    """K and V must be separate allocations, not halves of one tensor.
+
+    The paged-write custom op declares both as mutated. When they alias one
+    storage, inductor's reinplace pass can re-inplace only the first of the
+    two, and auto_functionalized_v2's clone of the entire second pool survives
+    into the compiled graph -- one full pool copied per compiled region per
+    denoising step.
+    """
+    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
+        num_blocks=4,
+        block_size=BLOCK,
+        num_layers=2,
+        num_kv_heads=4,
+        head_dim=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    seen = set()
+    for layer, (key_cache, value_cache) in enumerate(kv_pools):
+        k_storage = key_cache.untyped_storage().data_ptr()
+        v_storage = value_cache.untyped_storage().data_ptr()
+        assert k_storage != v_storage, f"layer {layer}: K and V share one allocation"
+        assert k_pools[layer].untyped_storage().data_ptr() == k_storage
+        assert v_pools[layer].untyped_storage().data_ptr() == v_storage
+        seen.update((k_storage, v_storage))
+    # Every layer's K and V are distinct allocations as well.
+    assert len(seen) == 2 * len(kv_pools)
+
+    # Negative control: writing V must not disturb K.
+    k_pools[0].fill_(0)
+    v_pools[0].fill_(5)
+    assert torch.equal(k_pools[0], torch.zeros_like(k_pools[0]))
+
+
+def _pool_sized_allocations(code: str, pool_numel: int) -> list[int]:
+    """Sizes of generated allocations at least as large as one pool."""
+    import re
+
+    sizes = []
+    for shape in re.findall(r"empty_strided_\w+\(\(([\d, ]*?)\)", code):
+        dims = [int(part) for part in shape.replace(" ", "").strip(",").split(",") if part]
+        numel = 1
+        for dim_size in dims:
+            numel *= dim_size
+        if dims and numel >= pool_numel:
+            sizes.append(numel)
+    return sizes
+
+
+def _generated_code_for_pools(key_cache, value_cache):
+    """Compile a paged-write-shaped op over these pools, return the generated code."""
+    from torch._inductor import config as inductor_config
+    from torch._inductor.utils import run_and_get_code
+
+    class _Holder(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.register_buffer("key_cache", key_cache)
+            self.register_buffer("value_cache", value_cache)
+
+        def forward(self, value):
+            return torch.ops.test_kv_pool_clone.write(self.key_cache, self.value_cache, value)
+
+    torch._dynamo.reset()
+    # Inductor caches compiled code across ``torch._dynamo.reset()``, so the
+    # second call in a process would otherwise return the first one's code and
+    # both branches of this test would describe the same graph.
+    with inductor_config.patch(force_disable_caches=True):
+        _, codes = run_and_get_code(torch.compile(_Holder(), backend="inductor", fullgraph=True), torch.ones(1))
+    return "\n".join(codes)
+
+
+def test_compiled_paged_write_does_not_clone_the_pool():
+    """The compiled graph must write the pools in place, not through a clone.
+
+    ``test_key_and_value_caches_do_not_share_storage`` pins the mechanism this
+    change introduces -- two allocations rather than one. This pins the effect
+    that mechanism buys, which is the reason the change exists: with K and V as
+    halves of one tensor, inductor's reinplace pass re-inplaces only the first,
+    and ``auto_functionalized_v2``'s clone of the pool survives into the
+    generated code as a pool-sized allocation plus a copy, once per compiled
+    region per denoising step.
+
+    The shared layout is compiled here as a positive control. Without it, a
+    codegen change that stops emitting ``empty_strided_*`` would leave the
+    detector matching nothing and this test passing for the wrong reason.
+    """
+    if "write" not in dir(getattr(torch.ops, "test_kv_pool_clone", object())):
+
+        @torch.library.custom_op("test_kv_pool_clone::write", mutates_args=("key_pool", "value_pool"))
+        def _write(key_pool: torch.Tensor, value_pool: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+            key_pool[0, 0, 0, 0] = value[0]
+            value_pool[0, 0, 0, 0] = value[0]
+            return value * 2
+
+        @_write.register_fake
+        def _(key_pool, value_pool, value):
+            return torch.empty_like(value)
+
+    # The stand-in models the real op's contract; that declaration is what
+    # 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"], (
+        f"the real paged-write op now mutates {real_mutated}; this test still models two pools"
+    )
+
+    num_blocks, heads, dim = 8, 4, 16
+    cache_shape = (num_blocks, BLOCK, heads, dim)
+    pool_numel = num_blocks * BLOCK * heads * dim
+
+    # Positive control: one allocation, K and V aliasing it.
+    shared = torch.empty(2, *cache_shape)
+    shared_allocations = _pool_sized_allocations(_generated_code_for_pools(shared[0], shared[1]), pool_numel)
+    assert shared_allocations, (
+        "the shared layout produced no pool-sized allocation, so this test can no "
+        "longer tell the two layouts apart -- inductor's codegen or the detector changed"
+    )
+
+    # The allocator under test.
+    kv_pools, _, _ = allocate_kv_pool_with_views(
+        num_blocks=num_blocks,
+        block_size=BLOCK,
+        num_layers=1,
+        num_kv_heads=heads,
+        head_dim=dim,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    key_cache, value_cache = kv_pools[0]
+    separate_allocations = _pool_sized_allocations(_generated_code_for_pools(key_cache, value_cache), pool_numel)
+    assert not separate_allocations, (
+        "the compiled graph allocates a pool-sized buffer, so the pool is being cloned "
+        f"rather than written in place: {separate_allocations} (control saw {shared_allocations})"
+    )
 
 
 def test_build_manager_allocate_free_roundtrip():

tests/diffusion/ar_diffusion/test_kv_cache_cuda.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CUDA-only regression test for the AR-Diffusion paged KV write op.
+
+Split out of ``test_kv_cache.py``: that module is swept into CI's
+``core_model and cpu`` job via a module-level ``pytest.mark.cpu``, and
+``hardware_test(..., num_cards=1)`` adds no ``skipif`` guard for a missing
+GPU, so this test was being collected (and would fail to import a CUDA
+device) on CPU-only runners. Keeping it in its own, non-``cpu`` module means
+it is only ever collected by a job that explicitly selects ``cuda`` for this
+path -- none does yet, so for now this is intentionally not wired into any
+CI job.
+"""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.experimental.ar_diffusion.kv_cache import allocate_kv_pool_with_views
+
+BLOCK = 16
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_real_paged_write_op_does_not_clone_the_pool_on_cuda():
+    """Same assertion as the CPU test, against the real op's schema and CUDA codegen.
+
+    The CPU test compiles a stand-in with the same ``mutates_args`` contract,
+    because the real op needs a device and an attention kernel. That leaves
+    one thing unverified: whether Triton codegen reinplaces the same way,
+    which is the path the reported regression was observed on.
+
+    This does not call the real op directly: ``ar_diffusion_paged_write_attn``
+    runs FlashAttention's paged ``flash_attn_varlen_func`` internally, and
+    compiling that under ``fullgraph=True`` can fail for reasons that have
+    nothing to do with cloning (a missing kernel for this toy shape, a graph
+    break, or an autotune illegal-memory-access on the control pool). Instead
+    this registers a second stand-in that copies the real op's full argument
+    list and ``mutates_args`` declaration -- so reinplace sees the identical
+    contract the real op presents -- but skips straight to the pool writes,
+    the same way the real op's own implementation does before it calls
+    attention. That isolates "does Inductor clone this schema on CUDA" from
+    "does FlashAttention compile", which is a separate, already-covered
+    concern.
+
+    Note the pools passed here are the flat views, since that is what the op
+    takes -- it indexes them by slot and unflattens internally.
+    """
+    import re
+
+    from torch._inductor import config as inductor_config
+    from torch._inductor.utils import run_and_get_code
+
+    device = torch.device("cuda")
+    num_blocks, heads, dim, num_tokens = 8, 4, 64, 8
+    pool_numel = num_blocks * BLOCK * heads * dim
+    # The size a genuine clone would produce: K and V pool-sized allocations,
+    # coalesced into one buffer the way `auto_functionalized_v2` was cloning
+    # them before this PR's fix. Matching this exact size (rather than "any
+    # allocation at least as large as one pool") avoids false positives from
+    # unrelated workspace/scratch buffers Inductor may emit in that range.
+    clone_numel = 2 * pool_numel
+
+    if "write" not in dir(getattr(torch.ops, "test_kv_pool_clone_schema", object())):
+
+        @torch.library.custom_op("test_kv_pool_clone_schema::write", mutates_args=("key_pool", "value_pool"))
+        def _write(
+            query: torch.Tensor,
+            k_curr: torch.Tensor,
+            v_curr: torch.Tensor,
+            k_act: torch.Tensor | None,
+            v_act: torch.Tensor | None,
+            key_pool: torch.Tensor,
+            value_pool: torch.Tensor,
+            block_size: int,
+            video_slots: torch.Tensor,
+            action_slots: torch.Tensor,
+            block_table: torch.Tensor,
+            query_start_loc: torch.Tensor,
+            seq_lens: torch.Tensor,
+            max_query_len: int,
+            max_seq_len: int,
+            softmax_scale: float,
+        ) -> torch.Tensor:
+            # Same pool-write side effect as `_paged_write_attn_impl`, minus
+            # the `ar_diffusion_paged_attention` call -- that is the part
+            # this test is not trying to exercise.
+            key_pool[video_slots] = k_curr.to(key_pool.dtype)
+            value_pool[video_slots] = v_curr.to(value_pool.dtype)
+            if k_act is not None and v_act is not None and k_act.shape[0] > 0:
+                key_pool[action_slots] = k_act.to(key_pool.dtype)
+                value_pool[action_slots] = v_act.to(value_pool.dtype)
+            return query * 2
+
+        @_write.register_fake
+        def _(
+            query,
+            k_curr,
+            v_curr,
+            k_act,
+            v_act,
+            key_pool,
+            value_pool,
+            block_size,
+            video_slots,
+            action_slots,
+            block_table,
+            query_start_loc,
+            seq_lens,
+            max_query_len,
+            max_seq_len,
+            softmax_scale,
+        ):
+            return torch.empty_like(query)
+
+    class _Holder(torch.nn.Module):
+        def __init__(self, key_flat, value_flat):
+            super().__init__()
+            self.register_buffer("key_flat", key_flat)
+            self.register_buffer("value_flat", value_flat)
+
+        def forward(self, query, k_curr, v_curr, video_slots, action_slots, block_table, query_start_loc, seq_lens):
+            return torch.ops.test_kv_pool_clone_schema.write(
+                query,
+                k_curr,
+                v_curr,
+                None,
+                None,
+                self.key_flat,
+                self.value_flat,
+                BLOCK,
+                video_slots,
+                action_slots,
+                block_table,
+                query_start_loc,
+                seq_lens,
+                num_tokens,
+                num_tokens,
+                dim**-0.5,
+            )
+
+    def compiled_pool_allocations(key_flat, value_flat) -> list[int]:
+        args = (
+            torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
+            torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
+            torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
+            torch.arange(num_tokens, device=device),
+            torch.empty(0, dtype=torch.long, device=device),
+            torch.arange(num_blocks, device=device, dtype=torch.int32).unsqueeze(0),
+            torch.tensor([0, num_tokens], device=device, dtype=torch.int32),
+            torch.tensor([num_tokens], device=device, dtype=torch.int32),
+        )
+        torch._dynamo.reset()
+        # See the CPU test: inductor's cache outlives torch._dynamo.reset().
+        with inductor_config.patch(force_disable_caches=True):
+            _, codes = run_and_get_code(
+                torch.compile(_Holder(key_flat, value_flat), backend="inductor", fullgraph=True), *args
+            )
+        sizes = []
+        for shape in re.findall(r"empty_strided_cuda\(\(([\d, ]*?)\)", "\n".join(codes)):
+            dims = [int(part) for part in shape.replace(" ", "").strip(",").split(",") if part]
+            numel = 1
+            for dim_size in dims:
+                numel *= dim_size
+            if dims and numel >= pool_numel:
+                sizes.append(numel)
+        return sizes
+
+    # Positive control: one allocation, K and V as its two halves. This must
+    # produce an allocation of exactly `clone_numel` -- that is the clone
+    # signature the assertion below checks for, not "no allocation this big".
+    shared = torch.empty(2, num_blocks * BLOCK, heads, dim, dtype=torch.bfloat16, device=device)
+    control = compiled_pool_allocations(shared[0], shared[1])
+    assert clone_numel in control, (
+        f"the shared layout produced no {clone_numel}-element allocation ({control}), so this test can no "
+        "longer tell the two layouts apart -- inductor's codegen or the detector changed"
+    )
+
+    _, k_pools, v_pools = allocate_kv_pool_with_views(
+        num_blocks=num_blocks,
+        block_size=BLOCK,
+        num_layers=1,
+        num_kv_heads=heads,
+        head_dim=dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    separate = compiled_pool_allocations(k_pools[0], v_pools[0])
+    assert clone_numel not in separate, (
+        f"the compiled graph allocates a {clone_numel}-element buffer, so the pool is being cloned "
+        f"rather than written in place: {separate} (control saw {control})"
+    )

tests/diffusion/attention/test_kernels_hub.py

@@ -13,6 +13,12 @@ from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.platforms import current_omni_platform
 
 
+def _clear_hub_module_cache(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.attention.backends import flash_attn_hub
+
+    monkeypatch.setattr(flash_attn_hub, "_hub_modules", {})
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
@@ -54,6 +60,74 @@ def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.Mo
         CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_explicit_kernels_hub_variant_miss_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
+    """Explicit Hub backends must fail loudly when no build variant resolves (#6971)."""
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
+    _clear_hub_module_cache(monkeypatch)
+
+    kernels_module = types.ModuleType("kernels")
+
+    def _missing_variant(repo_id, version=None, **kwargs):
+        raise FileNotFoundError(f"Cannot find a build variant for this system in {repo_id}")
+
+    kernels_module.get_kernel = _missing_variant  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "kernels", kernels_module)
+
+    with pytest.raises(RuntimeError, match="no compatible build"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+    with pytest.raises(RuntimeError, match="no compatible build"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_explicit_kernels_hub_preflight_uses_v2_when_v1_misses(monkeypatch: pytest.MonkeyPatch):
+    """Preflight must share _load_hub_module's (1, 2) policy, not pin version=1."""
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
+    _clear_hub_module_cache(monkeypatch)
+
+    kernels_module = types.ModuleType("kernels")
+    versions_tried: list[int | None] = []
+
+    def _v1_miss_v2_hit(repo_id, version=None, **kwargs):
+        versions_tried.append(version)
+        if version == 1:
+            raise FileNotFoundError(f"Cannot find a build variant for this system in {repo_id}")
+        return types.SimpleNamespace(name=f"{repo_id}@v{version}")
+
+    kernels_module.get_kernel = _v1_miss_v2_hit  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "kernels", kernels_module)
+
+    fa3_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+    assert fa3_path == DiffusionAttentionBackendEnum.FLASH_ATTN_3_HUB.get_path()
+    fa2_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
+    assert fa2_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
+    assert versions_tried == [1, 2, 1, 2]
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_explicit_flash_attention_unavailable_does_not_fallback(monkeypatch: pytest.MonkeyPatch):

tests/diffusion/attention/test_ulysses_uaa.py

@@ -27,6 +27,202 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_uaa_equal_rank_contract_skips_length_gather(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    class FakeGroup:
+        ulysses_group = object()
+        ulysses_world_size = 2
+        ulysses_rank = 0
+        ring_world_size = 1
+
+    observed = []
+
+    def fake_all_to_all(pg, tensor, **kwargs):
+        observed.append((kwargs["seq_lens"], kwargs["padded_head_cnt"]))
+        return tensor, tensor.shape[2]
+
+    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    monkeypatch.setattr(
+        ulysses,
+        "_all_gather_int",
+        lambda *args, **kwargs: pytest.fail("length gather must be skipped"),
+    )
+    strategy = ulysses.UlyssesParallelAttention(
+        FakeGroup(),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+
+    with set_forward_context():
+        get_forward_context()._sp_equal_pad_stack.append(True)
+        strategy.pre_attention(
+            torch.zeros(1, 3, 28, 4),
+            torch.zeros(1, 3, 7, 4),
+            torch.zeros(1, 3, 7, 4),
+            None,
+        )
+
+    assert observed == [([3, 3], 32), ([3, 3], 8), ([3, 3], 8)]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_uaa_hybrid_equal_rank_contract_skips_ulysses_gather_keeps_ring_check(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Contract + ring: the fast path skips only the Ulysses gather.
+
+    The ring equality check is a separate collective on ring_group and must
+    still run, receiving s_global derived from the contract.
+    """
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    class FakeGroup:
+        ulysses_group = object()
+        ring_group = object()
+        ulysses_world_size = 2
+        ulysses_rank = 0
+        ring_world_size = 2
+
+    sp_group = FakeGroup()
+    observed = []
+    ring_checks = []
+
+    def fake_all_to_all(pg, tensor, **kwargs):
+        observed.append((kwargs["seq_lens"], kwargs["padded_head_cnt"]))
+        return tensor, tensor.shape[2]
+
+    def fake_gather_int(pg, value, *, device):
+        if pg is sp_group.ulysses_group:
+            pytest.fail("Ulysses length gather must be skipped under the contract")
+        assert pg is sp_group.ring_group
+        ring_checks.append(int(value))
+        return [int(value), int(value)]
+
+    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
+    monkeypatch.setattr(ulysses, "_all_gather_int", fake_gather_int)
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    strategy = ulysses.UlyssesParallelAttention(
+        sp_group,
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+
+    with set_forward_context():
+        get_forward_context()._sp_equal_pad_stack.append(True)
+        strategy.pre_attention(
+            torch.zeros(1, 3, 28, 4),
+            torch.zeros(1, 3, 7, 4),
+            torch.zeros(1, 3, 7, 4),
+            None,
+        )
+
+    # Ring check ran: s_global = ulysses_world_size * local = 2 * 3.
+    assert ring_checks == [6]
+    assert observed == [([3, 3], 32), ([3, 3], 8), ([3, 3], 8)]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_uaa_without_contract_keeps_dynamic_length_gather(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """No auto_pad boundary means shards may differ, so the gather must stay."""
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    class FakeGroup:
+        ulysses_group = object()
+        ulysses_world_size = 2
+        ulysses_rank = 0
+        ring_world_size = 1
+
+    observed = []
+    monkeypatch.setattr(
+        ulysses,
+        "_all_gather_int",
+        lambda *args, **kwargs: [3, 4],
+    )
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+
+    def fake_all_to_all(pg, tensor, **kwargs):
+        observed.append(kwargs["seq_lens"])
+        return tensor, tensor.shape[2]
+
+    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
+    strategy = ulysses.UlyssesParallelAttention(
+        FakeGroup(),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+
+    with set_forward_context():
+        strategy.pre_attention(
+            torch.zeros(1, 3, 28, 4),
+            torch.zeros(1, 3, 7, 4),
+            torch.zeros(1, 3, 7, 4),
+            None,
+        )
+
+    assert observed == [[3, 4], [3, 4], [3, 4]]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_uaa_mixed_boundaries_keep_the_dynamic_length_gather(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """One manual boundary inside an auto_pad region voids the contract."""
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    class FakeGroup:
+        ulysses_group = object()
+        ulysses_world_size = 2
+        ulysses_rank = 0
+        ring_world_size = 1
+
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    gathered = []
+
+    def fake_gather(*args, **kwargs):
+        gathered.append(args)
+        return [3, 4]
+
+    monkeypatch.setattr(ulysses, "_all_gather_int", fake_gather)
+    monkeypatch.setattr(
+        ulysses,
+        "_ulysses_all_to_all_any_qkv",
+        lambda pg, tensor, **kwargs: (tensor, tensor.shape[2]),
+    )
+    strategy = ulysses.UlyssesParallelAttention(
+        FakeGroup(),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+
+    with set_forward_context():
+        ctx = get_forward_context()
+        ctx._sp_equal_pad_stack.extend([True, False])
+        assert not ctx.sp_rank_local_seq_lens_equal
+        strategy.pre_attention(
+            torch.zeros(1, 3, 28, 4),
+            torch.zeros(1, 3, 7, 4),
+            torch.zeros(1, 3, 7, 4),
+            None,
+        )
+
+    assert len(gathered) == 1
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_uaa_gqa_head_padding_preserves_the_query_to_kv_ratio(
@@ -634,6 +830,187 @@ def test_ulysses_uaa_matches_baseline(
                 pass
 
 
+def _run_contract_skip_case(
+    local_rank: int,
+    world_size: int,
+    init_method: str,
+    input_file: str,
+    num_heads: int,
+    head_size: int,
+    num_kv_heads: int | None,
+) -> None:
+    """Per-rank worker: slow path vs contract fast path on real rank pairs.
+
+    Phase 1 leaves the equal-pad stack empty so the Ulysses length all-gather
+    really runs (counted); phase 2 pushes the split hook's auto_pad marker and
+    replaces _all_gather_int with a raising guard. With ring_degree=1 any call
+    to it is the Ulysses length gather, so the guard proves the skip, and the
+    per-rank output comparison proves the fast path is numerically equivalent
+    to the slow path it replaces.
+    """
+    from vllm_omni.diffusion.attention.parallel import ulysses as ulysses_mod
+
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+
+    os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
+
+    init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method)
+    initialize_model_parallel(
+        data_parallel_size=1,
+        cfg_parallel_size=1,
+        sequence_parallel_size=world_size,
+        ulysses_degree=world_size,
+        ring_degree=1,
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+    )
+
+    parallel_config = DiffusionParallelConfig(
+        pipeline_parallel_size=1,
+        data_parallel_size=1,
+        tensor_parallel_size=1,
+        sequence_parallel_size=world_size,
+        ulysses_degree=world_size,
+        ring_degree=1,
+        cfg_parallel_size=1,
+        ulysses_mode="advanced_uaa",
+    )
+    od_config = OmniDiffusionConfig(model="test", dtype=torch.float32, parallel_config=parallel_config)
+
+    with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
+        attn = Attention(
+            num_heads=num_heads,
+            head_size=head_size,
+            causal=False,
+            softmax_scale=1.0 / (head_size**0.5),
+            num_kv_heads=num_kv_heads,
+        ).to(device=device, dtype=torch.float32)
+
+        with np.load(input_file, allow_pickle=False) as payload:
+            q_full = torch.from_numpy(payload["q"]).to(device=device)
+            k_full = torch.from_numpy(payload["k"]).to(device=device)
+            v_full = torch.from_numpy(payload["v"]).to(device=device)
+
+        def make_shards() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            # Fresh contiguous shards for each phase; seq_len is divisible by
+            # world_size, so every rank gets the same local length.
+            return (
+                torch.tensor_split(q_full, world_size, dim=1)[local_rank].contiguous(),
+                torch.tensor_split(k_full, world_size, dim=1)[local_rank].contiguous(),
+                torch.tensor_split(v_full, world_size, dim=1)[local_rank].contiguous(),
+            )
+
+        # The Attention layer only enables SP communication when
+        # ForwardContext.sp_active is True; in production the split hook sets
+        # this (and pushes the auto_pad equal-length marker onto
+        # _sp_equal_pad_stack). Here we drive it directly.
+        ctx = get_forward_context()
+        ctx._sp_shard_depth = 1
+
+        def make_sdp_ctx() -> object:
+            # Fresh instance per phase: the @contextmanager instance releases
+            # its state on first exit and cannot be re-entered.
+            return torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
+
+        # Phase 1: contract absent -> the length all-gather must really run.
+        real_gather = ulysses_mod._all_gather_int
+        gather_calls: list[object] = []
+
+        def counting_gather(pg: object, value: int, *, device: torch.device) -> list[int]:
+            gather_calls.append(pg)
+            return real_gather(pg, value, device=device)
+
+        ulysses_mod._all_gather_int = counting_gather
+        try:
+            q, k, v = make_shards()
+            with make_sdp_ctx():
+                out_slow = attn(q, k, v, attn_metadata=None).contiguous()
+        finally:
+            ulysses_mod._all_gather_int = real_gather
+        assert len(gather_calls) >= 1, "slow path must call the Ulysses length all-gather"
+
+        # Phase 2: replicate the split hook's auto_pad marker so the fast
+        # path fires; any _all_gather_int call in this phase is the Ulysses
+        # length gather and must not happen.
+        ctx._sp_equal_pad_stack.append(True)
+
+        def fail_on_gather(pg: object, value: int, *, device: torch.device) -> list[int]:
+            raise AssertionError("the Ulysses length all-gather must be skipped under the equal-pad contract")
+
+        ulysses_mod._all_gather_int = fail_on_gather
+        try:
+            q, k, v = make_shards()
+            with make_sdp_ctx():
+                out_fast = attn(q, k, v, attn_metadata=None).contiguous()
+        finally:
+            ulysses_mod._all_gather_int = real_gather
+        ctx._sp_equal_pad_stack.pop()
+
+        torch.testing.assert_close(out_fast, out_slow, atol=1e-5, rtol=1e-5)
+
+    destroy_distributed_env()
+
+
+@pytest.mark.parametrize(
+    "num_heads,num_kv_heads",
+    [
+        (3, None),  # MHA: heads padded 3 -> 4 by ulysses_degree=2
+        (28, 7),  # GQA: derived-Q padding 28/7 -> 32/8, the main case for the skip
+    ],
+)
+def test_ulysses_uaa_contract_skips_length_gather_matches_slow_path(num_heads: int, num_kv_heads: int | None) -> None:
+    """Real 2-rank parity: contract fast path == slow path, gather skipped.
+
+    The CPU-only contract tests replace the all-to-all with a fake; this test
+    runs the real strategy on two ranks. Each rank first computes the slow
+    path (empty equal-pad stack, real length all-gather), then the fast path
+    with the split hook's auto_pad marker pushed and _all_gather_int replaced
+    by a raising guard -- which also proves the collective was truly skipped,
+    not just silently replaced.
+    """
+    sp_world_size = 2
+    if current_omni_platform.get_device_count() < sp_world_size:
+        pytest.skip(f"Test requires {sp_world_size} GPUs")
+
+    batch_size = 2
+    head_size = 8
+    seq_len = 6  # divisible by 2 so both ranks hold equal-length shards
+
+    init_method = get_distributed_init_method()
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
+        input_file = f_in.name
+
+    try:
+        torch.manual_seed(0)
+        kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        np.savez(
+            input_file,
+            q=torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32).numpy(),
+            k=torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32).numpy(),
+            v=torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32).numpy(),
+        )
+
+        torch.multiprocessing.spawn(
+            _run_contract_skip_case,
+            args=(
+                sp_world_size,
+                init_method,
+                input_file,
+                num_heads,
+                head_size,
+                num_kv_heads,
+            ),
+            nprocs=sp_world_size,
+        )
+    finally:
+        try:
+            os.remove(input_file)
+        except OSError:
+            pass
+
+
 @pytest.mark.parametrize(
     "num_heads,num_kv_heads,joint_len",
     [

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -279,17 +279,17 @@ def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]
         group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None))
         return [[group], [group]], "cg-support", [4, 8]
 
+    # vLLM 0.29 dropped attn_groups and cache_dtype from init_kv_cache: the
+    # layout now rides on the resolved CacheConfig.kv_cache_layout instead.
     def init_cache(
         runner_kv_caches,
         forward_context,
         kv_cache_config,
-        attn_groups,
         device,
-        cache_dtype,
         kernel_block_sizes,
         vllm_config,
     ):
-        events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes))
+        events.append(("cache", kv_cache_config, kernel_block_sizes))
         runner_kv_caches.extend(["cache-0", "cache-1"])
         return {"layer-0": "tensor-0", "layer-1": "tensor-1"}
 

tests/diffusion/diffusion_kv/test_initialization.py

@@ -299,8 +299,10 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
     )
 
     assert vllm_config.model_config.original_max_model_len == -1
-    # Sixteen physical pages provide ``16 * block_size`` tokens.  The native
-    # backend chooses the block geometry, so keep this assertion backend
+    # Sixteen physical pages provide ``(16 - 1) * block_size`` tokens: vLLM 0.29
+    # subtracts the null block that BlockPool permanently holds back before
+    # sizing (check_enough_kv_cache_memory), so only 15 pages are usable.  The
+    # native backend chooses the block geometry, so keep this assertion backend
     # independent instead of assuming the generic 16-token block size.
-    expected_max_model_len = min(model_limit, spec.block_size * 16)
+    expected_max_model_len = min(model_limit, spec.block_size * (16 - 1))
     assert vllm_config.model_config.max_model_len == expected_max_model_len

tests/diffusion/diffusion_kv/test_layout.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Physical KV cache layout selection and propagation for the Diffusion cache."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.config import CacheConfig
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.kv_cache_layout import KVCacheLayout
+
+from vllm_omni.diffusion.diffusion_kv.layout import (
+    DEFAULT_DIFFUSION_KV_CACHE_LAYOUT,
+    adopt_kv_cache_layout,
+    assert_backend_layout_supported,
+    build_kv_cache_tensor,
+    resolve_diffusion_kv_cache_layout,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _config():
+    return SimpleNamespace(cache_config=CacheConfig())
+
+
+class _DenseBackend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return False
+
+
+class _BlockStrideBackend:
+    @classmethod
+    def indexes_kv_by_block_stride(cls) -> bool:
+        return True
+
+
+def test_default_layout_is_not_block_outermost() -> None:
+    """The default reproduces the pre-0.29 indexes_kv_by_block_stride=False."""
+    vllm_config = _config()
+    layout = resolve_diffusion_kv_cache_layout(vllm_config)
+
+    assert layout is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+    assert layout.is_block_outermost is False
+    assert vllm_config.cache_config.get_resolved_kv_cache_layout() is layout
+
+
+def test_block_stride_backend_selects_block_outermost_layout() -> None:
+    layout = resolve_diffusion_kv_cache_layout(_config(), indexes_kv_by_block_stride=True)
+
+    assert layout.is_block_outermost is True
+
+
+def test_resolution_is_idempotent() -> None:
+    vllm_config = _config()
+
+    assert resolve_diffusion_kv_cache_layout(vllm_config) is resolve_diffusion_kv_cache_layout(vllm_config)
+
+
+def test_contradicting_requirement_raises_instead_of_misaddressing() -> None:
+    vllm_config = _config()
+    resolve_diffusion_kv_cache_layout(vllm_config)
+
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True)
+
+
+def test_preset_layout_is_honoured() -> None:
+    vllm_config = _config()
+    vllm_config.cache_config.kv_cache_layout = KVCacheLayout.BLNHC.name
+
+    layout = resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True)
+
+    assert layout is KVCacheLayout.BLNHC
+
+
+def test_backend_assertion_is_quiet_until_the_layout_is_resolved() -> None:
+    """Workers collect specs before the control plane hands them the layout."""
+    vllm_config = _config()
+
+    assert_backend_layout_supported(vllm_config, _BlockStrideBackend)  # unresolved: nothing to check
+
+    resolve_diffusion_kv_cache_layout(vllm_config)
+    assert_backend_layout_supported(vllm_config, _DenseBackend)
+    assert_backend_layout_supported(vllm_config, None)
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        assert_backend_layout_supported(vllm_config, _BlockStrideBackend)
+
+
+def test_worker_adopts_the_layout_carried_on_the_kv_cache_config() -> None:
+    """Diffusion workers receive only a KVCacheConfig, so the name rides on it."""
+    worker = _config()
+    shipped = SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name)
+
+    assert adopt_kv_cache_layout(worker, shipped) is KVCacheLayout.BLNHC
+    assert worker.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.BLNHC
+
+
+def test_worker_without_a_carried_layout_falls_back_to_the_same_default() -> None:
+    worker = _config()
+
+    assert adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=None)) is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
+
+
+def test_worker_disagreeing_with_the_control_plane_raises() -> None:
+    worker = _config()
+    resolve_diffusion_kv_cache_layout(worker)
+
+    with pytest.raises(ValueError, match="disagrees with the layout"):
+        adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name))
+
+
+def test_block_stride_requirement_propagates_control_plane_to_worker() -> None:
+    """A block-stride backend must reach the worker as a block-outermost layout.
+
+    Covers the whole seam: the control plane resolves with the requirement, the
+    name rides on each rank-local KVCacheConfig, the worker adopts it, and the
+    backend contract check then passes instead of raising.
+    """
+    control_plane = _config()
+    layout = resolve_diffusion_kv_cache_layout(control_plane, indexes_kv_by_block_stride=True)
+    assert layout.is_block_outermost is True
+
+    shipped = SimpleNamespace(kv_cache_layout=layout.name)
+
+    worker = _config()
+    assert adopt_kv_cache_layout(worker, shipped) is layout
+    # The worker's own backend requirement now agrees with what it was handed.
+    assert_backend_layout_supported(worker, _BlockStrideBackend)
+    with pytest.raises(ValueError, match="contradicts the attention"):
+        assert_backend_layout_supported(worker, _DenseBackend)
+
+
+def test_tensor_strides_come_from_the_upstream_derivation() -> None:
+    spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
+
+    tensor = build_kv_cache_tensor(spec, 4, ["layer0", "layer1"])
+
+    assert tensor.layers == ["layer0", "layer1"]
+    assert tensor.size == spec.page_size_bytes * 2 * 4
+    # Layer-outermost: each layer owns a contiguous run of its own blocks.
+    assert tensor.layer_stride == spec.page_size_bytes * 4
+    assert tensor.block_stride == spec.page_size_bytes

tests/diffusion/diffusion_kv/test_manager.py

@@ -1,11 +1,12 @@
 # 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
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
+from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 
@@ -24,7 +25,7 @@ def _config(num_blocks: int) -> KVCacheConfig:
     )
     return KVCacheConfig(
         num_blocks=num_blocks,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
     )
 

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -15,6 +15,7 @@ import torch
 from torch import nn
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 
+from tests.helpers.kv_layout import layout_for_backend
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
@@ -1239,7 +1240,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
     assert native_layer.num_heads == 4
     assert native_layer.num_kv_heads == 2
     assert native_layer.spec.num_kv_heads == 2
-    assert native_layer.spec.indexes_kv_by_block_stride is True
+    # 0.29: the flag no longer rides on the spec, it selects the layout.
+    assert layout_for_backend(native_layer.attn_backend).is_block_outermost is True
     assert native_layer._q_scale_float == 1.0
     assert native_layer._k_scale_float == 1.0
     assert native_layer._v_scale_float == 1.0

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.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
 
@@ -10,12 +10,14 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
 from vllm.v1.worker.gpu.block_table import BlockTables
 
+from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout
 from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
     DiffusionPagedAttentionAdapter,
     DiffusionPagedAttentionLayerAdapter,
@@ -88,12 +90,7 @@ 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=[
-            KVCacheTensor(
-                size=canonical_spec.page_size_bytes * 4,
-                shared_by=[_LAYER_NAME],
-            )
-        ],
+        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME])],
         kv_cache_groups=[
             KVCacheGroupSpec(
                 layer_names=[_LAYER_NAME],
@@ -101,6 +98,9 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
             )
         ],
     )
+    # 0.29 reads the resolved physical layout when allocating; the diffusion
+    # control plane pins it, so this standalone harness must too.
+    resolve_diffusion_kv_cache_layout(vllm_config)
     with set_current_vllm_config(vllm_config):
         attn_groups, _, kernel_block_sizes = init_attn_backend(
             kv_cache_config,
@@ -108,13 +108,13 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
             device,
         )
         runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+        # vLLM 0.29 dropped attn_groups and cache_dtype here: the layout now
+        # comes from the resolved CacheConfig.kv_cache_layout inside the call.
         init_kv_cache(
             runner_kv_caches,
             vllm_config.compilation_config.static_forward_context,
             kv_cache_config,
-            attn_groups,
             device,
-            vllm_config.cache_config.cache_dtype,
             kernel_block_sizes,
             vllm_config,
         )

tests/diffusion/diffusion_kv/test_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
 
 import pytest
 import torch
@@ -10,10 +10,10 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
-    KVCacheTensor,
 )
 from vllm.v1.request import RequestStatus
 
+from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -119,7 +119,7 @@ def _manager(*, enable_caching: bool) -> KVCacheManager:
     group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["layer0"])],
         kv_cache_groups=[group],
     )
     return KVCacheManager(

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -8,9 +8,10 @@ from unittest.mock import Mock
 import pytest
 import torch
 from torch import nn
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
 import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
+from tests.helpers.kv_layout import build_kv_cache_tensor, layout_for_backend
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
@@ -79,7 +80,9 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None:
     assert spec.num_kv_heads == 2
     assert spec.head_size == 8
     assert spec.dtype is torch.bfloat16
-    assert spec.indexes_kv_by_block_stride is True
+    # vLLM 0.29 moved the block-stride decision off the spec onto the resolved
+    # physical layout; the backend's declaration must still select one.
+    assert layout_for_backend(_Backend).is_block_outermost is True
     assert spec.non_causal is True
 
 
@@ -112,7 +115,7 @@ def test_runner_retains_matching_rank_local_config() -> None:
     spec = runner.get_kv_cache_spec()["image_attention"]
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["image_attention"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
     )
 
@@ -126,7 +129,7 @@ def test_runner_rejects_rank_local_config_for_different_layers() -> None:
     spec = runner.get_kv_cache_spec()["image_attention"]
     config = KVCacheConfig(
         num_blocks=8,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["other_attention"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
     )
 

tests/diffusion/distributed/test_sp_plan_hooks.py

@@ -1113,6 +1113,120 @@ class TestStrictModeSplitValidation:
                 hook._prepare_sp_input(x, metadata["x"], (), {})
 
 
+@pytest.mark.cpu
+class TestEqualPadContract:
+    """The stack of per-boundary auto_pad flags that lets attention skip a gather.
+
+    A boundary is recorded as equal-length only when the framework padded it,
+    because only then is every rank's shard guaranteed the same size.
+    """
+
+    def _split_hook(self, *, auto_pad: bool, monkeypatch):
+        from vllm_omni.diffusion.attention import selector
+        from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelSplitHook,
+        )
+
+        class MaskBackend:
+            supports_attention_mask = True
+
+            @staticmethod
+            def get_name():
+                return "test"
+
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 1)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(selector, "get_attn_backend_for_role", lambda **_: (MaskBackend(), None))
+        # sp_shard, used by the manual path, binds these at import time.
+        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: 1)
+
+        hook = SequenceParallelSplitHook(
+            {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    expected_dims=3,
+                    split_output=True,
+                    auto_pad=auto_pad,
+                )
+            },
+            SequenceParallelConfig(ulysses_degree=2),
+        )
+        hook.initialize_hook(nn.Identity())
+        return hook
+
+    def test_auto_pad_split_claims_equal_rank_lengths(self, monkeypatch):
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            set_forward_context,
+        )
+
+        hook = self._split_hook(auto_pad=True, monkeypatch=monkeypatch)
+        with set_forward_context():
+            hook.pre_forward(nn.Identity())
+            shard = hook.post_forward(nn.Identity(), torch.arange(5).reshape(1, 5, 1))
+            ctx = get_forward_context()
+
+            # 5 tokens padded to 6, so both ranks hold 3.
+            assert shard.shape[1] == 3
+            assert ctx._sp_equal_pad_stack == [True]
+            assert ctx.sp_rank_local_seq_lens_equal
+
+    def test_manual_split_does_not_claim_equal_rank_lengths(self, monkeypatch):
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            set_forward_context,
+        )
+
+        hook = self._split_hook(auto_pad=False, monkeypatch=monkeypatch)
+        with set_forward_context():
+            hook.pre_forward(nn.Identity())
+            hook.post_forward(nn.Identity(), torch.arange(6).reshape(1, 6, 1))
+            ctx = get_forward_context()
+
+            assert ctx._sp_equal_pad_stack == [False]
+            assert not ctx.sp_rank_local_seq_lens_equal
+
+    def test_gather_pops_only_its_own_boundary(self, monkeypatch):
+        """A stack, not a counter: the gather must not resurrect the contract.
+
+        With one auto_pad and one manual boundary open, a counter could not
+        tell which one a gather closes and would wrongly report equal lengths.
+        """
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            set_forward_context,
+        )
+        from vllm_omni.diffusion.hooks import sequence_parallel
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelGatherHook,
+        )
+
+        monkeypatch.setattr(
+            sequence_parallel,
+            "sp_gather",
+            lambda tensor, dim, validate: torch.cat([tensor, tensor], dim=dim),
+        )
+        hook = SequenceParallelGatherHook(
+            SequenceParallelOutput(gather_dim=1, expected_dims=3),
+            SequenceParallelConfig(ulysses_degree=2),
+        )
+
+        with set_forward_context():
+            ctx = get_forward_context()
+            ctx._sp_shard_depth = 2
+            ctx._sp_equal_pad_stack.extend([False, True])
+            hook.post_forward(nn.Identity(), torch.zeros(1, 3, 4))
+
+            assert ctx._sp_shard_depth == 1
+            assert ctx._sp_equal_pad_stack == [False]
+            assert not ctx.sp_rank_local_seq_lens_equal
+
+
 @pytest.mark.cpu
 class TestSequenceParallelConfig:
     """Test SequenceParallelConfig dataclass."""

tests/diffusion/distributed/test_wan_chunked_decode.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import weakref
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from diffusers.models.autoencoders import AutoencoderKLWan
+
+from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+    DistributedAutoencoderKLWan,
+    OmniAutoencoderKLWan,
+)
+from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture
+def vae():
+    # Real causal convolutions and caches, with small random weights for CPU CI.
+    return OmniAutoencoderKLWan(
+        base_dim=4,
+        z_dim=2,
+        dim_mult=[1, 2],
+        num_res_blocks=1,
+        temperal_downsample=[True],
+        scale_factor_spatial=2,
+        scale_factor_temporal=2,
+    ).eval()
+
+
+@pytest.mark.parametrize("patch_size", [None, 2])
+@pytest.mark.parametrize("return_dict", [False, True])
+@torch.no_grad()
+def test_chunked_decode_matches_diffusers(vae, patch_size, return_dict):
+    if patch_size is not None:
+        vae = OmniAutoencoderKLWan.from_config(vae.config, patch_size=patch_size, out_channels=12).eval()
+    z = torch.randn(1, 2, 3, 2, 2)
+    expected = AutoencoderKLWan._decode(vae, z).sample
+    actual = vae.decode(z, return_dict=return_dict)
+    torch.testing.assert_close(actual.sample if return_dict else actual[0], expected, rtol=0, atol=0)
+    chunks: list[torch.Tensor] = []
+    assert supports_chunked_vae_decode(vae)
+    assert not supports_chunked_vae_decode(object())
+    assert vae.decode_with_chunks(z, on_chunk=chunks.append) is None
+    assert [chunk.shape[2] for chunk in chunks] == [1, 2, 2]
+    torch.testing.assert_close(torch.cat(chunks, dim=2), expected, rtol=0, atol=0)
+    assert all(entry is None for entry in vae._feat_map)
+
+
+@torch.no_grad()
+def test_slicing_and_forward_hooks(vae):
+    z = torch.randn(2, 2, 2, 2, 2)
+    vae.enable_slicing()
+    batches = []
+    handle = vae.post_quant_conv.register_forward_pre_hook(lambda module, args: batches.append(args[0].shape[0]))
+    vae._hf_hook = SimpleNamespace(pre_forward=Mock())
+    try:
+        assert vae.decode(z).sample.shape[0] == 2
+        assert batches == [1, 1]
+        vae._hf_hook.pre_forward.assert_called_once_with(vae)
+        batches.clear()
+        with pytest.raises(ValueError, match="use_slicing"):
+            vae.decode_with_chunks(z, on_chunk=Mock())
+        assert batches == []
+        vae._hf_hook.pre_forward.reset_mock()
+        vae.decode_with_chunks(z[:1], on_chunk=Mock())
+        assert batches == [1]
+        vae._hf_hook.pre_forward.assert_called_once_with(vae)
+    finally:
+        handle.remove()
+
+
+@torch.no_grad()
+def test_callback_failure_drains_and_clears_cache(vae):
+    z = torch.randn(1, 2, 4, 2, 2)
+    decoded = Mock(return_value=None)
+    handle = vae.decoder.register_forward_hook(lambda *args: decoded())
+    error = RuntimeError("consumer failed")
+    callback = Mock(side_effect=error)
+    try:
+        with pytest.raises(RuntimeError, match="consumer failed") as exc:
+            vae.decode_with_chunks(z, on_chunk=callback)
+        assert exc.value is error
+        assert callback.call_count == 1
+        assert decoded.call_count == z.shape[2]
+        assert all(entry is None for entry in vae._feat_map)
+    finally:
+        handle.remove()
+    # The same instance can decode again after the failure.
+    chunks: list[torch.Tensor] = []
+    vae.decode_with_chunks(z, on_chunk=chunks.append)
+    torch.testing.assert_close(torch.cat(chunks, 2), vae.decode(z).sample, rtol=0, atol=0)
+
+
+def test_chunked_decode_rejects_spatial_tiling(vae):
+    vae.enable_tiling(tile_sample_min_height=2, tile_sample_min_width=2)
+    with pytest.raises(ValueError, match="spatial tiling"):
+        vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=Mock())
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+@pytest.mark.parametrize("fail_callback", [False, True])
+def test_spatial_callback_ownership_and_bounded_retention(monkeypatch, rank, fail_callback):
+    # Isolate the temporal loop from halo kernels; weakrefs detect retention of
+    # decoded GPU-sized tensors even after the consumer has failed.
+    monkeypatch.setattr(wan_spatial_shard, "install_wan_spatial_shard_decode", lambda *args, **kwargs: None)
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2))
+    references: list[weakref.ReferenceType[torch.Tensor]] = []
+
+    def decode(x, **kwargs):
+        assert sum(ref() is not None for ref in references) <= 1
+        chunk = x.clone()
+        references.append(weakref.ref(chunk))
+        return chunk
+
+    vae = SimpleNamespace(
+        config=SimpleNamespace(patch_size=None),
+        clear_cache=Mock(),
+        post_quant_conv=lambda z: z,
+        decoder=decode,
+        _feat_map=[],
+    )
+    callback = Mock(side_effect=RuntimeError("consumer failed") if fail_callback else None)
+    z = torch.zeros(1, 3, 6, 2, 2)
+    if rank == 0 and fail_callback:
+        with pytest.raises(RuntimeError, match="consumer failed"):
+            wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback)
+    else:
+        assert wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback) is None
+    assert len(references) == z.shape[2]
+    assert callback.call_count == (0 if rank else 1 if fail_callback else z.shape[2])
+    assert vae.clear_cache.call_count == 2
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+@torch.no_grad()
+def test_distributed_small_latent_callback_ownership(vae, monkeypatch, rank):
+    vae.is_distributed_enabled = lambda: True
+    vae.distributed_executor = SimpleNamespace(group=None)
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda **kwargs: rank)
+    callback = Mock()
+    vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=callback)
+    assert callback.call_count == (2 if rank == 0 else 0)
+
+
+@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width", "tile"])
+def test_distributed_chunked_routing(monkeypatch, mode):
+    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
+    torch.nn.Module.__init__(vae)
+    vae.use_tiling = True
+    vae.use_slicing = False
+    vae.tile_sample_min_height = vae.tile_sample_min_width = 2
+    vae.spatial_compression_ratio = 2
+    vae.distributed_executor = SimpleNamespace(group=None, parallel_mode=mode, parallel_size=2)
+    monkeypatch.setattr(vae, "is_distributed_enabled", lambda: True)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: 2)
+    decode = Mock()
+    monkeypatch.setattr(wan_spatial_shard, "spatial_shard_decode", decode)
+    callback = Mock()
+    z = torch.zeros(1, 2, 2, 2, 2)
+    if mode == "tile":
+        with pytest.raises(ValueError, match="require spatial-shard"):
+            vae.decode_with_chunks(z, on_chunk=callback)
+        decode.assert_not_called()
+    else:
+        assert vae.decode_with_chunks(z, on_chunk=callback) is None
+        decode.assert_called_once_with(
+            vae, z, group=None, return_dict=True, split_dim=mode.removeprefix("spatial_shard_"), on_chunk=callback
+        )

tests/diffusion/layers/test_activation.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from tests.helpers.mark import hardware_test
+from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.layers.activation import SiluAndMul
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -19,9 +19,17 @@ def test_silu_and_mul_native_matches_packed_reference() -> None:
     torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up)
 
 
-@hardware_test(res={"npu": "A3"}, num_cards=1)
-def test_silu_and_mul_npu_matches_packed_reference() -> None:
-    packed = torch.randn(257, 256, device="npu", dtype=torch.bfloat16)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A2"}, num_cards=1), id="A2-bf16"),
+        pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A3"}, num_cards=1), id="A3-bf16"),
+        pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A5"}, num_cards=1), id="A5-bf16"),
+        pytest.param(torch.float16, marks=hardware_marks(res={"npu": "310P"}, num_cards=1), id="310P-fp16"),
+    ],
+)
+def test_silu_and_mul_npu_matches_packed_reference(dtype: torch.dtype) -> None:
+    packed = torch.randn(257, 256, device="npu", dtype=dtype)
     gate, up = packed.chunk(2, dim=-1)
 
     torch.testing.assert_close(

tests/diffusion/layers/test_adalayernorm.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 shared AdaLayerNorm layers used by FLUX and other models."""
 
 import os
@@ -46,7 +46,7 @@ def _force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/layers/test_rope_broadcast.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
 """
 Shape broadcasting regression tests for RotaryEmbedding.forward_native.
 
@@ -16,7 +16,11 @@ import types
 import pytest
 import torch
 
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rotary_emb_mindiesd
+from vllm_omni.diffusion.layers.rope import (
+    RotaryEmbedding,
+    RotaryEmbeddingWan,
+    apply_rotary_emb_mindiesd,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -239,6 +243,25 @@ class TestRotaryEmbeddingNativeShapeRegression:
         assert not torch.isnan(output).any()
 
 
+def test_cuda_paths_fall_back_without_vendored_rotary_kernel() -> None:
+    """Source/CPU installs need not contain the wheel-vendored CUDA kernel."""
+    x = torch.randn(1, 5, 2, 8)
+
+    rope = RotaryEmbedding(is_neox_style=True)
+    rope.apply_rotary_emb_vllm_flash_attn = None
+    output = rope.forward_cuda(x, torch.randn(5, 4), torch.randn(5, 4))
+    assert output.shape == x.shape
+
+    wan_rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
+    wan_rope.apply_rotary_emb_vllm_flash_attn = None
+    output = wan_rope.forward_cuda(
+        x,
+        torch.randn(1, 5, 1, 4),
+        torch.randn(1, 5, 1, 4),
+    )
+    assert output.shape == x.shape
+
+
 def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None:
     """H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout."""
     captured: dict[str, object] = {}

tests/diffusion/lora/helpers.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
 """Shared test helpers for diffusion LoRA tests."""
 
 from __future__ import annotations
@@ -26,15 +26,39 @@ class DummyBaseLayerWithLoRA(torch.nn.Module):
             tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
         ] = []
         self.reset_calls: int = 0
+        self.suspend_calls: int = 0
+        self.resume_calls: int = 0
+        self.active_slices: tuple[bool, ...] = ()
+        self.suspended_slices: tuple[bool, ...] | None = None
         self.create_calls: int = 0
 
     def set_lora(self, index: int, lora_a, lora_b):
         assert index == 0
         self.set_calls.append((lora_a, lora_b))
+        if isinstance(lora_b, list):
+            self.active_slices = tuple(b is not None for b in lora_b)
+        else:
+            self.active_slices = (True,)
+        self.suspended_slices = None
 
     def reset_lora(self, index: int):
         assert index == 0
         self.reset_calls += 1
+        self.active_slices = ()
+        self.suspended_slices = None
+
+    def suspend_lora(self) -> None:
+        if self.suspended_slices is not None:
+            return
+        self.suspended_slices = self.active_slices
+        self.active_slices = (False,) * len(self.active_slices)
+        self.suspend_calls += 1
+
+    def resume_lora(self) -> None:
+        if self.suspended_slices is not None:
+            self.active_slices = self.suspended_slices
+            self.suspended_slices = None
+        self.resume_calls += 1
 
     def create_lora_weights(self, max_loras, lora_config, model_config):
         self.create_calls += 1

tests/diffusion/lora/test_lora_manager.py

@@ -30,14 +30,38 @@ class _DummyLoRALayer:
             tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
         ] = []
         self.reset_calls: int = 0
+        self.suspend_calls: int = 0
+        self.resume_calls: int = 0
+        self.active_slices: tuple[bool, ...] = ()
+        self.suspended_slices: tuple[bool, ...] | None = None
 
     def set_lora(self, index: int, lora_a, lora_b):
         assert index == 0
         self.set_calls.append((lora_a, lora_b))
+        if isinstance(lora_b, list):
+            self.active_slices = tuple(b is not None for b in lora_b)
+        else:
+            self.active_slices = (True,)
+        self.suspended_slices = None
 
     def reset_lora(self, index: int):
         assert index == 0
         self.reset_calls += 1
+        self.active_slices = ()
+        self.suspended_slices = None
+
+    def suspend_lora(self) -> None:
+        if self.suspended_slices is not None:
+            return
+        self.suspended_slices = self.active_slices
+        self.active_slices = (False,) * len(self.active_slices)
+        self.suspend_calls += 1
+
+    def resume_lora(self) -> None:
+        if self.suspended_slices is not None:
+            self.active_slices = self.suspended_slices
+            self.suspended_slices = None
+        self.resume_calls += 1
 
 
 # Aliases for backward compatibility within this file
@@ -723,3 +747,104 @@ def test_lora_manager_discovers_unet_component(monkeypatch):
     assert "unet.down_block.proj" in manager._lora_modules
     # Verify the module was actually replaced in the tree (not just recorded)
     assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA)
+
+
+def _suspend_harness(monkeypatch, adapter_id: int = 7, rank: int = 2):
+    """Manager wired to a single dummy layer, ready to activate `adapter_id`."""
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    lora_model = _DummyLM(rank=rank)
+    manager = DiffusionLoRAManager(
+        pipeline=_DummyPipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+    monkeypatch.setattr(manager, "_load_adapter", lambda _req: (lora_model, type("PH", (), {"r": rank})()))
+    manager._registered_adapters = {adapter_id: lora_model}
+    manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
+    return manager, lora_model.transformer.foo
+
+
+def test_reactivating_same_adapter_resumes_without_rebinding(monkeypatch):
+    """activate -> deactivate -> activate at the same scale must re-arm the
+    saved mask instead of re-uploading every layer."""
+    adapter_id = 7
+    req = _dummy_lora_request(adapter_id)
+    manager, layer = _suspend_harness(monkeypatch, adapter_id)
+
+    manager.set_active_adapter(req, lora_scale=0.5)
+    assert len(layer.set_calls) == 1
+    assert layer.active_slices == (True,)
+
+    manager.set_active_adapter(None)
+    assert manager._suspended_adapter_id == adapter_id
+    assert layer.suspend_calls == 1
+    assert layer.active_slices == (False,)
+    # Suspending must not tear the upload down.
+    assert layer.reset_calls == 0
+
+    manager.set_active_adapter(req, lora_scale=0.5)
+    assert layer.resume_calls == 1
+    assert len(layer.set_calls) == 1, "resume must not rebind"
+    assert layer.active_slices == (True,)
+    assert manager._active_adapter_id == adapter_id
+    assert manager._suspended_adapter_id is None
+
+
+@pytest.mark.parametrize("second_scale", [0.5, 0.25])
+def test_scale_or_adapter_change_forces_rebind(monkeypatch, second_scale):
+    """A different scale, or a different adapter, must take the full bind
+    path rather than re-arming a stale mask."""
+    adapter_id = 7
+    manager, layer = _suspend_harness(monkeypatch, adapter_id)
+
+    manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
+    manager.set_active_adapter(None)
+    assert manager._suspended_adapter_id == adapter_id
+
+    # Same id at a new scale, or a different id entirely.
+    other_id = adapter_id if second_scale != 0.5 else adapter_id + 1
+    if other_id != adapter_id:
+        manager._registered_adapters[other_id] = manager._registered_adapters[adapter_id]
+
+    manager.set_active_adapter(_dummy_lora_request(other_id), lora_scale=second_scale)
+    assert len(layer.set_calls) == 2, "must rebind"
+    assert manager._suspended_adapter_id is None
+    assert manager._active_adapter_id == other_id
+
+
+def test_partial_packed_mask_survives_suspend_resume(monkeypatch):
+    """A packed layer where only some slices carry LoRA must come back with
+    the same mask, not an all-True one."""
+    adapter_id = 7
+    manager, layer = _suspend_harness(monkeypatch, adapter_id)
+
+    # Emulate a packed bind that left slice 1 empty.
+    layer.set_lora(0, [torch.ones(2, 2), None], [torch.ones(2, 2), None])
+    manager._active_adapter_id = adapter_id
+    manager._update_adapter_scale(adapter_id, 0.5)
+    assert layer.active_slices == (True, False)
+
+    manager.set_active_adapter(None)
+    assert layer.active_slices == (False, False)
+
+    manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
+    assert layer.active_slices == (True, False), "mask must survive unchanged"
+
+
+def test_removing_suspended_adapter_drops_the_upload(monkeypatch):
+    """A removed adapter can never be resumed, so its weights must not stay
+    in the stacked buffers."""
+    adapter_id = 7
+    manager, layer = _suspend_harness(monkeypatch, adapter_id)
+
+    manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
+    manager.set_active_adapter(None)
+    assert manager._suspended_adapter_id == adapter_id
+
+    manager.remove_adapter(adapter_id)
+    assert manager._suspended_adapter_id is None
+    assert layer.reset_calls == 1, "the upload must be torn down"
+    assert layer.suspended_slices is None

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -13,7 +13,6 @@ from unittest.mock import MagicMock
 import pytest
 import torch
 import torch.nn as nn
-from huggingface_hub import snapshot_download
 from safetensors.torch import save_file
 from vllm.config.load import LoadConfig
 
@@ -32,6 +31,7 @@ from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
 from vllm_omni.diffusion.registry import initialize_model
 from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+from vllm_omni.transformers_utils.repo_utils import hf_api
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -41,7 +41,7 @@ model_path = "hf-internal-testing/tiny-helios-modular-pipe"
 @pytest.fixture(scope="module")
 def prefetch_helios_model():
     """Downloads the tiny helios model prior to running a test."""
-    snapshot_download(model_path)
+    hf_api().snapshot_download(model_path)
 
 
 @pytest.fixture(scope="function")

tests/diffusion/model_loader/test_hub_prefetch.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 contextlib
 
@@ -14,10 +14,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 def test_prefetch_subfolders_propagates_revision(monkeypatch):
     calls = []
 
-    def fake_snapshot_download(**kwargs):
+    def fake_snapshot_download(self, **kwargs):
         calls.append(kwargs)
 
-    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", fake_snapshot_download)
     monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
 
     hub_prefetch.prefetch_subfolders(

tests/diffusion/models/boogu_image/test_boogu_image_attention_mask.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.utils import make_attention_mask
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("seq_lengths", [[16], [16, 16]])
+def test_attention_mask_is_skipped_for_dense_batch(seq_lengths):
+    hidden_states = torch.empty(len(seq_lengths), 16, 64)
+
+    assert make_attention_mask(hidden_states, seq_lengths) is None
+
+
+def test_attention_mask_preserves_variable_length_padding():
+    hidden_states = torch.empty(2, 16, 64)
+
+    attention_mask = make_attention_mask(hidden_states, [16, 12])
+
+    expected = torch.tensor(
+        [
+            [True] * 16,
+            [True] * 12 + [False] * 4,
+        ]
+    )
+    torch.testing.assert_close(attention_mask, expected)

tests/diffusion/models/boogu_image/test_boogu_image_registry.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""L1 tests for checkpoint-to-pipeline routing of Boogu-Image Turbo."""
+
+import json
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_turbo_registry_entry_loads_the_turbo_class():
+    from vllm_omni.diffusion.models.boogu_image import BooguImageTurboPipeline
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    assert DiffusionModelRegistry._try_load_model_cls("BooguImageTurboPipeline") is BooguImageTurboPipeline
+
+
+def test_turbo_entry_disables_request_batching():
+    from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline, BooguImageTurboPipeline
+
+    assert BooguImagePipeline.supports_request_batch
+    assert not BooguImageTurboPipeline.supports_request_batch
+
+
+def test_turbo_process_funcs_resolve_to_the_shared_boogu_helpers():
+    from vllm_omni.diffusion.models import boogu_image
+    from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS
+
+    module = boogu_image.pipeline_boogu_image
+    for funcs in (_DIFFUSION_PRE_PROCESS_FUNCS, _DIFFUSION_POST_PROCESS_FUNCS):
+        name = funcs["BooguImageTurboPipeline"]
+        assert name == funcs["BooguImagePipeline"]
+        assert callable(getattr(module, name))
+
+
+def test_turbo_inherits_the_base_multimodal_metadata():
+    from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
+
+    metadata = get_diffusion_model_metadata("BooguImageTurboPipeline")
+
+    assert metadata == get_diffusion_model_metadata("BooguImagePipeline")
+    assert metadata.supports_multimodal_inputs
+    assert metadata.max_multimodal_image_inputs == 1
+
+
+def test_turbo_model_index_resolves_and_enriches(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"_class_name": "BooguImageTurboPipeline"}),
+        encoding="utf-8",
+    )
+
+    assert resolve_model_class_name(str(tmp_path)) == "BooguImageTurboPipeline"
+
+    config = OmniDiffusionConfig(model=str(tmp_path))
+    config.enrich_config()
+
+    assert config.model_class_name == "BooguImageTurboPipeline"
+    assert config.supports_multimodal_inputs
+    assert config.max_multimodal_image_inputs == 1
+
+
+def test_edit_turbo_override_selects_the_turbo_class(tmp_path):
+    """Edit-Turbo publishes ``BooguImagePipeline``; ``--model-class-name`` must win."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"_class_name": "BooguImagePipeline"}),
+        encoding="utf-8",
+    )
+
+    config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="BooguImageTurboPipeline")
+    config.enrich_config()
+
+    assert config.model_class_name == "BooguImageTurboPipeline"
+    assert DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo
+
+
+def test_edit_turbo_without_override_keeps_the_base_class(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"_class_name": "BooguImagePipeline"}),
+        encoding="utf-8",
+    )
+
+    config = OmniDiffusionConfig(model=str(tmp_path))
+    config.enrich_config()
+
+    assert config.model_class_name == "BooguImagePipeline"
+    assert not DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -59,7 +59,7 @@ def _force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -1180,6 +1180,21 @@ def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
     assert cfg.fps == sp.resolved_frame_rate == 12.0
 
 
+@pytest.mark.parametrize("use_path", [False, True], ids=["image", "path"])
+def test_transfer_pil_conversion_returns_writable_array(tmp_path, use_path: bool) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transfer
+
+    image = Image.new("RGB", (5, 4), "red")
+    value = image
+    if use_path:
+        value = tmp_path / "control.png"
+        image.save(value)
+
+    array = transfer._pil_to_uint8_rgb(value)
+
+    assert array.flags.writeable
+
+
 def test_transfer_vae_executor_requires_distributed_vae() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
 
@@ -2386,6 +2401,42 @@ def test_forward_transfer_runs_multichunk_overlap_path(
     torch.testing.assert_close(captured["targets"][1][:, :, 0], torch.full((1, 3, 16, 16), -0.2))
 
 
+def test_forward_transfer_non_output_rank_uses_canonical_envelope(
+    make_cosmos3_pipeline,
+    sequential_cfg_parallel,
+) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.vae.distributed_executor = SimpleNamespace(rank=1)
+    pipeline.vae.is_distributed_enabled = lambda: True
+    pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1")
+    pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
+    pipeline._set_flow_shift = lambda *_args, **_kwargs: None
+    decoded = torch.zeros(1, 3, 1, 16, 16)
+    pipeline._decode_latents = lambda latents: decoded
+
+    request = SimpleNamespace(
+        prompts=[{"prompt": "transfer", "modalities": ["video"]}],
+        sampling_params=make_sampling_params(
+            height=16,
+            width=16,
+            num_inference_steps=1,
+            guidance_scale=1.0,
+            extra_args={
+                "edge": {"control": torch.zeros(3, 1, 16, 16, dtype=torch.uint8)},
+                "max_frames": 1,
+                "num_video_frames_per_chunk": 1,
+            },
+        ),
+    )
+
+    output = pipeline.forward(request)
+
+    assert set(output.output) == {"payload", "metadata"}
+    assert set(output.output["payload"]) == {"video"}
+    torch.testing.assert_close(output.output["payload"]["video"], decoded)
+    assert output.output["metadata"] == {"video": {"fps": 24.0}}
+
+
 def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
     """With cache-dit active the uncond pass is kept even outside the guidance
     interval (so cache-dit's has_separate_cfg parity stays in phase), and the

tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.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
 
@@ -81,9 +81,8 @@ def test_from_config_loads_local_diffusers_component(tmp_path, monkeypatch: pyte
 
 
 def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
-    import huggingface_hub
-
     from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
+    from vllm_omni.transformers_utils import repo_utils
 
     cache_dir = tmp_path / "hf"
     _write_component(cache_dir, checkpoint_name=DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME)
@@ -95,7 +94,7 @@ def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pyt
         calls.append((repo_id, revision, allow_patterns))
         return str(cache_dir)
 
-    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download)
 
     sound_tokenizer.Cosmos3SoundTokenizer.from_config(
         SimpleNamespace(

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -383,6 +383,117 @@ def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
         )
 
 
+def test_gen_sp_plan_auto_pads_hidden_states_and_rope_together() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+
+    prepare_plan = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"]
+    assert set(prepare_plan) == {0, 1, 2}
+    assert all(spec.auto_pad for spec in prepare_plan.values())
+    assert [prepare_plan[index].split_dim for index in range(3)] == [1, 1, 1]
+
+
+def test_gen_sp_prepare_resets_branch_local_padding_metadata() -> None:
+    from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare
+
+    ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=12121)
+    hidden = torch.zeros(1, 4, 8)
+    rope = torch.zeros(1, 4, 1, 2)
+    with override_forward_context(ctx):
+        output = Cosmos3GenSPPrepare()(hidden, rope, rope)
+
+    assert output[0] is hidden
+    assert output[1] is rope
+    assert output[2] is rope
+    assert ctx.sp_padding_size == 0
+    assert ctx.sp_original_seq_len is None
+
+
+def test_sp_attention_masks_padded_gen_and_joint_und_keys() -> None:
+    from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention
+
+    class RecordingAttention(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.metadata: Any | None = None
+
+        def forward(self, query, key, value, metadata):
+            del key, value
+            self.metadata = metadata
+            return query
+
+    module = object.__new__(Cosmos3CrossAttention)
+    nn.Module.__init__(module)
+    module.num_heads_local = 2
+    module.head_dim = 4
+    recording_attention = RecordingAttention()
+    module.attn = recording_attention
+
+    q = torch.zeros(1, 2, 2, 4)
+    k = torch.zeros_like(q)
+    v = torch.zeros_like(q)
+    k_und = torch.zeros(1, 3, 2, 4)
+    v_und = torch.zeros_like(k_und)
+    ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=3)
+    with override_forward_context(ctx):
+        output = module._forward_sp(q, k, v, k_und, v_und)
+
+    assert output.shape == (1, 2, 8)
+    assert recording_attention.metadata is not None
+    assert recording_attention.metadata.attn_mask.tolist() == [[True, True, True, False]]
+    assert recording_attention.metadata.joint_attn_mask.tolist() == [[True, True, True]]
+
+
+def test_ulysses_2d_mask_uses_key_length_with_kv_only_joint_context(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    def fake_all_to_all(
+        _group: object,
+        tensor: torch.Tensor,
+        _scatter_idx: int,
+        _gather_idx: int,
+        _use_sync: bool,
+    ) -> torch.Tensor:
+        return torch.cat([tensor, tensor], dim=1)[:, :, :1]
+
+    monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", staticmethod(fake_all_to_all))
+    strategy = ulysses.UlyssesParallelAttention(
+        sp_group=SimpleNamespace(
+            ulysses_group=object(),
+            ulysses_world_size=2,
+            ulysses_rank=0,
+            ring_world_size=1,
+        ),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+    query = torch.zeros(1, 2, 2, 4)
+    key = torch.zeros_like(query)
+    value = torch.zeros_like(query)
+    joint_query = torch.zeros(1, 0, 2, 4)
+    joint_key = torch.zeros(1, 3, 2, 4)
+    metadata = AttentionMetadata(
+        attn_mask=torch.tensor([[True, True, True, False]]),
+        joint_attn_mask=torch.ones(1, 3, dtype=torch.bool),
+        joint_query=joint_query,
+        joint_key=joint_key,
+        joint_value=torch.zeros_like(joint_key),
+        joint_strategy="front",
+    )
+
+    query_out, key_out, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
+
+    assert query_out.shape[1] == 4
+    assert key_out.shape[1] == 7
+    assert metadata_out is not None
+    assert metadata_out.attn_mask.tolist() == [[True, True, True, True, True, True, False]]
+
+
 def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 
@@ -721,7 +832,7 @@ def test_forward_returns_video_plus_optional_modality_predictions(
     assert [tuple(tensor.shape) for tensor in output] == expected_shapes
 
 
-def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_forward_with_sound_odd_sequence_uses_auto_padding_contract(monkeypatch: pytest.MonkeyPatch) -> None:
     import vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 as cosmos3_module
 
     model = cosmos3_module.Cosmos3VFMTransformer(
@@ -732,16 +843,17 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch
     )
     monkeypatch.setattr(cosmos3_module, "_get_ulysses_state", lambda: (2, 0, None))
 
-    with pytest.raises(ValueError, match=r"GEN sequence length \(3 = video tokens 2 \+ sound tokens 1\)"):
-        model(
-            hidden_states=torch.zeros(1, 2, 1, 1, 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, 1, 2),
-            fps=24.0,
-            sound_latents=torch.zeros(1, 3, 1),
-        )
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 1, 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, 1, 2),
+        fps=24.0,
+        sound_latents=torch.zeros(1, 3, 1),
+    )
+
+    assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 1, 2), (1, 3, 1)]
 
 
 def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None:

tests/diffusion/models/dreamzero/test_qkv_fusion.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
 
 """Equivalence test for DreamZero self-attn QKV fusion.
 
@@ -52,7 +52,7 @@ def _force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/models/glm_image/test_glm_image_quantization.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 GLM-Image quantization support (W4A16/AutoRound).
 
 These tests verify that the GLM-Image DiT transformer correctly accepts and uses
@@ -69,7 +69,7 @@ def _force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/models/gr00t/test_noise_seed.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The request's seed must reach GR00T's flow-matching noise.
+
+The runner builds ``sampling_params.generator`` from ``sampling_params.seed``, the pipeline hands
+that generator to the policy, and the action head draws its initial noise from it, so the same
+seed reproduces the same action chunk and different seeds give independent ones. These tests pin
+both ends of that plumbing.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+from transformers.feature_extraction_utils import BatchFeature
+
+from vllm_omni.diffusion.models.gr00t import pipeline_gr00t
+from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
+from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7ActionHead
+from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+ROBOT_OBS = {
+    "images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)},
+    "state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)},
+    "prompt": "pick the cube",
+}
+
+
+class _OptionsRecordingPolicy:
+    """Stand-in for ``Gr00tPolicy``; the pipeline only needs ``language_key`` and ``get_action``."""
+
+    def __init__(self, *, model_path: str, embodiment_tag: str, device: str, strict: bool) -> None:
+        del model_path, embodiment_tag, device, strict  # the pipeline's constructor contract; unused here
+        self.language_key = "annotation.language.language_instruction"
+        self.seen_options: dict[str, object] | None = None
+
+    def get_action(self, obs: dict[str, object], options: dict[str, object] | None = None):
+        self.seen_options = options
+        return {"arm": np.zeros((1, 2, 2), dtype=np.float32)}, {}
+
+
+@pytest.fixture
+def pipeline(monkeypatch) -> Gr00tN1d7Pipeline:
+    monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", _OptionsRecordingPolicy)
+    od_config = SimpleNamespace(
+        model="nvidia/GR00T-N1.7-3B",
+        model_config={"embodiment_tag": "LIBERO_PANDA", "strict": False},
+        custom_pipeline_args={},
+    )
+    return Gr00tN1d7Pipeline(od_config=od_config)
+
+
+def _request(**sampling_kwargs) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompt="pick",
+        request_id="req",
+        sampling_params=OmniDiffusionSamplingParams(extra_args={"robot_obs": ROBOT_OBS}, **sampling_kwargs),
+    )
+
+
+def test_runner_generator_is_passed_through_untouched(pipeline):
+    """In serving the runner has already built ``sampling_params.generator``; the head must get that object."""
+    generator = torch.Generator().manual_seed(99)
+
+    assert pipeline.forward(_request(generator=generator)).error is None
+
+    assert pipeline.policy.seen_options["generator"] is generator
+
+
+def test_seed_without_generator_builds_one_on_the_pipeline_device(pipeline):
+    """Callers that bypass the runner (direct use, tests) still get the seed honoured."""
+    pipeline.forward(_request(seed=1234))
+
+    generator = pipeline.policy.seen_options["generator"]
+    assert isinstance(generator, torch.Generator)
+    assert generator.initial_seed() == 1234
+    assert generator.device.type == pipeline.device
+
+
+def test_auto_assigned_seed_reaches_the_head(pipeline):
+    """``OmniDiffusionRequest`` seeds every request, so the head is never left on the global RNG."""
+    req = _request()
+    assert req.sampling_params.seed is not None
+
+    pipeline.forward(req)
+
+    assert pipeline.policy.seen_options["generator"].initial_seed() == req.sampling_params.seed
+
+
+def _captured_warnings(monkeypatch) -> list[tuple]:
+    """vLLM loggers do not propagate to the root logger, so ``caplog`` cannot see them; record calls instead."""
+    captured: list[tuple] = []
+    monkeypatch.setattr(pipeline_gr00t.logger, "warning", lambda *args, **kwargs: captured.append(args))
+    return captured
+
+
+def test_unexpected_generator_form_warns_and_falls_back_to_the_seed(pipeline, monkeypatch):
+    """GR00T draws one noise tensor per request, so a list of several generators cannot be honoured."""
+    warnings = _captured_warnings(monkeypatch)
+    generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
+
+    pipeline.forward(_request(generator=generators, seed=7))
+
+    assert len(warnings) == 1
+    assert "torch.Generator" in warnings[0][0]
+    assert pipeline.policy.seen_options["generator"].initial_seed() == 7
+
+
+def test_unexpected_generator_form_without_a_seed_warns_and_leaves_the_global_rng(pipeline, monkeypatch):
+    warnings = _captured_warnings(monkeypatch)
+    generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
+
+    pipeline.forward(_request(generator=generators))
+
+    assert len(warnings) == 1
+    assert pipeline.policy.seen_options is None
+
+
+def _noise_only_head(action_horizon: int, action_dim: int) -> Gr00tN1d7ActionHead:
+    """An action head whose single denoising step is the identity, so ``action_pred`` is the initial noise."""
+    config = Gr00tN1d7Config(
+        max_action_dim=action_dim,
+        action_horizon=action_horizon,
+        num_inference_timesteps=1,
+        add_pos_embed=False,
+        use_alternate_vl_dit=False,
+    )
+    head = Gr00tN1d7ActionHead.__new__(Gr00tN1d7ActionHead)
+    nn.Module.__init__(head)
+    head.config = config
+    head.action_dim = config.max_action_dim
+    head.action_horizon = config.action_horizon
+    head.num_inference_timesteps = config.num_inference_timesteps
+    head.num_timestep_buckets = config.num_timestep_buckets
+    head.action_encoder = lambda actions, timesteps, embodiment_id: actions
+    head.model = lambda hidden_states, encoder_hidden_states, timestep: hidden_states
+    head.action_decoder = lambda hidden_states, embodiment_id: torch.zeros_like(hidden_states)
+    return head
+
+
+def _predict(head: Gr00tN1d7ActionHead, options) -> torch.Tensor:
+    batch, seq, dim = 2, 3, head.action_dim
+    return head.get_action_with_features(
+        backbone_features=torch.zeros(batch, seq, dim),
+        state_features=torch.zeros(batch, 1, dim),
+        embodiment_id=torch.zeros(batch, dtype=torch.long),
+        backbone_output=BatchFeature(data={}),
+        action_input=BatchFeature(data={}),
+        options=options,
+    )["action_pred"]
+
+
+def test_head_draws_its_noise_from_the_request_generator():
+    head = _noise_only_head(action_horizon=4, action_dim=3)
+    expected = torch.randn(2, 4, 3, generator=torch.Generator().manual_seed(5))
+
+    assert torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(5)}), expected)
+    assert not torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(6)}), expected)
+
+
+def test_head_without_generator_stays_on_the_global_rng():
+    """Upstream Isaac-GR00T parity: no generator means a bare ``torch.randn`` on the global RNG."""
+    head = _noise_only_head(action_horizon=4, action_dim=3)
+
+    torch.manual_seed(5)
+    expected = torch.randn(2, 4, 3)
+    torch.manual_seed(5)
+
+    assert torch.equal(_predict(head, None), expected)

tests/diffusion/models/gr00t/test_pipeline.py

@@ -1,4 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
+from typing import Any
 
 import numpy as np
 import pytest
@@ -12,12 +16,12 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class FakeGr00tPolicy:
-    instances = []
+    instances: list["FakeGr00tPolicy"] = []
 
     def __init__(self, **kwargs):
         self.kwargs = kwargs
         self.reset_calls = 0
-        self.seen_obs = None
+        self.seen_obs: dict[str, Any] | None = None
         self.embodiment_tag = SimpleNamespace(value="fake_embodiment")
         self.language_key = "annotation.language.language_instruction"
         self.modality_configs = {
@@ -40,7 +44,7 @@ class FakeGr00tPolicy:
         )
         FakeGr00tPolicy.instances.append(self)
 
-    def get_action(self, obs):
+    def get_action(self, obs, options=None):
         self.seen_obs = obs
         return {
             "arm": np.array([[[1.0, 2.0]]], dtype=np.float64),
@@ -107,6 +111,7 @@ def test_forward_returns_dict_actions_in_output():
     assert actions["arm"].dtype == np.float32
     np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32))
     policy = FakeGr00tPolicy.instances[0]
+    assert policy.seen_obs is not None
     assert "video" in policy.seen_obs
     assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
     assert "images" not in policy.seen_obs

tests/diffusion/models/gr00t/test_processing_eval_image_transform.py

@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""GR00T-N1.7 eval image transform: chosen from ``processor_config.json`` and bit-exact with Isaac-GR00T.
+
+``GOLDEN`` holds sha256 digests of the ``(C, H, W)`` uint8 output of Isaac-GR00T's own eval transform
+(``build_image_transformations_albumentations(...)[1]`` applied through ``apply_with_replay``) for the
+synthetic frames built below, at Isaac-GR00T 51d4c89 / albumentations 1.4.18 / opencv 5.0.0.93. The frame
+generators are integer-only so they are identical on every platform.
+
+Regenerating a digest needs Isaac-GR00T installed, whose transformers pin conflicts with this repo's, so it
+cannot run here: in an Isaac-GR00T environment push each ``synthetic_frames()`` frame through that transform
+with ``letter_box_transform`` False and True and hash the output with ``digest()``.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+
+import numpy as np
+import pytest
+import torch
+import torchvision.transforms.v2 as transforms
+
+from vllm_omni.diffusion.models.gr00t.modeling import processing_gr00t_n1d7 as processing
+from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import (
+    AlbumentationsEvalImageTransform,
+    Gr00tN1d7Processor,
+    LetterBoxTransform,
+    _build_eval_image_transform,
+    _fractional_center_crop,
+    _smallest_edge_resize,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# nvidia/GR00T-N1.7-3B and nvidia/GR00T-N1.7-DROID both ship exactly these processor_kwargs.
+TARGET, CROP, EDGE, FRACTION = [256, 256], [230, 230], 256, 0.95
+MODALITY_CONFIGS = {
+    "new_embodiment": {
+        "video": {"delta_indices": [0], "modality_keys": ["cam"]},
+        "state": {"delta_indices": [0], "modality_keys": ["joint"]},
+        "action": {"delta_indices": [0, 1], "modality_keys": ["joint"]},
+    }
+}
+
+# Isaac-GR00T eval transform output for synthetic_frames(); key = "<frame>|letterbox=<flag>".
+GOLDEN: dict[str, tuple[tuple[int, int, int], str]] = {
+    "gradient_256x256|letterbox=False": (
+        (3, 256, 256),
+        "4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715",
+    ),
+    "checker_256x256|letterbox=False": (
+        (3, 256, 256),
+        "d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c",
+    ),
+    "noise_256x256|letterbox=False": (
+        (3, 256, 256),
+        "a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6",
+    ),
+    "noise_512x768|letterbox=False": (
+        (3, 256, 383),
+        "45bc3a309f3d06c3f31051911bc8a544ff3f067b023b999dd9c2b7e6029ecc93",
+    ),
+    "gradient_256x256|letterbox=True": (
+        (3, 256, 256),
+        "4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715",
+    ),
+    "checker_256x256|letterbox=True": (
+        (3, 256, 256),
+        "d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c",
+    ),
+    "noise_256x256|letterbox=True": ((3, 256, 256), "a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6"),
+    "noise_512x768|letterbox=True": ((3, 256, 256), "52a3db185e6ee4f00617b87711dd198b8f2e2ad53865713d34e56225a4ec61a2"),
+}
+
+
+def _hash_noise(height: int, width: int, seed: int) -> np.ndarray:
+    """Integer-only pseudo-noise (murmur-style mixing on uint32); no RNG or float, so identical everywhere."""
+    ys, xs = np.mgrid[0:height, 0:width].astype(np.uint32)
+    channels = []
+    for channel in range(3):
+        v = xs * np.uint32(73856093) ^ ys * np.uint32(19349663) ^ np.uint32((seed * 3 + channel) * 83492791)
+        v ^= v >> np.uint32(13)
+        v *= np.uint32(0x5BD1E995)
+        v ^= v >> np.uint32(15)
+        channels.append((v & np.uint32(0xFF)).astype(np.uint8))
+    return np.stack(channels, axis=-1)
+
+
+def synthetic_frames() -> dict[str, np.ndarray]:
+    """Deterministic HWC uint8 RGB frames: smooth, hard-edged, noisy, and a non-square noisy one.
+
+    The non-square frame is 512x768 so both INTER_AREA downscales (2x without letterbox, 3x with it)
+    are integer factors, which OpenCV computes with exact integer sums: the goldens do not depend on
+    the CPU architecture. The 243->256 upscale is a fixed-point bilinear on uint8, also exact.
+    """
+    ys, xs = np.mgrid[0:256, 0:256].astype(np.int64)
+    frames = {
+        "gradient_256x256": np.stack([xs % 256, ys % 256, (xs + ys) // 2 % 256], axis=-1).astype(np.uint8),
+        "checker_256x256": np.stack(
+            [((ys // 16 + xs // 16) % 2) * 255, ((ys // 8 + xs // 8) % 2) * 255, ((ys // 32) % 2) * 255],
+            axis=-1,
+        ).astype(np.uint8),
+        "noise_256x256": _hash_noise(256, 256, seed=1),
+        "noise_512x768": _hash_noise(512, 768, seed=2),
+    }
+    return frames
+
+
+def digest(tensor: torch.Tensor) -> str:
+    return hashlib.sha256(np.ascontiguousarray(tensor.numpy()).tobytes()).hexdigest()
+
+
+class _FakeTokenizer:
+    padding_side = "right"
+
+
+class _FakeVlmProcessor:
+    """The two attributes ``Gr00tN1d7Processor`` and its collator touch on the Qwen3-VL processor."""
+
+    def __init__(self) -> None:
+        self.tokenizer = _FakeTokenizer()
+
+    def apply_chat_template(self, *_args: object, **_kwargs: object) -> str:
+        return "<prompt>"
+
+
+@pytest.fixture(autouse=True)
+def fake_vlm_processor(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The real Qwen3-VL processor is a gated download and irrelevant to the image pipeline."""
+    monkeypatch.setattr(processing, "build_processor", lambda _model_name, _kwargs: _FakeVlmProcessor())
+
+
+def _processor(**kwargs) -> Gr00tN1d7Processor:
+    return Gr00tN1d7Processor(
+        modality_configs=MODALITY_CONFIGS, image_crop_size=CROP, image_target_size=TARGET, **kwargs
+    )
+
+
+def _has_letterbox(compose: transforms.Compose) -> bool:
+    return any(isinstance(step, LetterBoxTransform) for step in compose.transforms)
+
+
+def test_processor_builds_the_pipeline_named_by_processor_config():
+    albumentations = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION)
+    transform = albumentations.eval_image_transform
+    assert isinstance(transform, AlbumentationsEvalImageTransform)
+    assert (transform.max_size, transform.fraction, transform.letter_box_transform) == (EDGE, FRACTION, False)
+    assert albumentations.use_albumentations is True
+
+    torchvision = _processor(use_albumentations=False)
+    assert isinstance(torchvision.eval_image_transform, transforms.Compose)
+    assert not _has_letterbox(torchvision.eval_image_transform)
+
+    letterboxed = _processor(use_albumentations=False, letter_box_transform=True)
+    assert _has_letterbox(letterboxed.eval_image_transform)
+    assert _processor(use_albumentations=True, letter_box_transform=True).eval_image_transform.letter_box_transform
+
+
+def test_legacy_config_without_the_flags_keeps_the_torchvision_recipe():
+    """Checkpoints predating the flags behave like Isaac-GR00T's defaults; training-only keys stay ignored."""
+    processor = _processor(state_dropout_prob=0.2, extra_augmentation_config=None)
+    assert processor.use_albumentations is False
+    assert processor.letter_box_transform is False
+    assert isinstance(processor.eval_image_transform, transforms.Compose)
+    assert not _has_letterbox(processor.eval_image_transform)
+
+
+def test_fraction_and_max_size_fall_back_like_isaac_gr00t():
+    transform = _build_eval_image_transform(
+        TARGET, CROP, shortest_image_edge=None, crop_fraction=None, use_albumentations=True
+    )
+    assert isinstance(transform, AlbumentationsEvalImageTransform)
+    assert transform.fraction == CROP[0] / TARGET[0]
+    assert transform.max_size == TARGET[0]
+    with pytest.raises(ValueError, match="crop_fraction is None"):
+        _build_eval_image_transform(TARGET, None, crop_fraction=None, use_albumentations=True)
+
+
+def test_geometry_follows_isaac_rounding_rules():
+    frame = _hash_noise(540, 640, seed=3)
+    resized = _smallest_edge_resize(frame, EDGE)
+    assert resized.shape == (256, 303, 3)  # round() per dimension, shortest edge -> 256
+    cropped = _fractional_center_crop(resized, FRACTION)
+    assert cropped.shape == (243, 287, 3)  # int() truncation: 243.2 -> 243, 287.85 -> 287
+    assert np.array_equal(cropped, resized[6:249, 8:295])  # (256-243)//2, (303-287)//2
+    square = frame[:256, :256]
+    assert _smallest_edge_resize(square, EDGE) is square  # scale == 1.0 is a no-op, not a resample
+
+
+@pytest.mark.parametrize("case", sorted(GOLDEN))
+def test_albumentations_recipe_is_bit_exact_with_isaac_gr00t(case: str):
+    name, flag = case.split("|letterbox=")
+    transform = _build_eval_image_transform(
+        TARGET,
+        CROP,
+        shortest_image_edge=EDGE,
+        crop_fraction=FRACTION,
+        use_albumentations=True,
+        letter_box_transform=(flag == "True"),
+    )
+    out = transform(synthetic_frames()[name])
+    expected_shape, expected_digest = GOLDEN[case]
+    assert out.dtype == torch.uint8
+    assert tuple(out.shape) == expected_shape
+    assert digest(out) == expected_digest
+
+
+def test_letterbox_is_gated_on_both_recipes():
+    frame = _hash_noise(540, 640, seed=3)
+
+    def albumentations(letterbox: bool) -> torch.Tensor:
+        return _build_eval_image_transform(
+            TARGET,
+            CROP,
+            shortest_image_edge=EDGE,
+            crop_fraction=FRACTION,
+            use_albumentations=True,
+            letter_box_transform=letterbox,
+        )(frame)
+
+    padded, unpadded = albumentations(True), albumentations(False)
+    assert tuple(padded.shape) == (3, 256, 256)  # 540x640 -> 640x640 -> 256x256 -> 243x243 -> 256x256
+    assert tuple(unpadded.shape) == (3, 256, 302)  # aspect preserved: 256x303 -> 243x287 -> 256x302
+
+    tv_padded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=True)(frame)
+    tv_unpadded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=False)(frame)
+    assert tuple(tv_padded.shape) == tuple(tv_unpadded.shape) == (3, 256, 256)
+    assert not torch.equal(tv_padded, tv_unpadded)
+    assert (tv_padded[:, :2, :] == 0).all() and not (tv_unpadded[:, :2, :] == 0).all()  # black bars only when gated
+
+
+def test_albumentations_recipe_rejects_frames_outside_the_contract():
+    transform = _build_eval_image_transform(TARGET, CROP, use_albumentations=True)
+    with pytest.raises(ValueError, match="HWC uint8"):
+        transform(np.zeros((256, 256, 3), dtype=np.float32))
+    with pytest.raises(ValueError, match="HWC uint8"):
+        transform(np.zeros((3, 256, 256), dtype=np.uint8))
+    with pytest.raises(ValueError, match="HWC uint8"):
+        transform(np.zeros((256, 256), dtype=np.uint8))
+
+
+def test_get_vlm_inputs_runs_the_selected_transform_per_frame():
+    processor = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION)
+    frames = np.stack([_hash_noise(540, 640, seed=3)] * 2)  # (T, H, W, C) like policy.py hands over
+    vlm_inputs = processor._get_vlm_inputs(
+        image_keys=["cam"],
+        images={"cam": frames},
+        masks=None,
+        image_transform=processor.eval_image_transform,
+        language="x",
+    )
+    images = vlm_inputs["vlm_content"]["images"]
+    assert len(images) == 2
+    assert all(img.size == (302, 256) for img in images)  # PIL reports (W, H); no letterbox for this checkpoint
+
+
+def test_flags_round_trip_through_save_pretrained_and_from_pretrained(tmp_path):
+    saved = _processor(
+        use_albumentations=True, letter_box_transform=True, shortest_image_edge=EDGE, crop_fraction=FRACTION
+    )
+    saved.save_pretrained(tmp_path)
+    kwargs = json.loads((tmp_path / "processor_config.json").read_text())["processor_kwargs"]
+    assert kwargs["use_albumentations"] is True
+    assert kwargs["letter_box_transform"] is True
+
+    loaded = Gr00tN1d7Processor.from_pretrained(tmp_path)
+    assert loaded.use_albumentations is True
+    assert loaded.letter_box_transform is True
+    assert isinstance(loaded.eval_image_transform, AlbumentationsEvalImageTransform)
+    assert loaded.eval_image_transform.letter_box_transform is True
+
+
+def test_from_pretrained_defaults_missing_flags_to_isaac_gr00t_defaults(tmp_path):
+    (tmp_path / "processor_config.json").write_text(
+        json.dumps(
+            {
+                "processor_class": "Gr00tN1d7Processor",
+                "processor_kwargs": {
+                    "modality_configs": MODALITY_CONFIGS,
+                    "image_crop_size": CROP,
+                    "image_target_size": TARGET,
+                },
+            }
+        )
+    )
+    (tmp_path / "statistics.json").write_text("{}")
+    loaded = Gr00tN1d7Processor.from_pretrained(tmp_path)
+    assert loaded.use_albumentations is False
+    assert loaded.letter_box_transform is False
+    assert isinstance(loaded.eval_image_transform, transforms.Compose)

tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -13,6 +14,7 @@ from vllm_omni.diffusion.models.lingbot_world.actions import (
     LingBotCameraControlReducer,
     integrate_lingbot_camera_actions,
     parse_lingbot_camera_action_frames,
+    parse_lingbot_camera_action_script,
 )
 from vllm_omni.experimental.ar_diffusion.session import (
     ARDiffusionSession,
@@ -292,3 +294,8 @@ def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
         assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
 
     asyncio.run(exercise())
+
+
+def test_parse_camera_action_script_rejects_wrong_chunk_width() -> None:
+    with pytest.raises(ValueError, match="exactly 3 per-latent-frame"):
+        parse_lingbot_camera_action_script([[["w"], ["w"]]], frames_per_chunk=3)

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.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
 
@@ -15,8 +15,10 @@ from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
 from PIL import Image
 from torch import nn
 
+import vllm_omni.diffusion.models.lingbot_world.dmd_block as lingbot_dmd_block
 import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline
 from tests.diffusion.models.wan2_2.conftest import noop_progress_bar
+from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
 from vllm_omni.diffusion.models.lingbot_world.actions import (
     integrate_lingbot_camera_actions,
 )
@@ -24,6 +26,8 @@ from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as
 from vllm_omni.diffusion.models.lingbot_world.camera import (
     build_plucker_embedding as _real_build_plucker_embedding,
 )
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.utils import StepRequestState
 from vllm_omni.experimental.ar_diffusion.tick_protocol import (
     ARDiffusionControlInput,
     ARDiffusionTickRequest,
@@ -253,6 +257,8 @@ class _SamplingParams:
         if extra_args is not None:
             self.extra_args.update(extra_args)
         self.latents = None
+        self.guidance_scale = None
+        self.guidance_scale_2 = None
 
 
 class _RequestBatch:
@@ -345,7 +351,6 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
         "FlowUniPCMultistepScheduler": _FakeScheduler,
         "CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory,
         "get_local_device": lambda: torch.device("cpu"),
-        "set_forward_context_denoise_step_idx": lambda index: None,
         "prefetch_subfolders": prefetch_subfolders,
         "from_pretrained_with_prefetch": from_pretrained_with_prefetch,
         "load_transformer_config": load_transformer_config,
@@ -358,6 +363,9 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
     }
     for name, value in replacements.items():
         monkeypatch.setattr(lingbot_pipeline, name, value)
+    # The DMD block math lives in its own module; stub the symbols it imports.
+    monkeypatch.setattr(lingbot_dmd_block, "set_forward_context_denoise_step_idx", lambda index: None)
+    monkeypatch.setattr(lingbot_dmd_block, "randn_tensor", _diffusers_randn_tensor)
     monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False)
 
 
@@ -774,6 +782,7 @@ def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() ->
         return torch.zeros(shape, device=device, dtype=dtype)
 
     module.randn_tensor = randn
+    lingbot_dmd_block.randn_tensor = randn
     result = pipeline(_request())
 
     assert requested_noise_dtypes == [torch.float32] * 4
@@ -1139,7 +1148,7 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N
     module = _load_pipeline_module()
     transformer = _RecordingTransformer()
     pipeline = _pipeline(module, transformer=transformer)
-    module.randn_tensor = lambda shape, **kwargs: torch.full(
+    module.randn_tensor = lingbot_dmd_block.randn_tensor = lambda shape, **kwargs: torch.full(
         shape,
         -99.0,
         device=kwargs["device"],
@@ -1745,6 +1754,288 @@ def test_request_cache_becomes_unreachable_after_transformer_error() -> None:
     assert not hasattr(pipeline, "transformer_cache")
 
 
+class _FakeARState:
+    def __init__(self, session_id: str = "req-1") -> None:
+        self.session_id = session_id
+        self.commits: list[str] = []
+
+    def get_kv_caches(self, branch, *, seq_len, commit_current):
+        del branch, seq_len, commit_current
+        return [SimpleNamespace()]
+
+    def commit_paged_context(self, branch):
+        self.commits.append(branch)
+
+    def clear_cross_attention(self):
+        return None
+
+    def is_cross_attention_populated(self, branch, name):
+        del branch, name
+        return True
+
+    def get_cross_attention_kv(self, branch, name):
+        del branch, name
+        zeros = torch.zeros(1, 1, 2, 4)
+        return [{"k": zeros, "v": zeros} for _ in range(2)]
+
+
+def _empty_action_script(num_chunks: int):
+    return tuple(((), (), ()) for _ in range(num_chunks))
+
+
+def _stepwise_state(
+    *,
+    request_id: str = "req-1",
+    num_frames: int = 21,
+    script=None,
+    seed: int = 17,
+):
+    num_chunks = ((num_frames - 1) // 4 + 1) // 3
+    sampling = _SamplingParams(num_frames=num_frames, seed=seed)
+    sampling.extra_args["_lingbot_camera_trajectory"] = None
+    sampling.extra_args["_lingbot_camera_actions"] = None
+    sampling.extra_args["_lingbot_camera_action_script"] = (
+        script if script is not None else _empty_action_script(num_chunks)
+    )
+    return StepRequestState(
+        request_id=request_id,
+        sampling=sampling,
+        prompt=_prompt(),
+    )
+
+
+def _run_stepwise(pipeline, state):
+    outputs = []
+    pipeline.prepare_encode(state)
+    while not state.request_denoise_completed:
+        noise = pipeline.denoise_step(None, states=[state])
+        pipeline.step_scheduler(state, noise)
+        if state.chunk_denoise_completed:
+            outputs.append(pipeline.post_decode(state))
+    return outputs
+
+
+def test_pipeline_declares_step_execution_support() -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+
+    assert pipeline.supports_step_execution is True
+    assert isinstance(pipeline, SupportsStepExecution)
+    assert supports_step_execution(pipeline) is True
+
+
+def test_preprocess_materializes_camera_action_script_without_action_path() -> None:
+    module = _load_pipeline_module()
+    sampling = _SamplingParams(include_action=False)
+    sampling.extra_args["camera_action_script"] = [[["w"], ["w"], ["w"]], [["a"], [], []]]
+    request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=sampling, request_id="req-1")
+
+    result = module.get_lingbot_world_pre_process_func(_od_config())(request)
+
+    assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None
+    assert result.sampling_params.extra_args["_lingbot_camera_actions"] is None
+    assert result.sampling_params.extra_args["_lingbot_camera_action_script"] == (
+        (("w",), ("w",), ("w",)),
+        (("a",), (), ()),
+    )
+
+
+@pytest.mark.parametrize(
+    "extra_args",
+    [
+        pytest.param({}, id="request_mode"),
+        pytest.param(_tick_extra_args(chunk_index=0), id="tick_mode"),
+    ],
+)
+def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None:
+    """A script only steers step execution, so request mode must not drop it silently."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    sampling = _SamplingParams(
+        include_action=False,
+        extra_args={
+            **extra_args,
+            "_lingbot_camera_trajectory": None,
+            "_lingbot_camera_action_script": _empty_action_script(1),
+        },
+    )
+
+    with pytest.raises(ValueError, match="camera_action_script is read only by LingBot step execution"):
+        pipeline(_request(sampling=sampling))
+
+
+def test_stepwise_progress_metadata_and_commit_trace() -> None:
+    module = _load_pipeline_module()
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    assert pipeline._ar_diffusion_kv_state is None
+    assert len(outputs) == 2
+    assert [output.chunk_index for output in outputs] == [0, 1]
+    assert all(output.total_chunks == 2 for output in outputs)
+    assert outputs[-1].finished is True
+    assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2
+    assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
+    assert fake.commits == ["main", "main"]
+    for chunk_index, output in enumerate(outputs):
+        metadata = output.output["metadata"]["ar_diffusion"]
+        assert metadata == {
+            "session_id": "req-1",
+            "request_id": "req-1",
+            "chunk_index": chunk_index,
+            "applied_event_ids": [],
+        }
+        assert output.output["payload"]["latents"].shape == (1, 16, 3, 2, 2)
+
+
+def test_stepwise_matches_tick_transformer_trace_and_latents() -> None:
+    module = _load_pipeline_module()
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",)))
+    fake = _FakeARState("world-1")
+
+    with pipeline.bind_ar_diffusion_state("world-1", fake):
+        tick_latents = []
+        for chunk_index, frames in enumerate(script):
+            tick = ARDiffusionTickRequest(
+                session_id="world-1",
+                request_id="legacy-lingbot-request",
+                chunk_index=chunk_index,
+                controls=(
+                    ARDiffusionControlInput(
+                        track="camera",
+                        schema="lingbot.camera_actions.v1",
+                        data={"mode": "frames", "frames": [list(frame) for frame in frames]},
+                    ),
+                ),
+            )
+            sampling = _SamplingParams(extra_args=tick.to_extra_args(), seed=17)
+            sampling.extra_args["_lingbot_camera_trajectory"] = None
+            sampling.extra_args["_lingbot_camera_actions"] = frames
+            result = pipeline(_request(sampling=sampling))
+            tick_latents.append(result.output["payload"]["latents"].clone())
+    tick_calls = list(transformer.calls)
+    transformer.calls.clear()
+    pipeline.close_ar_diffusion_session("world-1")
+
+    stepwise_state = _stepwise_state(request_id="req-1", num_frames=21, script=script, seed=17)
+    stepwise_fake = _FakeARState(stepwise_state.request_id)
+    with pipeline.bind_ar_diffusion_state(stepwise_state.request_id, stepwise_fake):
+        stepwise_outputs = _run_stepwise(pipeline, stepwise_state)
+
+    assert [call["update_cache"] for call in transformer.calls] == [call["update_cache"] for call in tick_calls]
+    assert [call["start_frame"] for call in transformer.calls] == [call["start_frame"] for call in tick_calls]
+    for tick_latent, output in zip(tick_latents, stepwise_outputs, strict=True):
+        torch.testing.assert_close(output.output["payload"]["latents"], tick_latent)
+
+
+def test_stepwise_trajectory_camera_matches_request_mode_under_non_uniform_speed(monkeypatch) -> None:
+    """A trajectory that slows down between blocks must condition every stepwise
+    chunk exactly as request mode does. Request mode embeds the whole trajectory
+    once; embedding per chunk would re-normalize each block's framewise
+    translations by its own largest step and present a slow block as full-speed
+    motion."""
+    from vllm_omni.diffusion.models.lingbot_world import camera as camera_module
+
+    module = _load_pipeline_module()
+    # The fixture stubs the ray embedding with a pose-blind ramp; parity needs
+    # the real geometry so a scale mismatch actually shows up.
+    monkeypatch.setattr(module, "build_plucker_embedding", camera_module.build_plucker_embedding)
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+
+    # Six latent frames -> two 3-frame blocks. The first block moves at full
+    # speed and everything after it at a tenth of that, so a per-chunk
+    # normalization would rescale the second block by 10x.
+    steps = torch.tensor([1.0] * 3 + [0.1] * 18)
+    poses = torch.eye(4).repeat(21, 1, 1)
+    poses[:, 2, 3] = torch.cumsum(steps, dim=0) - steps[0]
+    trajectory = _CameraTrajectory(
+        poses=poses,
+        intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(21, 1),
+    )
+
+    sampling = _SamplingParams(num_frames=21)
+    sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
+    pipeline(_request(sampling=sampling))
+    request_cameras = [call["camera_hidden_states"] for call in transformer.calls]
+    transformer.calls.clear()
+
+    state = _stepwise_state(num_frames=21)
+    state.sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
+    state.sampling.extra_args.pop("_lingbot_camera_action_script")
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        _run_stepwise(pipeline, state)
+    stepwise_cameras = [call["camera_hidden_states"] for call in transformer.calls]
+
+    # Two blocks x (four probes + one commit), in the same order on both paths.
+    assert len(request_cameras) == len(stepwise_cameras) == 10
+    for stepwise, request in zip(stepwise_cameras, request_cameras, strict=True):
+        torch.testing.assert_close(stepwise, request)
+
+
+def test_stepwise_emits_latents_without_touching_the_vae() -> None:
+    """Latent mode must stay latent: streaming decode is opt-in per request."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module, transformer=_RecordingTransformer())
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    assert [set(output.output["payload"]) for output in outputs] == [{"latents"}, {"latents"}]
+    assert pipeline.vae.decode_inputs == []
+
+
+def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
+    """A non-latent request must stream pixels under the primary "video" key."""
+    import diffusers.video_processor as video_processor_module
+
+    module = _load_pipeline_module()
+    processed = []
+
+    class VideoProcessor:
+        def __init__(self, *, vae_scale_factor):
+            assert vae_scale_factor == 8
+
+        def postprocess_video(self, video, *, output_type):
+            processed.append((tuple(video.shape), output_type))
+            return f"frames-{len(processed)}"
+
+    monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
+    pipeline = _pipeline(module, transformer=_RecordingTransformer())
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    # One decode per AR block, each seeing only that block's three latent frames.
+    assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
+    assert [output_type for _, output_type in processed] == ["np", "np"]
+    assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
+    # Identity metadata survives the switch to pixel output.
+    assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1]
+
+
 def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
     module = _load_pipeline_module()
     resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
@@ -1755,7 +2046,7 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No
     assert resolved is module.LingBotWorldCausalDMDPipeline
     assert cache_acceleration_disabled
     assert preprocess_name == "get_lingbot_world_pre_process_func"
-    request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams())
+    request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=_SamplingParams(), request_id="req-1")
     assert preprocess(request) is request
     assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory)
     assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline"

tests/diffusion/models/ltx2/test_ltx25_pipeline.py

@@ -5,6 +5,7 @@
 
 import json
 from dataclasses import replace
+from threading import Lock
 from types import SimpleNamespace
 from typing import Any
 
@@ -185,7 +186,7 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc
         ltx2_components.initialize_pipeline_components(pipe, od_config)
 
     assert str(exc_info.value) == (
-        "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15."
+        "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.13.0,<5.15."
     )
 
 
@@ -386,7 +387,7 @@ def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch):
         MissingNattenProcessor,
     )
 
-    with pytest.raises(RuntimeError, match=r"kernels==0\.15\.2.*supported GPU.*allow Hub access"):
+    with pytest.raises(RuntimeError, match=r"kernels==0\.16\.1.*supported GPU.*allow Hub access"):
         ltx2_components._create_ltx25_natten_processor()
 
 
@@ -500,6 +501,32 @@ def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypa
     assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
 
 
+def test_ltx_runtime_profiles_phase_methods_separately(monkeypatch):
+    pipe = object.__new__(LTX2DistilledTwoStagePipeline)
+    pipe.enable_diffusion_pipeline_profiler = True
+    pipe._profiler_lock = Lock()
+    pipe._stage_durations = {}
+    monkeypatch.setattr(ltx2_runtime.current_omni_platform, "is_available", lambda: False)
+
+    with pipe._profile_phase_method("run_phase", "generate_lowres"):
+        pass
+    with pipe._profile_phase_method("run_phase", "refine"):
+        pass
+    with pipe._profile_phase_method("decode_phase"):
+        pass
+    output = DiffusionOutput()
+    assert pipe._attach_profiled_stage_durations(output) is output
+
+    expected_keys = {
+        "LTX2DistilledTwoStagePipeline.generate_lowres.run_phase",
+        "LTX2DistilledTwoStagePipeline.refine.run_phase",
+        "LTX2DistilledTwoStagePipeline.decode_phase",
+    }
+    assert set(pipe.stage_durations) == expected_keys
+    assert set(output.stage_durations) == expected_keys
+    assert all(duration >= 0.0 for duration in pipe.stage_durations.values())
+
+
 def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 

tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Native FNA against an independent, test-only explicit neighborhood mask."""
+
+import pytest
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from vllm_omni.diffusion.models.ltx2.ops import resolve_ltx2_vae_operators
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+
+
+def _selected_fna():
+    if not current_omni_platform.is_available():
+        pytest.skip("No accelerator is available")
+    operators = resolve_ltx2_vae_operators(current_omni_platform.get_torch_device())
+    if operators is None or operators.fna is None:
+        pytest.skip("LTX DiffVAE FNA is unavailable for this device or installation")
+    return operators.fna
+
+
+def _load_fna_ops():
+    return pytest.importorskip("vllm_omni.diffusion.models.ltx2.ops.fna", exc_type=ImportError)
+
+
+def _pack(x, tile):
+    b, t, h, w, heads, dim = x.shape
+    tt, th, tw = tile
+    nt, nh, nw = (t + tt - 1) // tt, (h + th - 1) // th, (w + tw - 1) // tw
+    padded = x.new_zeros((b, nt * tt, nh * th, nw * tw, heads, dim))
+    padded[:, :t, :h, :w] = x
+    # Outer and inner spatial tiles both flatten in W,H,T order.
+    return (
+        padded.view(b, nt, tt, nh, th, nw, tw, heads, dim)
+        .permute(0, 5, 3, 1, 6, 4, 2, 7, 8)
+        .reshape(b, nt * nh * nw * tt * th * tw, heads, dim)
+    )
+
+
+def _unpack(x, shape):
+    t, h, w = shape
+    b, _, heads, dim = x.shape
+    nt, nh, nw = ((n + 3) // 4 for n in shape)
+    return (
+        x.view(b, nw, nh, nt, 4, 4, 4, heads, dim)
+        .permute(0, 3, 6, 2, 5, 1, 4, 7, 8)
+        .reshape(b, nt * 4, nh * 4, nw * 4, heads, dim)[:, :t, :h, :w]
+    )
+
+
+@pytest.mark.parametrize("batch,shape", [(1, (11, 11, 11)), (2, (13, 12, 17))])
+@torch.inference_mode()
+def test_fna_matches_explicit_inward_shift_mask(batch, shape):
+    fna = _selected_fna()
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    q, k, v = [
+        torch.randn(batch, *shape, 4, 64, device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(3)
+    ]
+    q = q * 0.125
+    coordinates = torch.stack(
+        torch.meshgrid(*(torch.arange(n, device="cuda") for n in shape), indexing="ij"), dim=-1
+    ).reshape(-1, 3)
+    start = (coordinates - 5).clamp_min(0).minimum(torch.tensor(shape, device="cuda") - 11)
+    mask = ((coordinates[None] >= start[:, None]) & (coordinates[None] < start[:, None] + 11)).all(-1)
+    with sdpa_kernel(SDPBackend.MATH):
+        expected = (
+            torch.nn.functional.scaled_dot_product_attention(
+                *(x.flatten(1, 3).transpose(1, 2).float() for x in (q, k, v)), attn_mask=mask, scale=1.0
+            )
+            .transpose(1, 2)
+            .reshape(q.shape)
+        )
+    actual = _unpack(
+        fna(_pack(q, (4, 4, 4)), _pack(k, (4, 4, 8)), _pack(v, (4, 4, 8)), shape=shape, window=(11, 11, 11)),
+        shape,
+    ).float()
+    assert actual.isfinite().all()
+    assert float((actual - expected).norm() / expected.norm()) < 0.005
+    torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.025)
+
+
+def test_fna_rejects_cpu_without_fallback():
+    fna_ops = _load_fna_ops()
+    x = torch.zeros(1, 64, 4, 64, dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="CUDA BF16"):
+        fna_ops.fna3d_tilelang(x, x, x, shape=(11, 11, 11), window=(11, 11, 11))

tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py

@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm as qk_rms_norm_ops
+from vllm_omni.diffusion.models.ltx2.ops import residual_adaln as residual_adaln_ops
+from vllm_omni.diffusion.models.ltx2.ops import resolve_ltx2_vae_operators
+from vllm_omni.diffusion.models.ltx2.ops import swiglu as swiglu_ops
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+
+_DIM_SPLIT = (16, 24, 24)
+
+
+def _selected_operators():
+    if not current_omni_platform.is_available():
+        pytest.skip("No accelerator is available")
+    operators = resolve_ltx2_vae_operators(current_omni_platform.get_torch_device())
+    if operators is None:
+        pytest.skip("No LTX DiffVAE operators are available for this device")
+    return operators
+
+
+def _selected_fna():
+    operators = _selected_operators()
+    if operators.fna is None:
+        pytest.skip("LTX DiffVAE FNA is unavailable for this device or installation")
+    return operators.fna
+
+
+def _load_fna_ops():
+    return pytest.importorskip("vllm_omni.diffusion.models.ltx2.ops.fna", exc_type=ImportError)
+
+
+def test_diffvae_model_and_tests_do_not_import_tilelang_backend() -> None:
+    # Use a fresh interpreter: reloading the model in this process would leave
+    # already-imported distributed subclasses pointing at the old base class.
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys; "
+            "sys.modules['vllm_omni.diffusion.models.ltx2.ops.fna'] = None; "
+            "sys.modules['vllm.tilelang_utils'] = None; "
+            "from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel; "
+            "import tests.diffusion.models.ltx2.test_ltx2_diffvae_attention; "
+            "import tests.diffusion.models.ltx2.test_ltx2_diffvae_ops",
+        ],
+        capture_output=True,
+        text=True,
+        timeout=60,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr
+
+
+def test_diffvae_fna_tile_map_covers_exact_neighborhood() -> None:
+    fna_ops = _load_fna_ops()
+    problem = fna_ops.Problem(13, 14, 17, (11, 11, 11))
+    indices, counts = fna_ops.build_kv_tile_map(problem)
+    q_grid_t, q_grid_h, _ = problem.q_grid
+    kv_grid_t, kv_grid_h, _ = problem.kv_grid
+
+    for q_block in range(indices.shape[0]):
+        selected = set(indices[q_block, : counts[q_block]].tolist())
+        q_outer_t = q_block % q_grid_t
+        q_outer_h = (q_block // q_grid_t) % q_grid_h
+        q_outer_w = q_block // (q_grid_t * q_grid_h)
+        for q_inner in range(fna_ops.Q_TOKENS):
+            query = (
+                q_outer_t * 4 + q_inner % 4,
+                q_outer_h * 4 + (q_inner // 4) % 4,
+                q_outer_w * 4 + q_inner // 16,
+            )
+            if any(index >= length for index, length in zip(query, problem.shape, strict=True)):
+                continue
+            starts = tuple(
+                fna_ops._window_start(index, window, length)
+                for index, window, length in zip(query, problem.window, problem.shape, strict=True)
+            )
+            for key_t in range(starts[0], starts[0] + problem.window[0]):
+                for key_h in range(starts[1], starts[1] + problem.window[1]):
+                    for key_w in range(starts[2], starts[2] + problem.window[2]):
+                        key_block = ((key_w // 8) * kv_grid_h + key_h // 4) * kv_grid_t + key_t // 4
+                        assert key_block in selected
+
+
+def test_diffvae_fna_metadata_uses_static_kernel_shapes() -> None:
+    fna_ops = _load_fna_ops()
+    for shape in ((11, 11, 11), (13, 14, 17), (66, 96, 128), (79, 192, 192)):
+        problem = fna_ops.Problem(*shape, (11, 11, 11))
+        _, t_patterns, h_patterns, w_patterns = fna_ops.build_kv_tile_metadata(problem)
+        assert {table.numel() for table in (t_patterns, h_patterns, w_patterns)} == {fna_ops.PATTERN_TABLE_SIZE}
+        assert {limit for _, _, limit in fna_ops.build_kv_tile_buckets(problem)} <= set(fna_ops.KV_BUCKET_LIMITS)
+        for ids, metadata, limit in fna_ops.build_kv_tile_buckets(problem):
+            assert metadata.shape == (ids.numel(), limit)
+
+
+def test_diffvae_fna_rejects_oversized_packed_tile_grid() -> None:
+    fna_ops = _load_fna_ops()
+    with pytest.raises(ValueError, match="15-bit metadata"):
+        fna_ops.Problem(128, 256, 256, (11, 11, 11))
+
+
+def test_diffvae_fna_propagates_launch_failures(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fna_ops = _load_fna_ops()
+    fna = _selected_fna()
+    problem = fna_ops.Problem(11, 11, 11, (11, 11, 11))
+    query = torch.randn(1, problem.q_length, fna_ops.HEADS, 64, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(1, problem.kv_length, fna_ops.HEADS, 64, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    calls = 0
+
+    def fail_launch(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(fna_ops, "_fna3d_tilelang", fail_launch)
+    with torch.inference_mode():
+        for _ in range(2):
+            with pytest.raises(RuntimeError, match="injected launch failure"):
+                fna(
+                    query,
+                    key,
+                    value,
+                    shape=problem.shape,
+                    window=problem.window,
+                )
+    assert calls == 2
+
+
+@pytest.fixture
+def isolated_runtime_state():
+    targets = (
+        (qk_rms_norm_ops, "_FAILED_ROPE_KEYS"),
+        (qk_rms_norm_ops, "_VERIFIED_ROPE_KEYS"),
+        (swiglu_ops, "_FAILED_KEYS"),
+        (swiglu_ops, "_VERIFIED_KEYS"),
+        (residual_adaln_ops, "_FAILED_ADALN_KEYS"),
+        (residual_adaln_ops, "_VERIFIED_ADALN_KEYS"),
+        (residual_adaln_ops, "_FAILED_ADD_DEVICES"),
+        (residual_adaln_ops, "_VERIFIED_ADD_DEVICES"),
+    )
+    saved = [set(getattr(module, name)) for module, name in targets]
+    for module, name in targets:
+        getattr(module, name).clear()
+    try:
+        yield
+    finally:
+        for (module, name), state in zip(targets, saved, strict=True):
+            target = getattr(module, name)
+            target.clear()
+            target.update(state)
+
+
+@pytest.mark.parametrize(
+    ("dim", "rows"),
+    [
+        (256, 257),
+        (256, 16385),
+        (256, 131075),
+        (512, 65539),
+        (1024, 32771),
+        (2048, 16387),
+    ],
+)
+def test_diffvae_swiglu_tiled_is_bit_exact(dim: int, rows: int, isolated_runtime_state) -> None:
+    from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeSwiGLU
+
+    operators = _selected_operators()
+    generator = torch.Generator(device="cuda").manual_seed(20260829)
+    hidden_states = torch.randn((1, 1, 1, rows, dim), device="cuda", dtype=torch.bfloat16, generator=generator)
+    gate_weight = torch.randn((4 * dim, dim), device="cuda", dtype=torch.bfloat16, generator=generator)
+    up_weight = torch.randn_like(gate_weight)
+    down_weight = torch.randn((dim, 4 * dim), device="cuda", dtype=torch.bfloat16, generator=generator)
+    with torch.device("meta"):
+        reference = LTX2VideoVaeSwiGLU(dim, 4 * dim)
+    reference.w_gate.weight = torch.nn.Parameter(gate_weight, requires_grad=False)
+    reference.w_up.weight = torch.nn.Parameter(up_weight, requires_grad=False)
+    reference.w_down.weight = torch.nn.Parameter(down_weight, requires_grad=False)
+
+    with torch.inference_mode():
+        expected = reference(hidden_states)
+        actual = operators.swiglu(hidden_states, gate_weight, up_weight, down_weight)
+
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize(
+    ("dim", "expected_rows"),
+    [(256, 131072), (512, 65536), (1024, 32768), (2048, 16384)],
+)
+def test_diffvae_swiglu_tile_rows_keep_a_constant_hidden_budget(dim: int, expected_rows: int) -> None:
+    assert swiglu_ops._tile_rows(1_000_000, 4 * dim) == expected_rows
+
+
+def test_diffvae_swiglu_tiled_permanently_falls_back_after_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    hidden_states = torch.randn(1, 1, 1, 17, 256, device="cuda", dtype=torch.bfloat16)
+    gate_weight = torch.randn(1024, 256, device="cuda", dtype=torch.bfloat16)
+    up_weight = torch.randn_like(gate_weight)
+    down_weight = torch.randn(256, 1024, device="cuda", dtype=torch.bfloat16)
+    calls = 0
+
+    def fail_launch(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(swiglu_ops, "_launch", fail_launch)
+    with torch.inference_mode():
+        assert operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) is None
+        assert operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) is None
+
+    assert calls == 1
+
+
+@pytest.mark.parametrize("residual_count", [1, 2])
+def test_diffvae_residual_rms_norm_modulate_is_bit_exact(
+    residual_count: int,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    generator = torch.Generator(device="cuda").manual_seed(20260829)
+    shape = (2, 3, 7, 7, 256)
+    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator)
+    residual_a = torch.randn_like(x)
+    residual_b = torch.randn_like(x) if residual_count == 2 else None
+    norm_weight = torch.randn((256,), device="cuda", dtype=torch.bfloat16, generator=generator)
+    scale = torch.randn((2, 1, 1, 1, 256), device="cuda", dtype=torch.bfloat16, generator=generator)
+    shift = torch.randn_like(scale)
+    hidden_states = x + residual_a
+    if residual_b is not None:
+        hidden_states = hidden_states + residual_b
+    expected = torch.nn.functional.rms_norm(hidden_states, (256,), norm_weight, eps=1e-6) * (1 + scale) + shift
+
+    with torch.inference_mode():
+        actual = operators.residual_norm(
+            x,
+            residual_a,
+            residual_b,
+            norm_weight,
+            scale,
+            shift,
+            1e-6,
+        )
+
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+def test_diffvae_residual_add3_is_bit_exact(isolated_runtime_state) -> None:
+    operators = _selected_operators()
+    generator = torch.Generator(device="cuda").manual_seed(20260829)
+    tensors = [
+        torch.randn((2, 3, 7, 7, 256), device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(4)
+    ]
+    expected = tensors[0] + tensors[1]
+    expected = expected + tensors[2]
+    expected = expected + tensors[3]
+
+    with torch.inference_mode():
+        actual = operators.residual_add(*tensors)
+
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+def test_diffvae_residual_adaln_permanently_falls_back_after_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    shape = (1, 3, 7, 7, 256)
+    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    weight = torch.randn(256, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(1, 1, 1, 1, 256, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn_like(scale)
+    calls = 0
+
+    def fail_launch(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(residual_adaln_ops, "_launch_adaln", fail_launch)
+    with torch.inference_mode():
+        for _ in range(2):
+            assert operators.residual_norm(x, residual, None, weight, scale, shift, 1e-6) is None
+
+    assert calls == 1
+
+
+def _tables(
+    frames: int,
+    height: int,
+    width: int,
+    dim_split: tuple[int, int, int] = _DIM_SPLIT,
+    base: float = 10000.0,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
+    tables = []
+    for length, dim in zip((frames, height, width), dim_split, strict=True):
+        exponents = torch.arange(0, dim, 2, dtype=torch.float64, device="cuda") / dim
+        inv_freqs = (1.0 / base**exponents).to(torch.float32)
+        positions = torch.arange(length, dtype=torch.float32, device="cuda")
+        angles = positions[:, None] * inv_freqs[None, :]
+        tables.append((angles.cos(), angles.sin()))
+    return tuple(tables)
+
+
+def _apply_rope_reference(
+    hidden_states: torch.Tensor,
+    tables: tuple[tuple[torch.Tensor, torch.Tensor], ...],
+    dim_split: tuple[int, int, int] = _DIM_SPLIT,
+) -> torch.Tensor:
+    outputs = []
+    offset = 0
+    for axis, (dim, (cos, sin)) in enumerate(zip(dim_split, tables, strict=True), 1):
+        chunk = hidden_states[..., offset : offset + dim]
+        pairs = chunk.reshape(*chunk.shape[:-1], dim // 2, 2)
+        even = pairs[..., 0].float()
+        odd = pairs[..., 1].float()
+        shape = [1, 1, 1, 1, 1, dim // 2]
+        shape[axis] = cos.shape[0]
+        cos = cos.reshape(shape)
+        sin = sin.reshape(shape)
+        rotated = torch.stack([even * cos - odd * sin, even * sin + odd * cos], dim=-1)
+        outputs.append(rotated.reshape(chunk.shape).to(hidden_states.dtype))
+        offset += dim
+    return torch.cat(outputs, dim=-1)
+
+
+@pytest.mark.parametrize(
+    ("frames", "height", "width", "heads"),
+    [(3, 7, 7, 1), (31, 136, 192, 4)],
+)
+def test_diffvae_qk_rms_norm_scale_rope_3d_is_bit_exact(
+    frames: int,
+    height: int,
+    width: int,
+    heads: int,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    generator = torch.Generator(device="cuda").manual_seed(20260829)
+    shape = (1, frames, height, width, heads, 64)
+    query = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator)
+    key = torch.randn_like(query)
+    query_weight = torch.randn((64,), device="cuda", dtype=torch.bfloat16, generator=generator)
+    key_weight = torch.randn((64,), device="cuda", dtype=torch.bfloat16, generator=generator)
+    tables = _tables(frames, height, width)
+    expected = (
+        _apply_rope_reference(torch.nn.functional.rms_norm(query, (64,), query_weight, 1e-6) * 0.125, tables),
+        _apply_rope_reference(torch.nn.functional.rms_norm(key, (64,), key_weight, 1e-6), tables),
+    )
+
+    with torch.inference_mode():
+        actual = operators.qk_norm_rope(
+            query,
+            key,
+            query_weight,
+            key_weight,
+            1e-6,
+            0.125,
+            _DIM_SPLIT,
+            10000.0,
+        )
+
+    assert actual is not None
+    assert torch.equal(actual[0], expected[0])
+    assert torch.equal(actual[1], expected[1])
+
+
+def test_diffvae_qk_rms_norm_scale_rope_3d_permanently_falls_back_after_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    query = torch.randn(1, 3, 7, 7, 1, 64, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn_like(query)
+    weight = torch.ones(64, device="cuda", dtype=torch.bfloat16)
+    calls = 0
+
+    def fail_launch(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(qk_rms_norm_ops, "_launch_combined", fail_launch)
+    with torch.inference_mode():
+        for _ in range(2):
+            assert operators.qk_norm_rope(query, key, weight, weight, 1e-6, 0.125, _DIM_SPLIT, 10000.0) is None
+
+    assert calls == 1

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -103,7 +103,7 @@ def test_prepare_audio_latents_samples_directly_in_packed_token_space():
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)
 
 
-def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
+def test_prepare_audio_latents_pads_short_audio_for_sp():
     pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2)
 
     latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(

tests/diffusion/models/ltx2/test_ltx2_ops_platform.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+
+import pytest
+import torch
+from vllm.platforms.interface import DeviceCapability
+
+from vllm_omni.diffusion.models.ltx2 import ops as ltx2_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass(frozen=True)
+class _TensorProbe:
+    is_cuda: bool
+    device: torch.device
+
+
+@pytest.fixture(autouse=True)
+def clear_platform_cache():
+    ltx2_platform._cuda_compute_capability.cache_clear()
+    try:
+        yield
+    finally:
+        ltx2_platform._cuda_compute_capability.cache_clear()
+
+
+@pytest.mark.parametrize(
+    ("capability", "triton", "tilelang", "fusion_expected", "fna_expected"),
+    [
+        (80, True, True, False, False),
+        (90, True, True, True, True),
+        (90, True, False, True, False),
+        (90, False, True, False, False),
+        (100, True, True, True, False),
+        (100, True, False, True, False),
+        (103, True, True, True, False),
+        (103, True, False, True, False),
+        (110, True, True, False, False),
+        (120, True, True, False, False),
+        (121, True, True, False, False),
+    ],
+)
+def test_production_resolver_and_tensor_eligibility_agree(
+    monkeypatch, capability, triton, tilelang, fusion_expected, fna_expected
+):
+    monkeypatch.setattr(ltx2_platform, "HAS_TRITON", triton)
+    monkeypatch.setattr(ltx2_platform, "has_tilelang", lambda: tilelang)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True)
+    monkeypatch.setattr(
+        ltx2_platform.current_omni_platform,
+        "get_device_capability",
+        lambda device_id=0: DeviceCapability(major=capability // 10, minor=capability % 10),
+    )
+    # Resolution must be lazy even when FNA is selected.
+    monkeypatch.setitem(sys.modules, "vllm_omni.diffusion.models.ltx2.ops.fna", None)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
+    device = torch.device("cuda:0")
+    operators = ltx2_platform.resolve_ltx2_vae_operators(device)
+    assert (operators is not None) is fusion_expected
+    assert (operators is not None and operators.fna is not None) is fna_expected
+    if operators is not None:
+        from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm, residual_adaln, swiglu
+
+        assert operators.qk_norm_rope is qk_rms_norm.try_qk_rms_norm_scale_rope_3d_exact
+        assert operators.residual_norm is residual_adaln.try_residual_rms_norm_modulate_exact
+        assert operators.residual_add is residual_adaln.try_residual_add3_exact
+        assert operators.swiglu is swiglu.try_swiglu_tiled_exact
+
+    tensor = _TensorProbe(is_cuda=True, device=device)
+    with torch.inference_mode():
+        assert ltx2_platform.is_ltx2_fusion_eligible(tensor) is fusion_expected
+        assert ltx2_platform.is_ltx2_fna_eligible(tensor) is fna_expected
+
+
+@pytest.mark.parametrize(
+    ("device", "is_cuda", "available"),
+    [("cpu", True, True), ("cuda:0", False, True), ("cuda:0", True, False)],
+)
+def test_resolver_rejects_unsupported_platform_before_capability_lookup(monkeypatch, device, is_cuda, available):
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: is_cuda)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: available)
+
+    def unexpected_lookup(_index):
+        raise AssertionError("Unsupported platforms must not query CUDA capability")
+
+    monkeypatch.setattr(ltx2_platform, "_cuda_compute_capability", unexpected_lookup)
+    assert ltx2_platform.resolve_ltx2_vae_operators(torch.device(device)) is None
+
+
+@pytest.mark.parametrize(("compiling", "grad_enabled"), [(True, False), (False, True)])
+def test_runtime_guards_bypass_operator_resolution(monkeypatch, compiling, grad_enabled):
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: compiling)
+    monkeypatch.setattr(torch, "is_grad_enabled", lambda: grad_enabled)
+
+    def unexpected_resolve(_device):
+        raise AssertionError("Ineligible runtime must not select eager operators")
+
+    monkeypatch.setattr(ltx2_platform, "resolve_ltx2_vae_operators", unexpected_resolve)
+    tensor = _TensorProbe(is_cuda=True, device=torch.device("cuda:0"))
+    assert not ltx2_platform.is_ltx2_fusion_eligible(tensor)
+    assert not ltx2_platform.is_ltx2_fna_eligible(tensor)
+
+
+def test_device_capability_is_cached(monkeypatch):
+    calls = 0
+    monkeypatch.setattr(ltx2_platform, "HAS_TRITON", True)
+
+    def get_device_capability(device_id=0):
+        nonlocal calls
+        calls += 1
+        return DeviceCapability(major=9, minor=0)
+
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "get_device_capability", get_device_capability)
+    assert ltx2_platform._cuda_compute_capability(0) == 90
+    assert ltx2_platform._cuda_compute_capability(0) == 90
+    assert calls == 1

tests/diffusion/models/ltx2/test_ltx2_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
 
 """Unit tests for LTX guidance and forward-parallel behavior."""
 
@@ -85,14 +85,14 @@ class TestCFGParallelHelpers:
 
         LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size)
 
-    def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch):
+    def test_guidance_parallel_preserves_full_guidance_plan(self, monkeypatch):
         from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
         from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
         pipe = object.__new__(LTX2Pipeline)
-        req = SimpleNamespace(is_dummy_run=lambda: True)
+        req = SimpleNamespace()
         request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
 
         guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -10,6 +10,7 @@ from dataclasses import replace
 from types import SimpleNamespace
 from typing import Any
 
+import huggingface_hub
 import numpy as np
 import pytest
 import torch
@@ -175,7 +176,19 @@ def test_ltx_device_uint8_output_rejects_invalid_contract(video, error):
         _prepare_ltx2_video_output(video, do_normalize=True)
 
 
-def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
+@pytest.mark.parametrize(
+    "parallel_config",
+    [
+        None,
+        SimpleNamespace(ulysses_degree=4),
+        SimpleNamespace(tensor_parallel_size=2),
+        SimpleNamespace(use_hsdp=True, hsdp_shard_size=2),
+        SimpleNamespace(ulysses_degree=2, tensor_parallel_size=2),
+        SimpleNamespace(ulysses_degree=2, use_hsdp=True, hsdp_shard_size=2),
+    ],
+    ids=["single", "ulysses", "tp", "hsdp", "tp-ulysses", "hsdp-ulysses"],
+)
+def test_ltx_supported_parallel_configs_keep_output_profiling(tmp_path, monkeypatch, parallel_config):
     from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
 
     (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
@@ -193,7 +206,13 @@ def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
     monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
     monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", capture_profiler_setup)
 
-    LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
+    LTX2Pipeline(
+        od_config=SimpleNamespace(
+            model=str(tmp_path),
+            parallel_config=parallel_config,
+            enable_diffusion_pipeline_profiler=False,
+        )
+    )
 
     assert "video_processor.postprocess_video" in profiler_setup["profiler_targets"]
     assert "_prepare_video_output_for_transport" in profiler_setup["profiler_targets"]
@@ -284,11 +303,11 @@ def test_ltx_artifact_uses_source_revision_and_hub_fallback(
     filename = "ltx-sidecar.safetensors"
     calls = []
 
-    def fake_download(**kwargs):
+    def fake_download(self, **kwargs):
         calls.append(kwargs)
         return "/cache/ltx-sidecar.safetensors"
 
-    monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
+    monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download)
 
     assert (
         resolve_ltx_artifact(
@@ -313,7 +332,9 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
     filename = "ltx-sidecar.safetensors"
     expected = tmp_path / filename
     expected.write_bytes(b"sidecar")
-    monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
+    monkeypatch.setattr(
+        huggingface_hub.HfApi, "hf_hub_download", lambda *_args, **_kwargs: pytest.fail("unexpected Hub lookup")
+    )
 
     assert resolve_ltx_artifact(
         str(tmp_path),
@@ -327,11 +348,11 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
 def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
     calls = []
 
-    def fake_download(**kwargs):
+    def fake_download(self, **kwargs):
         calls.append(kwargs)
         return "/cache/ltx-sidecar.safetensors"
 
-    monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
+    monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download)
 
     assert (
         resolve_ltx_artifact(
@@ -383,14 +404,42 @@ def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expecte
         )
 
 
-def test_ltx_rejects_advanced_uaa_before_component_initialization():
+@pytest.mark.parametrize(
+    "pipeline_cls",
+    [
+        LTX2Pipeline,
+        LTX2DistilledOneStagePipeline,
+        LTX2TwoStagePipeline,
+        LTX2DistilledTwoStagePipeline,
+        LTX2T2VDMD2Pipeline,
+        LTX2I2VDMD2Pipeline,
+    ],
+)
+@pytest.mark.parametrize(
+    ("parallel_config", "error"),
+    [
+        (SimpleNamespace(ulysses_mode="advanced_uaa"), "does not support ulysses_mode='advanced_uaa'"),
+        (SimpleNamespace(ring_degree=2), "pure Ulysses sequence parallelism only"),
+        (SimpleNamespace(ulysses_degree=2, ring_degree=2), "pure Ulysses sequence parallelism only"),
+        (SimpleNamespace(allgather_degree=2), "pure Ulysses sequence parallelism only"),
+    ],
+    ids=["advanced-uaa", "ring", "hybrid", "allgather-kv"],
+)
+def test_ltx_rejects_unsupported_sp_before_component_initialization(pipeline_cls, parallel_config, error, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    def unexpected_model_access(*_args, **_kwargs):
+        pytest.fail("Unsupported SP must be rejected before accessing model files or initializing components")
+
+    monkeypatch.setattr(ltx2_runtime, "detect_ltx_model_version", unexpected_model_access)
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", unexpected_model_access)
     od_config = SimpleNamespace(
         model="unused",
-        parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
+        parallel_config=parallel_config,
     )
 
-    with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
-        LTX2Pipeline(od_config=od_config)
+    with pytest.raises(ValueError, match=error):
+        pipeline_cls(od_config=od_config)
 
 
 def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
@@ -1409,39 +1458,19 @@ def test_ltx_distilled_forward_rejects_fixed_recipe_overrides(direct_kwargs, sam
         pipeline.forward(req, **direct_kwargs)
 
 
-def test_ltx_distilled_dummy_run_uses_fixed_recipe_values():
-    from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
-    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-    req = DiffusionRequestBatch(
-        [
-            OmniDiffusionRequest(
-                prompt="dummy run",
-                sampling_params=OmniDiffusionSamplingParams(
-                    height=512,
-                    width=512,
-                    num_frames=LTX2DistilledPipeline.dummy_run_num_frames,
-                    num_inference_steps=1,
-                    guidance_scale=0.0,
-                ),
-                request_id=DUMMY_DIFFUSION_REQUEST_ID,
-            )
-        ]
-    )
-    pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline)
-    seen = {}
-
-    def fake_run_recipe(_req, request_inputs, **_kwargs):
-        seen["request_inputs"] = request_inputs
-        return "dummy-output"
-
-    object.__setattr__(pipeline, "_run_recipe", fake_run_recipe)
-
-    assert pipeline.forward(req) == "dummy-output"
-    assert seen["request_inputs"].num_frames == 1
-    assert seen["request_inputs"].num_inference_steps == 8
-    assert seen["request_inputs"].guidance == LTXGuidanceSpec.positive_only()
+@pytest.mark.parametrize(
+    "pipeline_cls",
+    [
+        LTX2Pipeline,
+        LTX2DistilledOneStagePipeline,
+        LTX2TwoStagePipeline,
+        LTX2DistilledTwoStagePipeline,
+        LTX2T2VDMD2Pipeline,
+        LTX2I2VDMD2Pipeline,
+    ],
+)
+def test_ltx_pipelines_disable_startup_dummy_run(pipeline_cls):
+    assert pipeline_cls.dummy_run_num_frames == 0
 
 
 def test_ltx_custom_sigmas_bypass_scheduler_shifting():

tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.models.ltx2.ltx2_sequence_parallel import (
+    LTX2VideoToAudioParallelAttention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _run_video_to_audio_parity(rank: int, world_size: int, master_port: int) -> None:
+    dist.init_process_group(
+        "gloo",
+        init_method=f"tcp://127.0.0.1:{master_port}",
+        rank=rank,
+        world_size=world_size,
+        timeout=timedelta(seconds=60),
+    )
+    try:
+        # Non-contiguous SP subgroups match TP-fastest rank ordering. Different
+        # seeds across groups catch collectives that accidentally use WORLD.
+        group_count = world_size // 2
+        ulysses_group = dist.group.WORLD
+        if group_count > 1:
+            for offset in range(group_count):
+                ranks = list(range(offset, world_size, group_count))
+                group = dist.new_group(ranks)
+                if rank in ranks:
+                    ulysses_group = group
+        ulysses_rank = dist.get_rank(ulysses_group)
+        sp_group = SimpleNamespace(
+            ring_world_size=1,
+            ulysses_group=ulysses_group,
+            ulysses_world_size=2,
+            ulysses_rank=ulysses_rank,
+        )
+        strategy = LTX2VideoToAudioParallelAttention(sp_group)
+
+        for batch_size, masked in [(1, False), (2, True)]:
+            torch.manual_seed(17 + rank % group_count)
+            query = torch.randn(batch_size, 3, 4, 8)
+            global_key = torch.randn(batch_size, 6, 4, 8)
+            global_value = torch.randn(batch_size, 6, 4, 8)
+            key = global_key.chunk(2, dim=1)[ulysses_rank].contiguous()
+            value = global_value.chunk(2, dim=1)[ulysses_rank].contiguous()
+            mask = None
+            if masked:
+                mask = torch.ones(batch_size, 1, 1, 6, dtype=torch.bool)
+                mask[0, ..., -1] = False
+                mask[1, ..., -2:] = False
+            metadata = AttentionMetadata(attn_mask=mask) if masked else None
+
+            local_query, local_key, local_value, local_metadata, ctx = strategy.pre_attention(
+                query, key, value, metadata
+            )
+            assert local_metadata is metadata
+            local_output = F.scaled_dot_product_attention(
+                local_query.transpose(1, 2),
+                local_key.transpose(1, 2),
+                local_value.transpose(1, 2),
+                attn_mask=local_metadata.attn_mask if local_metadata is not None else None,
+            ).transpose(1, 2)
+            actual = strategy.post_attention(local_output, ctx)
+            expected = F.scaled_dot_product_attention(
+                query.transpose(1, 2),
+                global_key.transpose(1, 2),
+                global_value.transpose(1, 2),
+                attn_mask=mask,
+            ).transpose(1, 2)
+
+            torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("world_size", [2, 4], ids=["world", "subgroups"])
+def test_ltx_video_to_audio_sp_matches_replicated_attention(unused_tcp_port, world_size):
+    torch.multiprocessing.spawn(
+        _run_video_to_audio_parity,
+        args=(world_size, unused_tcp_port),
+        nprocs=world_size,
+    )
+
+
+def test_ltx_video_to_audio_sp_rejects_ring_parallelism():
+    sp_group = SimpleNamespace(ring_world_size=2)
+
+    with pytest.raises(NotImplementedError, match="ring_degree must be 1"):
+        LTX2VideoToAudioParallelAttention(sp_group)

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -141,16 +141,54 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
 
 
 @pytest.mark.parametrize("rope_type", ["interleaved", "split"])
-def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
-    root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""]
+def test_ltx_sp_plan_only_shards_video_stream(rope_type):
+    plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)
+    root_plan = plan[""]
 
-    video_timestep = root_plan["timestep"]
-    audio_timestep = root_plan["audio_timestep"]
+    assert set(root_plan) == {"hidden_states", "keyframes_mask", "timestep"}
+    assert root_plan["hidden_states"].split_dim == 1
+    assert root_plan["keyframes_mask"].split_dim == 1
+    assert root_plan["timestep"].split_dim == 1
+    assert root_plan["timestep"].expected_dims == 2
+    assert not root_plan["timestep"].split_output
+    assert set(plan) == {"", "rope", "cross_attn_rope", "proj_out"}
 
-    assert video_timestep.split_dim == audio_timestep.split_dim == 1
-    assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
-    assert not video_timestep.split_output
-    assert not audio_timestep.split_output
+
+def test_ltx_attention_assigns_video_only_sp_modes(monkeypatch):
+    modes = {}
+
+    class FakeAttention(nn.Module):
+        def __init__(self, *, prefix, sequence_parallel_mode="local", **_kwargs):
+            super().__init__()
+            modes[prefix.rsplit(".", 1)[-1]] = sequence_parallel_mode
+
+    class FakeFeedForward(nn.Module):
+        def __init__(self, *_args, **_kwargs):
+            super().__init__()
+
+    monkeypatch.setattr(ltx2_transformer, "LTX2Attention", FakeAttention)
+    monkeypatch.setattr(ltx2_transformer, "LTX2FeedForward", FakeFeedForward)
+
+    LTX2VideoTransformerBlock(
+        dim=8,
+        num_attention_heads=2,
+        attention_head_dim=4,
+        cross_attention_dim=8,
+        audio_dim=8,
+        audio_num_attention_heads=2,
+        audio_attention_head_dim=4,
+        audio_cross_attention_dim=8,
+        prefix="transformer_blocks.0",
+    )
+
+    assert modes == {
+        "attn1": "video_self",
+        "audio_attn1": "local",
+        "attn2": "local",
+        "audio_attn2": "local",
+        "audio_to_video_attn": "local",
+        "video_to_audio_attn": "video_to_audio",
+    }
 
 
 def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):

tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Full-transformer numeric parity for LTX Ulysses sequence parallelism.
+
+This FP32/SDPA check verifies that Ulysses preserves the complete transformer
+mathematical path. It does not guarantee pixel-level E2E parity: production
+BF16 and compiled attention kernels can change rounding order, and diffusion
+can amplify that drift across denoising steps.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+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
+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.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.parallel]
+
+_FP32_REL_L2 = 1e-5
+_TINY_CONFIG = {
+    "in_channels": 4,
+    "out_channels": 4,
+    "num_attention_heads": 2,
+    "attention_head_dim": 4,
+    "cross_attention_dim": 8,
+    "vae_scale_factors": (1, 1, 1),
+    "pos_embed_max_pos": 4,
+    "base_height": 2,
+    "base_width": 2,
+    "audio_in_channels": 4,
+    "audio_out_channels": 4,
+    "audio_num_attention_heads": 2,
+    "audio_attention_head_dim": 4,
+    "audio_cross_attention_dim": 8,
+    "audio_scale_factor": 1,
+    "audio_pos_embed_max_pos": 4,
+    "audio_sampling_rate": 16,
+    "audio_hop_length": 4,
+    "num_layers": 1,
+    "caption_channels": 8,
+    "rope_type": "split",
+    "use_prompt_embeddings": False,
+    "use_keyframes_abs_pos_embedding": True,
+}
+
+
+def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float:
+    return ((actual - expected).double().norm() / expected.double().norm()).item()
+
+
+def _omni_config(sp_size: int) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig.from_kwargs(
+        model="test",
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            pipeline_parallel_size=1,
+            data_parallel_size=1,
+            tensor_parallel_size=1,
+            sequence_parallel_size=sp_size,
+            ulysses_degree=sp_size,
+            ring_degree=1,
+            cfg_parallel_size=1,
+        ),
+        diffusion_attention_backend="TORCH_SDPA",
+    )
+
+
+def _inputs(device: torch.device) -> dict[str, torch.Tensor | int | bool]:
+    generator = torch.Generator().manual_seed(123)
+    return {
+        "hidden_states": torch.randn(1, 8, 4, generator=generator, device="cpu").to(device),
+        "audio_hidden_states": torch.randn(1, 4, 4, generator=generator, device="cpu").to(device),
+        "encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
+        "audio_encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device),
+        "timestep": torch.linspace(0.1, 0.8, 8, device=device).view(1, 8),
+        "audio_timestep": torch.linspace(0.2, 0.5, 4, device=device).view(1, 4),
+        "keyframes_mask": torch.tensor([[[1.0], [0.0], [1.0], [0.0], [0.0], [1.0], [0.0], [0.0]]], device=device),
+        "sigma": torch.tensor([0.1], device=device),
+        "audio_sigma": torch.tensor([0.2], device=device),
+        "num_frames": 2,
+        "height": 2,
+        "width": 2,
+        "audio_num_frames": 4,
+        "return_dict": False,
+    }
+
+
+def _run_transformer_parity(rank: int, world_size: int, master_port: int) -> None:
+    device = torch.device(f"{current_omni_platform.device_type}:{rank}")
+    current_omni_platform.set_device(device)
+    os.environ.update(
+        {
+            "RANK": str(rank),
+            "LOCAL_RANK": str(rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(master_port),
+        }
+    )
+
+    try:
+        init_distributed_environment()
+        initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size)
+
+        sp_config = _omni_config(world_size)
+        torch.manual_seed(42)
+        with set_forward_context(omni_diffusion_config=sp_config), set_current_diffusion_config(sp_config):
+            model = LTX2VideoTransformer3DModel(**_TINY_CONFIG).to(device).eval()
+        with torch.no_grad():
+            for _, parameter in sorted(model.named_parameters()):
+                torch.nn.init.normal_(parameter, mean=0.0, std=0.02)
+
+        inputs = _inputs(device)
+        reference_config = _omni_config(1)
+        with torch.no_grad(), set_forward_context(omni_diffusion_config=reference_config):
+            expected_video, expected_audio = model(**inputs)
+
+        apply_sequence_parallel(
+            model,
+            SequenceParallelConfig(ulysses_degree=world_size, ring_degree=1),
+            model._sp_plan,
+        )
+        with torch.no_grad(), set_forward_context(omni_diffusion_config=sp_config):
+            get_forward_context().sp_plan_hooks_applied = True
+            actual_video, actual_audio = model(**inputs)
+
+        v2a_attention = model.transformer_blocks[0].video_to_audio_attn.attn
+        assert v2a_attention._video_to_audio_strategy is not None
+        assert torch.isfinite(expected_video).all()
+        assert torch.isfinite(expected_audio).all()
+        for modality, actual, expected in (
+            ("video", actual_video, expected_video),
+            ("audio", actual_audio, expected_audio),
+        ):
+            drift = _relative_l2(actual, expected)
+            assert drift <= _FP32_REL_L2, (
+                f"LTX SP{world_size} {modality} output drifted rel_l2={drift:.3e} from SP1 (bound {_FP32_REL_L2:.1e})"
+            )
+    finally:
+        destroy_distributed_env()
+
+
+@pytest.mark.full_model
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
+def test_ltx_ulysses_transformer_matches_sp1(unused_tcp_port) -> None:
+    """Exercise real plan hooks, RoPE/timestep sharding, and attention routing."""
+    world_size = 2
+    device_count = current_omni_platform.device_count()
+    assert device_count >= world_size, (
+        f"LTX Ulysses parity requires {world_size} accelerator devices; only {device_count} are visible"
+    )
+    torch.multiprocessing.spawn(
+        _run_transformer_parity,
+        args=(world_size, unused_tcp_port),
+        nprocs=world_size,
+    )

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -3,6 +3,8 @@
 
 """Unit tests for LTX video VAE tiling and distributed decode behavior."""
 
+from collections.abc import Callable
+from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
 
@@ -12,6 +14,45 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.parametrize("initial_deterministic", [False, True])
+@pytest.mark.parametrize("fail", [False, True])
+def test_ltx_vocoder_deterministic_context(monkeypatch, initial_deterministic, fail):
+    from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _deterministic_ltx_vocoder
+
+    monkeypatch.setattr(torch.backends.cudnn, "deterministic", initial_deterministic)
+    if torch.backends.cudnn.is_available():
+        monkeypatch.setattr(torch.backends.cudnn, "benchmark_limit", 17)
+    settings = {
+        name: getattr(torch.backends.cudnn, name) for name in ("enabled", "benchmark", "benchmark_limit", "allow_tf32")
+    }
+    try:
+        with _deterministic_ltx_vocoder():
+            assert torch.backends.cudnn.deterministic
+            assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings
+            if fail:
+                raise RuntimeError("injected failure")
+    except RuntimeError as exc:
+        assert fail and str(exc) == "injected failure"
+    assert torch.backends.cudnn.deterministic == initial_deterministic
+    assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings
+
+
+@dataclass(frozen=True)
+class _OperatorSetProbe:
+    fna: Callable[..., torch.Tensor] | None
+
+
+@dataclass
+class _AttentionProbe:
+    kernel_size: tuple[int, int, int]
+    head_dim: int
+    to_out: tuple[torch.nn.Module, ...]
+    projected: torch.Tensor
+
+    def project_qkv(self, _hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.projected, self.projected, self.projected
+
+
 def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
     import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
 
@@ -38,9 +79,164 @@ def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
 
 
 class TestLTXDiffusionDecoder:
+    def test_diffusion_decoder_reuses_diffusers_with_scoped_overrides(self, monkeypatch):
+        from diffusers.models.autoencoders.ltx2_diffusion_decoder import (
+            LTX2VideoDiffusionDecoderModel as DiffusersModel,
+        )
+
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import (
+            LTX2VideoDiffusionDecoderModel,
+            LTX2VideoVaeDiffusionNABlock,
+            LTX2VideoVaeNeighborhoodAttention,
+            LTX2VideoVaeSwiGLU,
+        )
+
+        original_init = DiffusersModel.__init__
+        original_parameters = {}
+
+        def capture_parameters(self, *args, **kwargs):
+            original_init(self, *args, **kwargs)
+            original_parameters.update(self.named_parameters())
+
+        with torch.device("meta"):
+            reference = DiffusersModel()
+            monkeypatch.setattr(DiffusersModel, "__init__", capture_parameters)
+            model = LTX2VideoDiffusionDecoderModel()
+        assert isinstance(model, DiffusersModel)
+        assert all(parameter is original_parameters[name] for name, parameter in model.named_parameters())
+        assert {name: tensor.shape for name, tensor in model.state_dict().items()} == {
+            name: tensor.shape for name, tensor in reference.state_dict().items()
+        }
+
+        modules = tuple(model.decoder.modules())
+        expected_blocks = sum(model.config.decoder_stage_depths)
+        assert sum(isinstance(module, LTX2VideoVaeNeighborhoodAttention) for module in modules) == expected_blocks
+        assert sum(isinstance(module, LTX2VideoVaeSwiGLU) for module in modules) == expected_blocks
+        assert (
+            sum(isinstance(module, LTX2VideoVaeDiffusionNABlock) for module in modules)
+            == model.config.decoder_stage_depths[-1]
+        )
+
+        assert type(model.decoder).__module__.startswith("vllm_omni.")
+        assert model.tiled_decode.__func__ is DiffusersModel.tiled_decode
+        assert type(reference.decoder) is type(model.decoder).__bases__[0]
+
+    def test_natten_processor_routes_stage5_through_tilelang_fna(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules
+
+        processor = object.__new__(diffvae_modules.LTX2VideoVaeNeighborhoodNattenProcessor)
+        processor.backend = None
+        fallback_calls = []
+        processor._na3d = lambda *args, **kwargs: fallback_calls.append((args, kwargs))
+
+        hidden_states = torch.randn(1, 11, 11, 11, 256, dtype=torch.bfloat16)
+        projected = hidden_states.view(1, 11, 11, 11, 4, 64)
+
+        class FakeAttention:
+            kernel_size = (11, 11, 11)
+            head_dim = 64
+            to_out = (torch.nn.Identity(),)
+
+            @staticmethod
+            def project_qkv(_hidden_states):
+                return projected, projected, projected
+
+        def token_permute(value, _tile, *, flip_tiled_dims):
+            assert flip_tiled_dims
+            return value, value.shape, None
+
+        def token_unpermute(value, *, token_layout_shape, tile_shape, flip_tiled_dims):
+            assert token_layout_shape == projected.shape
+            assert tile_shape == (4, 4, 4)
+            assert flip_tiled_dims
+            return value
+
+        fast_path_calls = []
+
+        def fast_path(query, key, value, *, shape, window):
+            fast_path_calls.append((query, key, value, shape, window))
+            return query
+
+        monkeypatch.setattr(diffvae_modules, "is_ltx2_fna_eligible", lambda _tensor: True)
+        monkeypatch.setattr(
+            diffvae_modules, "resolve_ltx2_vae_operators", lambda _device: _OperatorSetProbe(fna=fast_path)
+        )
+        monkeypatch.setattr(
+            diffvae_modules,
+            "_load_natten_token_permutation",
+            lambda: (token_permute, token_unpermute),
+        )
+
+        output = processor(FakeAttention(), hidden_states)
+
+        assert torch.equal(output, hidden_states)
+        assert len(fast_path_calls) == 1
+        assert fast_path_calls[0][3:] == ((11, 11, 11), (11, 11, 11))
+        assert not fallback_calls
+
+    @pytest.mark.parametrize("reason", ["device", "head_dim", "channels", "dtype", "tilelang_missing", "kernel"])
+    def test_natten_processor_keeps_natten_when_fast_path_does_not_apply(self, monkeypatch, reason):
+        from vllm_omni.diffusion.models.ltx2.vae import decoder as ops
+
+        processor = object.__new__(ops.LTX2VideoVaeNeighborhoodNattenProcessor)
+        processor.backend = None
+        channels = 128 if reason == "channels" else 256
+        head_dim = 32 if reason == "head_dim" else 64
+        x = torch.zeros(1, 11, 11, 11, channels, dtype=torch.float32 if reason == "dtype" else torch.bfloat16)
+        projected = x.view(1, 11, 11, 11, channels // head_dim, head_dim)
+        attn = _AttentionProbe(
+            kernel_size=(3, 5, 5) if reason == "kernel" else (11, 11, 11),
+            head_dim=head_dim,
+            to_out=(torch.nn.Identity(),),
+            projected=projected,
+        )
+        calls = []
+
+        def natten(q, k, v, **kwargs):
+            calls.append(kwargs)
+            return q
+
+        def unexpected_native_load(*_args, **_kwargs):
+            raise AssertionError("Native FNA must not be loaded for an unsupported configuration")
+
+        processor._na3d = natten
+        monkeypatch.setattr(ops, "is_ltx2_fna_eligible", lambda _tensor: reason != "device")
+        monkeypatch.setattr(
+            ops,
+            "resolve_ltx2_vae_operators",
+            lambda _device: _OperatorSetProbe(fna=None if reason == "tilelang_missing" else unexpected_native_load),
+        )
+        assert torch.equal(processor(attn, x), x)
+        assert calls == [{"kernel_size": attn.kernel_size, "scale": 1.0, "backend": None}]
+
+    def test_diffusion_decoder_noneligible_block_uses_diffusers_forward(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules
+
+        block = object.__new__(diffvae_modules.LTX2VideoVaeDiffusionNABlock)
+        torch.nn.Module.__init__(block)
+        expected = torch.randn(1, 1, 1, 1, 1)
+        calls = []
+
+        def upstream_forward(_self, hidden_states, latent_context, modulation, block_mask):
+            assert block_mask is None
+            calls.append((hidden_states, latent_context, modulation))
+            return expected
+
+        monkeypatch.setattr(
+            diffvae_modules.diffusers_decoder.LTX2VideoVaeDiffusionNABlock,
+            "forward",
+            upstream_forward,
+        )
+        hidden_states = torch.randn(1, 1, 1, 1, 1)
+        latent_context = torch.randn_like(hidden_states)
+        modulation = (torch.randn_like(hidden_states),)
+
+        assert block.forward(hidden_states, latent_context, modulation) is expected
+        assert calls == [(hidden_states, latent_context, modulation)]
+
     @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"])
     def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+        from vllm_omni.diffusion.models.ltx2.vae.distributed import (
             DistributedLTX2VideoDiffusionDecoderModel,
         )
 
@@ -52,7 +248,7 @@ class TestLTXDiffusionDecoder:
             model.set_parallel_size(2, mode=mode)
 
     def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+        from vllm_omni.diffusion.models.ltx2.vae.distributed import (
             DistributedLTX2VideoDiffusionDecoderModel,
         )
 
@@ -68,7 +264,7 @@ class TestLTXDiffusionDecoder:
         assert calls == [(4, "tile")]
 
     def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import (
             LTX2VideoDiffusionDecoder3d,
             LTX2VideoDiffusionDecoderModel,
         )
@@ -143,7 +339,7 @@ class TestLTXDiffusionDecoder:
             _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5")
 
     def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import (
             convert_ltx25_native_diffusion_decoder_state_dict,
         )
 
@@ -190,7 +386,7 @@ class TestLTXDiffusionDecoder:
         assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted)
 
     def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import (
             convert_ltx25_native_diffusion_decoder_state_dict,
         )
 
@@ -388,8 +584,8 @@ class TestLTXDiffusionDecoder:
         assert output.output[1].numel() == 0
 
     def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel
+        from vllm_omni.diffusion.models.ltx2.vae.distributed import (
             DistributedLTX2VideoDiffusionDecoderModel,
         )
 
@@ -424,7 +620,7 @@ class TestLTXDiffusionDecoder:
     def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch):
         from diffusers.utils.torch_utils import randn_tensor
 
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+        from vllm_omni.diffusion.models.ltx2.vae.distributed import (
             DistributedLTX2VideoDiffusionDecoderModel,
         )
 

tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""CUDA regression tests for LTX-2 BWE vocoder precision."""
+"""GPU regression tests for LTX-2 BWE vocoder precision."""
 
 import pytest
 import torch
@@ -14,10 +14,13 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 class _BWEConvVocoder(torch.nn.Module):
     def __init__(self):
         super().__init__()
+        # Match the production LTX vocoder's conv_pre geometry. The previous
+        # single-channel 1x1 probe intermittently had no MIOpen algorithm.
         self.conv = torch.nn.Conv1d(
-            1,
-            1,
-            kernel_size=1,
+            128,
+            1024,
+            kernel_size=7,
+            padding=3,
             bias=False,
             device="cuda",
             dtype=torch.bfloat16,
@@ -33,12 +36,12 @@ class _BWEConvVocoder(torch.nn.Module):
         return output
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32():
     from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder
 
     vocoder = _BWEConvVocoder()
-    generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16)
+    generated_mel = torch.ones((1, 128, 64), device="cuda", dtype=torch.bfloat16)
 
     output = _run_ltx_vocoder(vocoder, generated_mel)
 

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -3,6 +3,7 @@
 
 import json
 import sys
+from contextlib import contextmanager
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
 from typing import Any, TypeVar
@@ -13,6 +14,8 @@ import pytest
 import torch
 import torch.nn as nn
 
+from vllm_omni.diffusion.models.minimax_h3.lora import TurboSpec, parse_turbo_filename
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 _ItemT = TypeVar("_ItemT")
@@ -24,6 +27,13 @@ def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _
     return result
 
 
+def _turbo_spec(filename: str) -> TurboSpec:
+    """Build the spec a loaded artifact of this name would carry."""
+    spec = parse_turbo_filename(filename)
+    assert spec is not None
+    return spec
+
+
 def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -167,9 +177,9 @@ def _write_partition_index(path, *, partition, tasks):
 
 
 def test_modular_diffusers_index_is_resolved_generically(tmp_path):
+    from vllm_omni.config.resolver import resolve_omni_config
     from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
     from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
-    from vllm_omni.entrypoints.utils import resolve_model_config_path
 
     (tmp_path / "modular_model_index.json").write_text(
         json.dumps(
@@ -183,7 +193,16 @@ def test_modular_diffusers_index_is_resolved_generically(tmp_path):
 
     assert is_diffusion_model(str(tmp_path))
     assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline"
-    assert resolve_model_config_path(str(tmp_path)) is None
+    resolved = resolve_omni_config(
+        str(tmp_path),
+        trust_remote_code=False,
+        deploy_config_path=None,
+        cli_overrides=None,
+        stage_overrides=None,
+        strategy_config_path=None,
+    )
+    assert resolved.config_path is None
+    assert resolved.stage_by_id(0).engine_args.model_class_name == "MiniMaxH3ModularPipeline"
 
     config = OmniDiffusionConfig(model=str(tmp_path))
     config.enrich_config()
@@ -244,9 +263,12 @@ def test_combined_task_inference_and_transformer_routing():
     assert pipeline._resolve_task(None, {"image": object()}) == "fl2va"
     assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va"
     assert pipeline._resolve_task(None, {"video": object()}) == "ref2va"
-    assert pipeline._resolve_task("fl2va", {"image": object()}, has_turbo_lora=True) == "fl2va"
-    with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
-        pipeline._resolve_task("ref2va", {}, has_turbo_lora=True)
+    fl2v_spec = _turbo_spec("minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors")
+    ref2v_spec = _turbo_spec("minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors")
+    assert pipeline._resolve_task("fl2va", {"image": object()}, turbo_spec=fl2v_spec) == "fl2va"
+    with pytest.raises(OmniClientError, match="serves"):
+        pipeline._resolve_task("ref2va", {}, turbo_spec=fl2v_spec)
+    assert pipeline._resolve_task("ref2va", {"image": object()}, turbo_spec=ref2v_spec) == "ref2va"
     pipeline.partition = "ref2va"
     pipeline.supported_tasks = frozenset({"ref2va"})
     assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
@@ -265,14 +287,16 @@ def test_combined_task_inference_and_transformer_routing():
         "component_partition",
         "source_partitions",
         "expected_tasks",
+        "load_text_encoder",
     ),
     [
-        (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}),
-        ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}),
-        ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}),
+        (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}, True),
+        ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, True),
+        ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}, True),
+        ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, False),
     ],
 )
-def test_pipeline_loads_task_selected_dits_and_shared_components_once(
+def test_pipeline_loads_task_selected_components_with_encoder_ownership(
     monkeypatch,
     tmp_path,
     task_type,
@@ -281,6 +305,7 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
     component_partition,
     source_partitions,
     expected_tasks,
+    load_text_encoder,
 ):
     from vllm_omni.diffusion.data import (
         DiffusionParallelConfig,
@@ -392,6 +417,11 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
         revision=None,
         task_type=task_type,
         quantization_config=None,
+        model_loaded={
+            "transformer": True,
+            "vae": True,
+            "text_encoder": load_text_encoder,
+        },
         parallel_config=DiffusionParallelConfig(
             cfg_parallel_size=1,
             text_encoder_tp_size=1,
@@ -403,34 +433,54 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
     assert pipeline.supported_tasks == expected_tasks
     assert len(created["dit"]) == expected_dits
     component_path = tmp_path / component_partition
-    assert created["text_encoder"] == [str(component_path / "text_encoder")]
     assert created["video_vae"] == [str(component_path / "video_vae")]
     assert created["audio_vae"] == [str(component_path / "audio_vae")]
-    assert len(tokenizer_calls) == 1
-    assert len(processor_calls) == 1
     expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"]
     assert pipeline._dit_modules == expected_dit_modules
     assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2)
+    if load_text_encoder:
+        assert created["text_encoder"] == [str(component_path / "text_encoder")]
+        assert len(tokenizer_calls) == 1
+        assert len(processor_calls) == 1
+        assert pipeline.text_encoder is not None
+        assert pipeline.tokenizer is not None
+        assert pipeline.processor is not None
+        assert pipeline._encoder_modules == ["text_encoder"]
+        assert pipeline.text_encoder_tp_size == 1
+    else:
+        assert created["text_encoder"] == []
+        assert tokenizer_calls == []
+        assert processor_calls == []
+        assert pipeline.text_encoder is None
+        assert pipeline.tokenizer is None
+        assert pipeline.processor is None
+        assert pipeline._encoder_modules == []
+        assert pipeline.text_encoder_tp_size == 0
+        assert pipeline.text_encoder_group is None
     if expected_partition == "ref2va":
         assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
     assert [source.model_or_path for source in pipeline.weights_sources] == [
         *(str(tmp_path / partition_name) for partition_name in source_partitions),
-        str(component_path),
+        *((str(component_path),) if load_text_encoder else ()),
     ]
     assert [source.subfolder for source in pipeline.weights_sources] == [
         *("transformer" for _ in source_partitions),
-        "text_encoder",
+        *(("text_encoder",) if load_text_encoder else ()),
     ]
     expected_prefixes = ["transformer."]
     if expected_dits == 2:
         expected_prefixes.append("transformers_ref.")
-    expected_prefixes.append("text_encoder.")
+    if load_text_encoder:
+        expected_prefixes.append("text_encoder.")
     assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes
-    expected_patterns = (
-        pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
-        if expected_partition == "combined"
-        else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition]
-    )
+    if load_text_encoder:
+        expected_patterns = (
+            pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
+            if expected_partition == "combined"
+            else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition]
+        )
+    else:
+        expected_patterns = pipeline_module.MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[expected_partition]
     assert download_calls == [
         {
             "model_name_or_path": "MiniMaxAI/MiniMax-H3",
@@ -1510,6 +1560,97 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
     assert next(video_vae.parameters()).dtype == torch.float32
 
 
+def test_keyframe_encode_pins_and_restores_cudnn_settings():
+    from vllm_omni.diffusion.models.minimax_h3.vae import (
+        _minimax_h3_keyframe_encode_context,
+    )
+
+    cudnn = torch.backends.cudnn
+    original = (
+        cudnn.enabled,
+        cudnn.benchmark,
+        cudnn.deterministic,
+        cudnn.allow_tf32,
+    )
+    try:
+        cudnn.enabled = False
+        cudnn.benchmark = True
+        cudnn.deterministic = False
+        cudnn.allow_tf32 = False
+
+        with _minimax_h3_keyframe_encode_context(torch.device("cuda")):
+            assert cudnn.enabled
+            assert not cudnn.benchmark
+            assert cudnn.deterministic
+            assert cudnn.allow_tf32
+
+        assert not cudnn.enabled
+        assert cudnn.benchmark
+        assert not cudnn.deterministic
+        assert not cudnn.allow_tf32
+    finally:
+        (
+            cudnn.enabled,
+            cudnn.benchmark,
+            cudnn.deterministic,
+            cudnn.allow_tf32,
+        ) = original
+
+
+@pytest.mark.parametrize("fail_encode", [False, True])
+def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+    events = []
+    input_image = Image.new("RGB", (32, 32))
+
+    @contextmanager
+    def track_context(device):
+        assert device == torch.device("cpu")
+        events.append("enter")
+        try:
+            yield
+        finally:
+            events.append("exit")
+
+    class FakeModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.ones(1))
+            self.parallel_tiling = True
+
+        def encode_images(self, image, *, use_fp16_latent):
+            assert events == ["enter"]
+            assert image is input_image
+            assert use_fp16_latent
+            events.append("encode")
+            if fail_encode:
+                raise RuntimeError("encode failed")
+            return [torch.ones(1, 1, 2, 2)]
+
+    video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE)
+    torch.nn.Module.__init__(video_vae)
+    video_vae.model = FakeModel()
+    video_vae.config_dict = {
+        "latent_channels": 1,
+        "latents_mean": [0.0],
+        "latents_std": [1.0],
+    }
+    monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context)
+
+    if fail_encode:
+        with pytest.raises(RuntimeError, match="encode failed"):
+            video_vae.encode_image(input_image)
+    else:
+        rows = video_vae.encode_image(input_image)
+        torch.testing.assert_close(rows, torch.ones(1, 4))
+
+    assert events == ["enter", "encode", "exit"]
+    assert video_vae.model.parallel_tiling
+
+
 def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
 

tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py

@@ -303,7 +303,7 @@ def _pipeline_stub(fusion=None):
     torch.nn.Module.__init__(pipeline)
     pipeline.partition = "combined"
     pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._fasth3 = fusion
     pipeline.od_config = SimpleNamespace()
     pipeline.default_video_shift, pipeline.default_audio_shift = 12.0, 3.0

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -13,7 +13,9 @@ from vllm.lora.lora_weights import PackedLoRALayerWeights
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.models.minimax_h3 import lora as lora_module
 from vllm_omni.diffusion.models.minimax_h3.lora import (
+    TurboSpec,
     load_minimax_h3_turbo_lora,
+    parse_turbo_filename,
 )
 from vllm_omni.errors import OmniClientError
 from vllm_omni.lora.request import LoRARequest
@@ -46,7 +48,7 @@ def _request(path) -> LoRARequest:
 def _write_tiny_turbo(
     path,
     *,
-    alpha: str = "128",
+    alpha: str | None = "128",
     key_format: str = "minimax-h3-diffusers",
     omit_target: str | None = None,
     shape_overrides: dict[str, tuple[int, int]] | None = None,
@@ -88,11 +90,17 @@ def _write_tiny_turbo(
                         )
                 else:
                     tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank)))
-    save_file(
-        tensors,
-        str(path),
-        metadata={"alpha": alpha, "key_format": key_format},
-    )
+    metadata = {"key_format": key_format}
+    if alpha is not None:
+        metadata["alpha"] = alpha
+    save_file(tensors, str(path), metadata=metadata)
+
+
+def _spec(filename: str = "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors") -> TurboSpec:
+    """Build the spec a loaded artifact of this name would carry."""
+    spec = parse_turbo_filename(filename)
+    assert spec is not None
+    return spec
 
 
 def test_h3_turbo_loads_through_legacy_lora_model_and_swaps_ffn(tmp_path):
@@ -107,9 +115,12 @@ def test_h3_turbo_loads_through_legacy_lora_model_and_swaps_ffn(tmp_path):
     )
 
     assert loaded is not None
-    lora_model, peft_helper = loaded
+    lora_model, peft_helper, spec = loaded
     assert peft_helper.r == 128
     assert peft_helper.lora_alpha == 128
+    assert spec.denoise_steps == 4
+    assert spec.sigma_points == 5
+    assert spec.task_family == "fl2v"
     assert len(lora_model.loras) == 312
     assert "blocks.0.attn.to_q" in lora_model.loras
     assert "blocks.0.mlp.fc1" in lora_model.loras
@@ -152,7 +163,8 @@ def test_legacy_manager_uses_the_h3_model_loader_without_changing_its_interface(
 
     class _Pipeline:
         def _load_diffusion_lora_adapter(self, **kwargs):
-            return load_minimax_h3_turbo_lora(partition="fl2va", **kwargs)
+            lora_model, peft_helper, _spec_unused = load_minimax_h3_turbo_lora(partition="fl2va", **kwargs)
+            return lora_model, peft_helper
 
     manager = object.__new__(DiffusionLoRAManager)
     manager.pipeline = _Pipeline()
@@ -168,22 +180,30 @@ def test_legacy_manager_uses_the_h3_model_loader_without_changing_its_interface(
     assert "blocks.0.mlp.fc1" in lora_model.loras
 
 
-def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
-    wrong_alpha = tmp_path / "wrong_alpha" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    wrong_alpha.parent.mkdir()
-    _write_tiny_turbo(wrong_alpha, alpha="8")
-    with pytest.raises(ValueError, match="requires alpha=128"):
-        load_minimax_h3_turbo_lora(
-            partition="fl2va",
-            lora_request=_request(wrong_alpha),
-            lora_path=wrong_alpha,
-            dtype=torch.float32,
-        )
+def test_h3_turbo_honours_the_declared_alpha_and_task_family(tmp_path):
+    """Alpha comes from the artifact; the family decides which server serves it."""
+
+    eight_step_alpha = tmp_path / "alpha8" / "minimax_h3_fl2v_turbo_8step_v1.0_768p_bf16.safetensors"
+    eight_step_alpha.parent.mkdir()
+    _write_tiny_turbo(eight_step_alpha, alpha="8")
+    loaded = load_minimax_h3_turbo_lora(
+        partition="fl2va",
+        lora_request=_request(eight_step_alpha),
+        lora_path=eight_step_alpha,
+        dtype=torch.float32,
+    )
+    assert loaded is not None
+    _, peft_helper, spec = loaded
+    # The eight-step artifact is a 1/16-strength adapter; pinning alpha to the
+    # four-step value would over-drive it by 16x.
+    assert peft_helper.lora_alpha == 8
+    assert spec.alpha == 8
+    assert spec.sigma_points == 9
 
     valid = tmp_path / "valid" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
     valid.parent.mkdir()
     _write_tiny_turbo(valid)
-    with pytest.raises(ValueError, match="supports FL2VA/T2VA only"):
+    with pytest.raises(ValueError, match="Ref2VA-only server cannot serve it"):
         load_minimax_h3_turbo_lora(
             partition="ref2va",
             lora_request=_request(valid),
@@ -192,6 +212,43 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
         )
 
 
+@pytest.mark.parametrize("partition", ["fl2va", "combined"])
+def test_h3_turbo_ref2v_artifact_needs_a_ref2va_only_server(tmp_path, partition):
+    """A combined server routes ref2va to ``transformers_ref``, which the
+    adapter cannot bind to: the LoRA target pattern only injects into
+    ``transformer``. Admitting it there would run an undistilled DiT on the
+    artifact's few-step schedule with no error anywhere."""
+
+    path = tmp_path / "minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    with pytest.raises(ValueError, match="start the server with --task-type ref2va"):
+        load_minimax_h3_turbo_lora(
+            partition=partition,
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+        )
+
+
+def test_h3_turbo_ref2v_artifact_loads_on_a_ref2va_server(tmp_path):
+    path = tmp_path / "minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    loaded = load_minimax_h3_turbo_lora(
+        partition="ref2va",
+        lora_request=_request(path),
+        lora_path=path,
+        dtype=torch.float32,
+    )
+
+    assert loaded is not None
+    spec = loaded[2]
+    assert spec.task_family == "ref2v"
+    assert spec.supported_tasks == frozenset({"ref2va"})
+    assert spec.sigma_points == 9
+
+
 @pytest.mark.parametrize(
     "offload_mode",
     [
@@ -214,6 +271,8 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
 
 
 def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
+    """DLO keeps the LoRA buffers resident, so it is not an unsupported mode."""
+
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.models.minimax_h3 import (
         pipeline_minimax_h3 as pipeline_module,
@@ -222,7 +281,7 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.partition = "fl2va"
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._native_lora_adapter_ids = set()
     pipeline._lora_sigma_schedules = {}
     pipeline.od_config = SimpleNamespace(
@@ -231,10 +290,11 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
         enable_distributed_layerwise_offload=True,
     )
     captured = {}
+    spec = _spec()
 
     def load_turbo(**kwargs):
         captured.update(kwargs)
-        return object(), object()
+        return object(), object(), spec
 
     monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
     loaded = pipeline._load_diffusion_lora_adapter(
@@ -245,6 +305,10 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
 
     assert loaded is not None
     assert captured["unsupported_offload_mode"] is None
+    # The pipeline passes its own partition through and keeps the spec that
+    # decides the request contract.
+    assert captured["partition"] == "fl2va"
+    assert pipeline._turbo_lora_specs == {1: spec}
 
 
 @pytest.mark.parametrize(
@@ -268,7 +332,7 @@ def test_h3_turbo_offload_compatibility_follows_dit_selection(
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.partition = "fl2va"
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._native_lora_adapter_ids = set()
     pipeline._lora_sigma_schedules = {}
     pipeline.od_config = SimpleNamespace(
@@ -281,10 +345,11 @@ def test_h3_turbo_offload_compatibility_follows_dit_selection(
         pin_cpu_memory=True,
     )
     captured = {}
+    spec = _spec()
 
     def load_turbo(**kwargs):
         captured.update(kwargs)
-        return object(), object()
+        return object(), object(), spec
 
     monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
 
@@ -296,31 +361,84 @@ def test_h3_turbo_offload_compatibility_follows_dit_selection(
     assert captured["unsupported_offload_mode"] == expected_mode
 
 
-def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
-    unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
-    _write_tiny_turbo(unsupported)
+def test_h3_turbo_uses_the_reference_alpha_when_the_artifact_declares_none(tmp_path):
+    """``4step_v0.1`` ships no alpha. LightX2V's reference script applies
+    ``scale * alpha / rank`` with alpha defaulting to 8, so the fallback must be
+    8; alpha == rank would drive the adapter 16x too strongly."""
 
-    with pytest.raises(ValueError, match="supports only"):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v0.1.safetensors"
+    _write_tiny_turbo(path, alpha=None)
+
+    loaded = load_minimax_h3_turbo_lora(
+        partition="fl2va",
+        lora_request=_request(path),
+        lora_path=path,
+        dtype=torch.float32,
+    )
+
+    assert loaded is not None
+    _, peft_helper, spec = loaded
+    assert spec.alpha == 8
+    assert peft_helper.lora_alpha == 8
+    assert spec.alpha / spec.rank == 0.0625
+
+
+def test_h3_turbo_declines_the_comfyui_export(tmp_path):
+    """The ComfyUI exports fuse Q/K/V, so they are not this loader's format."""
+
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_comfyui_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    assert (
         load_minimax_h3_turbo_lora(
             partition="fl2va",
-            lora_request=_request(unsupported),
-            lora_path=unsupported,
+            lora_request=_request(path),
+            lora_path=path,
             dtype=torch.float32,
         )
+        is None
+    )
 
-    supported = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
-    _write_tiny_turbo(supported)
-    # Directory resolution selects the declared artifact even when another
-    # v1.0 checkpoint is present beside it.
-    assert (
+
+def test_h3_turbo_accepts_every_published_artifact_name(tmp_path):
+    """Each published filename loads and carries its own sampler contract."""
+
+    for index, (name, steps, shift) in enumerate(
+        (
+            ("minimax_h3_fl2v_turbo_4step_v0.1.safetensors", 4, 12.0),
+            ("minimax_h3_fl2v_turbo_4step_v1.1_768p_bf16.safetensors", 4, 6.0),
+            ("minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors", 8, 12.0),
+            ("minimax_h3_fl2v_turbo_8step_v1.0_768p_bf16.safetensors", 8, 6.0),
+        )
+    ):
+        directory = tmp_path / f"artifact{index}"
+        directory.mkdir()
+        path = directory / name
+        _write_tiny_turbo(path)
+        loaded = load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+        )
+        assert loaded is not None, name
+        spec = loaded[2]
+        assert spec.denoise_steps == steps
+        assert spec.sigma_points == steps + 1
+        assert spec.video_shift == shift
+
+
+def test_h3_turbo_directory_with_several_artifacts_is_ambiguous(tmp_path):
+    _write_tiny_turbo(tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors")
+    _write_tiny_turbo(tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_768p_bf16.safetensors")
+
+    with pytest.raises(ValueError, match="point --lora-path at one file"):
         load_minimax_h3_turbo_lora(
             partition="fl2va",
             lora_request=_request(tmp_path),
             lora_path=tmp_path,
             dtype=torch.float32,
         )
-        is not None
-    )
 
 
 def test_h3_turbo_rejects_a_truncated_declared_artifact(tmp_path):
@@ -352,7 +470,7 @@ def test_non_h3_checkpoint_falls_back_to_the_generic_peft_loader(tmp_path):
 
     declared = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
     _write_tiny_turbo(declared, key_format="other")
-    with pytest.raises(ValueError, match="requires safetensors metadata"):
+    with pytest.raises(ValueError, match="expected 'minimax-h3-diffusers'"):
         load_minimax_h3_turbo_lora(
             partition="fl2va",
             lora_request=_request(declared),
@@ -371,7 +489,7 @@ def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch):
     torch.nn.Module.__init__(pipeline)
     pipeline.partition = "combined"
     pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._native_lora_adapter_ids = set()
     pipeline._lora_sigma_schedules = {}
     request = _request("generic-peft")
@@ -388,14 +506,14 @@ def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch):
         return pipeline._resolve_task(
             "ref2va",
             {},
-            has_turbo_lora=pipeline._has_active_turbo_lora(sampling),
+            turbo_spec=pipeline._active_turbo_spec(sampling),
         )
 
-    recognized = (object(), object())
+    recognized = (object(), object(), _spec())
     monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: recognized)
-    assert load() is recognized
+    assert load() == recognized[:2]
     assert resolve(0.0) == "ref2va"
-    with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
+    with pytest.raises(OmniClientError, match="serves"):
         resolve(1.0)
 
     # Simulate manager eviction followed by a generic PEFT adapter reusing the
@@ -410,7 +528,7 @@ def test_h3_turbo_requires_all_loaded_targets_to_bind():
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
-    pipeline._turbo_lora_adapter_ids = {1}
+    pipeline._turbo_lora_specs = {1: _spec()}
     lora_model = SimpleNamespace(id=1, loras={"to_q": object(), "to_k": object()})
 
     pipeline._validate_diffusion_lora_binding(
@@ -448,7 +566,7 @@ def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args,
     )
 
     with pytest.raises(OmniClientError, match=error):
-        pipeline._validate_turbo_sampling(sampling)
+        pipeline._validate_turbo_sampling(sampling, _spec())
 
 
 def test_h3_turbo_accepts_five_sigma_points_for_four_nfe():
@@ -461,5 +579,6 @@ def test_h3_turbo_accepts_five_sigma_points_for_four_nfe():
         SimpleNamespace(
             num_inference_steps=5,
             extra_args={"flow_shift": 6.0, "audio_flow_shift": 3.0},
-        )
+        ),
+        _spec(),
     )

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -5,14 +5,154 @@ import pytest
 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 (
+    _MAX_1D_GRID_SIZE,
+    _iter_row_chunks,
+    _launch_row_chunks,
+    indexed_gate,
     indexed_gate_rms_norm_scale_shift,
+    indexed_scale_shift_,
     rms_norm_indexed_scale_shift,
 )
+from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("rows", "expected"),
+    [
+        (0, []),
+        (_MAX_1D_GRID_SIZE, [(0, _MAX_1D_GRID_SIZE)]),
+        (_MAX_1D_GRID_SIZE + 1, [(0, _MAX_1D_GRID_SIZE), (_MAX_1D_GRID_SIZE, 1)]),
+        (69888, [(0, _MAX_1D_GRID_SIZE), (_MAX_1D_GRID_SIZE, 4353)]),
+    ],
+)
+def test_row_chunks_respect_ascend_grid_limit(rows, expected):
+    assert list(_iter_row_chunks(rows)) == expected
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("device_type", "expected_grids_and_offsets"),
+    [
+        ("npu", [(_MAX_1D_GRID_SIZE, 0), (4353, _MAX_1D_GRID_SIZE)]),
+        ("cuda", [(69888, 0)]),
+    ],
+)
+def test_launch_row_chunks_is_npu_only(device_type, expected_grids_and_offsets):
+    class RecordingKernel:
+        def __init__(self):
+            self.calls = []
+
+        def __getitem__(self, grid):
+            def launch(*args, **kwargs):
+                self.calls.append((grid, args, kwargs))
+
+            return launch
+
+    kernel = RecordingKernel()
+    first_arg = object()
+    second_arg = object()
+
+    _launch_row_chunks(kernel, 69888, device_type, first_arg, second_arg, block_n=16)
+
+    assert kernel.calls == [
+        ((grid,), (first_arg, second_arg, row_offset), {"block_n": 16})
+        for grid, row_offset in expected_grids_and_offsets
+    ]
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+@pytest.mark.parametrize(
+    "device",
+    [pytest.param("npu", marks=hardware_marks(res={"npu": "A3"}, num_cards=1))],
+)
+def test_modulation_wrappers_match_reference_above_grid_limit(device):
+    rows = 69888
+    hidden_size = 16
+    num_conditions = 7
+    eps = 1e-6
+    dtype = torch.bfloat16
+
+    torch.manual_seed(0)
+    x = torch.randn(rows, hidden_size, dtype=dtype)
+    residual = torch.randn_like(x)
+    branch = torch.randn_like(x)
+    weight = torch.randn(hidden_size, dtype=dtype)
+    shift = torch.randn(num_conditions, hidden_size, dtype=dtype)
+    scale = torch.randn_like(shift)
+    gate = torch.randn_like(shift)
+    indices = torch.arange(rows) % num_conditions
+
+    indexed_shift = shift.index_select(0, indices)
+    indexed_scale = scale.index_select(0, indices)
+    indexed_gate_values = gate.index_select(0, indices)
+
+    x_float = x.float()
+    branch_float = branch.float()
+    indexed_shift_float = indexed_shift.float()
+    indexed_scale_float = indexed_scale.float()
+    indexed_gate_float = indexed_gate_values.float()
+
+    expected_scale_shift = (x_float * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype)
+    expected_gate = (x_float + indexed_gate_float * branch_float).to(dtype)
+
+    normalized = x_float
+    normalized = normalized * torch.rsqrt(normalized.pow(2).mean(-1, keepdim=True) + eps)
+    normalized = normalized * weight.float()
+    expected_rms = (normalized * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype)
+
+    updated_residual = residual.float() + indexed_gate_float * branch_float
+    expected_residual = updated_residual.to(dtype)
+    normalized_residual = updated_residual
+    normalized_residual = normalized_residual * torch.rsqrt(normalized_residual.pow(2).mean(-1, keepdim=True) + eps)
+    normalized_residual = normalized_residual * weight.float()
+    expected_modulated = (normalized_residual * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype)
+
+    npu_x = x.to(device)
+    npu_residual = residual.to(device)
+    npu_branch = branch.to(device)
+    npu_weight = weight.to(device)
+    npu_shift = shift.to(device)
+    npu_scale = scale.to(device)
+    npu_gate = gate.to(device)
+    npu_indices = indices.to(device)
+
+    actual_scale_shift = indexed_scale_shift_(npu_x.clone(), npu_shift, npu_scale, npu_indices)
+    actual_gate = indexed_gate(npu_x, npu_gate, npu_branch, npu_indices)
+    actual_rms = rms_norm_indexed_scale_shift(
+        npu_x,
+        npu_weight,
+        npu_shift,
+        npu_scale,
+        npu_indices,
+        eps,
+    )
+    actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift(
+        npu_residual,
+        npu_gate,
+        npu_branch,
+        npu_weight,
+        npu_shift,
+        npu_scale,
+        npu_indices,
+        eps,
+    )
+
+    for actual, expected in (
+        (actual_scale_shift, expected_scale_shift),
+        (actual_gate, expected_gate),
+        (actual_rms, expected_rms),
+        (actual_residual, expected_residual),
+        (actual_modulated, expected_modulated),
+    ):
+        torch.testing.assert_close(actual.cpu(), expected, atol=2e-2, rtol=2e-2)
+
+
+@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_modulation_preserves_bf16_residual_boundary() -> None:

tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py

@@ -309,7 +309,7 @@ def test_pipeline_native_schedule_and_task_validation(monkeypatch):
     torch.nn.Module.__init__(pipeline)
     pipeline.partition = "fl2va"
     pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._native_lora_adapter_ids = {7}
     pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
     pipeline._base_schedule_by_partition = {"fl2va": None}
@@ -330,7 +330,7 @@ def test_pipeline_native_schedule_and_task_validation(monkeypatch):
         pipeline._resolve_task(
             "fl2va",
             {},
-            has_turbo_lora=False,
+            turbo_spec=None,
             has_native_lora=True,
         )
 
@@ -418,7 +418,7 @@ def test_pipeline_replaces_native_classification_after_reload(monkeypatch, tmp_p
         enable_layerwise_offload=False,
         enable_distributed_layerwise_offload=False,
     )
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._native_lora_adapter_ids = set()
     pipeline._lora_sigma_schedules = {}
 
@@ -472,7 +472,7 @@ def test_h3_native_allows_distributed_layerwise_offload(monkeypatch):
         enable_layerwise_offload=False,
         enable_distributed_layerwise_offload=True,
     )
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._native_lora_adapter_ids = set()
     pipeline._lora_sigma_schedules = {}
     captured: dict[str, object] = {}
@@ -612,7 +612,7 @@ def test_legacy_manager_uses_native_loader(tmp_path):
                 enable_layerwise_offload=False,
                 enable_distributed_layerwise_offload=False,
             )
-            pipeline._turbo_lora_adapter_ids = set()
+            pipeline._turbo_lora_specs = {}
             pipeline._native_lora_adapter_ids = set()
             pipeline._lora_sigma_schedules = {}
             return pipeline._load_diffusion_lora_adapter(**kwargs)
@@ -634,7 +634,7 @@ def test_pipeline_schedule_inactive_when_scale_zero():
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
     pipeline.partition = "fl2va"
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._native_lora_adapter_ids = {7}
     pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
     pipeline._base_schedule_by_partition = {"fl2va": None}
@@ -664,7 +664,7 @@ def test_pipeline_schedule_falls_back_after_eviction(monkeypatch, tmp_path):
         enable_layerwise_offload=False,
         enable_distributed_layerwise_offload=False,
     )
-    pipeline._turbo_lora_adapter_ids = set()
+    pipeline._turbo_lora_specs = {}
     pipeline._native_lora_adapter_ids = set()
     pipeline._lora_sigma_schedules = {}
     pipeline._base_schedule_by_partition = {"fl2va": None}

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -13,12 +13,14 @@ from PIL import Image
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-def test_h3_profiler_targets_cover_aggregate_reference_encoders():
+def test_h3_profiler_targets_cover_reference_encoders_and_vae_decoders():
     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

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py

@@ -42,10 +42,10 @@ _QUALITY_CONFIG = QualityTestConfig(
 
 
 def _resolve_fl2va_model_ref() -> str:
-    from huggingface_hub import snapshot_download
+    from vllm_omni.transformers_utils.repo_utils import hf_api
 
     repo_root = Path(
-        snapshot_download(
+        hf_api().snapshot_download(
             repo_id=_MINIMAX_H3_REPO,
             revision=_MINIMAX_H3_REVISION,
             allow_patterns=["FL2VA/**"],
@@ -253,13 +253,13 @@ def test_resolve_fl2va_model_ref(tmp_path, monkeypatch):
     fl2va_root.mkdir()
     (fl2va_root / "model_index.json").write_text("{}", encoding="utf-8")
 
-    def fake_snapshot_download(*, repo_id, revision, allow_patterns):
+    def fake_snapshot_download(self, *, repo_id, revision, allow_patterns):
         assert repo_id == _MINIMAX_H3_REPO
         assert revision == _MINIMAX_H3_REVISION
         assert allow_patterns == ["FL2VA/**"]
         return str(tmp_path)
 
-    monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", fake_snapshot_download)
     assert _resolve_fl2va_model_ref() == str(fl2va_root)
 
 

tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py

@@ -49,3 +49,29 @@ def test_qwen3_vl_npu_swiglu_matches_two_projection_reference() -> None:
         atol=2e-2,
         rtol=2e-2,
     )
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_qwen3_vl_npu_swiglu_forward_uses_projection_module(mocker) -> None:
+    from vllm_omni.platforms.npu.models import minimax_h3
+    from vllm_omni.platforms.npu.models.minimax_h3 import (
+        _forward_minimax_h3_qwen3vl_text_mlp_npu,
+    )
+
+    x = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16)
+    packed = torch.randn(1, 16, 256, device="npu", dtype=torch.bfloat16)
+    activated = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16)
+    gate_up_proj = mocker.Mock(return_value=packed)
+    down_proj = mocker.Mock(side_effect=lambda value: value)
+    mlp = mocker.Mock(gate_up_proj=gate_up_proj, down_proj=down_proj)
+    swiglu = mocker.patch.object(minimax_h3, "npu_swiglu_from_packed", return_value=activated)
+
+    output = _forward_minimax_h3_qwen3vl_text_mlp_npu(mlp, x)
+
+    gate_up_proj.assert_called_once()
+    assert gate_up_proj.call_args.args[0] is x
+    swiglu.assert_called_once()
+    assert swiglu.call_args.args[0] is packed
+    down_proj.assert_called_once()
+    assert down_proj.call_args.args[0] is activated
+    assert output is activated

tests/diffusion/models/minimax_h3/test_minimax_h3_turbo_matrix.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Contract coverage for the LightX2V MiniMax-H3 Turbo artifact matrix."""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.diffusion.models.minimax_h3.lora import (
+    _TURBO_AUDIO_SHIFT,
+    _TURBO_VIDEO_SHIFT_544P,
+    _TURBO_VIDEO_SHIFT_768P,
+    TurboSpec,
+    parse_turbo_filename,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+# The eight Diffusers-layout artifacts published at lightx2v/Minimax-h3-Turbo,
+# with the contract each one implies. The ComfyUI export of an artifact is a
+# fused-QKV re-packing of the same weights and is not served.
+PUBLISHED = [
+    ("minimax_h3_fl2v_turbo_4step_v0.1.safetensors", "fl2v", 4, _TURBO_VIDEO_SHIFT_544P),
+    ("minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", "fl2v", 4, _TURBO_VIDEO_SHIFT_768P),
+    ("minimax_h3_fl2v_turbo_4step_v1.1_768p_bf16.safetensors", "fl2v", 4, _TURBO_VIDEO_SHIFT_768P),
+    ("minimax_h3_fl2v_turbo_4step_v1.2_768p_bf16.safetensors", "fl2v", 4, _TURBO_VIDEO_SHIFT_768P),
+    ("minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors", "fl2v", 8, _TURBO_VIDEO_SHIFT_544P),
+    ("minimax_h3_fl2v_turbo_8step_v1.0_768p_bf16.safetensors", "fl2v", 8, _TURBO_VIDEO_SHIFT_768P),
+    ("minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors", "ref2v", 4, _TURBO_VIDEO_SHIFT_544P),
+    ("minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors", "ref2v", 8, _TURBO_VIDEO_SHIFT_768P),
+]
+
+
+@pytest.mark.parametrize(("filename", "task", "steps", "video_shift"), PUBLISHED)
+def test_every_published_artifact_is_recognised(filename: str, task: str, steps: int, video_shift: float) -> None:
+    spec = parse_turbo_filename(filename)
+    assert spec is not None, filename
+    assert spec.filename == filename
+    assert spec.task_family == task
+    assert spec.denoise_steps == steps
+    assert spec.video_shift == video_shift
+    assert spec.audio_shift == _TURBO_AUDIO_SHIFT
+    # Until the loader reads metadata, alpha is LightX2V's reference default.
+    assert spec.rank == 128
+    assert spec.alpha == 8.0
+
+
+def test_all_eight_artifacts_are_covered() -> None:
+    assert len(PUBLISHED) == 8
+    assert len({name for name, *_ in PUBLISHED}) == 8
+
+
+@pytest.mark.parametrize(
+    "filename",
+    [
+        "adapter_model.safetensors",
+        "minimax_h3_fl2v_turbo_v1.0_768p_bf16.safetensors",
+        "minimax_h3_t2v_turbo_4step_v1.0.safetensors",
+        "minimax_h3_fl2v_turbo_0step_v1.0.safetensors",
+        # The ComfyUI exports are a different tensor layout, not a variant of
+        # the supported one; the loader refuses them by name.
+        "minimax_h3_fl2v_turbo_4step_v1.0_768p_comfyui_bf16.safetensors",
+        "minimax_h3_ref2v_turbo_8step_v1.0_768p_comfyui_bf16.safetensors",
+    ],
+)
+def test_unrelated_names_are_rejected(filename: str) -> None:
+    assert parse_turbo_filename(filename) is None
+
+
+def _spec(filename: str) -> TurboSpec:
+    spec = parse_turbo_filename(filename)
+    assert spec is not None
+    return spec
+
+
+def test_sigma_points_follow_the_step_count() -> None:
+    assert _spec("minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors").sigma_points == 5
+    assert _spec("minimax_h3_fl2v_turbo_8step_v1.0_768p_bf16.safetensors").sigma_points == 9
+
+
+def test_task_family_decides_supported_tasks() -> None:
+    fl2v = _spec("minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors")
+    ref2v = _spec("minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors")
+    assert fl2v.supported_tasks == frozenset({"t2va", "fl2va"})
+    assert ref2v.supported_tasks == frozenset({"ref2va"})

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py

@@ -0,0 +1,357 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import gc
+import tempfile
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.models.minimax_h3.chunked_decode import decode_h3_chunks
+from vllm_omni.diffusion.models.minimax_h3.temporal_chunks import decode_temporal_chunks
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(scope="module")
+def cpu_process_group():
+    if dist.is_initialized():
+        yield dist.group.WORLD
+        return
+
+    with tempfile.NamedTemporaryFile(prefix="h3_chunk_dist_") as rendezvous:
+        init_method = f"file://{rendezvous.name}"
+    dist.init_process_group("gloo", rank=0, world_size=1, init_method=init_method)
+    try:
+        yield dist.group.WORLD
+    finally:
+        dist.destroy_process_group()
+        gc.collect()
+
+
+class _FakeTemporalModel:
+    use_3d_conv = True
+    token_drop = 3
+    tokens_chunk_size = 5
+    token_overlap = 2
+    vae_ratio_t = 4
+    frame_pre_padding = 3
+    frame_overlap = 5
+    isolated_first_frame = False
+    isolated_last_frame = False
+
+    def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens):
+        del z, z_head, z_tail, num_chunks, pad_tokens
+        return 35, 0, 35
+
+    def _adaptive_decode(self, clip):
+        value = float(clip[:, :, 0].mean())
+        return torch.full((1, 1, 24, 2, 2), value)
+
+    @staticmethod
+    def blend(overlap, part, frame_overlap, dim):
+        del overlap, frame_overlap, dim
+        return part
+
+
+class _FakeHost:
+    def __init__(self):
+        self.model = _FakeTemporalModel()
+
+    @staticmethod
+    def _denormalize_latent(latent):
+        return latent
+
+    @staticmethod
+    def _normalize_decoded_frames(frames):
+        return frames.float()
+
+
+def test_temporal_chunks_emit_ordered_frames_and_collect_when_unconsumed():
+    model = _FakeTemporalModel()
+    latent = torch.arange(8, dtype=torch.float32).view(1, 1, 8, 1, 1)
+    chunks = []
+    marker = decode_temporal_chunks(model, latent, chunks.append)
+
+    assert marker.shape == (0,)
+    assert [chunk.shape[2] for chunk in chunks] == [17, 17, 1]
+    assert torch.equal(torch.cat(chunks, dim=2), decode_temporal_chunks(model, latent, None))
+
+
+def test_h3_callback_failure_is_deferred_until_temporal_decode_finishes():
+    host = _FakeHost()
+    latent = torch.zeros(1, 1, 8, 1, 1)
+    seen = []
+
+    def fail_once(frames):
+        seen.append(frames.shape[2])
+        raise RuntimeError("sink failed")
+
+    with pytest.raises(RuntimeError, match="sink failed"):
+        decode_h3_chunks(host, latent, fail_once, group=None)
+    assert seen == [17]
+
+
+def test_h3_without_callback_is_a_plain_full_decode_on_a_vae_group(cpu_process_group):
+    host = _FakeHost()
+    latent = torch.arange(8, dtype=torch.float32).view(1, 1, 8, 1, 1)
+
+    full = decode_h3_chunks(host, latent, None, group=cpu_process_group)
+
+    chunks = []
+    decode_h3_chunks(host, latent, chunks.append, group=cpu_process_group)
+    assert torch.equal(torch.cat(chunks, dim=2), full)
+
+
+class _SimulatedGroup:
+    """Stand in for a VAE group so one process can check every rank's verdict.
+
+    ``all_reduce`` is replaced with the sum the real collective would produce
+    for a group of ``world_size`` ranks, ``supplied`` of which handed in a
+    callback. That sum is what makes a rank-local verdict observable from a
+    single process.
+    """
+
+    def __init__(self, world_size: int, supplied: int):
+        self.world_size = world_size
+        self.supplied = supplied
+
+
+def _patch_simulated_ranks(monkeypatch, simulated: _SimulatedGroup, rank: int) -> None:
+    from vllm_omni.diffusion.models.minimax_h3 import chunked_decode as mod
+
+    def fake_all_reduce(tensor, group=None):
+        del group  # the simulated census below stands in for the real reduction
+        tensor.copy_(torch.tensor([simulated.supplied], dtype=tensor.dtype))
+
+    monkeypatch.setattr(mod.dist, "get_rank", lambda _group: rank)
+    monkeypatch.setattr(mod.dist, "get_world_size", lambda _group: simulated.world_size)
+    monkeypatch.setattr(mod.dist, "all_reduce", fake_all_reduce)
+    monkeypatch.setattr(mod.dist, "broadcast", lambda tensor, src=0, group=None: None)
+    monkeypatch.setattr(mod.dist, "get_global_rank", lambda _group, _rank: 0)
+
+
+@pytest.mark.parametrize("rank", [0, 1, 2])
+def test_h3_every_rank_rejects_a_partially_supplied_callback(monkeypatch, rank):
+    """Ranks that disagree must all refuse, not let one decode alone.
+
+    A rank-local verdict let a peer fall through into the decoder collectives
+    while its peers raised, hanging the stage instead of failing it.
+    """
+    host = _FakeHost()
+    latent = torch.zeros(1, 1, 8, 1, 1)
+    # Two ranks of three handed in a callback.
+    group = _SimulatedGroup(world_size=3, supplied=2)
+    _patch_simulated_ranks(monkeypatch, group, rank)
+
+    with pytest.raises(ValueError, match="every rank"):
+        decode_h3_chunks(host, latent, lambda _frames: None, group=object())
+
+
+@pytest.mark.parametrize("rank", [0, 1, 2])
+def test_h3_only_rank_zero_publishes_when_every_rank_supplies(monkeypatch, rank):
+    """Ownership is positional: all ranks stream, rank 0 publishes."""
+    host = _FakeHost()
+    latent = torch.zeros(1, 1, 8, 1, 1)
+    group = _SimulatedGroup(world_size=3, supplied=3)
+    _patch_simulated_ranks(monkeypatch, group, rank)
+
+    seen: list[int] = []
+    decode_h3_chunks(host, latent, lambda f: seen.append(int(f.shape[2])), group=object())
+
+    # Peers run the loop to stay in the collectives but publish nothing.
+    assert (seen != []) is (rank == 0)
+
+
+def _stub_video_vae(*, parallel_size: int, tile_count: int):
+    """A MiniMaxH3VideoVAE with only the pieces the decode entry points read."""
+    from contextlib import contextmanager
+
+    from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.parallel_size = parallel_size
+    vae.model = SimpleNamespace(_adaptive_decode=lambda clip: clip)
+    vae.entered_rank_local = False
+
+    @contextmanager
+    def rank_local():
+        vae.entered_rank_local = True
+        yield
+
+    vae._decoder_tile_count = lambda latent: tile_count
+    vae._rank_local_tiling = rank_local
+    vae.is_distributed_enabled = lambda: False
+    return vae
+
+
+def test_chunked_decode_falls_back_when_tiles_are_fewer_than_ranks(monkeypatch):
+    """Without the fallback, tileless ranks raise while their peers block."""
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_mod
+
+    vae = _stub_video_vae(parallel_size=4, tile_count=1)
+    monkeypatch.setattr(vae_mod, "decode_h3_chunks", lambda *args, **kwargs: torch.zeros(1))
+
+    vae.decode_with_chunks(torch.zeros(1, 1, 8, 1, 1), on_chunk=lambda _frames: None)
+
+    assert vae.entered_rank_local, "chunked decode must share decode_latent's too-few-tiles fallback"
+
+
+def test_chunked_decode_keeps_shared_tiling_when_every_rank_has_a_tile(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_mod
+
+    vae = _stub_video_vae(parallel_size=4, tile_count=8)
+    monkeypatch.setattr(vae_mod, "decode_h3_chunks", lambda *args, **kwargs: torch.zeros(1))
+
+    vae.decode_with_chunks(torch.zeros(1, 1, 8, 1, 1), on_chunk=lambda _frames: None)
+
+    assert not vae.entered_rank_local
+
+
+def test_h3_vae_declares_the_chunked_decode_capability():
+    """A consumer must be able to detect the capability and the pixel range."""
+    from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode
+    from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+    vae = object.__new__(MiniMaxH3VideoVAE)
+
+    assert supports_chunked_vae_decode(vae)
+    # H3 reverts through the checkpoint's processor, which lands in [0, 1];
+    # a Wan VAE publishes [-1, 1], so the range cannot be assumed.
+    assert vae.chunk_value_range == (0.0, 1.0)
+
+
+def _native_decode_temporal_streaming(model, z, z_head, z_tail, num_chunks, pad_tokens):
+    """The released ``_decode_temporal_streaming`` loop, as the parity oracle.
+
+    Transcribed from the checkpoint's ``AutoencoderKLLegacy`` so the fork can be
+    compared against the control flow it mirrors, including the isolated
+    head/tail extraction.
+    """
+    total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan(
+        z, z_head, z_tail, num_chunks, pad_tokens
+    )
+    chunk_dec = model.tokens_chunk_size * model.vae_ratio_t
+    split_count = int(model.token_drop > 0) + 1
+    dec = None
+    dec_overlap = None
+    write_pos = 0
+
+    def write_part(part):
+        nonlocal dec, write_pos
+        part_frames = int(part.shape[2])
+        if part_frames <= 0:
+            return
+        if dec is None:
+            out_shape = list(part.shape)
+            out_shape[2] = output_frames
+            dec = torch.empty(out_shape, dtype=part.dtype, device=part.device)
+        remaining = int(dec.shape[2]) - write_pos
+        copy_frames = min(part_frames, max(0, remaining))
+        if copy_frames > 0:
+            dec[:, :, write_pos : write_pos + copy_frames].copy_(part[:, :, :copy_frames])
+            write_pos += copy_frames
+
+    for i in range(num_chunks):
+        t_start = i * model.tokens_chunk_size
+        clip_z = z[:, :, t_start : t_start + model.tokens_chunk_size + model.token_overlap]
+        if i == 0 and z_head is not None:
+            clip_z = torch.cat([z_head, clip_z], dim=2)
+        if i == num_chunks - 1 and z_tail is not None:
+            clip_z = torch.cat([clip_z, z_tail], dim=2)
+
+        clip_dec = model._adaptive_decode(clip_z)
+
+        dec_tail = None
+        if i == 0 and z_head is not None:
+            write_part(clip_dec[:, :, model.vae_ratio_t - 1 : model.vae_ratio_t])
+            clip_dec = clip_dec[:, :, model.vae_ratio_t :]
+        if i == num_chunks - 1 and z_tail is not None:
+            dec_tail = clip_dec[:, :, -1:]
+            clip_dec = clip_dec[:, :, : -model.vae_ratio_t]
+
+        for j in range(split_count):
+            f_start = j * chunk_dec
+            f_end = min(f_start + chunk_dec, clip_dec.shape[2])
+            chunk = clip_dec[:, :, f_start:f_end][:, :, model.frame_pre_padding :]
+            if j == 0:
+                if dec_overlap is not None:
+                    chunk = model.blend(dec_overlap, chunk, model.frame_overlap, dim=-3)
+                    dec_overlap = None
+                write_part(chunk)
+            else:
+                dec_overlap = chunk.contiguous()
+
+        if i == num_chunks - 1:
+            if dec_overlap is not None:
+                write_part(dec_overlap)
+                dec_overlap = None
+            if dec_tail is not None:
+                write_part(dec_tail)
+
+    return dec
+
+
+class _IsolatedFrameModel(_FakeTemporalModel):
+    """Fake decoder with the isolated boundaries the released config can set."""
+
+    frame_pre_padding = 0
+
+    def __init__(self, *, first: bool, last: bool):
+        self.isolated_first_frame = first
+        self.isolated_last_frame = last
+
+    def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens):
+        del pad_tokens
+        frames = num_chunks * self.tokens_chunk_size * self.vae_ratio_t
+        frames += int(z_head is not None) + int(z_tail is not None)
+        del z
+        return frames, 0, frames
+
+    def _adaptive_decode(self, clip):
+        # Every latent token decodes to a distinct, ordered block, so a shifted
+        # or reordered emission is visible in the values themselves.
+        base = float(clip[:, :, 0].mean())
+        frames = int(clip.shape[2]) * self.vae_ratio_t
+        return torch.arange(frames, dtype=torch.float32).view(1, 1, frames, 1, 1) + base
+
+
+@pytest.mark.parametrize(
+    ("first", "last"),
+    [(True, False), (False, True), (True, True)],
+)
+def test_temporal_chunks_match_the_native_loop_at_isolated_boundaries(first, last):
+    """The fork must extract the isolated head/tail the way the native loop does."""
+    model = _IsolatedFrameModel(first=first, last=last)
+    latent = torch.arange(12, dtype=torch.float32).view(1, 1, 12, 1, 1)
+
+    chunks: list[torch.Tensor] = []
+    decode_temporal_chunks(model, latent, chunks.append)
+    streamed = torch.cat(chunks, dim=2)
+
+    # Feed the oracle exactly what the released ``decode_temporal`` would:
+    # split the isolated tokens off, then pad the remainder to a whole chunk.
+    isolated = int(first) + int(last)
+    pseudo_tokens = int(latent.shape[2]) - isolated + model.token_drop
+    remainder = pseudo_tokens % model.tokens_chunk_size
+    pad_tokens = (model.tokens_chunk_size - remainder) if remainder else 0
+    num_chunks = (pseudo_tokens + pad_tokens) // model.tokens_chunk_size - int(model.token_drop > 0)
+
+    body = latent
+    z_head = None
+    if first:
+        z_head = body[:, :, :1]
+        body = body[:, :, 1:]
+    z_tail = None
+    if last:
+        z_tail = body[:, :, -1:]
+        body = body[:, :, :-1]
+    if pad_tokens:
+        body = torch.cat([body, body[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)
+
+    expected = _native_decode_temporal_streaming(model, body, z_head, z_tail, num_chunks, pad_tokens)
+
+    assert streamed.shape == expected.shape
+    assert torch.equal(streamed, expected)

tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.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
 
 import pytest
@@ -64,6 +64,13 @@ def setup_tp_group(monkeypatch, mocker):
         lambda: 0,
     )
 
+    # vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever
+    # tp_size > 1; these tests drive CPU tensors, so keep the portable masked path.
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda",
+        lambda: False,
+    )
+
     mock_tp_group = mocker.MagicMock()
     mock_tp_group.world_size = 2
     mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group)
@@ -94,7 +101,7 @@ def setup_tp_group(monkeypatch, mocker):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
     with set_current_vllm_config(VllmConfig(device_config=device_config)):

tests/diffusion/models/qwen_image/test_qwen_image_pipelines_txt_seq_lens.py

@@ -0,0 +1,150 @@
+"""Behavioral regression tests for Qwen-Image Edit ``txt_seq_lens`` wiring.
+
+RoPE text length must follow padded embed width, not valid-token count from
+``mask.sum()``. This exercises the Edit pipeline call site: when masks are
+shorter than padded embeds, ``forward()`` must pass padded width into
+``diffuse()``.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import (
+    QwenImageEditPipeline,
+)
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+PADDED_WIDTH = 64
+VALID_TOKENS = 10
+EMBED_DIM = 8
+
+
+def _make_padded_prompt_pair():
+    prompt_embeds = torch.zeros(1, PADDED_WIDTH, EMBED_DIM)
+    prompt_embeds_mask = torch.zeros(1, PADDED_WIDTH, dtype=torch.bool)
+    prompt_embeds_mask[:, :VALID_TOKENS] = True
+    return prompt_embeds, prompt_embeds_mask
+
+
+def _make_edit_forward_batch(*, negative_prompt: str | None) -> DiffusionRequestBatch:
+    preprocessed = torch.zeros(3, 32, 32)
+    return DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="edit-txt-seq-lens",
+                prompt={
+                    "prompt": "edit prompt",
+                    "negative_prompt": negative_prompt,
+                    "additional_information": {
+                        "prompt_image": preprocessed,
+                        "preprocessed_image": preprocessed,
+                        "calculated_height": 32,
+                        "calculated_width": 32,
+                    },
+                },
+                sampling_params=SimpleNamespace(
+                    height=32,
+                    width=32,
+                    num_inference_steps=2,
+                    sigmas=None,
+                    max_sequence_length=1024,
+                    generator=None,
+                    true_cfg_scale=4.0,
+                    guidance_scale_provided=False,
+                    guidance_scale=1.0,
+                    num_outputs_per_prompt=1,
+                    latents=None,
+                    output_type="latent",
+                ),
+            )
+        ]
+    )
+
+
+def _make_stubbed_edit_pipeline():
+    pipeline = object.__new__(QwenImageEditPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+
+    pos_embeds, pos_mask = _make_padded_prompt_pair()
+    neg_embeds, neg_mask = _make_padded_prompt_pair()
+
+    pipeline.check_inputs = lambda *args, **kwargs: None
+    pipeline.check_cfg_parallel_validity = lambda *args, **kwargs: None
+    pipeline.vae_scale_factor = 8
+
+    class _GuidanceFreeTransformer:
+        guidance_embeds = False
+        in_channels = 16
+
+    pipeline.transformer = _GuidanceFreeTransformer()
+    pipeline._attention_kwargs = None
+
+    def _fake_encode_prompt(*, prompt_name="prompt", **kwargs):
+        if prompt_name == "negative_prompt":
+            return neg_embeds, neg_mask
+        return pos_embeds, pos_mask
+
+    pipeline.encode_prompt = _fake_encode_prompt
+
+    latents = torch.zeros(1, 4, 8, 8)
+    image_latents = torch.zeros(1, 4, 8, 8)
+    pipeline.prepare_latents = lambda *args, **kwargs: (latents, image_latents)
+    pipeline.prepare_timesteps = lambda *args, **kwargs: (torch.tensor([1.0, 0.5]), 2)
+
+    captured: dict[str, object] = {}
+
+    def _fake_diffuse(
+        _prompt_embeds,
+        _prompt_embeds_mask,
+        _negative_prompt_embeds,
+        _negative_prompt_embeds_mask,
+        latents_arg,
+        _img_shapes,
+        txt_seq_lens,
+        negative_txt_seq_lens,
+        _timesteps,
+        do_true_cfg,
+        _guidance,
+        _true_cfg_scale,
+        **kwargs,
+    ):
+        captured["txt_seq_lens"] = txt_seq_lens
+        captured["negative_txt_seq_lens"] = negative_txt_seq_lens
+        captured["do_true_cfg"] = do_true_cfg
+        return latents_arg
+
+    pipeline.diffuse = _fake_diffuse
+    return pipeline, captured, pos_mask, neg_mask
+
+
+def test_qwen_image_edit_forward_passes_padded_txt_seq_lens_to_diffuse():
+    pipeline, captured, pos_mask, neg_mask = _make_stubbed_edit_pipeline()
+    batch = _make_edit_forward_batch(negative_prompt="bad prompt")
+
+    pipeline.forward(batch)
+
+    assert pos_mask.sum(dim=1).tolist() == [VALID_TOKENS]
+    assert neg_mask.sum(dim=1).tolist() == [VALID_TOKENS]
+    assert captured["txt_seq_lens"] == [PADDED_WIDTH]
+    assert captured["negative_txt_seq_lens"] == [PADDED_WIDTH]
+    assert captured["do_true_cfg"] is True
+
+
+def test_qwen_image_edit_forward_skips_negative_txt_seq_lens_without_cfg():
+    pipeline, captured, pos_mask, _neg_mask = _make_stubbed_edit_pipeline()
+    batch = _make_edit_forward_batch(negative_prompt=None)
+
+    pipeline.forward(batch)
+
+    assert pos_mask.sum(dim=1).tolist() == [VALID_TOKENS]
+    assert captured["txt_seq_lens"] == [PADDED_WIDTH]
+    assert captured["negative_txt_seq_lens"] is None
+    assert captured["do_true_cfg"] is False

tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py

@@ -1,44 +0,0 @@
-"""Regression tests for #4443: padded prompt width vs RoPE table length.
-
-When ``max(txt_seq_lens) < prompt_embeds.shape[1]``, ``QwenEmbedRope`` builds a
-text frequency table that is too short for the padded encoder width.
-"""
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope
-from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-# Representative request-batch padding: valid tokens 10/20, shared width 32.
-PADDED_WIDTH = 32
-VALID_LENS = [10, 20]
-
-
-def test_pre_pr_txt_seq_lens_is_shorter_than_padded_width():
-    mask = torch.zeros(len(VALID_LENS), PADDED_WIDTH, dtype=torch.bool)
-    mask[0, : VALID_LENS[0]] = True
-    mask[1, : VALID_LENS[1]] = True
-
-    pre_pr_txt_seq_lens = mask.sum(dim=1).tolist()
-    assert pre_pr_txt_seq_lens == VALID_LENS
-    assert max(pre_pr_txt_seq_lens) < PADDED_WIDTH
-
-
-def test_post_pr_txt_seq_lens_matches_padded_width():
-    prompt_embeds = torch.zeros(len(VALID_LENS), PADDED_WIDTH, 8)
-    assert txt_seq_lens_from_embeds(prompt_embeds) == [PADDED_WIDTH, PADDED_WIDTH]
-
-
-def test_rope_table_length_matches_padded_width_after_fix():
-    rope = QwenEmbedRope(theta=10000, axes_dim=[16, 56, 56], scale_rope=True)
-    img_shapes = [[(1, 16, 16)]]
-
-    _, freqs_pre_pr = rope(img_shapes, VALID_LENS, device="cpu")
-    _, freqs_post_pr = rope(img_shapes, [PADDED_WIDTH, PADDED_WIDTH], device="cpu")
-
-    assert freqs_pre_pr.shape[0] == max(VALID_LENS)
-    assert freqs_pre_pr.shape[0] < PADDED_WIDTH
-    assert freqs_post_pr.shape[0] == PADDED_WIDTH

tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py

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

tests/diffusion/models/sana_video/test_cache_offload.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 dataclasses import dataclass
 from types import SimpleNamespace
@@ -155,6 +155,11 @@ def test_sana_video_rejects_unvalidated_cache_backends():
         _validate_cache_offload_parallelism(OmniDiffusionConfig(cache_backend="tea_cache"))
 
 
+def test_sana_video_accepts_none_cache_backend():
+    """An explicit ``cache_backend=None`` means "no cache", same as ``"none"``."""
+    _validate_cache_offload_parallelism(OmniDiffusionConfig(cache_backend=None))
+
+
 def test_sana_video_rejects_cache_dit_with_distributed_layerwise_offload():
     config = OmniDiffusionConfig(cache_backend="cache_dit", enable_distributed_layerwise_offload=True)
 

tests/diffusion/models/sana_video/test_transformer_sana_video.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 os
 
@@ -22,7 +22,7 @@ def _init_distributed(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
     os.environ.setdefault("MASTER_ADDR", "localhost")
     os.environ.setdefault("MASTER_PORT", "29501")

tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py

@@ -41,8 +41,8 @@ def tp1_group():
 
     os.environ.setdefault("MASTER_ADDR", "localhost")
     os.environ.setdefault("MASTER_PORT", "29501")
-    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
-    initialize_model_parallel()
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://", backend="gloo")
+    initialize_model_parallel(backend="gloo")
     yield
     cleanup_dist_env_and_memory()
 
@@ -54,7 +54,7 @@ def force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 
@@ -482,8 +482,12 @@ class _LockstepP2PDist:
 def test_tiny_transformer_sp2_lockstep_matches_dense(tp1_group, force_default_gemm, mocker, task) -> None:
     """Two lockstep rank threads over uneven frame shards must reproduce the
     dense tiny-transformer output, and must actually have sharded the work."""
+    from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
     from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
 
+    # These CPU tensor tests must not probe a CUDA device during construction.
+    mocker.patch("vllm_omni.diffusion.attention.selector._cached_get_backend_cls", return_value=SDPABackend)
+
     torch.manual_seed(3)
     model = SanaVideoTransformer3DModel(**_TINY_SP_CONFIG).eval()
     for _, param in sorted(model.named_parameters()):

tests/diffusion/models/sana_video/test_transformer_sana_video_tp.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
 
 """TP unit tests for SANA-Video: distributed RMSNorm, sharded attention and
 sharded GLUMB temporal conv.
@@ -51,7 +51,7 @@ def force_default_gemm(monkeypatch):
 
     monkeypatch.setattr(
         "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
+        lambda *_args, **_kwargs: default_unquantized_gemm,
     )
 
 

tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py

@@ -17,6 +17,7 @@ import torch
 from vllm.lora.layers import BaseLayerWithLoRA
 
 import vllm_omni.diffusion.models.sensenova_u1.paged_decode as paged_decode
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.models.sensenova_u1.paged_decode import (
     BLOCK_SIZE,
     BUCKETS,
@@ -169,7 +170,8 @@ def test_length_lives_in_device_tensors():
 
 
 # ---------------------------------------------------------------------------
-# The capture contract. These need CUDA and the bundled flash-attention kernel.
+# The capture contract. Graph tests need a torch.cuda-compatible accelerator;
+# paged-attention equivalence also needs the bundled kernel's paged API.
 # ---------------------------------------------------------------------------
 
 cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@@ -181,8 +183,7 @@ vllm_flash_attn_only = pytest.mark.skipif(
 
 @cuda_only
 @vllm_flash_attn_only
-@pytest.mark.cuda
-@pytest.mark.L4
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
 def test_paged_attention_matches_sdpa_over_the_same_buffer():
     import torch.nn.functional as F
 
@@ -223,8 +224,7 @@ def test_paged_attention_matches_sdpa_over_the_same_buffer():
 
 @cuda_only
 @vllm_flash_attn_only
-@pytest.mark.cuda
-@pytest.mark.L4
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
 def test_a_captured_graph_follows_a_later_length():
     """One capture has to serve every length in the bucket.
 
@@ -285,8 +285,7 @@ def test_a_captured_graph_follows_a_later_length():
 
 @cuda_only
 @vllm_flash_attn_only
-@pytest.mark.cuda
-@pytest.mark.L4
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
 def test_a_captured_graph_writes_each_step_to_its_own_slot():
     """The write slot has to be a device tensor, not a Python index.
 
@@ -347,8 +346,7 @@ def test_a_captured_graph_writes_each_step_to_its_own_slot():
 
 @cuda_only
 @vllm_flash_attn_only
-@pytest.mark.cuda
-@pytest.mark.L4
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
 def test_a_whole_decode_loop_matches_the_unpaged_path():
     """Per-step kernel equivalence is not loop equivalence.
 
@@ -413,8 +411,7 @@ class _StubLM(torch.nn.Module):
 
 
 @cuda_only
-@pytest.mark.cuda
-@pytest.mark.L4
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
 def test_the_decode_graph_captures_into_the_shared_platform_pool(monkeypatch):
     """``_decode_context`` builds a runner per request, so a private capture pool
     hands every request its own arena and never gives it back. Measured on a
@@ -556,10 +553,12 @@ def test_a_kernel_missing_any_kwarg_we_pass_is_not_supported(monkeypatch):
 def test_only_q_k_and_v_reach_the_kernel_positionally(monkeypatch):
     """A positional standard argument binds to whatever a future wheel puts in
     that slot, and the probe above only checks that the names exist."""
-    seen = {}
+    seen_args: tuple[object, ...] = ()
+    seen_kwargs: dict[str, object] = {}
 
-    def recorder(*args, **kwargs):
-        seen["args"], seen["kwargs"] = args, kwargs
+    def recorder(*args: object, **kwargs: object):
+        nonlocal seen_args, seen_kwargs
+        seen_args, seen_kwargs = args, kwargs
         return torch.zeros(1, N_HEADS, HEAD_DIM)
 
     monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: recorder)
@@ -572,8 +571,8 @@ def test_only_q_k_and_v_reach_the_kernel_positionally(monkeypatch):
         torch.zeros(1, KV_HEADS, 1, HEAD_DIM),
         1.0,
     )
-    assert len(seen["args"]) == 3, "only q, k and v may be positional"
-    assert set(seen["kwargs"]) == {
+    assert len(seen_args) == 3, "only q, k and v may be positional"
+    assert set(seen_kwargs) == {
         "max_seqlen_q",
         "cu_seqlens_q",
         "max_seqlen_k",
@@ -694,8 +693,7 @@ def test_no_capture_crosses_a_request_once_lora_wrappers_are_installed(monkeypat
 
 
 @cuda_only
-@pytest.mark.cuda
-@pytest.mark.L4
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
 def test_a_second_request_replays_the_first_capture(monkeypatch):
     """Reuse is only worth anything if the capture survives it."""
     dev = torch.device("cuda")

tests/diffusion/models/t5_encoder/test_t5_encoder_tp.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
 from transformers import T5Config
@@ -54,6 +54,13 @@ def setup_tp_group(monkeypatch, mocker):
         lambda: 0,
     )
 
+    # vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever
+    # tp_size > 1; these tests drive CPU tensors, so keep the portable masked path.
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda",
+        lambda: False,
+    )
+
     # TP group
     mock_tp_group = mocker.MagicMock()
     mock_tp_group.world_size = 2

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -52,6 +52,7 @@ from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
     get_offload_plan,
 )
+from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan
 from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
 from vllm_omni.host_weight_runtime import MappedHostRegion
 from vllm_omni.platforms import current_omni_platform
@@ -2813,18 +2814,21 @@ class TestDistributedComponentSelection:
 
     def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self):
         """Every rank must reject an unsafe encoder group, including stub ranks."""
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dp_size=2,
-                components=frozenset({"text_encoder"}),
-                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
-            ),
-            torch.device("cpu"),
+
+        class StubEncoderPipeline(nn.Module):
+            _offload_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)})
+
+            def __init__(self):
+                super().__init__()
+                self.text_encoder = nn.Module()
+
+        config = OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=False,
+            dp_size=2,
+            components=frozenset({"text_encoder"}),
+            dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
         )
-        backend.dp_group = object()
-        stub_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)})
 
         with pytest.raises(ValueError, match="not declared replicated"):
-            backend._try_layerwise_offload_encoder(nn.Module(), "text_encoder", stub_plan)
+            resolve_offload_plan(StubEncoderPipeline(), config)

tests/diffusion/offloader/test_layerwise_backend.py

@@ -488,7 +488,8 @@ class TestLayerwiseComponentSelection:
         with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"):
             backend.enable(pipeline)
 
-        assert not pipeline.text_encoder._omni_layerwise_enabled
+        # The plan resolver rejects the topology before any encoder is hooked.
+        assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
 
     def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime):
         pipeline = _LegacyComponentPipeline()

tests/diffusion/offloader/test_plan_resolver.py

@@ -0,0 +1,458 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Unit tests for the shared offload plan resolver."""
+
+from typing import ClassVar
+
+import pytest
+import torch
+from torch import nn
+
+from tests.diffusion.offloader.helpers import (
+    _DummyBlock,
+    _PlainEncoder,
+    _SingleBlockModel,
+    _StagedEncoder,
+    _StagedVAE,
+    patch_offload_runtime,
+)
+from vllm_omni.diffusion.offloader import layerwise_backend as layerwise_backend_module
+from vllm_omni.diffusion.offloader import plan_resolver as plan_resolver_module
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
+from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+CPU = torch.device("cpu")
+
+
+@pytest.fixture
+def patched_offload_runtime(monkeypatch):
+    patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform)
+
+
+def _config(
+    components: set[str] | None = None,
+    strategy: OffloadStrategy = OffloadStrategy.LAYER_WISE,
+    **kwargs,
+) -> OffloadConfig:
+    return OffloadConfig(
+        strategy=strategy,
+        pin_cpu_memory=False,
+        components=None if components is None else frozenset(components),
+        **kwargs,
+    )
+
+
+def _stack_of(component):
+    assert len(component.stacks) == 1
+    return component.stacks[0]
+
+
+class _WanStylePipeline(nn.Module):
+    """Streamable DiT plus a plan-declared, replicated text encoder."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("encoder.block",)},
+        encoder_dlo_weight_replication=frozenset({"text_encoder"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=4)
+        self.text_encoder = _PlainEncoder()
+        self.vae = _StagedVAE()
+
+
+class _StagedComponentPipeline(nn.Module):
+    """MiniMax-H3 shape: two encoder stacks plus pipeline-staged components."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
+        on_demand_component_paths=frozenset({"text_encoder", "vae"}),
+        resident_dit_paths=frozenset({"transformer"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=4)
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _LegacyScanPipeline(nn.Module):
+    """No SupportsComponentDiscovery declaration and no OffloadPlan."""
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=3)
+        self.text_encoder = _PlainEncoder()
+
+
+class TestProducerPrecedence:
+    def test_plan_block_attrs_win_over_class_attributes(self):
+        class TwoStackDiT(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([_DummyBlock() for _ in range(2)])
+                self.alt_blocks = nn.ModuleList([_DummyBlock() for _ in range(3)])
+
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(block_attrs={"transformer": ("alt_blocks",)})
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = TwoStackDiT()
+
+        pipeline = Pipeline()
+        resolved = resolve_offload_plan(pipeline, _config({"dit"}))
+
+        stack = _stack_of(resolved.dits[0])
+        assert stack.attrs == ("alt_blocks",)
+        assert list(stack.blocks) == list(pipeline.transformer.alt_blocks)
+
+    def test_class_attributes_are_used_without_a_plan(self):
+        pipeline = _LegacyScanPipeline()
+        resolved = resolve_offload_plan(pipeline, _config({"dit"}))
+
+        assert list(_stack_of(resolved.dits[0]).blocks) == list(pipeline.transformer.blocks)
+
+    def test_declared_component_type_wins_over_leaf_name(self):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["transformer"]
+            _encoder_modules: ClassVar[list[str]] = ["mllm"]
+            _vae_modules: ClassVar[list[str]] = []
+            _resident_modules: ClassVar[list[str]] = []
+            _offload_plan = OffloadPlan(
+                encoder_component_types={"mllm": "text_encoder"},
+                encoder_block_attrs={"mllm": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.mllm = _PlainEncoder()
+
+        resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"}))
+
+        encoder = resolved.encoders[0]
+        assert encoder.path == "mllm"
+        assert encoder.selected
+        assert len(_stack_of(encoder).blocks) == 2
+
+    def test_component_declaration_wins_over_attribute_scan(self):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["custom_dit"]
+            _encoder_modules: ClassVar[list[str]] = []
+            _vae_modules: ClassVar[list[str]] = []
+            _resident_modules: ClassVar[list[str]] = []
+
+            def __init__(self):
+                super().__init__()
+                self.custom_dit = _SingleBlockModel(num_blocks=2)
+                # Would be picked up by the deprecated attribute scan.
+                self.transformer = _SingleBlockModel(num_blocks=2)
+
+        resolved = resolve_offload_plan(Pipeline(), _config({"dit"}))
+
+        assert [component.path for component in resolved.dits] == ["custom_dit"]
+
+    def test_attribute_scan_warns_once_per_pipeline_class(self, monkeypatch):
+        warnings: list[str] = []
+
+        class _Recorder:
+            def warning(self, message, *args):
+                warnings.append(message % args if args else message)
+
+            def __getattr__(self, _name):
+                return lambda *args, **kwargs: None
+
+        monkeypatch.setattr(plan_resolver_module, "logger", _Recorder())
+
+        class UndeclaredPipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+
+        resolve_offload_plan(UndeclaredPipeline(), _config({"dit"}))
+        resolve_offload_plan(UndeclaredPipeline(), _config({"dit"}))
+
+        scan_warnings = [message for message in warnings if "SupportsComponentDiscovery" in message]
+        assert len(scan_warnings) == 1
+        assert "deprecated" in scan_warnings[0]
+
+
+class TestResolvedTopology:
+    def test_encoder_plan_resolves_one_stack_per_declared_path(self):
+        pipeline = _WanStylePipeline()
+        resolved = resolve_offload_plan(pipeline, _config({"dit", "text_encoder"}))
+
+        encoder = resolved.encoders[0]
+        assert encoder.selected and not encoder.on_demand
+        assert list(_stack_of(encoder).blocks) == list(pipeline.text_encoder.encoder.block)
+        assert [component.path for component in resolved.vaes] == ["vae"]
+
+    def test_staged_plan_resolves_every_declared_stack(self):
+        pipeline = _StagedComponentPipeline()
+        resolved = resolve_offload_plan(pipeline, _config({"text_encoder"}))
+
+        encoder = resolved.encoders[0]
+        assert encoder.on_demand
+        assert [len(stack.blocks) for stack in encoder.stacks] == [2, 2]
+        # An explicit selection never stages a VAE through the legacy plan.
+        assert not resolved.vaes[0].on_demand
+        # DiT is unselected, so it keeps no streaming topology.
+        assert resolved.dits[0].selected is False
+        assert resolved.dits[0].stacks == ()
+
+    def test_omitted_selector_keeps_the_legacy_staged_vae(self):
+        resolved = resolve_offload_plan(_StagedComponentPipeline(), _config())
+
+        assert resolved.vaes[0].on_demand
+        assert resolved.encoders[0].on_demand
+        assert resolved.dits[0].selected
+
+    def test_on_demand_only_plan_resolves_without_stacks(self):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["transformer"]
+            _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+            _vae_modules: ClassVar[list[str]] = ["vae"]
+            _resident_modules: ClassVar[list[str]] = []
+            _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder", "vae"}))
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.text_encoder = _StagedEncoder()
+                self.vae = _StagedVAE()
+
+        resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"}))
+
+        assert resolved.encoders[0].on_demand
+        assert resolved.encoders[0].stacks == ()
+
+    def test_module_strategy_resolves_selection_without_stacks(self):
+        resolved = resolve_offload_plan(
+            _WanStylePipeline(),
+            _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL),
+        )
+
+        assert [component.path for component in resolved.components] == ["transformer", "text_encoder", "vae"]
+        assert all(component.stacks == () for component in resolved.components)
+        assert resolved.encoders[0].selected
+
+    def test_module_strategy_leaves_staged_residency_to_the_pipeline(self):
+        resolved = resolve_offload_plan(
+            _StagedComponentPipeline(),
+            _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL),
+        )
+
+        assert not any(component.on_demand for component in resolved.components)
+
+
+class TestResidentLayers:
+    def _dlo_config(self, resident_layers: int, components: set[str] | None = None) -> OffloadConfig:
+        return _config(
+            components,
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            dlo_transfers={"dit": "rank-local", "text_encoder": "rank-local"},
+            dlo_resident_layers=resident_layers,
+        )
+
+    def test_resident_head_splits_the_declared_dit(self):
+        pipeline = _StagedComponentPipeline()
+        resolved = resolve_offload_plan(pipeline, self._dlo_config(2, {"dit"}))
+
+        stack = _stack_of(resolved.dits[0])
+        assert stack.resident_head == 2
+        assert list(stack.resident) == list(pipeline.transformer.blocks)[:2]
+        assert list(stack.streaming) == list(pipeline.transformer.blocks)[2:]
+
+    def test_resident_count_beyond_the_stack_keeps_every_block(self):
+        resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(9, {"dit"}))
+
+        stack = _stack_of(resolved.dits[0])
+        assert stack.resident_head == 4
+        assert stack.streaming == ()
+
+    def test_resident_head_leaving_one_streaming_block_is_rejected(self):
+        with pytest.raises(ValueError, match="leaves only one streaming block"):
+            resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3, {"dit"}))
+
+    def test_legacy_single_streaming_block_stays_resident(self):
+        resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3))
+
+        stack = _stack_of(resolved.dits[0])
+        assert stack.resident_head == 4
+        assert stack.streaming == ()
+
+    def test_resident_layers_without_a_declared_path_is_rejected(self):
+        with pytest.raises(ValueError, match="no matching resident_dit_paths"):
+            resolve_offload_plan(_WanStylePipeline(), self._dlo_config(2, {"dit"}))
+
+
+class TestValidation:
+    def test_selected_dit_requires_a_discovered_module(self):
+        pipeline = nn.Module()
+        pipeline.text_encoder = _PlainEncoder()
+
+        with pytest.raises(ValueError, match="No DiT/transformer modules found"):
+            resolve_offload_plan(pipeline, _config({"dit"}))
+
+    def test_model_level_selection_requires_both_swap_sides(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+
+        with pytest.raises(ValueError, match="requires an encoder execution stage"):
+            resolve_offload_plan(pipeline, _config({"dit"}, strategy=OffloadStrategy.MODEL_LEVEL))
+
+    def test_selected_text_encoder_requires_a_matching_module(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+
+        with pytest.raises(ValueError, match="No text encoder modules found"):
+            resolve_offload_plan(pipeline, _config({"text_encoder"}))
+
+    def test_selected_encoder_requires_a_streamable_or_staged_plan(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+        pipeline.text_encoder = _PlainEncoder()
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"):
+            resolve_offload_plan(pipeline, _config({"text_encoder"}))
+
+    def test_staged_component_requires_the_explicit_lifecycle(self):
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder"}))
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.text_encoder = _PlainEncoder()
+
+        with pytest.raises(ValueError, match="must implement load_to_device"):
+            resolve_offload_plan(Pipeline(), _config({"text_encoder"}))
+
+    def test_single_block_dit_is_rejected_for_layer_offload(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=1)
+
+        with pytest.raises(ValueError, match="requires at least two streamable"):
+            resolve_offload_plan(pipeline, _config({"dit"}))
+
+    def test_encoder_allgather_requires_declared_replication(self):
+        config = _config(
+            {"text_encoder"},
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            dp_size=2,
+            dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
+        )
+
+        with pytest.raises(ValueError, match="not declared replicated"):
+            resolve_offload_plan(_StagedComponentPipeline(), config)
+
+    def test_two_dits_cannot_own_the_same_blocks(self):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
+            _encoder_modules: ClassVar[list[str]] = []
+            _vae_modules: ClassVar[list[str]] = []
+            _resident_modules: ClassVar[list[str]] = []
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.transformer_2 = _SingleBlockModel(num_blocks=2)
+                self.transformer_2.blocks = self.transformer.blocks
+
+        with pytest.raises(ValueError, match="claimed by both"):
+            resolve_offload_plan(Pipeline(), _config({"dit"}))
+
+    def test_legacy_selection_keeps_unstreamable_components(self):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=1)
+        pipeline.text_encoder = _PlainEncoder()
+
+        resolved = resolve_offload_plan(pipeline, _config())
+
+        assert resolved.dits[0].stacks == ()
+        assert resolved.encoders[0].stacks == ()
+
+
+class TestPurity:
+    def _snapshot(self, pipeline: nn.Module):
+        return {
+            "devices": [(name, tensor.device) for name, tensor in pipeline.named_parameters()],
+            "attrs": {name: sorted(vars(module)) for name, module in pipeline.named_modules()},
+        }
+
+    def test_resolution_does_not_touch_the_model(self):
+        pipeline = _StagedComponentPipeline()
+        before = self._snapshot(pipeline)
+
+        resolve_offload_plan(pipeline, _config({"dit", "text_encoder"}))
+
+        assert self._snapshot(pipeline) == before
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.text_encoder.to_calls == 0
+        assert pipeline.vae.to_calls == 0
+
+    def test_resolution_is_repeatable(self):
+        pipeline = _WanStylePipeline()
+        config = _config({"dit", "text_encoder"})
+
+        first = resolve_offload_plan(pipeline, config)
+        second = resolve_offload_plan(pipeline, config)
+
+        assert [component.path for component in first.components] == [component.path for component in second.components]
+        for left, right in zip(first.components, second.components, strict=True):
+            assert left.selected == right.selected
+            assert left.on_demand == right.on_demand
+            assert [list(stack.blocks) for stack in left.stacks] == [list(stack.blocks) for stack in right.stacks]
+
+    def test_backend_enable_rejects_before_touching_earlier_components(self, patched_offload_runtime):
+        class Pipeline(nn.Module):
+            _dit_modules: ClassVar[list[str]] = ["transformer"]
+            _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"]
+            _vae_modules: ClassVar[list[str]] = ["vae"]
+            _resident_modules: ClassVar[list[str]] = []
+            _offload_plan = OffloadPlan(
+                encoder_component_types={
+                    "text_encoder": "text_encoder",
+                    "text_encoder_2": "text_encoder",
+                },
+                encoder_block_attrs={"text_encoder": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.text_encoder = _PlainEncoder()
+                self.text_encoder_2 = _PlainEncoder()
+                self.vae = _StagedVAE()
+
+        pipeline = Pipeline()
+        backend = LayerWiseOffloadBackend(_config({"text_encoder"}), CPU)
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
+        for block in pipeline.text_encoder.encoder.block:
+            assert getattr(block, "_hook_registry", None) is None
+        assert pipeline.vae.to_calls == 0

tests/diffusion/quantization/test_component_routing.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 component routing for quantization."""
 
 from unittest.mock import MagicMock
@@ -372,7 +372,8 @@ class TestComponentApplyVllmMapper:
                 "thinker.model": "thinker.language_model.model",
                 "talker.model": "talker.language_model.model",
             }
-        ).get_unstacked_mapper()
+            # vLLM 0.29 renamed get_unstacked_mapper to get_rename_mapper.
+        ).get_rename_mapper()
 
         config.apply_vllm_mapper(mapper)
 

tests/diffusion/quantization/test_quantization_fp8.py

@@ -300,7 +300,8 @@ def test_single_stage_zimage_fp8_uses_less_memory():
     assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
 
 
-@hardware_test(res={"cuda": "L4"})
+@hardware_test(res={"cuda": "L4", "xpu": "B60"})
+@pytest.mark.advanced_model
 def test_single_stage_qwen_image_fp8():
     """Qwen-Image (random weights) with FP8 generates valid images."""
     model = "riverclouds/qwen_image_random"

tests/diffusion/streaming_interactive_generation/test_prompt_interaction.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 midway prompt update across runner, batch, pipeline, and entrypoint layers."""
 
 from __future__ import annotations
@@ -7,6 +7,7 @@ from __future__ import annotations
 import asyncio
 import queue
 import time
+from contextlib import contextmanager
 from types import SimpleNamespace
 from typing import Any, cast
 from unittest.mock import MagicMock, patch
@@ -21,6 +22,11 @@ from tests.engine.test_orchestrator import (
 )
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator
+from vllm_omni.diffusion.interaction.modality_handlers.prompt import (
+    PromptInteractionHandler,
+    PromptSession,
+)
 from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.input_batch import InputBatch
@@ -54,30 +60,55 @@ def pipeline() -> HeliosPipeline:
             None,
         )
     )
-    pipeline._prepare_next_chunk = MagicMock()
+    pipeline.prepare_next_chunk = MagicMock()
+    od_config = SimpleNamespace(model_class_name="HeliosPipeline")
+    pipeline._interaction_coordinator = InteractionCoordinator.build(pipeline, od_config)
     return pipeline
 
 
-def _make_diffusion_request_state(*, request_id: str = "req-1") -> StepRequestState:
+def _make_diffusion_request_state(*, request_id: str = "req-1", fps: int | None = None) -> StepRequestState:
     state = StepRequestState(
         request_id=request_id,
-        sampling=SimpleNamespace(num_outputs_per_prompt=1, max_sequence_length=226),  # pyright: ignore[reportArgumentType]
+        sampling=SimpleNamespace(  # pyright: ignore[reportArgumentType]
+            num_outputs_per_prompt=1,
+            max_sequence_length=226,
+            fps=fps,
+        ),
         prompt="hello",
     )
     state.prompt_embeds = torch.zeros(1, 4, 2)
-    state.extra = {}
+    state.extra = {"window_num_frames": 8}
     return state
 
 
-def _make_diffusion_model_runner(*, pipeline, streaming_output: bool = True) -> DiffusionModelRunner:
+def _make_diffusion_model_runner(
+    *,
+    pipeline,
+    streaming_output: bool = True,
+    model_class_name: str = "HeliosPipeline",
+) -> DiffusionModelRunner:
     runner = object.__new__(DiffusionModelRunner)
     runner.pipeline = pipeline
     runner.state_cache = {}
     runner.od_config = SimpleNamespace(  # pyright: ignore[reportAttributeAccessIssue]
-        model_class_name="HeliosPipeline",
+        model_class_name=model_class_name,
         streaming_output=streaming_output,
+        cache_backend=None,
+        parallel_config=SimpleNamespace(use_hsdp=False),
     )
+    runner.vllm_config = object()
+    runner.device = torch.device("cpu")
+    runner.cache_backend = None
+    runner.offload_backend = None
+    runner.prompt_embed_cache = None
+    runner.kv_transfer_manager = SimpleNamespace(
+        receive_multi_kv_cache_distributed=lambda *a, **k: None,
+    )
+    # HeliosPipeline skeletons from object.__new__ are not full SupportsStepExecution
+    # instances; unit tests of submit_interaction stub this. Stepwise path tests
+    # rebind to the real method.
     runner._supports_step_mode = lambda: True
+    runner._interaction_coordinator = None
     return runner
 
 
@@ -108,10 +139,13 @@ class TestPromptUpdateExecution:
         runner.submit_interaction("req-1", _prompt_interaction(transition_chunks=2))
 
         pipeline.encode_prompt.assert_called_once()  # pyright: ignore[reportAttributeAccessIssue]
-        pending = state.extra["pending_prompt_update"]
-        assert pending["event_id"] == "ui-update-1"
-        assert pending["transition_chunks"] == 2
-        assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0))
+        session = state.interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        pending = session.pending_event
+        assert pending is not None
+        assert pending.event_id == "ui-update-1"
+        assert pending.transition_chunks == 2
+        assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
 
     def test_runner_prompt_update_rejects_unsupported_pipeline(self) -> None:
         """Runner rejects prompt updates when the pipeline lacks prompt-update support."""
@@ -135,6 +169,9 @@ class TestPromptUpdateExecution:
             {},
             {"multi_modal_data": {"camera": {"type": "pose"}}},
             {"event": {"prompt": "updated", "multi_modal_data": {"camera": {"type": "pose"}}}},
+            {"event": {"prompt": "updated", "multi_modal_data": {}}},
+            {"event": {"prompt": "updated", "multi_modal_data": None}},
+            {"event": {"prompt": "updated", "multi_modal_data": "bad"}},
         ],
     )
     def test_runner_interaction_rejects_structural_payloads_until_implemented(
@@ -149,6 +186,34 @@ class TestPromptUpdateExecution:
         with pytest.raises(NotImplementedError, match="Only text-only prompt update interactions"):
             runner.submit_interaction("req-1", cast(Any, interaction))
 
+    @pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HeliosPyramidPipeline"])
+    def test_coordinator_registers_prompt_for_helios_aliases(
+        self,
+        pipeline: HeliosPipeline,
+        model_class_name: str,
+    ) -> None:
+        """Both Helios architecture names resolve the prompt handler registry."""
+        od_config = SimpleNamespace(model_class_name=model_class_name)
+        coordinator = InteractionCoordinator.build(pipeline, od_config)
+        assert coordinator.has_modality("prompt")
+
+        runner = _make_diffusion_model_runner(pipeline=pipeline, model_class_name=model_class_name)
+        runner.state_cache["req-1"] = _make_diffusion_request_state()
+        runner.submit_interaction("req-1", _prompt_interaction())
+        session = runner.state_cache["req-1"].interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        assert session.pending_event is not None
+
+    def test_helios_peek_chunk_media_requires_positive_fps(self, pipeline: HeliosPipeline) -> None:
+        """Media-timeline peek must raise on missing/invalid fps instead of coercing to 0.0."""
+        state = _make_diffusion_request_state(fps=None)
+        with pytest.raises(ValueError, match="sampling.fps is required and must be > 0"):
+            pipeline.peek_chunk_media(state)
+
+        state.sampling.fps = 0  # pyright: ignore[reportAttributeAccessIssue]
+        with pytest.raises(ValueError, match="sampling.fps is required and must be > 0"):
+            pipeline.peek_chunk_media(state)
+
     def test_input_batch_refreshes_prompt_embeds_on_version_change(self) -> None:
         """InputBatch rebuilds prompt_embeds when prompt_update_version changes."""
         state = StepRequestState(
@@ -159,28 +224,37 @@ class TestPromptUpdateExecution:
         state.prompt_embeds = torch.zeros(1, 2, 3)
         state.latents = torch.zeros(1, 2)
         state.timesteps = torch.tensor([1.0])
-        state.extra["prompt_update_version"] = 0
+        state.interaction_sessions["prompt"] = PromptSession(version=0)
 
         batch = InputBatch.make_batch([state])
         assert torch.equal(batch.prompt_embeds, torch.zeros(1, 2, 3))  # pyright: ignore[reportArgumentType]
 
         state.prompt_embeds = torch.ones(1, 2, 3)
-        state.extra["prompt_update_version"] = 1
+        state.interaction_sessions["prompt"] = PromptSession(version=1)
         refreshed = InputBatch.make_batch([state], cached_batch=batch)
         assert torch.equal(refreshed.prompt_embeds, torch.ones(1, 2, 3))  # pyright: ignore[reportArgumentType]
 
-    def test_helios_prepare_prompt_update_queues_pending_target(self, pipeline: HeliosPipeline) -> None:
-        """HeliosPipeline queues target embeds without mutating current prompt_embeds."""
+    def test_helios_prompt_handler_enqueue_queues_pending_target(self, pipeline: HeliosPipeline) -> None:
+        """Prompt handler queues target embeds without mutating current prompt_embeds."""
         state = _make_diffusion_request_state()
 
-        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+        PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+            state,
+            event_id="ui-update-1",
+            received_at=0.0,
+            payload={"prompt": "new scene"},
+            transition_chunks=2,
+        )
 
-        pending = state.extra["pending_prompt_update"]
-        assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0))
-        assert pending["transition_chunks"] == 2
+        session = state.interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        pending = session.pending_event
+        assert pending is not None
+        assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
+        assert pending.transition_chunks == 2
         assert torch.equal(state.prompt_embeds, torch.zeros(1, 4, 2))  # pyright: ignore[reportArgumentType]
 
-    def test_helios_prepare_prompt_update_rejects_before_initial_generation(self, pipeline: HeliosPipeline) -> None:
+    def test_helios_prompt_handler_enqueue_rejects_before_initial_generation(self, pipeline: HeliosPipeline) -> None:
         """Reject prompt updates submitted before initial prompt embeds exist."""
         state = _make_diffusion_request_state()
         state.prompt_embeds = None
@@ -189,43 +263,209 @@ class TestPromptUpdateExecution:
             ValueError,
             match="prompt_update is not allowed before initial generation has started",
         ):
-            pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+            PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+                state,
+                event_id="ui-update-1",
+                received_at=0.0,
+                payload={"prompt": "new scene"},
+                transition_chunks=2,
+            )
 
-        assert "pending_prompt_update" not in state.extra
+        session = state.interaction_sessions.get("prompt")
+        assert session is None or (isinstance(session, PromptSession) and session.pending_event is None)
 
-    def test_helios_apply_prompt_update_at_chunk_boundary_starts_transition(self, pipeline: HeliosPipeline) -> None:
+    def test_helios_apply_interaction_at_chunk_boundary_starts_transition(self, pipeline: HeliosPipeline) -> None:
         """At chunk boundary, starts transition state and bumps prompt_update_version."""
-        state = _make_diffusion_request_state()
-        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+        state = _make_diffusion_request_state(fps=None)
+        PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+            state,
+            event_id="ui-update-1",
+            received_at=0.0,
+            payload={"prompt": "new scene"},
+            transition_chunks=2,
+        )
 
-        pipeline._apply_prompt_update_at_chunk_boundary(state)
+        pipeline.apply_interaction_at_chunk_boundary(state)
 
-        assert state.extra.get("prompt_update_state") is not None
+        session = state.interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        assert session.active_event is not None
         assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2))  # pyright: ignore[reportArgumentType]
-        assert state.extra["prompt_update_version"] == 1
-        assert state.extra["prompt_update_chunk_metadata"] == {
+        assert session.version == 1
+        assert state.interaction_chunk_metadata is not None
+        assert state.interaction_chunk_metadata.as_dict() == {
             "started_event_ids": ["ui-update-1"],
             "active_event_ids": ["ui-update-1"],
             "completed_event_ids": [],
         }
 
-    def test_helios_apply_prompt_update_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None:
+    def test_helios_apply_interaction_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None:
         """At chunk boundary, interpolates embeds until the target prompt is reached."""
-        state = _make_diffusion_request_state()
-        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=3)
-        pipeline._apply_prompt_update_at_chunk_boundary(state)
+        state = _make_diffusion_request_state(fps=None)
+        PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+            state,
+            event_id="ui-update-1",
+            received_at=0.0,
+            payload={"prompt": "new scene"},
+            transition_chunks=3,
+        )
+        pipeline.apply_interaction_at_chunk_boundary(state)
 
         assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0 / 3.0))  # pyright: ignore[reportArgumentType]
 
         for _ in range(2):
-            pipeline._apply_prompt_update_at_chunk_boundary(state)
+            pipeline.apply_interaction_at_chunk_boundary(state)
 
         assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0))  # pyright: ignore[reportArgumentType]
-        assert state.extra["prompt_update_version"] == 3
+        session = state.interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        assert session.version == 3
 
         # Further chunk boundaries after completion must not bump the version.
-        pipeline._apply_prompt_update_at_chunk_boundary(state)
-        assert state.extra["prompt_update_version"] == 3
+        pipeline.apply_interaction_at_chunk_boundary(state)
+        assert session.version == 3
+
+    def test_stepwise_runner_applies_prompt_update_and_acks(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Drive the production runner loop: submit → apply(no media args) → ACK on next chunk.
+
+        Also regresses ``sampling.fps is None`` on the prompt-only path (no peek).
+        ACK ids land on the *next* chunk's ``DiffusionOutput`` via
+        ``_attach_stepwise_metadata`` consuming the prior boundary metadata.
+        """
+        import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
+        from vllm_omni.diffusion.interaction.mixin import InteractionMixin
+        from vllm_omni.diffusion.interaction.types import ChunkMediaSpec
+
+        chunks = [
+            DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),  # pyright: ignore[reportArgumentType]
+            DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),  # pyright: ignore[reportArgumentType]
+        ]
+
+        class _InteractiveChunkPipeline(InteractionMixin):
+            device = torch.device("cpu")
+            supports_step_execution = True
+
+            def __init__(self) -> None:
+                self._outputs = chunks
+                self.decode_calls = 0
+                self.transformer = SimpleNamespace(dtype=torch.float32)
+                self.encode_prompt = MagicMock(return_value=(torch.full((1, 4, 2), 2.0), None))
+                self.prepare_next_chunk = MagicMock()
+                od_config = SimpleNamespace(model_class_name="HeliosPipeline")
+                self._interaction_coordinator = InteractionCoordinator.build(self, od_config)  # pyright: ignore[reportArgumentType]
+
+            def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec:
+                # Same contract as Helios: never coerce missing fps to 0.0.
+                fps = state.sampling.fps
+                if fps is None or float(fps) <= 0:
+                    raise ValueError(
+                        "sampling.fps is required and must be > 0 for interaction modalities "
+                        f"that use the chunk media timeline, got {fps!r} "
+                        f"(request_id={state.request_id!r})"
+                    )
+                return ChunkMediaSpec(num_frames=int(state.extra["window_num_frames"]), fps=float(fps))
+
+            def prepare_encode(self, state: StepRequestState) -> StepRequestState:
+                state.prompt_embeds = torch.zeros(1, 4, 2)
+                state.latents = torch.zeros(1, 1)
+                state.timesteps = torch.tensor([1.0, 0.0, 1.0, 0.0])
+                state.step_index = 0
+                state.step_in_chunk = 0
+                state.chunk_num_steps = 2
+                state.total_chunks = 2
+                state.extra = {"window_num_frames": 8}
+                return state
+
+            def denoise_step(self, input_batch: Any, states: Any) -> torch.Tensor:
+                del states
+                return torch.ones_like(input_batch.latents)
+
+            def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor | None) -> None:
+                state.latents = noise_pred
+                state.step_index += 1
+                state.step_in_chunk += 1
+
+            def post_decode(self, state: StepRequestState) -> DiffusionOutput:
+                output = self._outputs[self.decode_calls]
+                self.decode_calls += 1
+                state.chunk_index += 1
+                state.step_index = 0
+                state.step_in_chunk = 0
+                if not state.request_denoise_completed:
+                    state.latents = torch.zeros(1, 1)
+                return output
+
+        pipeline = _InteractiveChunkPipeline()
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        runner._supports_step_mode = lambda: DiffusionModelRunner._supports_step_mode(runner)
+        runner.od_config.streaming_output = True
+        runner.od_config.step_execution = True
+        runner.diffusion_kv_backend = SimpleNamespace(remove_diffusion_kv_requests=lambda *_: 0)  # pyright: ignore[reportAttributeAccessIssue]
+
+        sampling = SimpleNamespace(
+            generator=None,
+            seed=None,
+            generator_device=None,
+            num_inference_steps=4,
+            num_outputs_per_prompt=1,
+            max_sequence_length=226,
+            fps=None,
+        )
+        req = SimpleNamespace(
+            request_id="req",
+            prompt="a prompt",
+            sampling_params=sampling,
+            kv_sender_info=None,
+        )
+
+        @contextmanager
+        def _noop_forward_context(*args: Any, **kwargs: Any):
+            del args, kwargs
+            yield
+
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "is_available", lambda: False)
+
+        scheduler_output = SimpleNamespace(
+            finished_req_ids=set(),
+            scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req, diffusion_kv_metadata=None)],
+            scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+        )
+        # Admit request and run first denoise step (no chunk boundary yet).
+        DiffusionModelRunner.execute_stepwise(runner, scheduler_output)  # pyright: ignore[reportArgumentType]
+        assert "req" in runner.state_cache
+        assert runner.state_cache["req"].sampling.fps is None
+
+        runner.submit_interaction("req", _prompt_interaction("new scene", transition_chunks=2))
+
+        cached = SimpleNamespace(
+            finished_req_ids=set(),
+            scheduled_new_reqs=[],
+            scheduled_cached_reqs=SimpleNamespace(request_ids=["req"]),
+        )
+        # Second step completes chunk 0: apply runs (no peek), metadata staged.
+        chunk0 = DiffusionModelRunner.execute_stepwise(runner, cached).get_request_output("req")  # pyright: ignore[reportArgumentType]
+        assert chunk0 is not None and chunk0.result is not None
+        assert chunk0.result.started_event_ids == []
+        state = runner.state_cache["req"]
+        assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2))  # pyright: ignore[reportArgumentType]
+        assert state.interaction_chunk_metadata is not None
+        assert state.interaction_chunk_metadata.started_event_ids == ["ui-update-1"]
+        pipeline.prepare_next_chunk.assert_called()  # pyright: ignore[reportAttributeAccessIssue]
+
+        # Next chunk decode attaches prior-boundary ACK ids onto DiffusionOutput.
+        DiffusionModelRunner.execute_stepwise(runner, cached)  # pyright: ignore[reportArgumentType]
+        chunk1 = DiffusionModelRunner.execute_stepwise(runner, cached).get_request_output("req")  # pyright: ignore[reportArgumentType]
+        assert chunk1 is not None and chunk1.result is not None
+        assert chunk1.result.started_event_ids == ["ui-update-1"]
+        assert chunk1.result.active_event_ids == ["ui-update-1"]
+        assert chunk1.finished is True
 
 
 class TestPromptUpdateIntegration:
@@ -266,8 +506,8 @@ class TestPromptUpdateIntegration:
 
     @pytest.mark.asyncio
     async def test_runner_prompt_update_failure_surfaces_non_fatal_error(self, pipeline: HeliosPipeline) -> None:
-        """Runner-side prepare_prompt_update rejection is reported through the orchestrator."""
-        pipeline.prepare_prompt_update = MagicMock(  # pyright: ignore[reportAttributeAccessIssue, reportMethodAssignment]
+        """Runner-side prompt encode rejection is reported through the orchestrator."""
+        pipeline.encode_prompt = MagicMock(  # pyright: ignore[reportAttributeAccessIssue]
             side_effect=ValueError("prompt embeds are not ready")
         )
         runner = _make_diffusion_model_runner(pipeline=pipeline)
@@ -293,12 +533,7 @@ class TestPromptUpdateIntegration:
         finally:
             inline_client.shutdown()
 
-        pipeline.prepare_prompt_update.assert_called_once_with(  # pyright: ignore[reportAttributeAccessIssue]
-            runner.state_cache["req-1"],
-            "new scene",
-            "ui-update-1",
-            2,
-        )
+        pipeline.encode_prompt.assert_called_once()  # pyright: ignore[reportAttributeAccessIssue]
         msg = output_queue.get_nowait()
         assert msg == ErrorMessage(
             error="Failed interaction for request req-1: prompt embeds are not ready",
@@ -337,9 +572,10 @@ class TestPromptUpdateIntegration:
 
             outputs = await generate_task
 
-            pending = runner.state_cache[internal_request_id].extra["pending_prompt_update"]
-            assert pending["transition_chunks"] == 2
-            assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0))
+            pending = runner.state_cache[internal_request_id].interaction_sessions["prompt"].pending_event
+            assert pending is not None
+            assert pending.transition_chunks == 2
+            assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
             pipeline.encode_prompt.assert_called_once()  # pyright: ignore[reportAttributeAccessIssue]
             assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
         finally:
@@ -544,7 +780,10 @@ class TestPromptUpdateIntegration:
             ]
 
             deadline = time.monotonic() + 5.0
-            while "pending_prompt_update" not in state.extra:
+            while (
+                not isinstance(state.interaction_sessions.get("prompt"), PromptSession)
+                or state.interaction_sessions["prompt"].pending_event is None
+            ):
                 if time.monotonic() >= deadline:
                     raise TimeoutError("timed out waiting for prompt update during streaming")
                 await asyncio.sleep(0.01)

tests/diffusion/test_diffusion_config_propagation.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests that parallel_config survives the create_default_diffusion roundtrip.
 
 Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862
@@ -131,12 +131,21 @@ class TestCreateDefaultDiffusion:
 
     def test_dtype_serialized_as_string(self):
         stages = StageConfigFactory.create_default_diffusion({"dtype": torch.float16, "model": "x"})
-        assert stages[0]["engine_args"]["dtype"] == "torch.float16"
+        assert stages[0]["engine_args"]["dtype"] == "float16"
 
     def test_cache_backend_defaults_to_none(self):
         stages = StageConfigFactory.create_default_diffusion({"model": "x"})
         assert stages[0]["engine_args"]["cache_backend"] == "none"
 
+    def test_explicit_none_cache_backend_canonicalizes_to_none(self):
+        """Regression for #7032: ``cache_backend=None`` (e.g. an example CLI default)
+        must reach the pipeline as the canonical ``"none"`` string."""
+        od = _roundtrip_diffusion_config(model="x", cache_backend=None)
+        assert od.cache_backend == "none"
+
+        od = OmniDiffusionConfig.from_kwargs(model="x", cache_backend=None)
+        assert od.cache_backend == "none"
+
     def test_single_gpu_default_devices(self):
         stages = StageConfigFactory.create_default_diffusion({"model": "x"})
         assert stages[0]["runtime"]["devices"] == "0"

tests/diffusion/test_diffusion_engine_metrics.py

@@ -140,6 +140,7 @@ class TestMetricKeys:
                 raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
 
         assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
+        assert '"output_ready_wait_time_ms": output_ready_wait_time * 1000' in step_streaming_source
         # step_total_ms is no longer emitted as a metric (no family consumes
         # it); it lives only in the debug log breakdown below.
         assert '"diffusion_engine_total_time_ms"' not in step_streaming_source

tests/diffusion/test_diffusion_scheduler.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 asyncio
 import queue
@@ -10,8 +10,9 @@ import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from pytest_mock import MockerFixture
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
+from tests.helpers.kv_layout import build_kv_cache_tensor
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -104,7 +105,7 @@ def _initialize_paged_scheduler(
     )
     config = KVCacheConfig(
         num_blocks=num_blocks,
-        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
         kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
     )
     scheduler.initialize(
@@ -276,12 +277,16 @@ class TestGetRequestBatchSamplingParamsKey:
         generator: torch.Generator | None = None,
         extra_args: dict | None = None,
         condition_key: tuple | None = None,
+        guidance_scale: float | None = None,
+        guidance_scale_2: float | None = None,
     ) -> OmniDiffusionRequest:
         sp = OmniDiffusionSamplingParams(
             num_inference_steps=num_inference_steps,
             seed=seed,
             generator=generator,
             extra_args=extra_args or {},
+            guidance_scale=guidance_scale,
+            guidance_scale_2=guidance_scale_2,
         )
         return OmniDiffusionRequest(
             prompt="prompt",
@@ -356,6 +361,21 @@ class TestGetRequestBatchSamplingParamsKey:
             self._make(condition_key=("wan22_s2v_condition", True))
         ) != scheduler._build_sampling_params_key(self._make(condition_key=("wan22_s2v_condition", False)))
 
+    def test_distinguishes_explicit_guidance_scale_2(self) -> None:
+        # An omitted guidance_scale_2 is auto-filled from guidance_scale, so a
+        # request that omits it and one that passes the same value explicitly end
+        # up with an identical numeric guidance_scale_2 but different
+        # guidance_scale_2_provided. Pipelines read guidance_scale_2_provided from
+        # the batch's first request to gate image guidance, so the two must not
+        # share a request batch.
+        scheduler = RequestScheduler()
+        omitted = self._make(guidance_scale=2.0)
+        explicit = self._make(guidance_scale=2.0, guidance_scale_2=2.0)
+
+        assert omitted.sampling_params.guidance_scale_2 == explicit.sampling_params.guidance_scale_2
+        assert omitted.sampling_params.guidance_scale_2_provided != explicit.sampling_params.guidance_scale_2_provided
+        assert scheduler._build_sampling_params_key(omitted) != scheduler._build_sampling_params_key(explicit)
+
 
 class TestRequestScheduler:
     def setup_method(self) -> None:

tests/diffusion/test_diffusion_step_pipeline.py

@@ -15,6 +15,7 @@ from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.helpers.mark import hardware_test
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -46,7 +47,6 @@ from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunn
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -1182,7 +1182,7 @@ class TestSupportedPipelines:
     """Step-execution protocol checks for supported pipelines."""
 
     def test_default_stage_config_includes_step_execution(self):
-        stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        stage_cfg = StageConfigFactory.create_default_diffusion(
             {
                 "step_execution": True,
             }

tests/diffusion/utils/test_chunked_mp4.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import hashlib
+
+import av
+import numpy as np
+import pytest
+
+from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder, mux_av_video_audio_bytes
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _frames() -> np.ndarray:
+    rng = np.random.default_rng(6872)
+    return rng.integers(0, 256, size=(8, 16, 24, 3), dtype=np.uint8)
+
+
+def test_chunked_mp4_is_byte_identical_to_whole_mux() -> None:
+    frames = _frames()
+    baseline = mux_av_video_audio_bytes(
+        (av.VideoFrame.from_ndarray(frame, format="rgb24") for frame in frames),
+        width=24,
+        height=16,
+        fps=24,
+        video_codec_options={"threads": "1"},
+    )
+    encoder = ChunkedMP4Encoder(
+        width=24,
+        height=16,
+        fps=24,
+        max_pending=2,
+        video_codec_options={"threads": "1"},
+    )
+    for start in (0, 1, 5):
+        stop = {0: 1, 1: 5, 5: len(frames)}[start]
+        encoder.push(np.ascontiguousarray(frames[start:stop]))
+    chunked = encoder.finish()
+    assert chunked == baseline
+    assert hashlib.sha256(chunked).digest() == hashlib.sha256(baseline).digest()
+
+    audio = np.zeros((2, 320), dtype=np.float32)
+    baseline_audio = mux_av_video_audio_bytes(
+        (av.VideoFrame.from_ndarray(frame, format="rgb24") for frame in frames),
+        width=24,
+        height=16,
+        fps=24,
+        audio_waveform=audio,
+        audio_sample_rate=32000,
+        video_codec_options={"threads": "1"},
+    )
+    encoder_audio = ChunkedMP4Encoder(
+        width=24,
+        height=16,
+        fps=24,
+        audio_waveform=audio,
+        audio_sample_rate=32000,
+        video_codec_options={"threads": "1"},
+    )
+    encoder_audio.push(frames)
+    assert encoder_audio.finish() == baseline_audio
+
+
+def test_chunked_mp4_abort_joins_worker() -> None:
+    encoder = ChunkedMP4Encoder(width=24, height=16, fps=24)
+    encoder.push(_frames()[:1])
+    encoder.abort()
+    assert not encoder._thread.is_alive()
+    with pytest.raises(RuntimeError, match="already closed"):
+        encoder.push(_frames()[:1])
+
+
+def test_chunked_mp4_validates_shape_and_dtype() -> None:
+    encoder = ChunkedMP4Encoder(width=24, height=16, fps=24)
+    with pytest.raises(ValueError, match="shape"):
+        encoder.push(np.zeros((16, 24, 3), dtype=np.uint8))
+    with pytest.raises(ValueError, match="dtype"):
+        encoder.push(np.zeros((1, 16, 24, 3), dtype=np.float32))
+    encoder.abort()
+
+
+def test_chunked_mp4_close_returns_encoded_bytes() -> None:
+    import io
+
+    encoder = ChunkedMP4Encoder(width=24, height=16, fps=24)
+    encoder.push(_frames())
+    result = encoder.close()
+    assert result == encoder.finish() == encoder.close()
+    with av.open(io.BytesIO(result)) as container:
+        assert len(list(container.decode(video=0))) == len(_frames())
+
+
+def test_chunked_mp4_failure_during_push(monkeypatch) -> None:
+    from threading import Event
+
+    from vllm_omni.diffusion.utils import media_utils
+
+    fail = Event()
+    draining = Event()
+    error = RuntimeError("encoding failed")
+
+    def mux(*args, **kwargs):
+        assert fail.wait(5)
+        raise error
+
+    monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux)
+    drain = ChunkedMP4Encoder._drain_until_done
+
+    def notified_drain(self):
+        draining.set()
+        drain(self)
+
+    monkeypatch.setattr(ChunkedMP4Encoder, "_drain_until_done", notified_drain)
+    encoder = ChunkedMP4Encoder(width=24, height=16, fps=24)
+    put = encoder._queue.put
+
+    def fail_before_put(item):
+        fail.set()
+        assert draining.wait(5)
+        put(item)
+
+    monkeypatch.setattr(encoder._queue, "put", fail_before_put)
+    try:
+        with pytest.raises(RuntimeError, match="encoding failed") as exc:
+            encoder.push(_frames())
+        assert exc.value is error
+        with pytest.raises(RuntimeError, match="encoding failed"):
+            encoder.finish()
+    finally:
+        fail.set()
+        encoder.abort()
+    assert not encoder._thread.is_alive()
+
+
+def test_chunked_mp4_flush_failure_does_not_hang(monkeypatch) -> None:
+    from threading import Thread
+
+    from vllm_omni.diffusion.utils import media_utils
+
+    error = RuntimeError("flush failed")
+
+    def mux(frames, **kwargs):
+        list(frames)
+        raise error
+
+    monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux)
+    encoder = ChunkedMP4Encoder(width=24, height=16, fps=24)
+    encoder.push(_frames())
+    errors = []
+
+    def finish():
+        try:
+            encoder.finish()
+        except RuntimeError as exc:
+            errors.append(exc)
+
+    closer = Thread(target=finish, daemon=True)
+    closer.start()
+    closer.join(5)
+    try:
+        assert not closer.is_alive(), "finish hung after the worker consumed its sentinel"
+        assert errors == [error]
+    finally:
+        if closer.is_alive():
+            encoder._send_done()
+            closer.join(5)
+
+
+def test_chunked_mp4_abort_discards_pending_chunks(monkeypatch) -> None:
+    from threading import Event, Thread
+
+    from vllm_omni.diffusion.utils import media_utils
+
+    entered = Event()
+    release = Event()
+    done_sent = Event()
+    encoded = []
+
+    def mux(frames, **kwargs):
+        encoded.append(next(frames))
+        entered.set()
+        assert release.wait(5)
+        encoded.extend(frames)
+        return b"unused"
+
+    monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux)
+    encoder = ChunkedMP4Encoder(width=24, height=16, fps=24, max_pending=1)
+    send_done = encoder._send_done
+
+    def notified_done():
+        send_done()
+        done_sent.set()
+
+    monkeypatch.setattr(encoder, "_send_done", notified_done)
+    encoder.push(_frames()[:1])
+    assert entered.wait(5)
+    encoder.push(_frames()[1:])
+    closer = Thread(target=encoder.abort, daemon=True)
+    closer.start()
+    try:
+        assert done_sent.wait(5), "abort waited for encoding before discarding the pending queue"
+    finally:
+        release.set()
+        closer.join(5)
+    assert not closer.is_alive()
+    assert len(encoded) == 1
+    with pytest.raises(RuntimeError, match="aborted"):
+        encoder.finish()