[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

tests/diffusion/ar_diffusion/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

tests/diffusion/ar_diffusion/test_capability_runner.py

@@ -0,0 +1,317 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU contracts for capability-driven AR-Diffusion sessions."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.experimental.ar_diffusion.capability import (
+    ARDiffusionCrossAttentionKVSpec,
+    ARDiffusionKVBranchSpec,
+    ARDiffusionKVCacheSpec,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+BLOCK = 16
+POS = "positive"
+NEG = "negative"
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def lingbot_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
+    """Single-kv_branch causal DMD: 3 latent frames/block + sink/window + text KV."""
+    return ARDiffusionKVCacheSpec(
+        num_layers=2,
+        num_kv_heads=4,  # already TP-local
+        head_size=64,
+        tokens_per_frame=BLOCK,
+        frames_per_block=3,
+        window_frames=5,
+        sink_frames=1,
+        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+        session_capacity=capacity,
+        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 8),),
+    )
+
+
+def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
+    return ARDiffusionKVCacheSpec(
+        num_layers=2,
+        num_kv_heads=4,
+        head_size=64,
+        tokens_per_frame=BLOCK,
+        frames_per_block=4,
+        window_frames=6,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0), ARDiffusionKVBranchSpec(NEG, 1)),
+        session_capacity=capacity,
+        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 8),),
+    )
+
+
+class CapablePipeline:
+    def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
+        self.spec = spec
+        self.bound_state = None
+        self.binds: list[str] = []
+        self.resets: list[str] = []
+        self.closes: list[str] = []
+
+    def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec:
+        return self.spec
+
+    @contextmanager
+    def bind_ar_diffusion_state(self, session_id: str, state):
+        assert self.bound_state is None
+        self.bound_state = state
+        self.binds.append(session_id)
+        try:
+            yield
+        finally:
+            self.bound_state = None
+
+    def reset_ar_diffusion_session(self, session_id: str) -> None:
+        self.resets.append(session_id)
+
+    def close_ar_diffusion_session(self, session_id: str) -> None:
+        self.closes.append(session_id)
+
+
+class WarmupPipeline(CapablePipeline):
+    def __init__(self, spec: ARDiffusionKVCacheSpec, requests: list[object]) -> None:
+        super().__init__(spec)
+        self.requests = requests
+
+    def ar_diffusion_warmup_requests(self, session_id: str):
+        assert session_id == ARDiffusionModelRunner._WARMUP_SID
+        return iter(self.requests)
+
+
+class BatchCapablePipeline(CapablePipeline):
+    supports_request_batch = True
+
+
+def make_runner(
+    pipeline: object,
+    *,
+    available_bytes: int = 1 << 28,
+    step_execution: bool = False,
+) -> ARDiffusionModelRunner:
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.od_config = SimpleNamespace(
+        max_num_seqs=1,
+        dtype=torch.float32,
+        enforce_eager=True,
+        step_execution=step_execution,
+    )
+    runner.device = torch.device("cpu")
+    runner.pipeline = pipeline
+    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
+    runner.kv_cache = None
+    runner._ar_diffusion_capability = None
+    runner._ar_diffusion_kv_cache_spec = None
+    runner._sessions = OrderedDict()
+    runner._session_capacity = 0
+    runner._perf_e2e_times = []
+    runner._preallocate_kv_cache(available_bytes=available_bytes)
+    return runner
+
+
+def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
+    state = runner._get_or_create_session(session_id)
+    ctx = state.get_kv_caches(kv_branch, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    state.commit_paged_context(kv_branch)
+    return state
+
+
+def test_ar_runner_rejects_pipeline_without_capability():
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.od_config = SimpleNamespace(max_num_seqs=1)
+    runner.pipeline = object()
+    with pytest.raises(TypeError, match="SupportsARDiffusionPipeline"):
+        runner._preallocate_kv_cache(available_bytes=1 << 20)
+
+
+def test_lingbot_like_single_branch_session_reuse_reset_and_close():
+    pipeline = CapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline)
+    kv = runner.kv_cache
+    assert kv is not None
+    assert kv.num_local_kv_branches == 1
+    assert kv.frames_per_block == 3
+    assert kv.spec.window_chunks == 5
+    assert kv.spec.sink_chunks == 1
+    assert kv.cross_attention_lengths == {"text": 8}
+
+    first = commit_one_frame(runner, "s1", "main")
+    assert runner._get_or_create_session("s1") is first
+    k = torch.randn(1, 8, 4, 64)
+    v = torch.randn(1, 8, 4, 64)
+    first.populate_cross_attention("main", "text", [(k, v)] * first.num_layers)
+    assert first.get_cross_attention_kv("main", "text")[0]["k"].shape == k.shape
+
+    runner.reset_session("s1")
+    assert "s1" not in runner._sessions
+    assert "s1" not in kv._cross_sessions
+    assert pipeline.resets == ["s1"]
+    second = runner._get_or_create_session("s1")
+    assert second is not first
+
+    runner.close_session("s1")
+    assert "s1" not in runner._sessions
+    assert pipeline.closes == ["s1"]
+
+
+def test_lingbot_like_sink_survives_sliding_window_eviction():
+    runner = make_runner(CapablePipeline(lingbot_like_spec()))
+    kv = runner.kv_cache
+    assert kv is not None
+
+    for _ in range(8):
+        state = commit_one_frame(runner, "s1", "main")
+
+    table = kv.block_table(state.adapter("main"))
+    assert table[0] != kv.null_block_id
+    assert table[1] == kv.null_block_id
+
+    ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    visible, _ = ctx.video_block_table(torch.device("cpu"))
+    assert visible[0] == table[0]
+    assert len(visible) == 6  # sink + recent window, including current scratch
+
+
+def test_dreamzero_like_two_branches_are_independent():
+    runner = make_runner(CapablePipeline(dreamzero_like_spec()))
+    kv = runner.kv_cache
+    assert kv is not None and kv.num_local_kv_branches == 2
+    state = commit_one_frame(runner, "s1", POS)
+    assert len(kv.window_block_ids(state.adapter(POS))) == 1
+    assert kv.window_block_ids(state.adapter(NEG)) == []
+
+
+def test_lru_eviction_releases_blocks_and_notifies_pipeline():
+    pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
+    runner = make_runner(pipeline)
+    kv = runner.kv_cache
+    assert kv is not None
+    assert runner._session_capacity == 1
+    assert kv.session_capacity == 1
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+    old = commit_one_frame(runner, "old", "main")
+    k = torch.randn(1, 8, 4, 64)
+    old.populate_cross_attention("main", "text", [(k, k)] * old.num_layers)
+    assert kv.manager.block_pool.get_num_free_blocks() < free_total
+
+    runner._get_or_create_session("new")
+
+    assert tuple(runner._sessions) == ("new",)
+    assert pipeline.closes == ["old"]
+    assert "old" not in kv._cross_sessions
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
+    pipeline = CapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline)
+    kv = runner.kv_cache
+    assert kv is not None
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    def boom(self, req, kv_prefetch_job=None):
+        state = pipeline.bound_state
+        ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(torch.device("cpu"))
+        raise RuntimeError("layer exploded")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
+    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+
+    with pytest.raises(RuntimeError, match="layer exploded"):
+        runner.execute_model(request)
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["broken"]
+    assert not runner._sessions
+    assert not kv._adapters
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
+    pipeline = CapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline)
+    kv = runner.kv_cache
+    assert kv is not None
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    def return_after_allocation(self, req, kv_prefetch_job=None):
+        state = pipeline.bound_state
+        ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(torch.device("cpu"))
+        return object()
+
+    def synchronize_boom(device):
+        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")
+    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+
+    with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
+        runner.execute_model(request)
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["broken"]
+    assert not runner._sessions
+    assert not kv._adapters
+    assert not runner._perf_e2e_times
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_ar_runner_rejects_step_and_request_batch_modes():
+    with pytest.raises(ValueError, match="step_execution=True"):
+        make_runner(CapablePipeline(lingbot_like_spec()), step_execution=True)
+    with pytest.raises(ValueError, match="request-batch execution"):
+        make_runner(BatchCapablePipeline(lingbot_like_spec()))
+
+
+def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
+    runner = object.__new__(ARDiffusionModelRunner)
+    with pytest.raises(RuntimeError, match="request-batch execution"):
+        runner.execute_model_batch(None, None)
+    with pytest.raises(RuntimeError, match="step execution"):
+        runner.execute_stepwise(None)
+
+
+def test_model_specific_warmup_provider_is_consumed(monkeypatch):
+    requests = [object(), object()]
+    pipeline = WarmupPipeline(lingbot_like_spec(), requests)
+    runner = make_runner(pipeline)
+    seen: list[object] = []
+    monkeypatch.setattr(runner, "execute_model", seen.append)
+
+    runner._warmup_ar_rollout()
+
+    assert seen == requests
+    assert pipeline.closes == [runner._WARMUP_SID]
+
+
+def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch):
+    pipeline = CapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline)
+    execute = SimpleNamespace(called=False)
+
+    def fail_if_called(request):
+        execute.called = True
+
+    monkeypatch.setattr(runner, "execute_model", fail_if_called)
+    runner._warmup_ar_rollout()
+    assert execute.called is False

tests/diffusion/ar_diffusion/test_engine_routing.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the generic engine-backend dispatcher + AR-Diffusion runner wiring.
+
+``DiffusionEngine.resolve_engine_class`` is a generic dispatcher (``"default"`` / a
+``DiffusionEngine`` subclass / an import-path string). DreamZero selects the
+AR-Diffusion engine via its deploy config's ``engine_backend`` qualname — no
+DreamZero/ar_diffusion knowledge lives in the public base, so the routing check
+here is simply that the ``ARDiffusionEngine`` qualname resolves correctly.
+"""
+
+from dataclasses import fields
+from types import SimpleNamespace
+
+import pytest
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import (
+    AR_DIFFUSION_MODEL_RUNNER_CLS,
+    ARDiffusionEngine,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _cfg(backend):
+    return SimpleNamespace(engine_backend=backend)
+
+
+def test_default_resolves_to_diffusion_engine():
+    assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
+
+
+def test_subclass_type_is_returned():
+    assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine
+
+
+def test_ar_diffusion_qualname_resolves():
+    # How DreamZero's deploy config selects the engine: a full import-path string.
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
+    assert cls is ARDiffusionEngine
+
+
+def test_non_engine_type_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg(dict))
+
+
+def test_qualname_not_an_engine_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg("builtins.dict"))
+
+
+def test_bad_qualname_raises():
+    with pytest.raises(ValueError):
+        DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine"))
+
+
+def test_ar_diffusion_engine_is_diffusion_engine_subclass():
+    assert issubclass(ARDiffusionEngine, DiffusionEngine)
+
+
+def test_config_engine_backend_field_defaults_to_default():
+    field = {f.name: f for f in fields(OmniDiffusionConfig)}["engine_backend"]
+    assert field.default == "default"
+
+
+# --- AR-Diffusion worker/runner wiring -----------------------------------------------
+
+
+def test_ar_diffusion_model_runner_cls_resolves():
+    cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS)
+    assert cls.__name__ == "ARDiffusionModelRunner"
+
+
+def test_engine_class_declares_its_runner():
+    """Routing lives on the engine class (review: hsliuustc0106) — engines
+    never mutate od_config; the worker resolves the runner off the class."""
+    assert ARDiffusionEngine.default_diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
+    assert DiffusionEngine.default_diffusion_model_runner_cls is None
+
+
+def _select_runner(od_config, platform_default="PLATFORM"):
+    """Mirror of the worker's runner-selection chain (override > engine > platform)."""
+    runner_override = getattr(od_config, "diffusion_model_runner_cls", None)
+    if isinstance(runner_override, str) and runner_override:
+        return runner_override
+    engine_cls = DiffusionEngine.resolve_engine_class(od_config)
+    engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+    if isinstance(engine_runner, str) and engine_runner:
+        return engine_runner
+    return platform_default
+
+
+def test_runner_selection_chain():
+    ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+    # engine-declared runner wins over platform default
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar))
+        == AR_DIFFUSION_MODEL_RUNNER_CLS
+    )
+    # explicit override wins over the engine's declaration
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar))
+        == "my.custom.Runner"
+    )
+    # default engine -> platform default
+    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend="default")) == "PLATFORM"
+
+
+def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine():
+    """Review (hsliuustc0106): a stale config with engine_backend='default'
+    must fail at init, not mid-forward on the first KV access."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace(engine_backend="default"))
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace())
+
+
+def test_runner_rejects_multi_sequence_config():
+    """Review (hsliuustc0106): the paged path writes batch index 0 only, so
+    max_num_seqs > 1 must be rejected at init."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
+    with pytest.raises(ValueError, match="max_num_seqs=1"):
+        ARDiffusionModelRunner._preallocate_kv_cache(fake)
+
+
+def test_execute_model_accepts_kv_prefetch_job_kwarg():
+    """Regression: the worker passes ``kv_prefetch_job`` to every runner
+    (added with the KV prefetch path, #4448); the AR-Diffusion override must
+    accept and forward it or DreamZero crashes on the warm-up dummy run."""
+    import inspect
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    for cls in (DiffusionModelRunner, ARDiffusionModelRunner):
+        params = inspect.signature(cls.execute_model).parameters
+        assert "kv_prefetch_job" in params, cls

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
+
+Covers the request adapter, the chunk-window spec/manager (registration + the
+eviction policy), and the pool builder — exercised against the installed vLLM
+KV stack on CPU (block bookkeeping only, no GPU tensors).
+"""
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionRequestAdapter,
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    build_kv_manager,
+    compute_num_blocks,
+)
+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
+
+
+def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+    )
+
+
+# --- ChunkWindowSpec registration -------------------------------------------
+
+
+def test_spec_registration_resolves_to_chunk_window_manager():
+    # Without explicit registration the MRO walk would fall back to the parent
+    # SlidingWindowManager; assert the subclass manager wins.
+    spec = make_spec()
+    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+
+
+def test_spec_kind_is_sliding_window():
+    assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW
+
+
+def test_spec_rejects_inconsistent_window():
+    with pytest.raises(ValueError):
+        ChunkWindowSpec(
+            block_size=BLOCK,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float16,
+            sliding_window=99,  # != window_chunks * chunk_size
+            chunk_size=BLOCK,
+            window_chunks=2,
+        )
+
+
+# --- eviction policy (pure) -------------------------------------------------
+
+
+def test_sliding_replace_keeps_window():
+    # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
+    # the snap is to chunk boundaries.
+    def skip(n):
+        return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+
+    assert skip(32) == 0  # nothing past the window yet
+    assert skip(48) == 16  # one chunk fell out of the window
+    assert skip(64) == 32
+
+
+def test_sliding_replace_snaps_to_chunk_boundary():
+    # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
+    skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+    assert skip % 16 == 0 and skip == 16
+
+
+def test_sink_chunks_protected():
+    # sink = 1 chunk (16 tokens) is never skipped.
+    skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
+    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
+    assert skip == 32
+
+
+def test_reset_at_boundary_drops_completed_past_sink():
+    skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True)
+    # completed = 48; sink = 16 -> drop 32
+    assert skip == 32
+
+
+# --- ARDiffusionKVConfig ------------------------------------------------------------
+
+
+def test_kv_config_sliding_window_property():
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
+
+
+# --- ARDiffusionRequestAdapter ------------------------------------------------------
+
+
+def test_adapter_advances_per_chunk_not_per_step():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16)
+    assert a.num_computed_tokens == 0
+    assert a.num_tokens == 16  # in-flight chunk
+    a.on_chunk_committed()
+    assert a.num_computed_tokens == 16
+    assert a.num_tokens == 32
+
+
+def test_adapter_accounts_for_prefill_prefix():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
+    assert a.num_computed_tokens == 4
+    assert a.num_prompt_tokens == 4
+    assert a.num_tokens == 20
+
+
+def test_adapter_status_is_vllm_enum():
+    assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus)
+
+
+# --- pool / manager ---------------------------------------------------------
+
+
+def test_compute_num_blocks():
+    # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks.
+    assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
+
+
+def test_paged_pool_layout_exposes_flat_slot_views():
+    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
+        num_blocks=4,
+        block_size=BLOCK,
+        num_layers=1,
+        num_kv_heads=4,
+        head_dim=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
+    assert k_pools[0].shape == (4 * BLOCK, 4, 64)
+    assert v_pools[0].shape == (4 * BLOCK, 4, 64)
+
+    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])
+
+
+def test_build_manager_allocate_free_roundtrip():
+    """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.
+
+    This is the adapter conformance check — if the adapter were missing an
+    attribute the manager reads, allocate_slots/free would raise here.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
+    free_before = mgr.block_pool.get_num_free_blocks()
+
+    adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK)
+    blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
+    assert blocks is not None
+    assert mgr.block_pool.get_num_free_blocks() < free_before
+
+    mgr.free(adapter)
+    assert mgr.block_pool.get_num_free_blocks() == free_before
+
+
+def test_cross_attn_pool_deducted_from_self_attn_budget():
+    """The cross-attn pool is allocated directly; its bytes are subtracted from
+    the self-attn paged-pool budget so the two together stay within the free
+    memory budget (review: zwhzzz0821)."""
+    L = 512
+    avail = 1 << 30  # 1 GiB
+    kv = ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
+        num_layers=2,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=avail,
+        kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
+        session_capacity=1,
+        cross_attention_lengths={"text": L},
+        device=torch.device("cpu"),
+    )
+    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
+    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+    assert kv.num_blocks == expected
+    assert expected > 1 * (2 + 1) + 2  # above the local-slot floor, so the cross deduction is what's tested
+
+
+def _make_kv(
+    *,
+    local_branches,
+    num_frame_per_block=2,
+    window_chunks=9,
+    max_scratch_tokens_per_branch=0,
+):
+    kv_branches = (
+        (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0))
+        if local_branches == 1
+        else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
+    )
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+        num_layers=1,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 16,  # tiny -> the floor binds
+        device=torch.device("cpu"),
+        kv_branches=kv_branches,
+        session_capacity=1,
+        frames_per_block=num_frame_per_block,
+        max_scratch_tokens_per_branch=max_scratch_tokens_per_branch,
+    )
+
+
+def test_pool_floor_is_branch_aware():
+    """CFG-parallel rank (one local kv_branch) sizes for one window + in-flight chunk;
+    a single-process run (both kv_branches) sizes for two. Scratch scales the same way."""
+    one = _make_kv(local_branches=1)
+    two = _make_kv(local_branches=2)
+
+    assert one.managed_num_blocks == 1 * (9 + 2) + 2  # 13
+    assert two.managed_num_blocks == 2 * (9 + 2) + 2  # 24
+    assert one.scratch_num_blocks == one.scratch_blocks_per_kv_branch
+    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_kv_branch
+    assert one.scratch_blocks_per_kv_branch == 2
+    assert one.num_blocks_total == 13 + one.scratch_blocks_per_kv_branch
+
+
+def test_scratch_capacity_is_derived_from_declared_geometry(monkeypatch):
+    monkeypatch.delenv("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", raising=False)
+    kv = _make_kv(
+        local_branches=1,
+        num_frame_per_block=6,
+        max_scratch_tokens_per_branch=BLOCK + 1,
+    )
+    assert kv.scratch_blocks_per_kv_branch == 8  # six video blocks + two auxiliary blocks
+
+
+def test_scratch_env_override_cannot_reduce_declared_minimum(monkeypatch):
+    monkeypatch.setenv("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "1")
+    kv = _make_kv(
+        local_branches=1,
+        num_frame_per_block=6,
+        max_scratch_tokens_per_branch=BLOCK + 1,
+    )
+    assert kv.scratch_blocks_per_kv_branch == 8
+
+
+def test_scratch_maps_to_slot_zero_with_one_local_branch():
+    """A CFG-parallel rank runs exactly one kv_branch: whichever CFG side it is,
+    its scratch lands in the rank's single slot (no dead second slot)."""
+    one = _make_kv(local_branches=1)
+    assert one.scratch_block_ids("positive", 0, 2) == one.scratch_block_ids("negative", 0, 2)
+
+    two = _make_kv(local_branches=2)
+    assert two.scratch_block_ids("positive", 0, 2) != two.scratch_block_ids("negative", 0, 2)
+
+
+def test_scratch_exhaustion_still_raises():
+    one = _make_kv(local_branches=1)
+    cap = one.scratch_blocks_per_kv_branch
+    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
+        one.scratch_block_ids("positive", 0, cap + 1)
+
+
+def test_non_contiguous_branch_indices_rejected():
+    with pytest.raises(ValueError, match="contiguous"):
+        ARDiffusionKVCache(
+            ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2),
+            num_layers=1,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float32,
+            block_size=BLOCK,
+            max_model_len=4096,
+            available_bytes=1 << 16,
+            kv_branches=(ARDiffusionKVBranchSpec("main", 1),),
+            session_capacity=1,
+        )

tests/diffusion/ar_diffusion/test_kv_orchestration.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+DIMS = dict(
+    num_layers=2,
+    num_kv_heads=4,
+    head_size=64,
+    dtype=torch.float16,
+    block_size=16,
+    kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1)),
+    session_capacity=2,
+)
+
+
+def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+    return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
+
+
+def test_requires_enabled_config():
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_requires_bounded_window():
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None)
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_full_request_lifecycle_and_eviction():
+    """begin -> per-chunk allocate/slots/commit over a long rollout -> free.
+
+    Exercises the orchestrator end-to-end and asserts the chunk window bounds
+    memory (pool plateaus) and frees cleanly.
+    """
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    adapter = kv.begin_request("req-0")
+    free_after = []
+    for k in range(10):
+        block_table = kv.allocate_chunk(adapter)
+        slots = kv.chunk_write_slots(adapter)
+        # Slots target real blocks in this chunk's table; length == chunk_size.
+        assert len(slots) == kv.spec.chunk_size
+        used = {int(s) // kv.block_size for s in slots}
+        assert kv.null_block_id not in used
+        assert used <= set(block_table)
+        # Resident window stays bounded by window + the in-flight chunk.
+        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
+        free_after.append(kv.manager.block_pool.get_num_free_blocks())
+        kv.commit_chunk(adapter)
+
+    # Pool memory plateaus once the window is full (eviction recycles blocks).
+    assert free_after[-1] == free_after[kv.spec.window_chunks]
+
+    kv.end_request(adapter)
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_num_computed_advances_per_chunk():
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    a = kv.begin_request("r")
+    assert a.num_computed_tokens == 0
+    kv.allocate_chunk(a)
+    kv.commit_chunk(a)
+    assert a.num_computed_tokens == 16
+    kv.end_request(a)
+
+
+def test_state_close_frees_both_branch_blocks():
+    """ARDiffusionKVState.close() returns both CFG kv_branches' pool blocks to the pool.
+
+    This is the primitive the runner's LRU eviction relies on: when a session is
+    evicted, close() must free the blocks both adapters hold, or session churn
+    leaks pool ownership (review P1).
+    """
+    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    pos = kv.begin_request("bde__s")
+    neg = kv.begin_request("bde__s__neg")
+    state = ARDiffusionKVState(
+        kv,
+        "s",
+        {"positive": pos, "negative": neg},
+        num_layers=kv.num_layers,
+    )
+    for _ in range(3):
+        for adapter in (pos, neg):
+            kv.allocate_chunk(adapter)
+            kv.commit_chunk(adapter)
+    # Both kv_branches hold resident blocks now.
+    assert kv.manager.block_pool.get_num_free_blocks() < free_total
+
+    state.close()
+    # All blocks returned — no leak across an evicted session.
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for AR-Diffusion paged self-attention contexts."""
+
+from __future__ import annotations
+
+import subprocess
+from importlib.util import find_spec
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
+    ar_diffusion_paged_attention,
+    paged_write_attn,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+POS = "positive"
+NEG = "negative"
+
+
+def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=dtype,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 26,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0), ARDiffusionKVBranchSpec(NEG, 1)),
+        session_capacity=2,
+        frames_per_block=2,
+        max_scratch_tokens_per_branch=BLOCK,
+        device=device,
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, "s1", {POS: pos, NEG: neg}, num_layers=num_layers)
+
+
+def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5)
+    probs = torch.softmax(scores, dim=-1).to(value.dtype)
+    return torch.einsum("bhqk,bkhd->bqhd", probs, value)
+
+
+def _cuda_flash_attn_usable() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        spec = find_spec("vllm.vllm_flash_attn")
+        if spec is None or spec.origin is None:
+            return True
+        fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so"
+        linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5)
+    except Exception:
+        return True
+    if "libcudart.so.13" not in linked:
+        return True
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            text=True,
+            timeout=5,
+        )
+        driver_major = int(out.splitlines()[0].split(".")[0])
+    except Exception:
+        return True
+    return driver_major >= 580
+
+
+def _commit_video_span(
+    kv: ARDiffusionKVCache,
+    st: ARDiffusionKVState,
+    *,
+    kv_branch: str,
+    n_chunks: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
+    st.commit_paged_context(kv_branch)
+    return k, v
+
+
+def test_paged_context_allocates_lazily_and_commits_after_forward():
+    _, st = make_state()
+
+    contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+    assert st.adapter(POS).completed_chunks == 0
+    assert ctx.current_video_slot_mapping is None
+
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.adapter(POS).completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 1
+
+    st.commit_paged_context(POS)
+    assert st.adapter(POS).completed_chunks == 1
+    assert st._committed[POS] == BLOCK
+
+
+def test_scratch_video_and_action_blocks_do_not_commit():
+    kv, st = make_state()
+
+    ctx = st.get_kv_caches(POS, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    ctx.ensure_action_slots(3, torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(POS, 0, 2)
+    assert ctx.action_scratch_block_ids == kv.scratch_block_ids(POS, 2, 1)
+    st.commit_paged_context(POS)
+    assert st.adapter(POS).completed_chunks == 0
+    assert st._committed[POS] == 0
+
+
+def test_pipeline_kv_get_paged_path_has_no_gather_backend():
+    kv, st = make_state()
+    assert not hasattr(kv, "gather_window_all_layers")
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._ar_diffusion_kv_state = st
+    contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False)
+
+    assert len(contexts) == 1
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+
+
+@pytest.mark.parametrize("history_chunks", [0, 1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k_parts: list[torch.Tensor] = []
+    history_v_parts: list[torch.Tensor] = []
+    if history_chunks:
+        k, v = _commit_video_span(
+            kv,
+            st,
+            kv_branch=POS,
+            n_chunks=history_chunks,
+            dtype=dtype,
+            device=device,
+        )
+        history_k_parts.append(k)
+        history_v_parts.append(v)
+
+    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len,
+        query_len=query.shape[1],
+        device=device,
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    if history_k_parts:
+        history_k = torch.cat(history_k_parts, dim=1)
+        history_v = torch.cat(history_v_parts, dim=1)
+    else:
+        history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+    before = st.adapter(POS).completed_chunks
+    st.commit_paged_context(POS)
+    assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
+
+
+@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
+@pytest.mark.parametrize("history_chunks", [1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
+    pytest.importorskip("vllm.vllm_flash_attn")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k, history_v = _commit_video_span(
+        kv,
+        st,
+        kv_branch=POS,
+        n_chunks=history_chunks,
+        dtype=dtype,
+        device=device,
+    )
+
+    layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    action_k = action_v = None
+    if action_len:
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    # The production path: once-per-forward host prep, then the fused
+    # write+attend custom op consuming the NamedTuple payload.
+    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
+    inputs = layer_ctx.to_layer_inputs()
+    assert isinstance(inputs, ARDiffusionPagedLayerInputs)
+    paged = paged_write_attn(
+        inputs,
+        query[0],
+        current_k[0],
+        current_v[0],
+        action_k[0] if action_len else None,
+        action_v[0] if action_len else None,
+        HEAD_DIM**-0.5,
+    ).unsqueeze(0)
+
+    # Direct python-fn call on the same (already written) pools must be
+    # bit-exact: identical kernel, identical inputs.
+    direct = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=ctx.block_table,
+        query_start_loc=ctx.query_start_loc,
+        seq_lens=ctx.seq_lens,
+        max_query_len=ctx.max_query_len,
+        max_seq_len=ctx.max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    assert torch.equal(paged, direct)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
+def test_block_table_padded_to_fixed_width():
+    """Shapes must be constant across window growth: only values change."""
+    device = torch.device("cpu")
+    kv, st = make_state(window_chunks=2)
+
+    ctx1 = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(POS)
+
+    ctx2 = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(POS)
+
+    # 1-block vs 2-block visible history: same table width, same max_seq_len.
+    assert ctx1.block_table.shape == ctx2.block_table.shape
+    assert ctx1.max_seq_len == ctx2.max_seq_len
+    expected_width = kv.spec.sliding_window // kv.block_size + 1
+    assert ctx1.block_table.shape == (1, expected_width)
+    # Real lengths live in seq_lens, not the padded table.
+    assert int(ctx1.seq_lens[0]) == BLOCK
+    assert int(ctx2.seq_lens[0]) == 2 * BLOCK
+
+
+def test_prepare_is_idempotent_and_layers_share_metadata():
+    device = torch.device("cpu")
+    _, st = make_state(num_layers=2)
+    contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
+    fctx = contexts[0].forward_ctx
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    table = fctx.block_table
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    assert fctx.block_table is table  # memoized, not rebuilt
+
+    i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs()
+    # 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer
+    # value guards on the shared block code object.
+    assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0
+    assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1
+    # All layers share the same metadata tensor objects.
+    assert i0.block_table is i1.block_table
+    assert i0.seq_lens is i1.seq_lens
+    assert i0.video_slots is i1.video_slots
+
+
+def test_layer_inputs_before_prepare_raises():
+    _, st = make_state()
+    layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0]
+    with pytest.raises(RuntimeError, match="before prepare"):
+        layer_ctx.to_layer_inputs()
+
+
+def test_custom_op_registration_idempotent():
+    import importlib
+    import sys
+
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+    mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"
+    saved = sys.modules.pop(mod)
+    try:
+        importlib.import_module(mod)  # re-registration must not raise
+    finally:
+        sys.modules[mod] = saved
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+
+
+def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
+    """The op must trace as one opaque node: fullgraph OK, and changed tensor
+    VALUES (new slots / block ids) must not trigger recompilation."""
+    import torch._dynamo
+
+    device = torch.device("cpu")
+    kv, st = make_state(num_layers=2, window_chunks=2)
+
+    def run_one_forward(commit):
+        contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit)
+        fctx = contexts[0].forward_ctx
+        fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+        q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        # Both layers through ONE compiled fn: layer_idx is a tensor, so a
+        # different layer must NOT recompile (all 40 DiT blocks share the
+        # block-forward code object in production).
+        for layer_ctx in contexts:
+            out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
+        st.commit_paged_context(POS)
+        return out
+
+    torch._dynamo.reset()
+    try:
+        from torch._dynamo.testing import CompileCounter
+
+        counter = CompileCounter()
+
+        def fn(inputs, q, k, v):
+            return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
+
+        compiled = torch.compile(fn, backend=counter, fullgraph=True)
+
+        run_one_forward(commit=True)  # history grows between calls ->
+        run_one_forward(commit=True)  # block-table VALUES change, shapes don't
+        run_one_forward(commit=False)
+
+        assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
+    finally:
+        # Leave a clean dynamo state for later suites in the same pytest process
+        # (e.g. model_executor transformers models).
+        torch._dynamo.reset()

tests/diffusion/ar_diffusion/test_pipeline_bridge.py

@@ -0,0 +1,280 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+POS = "positive"
+NEG = "negative"
+
+
+def make_state(num_layers=1, window_chunks=4, cross_attn_length=0, shared_local_index=False):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 24,
+        kv_branches=(
+            ARDiffusionKVBranchSpec(POS, 0),
+            ARDiffusionKVBranchSpec(NEG, 0 if shared_local_index else 1),
+        ),
+        session_capacity=2,
+        cross_attention_lengths={"text": cross_attn_length} if cross_attn_length else None,
+        device=torch.device("cpu"),
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, "s1", {POS: pos, NEG: neg}, num_layers=num_layers)
+
+
+def _prepare_and_commit(st: ARDiffusionKVState, kv_branch: str, n_chunks: int) -> None:
+    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    st.commit_paged_context(kv_branch)
+
+
+def test_get_returns_paged_layer_contexts_when_nothing_committed():
+    _, st = make_state(num_layers=2)
+    for kv_branch in (POS, NEG):
+        contexts = st.get_kv_caches(kv_branch, seq_len=BLOCK, commit_current=False)
+        assert len(contexts) == 2
+        assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts)
+        assert contexts[0].history_block_ids == []
+
+
+def test_managed_paged_context_commits_only_after_forward():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(POS, seq_len=2 * BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+
+    assert st.adapter(POS).completed_chunks == 0
+    assert kv.window_block_ids(st.adapter(POS)) == []
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.adapter(POS).completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 2
+
+    st.commit_paged_context(POS)
+
+    assert st.adapter(POS).completed_chunks == 2
+    assert st._committed[POS] == 2 * BLOCK
+    assert len(kv.window_block_ids(st.adapter(POS))) == 2
+
+
+def test_scratch_paged_context_does_not_commit():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
+    ctx = contexts[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(POS, 0, 1)
+    st.commit_paged_context(POS)
+
+    assert st.adapter(POS).completed_chunks == 0
+    assert st._committed[POS] == 0
+    assert kv.window_block_ids(st.adapter(POS)) == []
+
+
+def test_eviction_bounds_resident_window():
+    kv, st = make_state(num_layers=1, window_chunks=3)
+    _prepare_and_commit(st, POS, 3)
+    assert len(kv.window_block_ids(st.adapter(POS))) == 3
+
+    _prepare_and_commit(st, POS, 2)
+    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    visible_blocks, video_len = ctx.video_block_table(torch.device("cpu"))
+    assert len(visible_blocks) == 3
+    assert video_len == 3 * BLOCK
+
+
+def test_branches_are_independent():
+    kv, st = make_state()
+    _prepare_and_commit(st, POS, 1)
+    assert len(kv.window_block_ids(st.adapter(POS))) == 1
+    assert kv.window_block_ids(st.adapter(NEG)) == []
+
+
+def test_reset_clears_session_window():
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, POS, 2)
+    _prepare_and_commit(st, NEG, 1)
+    assert st._committed[POS] == 2 * BLOCK and st._committed[NEG] == BLOCK
+    free_before = kv.manager.block_pool.get_num_free_blocks()
+
+    st.reset()
+
+    assert st._committed == {POS: 0, NEG: 0}
+    assert st._paged_pending == {POS: None, NEG: None}
+    assert kv.window_block_ids(st.adapter(POS)) == []
+    assert kv.window_block_ids(st.adapter(NEG)) == []
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before
+    _prepare_and_commit(st, POS, 1)
+    assert len(kv.window_block_ids(st.adapter(POS))) == 1
+
+
+def test_get_cross_attention_kv_raises_before_populate():
+    _, st = make_state(cross_attn_length=8)
+    for kv_branch in (POS, NEG):
+        with pytest.raises(RuntimeError, match="before it was populated"):
+            st.get_cross_attention_kv(kv_branch, "text")
+
+
+def test_get_cross_attention_kv_returns_session_pool_dicts_when_populated():
+    L = 8
+    kv, st = make_state(num_layers=2, cross_attn_length=L)
+    written = []
+    for _ in range(2):
+        k = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        v = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        written.append((k, v))
+    st.populate_cross_attention(POS, "text", written)
+
+    out = st.get_cross_attention_kv(POS, "text")
+    assert len(out) == 2
+    for i, (k, v) in enumerate(written):
+        assert out[i]["is_init"] is True
+        assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM)
+        assert torch.equal(out[i]["k"], k)
+        assert torch.equal(out[i]["v"], v)
+
+
+def test_window_reset_can_keep_named_cross_attention_cache():
+    kv, st = make_state(num_layers=1, cross_attn_length=8)
+    k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
+    v = torch.randn(1, 8, N_HEADS, HEAD_DIM)
+    st.populate_cross_attention(POS, "text", [(k, v)])
+    old_adapter = st.adapter(POS)
+
+    st.reset(keep_cross_attention=("text",))
+
+    assert st.adapter(POS) is not old_adapter
+    assert torch.equal(st.get_cross_attention_kv(POS, "text")[0]["k"], k)
+    assert "s1" in kv._cross_sessions
+
+    st.reset()
+    assert "s1" not in kv._cross_sessions
+    assert not st.is_cross_attention_populated(POS, "text")
+
+
+def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
+    _, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
+    assert st.kv_cache.num_local_kv_branches == 1
+    pos_layers = [
+        (torch.full((1, 8, N_HEADS, HEAD_DIM), 1.0), torch.full((1, 8, N_HEADS, HEAD_DIM), 2.0)) for _ in range(2)
+    ]
+    neg_layers = [
+        (torch.full((1, 8, N_HEADS, HEAD_DIM), 3.0), torch.full((1, 8, N_HEADS, HEAD_DIM), 4.0)) for _ in range(2)
+    ]
+
+    st.populate_cross_attention(POS, "text", pos_layers)
+    st.populate_cross_attention(NEG, "text", neg_layers)
+
+    for layer in st.get_cross_attention_kv(POS, "text"):
+        assert torch.count_nonzero(layer["k"] != 1.0) == 0
+        assert torch.count_nonzero(layer["v"] != 2.0) == 0
+    for layer in st.get_cross_attention_kv(NEG, "text"):
+        assert torch.count_nonzero(layer["k"] != 3.0) == 0
+        assert torch.count_nonzero(layer["v"] != 4.0) == 0
+
+
+def test_cross_attention_population_is_transactional():
+    kv, st = make_state(num_layers=2, cross_attn_length=8)
+    original = [
+        (torch.full((1, 8, N_HEADS, HEAD_DIM), value), torch.full((1, 8, N_HEADS, HEAD_DIM), -value))
+        for value in (1.0, 2.0)
+    ]
+    replacement = (
+        torch.full((1, 8, N_HEADS, HEAD_DIM), 9.0),
+        torch.full((1, 8, N_HEADS, HEAD_DIM), -9.0),
+    )
+
+    def interrupted_population():
+        yield replacement
+        raise RuntimeError("projection failed")
+
+    with pytest.raises(RuntimeError, match="projection failed"):
+        st.populate_cross_attention(POS, "text", interrupted_population())
+    assert not st.is_cross_attention_populated(POS, "text")
+    assert "s1" not in kv._cross_sessions
+
+    st.populate_cross_attention(POS, "text", original)
+
+    with pytest.raises(RuntimeError, match="projection failed"):
+        st.populate_cross_attention(POS, "text", interrupted_population())
+    out = st.get_cross_attention_kv(POS, "text")
+    for layer_index, layer in enumerate(out):
+        assert torch.equal(layer["k"], original[layer_index][0])
+        assert torch.equal(layer["v"], original[layer_index][1])
+
+    for invalid_layers in (original[:1], [*original, replacement]):
+        with pytest.raises(ValueError, match="expected 2 layers"):
+            st.populate_cross_attention(POS, "text", invalid_layers)
+        out = st.get_cross_attention_kv(POS, "text")
+        for layer_index, layer in enumerate(out):
+            assert torch.equal(layer["k"], original[layer_index][0])
+            assert torch.equal(layer["v"], original[layer_index][1])
+
+    invalid_shape = (torch.empty(1, 7, N_HEADS, HEAD_DIM), original[0][1])
+    with pytest.raises(ValueError, match="expected k/v shape"):
+        st.populate_cross_attention(POS, "text", [invalid_shape, original[1]])
+    out = st.get_cross_attention_kv(POS, "text")
+    assert torch.equal(out[0]["k"], original[0][0])
+
+
+def test_closed_state_cannot_repopulate_cross_attention():
+    kv, st = make_state(num_layers=1, cross_attn_length=8)
+    layer = (torch.ones(1, 8, N_HEADS, HEAD_DIM), torch.ones(1, 8, N_HEADS, HEAD_DIM))
+    st.close()
+
+    with pytest.raises(RuntimeError, match="is closed"):
+        st.populate_cross_attention(POS, "text", [layer])
+    assert "s1" not in kv._cross_sessions
+
+
+def test_kv_create_is_noop_engine_owns_allocation():
+    from unittest.mock import MagicMock
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    p._ar_diffusion_kv_state = object()
+    state = MagicMock()
+    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+    state.create_kv_caches.assert_not_called()
+
+
+def test_get_kv_cache_requires_frame_aligned_seqlen():
+    _, st = make_state()
+    with pytest.raises(AssertionError, match="frame-aligned"):
+        st.get_kv_caches(POS, seq_len=BLOCK + 1, commit_current=True)
+
+
+def test_prepare_one_branch_allocates_nothing_for_the_other():
+    """CFG-parallel laziness: a rank prepares only the kv_branch it runs."""
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    pos_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    neg_ctx = st.get_kv_caches(NEG, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+
+    pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK)
+
+    assert pos_ctx._allocated_video
+    assert not neg_ctx._allocated_video
+    assert kv.window_block_ids(st.adapter(NEG)) == []

tests/diffusion/ar_diffusion/test_slot_mapping.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionRequestAdapter,
+    ChunkWindowSpec,
+    build_kv_manager,
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    resident_block_ids,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+BLOCK = 16
+
+
+def make_spec(*, window_chunks=2):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * BLOCK,
+        chunk_size=BLOCK,
+        window_chunks=window_chunks,
+    )
+
+
+def test_slot_mapping_matches_block_offsets():
+    # block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4).
+    slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4)
+    assert slots.tolist() == [20, 21, 8, 9, 36, 37]
+
+
+def test_slot_mapping_rejects_bad_block_size():
+    with pytest.raises(ValueError):
+        compute_slot_mapping([1], [0], block_size=0)
+
+
+def test_chunk_slot_mapping_targets_current_chunk():
+    # After 2 committed chunks (num_computed=32), the in-flight chunk maps to the
+    # block covering positions 32..47 — i.e. block table index 2.
+    block_ids = [7, 4, 3]  # block index 2 -> physical block 3
+    slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16)
+    # positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51
+    assert slots.tolist() == [48, 49, 50, 51]
+
+
+def test_resident_block_ids_excludes_null():
+    assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3]
+
+
+def test_blocktable_build_single_request():
+    """Integration: a real KVCacheManager block table -> a valid slot mapping.
+
+    Every slot for the in-flight chunk must land inside a real (non-null) block
+    that the manager actually allocated for the request.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
+    null_id = mgr.block_pool.null_block.block_id
+    adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK)
+
+    mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
+    block_ids = mgr.get_block_ids(adapter.request_id)[0]
+    real = set(resident_block_ids(block_ids, null_id))
+    assert real, "expected at least one real block after allocation"
+
+    slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size)
+    # Each slot resolves to a real allocated block, none to null_block.
+    blocks_used = {int(s) // spec.block_size for s in slots}
+    assert blocks_used <= real
+    assert null_id not in blocks_used
+    assert len(slots) == spec.chunk_size

tests/diffusion/attention/test_attention_config.py

@@ -288,6 +288,13 @@ class TestBuildAttentionConfig:
 class TestOmniDiffusionConfigAttentionParsing:
     """Test OmniDiffusionConfig attention shorthand and structured config."""
 
+    @pytest.fixture(autouse=True)
+    def _clear_diffusion_attention_backend_env(self, monkeypatch):
+        # OmniDiffusionConfig.__post_init__ applies DIFFUSION_ATTENTION_BACKEND via
+        # build_attention_config(); clear it so these tests assert config defaults,
+        # not whatever the process inherited from CI / sibling tests.
+        monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False)
+
     def test_diffusion_attention_backend_sets_default(self):
         config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
         assert isinstance(config.diffusion_attention_config, AttentionConfig)

tests/diffusion/attention/test_attention_sp.py

@@ -49,6 +49,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.platforms import current_omni_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def update_environment_variables(envs_dict: dict[str, str]):
     """Update multiple environment variables with logging."""

tests/diffusion/attention/test_flash_attn.py

@@ -18,6 +18,9 @@ from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
 from vllm_omni.platforms import current_omni_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
 HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
 flash_attn_func = fa.flash_attn_func  # noqa: N813

tests/diffusion/attention/test_piecewise_attn.py

@@ -27,6 +27,9 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 DEVICE = torch.device("cpu")
 
 

tests/diffusion/attention/test_sage_attn3.py

@@ -9,6 +9,9 @@ import pytest
 import torch
 from vllm.platforms import current_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"
 
 

tests/diffusion/attention/test_ulysses_uaa.py

@@ -22,6 +22,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
 from vllm_omni.platforms import current_omni_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def _find_free_port() -> int:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:

tests/diffusion/distributed/test_comm.py

@@ -1,12 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for SeqAllToAll4D and SeqAllToAll5D communication primitives."""
+"""Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives.
+
+CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same
+checks on real multi-GPU NCCL collectives.
+"""
+
+from __future__ import annotations
 
 import os
+from typing import Literal
 
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -16,56 +24,51 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.platforms import current_omni_platform
 
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
 
-def update_environment_variables(envs_dict: dict[str, str]):
-    """Update multiple environment variables with logging."""
-    for k, v in envs_dict.items():
-        os.environ[k] = v
+DeviceKind = Literal["cpu", "cuda"]
 
 
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("seq_len_per_rank", [8])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [32])
-@pytest.mark.parametrize("use_sync", [False, True])
-def test_4d_identity(
+def _update_environment_variables(envs_dict: dict[str, str]) -> None:
+    for key, value in envs_dict.items():
+        os.environ[key] = value
+
+
+def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
+    if device_kind == "cpu":
+        return torch.device("cpu")
+    return torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+
+
+def _init_worker(
+    local_rank: int,
     world_size: int,
-    dtype: torch.dtype,
-    batch_size: int,
-    seq_len_per_rank: int,
-    num_heads: int,
-    head_size: int,
-    use_sync: bool,
-):
-    """Test that two consecutive all-to-all operations return the original input."""
-    # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
-    if available_gpus < world_size:
-        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
-    # Ensure num_heads is divisible by world_size
-    if num_heads % world_size != 0:
-        pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
-
-    # Run test with multiprocessing spawn
-    torch.multiprocessing.spawn(
-        _test_4d_identity_worker,
-        args=(
-            world_size,
-            dtype,
-            batch_size,
-            seq_len_per_rank,
-            num_heads,
-            head_size,
-            use_sync,
-        ),
-        nprocs=world_size,
+    master_port: int,
+    device_kind: DeviceKind,
+) -> None:
+    _update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(master_port),
+        }
     )
+    backend = "gloo" if device_kind == "cpu" else None
+    init_distributed_environment(backend=backend)
 
 
-def _test_4d_identity_worker(
+def _close_tolerance(dtype: torch.dtype) -> tuple[float, float]:
+    if dtype == torch.bfloat16:
+        return 1e-3, 1e-3
+    if dtype == torch.float16:
+        return 1e-3, 1e-3
+    return 1e-5, 1e-5
+
+
+def _run_4d_identity(
     local_rank: int,
     world_size: int,
     dtype: torch.dtype,
@@ -74,29 +77,17 @@ def _test_4d_identity_worker(
     num_heads: int,
     head_size: int,
     use_sync: bool,
-):
-    """Worker function for test_4d_identity."""
-    # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
 
-    # Set environment variables for distributed training
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "29500",
-        }
-    )
+    _init_worker(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(ulysses_degree=world_size)
+    sp_group = get_sp_group().ulysses_group
 
-    # Initialize distributed environment
-    init_distributed_environment()
-    initialize_model_parallel(ulysses_degree=world_size)  # test ulysses sp by default
-    sp_group = get_sp_group().ulysses_group  # get ulysses sp group not ring sp group
-
-    # Create input tensor: (bs, seqlen/P, hc, hs)
     torch.manual_seed(42 + local_rank)
     input_tensor = torch.randn(
         batch_size,
@@ -106,100 +97,36 @@ def _test_4d_identity_worker(
         dtype=dtype,
         device=device,
     )
-
-    # Save original input for comparison
     original_input = input_tensor.clone()
 
-    # First all-to-all: (bs, seqlen/P, hc, hs) -> (bs, seqlen, hc/P, hs)
-    intermediate = SeqAllToAll4D.apply(
-        sp_group,
-        input_tensor,
-        2,  # scatter head dimension
-        1,  # gather sequence dimension
-        use_sync,
-    )
-
-    # Verify intermediate shape
-    expected_shape = (
+    intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, use_sync)
+    expected_intermediate_shape = (
         batch_size,
         seq_len_per_rank * world_size,
         num_heads // world_size,
         head_size,
     )
-    assert intermediate.shape == expected_shape, (
-        f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
+    assert intermediate.shape == expected_intermediate_shape, (
+        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
     )
 
-    # Second all-to-all: (bs, seqlen, hc/P, hs) -> (bs, seqlen/P, hc, hs)
-    output = SeqAllToAll4D.apply(
-        sp_group,
-        intermediate,
-        1,  # scatter sequence dimension
-        2,  # gather head dimension
-        use_sync,
-    )
-
-    # Verify output shape matches input
+    output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, use_sync)
     assert output.shape == original_input.shape, (
         f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
     )
 
-    # Verify output matches original input
+    rtol, atol = _close_tolerance(dtype)
     torch.testing.assert_close(
         output,
         original_input,
-        rtol=1e-5,
-        atol=1e-5,
+        rtol=rtol,
+        atol=atol,
         msg="Output does not match original input after two all-to-all operations",
     )
-
-    # Cleanup distributed environment
     destroy_distributed_env()
 
 
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("seq_len_per_rank", [8])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [32])
-@pytest.mark.parametrize("use_sync", [False, True])
-def test_5d_identity(
-    world_size: int,
-    dtype: torch.dtype,
-    batch_size: int,
-    seq_len_per_rank: int,
-    num_heads: int,
-    head_size: int,
-    use_sync: bool,
-):
-    """Test that two consecutive all-to-all operations return the original input."""
-    # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
-    if available_gpus < world_size:
-        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
-    # Ensure num_heads is divisible by world_size
-    if num_heads % world_size != 0:
-        pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
-
-    # Run test with multiprocessing spawn
-    torch.multiprocessing.spawn(
-        _test_5d_identity_worker,
-        args=(
-            world_size,
-            dtype,
-            batch_size,
-            seq_len_per_rank,
-            num_heads,
-            head_size,
-            use_sync,
-        ),
-        nprocs=world_size,
-    )
-
-
-def _test_5d_identity_worker(
+def _run_5d_identity(
     local_rank: int,
     world_size: int,
     dtype: torch.dtype,
@@ -208,199 +135,329 @@ def _test_5d_identity_worker(
     num_heads: int,
     head_size: int,
     use_sync: bool,
-):
-    """Worker function for test_5d_identity."""
-    # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
 
-    # Set environment variables for distributed training
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "29500",
-        }
-    )
+    _init_worker(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(ulysses_degree=world_size)
+    sp_group = get_sp_group().ulysses_group
 
-    # Initialize distributed environment
-    init_distributed_environment()
-    initialize_model_parallel(ulysses_degree=world_size)  # test ulysses sp by default
-    sp_group = get_sp_group().ulysses_group  # get ulysses sp group not ring sp group
-
-    # Create input tensor: (bs, seqlen/P, 3, hc, hs)
-    # The '3' dimension is for Q, K, V
     torch.manual_seed(42 + local_rank)
     input_tensor = torch.randn(
         batch_size,
         seq_len_per_rank,
-        3,  # Q, K, V
+        3,
         num_heads,
         head_size,
         dtype=dtype,
         device=device,
     )
-
-    # Save original input for comparison
     original_input = input_tensor.clone()
 
-    # First all-to-all: (bs, seqlen/P, 3, hc, hs) -> (bs, seqlen, 3, hc/P, hs)
-    intermediate = SeqAllToAll5D.apply(
-        sp_group,
-        input_tensor,
-        3,  # scatter head dimension
-        1,  # gather sequence dimension
-        use_sync,
-    )
-
-    # Verify intermediate shape
-    expected_shape = (
+    intermediate = SeqAllToAll5D.apply(sp_group, input_tensor, 3, 1, use_sync)
+    expected_intermediate_shape = (
         batch_size,
         seq_len_per_rank * world_size,
         3,
         num_heads // world_size,
         head_size,
     )
-    assert intermediate.shape == expected_shape, (
-        f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
+    assert intermediate.shape == expected_intermediate_shape, (
+        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
     )
 
-    # Second all-to-all: (bs, seqlen, 3, hc/P, hs) -> (bs, seqlen/P, 3, hc, hs)
-    output = SeqAllToAll5D.apply(
-        sp_group,
-        intermediate,
-        1,  # scatter sequence dimension
-        3,  # gather head dimension
-        use_sync,
-    )
-
-    # Verify output shape matches input
+    output = SeqAllToAll5D.apply(sp_group, intermediate, 1, 3, use_sync)
     assert output.shape == original_input.shape, (
         f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
     )
 
-    # Verify output matches original input
+    rtol, atol = _close_tolerance(dtype)
     torch.testing.assert_close(
         output,
         original_input,
-        rtol=1e-5,
-        atol=1e-5,
+        rtol=rtol,
+        atol=atol,
         msg="Output does not match original input after two all-to-all operations",
     )
-
-    # Cleanup distributed environment
     destroy_distributed_env()
 
 
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [128])
-def test_ring_p2p(
-    world_size: int,
-    dtype: torch.dtype,
-    batch_size: int,
-    num_heads: int,
-    head_size: int,
-):
-    """Test Ring P2P communication (send_recv)."""
-    # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
-    if available_gpus < world_size:
-        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
-    torch.multiprocessing.spawn(
-        _test_ring_p2p_worker,
-        args=(world_size, dtype, batch_size, num_heads, head_size),
-        nprocs=world_size,
-    )
-
-
-def _test_ring_p2p_worker(
+def _run_ring_p2p(
     local_rank: int,
     world_size: int,
     dtype: torch.dtype,
     batch_size: int,
     num_heads: int,
     head_size: int,
-):
-    """Worker for Ring P2P test."""
-    import sys
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
 
-    # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
-
-    # Set env vars
-    # Use a different port to avoid conflict with other tests if run in parallel
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "29501",
-        }
-    )
-
-    # Init distributed
     try:
-        init_distributed_environment()
-        # Ring degree = world_size to test ring group
+        _init_worker(local_rank, world_size, master_port, device_kind)
         initialize_model_parallel(ring_degree=world_size)
         sp_group = get_sp_group()
-
-        print(f"[Rank {local_rank}] Initialized. Ring group size: {sp_group.ring_group.size()}")
-        sys.stdout.flush()
-
-        # Create RingComm
         comm = RingComm(sp_group.ring_group)
 
-        # Create tensor: rank-specific data
-        # (batch, num_heads, head_size)
-        # Fill with rank value + 1 to avoid 0 and make verification easy
         input_tensor = torch.full(
-            (batch_size, num_heads, head_size), fill_value=float(local_rank + 1), dtype=dtype, device=device
+            (batch_size, num_heads, head_size),
+            fill_value=float(local_rank + 1),
+            dtype=dtype,
+            device=device,
         )
-
-        print(f"[Rank {local_rank}] Input sum: {input_tensor.sum().item()}")
-        sys.stdout.flush()
-
-        # Send input, receive from prev
-        # RingComm.send_recv sends to next, receives from prev
-        t0 = __import__("time").time()
         recv_tensor = comm.send_recv(input_tensor)
         comm.commit()
         comm.wait()
-        t1 = __import__("time").time()
 
-        print(f"[Rank {local_rank}] Communication done in {t1 - t0:.4f}s")
-
-        # Verify
-        # Expected value: from (rank - 1) % world_size
         prev_rank = (local_rank - 1 + world_size) % world_size
         expected_value = float(prev_rank + 1)
-
-        recv_sum = recv_tensor.sum().item()
-        print(f"[Rank {local_rank}] Received sum: {recv_sum}, Expected value: {expected_value}")
-        sys.stdout.flush()
-
         expected_tensor = torch.full_like(recv_tensor, fill_value=expected_value)
-
-        # Use a slightly loose tolerance for bfloat16
+        rtol, atol = _close_tolerance(dtype)
         torch.testing.assert_close(
-            recv_tensor, expected_tensor, rtol=1e-3, atol=1e-3, msg=f"[Rank {local_rank}] Data mismatch!"
+            recv_tensor,
+            expected_tensor,
+            rtol=rtol,
+            atol=atol,
+            msg=f"[Rank {local_rank}] Ring P2P data mismatch",
         )
-        print(f"[Rank {local_rank}] Verification PASSED")
-
-    except Exception as e:
-        print(f"[Rank {local_rank}] FAILED with error: {e}")
-        import traceback
-
-        traceback.print_exc()
-        raise e
     finally:
         destroy_distributed_env()
+
+
+def _spawn_4d_identity(
+    *,
+    world_size: int,
+    dtype: torch.dtype,
+    batch_size: int,
+    seq_len_per_rank: int,
+    num_heads: int,
+    head_size: int,
+    use_sync: bool,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    torch.multiprocessing.spawn(
+        _run_4d_identity,
+        args=(
+            world_size,
+            dtype,
+            batch_size,
+            seq_len_per_rank,
+            num_heads,
+            head_size,
+            use_sync,
+            device_kind,
+            master_port,
+        ),
+        nprocs=world_size,
+    )
+
+
+def _spawn_5d_identity(
+    *,
+    world_size: int,
+    dtype: torch.dtype,
+    batch_size: int,
+    seq_len_per_rank: int,
+    num_heads: int,
+    head_size: int,
+    use_sync: bool,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    torch.multiprocessing.spawn(
+        _run_5d_identity,
+        args=(
+            world_size,
+            dtype,
+            batch_size,
+            seq_len_per_rank,
+            num_heads,
+            head_size,
+            use_sync,
+            device_kind,
+            master_port,
+        ),
+        nprocs=world_size,
+    )
+
+
+def _spawn_ring_p2p(
+    *,
+    world_size: int,
+    dtype: torch.dtype,
+    batch_size: int,
+    num_heads: int,
+    head_size: int,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    torch.multiprocessing.spawn(
+        _run_ring_p2p,
+        args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
+        nprocs=world_size,
+    )
+
+
+def _require_heads_divisible(num_heads: int, world_size: int) -> None:
+    if num_heads % world_size != 0:
+        pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
+
+
+def _require_gpus(world_size: int) -> None:
+    available_gpus = current_omni_platform.get_device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
+
+# ---------------------------------------------------------------------------
+# CPU: gloo collectives on CPU tensors (no GPU)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_4d_identity(world_size: int, use_sync: bool):
+    _require_heads_divisible(8, world_size)
+    _spawn_4d_identity(
+        world_size=world_size,
+        dtype=torch.float32,
+        batch_size=2,
+        seq_len_per_rank=8,
+        num_heads=8,
+        head_size=32,
+        use_sync=use_sync,
+        device_kind="cpu",
+        master_port=29610,
+    )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_5d_identity(world_size: int, use_sync: bool):
+    _require_heads_divisible(8, world_size)
+    _spawn_5d_identity(
+        world_size=world_size,
+        dtype=torch.float32,
+        batch_size=2,
+        seq_len_per_rank=8,
+        num_heads=8,
+        head_size=32,
+        use_sync=use_sync,
+        device_kind="cpu",
+        master_port=29611,
+    )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_ring_p2p(world_size: int):
+    _spawn_ring_p2p(
+        world_size=world_size,
+        dtype=torch.float32,
+        batch_size=2,
+        num_heads=8,
+        head_size=128,
+        device_kind="cpu",
+        master_port=29612,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Nightly: same checks on real multi-GPU NCCL collectives
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "world_size",
+    [
+        pytest.param(2, marks=_L4_TWO_GPU),
+        pytest.param(4, marks=_L4_FOUR_GPU),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_4d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
+    _require_gpus(world_size)
+    _require_heads_divisible(8, world_size)
+    _spawn_4d_identity(
+        world_size=world_size,
+        dtype=dtype,
+        batch_size=2,
+        seq_len_per_rank=8,
+        num_heads=8,
+        head_size=32,
+        use_sync=use_sync,
+        device_kind="cuda",
+        master_port=29500,
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "world_size",
+    [
+        pytest.param(2, marks=_L4_TWO_GPU),
+        pytest.param(4, marks=_L4_FOUR_GPU),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_5d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
+    _require_gpus(world_size)
+    _require_heads_divisible(8, world_size)
+    _spawn_5d_identity(
+        world_size=world_size,
+        dtype=dtype,
+        batch_size=2,
+        seq_len_per_rank=8,
+        num_heads=8,
+        head_size=32,
+        use_sync=use_sync,
+        device_kind="cuda",
+        master_port=29502,
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "world_size",
+    [
+        pytest.param(2, marks=_L4_TWO_GPU),
+        pytest.param(4, marks=_L4_FOUR_GPU),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
+    _require_gpus(world_size)
+    _spawn_ring_p2p(
+        world_size=world_size,
+        dtype=dtype,
+        batch_size=2,
+        num_heads=8,
+        head_size=128,
+        device_kind="cuda",
+        master_port=29501,
+    )

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -1,8 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for PipelineParallelMixin and related PP communication helpers.
+
+CPU unit tests exercise mixin logic without a process group. Distributed CPU
+tests use gloo on CPU tensors; nightly parity tests run the same checks on
+real multi-GPU NCCL collectives.
+"""
+
+from __future__ import annotations
 
 import os
 import socket
+from typing import Literal
 
 import pytest
 import torch
@@ -11,6 +20,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
 import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -22,7 +32,12 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.parallel]
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+
+DeviceKind = Literal["cpu", "cuda"]
+_UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def _find_free_port() -> str:
@@ -36,6 +51,22 @@ def update_environment_variables(envs_dict: dict[str, str]) -> None:
         os.environ[k] = v
 
 
+def _cleanup_distributed() -> None:
+    """Destroy omni distributed state and clear test-process leftovers.
+
+    ``destroy_distributed_env`` leaves ``vllm.distributed.parallel_state._PP``
+    aliased to the destroyed group. Parent-process baselines must clear that
+    dangling reference (and RANK/MASTER_* env) so later unit tests that call
+    ``vllm.distributed.initialize_model_parallel`` do not see stale state.
+    """
+    destroy_distributed_env()
+    import vllm.distributed.parallel_state as vllm_parallel_state
+
+    vllm_parallel_state._PP = None
+    for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
+        os.environ.pop(key, None)
+
+
 # ---------------------------------------------------------------------------
 # Shared stubs used by both unit and distributed tests
 # ---------------------------------------------------------------------------
@@ -143,7 +174,7 @@ class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
 class TestAsyncLatents:
     """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
 
-    pytestmark = [pytest.mark.cpu]
+    pytestmark = _UNIT_MARKS
 
     def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
         return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
@@ -214,7 +245,7 @@ class TestAsyncLatents:
 class TestSyncPPSend:
     """Verifies PipelineParallelMixin's internal PP-send flush."""
 
-    pytestmark = [pytest.mark.cpu]
+    pytestmark = _UNIT_MARKS
 
     @staticmethod
     def _make_pipeline() -> PipelineParallelMixin:
@@ -247,7 +278,7 @@ class TestSyncPPSend:
 class TestDiffuseWrapper:
     """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
 
-    pytestmark = [pytest.mark.cpu]
+    pytestmark = _UNIT_MARKS
 
     def test_diffuse_flushes_pending_sends_on_success(self):
         work = FakeWork()
@@ -289,7 +320,7 @@ class TestDiffuseWrapper:
 
 
 class TestVaeDecodeGuard:
-    pytestmark = [pytest.mark.cpu]
+    pytestmark = _UNIT_MARKS
 
     @staticmethod
     def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
@@ -351,6 +382,8 @@ class TestVaeDecodeGuard:
         assert pipeline.vae.decode.__doc__ == "Original decode docstring."
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
 @pytest.mark.cpu
 def test_pipeline_parallel_requires_cfg_mixin():
     with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
@@ -359,6 +392,8 @@ def test_pipeline_parallel_requires_cfg_mixin():
             pass
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
 @pytest.mark.cpu
 def test_pipeline_parallel_requires_mro_before_cfg_mixin():
     with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
@@ -372,10 +407,22 @@ def test_pipeline_parallel_requires_mro_before_cfg_mixin():
 # ---------------------------------------------------------------------------
 
 
-def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
+def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
+    if device_kind == "cpu":
+        return torch.device("cpu")
+    return torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+
+
+def init_dist(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    device_kind: DeviceKind,
+) -> torch.device:
     """Initialise the distributed environment for a spawned worker."""
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
     update_environment_variables(
         {
             "RANK": str(local_rank),
@@ -385,13 +432,31 @@ def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.devic
             "MASTER_PORT": master_port,
         }
     )
-    init_distributed_environment()
+    backend = "gloo" if device_kind == "cpu" else None
+    init_distributed_environment(backend=backend)
     return device
 
 
+def _mp_backend(device_kind: DeviceKind) -> str | None:
+    """Backend for model-parallel process groups.
+
+    ``GroupCoordinator.all_gather`` always uses ``device_group``. CPU tensors
+    therefore require gloo device groups; otherwise a CUDA host's default NCCL
+    backend breaks PP+CFG CPU tests at CFG all_gather (pp2-cfg2).
+    isend/irecv of CPU tensors can already fall back to ``cpu_group``.
+    """
+    return "gloo" if device_kind == "cpu" else None
+
+
+def _seed_device(input_seed: int, device: torch.device) -> None:
+    torch.manual_seed(input_seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(input_seed)
+
+
 def make_pipeline_and_inputs(
     test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
-) -> tuple["MockPipelineParallel", dict, dict | None]:
+) -> tuple[MockPipelineParallel, dict, dict | None]:
     """Create a MockPipelineParallel and seeded inputs from a test_config dict.
 
     Must be called after ``initialize_model_parallel`` so that ``make_layers``
@@ -408,16 +473,12 @@ def make_pipeline_and_inputs(
         dtype=dtype,
     )
 
-    torch.manual_seed(test_config["input_seed"])
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(test_config["input_seed"])
+    _seed_device(test_config["input_seed"], device)
     pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
 
     negative_kwargs = None
     if do_true_cfg:
-        torch.manual_seed(test_config["input_seed"] + 1)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
+        _seed_device(test_config["input_seed"] + 1, device)
         neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
         negative_kwargs = {"x": neg_x}
 
@@ -425,19 +486,23 @@ def make_pipeline_and_inputs(
 
 
 # ---------------------------------------------------------------------------
-# 3.  isend_tensor_dict / irecv_tensor_dict  (2 GPUs)
+# 3.  isend_tensor_dict / irecv_tensor_dict
 # ---------------------------------------------------------------------------
 
 
-def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
-    device = init_dist(local_rank, world_size, master_port)
-    initialize_model_parallel(pipeline_parallel_size=world_size)
+def isend_irecv_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    device_kind: DeviceKind,
+    result_queue,
+):
+    device = init_dist(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind))
     pp_group = get_pp_group()
 
     if pp_group.is_first_rank:
-        torch.manual_seed(77)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(77)
+        _seed_device(77, device)
         tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
         handles = pp_group.isend_tensor_dict({"t": tensor})
         for h in handles:
@@ -449,27 +514,42 @@ def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, resul
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    destroy_distributed_env()
+    _cleanup_distributed()
 
 
-@pytest.mark.gpu
-@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
-@pytest.mark.parametrize("pp_size", [2])
-def test_isend_irecv_tensor_dict(pp_size: int):
-    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None:
     mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     q = manager.Queue()
-
-    port = _find_free_port()
-    torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
-
+    torch.multiprocessing.spawn(
+        isend_irecv_worker,
+        args=(pp_size, master_port, device_kind, q),
+        nprocs=pp_size,
+    )
     results = {label: tensor for label, tensor in [q.get(), q.get()]}
     torch.testing.assert_close(
         results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
     )
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("pp_size", [2])
+def test_isend_irecv_tensor_dict(pp_size: int):
+    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+    _run_isend_irecv(pp_size, device_kind="cpu", master_port=_find_free_port())
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize("pp_size", [pytest.param(2, marks=_L4_TWO_GPU)])
+def test_isend_irecv_tensor_dict_parity(pp_size: int):
+    """Nightly: isend/irecv on real multi-GPU NCCL."""
+    _run_isend_irecv(pp_size, device_kind="cuda", master_port=_find_free_port())
+
+
 # ---------------------------------------------------------------------------
 # 4.  predict_noise_maybe_with_cfg
 # ---------------------------------------------------------------------------
@@ -477,12 +557,13 @@ def test_isend_irecv_tensor_dict(pp_size: int):
 _baseline_cache: dict[tuple, torch.Tensor] = {}
 
 
-def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
-    """Compute expected single-GPU output using the same MockPipelineParallel.
-
-    Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
-    Results are cached so identical configs are only computed once.
-    """
+def compute_single_gpu_baseline(
+    test_config: dict,
+    dtype: torch.dtype,
+    do_true_cfg: bool,
+    device_kind: DeviceKind,
+) -> torch.Tensor:
+    """Compute expected single-process output using the same MockPipelineParallel."""
     key = (
         test_config["num_layers"],
         test_config["dim"],
@@ -492,27 +573,29 @@ def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_c
         test_config["cfg_scale"],
         dtype,
         do_true_cfg,
+        device_kind,
     )
     if key in _baseline_cache:
         return _baseline_cache[key]
 
-    device = init_dist(0, 1, _find_free_port())
-    initialize_model_parallel(pipeline_parallel_size=1)
+    device = init_dist(0, 1, _find_free_port(), device_kind)
+    try:
+        initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
 
-    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
-        test_config, dtype, device, do_true_cfg=do_true_cfg
-    )
-
-    with torch.inference_mode():
-        noise_pred = pipeline.predict_noise_maybe_with_cfg(
-            do_true_cfg=do_true_cfg,
-            true_cfg_scale=test_config["cfg_scale"],
-            positive_kwargs=positive_kwargs,
-            negative_kwargs=negative_kwargs,
-            cfg_normalize=False,
+        pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+            test_config, dtype, device, do_true_cfg=do_true_cfg
         )
 
-    destroy_distributed_env()
+        with torch.inference_mode():
+            noise_pred = pipeline.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=test_config["cfg_scale"],
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+                cfg_normalize=False,
+            )
+    finally:
+        _cleanup_distributed()
 
     _baseline_cache[key] = noise_pred.cpu()
     return _baseline_cache[key]
@@ -527,11 +610,16 @@ def predict_noise_worker(
     do_true_cfg: bool,
     dtype: torch.dtype,
     test_config: dict,
+    device_kind: DeviceKind,
     result_queue,
 ):
     """Generic predict-noise worker parameterized by PP and CFG topology."""
-    device = init_dist(local_rank, world_size, master_port)
-    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+    device = init_dist(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(
+        pipeline_parallel_size=pp_size,
+        cfg_parallel_size=cfg_size,
+        backend=_mp_backend(device_kind),
+    )
 
     pp_group = get_pp_group()
     cfg_rank = get_classifier_free_guidance_rank()
@@ -548,8 +636,6 @@ def predict_noise_worker(
             negative_kwargs=negative_kwargs,
             cfg_normalize=False,
         )
-    # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
-    # Flush the non-last PP rank's async send before barrier / process teardown.
     pipeline._sync_pp_send()
 
     if pp_group.is_last_rank:
@@ -561,77 +647,22 @@ def predict_noise_worker(
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    destroy_distributed_env()
+    _cleanup_distributed()
 
 
-@pytest.mark.gpu
-@pytest.mark.parametrize(
-    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
-    [
-        pytest.param(
-            2,
-            1,
-            False,
-            torch.float32,
-            4,
-            100,
-            1e-5,
-            1e-5,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
-            id="pp2-no_cfg-float32",
-        ),
-        pytest.param(
-            2,
-            1,
-            False,
-            torch.bfloat16,
-            4,
-            100,
-            1e-2,
-            1e-2,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
-            id="pp2-no_cfg-bfloat16",
-        ),
-        pytest.param(
-            2,
-            1,
-            True,
-            torch.bfloat16,
-            4,
-            100,
-            1e-2,
-            1e-2,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
-            id="pp2-seq_cfg-bfloat16",
-        ),
-        pytest.param(
-            2,
-            2,
-            True,
-            torch.bfloat16,
-            4,
-            100,
-            1e-2,
-            1e-2,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
-            id="pp2-cfg2-bfloat16",
-        ),
-        pytest.param(
-            3,
-            1,
-            False,
-            torch.bfloat16,
-            6,
-            100,
-            1e-2,
-            1e-2,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
-            id="pp3-no_cfg-bfloat16",
-        ),
-    ],
-)
-def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
-    """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
+def _run_predict_noise(
+    *,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    dtype: torch.dtype,
+    num_layers: int,
+    input_seed: int,
+    rtol: float,
+    atol: float,
+    device_kind: DeviceKind,
+    master_port: str,
+) -> None:
     test_config = {
         "num_layers": num_layers,
         "dim": 64,
@@ -641,22 +672,20 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_
         "input_seed": input_seed,
     }
 
-    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
+    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
 
     mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     pp_q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    port = _find_free_port()
     torch.multiprocessing.spawn(
         predict_noise_worker,
-        args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
+        args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
         nprocs=world_size,
     )
 
     pp_out = pp_q.get()
-
     assert baseline_out.shape == pp_out.shape
     torch.testing.assert_close(
         pp_out,
@@ -667,35 +696,157 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_
     )
 
 
+_CPU_PREDICT_NOISE_CASES = [
+    pytest.param(2, 1, False, 4, 100, 1e-5, 1e-5, id="pp2-no_cfg"),
+    pytest.param(2, 1, True, 4, 100, 1e-5, 1e-5, id="pp2-seq_cfg"),
+    pytest.param(2, 2, True, 4, 100, 1e-5, 1e-5, id="pp2-cfg2"),
+    pytest.param(3, 1, False, 6, 100, 1e-5, 1e-5, id="pp3-no_cfg"),
+]
+
+_GPU_PREDICT_NOISE_CASES = [
+    pytest.param(
+        2,
+        1,
+        False,
+        torch.float32,
+        4,
+        100,
+        1e-5,
+        1e-5,
+        marks=_L4_TWO_GPU,
+        id="pp2-no_cfg-float32",
+    ),
+    pytest.param(
+        2,
+        1,
+        False,
+        torch.bfloat16,
+        4,
+        100,
+        1e-2,
+        1e-2,
+        marks=_L4_TWO_GPU,
+        id="pp2-no_cfg-bfloat16",
+    ),
+    pytest.param(
+        2,
+        1,
+        True,
+        torch.bfloat16,
+        4,
+        100,
+        1e-2,
+        1e-2,
+        marks=_L4_TWO_GPU,
+        id="pp2-seq_cfg-bfloat16",
+    ),
+    pytest.param(
+        2,
+        2,
+        True,
+        torch.bfloat16,
+        4,
+        100,
+        1e-2,
+        1e-2,
+        marks=_L4_FOUR_GPU,
+        id="pp2-cfg2-bfloat16",
+    ),
+    pytest.param(
+        3,
+        1,
+        False,
+        torch.bfloat16,
+        6,
+        100,
+        1e-2,
+        1e-2,
+        marks=_L4_THREE_GPU,
+        id="pp3-no_cfg-bfloat16",
+    ),
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol",
+    _CPU_PREDICT_NOISE_CASES,
+)
+def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol):
+    """predict_noise_maybe_with_cfg output matches baseline across PP / CFG topologies."""
+    _run_predict_noise(
+        pp_size=pp_size,
+        cfg_size=cfg_size,
+        do_true_cfg=do_true_cfg,
+        dtype=torch.float32,
+        num_layers=num_layers,
+        input_seed=input_seed,
+        rtol=rtol,
+        atol=atol,
+        device_kind="cpu",
+        master_port=_find_free_port(),
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
+    _GPU_PREDICT_NOISE_CASES,
+)
+def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
+    """Nightly: predict_noise parity on real multi-GPU NCCL."""
+    _run_predict_noise(
+        pp_size=pp_size,
+        cfg_size=cfg_size,
+        do_true_cfg=do_true_cfg,
+        dtype=dtype,
+        num_layers=num_layers,
+        input_seed=input_seed,
+        rtol=rtol,
+        atol=atol,
+        device_kind="cuda",
+        master_port=_find_free_port(),
+    )
+
+
 # ---------------------------------------------------------------------------
 # 5.  scheduler_step_maybe_with_cfg
 # ---------------------------------------------------------------------------
 
 
-def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
-    """Single-GPU reference: predict_noise + scheduler_step."""
-    device = init_dist(0, 1, _find_free_port())
-    initialize_model_parallel(pipeline_parallel_size=1)
+def compute_scheduler_step_baseline(
+    test_config: dict,
+    do_true_cfg: bool,
+    device_kind: DeviceKind,
+) -> torch.Tensor:
+    """Single-process reference: predict_noise + scheduler_step."""
+    device = init_dist(0, 1, _find_free_port(), device_kind)
+    try:
+        initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
 
-    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
-        test_config, torch.float32, device, do_true_cfg=do_true_cfg
-    )
-    latents = positive_kwargs["x"]
-    t = torch.tensor(500, device=device)
-
-    with torch.inference_mode():
-        noise_pred = pipeline.predict_noise_maybe_with_cfg(
-            do_true_cfg=do_true_cfg,
-            true_cfg_scale=test_config["cfg_scale"],
-            positive_kwargs=positive_kwargs,
-            negative_kwargs=negative_kwargs,
-            cfg_normalize=False,
-        )
-        result = pipeline.scheduler_step_maybe_with_cfg(
-            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+            test_config, torch.float32, device, do_true_cfg=do_true_cfg
         )
+        latents = positive_kwargs["x"]
+        t = torch.tensor(500, device=device)
 
-    destroy_distributed_env()
+        with torch.inference_mode():
+            noise_pred = pipeline.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=test_config["cfg_scale"],
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+                cfg_normalize=False,
+            )
+            result = pipeline.scheduler_step_maybe_with_cfg(
+                noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+            )
+    finally:
+        _cleanup_distributed()
     return result.cpu()
 
 
@@ -707,10 +858,15 @@ def scheduler_step_worker(
     cfg_size: int,
     do_true_cfg: bool,
     test_config: dict,
+    device_kind: DeviceKind,
     result_queue,
 ):
-    device = init_dist(local_rank, world_size, master_port)
-    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+    device = init_dist(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(
+        pipeline_parallel_size=pp_size,
+        cfg_parallel_size=cfg_size,
+        backend=_mp_backend(device_kind),
+    )
 
     pp_group = get_pp_group()
     cfg_rank = get_classifier_free_guidance_rank()
@@ -732,43 +888,25 @@ def scheduler_step_worker(
         latents = pipeline.scheduler_step_maybe_with_cfg(
             noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
         )
-    # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
-    # Flush the last PP rank's async latent send before barrier / process teardown.
     pipeline._sync_pp_send()
 
     if pp_group.is_first_rank and cfg_rank == 0:
-        resolved = latents.contiguous()
-        result_queue.put(resolved.cpu())
+        result_queue.put(latents.contiguous().cpu())
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    destroy_distributed_env()
+    _cleanup_distributed()
 
 
-@pytest.mark.gpu
-@pytest.mark.parametrize(
-    "pp_size, cfg_size, do_true_cfg, input_seed",
-    [
-        pytest.param(
-            2,
-            1,
-            False,
-            300,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
-            id="pp2-no_cfg",
-        ),
-        pytest.param(
-            2,
-            2,
-            True,
-            600,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
-            id="pp2-cfg2-true_cfg",
-        ),
-    ],
-)
-def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
-    """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
+def _run_scheduler_step(
+    *,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    input_seed: int,
+    device_kind: DeviceKind,
+    master_port: str,
+) -> None:
     test_config = {
         "num_layers": 4,
         "dim": 64,
@@ -778,17 +916,16 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
         "input_seed": input_seed,
     }
 
-    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
+    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
 
     mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     q = manager.Queue()
 
-    port = _find_free_port()
     world_size = pp_size * cfg_size
     torch.multiprocessing.spawn(
         scheduler_step_worker,
-        args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
+        args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
         nprocs=world_size,
     )
 
@@ -800,3 +937,47 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
         atol=0,
         msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
     )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, input_seed",
+    [
+        pytest.param(2, 1, False, 300, id="pp2-no_cfg"),
+        pytest.param(2, 2, True, 600, id="pp2-cfg2-true_cfg"),
+    ],
+)
+def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
+    """Rank 0 latents after scheduler_step match baseline across PP / CFG topologies."""
+    _run_scheduler_step(
+        pp_size=pp_size,
+        cfg_size=cfg_size,
+        do_true_cfg=do_true_cfg,
+        input_seed=input_seed,
+        device_kind="cpu",
+        master_port=_find_free_port(),
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, input_seed",
+    [
+        pytest.param(2, 1, False, 300, marks=_L4_TWO_GPU, id="pp2-no_cfg"),
+        pytest.param(2, 2, True, 600, marks=_L4_FOUR_GPU, id="pp2-cfg2-true_cfg"),
+    ],
+)
+def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed):
+    """Nightly: scheduler_step parity on real multi-GPU NCCL."""
+    _run_scheduler_step(
+        pp_size=pp_size,
+        cfg_size=cfg_size,
+        do_true_cfg=do_true_cfg,
+        input_seed=input_seed,
+        device_kind="cuda",
+        master_port=_find_free_port(),
+    )

tests/diffusion/distributed/test_tensor_parallel.py

@@ -1,16 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Z-Image end-to-end tests for diffusion parallelism.
+"""DiT tensor-parallel parity and perf checks for diffusion models.
 
-This file currently covers:
-- DiT tensor parallelism (TP=2) vs TP=1.
-- VAE patch parallelism (vae_patch_parallel_size=2) vs baseline on TP=2.
-
-Note: CUDA-only (>=2 GPUs). We use `enforce_eager=False` (default) to enable
-`torch.compile`.
+Compares TP=2 against TP=1 on output correctness, latency, and peak memory.
+Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
 """
 
+from __future__ import annotations
+
 import os
 import time
 from pathlib import Path
@@ -31,22 +29,15 @@ from vllm_omni.platforms import current_omni_platform
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
-
-
-def _get_zimage_model() -> str:
-    # Allow overriding the model for local/offline environments.
-    # Can be either a HuggingFace repo id or a local path.
-    return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")
+DEFAULT_MODEL = "Tongyi-MAI/Z-Image-Turbo"
 
 
 def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
-    """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
     arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
     return torch.from_numpy(arr)
 
 
 def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
-    """Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1]."""
     ta = _pil_to_float_rgb_tensor(a)
     tb = _pil_to_float_rgb_tensor(b)
     assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
@@ -71,7 +62,7 @@ def _extract_single_image(outputs) -> Image.Image:
     return images[0]
 
 
-def _run_zimage_generate(
+def _run_generate(
     *,
     tp_size: int,
     height: int,
@@ -79,8 +70,6 @@ def _run_zimage_generate(
     num_inference_steps: int,
     seed: int,
     enforce_eager: bool,
-    vae_use_tiling: bool = False,
-    vae_patch_parallel_size: int = 1,
     num_requests: int = 4,
 ) -> tuple[Image.Image, float, float]:
     if num_requests < 2:
@@ -91,25 +80,11 @@ def _run_zimage_generate(
     monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
     try:
-        # Each run needs a distinct DiffusionParallelConfig; use OmniRunner per call (not the
-        # parametrized omni_runner fixture, which is fixed per module).
         with OmniRunner(
-            _get_zimage_model(),
-            parallel_config=DiffusionParallelConfig(
-                tensor_parallel_size=tp_size,
-                vae_patch_parallel_size=vae_patch_parallel_size,
-            ),
+            DEFAULT_MODEL,
+            parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
             enforce_eager=enforce_eager,
-            vae_use_tiling=vae_use_tiling,
         ) as runner:
-            # NOTE: Omni closes itself when a generate() call is exhausted.
-            # To avoid measuring teardown time (process shutdown, memory cleanup),
-            # we measure the latency to produce *subsequent* outputs within a single
-            # generator run.
-            #
-            # This also serves as a warmup: the first output may include extra
-            # compilation/caching overhead, while later outputs are closer to
-            # steady-state inference.
             gen = runner.omni.generate(
                 [PROMPT] * num_requests,
                 OmniDiffusionSamplingParams(
@@ -134,14 +109,11 @@ def _run_zimage_generate(
                 per_request_times_s.append(t_now - t_prev)
                 t_prev = t_now
 
-            # Ensure the generator is fully consumed so it can clean up.
             for _ in gen:
                 pass
 
             median_time_s = float(np.median(per_request_times_s))
-
             peak_memory_mb = monitor.peak_used_mb
-
             return _extract_single_image([last_output]), median_time_s, peak_memory_mb
     finally:
         monitor.stop()
@@ -151,20 +123,20 @@ def _run_zimage_generate(
 @pytest.mark.diffusion
 @pytest.mark.parallel
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
-def test_zimage_tensor_parallel_tp2(tmp_path: Path):
+def test_tensor_parallel_tp2(tmp_path: Path):
     if current_omni_platform.is_npu():
-        pytest.skip("Z-Image TP e2e test is only supported on CUDA and ROCm for now.")
+        pytest.skip("Tensor parallel parity test is only supported on CUDA and ROCm for now.")
     if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
-        pytest.skip("Z-Image TP=2 requires >= 2 devices.")
+        pytest.skip("Tensor parallel TP=2 requires >= 2 devices.")
 
+    model_name = DEFAULT_MODEL
     enforce_eager = False
-
     height = 512
     width = 512
     num_inference_steps = 2
     seed = 42
 
-    tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
+    tp1_img, tp1_time_s, tp1_peak_mem = _run_generate(
         tp_size=1,
         height=height,
         width=width,
@@ -172,7 +144,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
         seed=seed,
         enforce_eager=enforce_eager,
     )
-    tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
+    tp2_img, tp2_time_s, tp2_peak_mem = _run_generate(
         tp_size=2,
         height=height,
         width=width,
@@ -181,8 +153,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
         enforce_eager=enforce_eager,
     )
 
-    tp1_path = tmp_path / "zimage_tp1.png"
-    tp2_path = tmp_path / "zimage_tp2.png"
+    tp1_path = tmp_path / "tp1.png"
+    tp2_path = tmp_path / "tp2.png"
     tp1_img.save(tp1_path)
     tp2_img.save(tp2_path)
 
@@ -193,7 +165,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
     mean_threshold = 3e-2
     p99_threshold = 2.5e-1
     print(
-        "Z-Image TP image diff stats (TP=1 vs TP=2): "
+        f"{model_name} TP image diff (TP=1 vs TP=2): "
         f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
         f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
         f"tp1_img={tp1_path}, tp2_img={tp2_path}"
@@ -203,73 +175,11 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
         f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
     )
 
-    print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
-    # ROCm is not optimized TP2 can be slower than TP1
+    print(f"TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
     if not current_omni_platform.is_rocm():
         assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
 
-    print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
+    print(f"TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
     assert tp2_peak_mem < tp1_peak_mem, (
         f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
     )
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
-def test_zimage_vae_patch_parallel_tp2(tmp_path: Path):
-    if current_omni_platform.is_npu():
-        pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA and ROCm for now.")
-    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
-        pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.")
-
-    enforce_eager = False
-
-    # Use a larger image to ensure there are multiple VAE tiles.
-    height = 1152
-    width = 1152
-    num_inference_steps = 2
-    seed = 42
-
-    baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate(
-        tp_size=2,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        seed=seed,
-        enforce_eager=enforce_eager,
-        vae_use_tiling=True,
-        vae_patch_parallel_size=1,
-        num_requests=2,
-    )
-    pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate(
-        tp_size=2,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        seed=seed,
-        enforce_eager=enforce_eager,
-        vae_use_tiling=True,
-        vae_patch_parallel_size=2,
-        num_requests=2,
-    )
-
-    baseline_path = tmp_path / "zimage_tp2_vae_pp1.png"
-    pp2_path = tmp_path / "zimage_tp2_vae_pp2.png"
-    baseline_img.save(baseline_path)
-    pp2_img.save(pp2_path)
-
-    mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
-    mean_threshold = 5e-3
-    p99_threshold = 1e-1
-    print(
-        "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
-        f"pp1_img={baseline_path}, pp2_img={pp2_path}"
-    )
-    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
-    )

tests/diffusion/distributed/test_vae_decode_parallelism.py

@@ -1,151 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""VAE patch-parallel decode parity tests for diffusion models.
+
+Compares ``vae_patch_parallel_size=2`` against ``vae_patch_parallel_size=1`` on
+TP=2 with VAE tiling enabled. Covers one video model (Wan T2V) and one image
+model (Qwen-Image). CUDA/ROCm-only (>=2 devices).
+"""
+
+from __future__ import annotations
+
 import os
-import sys
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import pytest
 import torch
 from PIL import Image
-from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-import time
-
-from tests.helpers.runtime import OmniRunner
 from vllm_omni.platforms import current_omni_platform
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-all_model_configs = [
-    {
-        "model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
-        "out_width": 1280,
-        "out_height": 720,
-        "out_frames": 16,
-    },
-    {
-        "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-        "out_width": 1280,
-        "out_height": 720,
-        "out_frames": 16,
-    },
-    {
-        "model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-        "out_width": 1280,
-        "out_height": 720,
-        "out_frames": 16,
-    },
-    {
-        "model_name": "Qwen/Qwen-Image",
-        "out_width": 3072,
-        "out_height": 3072,
-        "out_frames": 1,
-    },
-    # {
-    #     "model_name": "stepfun-ai/NextStep-1.1",
-    #     "out_width": 3072,
-    #     "out_height": 3072,
-    #     "out_frames": 1,
-    # },
+PROMPT = "A cat sitting on a table"
+
+MODEL_CASES = [
+    pytest.param(
+        {
+            "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            "height": 512,
+            "width": 512,
+            "num_frames": 5,
+            "needs_image": False,
+            "is_moe": False,
+            "mean_threshold": 3e-2,
+            "p99_threshold": 3e-2,
+        },
+        id="wan22_t2v_a14b",
+    ),
+    pytest.param(
+        {
+            "model_name": "Qwen/Qwen-Image",
+            "height": 1152,
+            "width": 1152,
+            "num_frames": 1,
+            "needs_image": False,
+            "is_moe": False,
+            "mean_threshold": 5e-3,
+            "p99_threshold": 1e-1,
+        },
+        id="qwen_image",
+    ),
 ]
 
 
-def is_nextstep_model(model_name: str) -> bool:
-    """Check if the model is a NextStep model by reading its config."""
-    from vllm.transformers_utils.config import get_hf_file_to_dict
-
-    try:
-        cfg = get_hf_file_to_dict("config.json", model_name)
-        if cfg and cfg.get("model_type") == "nextstep":
-            return True
-    except Exception:
-        pass
-    return False
+def _to_float_array(data: Any) -> np.ndarray:
+    if isinstance(data, Image.Image):
+        return np.asarray(data.convert("RGB"), dtype=np.float32) / 255.0
+    if isinstance(data, torch.Tensor):
+        arr = data.detach().cpu().numpy()
+    else:
+        arr = np.asarray(data)
+    if arr.dtype == np.uint8:
+        return arr.astype(np.float32) / 255.0
+    return arr.astype(np.float32)
 
 
-def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
+def _extract_output_array(output: Any) -> np.ndarray:
+    req_out = output.request_output
+    payload = req_out.images[0]
+    return _to_float_array(payload)
+
+
+def _diff_metrics(a: np.ndarray, b: np.ndarray) -> tuple[float, float]:
+    assert a.shape == b.shape, f"Output shapes differ: {a.shape} vs {b.shape}"
+    abs_diff = np.abs(a - b)
+    return float(abs_diff.mean()), float(np.quantile(abs_diff, 0.99))
+
+
+def _save_preview(arr: np.ndarray, path: Path) -> None:
+    """Save a single RGB frame preview; video tensors use the first frame."""
+    preview = np.asarray(arr)
+    # (B, T, H, W, C) video
+    if preview.ndim == 5:
+        preview = preview[0, 0]
+    # (T, H, W, C) video or (B, H, W, C) image batch
+    elif preview.ndim == 4:
+        preview = preview[0]
+    if preview.ndim != 3 or preview.shape[-1] not in (3, 4):
+        raise ValueError(f"Expected HWC RGB(A) preview, got shape {arr.shape} -> {preview.shape}")
+    preview = np.clip(preview[..., :3], 0.0, 1.0)
+    Image.fromarray((preview * 255.0).astype(np.uint8)).save(path)
+
+
+def _build_sampling_params(model_case: dict[str, Any], seed: int) -> OmniDiffusionSamplingParams:
+    kwargs: dict[str, Any] = {
+        "height": model_case["height"],
+        "width": model_case["width"],
+        "num_frames": model_case["num_frames"],
+        "num_inference_steps": 2,
+        "seed": seed,
+    }
+    if model_case["model_name"].startswith("Wan-AI/"):
+        kwargs["guidance_scale"] = 4.0
+    if model_case["is_moe"]:
+        kwargs["guidance_scale_2"] = 1.0
+        kwargs["boundary_ratio"] = 0.5
+    return OmniDiffusionSamplingParams(**kwargs)
+
+
+def _run_generate(
+    *,
+    model_case: dict[str, Any],
+    vae_patch_parallel_size: int,
+    seed: int,
+) -> np.ndarray:
+    current_omni_platform.empty_cache()
     parallel_config = DiffusionParallelConfig(
-        tensor_parallel_size=tp,
+        tensor_parallel_size=2,
         vae_patch_parallel_size=vae_patch_parallel_size,
     )
-
-    omni_kwargs = {
-        "vae_use_tiling": using_tile,
-        "parallel_config": parallel_config,
-    }
-    use_nextstep = is_nextstep_model(model_configs["model_name"])
-    if use_nextstep:
-        # NextStep-1.1 requires explicit pipeline class
-        omni_kwargs["model_class_name"] = "NextStep11Pipeline"
-
-    with OmniRunner(model_configs["model_name"], **omni_kwargs) as runner:
-        m = runner.omni
-        image = Image.new("RGB", (out_width, out_height), (0, 0, 0))
-        start = time.perf_counter()
-        outputs = m.generate(
-            {
-                "prompt": "A cat sitting on a table",
-                "multi_modal_data": {"image": image},
-            },
-            sampling_params_list=OmniDiffusionSamplingParams(
-                height=model_configs["out_height"],
-                width=model_configs["out_width"],
-                num_frames=model_configs["out_frames"],
-                num_inference_steps=2,
-                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
-            ),
+    with OmniRunner(
+        model_case["model_name"],
+        parallel_config=parallel_config,
+        vae_use_tiling=True,
+        enforce_eager=False,
+    ) as runner:
+        request: dict[str, Any] = {"prompt": model_case.get("prompt", PROMPT)}
+        if model_case["needs_image"]:
+            request["multi_modal_data"] = {
+                "image": Image.new("RGB", (model_case["width"], model_case["height"]), (0, 0, 0))
+            }
+        outputs = runner.omni.generate(
+            request,
+            sampling_params_list=_build_sampling_params(model_case, seed),
         )
-        end = time.perf_counter()
-        first_output = outputs[0]
-        req_out = first_output.request_output
-        frames = req_out.images[0]
-        if isinstance(frames, torch.Tensor):
-            frames = frames.detach().cpu().numpy()
-        elif isinstance(frames, Image.Image):
-            frames = np.array(frames)
-        # frames shape: (batch, num_frames, height, width, channels)
-        cost = (end - start) * 1000
-        return frames, cost
-        cleanup_dist_env_and_memory()
+        return _extract_output_array(outputs[0])
 
 
-@pytest.mark.parametrize("model_configs", all_model_configs)
-def test_vae_parallel_model(model_configs: str):
-    parallel_result, parallel_time = model_run(
-        model_configs=model_configs,
-        tp=2,
-        out_width=1280,
-        out_height=704,
-        out_frames=5,
-        using_tile=True,
-        vae_patch_parallel_size=2,
-    )
-    non_parallel_result, non_parallel_time = model_run(
-        model_configs=model_configs,
-        tp=2,
-        out_width=1280,
-        out_height=704,
-        out_frames=5,
-        using_tile=True,
-        vae_patch_parallel_size=1,
-    )
-    result_diff = np.abs(non_parallel_result - parallel_result)
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+@pytest.mark.parametrize("model_case", MODEL_CASES)
+def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
+    if current_omni_platform.is_npu():
+        pytest.skip("VAE patch parallel test is only supported on CUDA and ROCm for now.")
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("VAE patch parallel TP=2 requires >= 2 devices.")
 
-    mean_threshold = 3e-2
-    max_threshold = 3e-2  # they should be totally same
+    seed = 42
+    model_id = model_case["model_name"].split("/")[-1]
+
+    baseline = _run_generate(model_case=model_case, vae_patch_parallel_size=1, seed=seed)
+    parallel = _run_generate(model_case=model_case, vae_patch_parallel_size=2, seed=seed)
+
+    baseline_path = tmp_path / f"{model_id}_tp2_vae_pp1.png"
+    parallel_path = tmp_path / f"{model_id}_tp2_vae_pp2.png"
+    _save_preview(baseline, baseline_path)
+    _save_preview(parallel, parallel_path)
+
+    mean_abs_diff, p99_abs_diff = _diff_metrics(baseline, parallel)
+    mean_threshold = model_case["mean_threshold"]
+    p99_threshold = model_case["p99_threshold"]
     print(
-        f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): "
-        f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
-        f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms"
+        f"{model_case['model_name']} VAE patch parallel diff (TP=2, pp=1 vs pp=2): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
+        f"pp1_preview={baseline_path}, pp2_preview={parallel_path}"
+    )
+    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+        f"Output diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, "
+        f"p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
     )
-    # assert non_parallel_time > parallel_time
-    assert result_diff.mean() < mean_threshold
-    assert result_diff.max() < max_threshold

tests/diffusion/hooks/test_hook_registry.py

@@ -12,6 +12,9 @@ from torch import nn
 
 from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 DEFAULT_OUT = "ECHO"
 OVERRIDE_OUT = "OVERRIDE"
 INPUT_KWARG = "inp"

tests/diffusion/layers/test_rotary_emb_equivalence.py

@@ -13,6 +13,8 @@ from __future__ import annotations
 import pytest
 import torch
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def _apply_rotary_emb_helios_original(
     hidden_states: torch.Tensor,

tests/diffusion/models/dreamzero/test_crossattn_cache.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests that DreamZero's cross-attention KV cache is actually init'd and reused.
+
+DreamZero's I2V cross-attention caches the text k/v and image k_img/v_img in a
+per-layer ``crossattn_cache`` dict: they are computed on the first forward of a
+session (``is_init`` False -> True) and reused on every later step, because the
+text/image context is session-invariant. Only the query (which depends on the
+per-step hidden state) is recomputed.
+
+These tests lock in that ``WanI2VCrossAttention.forward`` initialises the cache
+on the first call and reuses it (not recompute) on later calls, and recomputes
+every call when no cache is provided.
+
+CPU-only: ``WanI2VCrossAttention.__init__`` needs a tensor-parallel group (it
+builds ColumnParallel/RowParallel linears), so we bypass ``__init__`` and stub
+the projections/attention to count calls and exercise only the caching logic.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.causal_wan_model import WanI2VCrossAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_CLIP_LEN = 257  # image tokens the I2V cross-attn splits off the front of context
+
+
+def _make_cross_attn(n_heads: int = 2, head_dim: int = 4) -> WanI2VCrossAttention:
+    """A WanI2VCrossAttention whose submodules are call-counting stubs.
+
+    ``__init__`` is bypassed (it needs a TP group); we only wire what
+    ``forward`` touches. Each projection returns a [B, S, n*head_dim] tensor so
+    the subsequent ``unflatten(2, (n, head_dim))`` works, and is a MagicMock so
+    we can assert how many times it was invoked.
+    """
+    inner = n_heads * head_dim
+    module = WanI2VCrossAttention.__new__(WanI2VCrossAttention)
+    module.tp_num_heads = n_heads
+    module.head_dim = head_dim
+
+    def _proj() -> MagicMock:
+        return MagicMock(side_effect=lambda t: torch.zeros(t.shape[0], t.shape[1], inner))
+
+    module.q = _proj()
+    module.k = _proj()
+    module.v = _proj()
+    module.k_img = _proj()
+    module.v_img = _proj()
+    module.norm_q = lambda t: t
+    module.norm_k = lambda t: t
+    module.norm_k_img = lambda t: t
+    # attn returns the query unchanged ([B, S, n, d]); flatten(2) downstream is happy.
+    module.attn = lambda q, k, v: q
+    module.o = lambda t: t
+    return module
+
+
+def _fresh_cache() -> dict:
+    return {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None}
+
+
+def _inputs(inner: int = 8):
+    x = torch.randn(1, 3, inner)  # B=1, query_len=3
+    context = torch.randn(1, _CLIP_LEN + 5, inner)  # 257 image tokens + 5 text tokens
+    return x, context
+
+
+def test_cache_inits_on_first_call_and_reuses_on_second() -> None:
+    module = _make_cross_attn()
+    cache = _fresh_cache()
+    x, context = _inputs()
+
+    # First call: cache is initialised and every k/v projection runs once.
+    module.forward(x, context, crossattn_cache=cache)
+    assert cache["is_init"] is True
+    for key in ("k", "v", "k_img", "v_img"):
+        assert cache[key] is not None, f"{key} should be cached after first call"
+    assert module.k.call_count == 1
+    assert module.v.call_count == 1
+    assert module.k_img.call_count == 1
+    assert module.v_img.call_count == 1
+
+    cached = {key: cache[key] for key in ("k", "v", "k_img", "v_img")}
+
+    # Second call: k/v/k_img/v_img are read from cache, NOT recomputed.
+    module.forward(x, context, crossattn_cache=cache)
+    assert module.k.call_count == 1
+    assert module.v.call_count == 1
+    assert module.k_img.call_count == 1
+    assert module.v_img.call_count == 1
+    for key, tensor in cached.items():
+        assert cache[key] is tensor, f"{key} should be the same cached object"
+
+    # The query always depends on the per-step hidden state, so it is recomputed.
+    assert module.q.call_count == 2
+
+
+def test_no_cache_recomputes_every_call() -> None:
+    module = _make_cross_attn()
+    x, context = _inputs()
+
+    module.forward(x, context, crossattn_cache=None)
+    module.forward(x, context, crossattn_cache=None)
+
+    # Without a cache, the image/text k/v are recomputed on every call.
+    assert module.k_img.call_count == 2
+    assert module.v_img.call_count == 2
+    assert module.k.call_count == 2
+    assert module.v.call_count == 2

tests/diffusion/models/dreamzero/test_deploy_configs.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Every shipped DreamZero deploy config must select the AR-Diffusion engine.
+
+``DreamZeroPipeline.__init__`` rejects any other engine, so a config missing
+``engine_backend`` fails at worker startup rather than at request time. That
+only surfaces in slow multi-GPU e2e serving tests (see issue #5046), so assert
+it here on CPU instead.
+"""
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import yaml
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Discover siblings of a known config so a future dreamzero*.yaml is covered
+# without editing this test.
+DEPLOY_DIR = Path(get_deploy_config_path("dreamzero.yaml")).parent
+
+
+def _dreamzero_stages() -> list[tuple[str, int, dict]]:
+    stages = []
+    for name in sorted(path.name for path in DEPLOY_DIR.glob("dreamzero*.yaml")):
+        config = yaml.safe_load(Path(get_deploy_config_path(name)).read_text())
+        for index, stage in enumerate(config.get("stages") or []):
+            stages.append((name, index, stage))
+    return stages
+
+
+def test_dreamzero_deploy_configs_exist() -> None:
+    assert _dreamzero_stages(), f"no DreamZero deploy configs found under {DEPLOY_DIR}"
+
+
+@pytest.mark.parametrize(
+    ("config_name", "stage_index", "stage"),
+    _dreamzero_stages(),
+    ids=lambda value: value if isinstance(value, str) else None,
+)
+def test_dreamzero_stage_selects_ar_diffusion_engine(config_name: str, stage_index: int, stage: dict) -> None:
+    backend = stage.get("engine_backend")
+    assert backend, (
+        f"{config_name} stage {stage_index} does not set engine_backend; "
+        "DreamZeroPipeline requires the AR-Diffusion engine"
+    )
+    # Resolving also proves the import path is real, not just a lookalike string.
+    engine_class = DiffusionEngine.resolve_engine_class(SimpleNamespace(engine_backend=backend))
+    assert issubclass(engine_class, ARDiffusionEngine), (
+        f"{config_name} stage {stage_index} resolves to {engine_class.__name__}, not an AR-Diffusion engine"
+    )

tests/diffusion/models/dreamzero/test_dreamzero_utils.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.models.dreamzero.utils import (
+    DEFAULT_CFG_SCALE,
+    DEFAULT_EMBODIMENT_NAME_TO_ID,
+    DEFAULT_NEGATIVE_PROMPT,
+    DEFAULT_NUM_INFERENCE_STEPS,
+    DEFAULT_SEED,
+    DEFAULT_SIGMA_SHIFT,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dreamzero_default_constants_match_source_baseline():
+    assert DEFAULT_NUM_INFERENCE_STEPS == 16
+    assert DEFAULT_CFG_SCALE == 5.0
+    assert DEFAULT_SIGMA_SHIFT == 5.0
+    assert DEFAULT_SEED == 1140
+    assert "worst quality" in DEFAULT_NEGATIVE_PROMPT
+    assert DEFAULT_EMBODIMENT_NAME_TO_ID["oxe_droid"] == 17

tests/diffusion/models/dreamzero/test_fused_qk_rms_norm.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for DreamZero's fused q/k qk-norm TP all-reduce.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the TP world size and the collective are mocked, so a single process
+simulates one rank (all-reduce is identity). This exercises the cat/split and
+count-scaling logic that a real multi-rank run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.causal_wan_model import (
+    DistributedRMSNorm,
+    fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.dreamzero.causal_wan_model"
+
+
+def _make_norms(dim: int) -> tuple[DistributedRMSNorm, DistributedRMSNorm]:
+    """Two norms with distinct random weights so a swapped q/k slice is caught."""
+    torch.manual_seed(0)
+    norm_q = DistributedRMSNorm(dim)
+    norm_k = DistributedRMSNorm(dim)
+    norm_q.weight.data.normal_()
+    norm_k.weight.data.normal_()
+    return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+    """Single-rank all-reduce: the reduction over one rank is the identity."""
+    return tensor
+
+
+def test_fused_matches_separate_tp1() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_matches_separate_tp_gt1() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    # tp_size=2 exercises the reduce branch and the `global_count *= tp_size`
+    # path in both the per-norm forward and the fused helper.
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce),
+    ):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_single_all_reduce_over_packed_pair() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch(
+            f"{_MODULE}.tensor_model_parallel_all_reduce",
+            side_effect=_identity_all_reduce,
+        ) as mock_ar,
+    ):
+        fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    # Exactly one collective, over a tensor whose last dim packs the q+k pair.
+    assert mock_ar.call_count == 1
+    packed = mock_ar.call_args[0][0]
+    assert packed.shape[-1] == 2
+
+
+def test_fused_identity_fallback_when_qk_norm_disabled() -> None:
+    # qk_norm=False makes the norms nn.Identity; the helper must fall back to
+    # plain passthrough without touching the collective.
+    norm_q, norm_k = torch.nn.Identity(), torch.nn.Identity()
+    q = torch.randn(2, 5, 64)
+    k = torch.randn(2, 5, 64)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce") as mock_ar:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, k, rtol=0, atol=0)
+    mock_ar.assert_not_called()

tests/diffusion/models/dreamzero/test_pipeline_state.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import OrderedDict
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _empty_pipeline() -> DreamZeroPipeline:
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._states = OrderedDict()
+    return pipeline
+
+
+def test_dreamzero_pipeline_state_is_session_keyed() -> None:
+    pipeline = _empty_pipeline()
+
+    session_a = pipeline._get_or_create_state("session-a")
+    session_b = pipeline._get_or_create_state("session-b")
+    session_a.call_count = 7
+    session_b.call_count = 3
+
+    assert pipeline._get_or_create_state("session-a") is session_a
+    assert pipeline._get_or_create_state("session-b") is session_b
+    assert session_a.call_count == 7
+    assert session_b.call_count == 3
+
+
+def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
+    pipeline = _empty_pipeline()
+
+    session_a = pipeline._get_or_create_state("session-a")
+    session_b = pipeline._get_or_create_state("session-b")
+    pipeline.state = session_b
+
+    pipeline.close_ar_diffusion_session("session-a")
+    assert pipeline.state is session_b
+    pipeline.reset_ar_diffusion_session("session-b")
+
+    assert not pipeline._states
+    assert pipeline.state is None
+    assert pipeline._get_or_create_state("session-a") is not session_a
+    assert pipeline._get_or_create_state("session-b") is not session_b
+
+
+def test_close_session_clears_active_state_alias() -> None:
+    pipeline = _empty_pipeline()
+    session_a = pipeline._get_or_create_state("session-a")
+    pipeline.state = session_a
+
+    pipeline.close_ar_diffusion_session("session-a")
+
+    assert not pipeline._states
+    assert pipeline.state is None
+
+
+def test_dreamzero_warmup_provider_builds_session_scoped_requests() -> None:
+    provider = SimpleNamespace(
+        ar_diffusion_kv_cache_spec=lambda: SimpleNamespace(window_frames=5, frames_per_block=4),
+        od_config=SimpleNamespace(
+            ar_diffusion_kv_config=None,
+            model_config={"policy_server_config": {"image_resolution": [8, 16]}},
+        ),
+        _ar_warmup_robot_obs=DreamZeroPipeline._ar_warmup_robot_obs,
+    )
+
+    requests = list(DreamZeroPipeline.ar_diffusion_warmup_requests(provider, "warmup-session"))
+
+    assert [request.prompt for request in requests] == ["warmup", "warmup"]
+    assert all(request.sampling_params.extra_args["session_id"] == "warmup-session" for request in requests)
+    assert requests[0].sampling_params.extra_args["robot_obs"]["observation/exterior_image_0_left"].shape == (
+        8,
+        16,
+        3,
+    )
+    assert requests[1].sampling_params.extra_args["robot_obs"]["observation/exterior_image_0_left"].shape == (
+        4,
+        8,
+        16,
+        3,
+    )
+
+
+def test_dreamzero_state_owns_no_kv_caches() -> None:
+    """KV (self- and cross-attention) is engine-owned since the AR-Diffusion
+    paged backend: the model-local cache accessors must be gone so nothing can
+    silently bypass the engine pool."""
+    state = DreamZeroState()
+
+    for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"):
+        assert not hasattr(state, removed)

tests/diffusion/models/dreamzero/test_qkv_fusion.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence test for DreamZero self-attn QKV fusion.
+
+CausalWanSelfAttention now projects q/k/v with a single QKVParallelLinear and
+splits the result, instead of three separate ColumnParallelLinear. The DreamZero
+checkpoint still ships separate q/k/v weights, so load_weights routes each into
+the packed `qkv` param with its shard id. This test proves the chain
+"load separate q/k/v via shard id -> single GEMM -> split" reproduces the three
+independent projections, given identical weights.
+
+CPU-only: a single-rank (world_size=1) tensor-parallel group is initialized and
+the GEMM dispatcher is forced to the default (CPU-compatible) path, following
+tests/diffusion/layers/test_adalayernorm.py.
+"""
+
+import os
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Minimal world_size=1 TP group required to build QKVParallelLinear."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29507")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for CPU test tensors."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+def test_fused_qkv_load_and_split_matches_separate_projections():
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    torch.manual_seed(0)
+    dim, head_dim, num_heads = 16, 4, 4  # inner_dim == dim, no GQA
+
+    qkv = QKVParallelLinear(
+        hidden_size=dim,
+        head_size=head_dim,
+        total_num_heads=num_heads,
+        bias=True,
+    )
+
+    # Distinct random weight + bias per projection, loaded the exact way
+    # DreamZero.load_weights does: param.weight_loader(param, tensor, shard_id).
+    refs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+    for shard_id in ("q", "k", "v"):
+        w = torch.randn(dim, dim)
+        b = torch.randn(dim)
+        refs[shard_id] = (w, b)
+        qkv.weight.weight_loader(qkv.weight, w, shard_id)
+        qkv.bias.weight_loader(qkv.bias, b, shard_id)
+
+    x = torch.randn(2, 3, dim)
+    fused, _ = qkv(x)
+    q, k, v = fused.split([dim, dim, dim], dim=-1)
+
+    for got, shard_id in zip((q, k, v), ("q", "k", "v"), strict=True):
+        w, b = refs[shard_id]
+        ref = torch.nn.functional.linear(x, w, b)
+        torch.testing.assert_close(got, ref)
+
+
+def test_fused_qkv_shard_ids_are_not_transposed():
+    # Guard against q/k/v being routed to the wrong slice: give each projection a
+    # constant, distinct weight so a swapped shard id would surface immediately.
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    dim, head_dim, num_heads = 16, 4, 4
+    qkv = QKVParallelLinear(
+        hidden_size=dim,
+        head_size=head_dim,
+        total_num_heads=num_heads,
+        bias=True,
+    )
+
+    scale = {"q": 1.0, "k": 2.0, "v": 3.0}
+    for shard_id, s in scale.items():
+        qkv.weight.weight_loader(qkv.weight, torch.full((dim, dim), s), shard_id)
+        qkv.bias.weight_loader(qkv.bias, torch.zeros(dim), shard_id)
+
+    x = torch.ones(1, dim)
+    fused, _ = qkv(x)
+    q, k, v = fused.split([dim, dim, dim], dim=-1)
+
+    # F.linear(ones, full(s)) == s * dim per element.
+    torch.testing.assert_close(q, torch.full((1, dim), scale["q"] * dim))
+    torch.testing.assert_close(k, torch.full((1, dim), scale["k"] * dim))
+    torch.testing.assert_close(v, torch.full((1, dim), scale["v"] * dim))

tests/diffusion/models/dreamzero/test_session_state_accounting.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-device byte accounting and the bounded VAE encoder ring (RFC #4480).
+
+Two things are pinned here. First, that a session reports the bytes it actually
+holds: both buckets walked, storages rather than logical element counts, and
+aliases counted once. Second, that bounding the adapter's VAE encoder history
+changes no numbers -- ``quant_conv`` is pointwise in time, so the frames the
+reader keeps do not depend on the ones the ring drops.
+
+All CPU, tiny tensors, no model.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import (
+    DreamZeroStateAdapter,
+)
+from vllm_omni.experimental.world_models.session_state import (
+    LatentBuffer,
+    SessionStateManager,
+    device_bytes,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _bytes(value: object) -> dict[str, int]:
+    return device_bytes(value, set())
+
+
+# -- device_bytes ------------------------------------------------------------
+
+
+def test_counts_tensor_storage_under_its_device() -> None:
+    tensor = torch.zeros(64, dtype=torch.float32)
+    assert _bytes(tensor) == {"cpu": 256}
+
+
+def test_views_on_one_storage_counted_once() -> None:
+    tensor = torch.zeros(64, dtype=torch.float32)
+    seen: set[tuple[str, int]] = set()
+    first = device_bytes(tensor, seen)
+    second = device_bytes(tensor[:8], seen)
+    assert first == {"cpu": 256}
+    assert second == {}
+
+
+def test_slice_reports_the_storage_it_pins() -> None:
+    # The whole point of counting storages: this slice keeps all 256 bytes
+    # alive, and ``numel() * element_size()`` would claim 32.
+    sliced = torch.zeros(64, dtype=torch.float32)[:8]
+    assert _bytes(sliced) == {"cpu": 256}
+
+
+def test_containers_recurse_and_scalars_contribute_nothing() -> None:
+    payload = {
+        "tensors": [torch.zeros(16, dtype=torch.float32), None],
+        "count": 7,
+        "flag": True,
+        "name": "not-a-buffer",
+    }
+    assert _bytes(payload) == {"cpu": 64}
+
+
+def test_numpy_arrays_land_on_the_host() -> None:
+    assert _bytes(np.zeros(32, dtype=np.float32)) == {"cpu": 128}
+
+
+# -- SessionState / manager --------------------------------------------------
+
+
+def test_session_counts_both_buckets_and_dedupes_across_them() -> None:
+    manager = SessionStateManager()
+    session = manager.get_or_create_session("s")
+
+    buffer: LatentBuffer[torch.Tensor] = LatentBuffer()
+    buffer.allocate(maxlen=None)
+    shared = torch.zeros(64, dtype=torch.float32)  # 256 bytes
+    buffer.append(shared)
+    session.put("latents", buffer)
+    session.attrs["conditioning"] = torch.zeros(16, dtype=torch.float32)  # 64 bytes
+
+    assert session.nbytes_by_device() == {"cpu": 320}
+
+    # The same storage reachable from both buckets is one allocation, not two.
+    session.attrs["alias"] = shared
+    assert session.nbytes_by_device() == {"cpu": 320}
+
+
+def test_stats_reports_per_device_totals() -> None:
+    manager = SessionStateManager()
+    manager.get_or_create_session("s").attrs["x"] = torch.zeros(64, dtype=torch.float32)
+
+    stats = manager.stats()
+    assert stats["nbytes:cpu"] == 256
+    assert stats["total_nbytes"] == 256
+
+
+def test_attribute_bytes_were_previously_invisible() -> None:
+    """The defect behind tzhouam's review comment on PR #4487.
+
+    Every byte here lives in ``attrs``, and the old ``nbytes`` walked only the
+    typed objects -- so a session holding real memory reported zero.
+    """
+    manager = SessionStateManager()
+    session = manager.get_or_create_session("s")
+    session.attrs["conv_cache"] = [torch.zeros(1024, dtype=torch.float32)]
+
+    assert session.nbytes_by_device() == {"cpu": 4096}
+
+
+# -- the largest value DreamZero holds ---------------------------------------
+
+
+def test_feat_map_bytes_are_counted() -> None:
+    """``vae_enc_feat_map`` is the largest thing a DreamZero session holds --
+    603 MiB at the shipped config -- and it stays an attribute, so this is
+    exactly the memory the old accounting missed."""
+    manager = SessionStateManager()
+    adapter = DreamZeroStateAdapter("s", manager)
+    adapter.vae_enc_feat_map = [torch.zeros(1024, dtype=torch.float32), None]
+
+    assert manager.stats()["nbytes:cpu"] == 4096
+
+    # Freeing it piecemeal is not something the manager can do: the cache has no
+    # recompute source and the stream breaks without it. Ending the stream (or
+    # the session) is the only release, and that reclaims attributes too.
+    adapter.reset_vae_encoder_stream()
+    assert manager.stats().get("nbytes:cpu", 0) == 0
+
+
+# -- bounded VAE encoder ring ------------------------------------------------
+
+
+def _quantize(encoder_out: torch.Tensor) -> torch.Tensor:
+    """Stand-in for ``_vae_quantize_encoder_out``.
+
+    The real one is ``WanCausalConv3d(z_dim*2, z_dim*2, 1)`` followed by a
+    per-channel affine rescale. Every operation is pointwise in time, which is
+    the property the ring relies on; a 1x1x1 convolution reproduces it.
+    """
+    weight = torch.arange(1, encoder_out.shape[1] + 1, dtype=encoder_out.dtype).view(1, -1, 1, 1, 1)
+    return encoder_out * weight + 0.5
+
+
+def test_ring_is_bounded() -> None:
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    adapter.vae_encoder_out = torch.randn(1, 4, 1, 2, 2)
+    for _ in range(10):
+        adapter.append_vae_encoder_chunk(torch.randn(1, 4, 1, 2, 2))
+
+    assert adapter.vae_encoder_out is not None
+    assert adapter.vae_encoder_out.shape[2] == 2
+
+
+def test_ring_matches_unbounded_history_for_the_frames_that_are_read() -> None:
+    window = 2
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=window)
+    seed = torch.randn(1, 4, 1, 2, 2)
+    chunks = [torch.randn(1, 4, 1, 2, 2) for _ in range(9)]
+
+    adapter.vae_encoder_out = seed
+    unbounded = seed
+    for chunk in chunks:
+        adapter.append_vae_encoder_chunk(chunk)
+        unbounded = torch.cat([unbounded, chunk], dim=2)
+
+    # What the reader actually returns: quantise, then keep the last `window`.
+    from_ring = _quantize(adapter.vae_encoder_out)[:, :, -window:]
+    from_history = _quantize(unbounded)[:, :, -window:]
+    torch.testing.assert_close(from_ring, from_history, rtol=0, atol=0)
+
+
+def test_ring_holds_a_single_frame_before_the_first_body_chunk() -> None:
+    # After `_vae_stream_seed` there is exactly one latent frame, which is what
+    # sends the reader down its padding branch.
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    seed = torch.randn(1, 4, 1, 2, 2)
+    adapter.vae_encoder_out = seed
+
+    assert adapter.vae_encoder_out is not None
+    assert adapter.vae_encoder_out.shape[2] == 1
+    torch.testing.assert_close(adapter.vae_encoder_out, seed, rtol=0, atol=0)
+
+
+def test_resetting_the_stream_clears_the_ring() -> None:
+    adapter = DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=2)
+    adapter.vae_encoder_out = torch.randn(1, 4, 1, 2, 2)
+    adapter.reset_vae_encoder_stream()
+
+    assert adapter.vae_encoder_out is None
+
+
+def test_window_must_be_positive() -> None:
+    with pytest.raises(ValueError, match="vae_encoder_window"):
+        DreamZeroStateAdapter("s", SessionStateManager(), vae_encoder_window=0)

tests/diffusion/models/dreamzero/test_session_state_equivalence.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bit-equivalence tests: DreamZeroStateAdapter vs bespoke DreamZeroState.
+
+These drive both state objects directly with tiny CPU tensors -- no model, no
+GPU -- and assert the adapter (backed by the SessionStateManager) stores and
+returns exactly what the bespoke DreamZeroState does. This gates routing
+DreamZero through the shared session state manager (RFC #4480).
+
+DreamZero's attention KV is engine-owned (AR-Diffusion paged pool, PR #4534)
+and out of the adapter's scope; the state surface tested here is the model's
+non-KV session state.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import (
+    FRAMES_PER_CHUNK,
+    DreamZeroState,
+)
+from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import DreamZeroStateAdapter
+from vllm_omni.experimental.world_models.session_state import SessionStateManager
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _create_both() -> tuple[DreamZeroState, DreamZeroStateAdapter]:
+    return DreamZeroState(), DreamZeroStateAdapter("session-0", SessionStateManager())
+
+
+def test_accumulate_frames_and_call_count_match() -> None:
+    bespoke, adapter = _create_both()
+    rng = np.random.default_rng(0)
+    for _ in range(2 * FRAMES_PER_CHUNK + 1):
+        frame = rng.integers(0, 255, size=(5, 6, 3), dtype=np.uint8)
+        b_out = bespoke.accumulate_frames(frame)
+        a_out = adapter.accumulate_frames(frame)
+        np.testing.assert_array_equal(b_out, a_out)
+        assert bespoke.call_count == adapter.call_count
+        assert len(bespoke.stitched_buffer) == len(adapter.stitched_buffer)
+    # Deque rollover capped at FRAMES_PER_CHUNK.
+    assert len(adapter.stitched_buffer) == FRAMES_PER_CHUNK
+
+
+def test_accumulate_frames_4d_extend_matches() -> None:
+    bespoke, adapter = _create_both()
+    clip = np.random.default_rng(1).integers(0, 255, size=(3, 5, 6, 3), dtype=np.uint8)
+    np.testing.assert_array_equal(bespoke.accumulate_frames(clip), adapter.accumulate_frames(clip))
+    assert len(bespoke.stitched_buffer) == len(adapter.stitched_buffer)
+
+
+def test_video_latents_roundtrip_matches() -> None:
+    bespoke, adapter = _create_both()
+    assert bespoke.get_concatenated_video_latents() is None
+    assert adapter.get_concatenated_video_latents() is None
+
+    rng = torch.Generator().manual_seed(0)
+    for _ in range(3):
+        # (B, T, C, H, W) as returned by the denoise loop.
+        chunk = torch.randn(1, 2, 3, 4, 4, generator=rng)
+        bespoke.append_video_latents(chunk)
+        adapter.append_video_latents(chunk)
+        torch.testing.assert_close(
+            bespoke.get_concatenated_video_latents(),
+            adapter.get_concatenated_video_latents(),
+        )
+
+    bespoke.clear_video_latents()
+    adapter.clear_video_latents()
+    assert bespoke.get_concatenated_video_latents() is None
+    assert adapter.get_concatenated_video_latents() is None
+
+
+def test_vae_stream_fields_match_bespoke() -> None:
+    bespoke, adapter = _create_both()
+    for state in (bespoke, adapter):
+        assert state.vae_stream_initialized is False
+        assert state.vae_enc_feat_map is None
+        assert state.vae_encoder_out is None
+        assert state.vae_pending_body_frames is None
+
+    feat_map: list[torch.Tensor | None] = [torch.ones(1, 2), None]
+    encoder_out = torch.randn(1, 4, 2, 3, 3)
+    pending = torch.randn(1, 3, 2, 8, 8)
+    for state in (bespoke, adapter):
+        state.vae_enc_feat_map = feat_map
+        state.vae_encoder_out = encoder_out
+        state.vae_pending_body_frames = pending
+        state.vae_stream_initialized = True
+
+    assert adapter.vae_stream_initialized is bespoke.vae_stream_initialized is True
+    torch.testing.assert_close(adapter.vae_encoder_out, bespoke.vae_encoder_out)
+    torch.testing.assert_close(adapter.vae_pending_body_frames, bespoke.vae_pending_body_frames)
+    assert adapter.vae_enc_feat_map is feat_map
+
+    for state in (bespoke, adapter):
+        state.reset_vae_encoder_stream()
+        assert state.vae_stream_initialized is False
+        assert state.vae_enc_feat_map is None
+        assert state.vae_encoder_out is None
+        assert state.vae_pending_body_frames is None
+
+
+def test_reset_clears_state_and_should_reset_matches() -> None:
+    bespoke, adapter = _create_both()
+    tokens = torch.tensor([1, 2, 3])
+    # language unset -> both want reset.
+    assert bespoke.should_reset(tokens, 4, -1) is adapter.should_reset(tokens, 4, -1) is True
+
+    for state in (bespoke, adapter):
+        state.language = tokens
+        state.current_start_frame = 2
+    # same language -> no reset; both agree.
+    assert bespoke.should_reset(tokens, 4, -1) == adapter.should_reset(tokens, 4, -1) is False
+    # language change -> reset.
+    other = torch.tensor([9, 9])
+    assert bespoke.should_reset(other, 4, -1) == adapter.should_reset(other, 4, -1) is True
+    # local_attn_size exceeded -> reset ("inference", not "session").
+    assert bespoke.reset_reason(tokens, 4, 1) == adapter.reset_reason(tokens, 4, 1) == "inference"
+
+    # A session reset clears everything, including the prompt-embed cache and
+    # the VAE encoder stream (mirrors DreamZeroState.reset()).
+    for state in (bespoke, adapter):
+        state.prompt_embeds = torch.ones(1, 2)
+        state.vae_encoder_out = torch.ones(1, 4, 1, 2, 2)
+        state.vae_stream_initialized = True
+        state.reset()
+        assert state.language is None
+        assert state.prompt_embeds is None
+        assert state.current_start_frame == 0
+        assert state.call_count == 0
+        assert state.vae_stream_initialized is False
+        assert state.vae_encoder_out is None
+
+
+def test_window_reset_preserves_latents_prompt_and_vae_stream() -> None:
+    # reset(clear_video_latents=False) is the window ("inference") reset: the
+    # prompt is unchanged and the VAE feat_cache history is independent of the
+    # DiT attention window, so both survive alongside the video latents.
+    bespoke, adapter = _create_both()
+    tokens = torch.tensor([1, 2, 3])
+    embeds = torch.randn(1, 4)
+    encoder_out = torch.randn(1, 4, 2, 3, 3)
+    chunk = torch.randn(1, 2, 3, 4, 4)
+    for state in (bespoke, adapter):
+        state.language = tokens
+        state.prompt_embeds = embeds
+        state.vae_encoder_out = encoder_out
+        state.vae_stream_initialized = True
+        state.current_start_frame = 7
+        state.append_video_latents(chunk)
+
+        state.reset_inference_state()
+
+        assert state.current_start_frame == 0
+        assert state.call_count == 0
+        assert state.clip_feas is None
+        torch.testing.assert_close(state.language, tokens)
+        torch.testing.assert_close(state.prompt_embeds, embeds)
+        assert state.vae_stream_initialized is True
+        torch.testing.assert_close(state.vae_encoder_out, encoder_out)
+        assert state.get_concatenated_video_latents() is not None
+
+    torch.testing.assert_close(
+        bespoke.get_concatenated_video_latents(),
+        adapter.get_concatenated_video_latents(),
+    )
+
+
+def test_metadata_fields_persist_across_adapter_instances() -> None:
+    # A fresh adapter for the same session must see prior metadata (the manager
+    # is the single source of truth).
+    manager = SessionStateManager()
+    first = DreamZeroStateAdapter("s", manager)
+    first.current_start_frame = 5
+    first.clip_feas = torch.ones(2, 2)
+    first.append_video_latents(torch.ones(1, 2, 3, 4, 4))
+
+    second = DreamZeroStateAdapter("s", manager)
+    assert second.current_start_frame == 5
+    torch.testing.assert_close(second.clip_feas, torch.ones(2, 2))
+    # Video latents appended via `first` are visible via `second`.
+    assert second.get_concatenated_video_latents() is not None
+
+
+def test_session_lru_caps_retained_sessions() -> None:
+    manager = SessionStateManager(max_sessions=3)
+    for i in range(5):
+        DreamZeroStateAdapter(f"s{i}", manager)
+    assert len(manager) == 3
+    # Oldest two evicted; newest three retained.
+    assert "s0" not in manager
+    assert "s1" not in manager
+    for i in (2, 3, 4):
+        assert f"s{i}" in manager
+
+
+def test_adapter_keeps_state_when_its_session_is_evicted() -> None:
+    # Regression: an adapter mid-generation must not lose its data when the
+    # manager evicts its session to make room for newer ones. The adapter pins
+    # the session, so its stored latents survive eviction from the lookup table.
+    manager = SessionStateManager(max_sessions=2)
+    adapter = DreamZeroStateAdapter("active", manager)
+    chunk = torch.randn(1, 2, 3, 4, 4)
+    adapter.append_video_latents(chunk)
+    adapter.language = torch.tensor([1, 2, 3])
+
+    # Push "active" out of the manager's table with newer sessions.
+    for i in range(3):
+        DreamZeroStateAdapter(f"other{i}", manager)
+    assert "active" not in manager
+
+    # The in-use adapter still returns its own data.
+    assert adapter.get_concatenated_video_latents() is not None
+    torch.testing.assert_close(adapter.language, torch.tensor([1, 2, 3]))
+
+
+def test_session_reset_clears_metadata() -> None:
+    # Regression: resetting a session must not leave stale metadata behind, so a
+    # fresh adapter for the same session sees defaults, not the old values.
+    manager = SessionStateManager()
+    adapter = DreamZeroStateAdapter("s", manager)
+    adapter.language = torch.tensor([1, 2, 3])
+    adapter.current_start_frame = 5
+    adapter.clip_feas = torch.ones(2, 2)
+    adapter.prompt_embeds = torch.ones(1, 2)
+
+    adapter.reset()
+
+    fresh = DreamZeroStateAdapter("s", manager)
+    assert fresh.language is None
+    assert fresh.current_start_frame == 0
+    assert fresh.clip_feas is None
+    assert fresh.prompt_embeds is None
+
+
+def test_bespoke_path_without_memory_manager() -> None:
+    # Regression: a pipeline built via __new__ (lightweight fixtures) seeds only
+    # the bespoke fields and never sets _memory_manager; _get_or_create_state
+    # must fall through to the bespoke DreamZeroState path without AttributeError.
+    from collections import OrderedDict
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
+        MAX_DREAMZERO_SESSIONS,
+        DreamZeroPipeline,
+    )
+
+    pipe = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipe._states = OrderedDict()
+    pipe._max_session_states = MAX_DREAMZERO_SESSIONS
+
+    state = DreamZeroPipeline._get_or_create_state(pipe, "x")
+    assert isinstance(state, DreamZeroState)

tests/diffusion/models/dreamzero/test_session_state_lifecycle.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Session end-of-life routing (RFC #4480, adapting to #5271).
+
+The AR-Diffusion runner raises an explicit end-of-session signal on reset,
+close, eviction, and failed forwards (``reset_ar_diffusion_session`` /
+``close_ar_diffusion_session``). On the bespoke path those pop ``_states``; on
+the opt-in manager path the session lives in the manager instead, so the hooks
+must release it there or every closed session leaks its buffers.
+
+All CPU, tiny tensors, no model.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
+from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import (
+    DreamZeroStateAdapter,
+)
+from vllm_omni.experimental.world_models.session_state import (
+    LatentBuffer,
+    SessionStateManager,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# -- SessionStateManager.drop_session ---------------------------------------
+
+
+def test_drop_session_absent_returns_false() -> None:
+    manager = SessionStateManager()
+    assert manager.drop_session("missing") is False
+
+
+def test_drop_session_removes_and_frees() -> None:
+    manager = SessionStateManager()
+    session = manager.get_or_create_session("s")
+    buffer: LatentBuffer[torch.Tensor] = LatentBuffer()
+    buffer.allocate(maxlen=None)
+    buffer.append(torch.zeros(256, dtype=torch.float32))
+    session.put("payload", buffer)
+    assert manager.nbytes_by_device().get("cpu", 0) >= 1024
+
+    assert manager.drop_session("s") is True
+    assert "s" not in manager
+    assert len(manager) == 0
+    # The session is gone from the table and its buffers were reset, so the
+    # manager reports no held bytes.
+    assert manager.nbytes_by_device() == {}
+
+
+def test_drop_session_recreates_fresh() -> None:
+    manager = SessionStateManager()
+    manager.get_or_create_session("s").attrs["k"] = 1
+    manager.drop_session("s")
+    assert manager.get_or_create_session("s").attrs == {}
+
+
+# -- pipeline hook routing ---------------------------------------------------
+
+
+def _manager_pipe(manager: SessionStateManager, session_id: str) -> DreamZeroPipeline:
+    """A pipeline built via __new__ whose state is a manager-backed adapter."""
+    pipe = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipe._states = OrderedDict()
+    pipe._memory_manager = manager
+    pipe.state = DreamZeroStateAdapter(session_id, manager, vae_encoder_window=2)
+    return pipe
+
+
+@pytest.mark.parametrize("hook", ["close_ar_diffusion_session", "reset_ar_diffusion_session"])
+def test_hook_releases_manager_session_and_clears_alias(hook: str) -> None:
+    manager = SessionStateManager()
+    pipe = _manager_pipe(manager, "sess")
+    assert "sess" in manager
+
+    getattr(DreamZeroPipeline, hook)(pipe, "sess")
+
+    assert "sess" not in manager
+    # The alias viewed the released session, so it is dropped and rebuilt lazily.
+    assert pipe.state is None
+
+
+def test_hook_keeps_alias_for_a_different_session() -> None:
+    manager = SessionStateManager()
+    pipe = _manager_pipe(manager, "default")
+    manager.get_or_create_session("other")
+
+    DreamZeroPipeline.close_ar_diffusion_session(pipe, "other")
+
+    assert "other" not in manager
+    # The alias views "default", not the closed session, so it survives.
+    assert isinstance(pipe.state, DreamZeroStateAdapter)
+    assert pipe.state.session_id == "default"
+
+
+def test_bespoke_hook_pops_states_and_clears_alias() -> None:
+    pipe = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipe._states = OrderedDict()
+    pipe._memory_manager = None
+    state = DreamZeroState()
+    pipe._states["sess"] = state
+    pipe.state = state
+
+    DreamZeroPipeline.close_ar_diffusion_session(pipe, "sess")
+
+    assert "sess" not in pipe._states
+    assert pipe.state is None

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -9,6 +9,8 @@ from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransforme
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def test_ltx_rms_norm_no_affine_matches_official_torch_module():
     norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)

tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py

@@ -6,6 +6,8 @@ import torch
 
 from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def _pipeline_without_init() -> SenseNovaU1Pipeline:
     return object.__new__(SenseNovaU1Pipeline)

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -5,7 +5,6 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
@@ -16,32 +15,26 @@ AUDIO_MODEL = {
     "stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None},
 }
 
-IMAGE_VIDEO_MODELS = {
+IMAGE_MODELS = {
     "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100},
-    "Tongyi-MAI/Z-Image-Turbo": {"cuda": 2500, "rocm": 2100},
-    "OmniGen2/OmniGen2": {"cuda": 2500, "rocm": 2100},
 }
 
-MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
+MODELS = {**AUDIO_MODEL, **IMAGE_MODELS}
 
 MODEL_MARKS = {
     "riverclouds/qwen_image_random": pytest.mark.core_model,
-    "Tongyi-MAI/Z-Image-Turbo": pytest.mark.advanced_model,
     "stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
-    "OmniGen2/OmniGen2": pytest.mark.full_model,
 }
 
 _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"}
 
-# Aliased for backward compatibility (imported by test_diffusion_layerwise_offload.py).
-_skip_if_gated_repo_inaccessible = skip_if_gated_repo_inaccessible
 
 AUDIO_MODEL_PARAMS = {
     "runner_params": {},
     "sampler_params": {},
 }
 
-IMAGE_VIDEO_MODELS_PARAMS = {
+IMAGE_MODEL_PARAMS = {
     "runner_params": {},
     "sampler_params": {
         "height": 256,
@@ -60,7 +53,7 @@ def inference(model_name: str, offload: bool = True):
     if model_name in AUDIO_MODEL:
         params = AUDIO_MODEL_PARAMS
     else:
-        params = IMAGE_VIDEO_MODELS_PARAMS
+        params = IMAGE_MODEL_PARAMS
 
     with OmniRunner(
         model_name,
@@ -111,19 +104,12 @@ def check_audio_determinism(audio1, audio2, atol=1e-2):
     [pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS],
 )
 def test_cpu_offload_diffusion_model(model_name: str):
-    if model_name == "OmniGen2/OmniGen2":
-        pytest.skip("issue #4537")
-    if model_name in _GATED_MODELS:
-        _skip_if_gated_repo_inaccessible(model_name)
     try:
         offload_peak_memory, output_offload = inference(model_name, offload=True)
         cleanup_dist_env_and_memory()
         no_offload_peak_memory, output_no_offload = inference(model_name, offload=False)
     except ValueError as exc:
-        # omni_snapshot_download wraps GatedRepoError in a ValueError.
-        # If the pre-flight guard above did not catch it (e.g. partial
-        # HF_TOKEN where config.json is accessible but weight shards are
-        # blocked), skip instead of failing.
+        # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing.
         if "Access to model" in str(exc) and "is restricted" in str(exc):
             pytest.skip(
                 f"Skipping: gated HF repo {model_name!r} inaccessible "

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -3,11 +3,8 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.diffusion.offloader.test_diffusion_cpu_offload import (
-    _GATED_MODELS,
-    _skip_if_gated_repo_inaccessible,
-)
 from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -23,6 +20,19 @@ IMAGE_VIDEO_MODELS = {
 
 MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
 
+LAYERWISE_MODEL_CASES = [
+    pytest.param(
+        "stabilityai/stable-audio-open-1.0",
+        marks=[pytest.mark.full_model, pytest.mark.diffusion],
+        id="stable_audio",
+    ),
+    pytest.param(
+        "riverclouds/qwen_image_random",
+        marks=[pytest.mark.core_model, pytest.mark.diffusion],
+        id="qwen_image_random",
+    ),
+]
+
 AUDIO_MODEL_PARAMS = {
     "runner_params": {},
     "sampler_params": {},
@@ -91,7 +101,8 @@ def run_inference(
     return peak, output
 
 
-@pytest.mark.parametrize("model_name", list(MODELS.keys()))
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("model_name", LAYERWISE_MODEL_CASES)
 def test_layerwise_offload_diffusion_model(model_name: str):
     """Test that layerwise offloading reduces GPU memory usage.
 
@@ -100,8 +111,6 @@ def test_layerwise_offload_diffusion_model(model_name: str):
     offloader keeps only a single transformer block on GPU at a time, with
     prefetching for compute-memory overlap.
     """
-    if model_name in _GATED_MODELS:
-        _skip_if_gated_repo_inaccessible(model_name)
     try:
         # Run without layerwise offloading (baseline)
         no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False)
@@ -111,10 +120,7 @@ def test_layerwise_offload_diffusion_model(model_name: str):
         layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True)
         cleanup_dist_env_and_memory()
     except ValueError as exc:
-        # omni_snapshot_download wraps GatedRepoError in a ValueError.
-        # If the pre-flight guard above did not catch it (e.g. partial
-        # HF_TOKEN where config.json is accessible but weight shards are
-        # blocked), skip instead of failing.
+        # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing.
         if "Access to model" in str(exc) and "is restricted" in str(exc):
             pytest.skip(
                 f"Skipping: gated HF repo {model_name!r} inaccessible "

tests/diffusion/quantization/test_gguf_diffusion.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end tests for GGUF quantization on diffusion models.
+
+Validates that GGUF-quantized diffusion models generate valid images and
+use less peak GPU memory than BF16 baseline.
+
+Requires vllm-omni to be installed alongside the vllm-gguf-plugin package.
+
+Usage:
+    pytest tests/diffusion/quantization/test_gguf_diffusion.py -v
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+class DiffusionGGUFTestConfig(NamedTuple):
+    artifact_prefix: str
+    gguf_model: str
+    hf_model: str
+    min_cosine_similarity: float = 0.95
+
+
+Z_IMAGE_CONFIG = DiffusionGGUFTestConfig(
+    artifact_prefix="zimage",
+    hf_model="Tongyi-MAI/Z-Image-Turbo",
+    gguf_model="unsloth/Z-Image-Turbo-GGUF:Q4_0",
+)
+
+
+FLUX_CONFIG = DiffusionGGUFTestConfig(
+    artifact_prefix="flux2_klein",
+    hf_model="black-forest-labs/FLUX.2-klein-4B",
+    gguf_model="unsloth/FLUX.2-klein-4B-GGUF:Q8_0",
+)
+
+
+def _image_cosine_similarity(hf_image, gguf_image) -> float:
+    hf_tensor = torch.as_tensor(np.array(hf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
+    gguf_tensor = torch.as_tensor(np.array(gguf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
+    return torch.nn.functional.cosine_similarity(hf_tensor, gguf_tensor, dim=0).item()
+
+
+def _generate_single_stage_image(
+    model: str,
+    height: int = 256,
+    width: int = 256,
+    num_inference_steps: int = 20,
+    seed: int = 42,
+    **extra_kwargs,
+) -> tuple[list, float]:
+    """Generate an image with a single-stage diffusion model.
+
+    Returns (images, peak_memory_gib).
+    """
+    omni_kwargs = dict(extra_kwargs)
+
+    memory_monitor = DeviceMemoryMonitor(device_index=0)
+    memory_monitor.start()
+    try:
+        with OmniRunner(model, **omni_kwargs) as runner:
+            generator = torch.Generator(
+                device=current_omni_platform.device_type,
+            ).manual_seed(seed)
+            outputs = runner.omni.generate(
+                "a photo of a cat sitting on a laptop keyboard",
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=0.0,
+                    generator=generator,
+                ),
+            )
+    finally:
+        memory_monitor.stop()
+    peak_mem = memory_monitor.peak_used_mb / 1024
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    if hasattr(first_output, "images") and first_output.images:
+        images = first_output.images
+    else:
+        assert hasattr(first_output, "request_output") and first_output.request_output
+        request_output = first_output.request_output
+        if isinstance(request_output, list):
+            req_out = request_output[0]
+        else:
+            req_out = request_output
+        assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+        images = req_out.images
+    assert len(images) >= 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+    return images, peak_mem
+
+
+@hardware_test(res={"cuda": "L4"})
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"])
+def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Z-Image-Turbo GGUF generates valid images and uses less memory than BF16."""
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+    # BF16 baseline
+    hf_images, mem_bf16 = _generate_single_stage_image(
+        model=model.hf_model,
+    )
+
+    # GGUF
+    images, mem_gguf = _generate_single_stage_image(
+        model=model.hf_model,
+        diffusion_quantization_config={
+            "method": "gguf",
+            "gguf_model": model.gguf_model,
+        },
+    )
+
+    hf_image_path = f"test_{model.artifact_prefix}_hf.png"
+    gguf_image_path = f"test_{model.artifact_prefix}_gguf.png"
+    assert len(hf_images) >= 1
+    hf_images[0].save(hf_image_path)
+    assert len(images) >= 1
+    images[0].save(gguf_image_path)
+    print(f"Saved HF image: {hf_image_path}")
+    print(f"Saved GGUF image: {gguf_image_path}")
+
+    image_cosine_similarity = _image_cosine_similarity(hf_images[0], images[0])
+    print(f"{model.artifact_prefix} image cosine similarity: {image_cosine_similarity:.4f}")
+    assert image_cosine_similarity >= model.min_cosine_similarity, (
+        f"GGUF image cosine similarity ({image_cosine_similarity:.4f}) should be >= {model.min_cosine_similarity:.4f}"
+    )
+
+    print(f"{model.artifact_prefix} BF16 peak VRAM delta: {mem_bf16:.2f} GiB")
+    print(f"{model.artifact_prefix} GGUF peak VRAM delta: {mem_gguf:.2f} GiB")
+    reduction = (mem_bf16 - mem_gguf) / mem_bf16 * 100
+    print(f"VRAM reduction: {reduction:.1f}%")
+    assert mem_gguf < mem_bf16, f"GGUF ({mem_gguf:.2f} GiB) should use less VRAM than BF16 ({mem_bf16:.2f} GiB)"

tests/diffusion/quantization/test_quantization_quality.py

@@ -431,6 +431,8 @@ def _free_gpu_memory():
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch):
     from benchmarks.diffusion.quantization_quality import _generate_image as benchmark_generate_image
     from vllm_omni.outputs import OmniRequestOutput

tests/diffusion/quantization/test_trajectory_similarity_tool.py

@@ -22,6 +22,8 @@ from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import
     summarize_output_image_metrics,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def test_compute_tensor_metrics_identical_tensors():
     metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3))

tests/diffusion/test_diffusion_config_fields.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ensure diffusion stage YAML configs only use valid OmniDiffusionConfig fields.
+
+Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
+"""
+
+from dataclasses import fields
+from pathlib import Path
+
+import pytest
+import yaml
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+try:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+except Exception:
+    OmniDiffusionConfig = None
+
+
+@pytest.mark.skipif(
+    OmniDiffusionConfig is None,
+    reason="OmniDiffusionConfig could not be imported (missing torch?)",
+)
+def test_diffusion_stage_configs_only_contain_valid_fields():
+    """Diffusion stage engine_args must only contain OmniDiffusionConfig fields.
+
+    Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
+    """
+    # Scan both main configs and test configs
+    repo_root = Path(__file__).parent.parent.parent
+    config_dirs = [
+        repo_root / "vllm_omni" / "model_executor" / "stage_configs",
+    ]
+    # Also scan test directories recursively
+    test_dir = repo_root / "tests"
+
+    yaml_paths: list[Path] = []
+    for config_dir in config_dirs:
+        yaml_paths.extend(sorted(config_dir.glob("*.yaml")))
+    yaml_paths.extend(sorted(test_dir.rglob("*.yaml")))
+
+    valid_fields = {f.name for f in fields(OmniDiffusionConfig)}
+    # model_stage is consumed by the stage init layer, not OmniDiffusionConfig
+    valid_fields.add("model_stage")
+    # model_arch is consumed by the stage init layer for diffusion model class resolution
+    valid_fields.add("model_arch")
+    # "quantization" is mapped to "quantization_config" by from_kwargs() backwards-compat
+    valid_fields.add("quantization")
+
+    invalid_entries: list[tuple[str, set[str]]] = []
+    for yaml_path in yaml_paths:
+        with open(yaml_path) as fh:
+            config = yaml.safe_load(fh)
+
+        stages = config.get("stage_args", config.get("stages", []))
+        for stage in stages:
+            if stage.get("stage_type") != "diffusion":
+                continue
+            engine_args = stage.get("engine_args", {})
+            invalid = set(engine_args.keys()) - valid_fields
+            if invalid:
+                invalid_entries.append((yaml_path.relative_to(repo_root), invalid))
+
+    assert not invalid_entries, "Diffusion stage configs contain fields not in OmniDiffusionConfig:\n" + "\n".join(
+        f"  {name}: {sorted(bad)}" for name, bad in invalid_entries
+    )

tests/diffusion/test_diffusion_config_propagation.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests that parallel_config survives the create_default_diffusion roundtrip.
+
+Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862
+"""
+
+from collections.abc import Mapping
+
+import pytest
+import torch
+
+from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.diffusion.data import (
+    DiffusionParallelConfig,
+    OmniDiffusionConfig,
+)
+from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
+    """Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
+
+    Does NOT manually reconstruct parallel_config — relies on
+    OmniDiffusionConfig.__post_init__ to handle the dict, just like
+    the production code path does.
+    """
+    stages = StageConfigFactory.create_default_diffusion(kwargs)
+    engine_args = dict(stages[0]["engine_args"])
+    return OmniDiffusionConfig.from_kwargs(**engine_args)
+
+
+class TestParallelConfigPropagation:
+    """Core regression tests: parallel_config must survive serialization."""
+
+    def test_tp2_roundtrip(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=2)
+        od = _roundtrip_diffusion_config(model="test-model", parallel_config=pc)
+        assert od.parallel_config.tensor_parallel_size == 2
+        assert od.parallel_config.world_size == 2
+
+    def test_tp4_devices_and_config(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=4)
+        stages = StageConfigFactory.create_default_diffusion({"parallel_config": pc, "model": "x"})
+        assert stages[0]["runtime"]["devices"] == "0,1,2,3"
+
+        # Let __post_init__ reconstruct from dict (real code path)
+        ea = dict(stages[0]["engine_args"])
+        od = OmniDiffusionConfig.from_kwargs(**ea)
+        assert od.parallel_config.tensor_parallel_size == 4
+        assert od.parallel_config.world_size == 4
+
+    def test_sp_config_roundtrip(self):
+        pc = DiffusionParallelConfig(
+            tensor_parallel_size=2,
+            ulysses_degree=2,
+            ring_degree=1,
+        )
+        od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
+        assert od.parallel_config.ulysses_degree == 2
+        assert od.parallel_config.ring_degree == 1
+
+    def test_mask_sp_padding_roundtrip(self):
+        pc = DiffusionParallelConfig(ulysses_degree=2, mask_sp_padding=True)
+        od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
+        assert od.parallel_config.mask_sp_padding is True
+
+    def test_mask_sp_padding_defaults_false(self):
+        pc = DiffusionParallelConfig(ulysses_degree=2)
+        od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
+        assert od.parallel_config.mask_sp_padding is False
+
+    def test_cfg_parallel_roundtrip(self):
+        pc = DiffusionParallelConfig(cfg_parallel_size=2)
+        od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
+        assert od.parallel_config.cfg_parallel_size == 2
+        assert od.parallel_config.world_size == 2
+
+    def test_no_parallel_config_defaults_to_tp1(self):
+        od = _roundtrip_diffusion_config(model="x")
+        assert od.parallel_config.tensor_parallel_size == 1
+        assert od.parallel_config.world_size == 1
+
+    def test_num_gpus_derived_from_world_size(self):
+        pc = DiffusionParallelConfig(tensor_parallel_size=2)
+        od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
+        assert od.num_gpus == 2
+
+
+class TestCreateDefaultDiffusion:
+    """Verify engine_args structure from create_default_diffusion."""
+
+    def test_parallel_config_serialized_as_dict(self):
+        """The key fix: parallel_config must appear in engine_args as a dict."""
+        pc = DiffusionParallelConfig(tensor_parallel_size=2)
+        stages = StageConfigFactory.create_default_diffusion({"model": "x", "parallel_config": pc})
+        ea = stages[0]["engine_args"]
+        assert "parallel_config" in ea
+        assert isinstance(ea["parallel_config"], Mapping)
+        assert ea["parallel_config"]["tensor_parallel_size"] == 2
+
+    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"
+
+    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_single_gpu_default_devices(self):
+        stages = StageConfigFactory.create_default_diffusion({"model": "x"})
+        assert stages[0]["runtime"]["devices"] == "0"
+
+    def test_extra_kwargs_forwarded(self):
+        stages = StageConfigFactory.create_default_diffusion(
+            {"model": "x", "enforce_eager": True, "lora_path": "/tmp/lora"}
+        )
+        ea = stages[0]["engine_args"]
+        assert ea["enforce_eager"] is True
+        assert ea["lora_path"] == "/tmp/lora"
+
+
+def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
+    od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
+
+    od_config.update_multimodal_support()
+
+    assert od_config.supports_multimodal_inputs is True
+    assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
+
+
+def test_additional_config_roundtrip():
+    additional_config = {"torchair_graph_config": {"enabled": True}}
+    od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)
+    assert od.additional_config == additional_config

tests/diffusion/test_diffusion_engine.py

@@ -14,6 +14,7 @@ import torch
 from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import (
     DiffusionEngine,
@@ -30,6 +31,8 @@ from vllm_omni.diffusion.sched.interface import (
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
 
 @dataclass
 class DiffusionSchedulerOutput:
@@ -56,6 +59,8 @@ class MockScheduler:
     def __init__(self):
         self._waiting_queue = []
         self._step_id = 0
+        # Match RequestScheduler API used by request-batch admission wait.
+        self.max_num_running_reqs = 5
 
     def add_request(self, request):
         self._waiting_queue.append(request)
@@ -64,6 +69,12 @@ class MockScheduler:
     def has_requests(self):
         return len(self._waiting_queue) > 0
 
+    def num_waiting_requests(self) -> int:
+        return len(self._waiting_queue)
+
+    def num_running_requests(self) -> int:
+        return 0
+
     def schedule(self) -> DiffusionSchedulerOutput:
         if not self._waiting_queue:
             return DiffusionSchedulerOutput(step_id=self._step_id)
@@ -582,9 +593,7 @@ def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None:
     assert moved is value
 
 
-@pytest.mark.diffusion
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
     tensor = torch.arange(8, dtype=torch.float32, device="cuda")
     other = torch.arange(4, dtype=torch.int64, device="cuda")
@@ -610,6 +619,7 @@ async def _consume_final_output(generator):
     return final_output
 
 
+@pytest.mark.cpu
 @pytest.mark.asyncio
 async def test_async_add_req_and_stream_response():
     engine = object.__new__(DiffusionEngine)
@@ -621,13 +631,24 @@ async def test_async_add_req_and_stream_response():
     engine._cv = threading.Condition(engine._rpc_lock)
     engine._init_lock = asyncio.Lock()
     engine._closed = False
-    engine.od_config = SimpleNamespace(streaming_output=False)
+    # Enable admission wait so concurrent adds land in one schedule wave;
+    # otherwise the first request can execute alone (~1s) and the rest in a
+    # second wave (~2s), failing the latency-spread assertion.
+    engine.od_config = SimpleNamespace(
+        streaming_output=False,
+        request_batch_max_wait_ms=500.0,
+    )
     engine._loop_started = False
     engine.main_loop = None
-    engine.supports_request_batch = False
+    engine.supports_request_batch = True
+    engine.step_execution = False
     engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
 
-    engine._finalize_finished_request = lambda rid, out, err: out.result
+    def _finalize(rid, out, err=None):
+        # Stream consumers stop on ``finished``; keep result_data for assertions.
+        return SimpleNamespace(result_data=out.result.result_data, finished=True)
+
+    engine._finalize_finished_request = _finalize
 
     def mock_execute_batch(sched_output):
         request_ids = sched_output.scheduled_request_ids

tests/diffusion/test_diffusion_engine_metrics.py

@@ -12,6 +12,11 @@ from __future__ import annotations
 import ast
 import os
 
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 _ENGINE_PATH = os.path.normpath(
     os.path.join(
         os.path.dirname(__file__),