[Bugfix] Fix tokenizer config for higgs-audio-v3 test error (#4918)

Signed-off-by: herotai214 <herotai214@gmail.com>

.gitignore

@@ -208,6 +208,7 @@ cache/
 .cache/
 diffusion_cache/
 kv_cache/
+!vllm_omni/experimental/ar_diffusion/kv_cache/
 
 # Logs
 logs/

tests/ar_diffusion/__init__.py

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

tests/ar_diffusion/test_engine_routing.py

@@ -0,0 +1,136 @@
+# 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,
+)
+
+
+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)

tests/ar_diffusion/test_kv_cache.py

@@ -0,0 +1,262 @@
+# 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.request import RequestStatus
+
+try:
+    from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+except ModuleNotFoundError:
+    KVCacheSpecRegistry = None
+    from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
+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
+
+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()
+    if KVCacheSpecRegistry is not None:
+        assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+    else:
+        assert spec_manager_map[type(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,
+        cross_attn_length=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 > 2 * (2 + 1) + 2  # above the min-blocks floor, so the deduction is what's tested
+
+
+def _make_kv(*, local_branches, num_frame_per_block=2, window_chunks=9):
+    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"),
+        local_branches=local_branches,
+        num_frame_per_block=num_frame_per_block,
+    )
+
+
+def test_pool_floor_is_branch_aware():
+    """CFG-parallel rank (one local branch) sizes for one window + in-flight chunk;
+    a single-process run (both 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_branch
+    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_branch
+    assert one.num_blocks_total == 13 + one.scratch_blocks_per_branch
+
+
+def test_scratch_maps_to_slot_zero_with_one_local_branch():
+    """A CFG-parallel rank runs exactly one 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(True, 0, 2) == one.scratch_block_ids(False, 0, 2)
+
+    two = _make_kv(local_branches=2)
+    assert two.scratch_block_ids(True, 0, 2) != two.scratch_block_ids(False, 0, 2)
+
+
+def test_scratch_exhaustion_still_raises():
+    one = _make_kv(local_branches=1)
+    cap = one.scratch_blocks_per_branch
+    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
+        one.scratch_block_ids(False, 0, cap + 1)
+
+
+def test_invalid_local_branches_rejected():
+    with pytest.raises(ValueError, match="local_branches"):
+        _make_kv(local_branches=3)

tests/ar_diffusion/test_kv_orchestration.py

@@ -0,0 +1,93 @@
+# 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.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+
+DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16)
+
+
+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 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, pos, neg, num_layers=kv.num_layers)
+    for _ in range(3):
+        for adapter in (pos, neg):
+            kv.allocate_chunk(adapter)
+            kv.commit_chunk(adapter)
+    # Both 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/ar_diffusion/test_paged_attention.py

@@ -0,0 +1,394 @@
+# 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.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
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+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,
+        device=device,
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, 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,
+    *,
+    is_negative: bool,
+    n_chunks: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ctx = st.get_kv_caches(is_negative, 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(is_negative)
+    return k, v
+
+
+def test_paged_context_allocates_lazily_and_commits_after_forward():
+    _, st = make_state()
+
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+    assert st.pos.completed_chunks == 0
+    assert ctx.current_video_slot_mapping is None
+
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 1
+
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 1
+    assert st._committed[False] == BLOCK
+
+
+def test_scratch_video_and_action_blocks_do_not_commit():
+    kv, st = make_state()
+
+    ctx = st.get_kv_caches(False, 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(False, 0, 2)
+    assert ctx.action_scratch_block_ids == kv.scratch_block_ids(False, 2, 1)
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 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,
+            is_negative=False,
+            n_chunks=history_chunks,
+            dtype=dtype,
+            device=device,
+        )
+        history_k_parts.append(k)
+        history_v_parts.append(v)
+
+    ctx = st.get_kv_caches(False, 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.pos.completed_chunks
+    st.commit_paged_context(False)
+    assert st.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,
+        is_negative=False,
+        n_chunks=history_chunks,
+        dtype=dtype,
+        device=device,
+    )
+
+    layer_ctx = st.get_kv_caches(False, 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(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    ctx2 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    # 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(False, 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(False, 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(False, 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(False)
+        return out
+
+    torch._dynamo.reset()
+    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}"

tests/ar_diffusion/test_pipeline_bridge.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
+    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,
+        cross_attn_length=cross_attn_length,
+        device=torch.device("cpu"),
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def _prepare_and_commit(st: ARDiffusionKVState, is_negative: bool, n_chunks: int) -> None:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    st.commit_paged_context(is_negative)
+
+
+def test_get_returns_paged_layer_contexts_when_nothing_committed():
+    _, st = make_state(num_layers=2)
+    for neg in (False, True):
+        contexts = st.get_kv_caches(neg, 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(False, seq_len=2 * BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+
+    assert st.pos.completed_chunks == 0
+    assert kv.window_block_ids(st.pos) == []
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 2
+
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 2
+    assert st._committed[False] == 2 * BLOCK
+    assert len(kv.window_block_ids(st.pos)) == 2
+
+
+def test_scratch_paged_context_does_not_commit():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, 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(False, 0, 1)
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+    assert kv.window_block_ids(st.pos) == []
+
+
+def test_eviction_bounds_resident_window():
+    kv, st = make_state(num_layers=1, window_chunks=3)
+    _prepare_and_commit(st, False, 3)
+    assert len(kv.window_block_ids(st.pos)) == 3
+
+    _prepare_and_commit(st, False, 2)
+    ctx = st.get_kv_caches(False, 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, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+    assert kv.window_block_ids(st.neg) == []
+
+
+def test_reset_clears_session_window():
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)
+    _prepare_and_commit(st, True, 1)
+    assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK
+    free_before = kv.manager.block_pool.get_num_free_blocks()
+
+    st.reset()
+
+    assert st._committed == {False: 0, True: 0}
+    assert st._paged_pending == {False: None, True: None}
+    assert kv.window_block_ids(st.pos) == []
+    assert kv.window_block_ids(st.neg) == []
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+
+
+def test_get_cross_kv_caches_raises_before_populate():
+    """Engine ownership guard: a cross-attn read before _kv_populate_cross must fail loud."""
+    _, st = make_state()  # cross_attn_length == 0, nothing populated
+    for neg in (False, True):
+        with pytest.raises(RuntimeError, match="cross-attn read before"):
+            st.get_cross_kv_caches(neg)
+    # Guards the full AND: the populated flag alone is not enough without a pool.
+    st._cross_text_populated[False] = True
+    with pytest.raises(RuntimeError, match="cross-attn read before"):
+        st.get_cross_kv_caches(False)
+
+
+def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
+    L = 8
+    kv, st = make_state(num_layers=2, cross_attn_length=L)
+    written = []
+    for i in range(2):
+        k = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        v = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        kv.write_cross_kv(i, False, k, v)
+        written.append((k, v))
+    st._cross_text_populated[False] = True
+
+    out = st.get_cross_kv_caches(False)
+    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_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(False, 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 branch it runs."""
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    pos_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    neg_ctx = st.get_kv_caches(True, 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.neg) == []
+
+
+def test_failed_forward_tears_down_session(monkeypatch):
+    """Review follow-up (hsliuustc0106): a forward that dies partway must not
+    leave allocated-but-uncommitted KV behind — the session is torn down (pool
+    blocks freed, model-local state dropped) and the error propagates."""
+    from types import SimpleNamespace
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)  # session owns pool blocks
+    free_before_failure = kv.manager.block_pool.get_num_free_blocks()
+
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.kv_cache = kv
+    runner._ar_diffusion_states = __import__("collections").OrderedDict({"s1": st})
+    runner._max_ar_diffusion_states = 4
+    runner.pipeline = SimpleNamespace(_states={"s1": object()}, _ar_diffusion_kv_state=None)
+    runner.device = None
+    runner._perf_e2e_times = []
+
+    def boom(self, req):
+        raise RuntimeError("layer 17 exploded")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
+    req = SimpleNamespace(
+        request_id="r0",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "s1"}),
+    )
+
+    with pytest.raises(RuntimeError, match="layer 17 exploded"):
+        ARDiffusionModelRunner.execute_model(runner, req)
+
+    assert "s1" not in runner._ar_diffusion_states
+    assert "s1" not in runner.pipeline._states
+    assert runner.pipeline._ar_diffusion_kv_state is None
+    # close() freed the session's resident blocks.
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before_failure

tests/ar_diffusion/test_slot_mapping.py

@@ -0,0 +1,76 @@
+# 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,
+)
+
+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/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -6,6 +6,7 @@ import pytest
 from vllm.v1.request import RequestStatus
 
 import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod
+import vllm_omni.model_executor.models.voxcpm2.scheduler as voxcpm2_scheduler_mod
 from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -51,14 +52,30 @@ class _MockRequest:
         return RequestStatus.is_finished(self.status)
 
 
-def _make_scheduler(*, enable_unified_decode_graph: bool = True) -> VoxCPM2OmniARAsyncScheduler:
+@pytest.fixture(autouse=True)
+def _mock_cuda_graph_platform(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: True)
+
+
+def _make_scheduler(
+    *,
+    enable_unified_decode_graph: bool | None = True,
+    deterministic_cfm_noise: bool = False,
+    runtime_config_on_model_config: bool = False,
+) -> VoxCPM2OmniARAsyncScheduler:
     sched = VoxCPM2OmniARAsyncScheduler.__new__(VoxCPM2OmniARAsyncScheduler)
-    runtime_config = SimpleNamespace(enable_unified_decode_graph=enable_unified_decode_graph)
-    sched.vllm_config = SimpleNamespace(
-        model_config=SimpleNamespace(
-            hf_config=SimpleNamespace(voxcpm2_runtime_config=runtime_config),
+    hf_config = SimpleNamespace()
+    model_config = SimpleNamespace(hf_config=hf_config)
+    if enable_unified_decode_graph is not None or deterministic_cfm_noise:
+        runtime_config = SimpleNamespace(
+            enable_unified_decode_graph=True if enable_unified_decode_graph is None else enable_unified_decode_graph,
+            deterministic_cfm_noise=deterministic_cfm_noise,
         )
-    )
+        if runtime_config_on_model_config:
+            model_config.voxcpm2_runtime_config = runtime_config
+        else:
+            hf_config.voxcpm2_runtime_config = runtime_config
+    sched.vllm_config = SimpleNamespace(model_config=model_config)
     return sched
 
 
@@ -86,6 +103,39 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None:
     assert not scheduler._should_defer_waiting_for_unified_decode_graph()
 
 
+def test_voxcpm2_unified_decode_graph_uses_model_runtime_defaults() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=None)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_reads_model_config_runtime_config() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=True, runtime_config_on_model_config=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_with_deterministic_noise() -> None:
+    scheduler = _make_scheduler(deterministic_cfm_noise=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: False)
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
 def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
     scheduler = _make_scheduler()
     scheduler.running = [_MockRequest("decode")]

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -19,6 +19,10 @@ class _FakeGroup:
         self.group_ranks = group_ranks
         self.parallel_mode = parallel_mode
         self.device_group = object()
+        self.device_communicator = kwargs.get("device_communicator")
+        reduce_scatter = kwargs.get("reduce_scatter")
+        if reduce_scatter is not None:
+            self.reduce_scatter = reduce_scatter
         self.ulysses_group = kwargs.get("ulysses_group")
         self.ring_group = kwargs.get("ring_group")
         self.local_group = next(group for group in group_ranks if local_rank in group)
@@ -55,6 +59,23 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         created_groups.append(group)
         return group
 
+    def fake_init_vllm_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        group_name,
+    ):
+        del backend
+        group = _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            f"vllm_{group_name}",
+            device_communicator=object(),
+            reduce_scatter=lambda tensor, **kwargs: tensor,
+        )
+        created_groups.append(group)
+        return group
+
     fake_world_group = SimpleNamespace(
         rank_in_group=local_rank,
         local_rank=local_rank,
@@ -68,6 +89,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
     monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
     for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
@@ -87,7 +109,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         backend="gloo",
     )
 
-    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP is not parallel_state._SP
     assert parallel_state.vllm_parallel_state._PCP.world_size == 2
     assert parallel_state._DP.world_size == 2
     assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP
@@ -96,6 +118,13 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     assert parallel_state.vllm_parallel_state._TP.world_size == 2
     assert parallel_state._PP.world_size == 2
 
+    assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None
+    assert parallel_state.vllm_parallel_state._DP.device_communicator is not None
+    assert parallel_state.vllm_parallel_state._EP.device_communicator is not None
+    assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter")
+    assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter")
+    assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter")
+
     assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2]
     assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24]
     assert parallel_state.vllm_parallel_state._EP.local_group == [
@@ -155,7 +184,9 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         ],
     ]
 
-    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
+    vllm_group_names = [group.parallel_mode for group in created_groups if group.parallel_mode.startswith("vllm_")]
+    assert vllm_group_names == ["vllm_pcp", "vllm_tp", "vllm_dp", "vllm_ep"]
+    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "vllm_ep"]
     assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
 
 

tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py

@@ -50,6 +50,48 @@ class TestSetForwardContextNumTokens:
         mock_get.assert_not_called()
 
 
+class TestSetForwardContextDPMetadata:
+    """Test diffusion MoE DP metadata setup for vLLM 0.24 all2all dispatch."""
+
+    def test_sets_dp_metadata_from_vllm_dp_group(self, mocker):
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = None
+        fake_group = mocker.MagicMock()
+        fake_group.world_size = 2
+        fake_group.cpu_group = object()
+
+        def _all_gather_object(output, obj, group):
+            assert group is fake_group.cpu_group
+            output[:] = [obj, obj + 1]
+
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mocker.patch.object(hunyuan_moe._vllm_ps, "_DP", fake_group, create=True)
+        mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object)
+
+        hunyuan_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is not None
+        assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025]
+
+    def test_keeps_existing_dp_metadata(self, mocker):
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
+
+        existing_metadata = object()
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = existing_metadata
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mock_all_gather = mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object")
+
+        hunyuan_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is existing_metadata
+        mock_all_gather.assert_not_called()
+
+
 class TestHunyuanFusedMoEPlatformDispatch:
     """Test platform dispatch via platform qualname hooks."""
 

tests/dreamzero/test_pipeline_state.py

@@ -4,7 +4,6 @@
 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
@@ -46,14 +45,11 @@ def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None:
     assert "session-b" not in pipeline._states
 
 
-def test_dreamzero_state_cache_access_requires_initialization() -> None:
+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()
 
-    with pytest.raises(RuntimeError, match="KV caches not initialized"):
-        state.get_kv_caches()
-
-    with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"):
-        state.get_crossattn_caches()
-
-    with pytest.raises(RuntimeError, match="create_kv_caches first"):
-        state.update_kv_cache(0, torch.empty(0))
+    for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"):
+        assert not hasattr(state, removed)

tests/e2e/offline_inference/test_magi_human_expansion.py

@@ -22,6 +22,7 @@ _OMNI_RUNNER_PARAM = (
     {
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
+        "vae_use_tiling": True,
     },
 )
 
@@ -32,6 +33,7 @@ _OMNI_RUNNER_PARAM_OFFLOAD = (
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
         "enable_cpu_offload": True,
+        "vae_use_tiling": True,
     },
 )
 
@@ -45,10 +47,6 @@ pytestmark = [
     ),
 ]
 
-_SKIP_ISSUE_3236 = pytest.mark.skip(
-    reason="https://github.com/vllm-project/vllm-omni/issues/3236",
-)
-
 
 def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     """Validate that the MP4 contains meaningful video and audio tracks."""
@@ -74,7 +72,6 @@ def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     container.close()
 
 
-@_SKIP_ISSUE_3236
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 def test_magi_human_e2e(omni_runner: OmniRunner):
     """End-to-end test for MagiHuman generating video and audio."""

tests/e2e/online_serving/test_higgs_audio_v3.py

@@ -127,7 +127,7 @@ class TestHiggsAudioV3OnlineHappyPath:
         windows stitch into a coherent PCM stream. The byte-count gate is the same as
         the sync paths; per-chunk audio content is verified offline against Whisper.
 
-        NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
+        NOTE: ``min_hnr_db=0.0`` sit below the typical speech-noise floor so the
         check still catches catastrophic codec failure (silence, white noise, sample
         scramble all give HNR << 0) while allowing for the sliding-window codec's
         slightly-noisier-than-sync output. The default 1.0 dB threshold has only

tests/entrypoints/openai_api/test_serving_speech.py

@@ -3344,6 +3344,27 @@ class TestTTSAsyncOffloading:
         qwen3_tts_server._build_tts_params.assert_called_once()
         qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
 
+    def test_prepare_speech_generation_qwen3_default_seed_sets_tts_local_seed(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """Deploy default seed should seed Qwen3 TTS residual MTP sampling."""
+        qwen3_tts_server.engine_client.default_sampling_params_list = [
+            SimpleNamespace(max_tokens=2048, seed=42, extra_args=None)
+        ]
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = mocker.MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+        request = OpenAICreateSpeechRequest(input="hello")
+
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        stage0_params = qwen3_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
+        assert stage0_params.seed == 42
+        assert stage0_params.extra_args["tts_local_seed"] == 42
+        assert qwen3_tts_server.engine_client.default_sampling_params_list[0].extra_args is None
+
     def test_prepare_speech_generation_uses_adapter_model_type_label(
         self,
         voxtral_server,
@@ -3429,6 +3450,48 @@ class TestTTSAsyncOffloading:
 
         mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True)
 
+    def test_prepare_speech_generation_no_async_chunk_stream_uses_final_only(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """Full-payload TTS streaming should not request delta multimodal outputs."""
+        qwen3_tts_server.engine_client.model_config.async_chunk = False
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = mocker.MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+        mock_coerce = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
+            return_value=qwen3_tts_server.engine_client.default_sampling_params_list,
+        )
+        request = OpenAICreateSpeechRequest(input="hello", stream_format="audio", response_format="pcm")
+
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, False)
+
+    def test_prepare_speech_generation_no_async_chunk_stream_keeps_delta_for_non_qwen3(
+        self, voxtral_server, mocker: MockerFixture
+    ):
+        """FINAL_ONLY streaming for async_chunk=False is scoped to qwen3_tts only."""
+        voxtral_server.engine_client.model_config.async_chunk = False
+        mocker.patch.object(voxtral_server._get_tts_adapter(), "validate", return_value=None)
+        voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock(
+            return_value={
+                "prompt_token_ids": [1, 2, 3],
+                "additional_information": {"voice": ["test"]},
+            }
+        )
+        mock_coerce = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
+            return_value=voxtral_server.engine_client.default_sampling_params_list,
+        )
+        request = OpenAICreateSpeechRequest(input="hello", voice="test", stream_format="audio", response_format="pcm")
+
+        asyncio.run(voxtral_server._prepare_speech_generation(request))
+
+        mock_coerce.assert_called_once_with(voxtral_server.engine_client.default_sampling_params_list, True)
+
     def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false(
         self, qwen3_tts_server, mocker: MockerFixture
     ):

tests/entrypoints/test_omni_sleep_mode.py

@@ -179,7 +179,7 @@ class TestOmniSleepMode:
         try:
             device_id = 0
             used_before = get_device_global_memory_used_gib(device_id)
-            acks = await llm_engine.sleep(stage_ids=[0], level=2)
+            acks = await llm_engine.sleep(stage_ids=[0], level=1)
             await asyncio.sleep(1.5)
             used_after = get_device_global_memory_used_gib(device_id)
             drop_gib = used_before - used_after
@@ -209,7 +209,7 @@ class TestOmniSleepMode:
         """Diffusion Worker stage signal loop"""
         try:
             logger.info("Starting Diffusion Worker Handshake Test")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
 
             def _get_status(ack):
                 return ack.status if hasattr(ack, "status") else ack.get("status")
@@ -264,9 +264,9 @@ class TestOmniSleepMode:
                 base_output = output
             assert base_output is not None and len(base_output.images) > 0
             logger.info("Baseline Generation successful.")
-            # Sleep Level 2
+            # Sleep Level 1
             logger.info("Entering Deep Sleep (VRAM Scavenging)...")
-            await diffusion_engine.sleep(stage_ids=[0], level=2)
+            await diffusion_engine.sleep(stage_ids=[0], level=1)
             # Wake-up
             logger.info("Waking up (Reloading Weights)...")
             await diffusion_engine.wake_up(stage_ids=[0])
@@ -353,8 +353,8 @@ class TestOmniSleepMode:
             logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB")
 
             # Sleep
-            logger.info("Triggering Level 2 Deep Sleep (Full Weight Offloading)...")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+            logger.info("Triggering Level 1 Deep Sleep (Partial Weight Offloading)...")
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
 
             reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks)
             reported_freed_gib = reported_freed_bytes / 1024**3
@@ -407,6 +407,47 @@ class TestOmniSleepMode:
             diffusion_engine.shutdown()
             await asyncio.sleep(1)
 
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_level2_sleep_wake_raises(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro A: wake_up() after sleep(level=2) must raise
+        NotImplementedError instead of silently producing corrupted output."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=2)
+            with pytest.raises(NotImplementedError, match="sleep\\(level=2\\)"):
+                await llm_engine.wake_up(stage_ids=[0])
+        finally:
+            llm_engine.shutdown()
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro B: generate() must be rejected if kv_cache
+        is still asleep after wake_up(tags=["weights"]), instead of crashing with
+        CUDA illegal memory access."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=1)
+            await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
+            with pytest.raises(RuntimeError, match="partially or fully asleep"):
+                async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
+                    pass
+        finally:
+            llm_engine.shutdown()
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a
+        safe no-op instead of raising a cumem CUDA invalid argument error."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=1)
+            first_acks = await llm_engine.wake_up(stage_ids=[0])
+            assert len(first_acks) > 0, "First wake_up() should return ACKs"
+            second_acks = await llm_engine.wake_up(stage_ids=[0])
+            assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}"
+        finally:
+            llm_engine.shutdown()
+
 
 # ---------------------------------------------------------------------------
 # H100 (or MI325) e2e: BAGEL / pure diffusion, multi-TP
@@ -444,7 +485,7 @@ async def test_diffusion_model_sleep_tp(tp_size: int):
         async for _ in engine.generate("test", sampling_params_list=[llm_sp, diff_sp]):
             pass
 
-        acks = await engine.sleep(level=2)
+        acks = await engine.sleep(level=1)
         statuses = [get_ack_info(ack, "status") for ack in acks]
         assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
 
@@ -500,7 +541,7 @@ async def test_multistage_sleep_h100(tp_size: int):
         async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]):
             pass
 
-        acks = await engine.sleep(stage_ids=[0, 1], level=2)
+        acks = await engine.sleep(stage_ids=[0, 1], level=1)
         assert len(acks) == 2
 
         await engine.wake_up(stage_ids=[0, 1])

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -339,6 +339,91 @@ class TestCodePredictorDtypeAlignment:
         assert not torch.equal(first[:, 1:], different[:, 1:])
 
 
+class TestCodePredictorPerRowGenerators:
+    """Seeded requests must stay deterministic inside a multi-row batch (#4883)."""
+
+    def _make_predictor(self, mocker: MockerFixture, loaded_target_classes):
+        _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+        cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
+        predictor = code_predictor_wrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+        predictor._wrapper_config.use_cuda_graphs = False
+        return predictor, talker_config
+
+    def test_multinomial_per_row_matches_single_row_draws(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        multinomial = common_mod.CodePredictorWrapper._multinomial
+        torch.manual_seed(7)
+        probs = torch.softmax(torch.randn(3, 64), dim=-1)
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=probs.device)
+            generator.manual_seed(seed)
+            return generator
+
+        batched = multinomial(probs, None, [seeded(11), None, seeded(22)])
+        assert batched.shape == (3, 1)
+        # Each seeded row must reproduce a standalone draw from the same seed,
+        # independent of what the other rows in the batch consume.
+        assert torch.equal(batched[0:1], torch.multinomial(probs[0:1], num_samples=1, generator=seeded(11)))
+        assert torch.equal(batched[2:3], torch.multinomial(probs[2:3], num_samples=1, generator=seeded(22)))
+
+    def test_forward_per_row_generators_are_row_independent(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        torch.manual_seed(123)
+        row_embed = torch.randn(1, hidden)
+        row_hidden = torch.randn(1, hidden)
+        # Three identical rows: two share a seed, one differs.
+        layer0_code = torch.zeros(3, dtype=torch.long)
+        layer0_embed = row_embed.expand(3, hidden).contiguous()
+        last_talker_hidden = row_hidden.expand(3, hidden).contiguous()
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=layer0_code.device)
+            generator.manual_seed(seed)
+            return generator
+
+        def run(generators):
+            return predictor(
+                layer0_code=layer0_code,
+                layer0_embed=layer0_embed,
+                last_talker_hidden=last_talker_hidden,
+                do_sample=True,
+                temperature=0.9,
+                top_k=50,
+                top_p=1.0,
+                generators=generators,
+            )
+
+        first = run([seeded(1234), seeded(1234), seeded(4321)])
+        second = run([seeded(1234), seeded(1234), seeded(4321)])
+
+        # Same per-row seeds -> the whole batch reproduces across calls.
+        assert torch.equal(first, second)
+        # Identical inputs + identical seeds -> identical rows within a call.
+        assert torch.equal(first[0], first[1])
+        # A different seed on the same inputs must diverge in the residual layers.
+        assert not torch.equal(first[0, 1:], first[2, 1:])
+
+    def test_forward_rejects_mismatched_generators_length(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        with pytest.raises(ValueError, match="one entry per row"):
+            predictor(
+                layer0_code=torch.zeros(2, dtype=torch.long),
+                layer0_embed=torch.randn(2, hidden),
+                last_talker_hidden=torch.randn(2, hidden),
+                generators=[torch.Generator()],
+            )
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 

tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py

@@ -12,6 +12,7 @@ REPO_ROOT = Path(__file__).resolve().parents[4]
 RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py"
 RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py"
 TALKER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "voxcpm2_talker.py"
+VOXCPM2_RUNTIME_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "runtime_config.py"
 VOXCPM2_PIPELINE = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "pipeline.py"
 VOXCPM2_SCHEDULER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "scheduler.py"
 OMNI_AR_SCHEDULER = REPO_ROOT / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py"
@@ -74,6 +75,7 @@ def test_ar_runner_without_model_hook_stays_on_normal_path():
 
 def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
     source = TALKER.read_text()
+    runtime_config_source = VOXCPM2_RUNTIME_CONFIG.read_text()
     compact_source = "".join(source.split())
 
     assert "_voxcpm2_compile_without_inductor_cudagraphs" in source
@@ -86,8 +88,8 @@ def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
     assert '"triton.cudagraphs": False' in compile_source
     assert '"triton.cudagraph_trees": False' in compile_source
 
-    assert "self._enable_unified_decode_graph=(use_cuda_graph" in compact_source
-    assert "andnotself._deterministic_cfm_noise" in compact_source
+    assert "self._enable_unified_decode_graph=self._runtime_config.unified_decode_graph_available(" in compact_source
+    assert "andnotself.deterministic_cfm_noise" in "".join(runtime_config_source.split())
     assert "decode_tail_graph" not in source
 
     unified_source = source[source.index("def _forward_unified_decode") :]
@@ -183,7 +185,9 @@ def test_voxcpm2_scheduler_policy_stays_model_local():
     assert "def _should_defer_waiting_admission(self) -> bool:" in common_source
 
     assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source
-    assert "voxcpm2_runtime_config" in voxcpm2_scheduler_source
+    assert "from .runtime_config import _VoxCPM2RuntimeConfig" in voxcpm2_scheduler_source
+    assert "_VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)" in voxcpm2_scheduler_source
+    assert "unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())" in voxcpm2_scheduler_source
     assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source
     assert "def schedule(" not in voxcpm2_scheduler_source
     assert "create_request_queue" not in voxcpm2_scheduler_source

tests/unit/higgs_audio_v3/test_higgs_audio_v3.py

@@ -7,6 +7,7 @@ without requiring the actual checkpoint or GPU.
 """
 
 import asyncio
+import json
 import time
 from collections import OrderedDict, defaultdict
 from types import SimpleNamespace
@@ -1309,7 +1310,10 @@ class TestVoiceCloneReferenceCache:
         parent = tmp_path / "OmniVoice"
         subdir = parent / "audio_tokenizer"
         subdir.mkdir(parents=True)
-        (subdir / "config.json").write_text("{}", encoding="utf-8")
+        (subdir / "config.json").write_text(
+            json.dumps({"model_type": "higgs_audio_v2_tokenizer"}),
+            encoding="utf-8",
+        )
 
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
@@ -1317,6 +1321,32 @@ class TestVoiceCloneReferenceCache:
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
 
+    def test_audio_tokenizer_dir_rejects_omnivoice_config(self, tmp_path, monkeypatch):
+        import huggingface_hub
+        import huggingface_hub.constants as hub_constants
+
+        from vllm_omni.model_executor.models.higgs_audio_v2 import higgs_audio_v2_tokenizer as tok
+
+        parent = tmp_path / "OmniVoice"
+        subdir = parent / "audio_tokenizer"
+        subdir.mkdir(parents=True)
+        (subdir / "config.json").write_text(
+            json.dumps({"model_type": "omnivoice"}),
+            encoding="utf-8",
+        )
+
+        monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
+        monkeypatch.setenv("HIGGS_AUDIO_V2_TOKENIZER_PATH", "")
+        monkeypatch.setattr(
+            huggingface_hub,
+            "try_to_load_from_cache",
+            lambda **_: str(subdir / "config.json"),
+        )
+        monkeypatch.setattr(hub_constants, "HF_HUB_CACHE", str(tmp_path / "empty_cache"))
+
+        assert tok._normalize_audio_tokenizer_dir(str(parent)) is None
+        assert tok._resolve_audio_tokenizer_dir() is None
+
     def test_higgs_v3_ref_code_cache_returns_clone(self):
         from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 

tests/worker/test_gpu_ar_model_runner.py

@@ -266,7 +266,6 @@ def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monke
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=["r2"],
-        seq_len=3,
         num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )
@@ -312,7 +311,6 @@ def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatc
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=6,
         num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long),
     )
@@ -366,7 +364,6 @@ def test_process_additional_information_uses_snapshot_request_order(monkeypatch)
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=3,
         num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )
@@ -427,7 +424,6 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=2,
         num_scheduled_tokens_np=np.array([2], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0], dtype=torch.long),
     )
@@ -439,6 +435,97 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
     assert output.multimodal_outputs is None
 
 
+def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_only(monkeypatch):
+    runner = _make_async_output_runner(engine_output_type="latent")
+    runner.model.omni_pooler_payload_include_hidden = False
+    runner._async_chunk = False
+    runner.requests = {"r1": object(), "r2": object(), "r3": object()}
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+
+    codes = torch.arange(48, dtype=torch.long).reshape(3, 16)
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=3,
+            num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
+        ),
+        hidden_states=torch.tensor([[1.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"codes": {"audio": codes}},
+        req_ids_output_copy=["r1", "r2", "r3"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
+        valid_sampled_token_ids=[[101], [102], [103]],
+        logprobs_lists=None,
+        prompt_logprobs_dict={},
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        ec_connector_output=None,
+        cudagraph_stats=None,
+        kv_extracted_req_ids=None,
+        num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
+    )
+
+    assert output.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1])
+    assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2])
+    assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3])
+
+
+def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monkeypatch):
+    """Thinker mm rows align to hidden_states.shape[0], not padded scheduled count."""
+    runner = _make_async_output_runner(engine_output_type="latent")
+    runner.model.omni_pooler_payload_include_hidden = False
+    runner._async_chunk = False
+    runner.requests = {"r1": object(), "r2": object(), "r3": object()}
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+
+    layers = torch.arange(96, dtype=torch.long).reshape(3, 32)
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=5,
+            num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
+        ),
+        hidden_states=torch.randn(3, 4),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"hidden_states": {"layers": {0: layers}}},
+        req_ids_output_copy=["r1", "r2", "r3"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
+        valid_sampled_token_ids=[[101], [102], [103]],
+        logprobs_lists=None,
+        prompt_logprobs_dict={},
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        ec_connector_output=None,
+        cudagraph_stats=None,
+        kv_extracted_req_ids=None,
+        num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
+    )
+
+    assert output.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["hidden_states.layer_0"], layers[0:1])
+    assert torch.equal(output.inter_stage_outputs[1]["hidden_states.layer_0"], layers[1:2])
+    assert torch.equal(output.inter_stage_outputs[2]["hidden_states.layer_0"], layers[2:3])
+
+
 def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
     runner = _make_async_output_runner()
     runner.model.omni_pooler_payload_include_hidden = False
@@ -627,7 +714,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=2,
         num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )

tests/worker/test_omni_gpu_model_runner.py

@@ -81,6 +81,7 @@ class CaptureTalkerMTP(torch.nn.Module):
         top_k=None,
         top_p=None,
         generator=None,
+        generators=None,
     ):
         self.calls.append(
             {
@@ -90,6 +91,7 @@ class CaptureTalkerMTP(torch.nn.Module):
                 "top_k": top_k,
                 "top_p": top_p,
                 "generator": generator,
+                "generators": generators,
             }
         )
         codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64)
@@ -306,6 +308,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
             "top_k": 9,
             "top_p": 0.55,
             "generator": runner.talker_mtp.calls[0]["generator"],
+            "generators": None,
         }
     ]
     assert runner.talker_mtp.calls[0]["generator"] is not None
@@ -350,6 +353,57 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds)
 
 
+def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypatch):
+    """Models with talker_mtp_accepts_per_row_generators get one batched call (#4883)."""
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.requests["r1"].sampling_params = SimpleNamespace(
+        seed=11,
+        extra_args={"tts_local_seed": 11},
+    )
+    runner.requests["r2"].sampling_params = SimpleNamespace(
+        seed=22,
+        extra_args={"tts_local_seed": 22},
+    )
+    runner.talker_mtp = CaptureTalkerMTP()
+    runner.model = SimpleNamespace(
+        talker_mtp_output_key=("codes", "audio"),
+        talker_mtp_accepts_per_row_generators=True,
+    )
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
+
+    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
+        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
+        return (False, batch_desc, None, None, None)
+
+    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
+
+    inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+
+    # One batched call with distinct per-row generators, not two scalar calls.
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    row_generators = runner.talker_mtp.calls[0]["generators"]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    assert len(row_generators) == 2
+    assert all(generator is not None for generator in row_generators)
+    assert row_generators[0] is not row_generators[1]
+
+    # The per-request generator stream persists across steps...
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+    assert runner.talker_mtp.calls[1]["generators"][0] is row_generators[0]
+    assert runner.talker_mtp.calls[1]["generators"][1] is row_generators[1]
+
+    # ...and is evicted once its request finishes.
+    del runner.requests["r2"]
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
+    assert set(runner._talker_mtp_generators) == {"r1"}
+    assert runner.talker_mtp.calls[2]["generator"] is row_generators[0]
+
+
 def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
     """Validate that _update_intermediate_buffer writes to model_intermediate_buffer
     (forward path) and mirrors to additional_information_cpu setattr (backward compat)."""

vllm_omni/deploy/dreamzero.yaml

@@ -14,6 +14,8 @@ stages:
     max_num_seqs: 1
     enforce_eager: false
     model_class_name: DreamZeroPipeline
+    # DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache).
+    engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine
     model_config:
       default_robot_embodiment: roboarena
       policy_server_config:

vllm_omni/deploy/qwen3_tts_high_concurrency.yaml

@@ -60,7 +60,6 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      seed: 42
       repetition_penalty: 1.05
     subtalker_sampling_params:
       do_sample: true
@@ -85,7 +84,6 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
-      seed: 42
       repetition_penalty: 1.0
 
 platforms:

vllm_omni/deploy/voxcpm2.yaml

@@ -28,20 +28,21 @@ stages:
     max_num_batched_tokens: 4096
     max_model_len: 4096
     devices: "0"
-    hf_overrides:
-      voxcpm2_runtime_config:
-        audio_emit_every: 1
-        vae_decode_every: 3
-        enable_batched_cfm: true
-        enable_batched_vae_decode: true
-        coalesce_audio_d2h: true
-        enable_loc_dit_fused_qkv: true
-        enable_loc_dit_fused_mlp: true
-        enable_loc_dit_zero_dt_cache: true
-        enable_loc_dit_skip_qkv_contig: true
-        enable_loc_dit_reduce_overhead_no_cg: true
-        enable_unified_decode_graph: true
-        unified_decode_graph_max_batch_size: 8
+    engine_extras:
+      hf_overrides:
+        voxcpm2_runtime_config:
+          audio_emit_every: 1
+          vae_decode_every: 3
+          enable_batched_cfm: true
+          enable_batched_vae_decode: true
+          coalesce_audio_d2h: true
+          enable_loc_dit_fused_qkv: true
+          enable_loc_dit_fused_mlp: true
+          enable_loc_dit_zero_dt_cache: true
+          enable_loc_dit_skip_qkv_contig: true
+          enable_loc_dit_reduce_overhead_no_cg: true
+          enable_unified_decode_graph: true
+          unified_decode_graph_max_batch_size: 8
     default_sampling_params:
       temperature: 0.0
       top_p: 1.0

vllm_omni/diffusion/data.py

@@ -6,6 +6,7 @@ import os
 import random
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field, fields
+from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 import diffusers
@@ -613,6 +614,20 @@ class OmniDiffusionConfig:
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None
 
+    # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class``
+    # (mirrors ``DiffusionExecutor.get_class``). Config files use a string:
+    # "default" -> DiffusionEngine, or an import path (set e.g. by a deploy
+    # config). Programmatic callers may pass a DiffusionEngine subclass
+    # directly — hence ``str | type`` (structured-config mirrors should expose
+    # the string form only).
+    engine_backend: str | type = "default"
+
+    # Optional override for the diffusion model runner class (import path).
+    # Precedence in the worker: this override > the runner declared by the
+    # selected engine class (``default_diffusion_model_runner_cls``) > the
+    # platform default. Never mutated by engines.
+    diffusion_model_runner_cls: str | None = None
+
     # HuggingFace specific parameters
     trust_remote_code: bool = False
     revision: str | None = None
@@ -1173,7 +1188,8 @@ class OmniDiffusionConfig:
         valid_fields = {f.name for f in fields(cls)}
         filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
 
-        return cls(**filtered_kwargs)
+        instance = cls(**filtered_kwargs)
+        return instance
 
 
 @dataclass
@@ -1473,5 +1489,12 @@ class OmniWakeTask:
     tags: list[str] | None = None
 
 
+class CuMemTag(str, Enum):
+    """Tags representing specific CuMem allocations for sleep/wake state tracking."""
+
+    WEIGHTS = "weights"
+    KV_CACHE = "kv_cache"
+
+
 # Special message broadcast via scheduler queues to signal worker shutdown.
 SHUTDOWN_MESSAGE = {"type": "shutdown"}

vllm_omni/diffusion/diffusion_engine.py

@@ -134,6 +134,14 @@ class _RpcTask:
 class DiffusionEngine:
     """The diffusion engine for vLLM-Omni diffusion models."""
 
+    #: Import path of the model runner this engine's workers should build, or
+    #: ``None`` for the platform default. Subclasses declare their runner here
+    #: (e.g. the AR-Diffusion engine), the worker resolves it through
+    #: :meth:`resolve_engine_class` — so engine->runner routing lives on the
+    #: engine class itself, and ``od_config.diffusion_model_runner_cls``
+    #: remains a pure explicit user override (never mutated by engines).
+    default_diffusion_model_runner_cls: str | None = None
+
     def __init__(
         self,
         od_config: OmniDiffusionConfig,
@@ -609,19 +617,55 @@ class DiffusionEngine:
             self._put_streaming_output_with_cv(request_id, out)
 
     @staticmethod
-    def make_engine(
-        config: OmniDiffusionConfig,
-        scheduler: SchedulerInterface | None = None,
-    ) -> DiffusionEngine:
-        """Factory method to create a DiffusionEngine instance.
+    def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
+        """Resolve the engine class selected by ``config.engine_backend``.
+
+        Mirrors ``DiffusionExecutor.get_class``: accepts ``"default"``, a
+        ``DiffusionEngine`` subclass, or an import-path string (e.g. a deploy
+        config's ``engine_backend``). Kept separate from :meth:`make_engine` so the
+        selection is testable without constructing an engine (which runs a dummy forward).
 
         Args:
             config: The configuration for the diffusion engine.
 
         Returns:
-            An instance of DiffusionEngine.
+            The ``DiffusionEngine`` (sub)class to instantiate.
         """
-        return DiffusionEngine(config, scheduler=scheduler)
+        backend = getattr(config, "engine_backend", "default") or "default"
+
+        if isinstance(backend, type):
+            if not issubclass(backend, DiffusionEngine):
+                raise TypeError(f"engine_backend must be a DiffusionEngine subclass. Got {backend}.")
+            return backend
+        if backend == "default":
+            return DiffusionEngine
+        if isinstance(backend, str):
+            try:
+                engine_class = resolve_obj_by_qualname(backend)
+            except (ImportError, ValueError) as e:
+                raise ValueError(
+                    f"Failed to load engine_backend '{backend}'. Ensure it is a valid python path. Error: {e}"
+                ) from e
+            if not issubclass(engine_class, DiffusionEngine):
+                raise TypeError(f"engine_backend must resolve to a DiffusionEngine subclass. Got {engine_class}.")
+            return engine_class
+        raise ValueError(f"Unknown engine_backend: {backend!r}")
+
+    @staticmethod
+    def make_engine(
+        config: OmniDiffusionConfig,
+        scheduler: SchedulerInterface | None = None,
+    ) -> DiffusionEngine:
+        """Factory method to create the engine selected by ``config.engine_backend``.
+
+        Args:
+            config: The configuration for the diffusion engine.
+
+        Returns:
+            An instance of the resolved ``DiffusionEngine`` (sub)class.
+        """
+        engine_class = DiffusionEngine.resolve_engine_class(config)
+        return engine_class(config, scheduler=scheduler)
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         with self._cv:

vllm_omni/diffusion/distributed/parallel_state.py

@@ -512,6 +512,21 @@ def init_model_parallel_group(
         )
 
 
+def init_vllm_model_parallel_group(
+    group_ranks: list[list[int]],
+    local_rank: int,
+    backend: str,
+    group_name: str,
+) -> vllm_parallel_state.GroupCoordinator:
+    return vllm_parallel_state.init_model_parallel_group(
+        group_ranks=group_ranks,
+        local_rank=local_rank,
+        backend=backend,
+        group_name=group_name,
+        use_device_communicator=True,
+    )
+
+
 def init_dit_group(
     dit_parallel_size: int,
     backend: str,
@@ -838,23 +853,40 @@ def initialize_model_parallel(
     if use_moe_parallel_mapping:
         # Diffusion normally uses its own SP group. Map it to vLLM PCP only for
         # expert-parallel runtimes that rely on vLLM FusedMoE group semantics.
-        vllm_parallel_state._PCP = _SP
+        # vLLM 0.24 MoE kernels require GroupCoordinator.device_communicator
+        # and reduce_scatter(), which the diffusion SP coordinator intentionally
+        # does not own. Keep the rank layout but build a vLLM coordinator.
+        vllm_parallel_state._PCP = init_vllm_model_parallel_group(
+            group_ranks=sp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            group_name="pcp",
+        )
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
-    vllm_parallel_state._TP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("tp"),
-        local_rank=get_world_group().local_rank,
-        backend=backend,
-        parallel_mode="tensor",
-    )
-    if use_moe_parallel_mapping and cfg_parallel_size > 1:
+    tp_group_ranks = rank_generator.get_ranks("tp")
+    if use_moe_parallel_mapping:
+        vllm_parallel_state._TP = init_vllm_model_parallel_group(
+            group_ranks=tp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            group_name="tp",
+        )
+    else:
+        vllm_parallel_state._TP = init_model_parallel_group(
+            group_ranks=tp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            parallel_mode="tensor",
+        )
+    if use_moe_parallel_mapping:
         # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
         # only when constructing the vLLM EP layout for expert-parallel paths.
-        vllm_parallel_state._DP = init_model_parallel_group(
+        vllm_parallel_state._DP = init_vllm_model_parallel_group(
             group_ranks=rank_generator.get_ranks("cfg-dp"),
             local_rank=get_world_group().local_rank,
             backend=backend,
-            parallel_mode="data",
+            group_name="dp",
         )
 
     global _FS, _EXPERT_PARALLEL_GROUP_RANKS
@@ -869,11 +901,11 @@ def initialize_model_parallel(
     _EXPERT_PARALLEL_GROUP_RANKS = None
     if use_moe_parallel_mapping:
         ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp")
-        vllm_parallel_state._EP = init_model_parallel_group(
+        vllm_parallel_state._EP = init_vllm_model_parallel_group(
             group_ranks=ep_group_ranks,
             local_rank=get_world_group().local_rank,
             backend=backend,
-            parallel_mode="expert",
+            group_name="ep",
         )
         _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks
 

vllm_omni/diffusion/models/dreamzero/causal_wan_model.py

@@ -37,6 +37,20 @@ from vllm_omni.diffusion.models.dreamzero.action_encoder import (
     MultiEmbodimentActionEncoder,
 )
 
+# AR-Diffusion paged self-attention (in-tree experimental engine). Import at
+# module level so the isinstance check + custom-op call trace cleanly inside
+# the fullgraph-compiled DiT block (an import inside the traced region would
+# graph-break). The model still works without the engine: the payload type is
+# only ever constructed by the AR-Diffusion runner.
+try:
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+        ARDiffusionPagedLayerInputs,
+        paged_write_attn,
+    )
+except ImportError:  # pragma: no cover - experimental package always ships in-tree
+    ARDiffusionPagedLayerInputs = None
+    paged_write_attn = None
+
 # ── RoPE utilities ──────────────────────────────────────────────────
 
 
@@ -510,7 +524,7 @@ class CausalWanSelfAttention(nn.Module):
         freqs_action: torch.Tensor,
         freqs_state: torch.Tensor,
         action_register_length: int | None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -567,24 +581,44 @@ class CausalWanSelfAttention(nn.Module):
             action_v = v[:, -action_register_length:]
             v = v[:, :-action_register_length]
 
-        updated_k = kv_cache[0]
-        updated_v = kv_cache[1]
-        new_k = torch.cat([updated_k, roped_key], dim=1)
-        new_v = torch.cat([updated_v, v], dim=1)
-        new_k = new_k[:, -self.max_attention_size :]
-        new_v = new_v[:, -self.max_attention_size :]
-
-        if action_register_length is not None:
-            q_cat = torch.cat([roped_query, roped_action_query], dim=1)
-            k_cat = torch.cat([new_k, roped_action_key], dim=1)
-            v_cat = torch.cat([new_v, action_v], dim=1)
+        if ARDiffusionPagedLayerInputs is not None and isinstance(kv_cache, ARDiffusionPagedLayerInputs):
+            # Fused write+attend custom op: one opaque node in the compiled
+            # graph (slot writes + FlashAttention block-table kernel inside).
+            # Metadata tensors were prepared once per forward in _forward_blocks.
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_act, v_act = roped_action_key[0], action_v[0]
+            else:
+                q_cat = roped_query
+                k_act = v_act = None
+            x = paged_write_attn(
+                kv_cache,
+                q_cat[0],
+                roped_key[0],
+                v[0],
+                k_act,
+                v_act,
+                self.head_dim**-0.5,
+            ).unsqueeze(0)
         else:
-            q_cat = roped_query
-            k_cat = new_k
-            v_cat = new_v
+            updated_k = kv_cache[0]
+            updated_v = kv_cache[1]
+            new_k = torch.cat([updated_k, roped_key], dim=1)
+            new_v = torch.cat([updated_v, v], dim=1)
+            new_k = new_k[:, -self.max_attention_size :]
+            new_v = new_v[:, -self.max_attention_size :]
 
-        x = self.attn(q_cat, k_cat, v_cat)
-        updated_kv_cache = torch.stack([new_k, new_v], dim=0)
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_cat = torch.cat([new_k, roped_action_key], dim=1)
+                v_cat = torch.cat([new_v, action_v], dim=1)
+            else:
+                q_cat = roped_query
+                k_cat = new_k
+                v_cat = new_v
+
+            x = self.attn(q_cat, k_cat, v_cat)
+            updated_kv_cache = torch.stack([new_k, new_v], dim=0)
 
         x = x.flatten(2)
         x = self.o(x)
@@ -646,7 +680,7 @@ class CausalWanAttentionBlock(nn.Module):
         freqs_state: torch.Tensor,
         context: torch.Tensor,
         action_register_length: int | None = None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         crossattn_cache: dict | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
@@ -925,10 +959,10 @@ class CausalWanModel(nn.Module):
         action: torch.Tensor | None,
         timestep_action: torch.Tensor | None,
         state: torch.Tensor | None,
-        kv_cache: list[torch.Tensor],
-        crossattn_cache: list[dict] | None,
+        kv_cache: list[Any],
         current_start_frame: int,
-    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+        crossattn_cache: list[dict] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor | None]]:
         x = x.flatten(start_dim=2).transpose(1, 2)
         B = x.shape[0]
         F_t = timestep.shape[1]
@@ -966,7 +1000,28 @@ class CausalWanModel(nn.Module):
             clip_embedding = self.img_emb(clip_feature)
             context = torch.cat([clip_embedding, context], dim=1)
 
-        updated_kv_caches: list[torch.Tensor] = []
+        # AR-Diffusion paged path: host-side prep ONCE per branch forward, outside
+        # the compiled blocks — allocate current video/action slots, build the
+        # padded block-table metadata all 40 layers share, and hand the compiled
+        # region plain tensors (ARDiffusionPagedLayerInputs) instead of a Python
+        # context object. Lazy-allocation contract preserved: only the branch this
+        # CFG-parallel rank executes reaches its _forward_blocks.
+        if kv_cache and getattr(kv_cache[0], "is_ar_diffusion_paged_context", False):
+            if B != 1:
+                raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1")
+            fctx = kv_cache[0].forward_ctx
+            if seq_len != fctx.seq_len:
+                raise RuntimeError(
+                    f"AR-Diffusion paged context seq_len={fctx.seq_len} but current video KV has {seq_len} tokens"
+                )
+            fctx.prepare(
+                device=x.device,
+                action_len=int(action_register_length or 0),
+                query_len=int(x.shape[1]),
+            )
+            kv_cache = [c.to_layer_inputs() for c in kv_cache]
+
+        updated_kv_caches: list[torch.Tensor | None] = []
         for block_index, block in enumerate(self.blocks):
             x, updated_kv_cache = block(
                 x=x,
@@ -1000,9 +1055,9 @@ class CausalWanModel(nn.Module):
         timestep: torch.Tensor,
         context: torch.Tensor,
         seq_len: int,
-        kv_cache: list[torch.Tensor],
-        crossattn_cache: list[torch.Tensor],
+        kv_cache: list[Any],
         current_start_frame: int,
+        crossattn_cache: list[dict] | None = None,
         y: torch.Tensor | None = None,
         clip_feature: torch.Tensor | None = None,
         action: torch.Tensor | None = None,

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -3,7 +3,7 @@
 
 """DreamZero pipeline for vllm-omni.
 
-Entry point for DiffusionEngine.step() → pipeline.forward(req)
+Entry point for DiffusionEngine.step() -> pipeline.forward(req)
 """
 
 from __future__ import annotations
@@ -93,17 +93,102 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     Multi-output: predict_noise() returns (video_pred, action_pred).
     CFG: video gets standard CFG, action takes positive branch only.
-    State: DreamZeroState manages KV cache + frame buffer across forward() calls.
+
+    KV is managed by the AR-Diffusion engine: ``self._ar_diffusion_kv_state`` is set by the
+    runner before ``forward()`` and the pipeline routes all KV access (get / update /
+    commit / reset) through the pool-backed state. Purely duck-typed (no engine import).
     """
 
+    _ar_diffusion_kv_state = None  # set by the runner before each forward
+
+    def _kv_get(self, state, is_negative, seq_len=None, update_kv_cache=False):
+        return self._ar_diffusion_kv_state.get_kv_caches(
+            is_negative,
+            seq_len=seq_len,
+            commit_current=update_kv_cache,
+        )
+
+    def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
+        # The engine owns all KV allocation: self-attn is allocated lazily from
+        # paged contexts, and cross-attn is populated eagerly in _kv_populate_cross.
+        # Nothing is created model-side.
+        return
+
+    def _kv_commit(self, is_negative: bool):
+        self._ar_diffusion_kv_state.commit_paged_context(is_negative)
+
+    def _kv_get_cross(self, state, is_negative):
+        """Cross-attn cache from the engine pool (text k/v + I2V image k_img/v_img)."""
+        return self._ar_diffusion_kv_state.get_cross_kv_caches(is_negative)
+
+    def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None:
+        """Eagerly project cross-attn K/V for all layers into the AR-Diffusion pool.
+
+        Caches the session-invariant cross-attn projections once, per half: the text
+        ``k``/``v`` from ``text_embedding(context)`` survive window-boundary resets
+        (prompt unchanged within a session — only session resets clear them), while
+        the I2V image-token ``k_img``/``v_img`` from ``img_emb(clip_feature)`` (the
+        257 image tokens the forward splits off, cached model-side by #4154) are
+        re-projected on every window restart from the fresh CLIP features. Must run
+        after the image is encoded so ``clip_feature`` is available.
+        """
+        s = self._ar_diffusion_kv_state
+        need_text = not s._cross_text_populated.get(is_negative, False)
+        need_img = (
+            clip_feature is not None
+            and getattr(self.transformer, "model_type", "t2v") == "i2v"
+            and not s._cross_img_populated.get(is_negative, False)
+        )
+        if not need_text and not need_img:
+            return
+        projected = self.transformer.text_embedding(context) if need_text else None
+        img_ctx = self.transformer.img_emb(clip_feature) if need_img else None
+        for i, block in enumerate(self.transformer.blocks):
+            ca = block.cross_attn
+            n, d = ca.tp_num_heads, ca.head_dim
+            k = v = None
+            if projected is not None:
+                k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
+                v = ca.v(projected).unflatten(2, (n, d))
+            k_img = v_img = None
+            if img_ctx is not None:
+                k_img = ca.norm_k_img(ca.k_img(img_ctx)).unflatten(2, (n, d))
+                v_img = ca.v_img(img_ctx).unflatten(2, (n, d))
+            s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img)
+        if need_text:
+            s._cross_text_populated[is_negative] = True
+        if need_img:
+            s._cross_img_populated[is_negative] = True
+        logger.info(
+            "AR-Diffusion CROSS POPULATE [%s]: %d layers, text=%s img=%s",
+            "neg" if is_negative else "pos",
+            len(self.transformer.blocks),
+            "kept" if projected is None else tuple(context.shape),
+            None if img_ctx is None else tuple(img_ctx.shape),
+        )
+
+    def _kv_reset(self, state, *, clear_video_latents: bool = True):
+        """Reset the engine's pooled session window plus the model's non-KV state.
+
+        DreamZero resets at the attention-window boundary; the engine pool drops the
+        same window so the next forward starts fresh. ``clear_video_latents=False``
+        keeps the accumulated video latents for export.
+
+        ``clear_video_latents=False`` also marks a window ("inference") reset: the
+        prompt is unchanged, so the pool keeps the text cross-attn K/V and only the
+        image half repopulates on the restart forward.
+        """
+        state.reset(clear_video_latents=clear_video_latents)
+        self._ar_diffusion_kv_state.reset(keep_cross_text=not clear_video_latents)
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.
 
         DreamZero root checkpoint layout (GEAR-Dreams/DreamZero-DROID):
-          config.json                     — root config (action_head_cfg, architectures, etc.)
-          model-*.safetensors             — all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
-          experiment_cfg/metadata.json    — per-embodiment action normalization stats
-          vae/                            — symlink to Wan2.1 VAE (diffusers-compatible)
+          config.json                     -- root config (action_head_cfg, architectures, etc.)
+          model-*.safetensors             -- all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
+          experiment_cfg/metadata.json    -- per-embodiment action normalization stats
+          vae/                            -- symlink to Wan2.1 VAE (diffusers-compatible)
 
         Components are instantiated from config (not from_pretrained), then filled
         by load_weights() which reads root safetensors and remaps key prefixes.
@@ -118,6 +203,18 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         """
         super().__init__()
 
+        # DreamZero is engine-only: every KV access in forward() routes through
+        # the AR-Diffusion engine's pool-backed state. Fail fast here — a stale
+        # or programmatic config that leaves engine_backend="default" would
+        # otherwise only crash mid-forward on the first KV access.
+        engine_backend = str(getattr(od_config, "engine_backend", "") or "")
+        if "ar_diffusion" not in engine_backend.lower().replace("-", "_"):
+            raise ValueError(
+                "DreamZeroPipeline requires the AR-Diffusion engine; set "
+                "engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine "
+                f"in the deploy config (got engine_backend={engine_backend!r})."
+            )
+
         model_path = od_config.model
         model_config = od_config.model_config
         local_files_only = os.path.exists(model_path)
@@ -228,8 +325,11 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         self.max_action_dim: int = ah_config["max_action_dim"]
 
         self.negative_prompt: str = model_config.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT)
+        # The negative prompt is a model constant: encode it once, lazily, and
+        # reuse across every forward/session (UMT5 encode is deterministic).
+        self._negative_prompt_embeds_cache: torch.Tensor | None = None
 
-        # Embodiment name → numeric ID mapping (model knowledge)
+        # Embodiment name -> numeric ID mapping (model knowledge)
         self.embodiment_name_to_id: dict[str, int] = model_config.get(
             "embodiment_name_to_id",
             DEFAULT_EMBODIMENT_NAME_TO_ID,
@@ -344,11 +444,17 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             state=kwargs.get("state_features"),
             embodiment_id=kwargs.get("embodiment_id"),
         )
-        if kwargs.get("update_kv_cache", False) and updated_kv_caches:
-            state = kwargs.get("dreamzero_state", self.state)
+        if kwargs.get("update_kv_cache", False):
             is_neg = kwargs.get("is_negative", False)
-            for i, kv in enumerate(updated_kv_caches):
-                state.update_kv_cache(i, kv, is_negative=is_neg)
+            logger.debug(
+                "AR-Diffusion pipeline predict_noise -> commit paged context: "
+                "is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
+                is_neg,
+                kwargs.get("seq_len"),
+                kwargs.get("current_start_frame"),
+                len(updated_kv_caches),
+            )
+            self._kv_commit(is_neg)
 
         return video_pred, action_pred
 
@@ -362,6 +468,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False``
         on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and
         compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop).
+        Incremental VAE encode (``_vae_encode_encoder_chunk``) stays eager because
+        Wan ``feat_cache`` mutation is incompatible with CUDAGraph capture.
         DiT blocks use per-block ``fullgraph=True``.
         """
         if not torch.cuda.is_available():
@@ -376,6 +484,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
         # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
+        # The AR-Diffusion paged self-attention is a registered custom op, so the
+        # block stays fullgraph even on that path.
         dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
 
         logger.info(
@@ -396,9 +506,9 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
 
         try:
-            self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
+            self.vae._encode = torch.compile(self.vae._encode, **compile_ro)
         except Exception as exc:
-            logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
+            logger.warning("DreamZero: vae._encode compile failed (%s); skipping.", exc)
 
         compiled_blocks = 0
         for block in self.transformer.blocks:
@@ -437,10 +547,24 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
             try:
                 image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
-                self._encode_image(image, self.num_frames, 180, 320)
+                self._encode_image(image, self.num_frames, 180, 320, state=self.state)
             except Exception as exc:
                 logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
 
+            try:
+                self.state.reset_vae_encoder_stream()
+                dummy_video = torch.zeros(1, 3, 1, 180, 320, dtype=torch.bfloat16, device=device)
+                self._vae_stream_seed(self.state, dummy_video[:, :, :1])
+                for _ in range(4):
+                    self._vae_stream_append_frame(self.state, dummy_video[:, :, :1])
+                self._vae_stream_get_observation_latents(
+                    self.state,
+                    self.num_frame_per_block,
+                    dtype=torch.bfloat16,
+                )
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (vae encode stream) skipped: %s", exc)
+
             try:
                 latent_h, latent_w = 180 // 8, 320 // 8
                 dummy_latent = torch.zeros(
@@ -499,7 +623,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     # -----------------------------------------------------------------------
 
     def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor:
-        """uint8 [B,T,H,W,C] → bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
+        """uint8 [B,T,H,W,C] -> bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
         videos = videos.permute(0, 4, 1, 2, 3)
         if videos.dtype == torch.uint8:
             videos = videos.float() / 255.0
@@ -546,6 +670,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         num_frames: int,
         height: int,
         width: int,
+        *,
+        state: DreamZeroState | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Encode first frame via CLIP + VAE.
         Returns: (clip_feas, ys, image_latent)
@@ -581,11 +707,147 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             new_image = y[:, :, 0:1]
             y = torch.concat([msk, y], dim=1)
 
+            if state is not None:
+                if not state.vae_stream_initialized:
+                    # Seed the AR streaming encoder after the compiled full encode above;
+                    # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph.
+                    self._vae_stream_seed(state, image_input[:, :, :1])
+                else:
+                    # Window ("inference") restart with a live stream: keep the real
+                    # frame history in the Wan feat_cache and append the restart
+                    # observation instead of reseeding from scratch.
+                    self._cudagraph_mark_step_begin()
+                    self._vae_stream_append_frame(state, image_input[:, :, :1])
+
         return clip_context, y, new_image
 
+    def _vae_patchify(self, videos: torch.Tensor) -> torch.Tensor:
+        if self.vae.config.patch_size is not None:
+            from diffusers.models.autoencoders.autoencoder_kl_wan import patchify
+
+            return patchify(videos, patch_size=self.vae.config.patch_size)
+        return videos
+
+    @staticmethod
+    def _vae_clone_feat_map(feat_map: list[torch.Tensor | None]) -> list[torch.Tensor | None]:
+        return [entry.clone() if isinstance(entry, torch.Tensor) else entry for entry in feat_map]
+
+    def _vae_init_enc_feat_map(self) -> list[torch.Tensor | None]:
+        self.vae.clear_cache()
+        return self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_encode_encoder_chunk(
+        self,
+        chunk: torch.Tensor,
+        feat_map: list[torch.Tensor | None],
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """Run one Wan encoder chunk while mutating ``feat_map`` in place.
+
+        Must stay eager (not ``torch.compile``): Wan causal ``feat_cache`` updates
+        conflict with ``reduce-overhead`` CUDAGraph capture.
+        """
+        self.vae._enc_feat_map = feat_map
+        self.vae._enc_conv_idx = [0]
+        out = self.vae.encoder(
+            chunk,
+            feat_cache=self.vae._enc_feat_map,
+            feat_idx=self.vae._enc_conv_idx,
+        )
+        return out, self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_quantize_encoder_out(self, encoder_out: torch.Tensor) -> torch.Tensor:
+        enc = self.vae.quant_conv(encoder_out)
+        mu, _ = enc.chunk(2, dim=1)
+        mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
+        inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype)
+        return (mu - mean) * inv_std
+
+    def _vae_stream_seed(self, state: DreamZeroState, first_frame: torch.Tensor) -> None:
+        """Seed incremental VAE encode with the first observation frame."""
+        state.reset_vae_encoder_stream()
+        feat_map = self._vae_init_enc_feat_map()
+        chunk = self._vae_patchify(first_frame.to(dtype=self.vae.dtype))
+        self._cudagraph_mark_step_begin()
+        encoder_out, feat_map = self._vae_encode_encoder_chunk(chunk[:, :, :1], feat_map)
+        state.vae_enc_feat_map = feat_map
+        state.vae_encoder_out = encoder_out
+        state.vae_pending_body_frames = None
+        state.vae_stream_initialized = True
+
+    def _vae_stream_append_frame(self, state: DreamZeroState, new_frame: torch.Tensor) -> None:
+        """Append one pixel frame and encode a 4-frame body chunk when ready."""
+        if not state.vae_stream_initialized or state.vae_enc_feat_map is None or state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream is not initialized.")
+
+        frame = new_frame.to(dtype=self.vae.dtype)
+        if state.vae_pending_body_frames is None:
+            state.vae_pending_body_frames = frame
+        else:
+            state.vae_pending_body_frames = torch.cat([state.vae_pending_body_frames, frame], dim=2)
+
+        while state.vae_pending_body_frames is not None and state.vae_pending_body_frames.shape[2] >= 4:
+            body = state.vae_pending_body_frames[:, :, :4]
+            if state.vae_pending_body_frames.shape[2] > 4:
+                state.vae_pending_body_frames = state.vae_pending_body_frames[:, :, 4:]
+            else:
+                state.vae_pending_body_frames = None
+            chunk = self._vae_patchify(body)
+            encoder_chunk, feat_map = self._vae_encode_encoder_chunk(chunk, state.vae_enc_feat_map)
+            state.vae_encoder_out = torch.cat([state.vae_encoder_out, encoder_chunk], dim=2)
+            state.vae_enc_feat_map = feat_map
+
+    def _vae_stream_get_observation_latents(
+        self,
+        state: DreamZeroState,
+        num_latent_frames: int,
+        *,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream has no accumulated encoder output.")
+        latents = self._vae_quantize_encoder_out(state.vae_encoder_out).to(dtype=dtype)
+        if latents.shape[2] >= num_latent_frames:
+            return latents[:, :, -num_latent_frames:]
+        pad_count = num_latent_frames - latents.shape[2]
+        pad = latents[:, :, -1:].expand(-1, -1, pad_count, -1, -1)
+        return torch.cat([pad, latents], dim=2)
+
+    def _preprocess_vae_observation_window(self, videos: torch.Tensor) -> torch.Tensor:
+        _, _, num_frames_raw, _, _ = videos.shape
+        if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
+            return videos
+        if num_frames_raw // 4 != self.num_frame_per_block:
+            repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
+            videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
+            first_frame = videos[:, :, 0:1]
+            return torch.cat([first_frame, videos], dim=2)
+        first_frame = videos[:, :, 0:1]
+        return torch.cat([first_frame, videos], dim=2)
+
+    def _encode_observation_latents(
+        self,
+        state: DreamZeroState,
+        videos: torch.Tensor,
+        *,
+        latent_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Encode current robot observation into normalized VAE latents."""
+        if state.vae_stream_initialized:
+            self._cudagraph_mark_step_begin()
+            self._vae_stream_append_frame(state, videos[:, :, -1:])
+            return self._vae_stream_get_observation_latents(
+                state,
+                self.num_frame_per_block,
+                dtype=latent_dtype,
+            )
+
+        videos = self._preprocess_vae_observation_window(videos)
+        return self._encode_vae_latents(videos).to(dtype=latent_dtype)
+
     def _encode_vae_latents(self, videos: torch.Tensor) -> torch.Tensor:
         """Encode videos into normalized VAE latents."""
         input_dtype = videos.dtype
+        self._cudagraph_mark_step_begin()
         hidden = self.vae._encode(videos.to(dtype=self.vae.dtype))
         mu, _ = hidden.chunk(2, dim=1)
         mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
@@ -635,7 +897,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     ) -> None:
         """Prefill KV cache with first frame and/or current observation.
 
-        Uses predict_noise_maybe_with_cfg() for CFG parallel — same path as
+        Uses predict_noise_maybe_with_cfg() for CFG parallel -- same path as
         the denoise loop. The mixin handles rank dispatch automatically.
         KV cache update happens as a side effect inside predict_noise().
         """
@@ -646,7 +908,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         head_dim = self.transformer.dim // self.transformer.num_heads
 
         if state.current_start_frame == 0:
-            state.create_kv_caches(
+            self._kv_create(
+                state,
                 batch_size,
                 dtype,
                 device,
@@ -671,16 +934,16 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                kv_cache=self._kv_get(state, False, seq_len=frame_seqlen, update_kv_cache=True),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    kv_cache=self._kv_get(state, True, seq_len=frame_seqlen, update_kv_cache=True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -721,16 +984,16 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=True),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -763,8 +1026,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         For each timestep:
           1. Build positive_kwargs / negative_kwargs
-          2. predict_noise_maybe_with_cfg()    → (video_pred, action_pred)
-          3. scheduler_step_maybe_with_cfg()   → VideoActionScheduler
+          2. predict_noise_maybe_with_cfg()    -> (video_pred, action_pred)
+          3. scheduler_step_maybe_with_cfg()   -> VideoActionScheduler
           4. _synchronize_cfg_parallel_step_output()
         """
         seq_len = kwargs["seq_len"]
@@ -823,8 +1086,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                     hidden_states=noisy_input.transpose(1, 2),
                     timestep_video=timestep,
                     encoder_hidden_states=prompt_embeds,
-                    kv_cache=state.get_kv_caches(False),
-                    crossattn_cache=state.get_crossattn_caches(False),
+                    kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=False),
+                    crossattn_cache=self._kv_get_cross(state, False),
                     y=y,
                     clip_feature=state.clip_feas,
                     action=noisy_input_action,
@@ -838,8 +1101,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                         hidden_states=noisy_input.transpose(1, 2),
                         timestep_video=timestep,
                         encoder_hidden_states=negative_prompt_embeds,
-                        kv_cache=state.get_kv_caches(True),
-                        crossattn_cache=state.get_crossattn_caches(True),
+                        kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=False),
+                        crossattn_cache=self._kv_get_cross(state, True),
                         y=y,
                         clip_feature=state.clip_feas,
                         action=noisy_input_action,
@@ -936,7 +1199,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             device=device,
         )
 
-        # State: raw from transform → pad to (B, state_horizon=1, max_state_dim)
+        # State: raw from transform -> pad to (B, state_horizon=1, max_state_dim)
         raw_state = unified_obs["state"]
         state_for_postprocess = None
         if raw_state is not None:
@@ -976,38 +1239,66 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # Explicit reset from OpenPI serving is carried by `extra_args["reset"]`
         # on the next inference request after websocket reset/session switch.
         if extra_args.get("reset", False):
-            state.reset()
+            self._kv_reset(state)
         else:
+            # Auto-reset based on model state (before accumulation). Both a "session"
+            # reset and an "inference" (window-boundary) reset clear the model-local
+            # KV window, so route both through _kv_reset to drop the AR-Diffusion pool window
+            # too and stay parity-exact; "inference" keeps the accumulated video
+            # latents (reset_inference_state == reset(clear_video_latents=False)).
             reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
             if reset_reason == "session":
-                state.reset()
+                self._kv_reset(state)
             elif reset_reason == "inference":
-                state.reset_inference_state()
+                self._kv_reset(state, clear_video_latents=False)
         state.language = text_tokens
 
-        # Frame accumulation: stitched single frame → multi-frame video
+        # Frame accumulation: stitched single frame -> multi-frame video
         video_frames = state.accumulate_frames(stitched)  # (T, H, W, C)
         videos = torch.from_numpy(video_frames).unsqueeze(0).to(device)  # (B=1, T, H, W, C)
 
-        videos = self._preprocess_video(videos)  # → [B,C,T,H,W] bf16
+        videos = self._preprocess_video(videos)  # -> [B,C,T,H,W] bf16
         _, _, num_frames_raw, height, width = videos.shape
 
-        prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        # Optional phase timing (DZ_PHASE_TIMING=1): logs per-forward stage costs
+        # (text encode / obs VAE encode / KV prefill / denoise) at INFO. Each mark
+        # synchronizes CUDA, so leave it off for timed benchmark runs.
+        _pt = None
+        if os.environ.get("DZ_PHASE_TIMING"):
+            import time as _time
+
+            torch.accelerator.synchronize()
+            _pt = {"time": _time, "t0": _time.perf_counter(), "marks": []}
+
+        def _pt_mark(name: str) -> None:
+            if _pt is not None:
+                torch.accelerator.synchronize()
+                _pt["marks"].append((name, _pt["time"].perf_counter()))
+
+        # Prompt embeds are constant within a session (a prompt change triggers a
+        # "session" reset above, which clears this cache alongside state.language).
+        if state.prompt_embeds is None:
+            state.prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        prompt_embeds = state.prompt_embeds
         # Negative prompt for CFG uncond branch (model constant)
         negative_prompt_embeds = None
         if self.cfg_scale > 1.0:
-            neg_inputs = self.tokenizer(
-                self.negative_prompt,
-                max_length=512,
-                padding="max_length",
-                truncation=True,
-                return_tensors="pt",
-                add_special_tokens=True,
-            )
-            negative_prompt_embeds = self._encode_text(
-                neg_inputs["input_ids"].to(device),
-                neg_inputs["attention_mask"].to(device),
-            )
+            if self._negative_prompt_embeds_cache is None:
+                neg_inputs = self.tokenizer(
+                    self.negative_prompt,
+                    max_length=512,
+                    padding="max_length",
+                    truncation=True,
+                    return_tensors="pt",
+                    add_special_tokens=True,
+                )
+                self._negative_prompt_embeds_cache = self._encode_text(
+                    neg_inputs["input_ids"].to(device),
+                    neg_inputs["attention_mask"].to(device),
+                )
+            negative_prompt_embeds = self._negative_prompt_embeds_cache
+
+        _pt_mark("text_encode")
 
         # Extract first/last frame for CLIP + VAE encoding
         if num_frames_raw == 4 or num_frames_raw == 9:
@@ -1021,27 +1312,24 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 self.num_frames,
                 height,
                 width,
+                state=state,
             )
             state.clip_feas = clip_feas.to(dtype=image.dtype)
             state.ys = ys.to(dtype=image.dtype)
 
-        if state.current_start_frame != 0:
-            # Subsequent calls: encode current observation via VAE
-            if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
-                pass
-            elif num_frames_raw // 4 != self.num_frame_per_block:
-                repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
-                videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
-            else:
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
+            # Eager cross-attn population (AR-Diffusion only): cache text + image-token K/V
+            # into the pool now that the image is encoded (clip_feas available).
+            # Runs on the first forward of a session and after each window-boundary
+            # reset (current_start_frame returns to 0); the populate guards (text/img halves) gate re-entry.
+            self._kv_populate_cross(prompt_embeds, state.clip_feas, is_negative=False)
+            if negative_prompt_embeds is not None:
+                self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True)
 
+        if state.current_start_frame != 0:
             latent_dtype = videos.dtype
             with torch.no_grad():
-                image = self._encode_vae_latents(videos)
-            image = image.to(dtype=latent_dtype)
+                image = self._encode_observation_latents(state, videos, latent_dtype=latent_dtype)
+        _pt_mark("obs_vae_encode")
 
         batch_size = image.shape[0]
         generator = torch.Generator(device=device).manual_seed(self.seed)
@@ -1082,6 +1370,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             do_true_cfg,
             state,
         )
+        _pt_mark("prefill_kv")
 
         sample_scheduler = copy.deepcopy(self.scheduler)
         sample_scheduler_action = copy.deepcopy(self.scheduler)
@@ -1123,6 +1412,19 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             state_features=state_features,
             embodiment_id=embodiment_id,
         )
+        if _pt is not None:
+            _pt_mark("diffuse")
+            prev_t = _pt["t0"]
+            parts = []
+            for name, t in _pt["marks"]:
+                parts.append(f"{name}={1000 * (t - prev_t):.1f}ms")
+                prev_t = t
+            logger.info(
+                "DZ_PHASE_TIMING csf=%s total=%.1fms %s",
+                state.current_start_frame,
+                1000 * (prev_t - _pt["t0"]),
+                " ".join(parts),
+            )
 
         if state.current_start_frame == 1:
             video_out = torch.cat([image, video_out], dim=1)
@@ -1133,7 +1435,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # q99 denorm: [-1,1] → real values
         action_out = self._denormalize_action(action_out.float(), embodiment_name)
 
-        # Relative → absolute: only for relative_action_keys (joint_position only)
+        # Relative -> absolute: only for relative_action_keys (joint_position only)
         # gripper_position is NOT relative, so don't add state back to it
         if self.relative_action and state_for_postprocess is not None:
             n_relative = self.relative_action_dim  # 7 for DROID (joint only)
@@ -1145,7 +1447,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 action_out[..., :n_relative] + last_state.unsqueeze(1)  # broadcast over horizon
             )
 
-        # Squeeze batch dim for output: (B, horizon, dim) → (horizon, dim)
+        # Squeeze batch dim for output: (B, horizon, dim) -> (horizon, dim)
         actions_np = action_out.squeeze(0).float().cpu().numpy()  # (horizon, max_action_dim)
         actions_np = transform.transform_action_output(actions_np)
 
@@ -1245,7 +1547,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         q01 = stats["q01"].to(device=action.device, dtype=action.dtype)
         q99 = stats["q99"].to(device=action.device, dtype=action.dtype)
         # action shape: (B, horizon, action_dim) or (B, horizon, max_action_dim)
-        # q01/q99 shape: (actual_action_dim,) — only denorm actual dims
+        # q01/q99 shape: (actual_action_dim,) -- only denorm actual dims
         actual_dim = q01.shape[0]
         action_real = action.clone()
         action_real[..., :actual_dim] = (action[..., :actual_dim] + 1) / 2 * (q99 - q01) + q01

vllm_omni/diffusion/models/dreamzero/state_dreamzero.py

@@ -115,17 +115,24 @@ class DreamZeroState:
         # Concatenated along T before decode (matches upstream video_across_time).
         self.video_latents_across_time = saved_video_latents
 
-        # KV cache once robot-policy diffusion supports that integration.
-        self.kv_cache = None
-        self.kv_cache_neg = None
-        self.crossattn_cache = None
-        self.crossattn_cache_neg = None
         self.current_start_frame = 0
 
         self.clip_feas = None
         self.ys = None
         if clear_video_latents:
+            # Session reset: drop the prompt-embed cache and VAE encoder stream.
             self.language = None
+            self.prompt_embeds = None
+            self.reset_vae_encoder_stream()
+        # Window ("inference") resets keep both: the prompt is unchanged, and the
+        # Wan feat_cache history is independent of the DiT attention window.
+
+    def reset_vae_encoder_stream(self) -> None:
+        """Clear incremental Wan VAE encoder state used across AR steps."""
+        self.vae_stream_initialized = False
+        self.vae_enc_feat_map: list[torch.Tensor | None] | None = None
+        self.vae_encoder_out: torch.Tensor | None = None
+        self.vae_pending_body_frames: torch.Tensor | None = None
 
     def reset_inference_state(self) -> None:
         """Reset KV/frame state after local attention rolls without dropping video latents."""
@@ -166,57 +173,3 @@ class DreamZeroState:
     def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
         """Determine if state should be reset before this forward()."""
         return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None
-
-    # ------------------------------------------------------------------
-    # KV cache management
-    # ------------------------------------------------------------------
-
-    def create_kv_caches(
-        self,
-        batch_size: int,
-        dtype: torch.dtype,
-        device: torch.device,
-        num_layers: int,
-        num_heads: int,
-        head_dim: int,
-    ) -> None:
-        """Initialize empty KV caches and cross-attention caches."""
-        self.kv_cache = [
-            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
-        ]
-        self.kv_cache_neg = [
-            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
-        ]
-
-        self.crossattn_cache = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-        self.crossattn_cache_neg = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-
-    def update_kv_cache(
-        self,
-        layer_index: int,
-        updated_kv: torch.Tensor,
-        is_negative: bool = False,
-    ) -> None:
-        """Update a single layer's KV cache after prefill."""
-        cache = self.kv_cache_neg if is_negative else self.kv_cache
-        if cache is None:
-            raise RuntimeError("KV caches not initialized, call create_kv_caches first.")
-        cache[layer_index] = updated_kv.clone()
-
-    def get_kv_caches(self, is_negative: bool = False) -> list[torch.Tensor]:
-        """Get KV caches for the specified branch."""
-        cache = self.kv_cache_neg if is_negative else self.kv_cache
-        if cache is None:
-            raise RuntimeError("KV caches not initialized.")
-        return cache
-
-    def get_crossattn_caches(self, is_negative: bool = False) -> list[dict[str, bool | torch.Tensor | None]]:
-        """Get cross-attention caches for the specified branch."""
-        cache = self.crossattn_cache_neg if is_negative else self.crossattn_cache
-        if cache is None:
-            raise RuntimeError("Cross-attn caches not initialized.")
-        return cache

vllm_omni/diffusion/models/dreamzero/video_export_worker.py

@@ -9,6 +9,35 @@ import torch
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
+    def gpu_mem_stats(self) -> dict:
+        """Peak GPU memory (GiB) for this worker's device, for profiling.
+
+        ``max_memory_reserved`` is the caching-allocator high-water mark (the best
+        in-process proxy for "peak VRAM"); ``max_memory_allocated`` is the live
+        tensor high-water mark. Both are monotonic from process start, so the
+        end-of-run values capture the whole rollout's peak.
+        """
+        dev = torch.accelerator.current_device_index()
+        return {
+            "device": int(dev),
+            "peak_reserved_gib": torch.accelerator.max_memory_reserved(dev) / (1024**3),
+            "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3),
+        }
+
+    def ar_diffusion_perf_stats(self, reset: bool = True) -> dict:
+        """Per-request server-side forward E2E times (s) recorded by the AR-Diffusion runner.
+
+        The worker-side compute time per ``execute_model`` — the analog of the
+        upstream server's per-request E2E — which excludes the engine<->worker IPC
+        that the client's ``omni.generate`` wall time includes. Returns ``[]`` when
+        the runner is the base (KV-disabled) runner that does not record timings.
+        """
+        runner = self.model_runner
+        times = list(getattr(runner, "_perf_e2e_times", []) or [])
+        if reset and hasattr(runner, "_perf_e2e_times"):
+            runner._perf_e2e_times = []
+        return {"server_e2e_s": times}
+
     @staticmethod
     def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
         decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py

@@ -3,6 +3,8 @@
 
 from typing import Any
 
+import torch
+import vllm.distributed.parallel_state as _vllm_ps
 import vllm.forward_context as _vllm_fc
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
@@ -24,6 +26,30 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None:
         forward_context.in_profile_run = False
 
 
+def _set_forward_context_dp_metadata(num_tokens: int) -> None:
+    """Populate vLLM MoE DP token metadata for diffusion-only DP mappings."""
+    if not _vllm_fc.is_forward_context_available():
+        return
+    forward_context = _vllm_fc.get_forward_context()
+    if getattr(forward_context, "dp_metadata", None) is not None:
+        return
+
+    dp_group = getattr(_vllm_ps, "_DP", None)
+    if dp_group is None or getattr(dp_group, "world_size", 1) <= 1:
+        return
+
+    gathered_num_tokens: list[int | None] = [None] * dp_group.world_size
+    torch.distributed.all_gather_object(
+        gathered_num_tokens,
+        int(num_tokens),
+        group=dp_group.cpu_group,
+    )
+    if any(count is None for count in gathered_num_tokens):
+        raise RuntimeError(f"Failed to gather MoE DP token counts: {gathered_num_tokens}")
+
+    forward_context.dp_metadata = _vllm_fc.DPMetadata(torch.tensor(gathered_num_tokens, dtype=torch.int64))
+
+
 class HunyuanFusedMoEDefault:
     """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3.
 
@@ -62,6 +88,7 @@ class HunyuanFusedMoEDefault:
                 hidden_states = args[0]
             if hidden_states is not None:
                 _set_forward_context_num_tokens(hidden_states.shape[0])
+                _set_forward_context_dp_metadata(hidden_states.shape[0])
 
         moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True)
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -2235,13 +2235,8 @@ class HunyuanImage3Model(nn.Module):
             # processed with quantization, LoRA, fine-tuning, etc.
             if self.config.tie_word_embeddings and "lm_head.weight" in name:
                 continue
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache scales for compressed-tensors quantization
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                continue
+            # KV-cache scales are renamed via maybe_remap_kv_scale_name below;
+            # quant_config.get_cache_scale was removed in vLLM v0.23.0 (see #4810).
 
             is_found = False
             for param_name, weight_name, shard_id in stacked_params_mapping:

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -213,8 +213,32 @@ class DiffusionWorker:
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
-        # Create model runner using the platform-specified class
-        model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls()
+        # Create model runner — one decision chain, in precedence order:
+        #   1. explicit od_config.diffusion_model_runner_cls (user override),
+        #   2. the runner declared by the engine class that engine_backend
+        #      selects (e.g. ARDiffusionEngine -> ARDiffusionModelRunner),
+        #   3. the platform default.
+        # Routing policy therefore lives on the engine class / config surface;
+        # engines never mutate od_config. Overrides must be import-path
+        # strings — guard with isinstance so a non-string (e.g. a Mock
+        # od_config in tests) doesn't shadow the platform hook.
+        runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None)
+        engine_runner = None
+        if not (isinstance(runner_override, str) and runner_override):
+            try:
+                from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+                engine_cls = DiffusionEngine.resolve_engine_class(self.od_config)
+                engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+            except Exception:
+                logger.warning("Worker %s: engine_backend resolution failed; using platform runner", self.rank)
+                engine_runner = None
+        if isinstance(runner_override, str) and runner_override:
+            model_runner_cls_path = runner_override
+        elif isinstance(engine_runner, str) and engine_runner:
+            model_runner_cls_path = engine_runner
+        else:
+            model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls()
         model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path)
         self.model_runner = model_runner_cls(
             vllm_config=self.vllm_config,

vllm_omni/entrypoints/async_omni.py

@@ -26,7 +26,7 @@ from vllm.tasks import SupportedTask
 from vllm.utils import random_uuid
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
+from vllm_omni.diffusion.data import CuMemTag, OmniACK, OmniSleepTask, OmniWakeTask
 from vllm_omni.engine.messages import ErrorMessage, OutputMessage
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import (
@@ -141,7 +141,8 @@ class AsyncOmni(EngineClient, OmniBase):
         OmniBase.__init__(self, model=model, **kwargs)
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
-        self._is_sleeping: bool = False
+        self._sleeping_tags: set[str] = set()
+        self._level2_sleeping: bool = False
         self.final_output_task: asyncio.Task | None = None
         self.event_resolver = AsyncEventResolver(orchestrator=self)
         self.config_path = self.engine.config_path
@@ -208,10 +209,13 @@ class AsyncOmni(EngineClient, OmniBase):
 
     def get_diffusion_od_config(self) -> Any | None:
         """Return the diffusion-stage config when the pipeline has one."""
+        saw_diffusion_stage = False
         for stage_client in self.engine.stage_clients:
             if getattr(stage_client, "stage_type", None) != "diffusion":
                 continue
 
+            saw_diffusion_stage = True
+
             od_config = getattr(stage_client, "od_config", None)
             if od_config is not None:
                 return od_config
@@ -221,6 +225,11 @@ class AsyncOmni(EngineClient, OmniBase):
             if od_config is not None:
                 return od_config
 
+        # Out-of-process diffusion clients don't carry od_config (it lives in the
+        # worker); fall back to the engine's model_class_name resolution.
+        if saw_diffusion_stage:
+            return self.engine.get_diffusion_od_config()
+
         return None
 
     @property
@@ -306,6 +315,14 @@ class AsyncOmni(EngineClient, OmniBase):
 
         logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}")
 
+        _sleeping_tags = getattr(self, "_sleeping_tags", None)
+        if _sleeping_tags:
+            raise RuntimeError(
+                f"Generation rejected: Engine is partially or fully asleep. "
+                f"Currently sleeping tags: {list(_sleeping_tags)}. "
+                f"Please perform a full wake_up before generating."
+            )
+
         # Reject diffusion list-prompt early with a clear API error.
         if isinstance(prompt, list) and any(
             getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", [])
@@ -922,11 +939,32 @@ class AsyncOmni(EngineClient, OmniBase):
                 if ack is not None:
                     await self.event_resolver.resolve(ack)
                     final_acks.append(ack)
-        self._is_sleeping = True
+        if not hasattr(self, "_sleeping_tags"):
+            self._sleeping_tags = set()
+        self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value])
+        if level == 2:
+            self._level2_sleeping = True
         return final_acks
 
     async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]:
         self._final_output_handler()
+
+        if getattr(self, "_level2_sleeping", False):
+            raise NotImplementedError(
+                "wake_up() after sleep(level=2) is not yet implemented: weights were "
+                "discarded from GPU and reloading from disk is not yet supported. "
+                "Use sleep(level=1) instead, which offloads weights to CPU RAM "
+                "and supports fast DMA restore."
+            )
+        _current_tags = getattr(self, "_sleeping_tags", set())
+        if tags is None:
+            requested_tags = list(_current_tags)
+        else:
+            requested_tags = [t for t in tags if t in _current_tags]
+        if not requested_tags:
+            logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.")
+            return []
+
         if stage_ids is None:
             stage_ids = list(range(len(self.engine.stage_clients)))
         total_workers = 0
@@ -940,7 +978,7 @@ class AsyncOmni(EngineClient, OmniBase):
         task_id = str(uuid.uuid4())
         self.event_resolver.watch_task(task_id, expected_count=total_workers)
         logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
-        task = OmniWakeTask(tags=tags, task_id=task_id)
+        task = OmniWakeTask(tags=requested_tags, task_id=task_id)
         rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
         final_acks = []
         for stage_res in rpc_results:
@@ -951,7 +989,14 @@ class AsyncOmni(EngineClient, OmniBase):
                     final_acks.append(ack)
         current_omni_platform.synchronize()
         await asyncio.sleep(0.1)
-        self._is_sleeping = False
+
+        for t in requested_tags:
+            if hasattr(self, "_sleeping_tags"):
+                self._sleeping_tags.discard(t)
+        # Only clear the level-2 flag once all tags are warm, in case partial
+        # wake support (e.g. tags=["kv_cache"] only) is added in the future.
+        if not getattr(self, "_sleeping_tags", None):
+            self._level2_sleeping = False
         logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.")
         return final_acks
 
@@ -961,7 +1006,7 @@ class AsyncOmni(EngineClient, OmniBase):
         TODO(AsyncOmni): query the orchestrator once all stage backends expose
         a real sleeping-state RPC. For now we track the requested state locally.
         """
-        return self._is_sleeping
+        return bool(getattr(self, "_sleeping_tags", None))
 
     async def add_lora(self, lora_request: LoRARequest) -> bool:
         """Load a new LoRA adapter into all stages.

vllm_omni/entrypoints/openai/serving_speech.py

@@ -3495,12 +3495,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_id = request_id or f"speech-{random_uuid()}"
         qwen3_ref_audio_warmup_artifact_key: str | None = None
 
-        # If this is a streaming request, we need to coerce
-        # cumulative outputs to delta outputs; this ensures
-        # we don't emit redundant MM data & drain after emitting.
+        # If this is a streaming request with real async chunks, we need to
+        # coerce cumulative outputs to delta outputs; this ensures we don't
+        # emit redundant MM data & drain after emitting. Qwen3-TTS full-payload
+        # (async_chunk=False) has no incremental audio chunks, so keep
+        # FINAL_ONLY semantics and let the streaming response send the final
+        # waveform once. Scoped to qwen3_tts: other async_chunk=False models
+        # keep the DELTA coercion they stream with today.
         # list() makes a copy to avoid mutating the params.
         sampling_params_list = list(self.engine_client.default_sampling_params_list)
-        is_streaming_request = request.is_streaming()
+        async_chunk = getattr(self.model_config, "async_chunk", True)
+        qwen3_full_payload = self._tts_model_type == "qwen3_tts" and not bool(async_chunk)
+        is_streaming_request = request.is_streaming() and not qwen3_full_payload
         sampling_params_list = coerce_param_message_types(sampling_params_list, is_streaming_request)
 
         # Build prompt + tts_params via the per-model adapter (RFC #4327). Every
@@ -3687,6 +3693,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 stage0_params.extra_args = {}
             stage0_params.extra_args["tts_local_seed"] = request.seed
 
+        if self._tts_model_type == "qwen3_tts" and sampling_params_list:
+            stage0_params = sampling_params_list[0]
+            default_seed = getattr(stage0_params, "seed", None)
+            if default_seed is not None:
+                import copy
+
+                sampling_params_list = copy.deepcopy(sampling_params_list)
+                stage0_params = sampling_params_list[0]
+                if stage0_params.extra_args is None:
+                    stage0_params.extra_args = {}
+                stage0_params.extra_args.setdefault("tts_local_seed", int(default_seed))
+
         generator = self.engine_client.generate(
             prompt=prompt,
             request_id=request_id,

vllm_omni/experimental/__init__.py

@@ -1,3 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Experimental vLLM-Omni features (APIs may change without notice)."""
+"""Experimental vLLM-Omni subsystems.
+
+Modules under ``vllm_omni.experimental`` are in active development and not yet
+covered by the stability guarantees of the rest of the package — their APIs,
+config surfaces, and defaults may change without notice. The AR-Diffusion
+Engine (:mod:`vllm_omni.experimental.ar_diffusion`) lives here.
+"""

vllm_omni/experimental/ar_diffusion/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AR-Diffusion Engine (AR-Diffusion).
+
+The AR-Diffusion engine: a ``DiffusionEngine`` subclass that adds engine-level KV
+cache management for autoregressive / chunked "world-model" diffusion models
+(DreamZero and the AR-DiT family). Selected via
+``OmniDiffusionConfig.engine_backend = "ar_diffusion"``.
+"""
+
+from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
+
+__all__ = ["ARDiffusionEngine"]

vllm_omni/experimental/ar_diffusion/engine.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The AR-Diffusion Engine (AR-Diffusion)."""
+
+from __future__ import annotations
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+logger = init_logger(__name__)
+
+#: Import path of the runner the AR-Diffusion engine routes its workers to.
+AR_DIFFUSION_MODEL_RUNNER_CLS = "vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner"
+
+
+class ARDiffusionEngine(DiffusionEngine):
+    """AR-Diffusion engine with engine-level KV cache management.
+
+    AR-Diffusion serves autoregressive / chunked blockwise-causal diffusion models
+    (world models, AR-DiT) that materialize persistent attention KV. It reuses
+    vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` / ``BlockTables``)
+    as a library, driven from the engine rather than hand-rolled inside each
+    model. See ``BDE_doc/diffusion_kv_cache_management.md`` for the design and
+    ``BDE_doc/dreamzero_kv_phase1_plan.md`` for the rollout.
+
+    It is selected per model via ``OmniDiffusionConfig.engine_backend = "ar_diffusion"``
+    (resolved by :meth:`DiffusionEngine.make_engine`), so models that do not opt
+    in keep using the base ``DiffusionEngine`` unchanged.
+
+    Architecture note: in the multiproc setup the KV cache lives in the worker /
+    runner process (GPU side), co-located with the model and KV tensors — so the
+    actual KV *body* is :class:`~vllm_omni.experimental.ar_diffusion.kv_cache.manager.ARDiffusionKVCache`, owned
+    by :class:`~vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner`. ``ARDiffusionEngine`` itself is the
+    thin selection / injection seam; it wires the AR-Diffusion executor → worker → runner
+    so DreamZero's rollout runs against the runner-owned KV cache.
+    """
+
+    #: Workers of this engine build the AR-Diffusion runner (unless the config
+    #: sets an explicit ``diffusion_model_runner_cls`` override). Declared on
+    #: the class so routing needs no od_config mutation at engine init.
+    default_diffusion_model_runner_cls = AR_DIFFUSION_MODEL_RUNNER_CLS

vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AR-Diffusion engine-level KV cache helpers.
+
+Thin glue over vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` /
+``SlidingWindowManager``), used by the AR-Diffusion Engine to manage KV for
+AR-diffusion models. See ``BDE_doc/dreamzero_kv_phase1_plan.md``.
+
+Layout: ``config`` (public knob) · ``paged`` (engine-generic paging mechanics +
+chunk-window eviction spec) · ``manager`` (the ARDiffusionKVCache orchestrator + its
+request adapter / pool builders) · ``state`` (the model-facing ARDiffusionKVState bridge).
+"""
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.manager import (
+    ARDiffusionKVCache,
+    ARDiffusionRequestAdapter,
+    build_kv_manager,
+    compute_num_blocks,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    pool_write_chunk,
+    resident_block_ids,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+    ARDiffusionPagedForwardContext,
+    ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
+    ar_diffusion_paged_attention,
+    paged_write_attn,
+    set_current_paged_kv_cache,
+)
+
+__all__ = [
+    "ARDiffusionKVCache",
+    "ARDiffusionKVConfig",
+    "ARDiffusionPagedForwardContext",
+    "ARDiffusionPagedLayerContext",
+    "ARDiffusionPagedLayerInputs",
+    "ARDiffusionRequestAdapter",
+    "ChunkWindowManager",
+    "ChunkWindowSpec",
+    "allocate_kv_pool_with_views",
+    "ar_diffusion_paged_attention",
+    "build_kv_manager",
+    "chunk_slot_mapping",
+    "compute_num_blocks",
+    "compute_slot_mapping",
+    "paged_write_attn",
+    "pool_write_chunk",
+    "resident_block_ids",
+    "set_current_paged_kv_cache",
+]

vllm_omni/experimental/ar_diffusion/kv_cache/config.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration for AR-Diffusion engine-level KV cache management."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ARDiffusionKVConfig:
+    """Settings for the AR-Diffusion paged KV cache.
+
+    Disabled by default: when ``enable`` is False the AR-Diffusion engine behaves exactly
+    like the base ``DiffusionEngine`` (no pool, no paged KV).
+    """
+
+    enable: bool = False
+    # Persistent KV tokens materialized per AR chunk. For DreamZero this is
+    # ``num_frame_per_block * frame_seqlen``.
+    chunk_size: int = 0
+    # Resident window in chunks. ``None`` means full attention (no eviction).
+    window_chunks: int | None = None
+    # Protected leading chunks (attention sink); never evicted.
+    sink_chunks: int = 0
+    # DreamZero-style window reset vs. VGGT-style sliding replace.
+    reset_at_boundary: bool = False
+    # Fraction of free device memory budgeted for the AR-Diffusion KV pool.
+    gpu_memory_fraction: float = 0.1
+    # When CUDA graph / torch.compile is on (not enforce_eager), pre-capture the
+    # DiT graphs for every window-fill shape at load time via a synthetic rollout,
+    # so the serving run is fast from the first chunk. No effect when eager.
+    warmup_cudagraph: bool = True
+    # Also capture the post-window-boundary (reset-cycle) forward during warm-up.
+    warmup_capture_reset: bool = False
+
+    @property
+    def sliding_window(self) -> int | None:
+        """Window size in tokens, or ``None`` for full attention."""
+        if self.window_chunks is None:
+            return None
+        return self.window_chunks * self.chunk_size

vllm_omni/experimental/ar_diffusion/kv_cache/manager.py

@@ -0,0 +1,497 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ARDiffusionKVCache — the engine-level KV cache orchestrator for one AR-Diffusion model.
+
+This is the *body* of AR-Diffusion's KV management: it owns a vLLM ``KVCacheManager`` (a
+single chunk-window group) and the per-request adapter lifecycle, and exposes the
+per-chunk operations a rollout needs — allocate, slot mapping, commit, window
+lookup, free. It lives in the model runner (worker / GPU side), co-located with
+the model and the KV tensors; the DreamZero pipeline calls these methods during a
+rollout. The main-process ``ARDiffusionEngine`` only selects the engine and is otherwise
+thin.
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+from collections.abc import Sequence
+
+import torch
+from vllm.logger import init_logger
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheSpec,
+    KVCacheTensor,
+)
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    chunk_slot_mapping,
+    pool_write_chunk,
+    resident_block_ids,
+)
+
+_log = init_logger(__name__)
+
+
+class ARDiffusionRequestAdapter:
+    """Duck-types the subset of ``vllm.v1.request.Request`` that the
+    ``KVCacheManager`` reads (``allocate_slots`` / ``get_computed_blocks`` /
+    ``free`` and the coordinator they call into).
+
+    It is intentionally NOT a full ``Request``. The conformance test exercises a
+    real ``KVCacheManager`` against this adapter so the surface cannot silently
+    drift across vLLM versions.
+
+    A AR-Diffusion request advances one *chunk* at a time: ``allocate_slots`` is called
+    once per chunk and ``num_computed_tokens`` advances only when a chunk is
+    committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk
+    reuse the same slots.
+    """
+
+    def __init__(
+        self,
+        request_id: str,
+        *,
+        chunk_size: int,
+        prefill_prefix_tokens: int = 0,
+    ) -> None:
+        self.request_id = request_id
+        self._chunk_size = chunk_size
+        self._prefill = prefill_prefix_tokens
+        self._completed_chunks = 0
+        # Filled only when cross-request prefix reuse is enabled (Phase 3).
+        self.block_hashes: list = []
+        self.skip_reading_prefix_cache = True
+        self.num_preemptions = 0
+        # vLLM watermark gate reads this; map the request lifecycle onto it.
+        self.status = RequestStatus.WAITING
+
+    @property
+    def num_computed_tokens(self) -> int:
+        """Persistent KV already materialized (committed chunks + prefill)."""
+        return self._prefill + self._completed_chunks * self._chunk_size
+
+    @property
+    def num_tokens(self) -> int:
+        """Total tokens once the in-flight chunk is committed."""
+        return self._prefill + (self._completed_chunks + 1) * self._chunk_size
+
+    @property
+    def num_prompt_tokens(self) -> int:
+        """The prefill prefix length (read by ``cache_blocks`` when caching)."""
+        return self._prefill
+
+    @property
+    def completed_chunks(self) -> int:
+        return self._completed_chunks
+
+    def on_chunk_committed(self) -> None:
+        """Advance by one chunk. Call once per chunk, not per denoise step."""
+        self._completed_chunks += 1
+
+
+def compute_num_blocks(
+    available_bytes: int,
+    gpu_memory_fraction: float,
+    page_size_bytes: int,
+) -> int:
+    """Number of KV blocks that fit in ``fraction`` of the memory budget."""
+    if page_size_bytes <= 0:
+        raise ValueError(f"page_size_bytes must be positive, got {page_size_bytes}")
+    if not 0.0 < gpu_memory_fraction <= 1.0:
+        raise ValueError(f"gpu_memory_fraction must be in (0, 1], got {gpu_memory_fraction}")
+    budget = int(available_bytes * gpu_memory_fraction)
+    return max(0, budget // page_size_bytes)
+
+
+def build_kv_manager(
+    spec: KVCacheSpec,
+    layer_names: Sequence[str],
+    num_blocks: int,
+    max_model_len: int,
+    *,
+    enable_caching: bool = False,
+) -> KVCacheManager:
+    """Build a ``KVCacheManager`` with a single KV cache group for ``spec``.
+
+    Args:
+        spec: The KV cache spec for the group (e.g. a ``ChunkWindowSpec``).
+        layer_names: Attention layers sharing this group's block table.
+        num_blocks: Total physical blocks in the pool.
+        max_model_len: Upper bound on a request's sequence length.
+        enable_caching: Cross-request prefix caching (Phase 3); off in Phase 1.
+    """
+    layer_names = list(layer_names)
+    group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
+    tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)]
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=tensors,
+        kv_cache_groups=[group],
+    )
+    kwargs = dict(max_model_len=max_model_len, hash_block_size=spec.block_size, enable_caching=enable_caching)
+    params = inspect.signature(KVCacheManager).parameters
+    if "scheduler_block_size" in params:
+        kwargs["scheduler_block_size"] = spec.block_size
+    if "max_num_batched_tokens" in params:
+        kwargs["max_num_batched_tokens"] = max_model_len
+    return KVCacheManager(config, **kwargs)
+
+
+class ARDiffusionKVCache:
+    """Owns the paged KV pool + per-request lifecycle for a AR-Diffusion model.
+
+    Build once per loaded model (dimensions known); then per request:
+    ``begin_request`` → per chunk (``allocate_chunk`` → ``chunk_write_slots`` →
+    [model writes K/V] → ``commit_chunk``) → ``end_request``.
+    """
+
+    def __init__(
+        self,
+        config: ARDiffusionKVConfig,
+        *,
+        num_layers: int,
+        num_kv_heads: int,
+        head_size: int,
+        dtype: torch.dtype,
+        block_size: int,
+        max_model_len: int,
+        available_bytes: int,
+        cross_attn_length: int = 0,
+        cross_attn_img_length: int = 0,
+        device: torch.device | None = None,
+        local_branches: int = 2,
+        num_frame_per_block: int = 1,
+    ) -> None:
+        if not config.enable:
+            raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig")
+        if config.window_chunks is None:
+            raise ValueError("Phase 1 requires a bounded window (window_chunks)")
+        if config.chunk_size <= 0:
+            raise ValueError("ARDiffusionKVConfig.chunk_size must be set (> 0)")
+        if local_branches not in (1, 2):
+            raise ValueError(f"local_branches must be 1 or 2, got {local_branches}")
+
+        self.config = config
+        # How many CFG branches allocate from THIS rank's pool. Under
+        # CFG-parallel x2 each rank executes exactly one branch (rank0 pos,
+        # rank1 neg) and the other branch's lazy contexts never allocate, so
+        # sizing for both would leave ~half the pool as idle capacity. A
+        # single-process run (TP1 / offline) executes both branches -> 2.
+        self.local_branches = local_branches
+        self.num_frame_per_block = max(1, int(num_frame_per_block))
+        self.block_size = block_size
+        self.num_layers = num_layers
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.dtype = dtype
+        self.cross_attn_length = cross_attn_length
+        self.device = device or torch.device("cpu")
+        self._adapters: dict[str, ARDiffusionRequestAdapter] = {}
+
+        # -- cross-attention pool (static, one-time-fill) -----------------------
+        # Cross-attn KV is computed once from the text encoder and never changes.
+        # Allocate a per-layer contiguous tensor for each branch: (2, text_len,
+        # kv_heads, head_dim) where dim-0 = [pos, neg].  Allocated separately
+        # from the self-attn block pool (static, small — ~400 MiB for
+        # DreamZero I2V).  Both pools draw from the same GPU free memory but
+        # the cross-attn pool is sized directly rather than via the block pool
+        # budget (no eviction/chunk lifecycle needed).
+        self.cross_attn_img_length = cross_attn_img_length
+        self._cross_k: list[torch.Tensor] = []
+        self._cross_v: list[torch.Tensor] = []
+        # I2V image-token cross-attn pool. Like the text k/v, the image-token
+        # k_img/v_img are session-invariant (the conditioning image doesn't change),
+        # so they are cached once and read every denoise step — see
+        # WanI2VCrossAttention (#4154 caches these model-side too). Empty for T2V.
+        self._cross_k_img: list[torch.Tensor] = []
+        self._cross_v_img: list[torch.Tensor] = []
+
+        # Bytes the cross-attn pools consume directly (K+V, pos+neg, all layers).
+        # Deducted from the self-attn paged-pool budget below so the two
+        # allocations together stay within the GPU free-memory budget.
+        def _cross_pool_bytes(length: int) -> int:
+            return 2 * 2 * length * num_kv_heads * head_size * dtype.itemsize * num_layers
+
+        cross_total_bytes = 0
+        if device is not None and cross_attn_length > 0:
+            cross_shape = (2, cross_attn_length, num_kv_heads, head_size)
+            cross_bytes = _cross_pool_bytes(cross_attn_length)
+            cross_total_bytes += cross_bytes
+            for _ in range(num_layers):
+                self._cross_k.append(torch.empty(cross_shape, dtype=dtype, device=device))
+                self._cross_v.append(torch.empty(cross_shape, dtype=dtype, device=device))
+            _log.info(
+                "AR-Diffusion cross-attn pool: %d layers × (%d tok × %d heads × %d) = %.1f MiB",
+                num_layers,
+                cross_attn_length,
+                num_kv_heads,
+                head_size,
+                cross_bytes / (1024 * 1024),
+            )
+            if cross_attn_img_length > 0:
+                img_shape = (2, cross_attn_img_length, num_kv_heads, head_size)
+                cross_total_bytes += _cross_pool_bytes(cross_attn_img_length)
+                for _ in range(num_layers):
+                    self._cross_k_img.append(torch.empty(img_shape, dtype=dtype, device=device))
+                    self._cross_v_img.append(torch.empty(img_shape, dtype=dtype, device=device))
+                _log.info(
+                    "AR-Diffusion cross-attn IMG pool: %d layers × %d img-tok (I2V)",
+                    num_layers,
+                    cross_attn_img_length,
+                )
+
+        self.spec = ChunkWindowSpec(
+            block_size=block_size,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=dtype,
+            sliding_window=config.window_chunks * config.chunk_size,
+            chunk_size=config.chunk_size,
+            window_chunks=config.window_chunks,
+            sink_chunks=config.sink_chunks,
+            reset_at_boundary=config.reset_at_boundary,
+        )
+        # Each pool block spans all layers' K/V, so size against the per-layer
+        # page size times the layer count.
+        # Size the self-attn pool against the memory left after the cross-attn
+        # pools (allocated above), so cross + self-attn stays within the budget.
+        num_blocks = compute_num_blocks(
+            max(0, available_bytes - cross_total_bytes),
+            config.gpu_memory_fraction,
+            self.spec.page_size_bytes * num_layers,
+        )
+        # Floor: one forward needs the resident window plus the in-flight chunk
+        # (num_frame_per_block frame-blocks) for every branch THIS rank runs,
+        # with a little eviction-transient headroom. The memory-fraction
+        # heuristic can under-size this once block_size grows — e.g. frame-granular
+        # paging at the true frame_seqlen makes each block larger and the pool
+        # fewer-blocks — so guarantee the minimum the rollout cannot run without,
+        # otherwise allocate_chunk hits an exhausted pool mid-forward.
+        min_blocks = self.local_branches * (config.window_chunks + self.num_frame_per_block) + 2
+        if num_blocks < min_blocks:
+            _log.warning(
+                "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
+                "floor (%d local CFG branch(es) x (window_chunks=%d + num_frame_per_block=%d) + 2 headroom)",
+                num_blocks,
+                min_blocks,
+                self.local_branches,
+                config.window_chunks,
+                self.num_frame_per_block,
+            )
+            num_blocks = min_blocks
+        layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
+        self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
+        self.managed_num_blocks = num_blocks
+        self.num_blocks = num_blocks
+        # Scratch blocks are outside KVCacheManager ownership. They hold current
+        # denoise-step video KV when update_kv_cache=False and action/state KV in
+        # both modes; they are reused every forward and never committed.
+        # Default 4: measured high-water is num_frame_per_block video blocks (2)
+        # + 1 action block per branch per forward, +1 margin.
+        scratch_per_branch = int(os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "4"))
+        if scratch_per_branch <= 0:
+            raise ValueError("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH must be positive")
+        self.scratch_blocks_per_branch = scratch_per_branch
+        self.scratch_num_blocks = self.local_branches * scratch_per_branch
+        self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks
+        self.null_block_id = self.manager.block_pool.null_block.block_id
+
+        # Allocate the per-layer paged K/V pools on the given device.
+        self._kv_pools: list[torch.Tensor] = []
+        self._k_pools: list[torch.Tensor] = []
+        self._v_pools: list[torch.Tensor] = []
+        if device is not None:
+            self._kv_pools, self._k_pools, self._v_pools = allocate_kv_pool_with_views(
+                self.num_blocks_total,
+                block_size,
+                num_layers,
+                num_kv_heads,
+                head_size,
+                dtype,
+                device,
+            )
+
+    # -- cross-attention pool access -------------------------------------------
+    # Cross-attn KV is static once populated — write once (from text encoder),
+    # read many (every denoising step). Not managed through the paged block pool.
+
+    def write_cross_kv(
+        self,
+        layer_idx: int,
+        is_negative: bool,
+        k: torch.Tensor | None,
+        v: torch.Tensor | None,
+        k_img: torch.Tensor | None = None,
+        v_img: torch.Tensor | None = None,
+    ) -> None:
+        """Write one layer's cross-attn K/V into the pool.
+
+        ``k`` / ``v``: ``(B, text_len, tp_num_heads, head_dim)``; only batch-0
+        is copied (B=1 for inference). ``None`` skips the text half (window
+        restarts keep the still-valid text K/V and rewrite only the image half).
+        The ``is_negative`` flag selects the correct CFG branch slot. ``k_img`` /
+        ``v_img`` (I2V image tokens, ``(B, 257, ...)``) are written when the
+        image pool is allocated.
+        """
+        branch = 1 if is_negative else 0
+        if k is not None:
+            self._cross_k[layer_idx][branch].copy_(k[0])
+            self._cross_v[layer_idx][branch].copy_(v[0])
+        if k_img is not None and self._cross_k_img:
+            self._cross_k_img[layer_idx][branch].copy_(k_img[0])
+            self._cross_v_img[layer_idx][branch].copy_(v_img[0])
+
+    def read_cross_kv(self, layer_idx: int, is_negative: bool) -> dict:
+        """Return a pool-backed cross-attn cache dict for one layer.
+
+        The dict matches the ``{"is_init": True, "k": Tensor, "v": Tensor}``
+        convention the cross-attention module expects — it reads from the pool
+        slice rather than from the lazy-initialised model-local dict. For I2V,
+        ``k_img`` / ``v_img`` are added so the image-token cache (added by #4154)
+        reads from the pool too.
+        """
+        branch = 1 if is_negative else 0
+        k = self._cross_k[layer_idx][branch].unsqueeze(0)  # (1, L, heads, dim)
+        v = self._cross_v[layer_idx][branch].unsqueeze(0)
+        cache = {"is_init": True, "k": k, "v": v}
+        if self._cross_k_img:
+            cache["k_img"] = self._cross_k_img[layer_idx][branch].unsqueeze(0)
+            cache["v_img"] = self._cross_v_img[layer_idx][branch].unsqueeze(0)
+        return cache
+
+    # -- request lifecycle ---------------------------------------------------
+
+    def begin_request(self, request_id: str, *, prefill_prefix_tokens: int = 0) -> ARDiffusionRequestAdapter:
+        adapter = ARDiffusionRequestAdapter(
+            request_id,
+            chunk_size=self.spec.chunk_size,
+            prefill_prefix_tokens=prefill_prefix_tokens,
+        )
+        self._adapters[request_id] = adapter
+        _log.debug("AR-Diffusion begin_request: req=%s prefill=%d", request_id, prefill_prefix_tokens)
+        return adapter
+
+    def end_request(self, adapter: ARDiffusionRequestAdapter) -> None:
+        _log.debug(
+            "AR-Diffusion end_request: req=%s chunks=%d free=%d",
+            adapter.request_id,
+            adapter.completed_chunks,
+            self.manager.block_pool.get_num_free_blocks(),
+        )
+        self.manager.free(adapter)
+        self._adapters.pop(adapter.request_id, None)
+
+    # -- per-chunk operations ------------------------------------------------
+
+    def allocate_chunk(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
+        """Allocate a chunk's blocks (evicting out-of-window blocks first).
+
+        Returns the request's full block table (incl. null_block placeholders).
+        """
+        blocks = self.manager.allocate_slots(adapter, num_new_tokens=self.spec.chunk_size)
+        if blocks is None:
+            raise RuntimeError("AR-Diffusion KV pool exhausted while allocating a chunk")
+        table = self.block_table(adapter)
+        resident = resident_block_ids(table, self.null_block_id)
+        _log.debug(
+            "AR-Diffusion allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d",
+            adapter.request_id,
+            adapter.completed_chunks,
+            len(table),
+            len(resident),
+            self.manager.block_pool.get_num_free_blocks(),
+        )
+        return table
+
+    def allocate_token_slots(self, adapter: ARDiffusionRequestAdapter, num_tokens: int) -> list[int]:
+        """Allocate managed blocks for an in-flight video span without committing it."""
+        if num_tokens <= 0:
+            raise ValueError(f"num_tokens must be positive, got {num_tokens}")
+        blocks = self.manager.allocate_slots(adapter, num_new_tokens=num_tokens)
+        if blocks is None:
+            raise RuntimeError("AR-Diffusion KV pool exhausted while allocating paged attention slots")
+        return self.block_table(adapter)
+
+    def block_table(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
+        return list(self.manager.get_block_ids(adapter.request_id)[0])
+
+    def chunk_write_slots(self, adapter: ARDiffusionRequestAdapter) -> torch.Tensor:
+        """Slot mapping for the in-flight chunk — the K/V write target."""
+        return chunk_slot_mapping(
+            self.block_table(adapter),
+            adapter.num_computed_tokens,
+            self.spec.chunk_size,
+            self.block_size,
+        )
+
+    def scratch_block_ids(self, is_negative: bool, start: int, count: int) -> list[int]:
+        """Return branch-local scratch block ids outside manager ownership."""
+        if count < 0 or start < 0:
+            raise ValueError(f"scratch start/count must be non-negative, got start={start}, count={count}")
+        if start + count > self.scratch_blocks_per_branch:
+            raise RuntimeError(
+                "AR-Diffusion paged attention scratch blocks exhausted: "
+                f"need [{start}, {start + count}) of {self.scratch_blocks_per_branch}. "
+                "Increase AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH."
+            )
+        # With one local branch (CFG-parallel x2) the rank's only branch maps to
+        # slot 0 whichever CFG side it is; the other branch never allocates here.
+        branch_offset = self.scratch_blocks_per_branch if (is_negative and self.local_branches == 2) else 0
+        base = self.managed_num_blocks + branch_offset + start
+        return list(range(base, base + count))
+
+    def key_cache(self, layer_idx: int) -> torch.Tensor:
+        return self._kv_pools[layer_idx][0]
+
+    def value_cache(self, layer_idx: int) -> torch.Tensor:
+        return self._kv_pools[layer_idx][1]
+
+    def window_block_ids(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
+        """Resident (non-null) managed blocks visible to paged attention."""
+        return resident_block_ids(self.block_table(adapter), self.null_block_id)
+
+    def commit_chunk(self, adapter: ARDiffusionRequestAdapter) -> None:
+        """Advance the adapter by one chunk after its K/V is written.
+
+        This standalone primitive is used by low-level manager tests. In the
+        DreamZero paged-attention path, :meth:`ARDiffusionKVState.commit_paged_context`
+        advances the adapter only after the forward succeeds. Call once per
+        committed chunk, not per denoise step.
+        """
+        _log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
+        adapter.on_chunk_committed()
+        _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
+
+    # -- pool-backed K/V access --------------------------------------------
+
+    def write_chunk_kv(
+        self,
+        layer_index: int,
+        new_k: torch.Tensor,
+        new_v: torch.Tensor,
+        adapter: ARDiffusionRequestAdapter,
+    ) -> None:
+        """Write one layer's committed-chunk K/V into the pool."""
+        slots = self.chunk_write_slots(adapter)
+        _log.debug(
+            "AR-Diffusion write: req=%s layer=%d chunk=%d shapes=%s dev=%s",
+            adapter.request_id,
+            layer_index,
+            adapter.completed_chunks,
+            (tuple(new_k.shape), tuple(new_v.shape)),
+            slots.device,
+        )
+        pool_write_chunk(
+            self._k_pools[layer_index],
+            self._v_pools[layer_index],
+            new_k,
+            new_v,
+            slots,
+        )

vllm_omni/experimental/ar_diffusion/kv_cache/paged.py

@@ -0,0 +1,223 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generic paging mechanics + chunk-window eviction spec for the AR-Diffusion engine.
+
+Engine-generic, model-agnostic primitives — the layer a second model (e.g. the
+Cosmos port) reuses unchanged. Three concerns live here:
+
+* **Slot mapping** — absolute token positions → physical KV-cache slots, the
+  standard PagedAttention layout ``slot(pos) = block_id(pos) * block_size +
+  (pos % block_size)``.
+* **Pool I/O** — allocate per-layer FlashAttention-compatible paged K/V pools
+  plus flat compatibility views for direct slot writes.
+* **Chunk-window eviction** — a ``SlidingWindowSpec`` subclass whose unit is a
+  *chunk* (``sliding_window = window_chunks * chunk_size``) plus a manager that
+  evicts at chunk boundaries. Memory policy / refcounting / ``null_block``
+  replacement stay in vLLM's ``BlockPool``; only the token-skip math is here.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import torch
+from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
+from vllm.v1.kv_cache_interface import SlidingWindowSpec
+
+try:
+    from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+except ModuleNotFoundError:
+
+    def register_kv_cache_spec(*, manager_class, uniform_type_base_spec=None):
+        """Compatibility shim for vLLM versions without kv_cache_spec_registry."""
+
+        def decorator(spec_class):
+            from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
+            spec_manager_map[spec_class] = manager_class
+            return spec_class
+
+        return decorator
+
+
+# ── Slot mapping ────────────────────────────────────────────────────────────
+
+
+def compute_slot_mapping(
+    block_ids: Sequence[int],
+    positions: torch.Tensor | Sequence[int],
+    block_size: int,
+) -> torch.Tensor:
+    """Map absolute token positions to physical KV-cache slots.
+
+    Args:
+        block_ids: physical block id per block index (the request's block table).
+        positions: absolute token positions to map (1-D).
+        block_size: tokens per block.
+
+    Returns:
+        ``LongTensor`` of physical slots, one per position.
+    """
+    if block_size <= 0:
+        raise ValueError(f"block_size must be positive, got {block_size}")
+    table = torch.as_tensor(block_ids, dtype=torch.long)
+    pos = torch.as_tensor(positions, dtype=torch.long)
+    block_index = torch.div(pos, block_size, rounding_mode="floor")
+    offset = pos % block_size
+    return table[block_index] * block_size + offset
+
+
+def chunk_slot_mapping(
+    block_ids: Sequence[int],
+    num_computed_tokens: int,
+    chunk_size: int,
+    block_size: int,
+) -> torch.Tensor:
+    """Slot mapping for the in-flight chunk's tokens (the commit write target).
+
+    The chunk occupies absolute positions
+    ``[num_computed_tokens, num_computed_tokens + chunk_size)``.
+    """
+    positions = torch.arange(
+        num_computed_tokens,
+        num_computed_tokens + chunk_size,
+        dtype=torch.long,
+    )
+    return compute_slot_mapping(block_ids, positions, block_size)
+
+
+def resident_block_ids(block_ids: Sequence[int], null_block_id: int) -> list[int]:
+    """Real (non-null) blocks currently resident, in table order.
+
+    These are the blocks the read path gathers the attention window from;
+    out-of-window positions are the shared ``null_block`` and are excluded.
+    """
+    return [int(b) for b in block_ids if int(b) != null_block_id]
+
+
+# ── Pool I/O (allocate / write) ─────────────────────────────────────────────
+
+
+def allocate_kv_pool_with_views(
+    num_blocks: int,
+    block_size: int,
+    num_layers: int,
+    num_kv_heads: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    """Allocate vLLM-style paged KV pools plus flat slot-write views.
+
+    The owning tensor follows FlashAttention's block-table cache layout:
+    ``(2, num_blocks, block_size, num_kv_heads, head_dim)``, where dim-0 is
+    ``[K, V]``.  The flat K/V views keep ``slot = block_id * block_size +
+    offset`` writes simple without changing the kernel-facing cache layout.
+    """
+    kv_pools: list[torch.Tensor] = []
+    k_pools: list[torch.Tensor] = []
+    v_pools: list[torch.Tensor] = []
+    for _ in range(num_layers):
+        kv = torch.empty(2, num_blocks, block_size, num_kv_heads, head_dim, dtype=dtype, device=device)
+        kv_pools.append(kv)
+        k_pools.append(kv[0].reshape(num_blocks * block_size, num_kv_heads, head_dim))
+        v_pools.append(kv[1].reshape(num_blocks * block_size, num_kv_heads, head_dim))
+    return kv_pools, k_pools, v_pools
+
+
+def pool_write_chunk(
+    k_pool: torch.Tensor,
+    v_pool: torch.Tensor,
+    new_k: torch.Tensor,
+    new_v: torch.Tensor,
+    slot_mapping: torch.Tensor,
+) -> None:
+    """Write the committed chunk's K/V into pool slots (in place).
+
+    ``new_k`` / ``new_v`` shape: ``(batch, chunk_size, num_kv_heads, head_dim)``.
+    ``slot_mapping`` shape: ``(chunk_size,)`` — one physical slot per token position.
+    The batch dim is written to the same slots (multi-batch writes the same
+    positions; the caller is responsible for per-sequence bookkeeping).
+
+    This is the *commit* write — called once per chunk, after the last denoise step.
+    """
+    k_pool[slot_mapping] = new_k[0]  # batch index 0
+    v_pool[slot_mapping] = new_v[0]
+
+
+def _drop_batch(t: torch.Tensor) -> torch.Tensor:
+    """``(1, L, n_heads, head_dim)`` -> ``(L, n_heads, head_dim)``; pass 3-D through."""
+    if t.dim() == 4 and t.shape[0] == 1:
+        return t[0]
+    return t
+
+
+# ── Chunk-window eviction (spec + manager) ──────────────────────────────────
+
+
+def chunk_window_skipped_tokens(
+    num_computed_tokens: int,
+    *,
+    chunk_size: int,
+    sliding_window: int,
+    sink_chunks: int,
+    reset_at_boundary: bool,
+) -> int:
+    """Tokens outside the resident chunk window, snapped to a chunk boundary.
+
+    Pure function so the eviction policy is unit-testable without constructing a
+    manager. Two strategies:
+
+    - ``reset_at_boundary`` (DreamZero): at each chunk boundary everything past
+      the sink is dropped.
+    - otherwise (VGGT-style sliding replace): keep the last ``window`` tokens
+      (plus the sink); the skip count snaps down to a chunk boundary so a chunk
+      is never half-evicted.
+    """
+    sink = sink_chunks * chunk_size
+    if reset_at_boundary:
+        completed = (num_computed_tokens // chunk_size) * chunk_size
+        return max(0, completed - sink)
+    skipped = max(0, num_computed_tokens - sliding_window - sink)
+    return (skipped // chunk_size) * chunk_size
+
+
+class ChunkWindowManager(SlidingWindowManager):
+    """``SlidingWindowManager`` that evicts at chunk boundaries.
+
+    ``self.sliding_window`` is set by the base ``__init__``; the chunk fields are
+    read from ``self.kv_cache_spec`` (a :class:`ChunkWindowSpec`).
+    """
+
+    def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
+        spec = self.kv_cache_spec
+        return chunk_window_skipped_tokens(
+            num_computed_tokens,
+            chunk_size=spec.chunk_size,
+            sliding_window=self.sliding_window,
+            sink_chunks=spec.sink_chunks,
+            reset_at_boundary=spec.reset_at_boundary,
+        )
+
+
+# Register so KVCacheManager resolves ChunkWindowSpec to ChunkWindowManager.
+# Dispatch walks the spec's MRO, so without explicit registration the subclass
+# would silently fall back to the parent SlidingWindowManager (override ignored).
+# uniform_type_base_spec=None => its own KV cache group.
+@register_kv_cache_spec(manager_class=ChunkWindowManager, uniform_type_base_spec=None)
+@dataclass(frozen=True, kw_only=True)
+class ChunkWindowSpec(SlidingWindowSpec):
+    # sliding_window (inherited) MUST equal window_chunks * chunk_size.
+    chunk_size: int
+    window_chunks: int
+    sink_chunks: int = 0
+    reset_at_boundary: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.sliding_window != self.window_chunks * self.chunk_size:
+            raise ValueError(
+                "ChunkWindowSpec.sliding_window must equal "
+                f"window_chunks * chunk_size ({self.window_chunks} * "
+                f"{self.chunk_size}), got {self.sliding_window}"
+            )

vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py

@@ -0,0 +1,550 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Paged self-attention helpers for AR-Diffusion DreamZero KV reuse."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, ClassVar, NamedTuple
+
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
+
+# Set by ARDiffusionPagedForwardContext.prepare() before each branch forward and
+# read by the fused write+attend custom op below. The pools are process-lifetime
+# allocations owned by one ARDiffusionKVCache per worker process (CFG-parallel /
+# TP ranks are separate processes), mirroring vLLM's forward-context pattern for
+# keeping multi-GiB cache tensors out of the compiled graph's inputs.
+_CURRENT_PAGED_KV_CACHE: Any = None
+
+
+def set_current_paged_kv_cache(kv_cache: Any) -> None:
+    global _CURRENT_PAGED_KV_CACHE
+    _CURRENT_PAGED_KV_CACHE = kv_cache
+
+
+_LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {}
+
+
+def _layer_idx_tensor(layer_idx: int) -> torch.Tensor:
+    t = _LAYER_IDX_TENSORS.get(layer_idx)
+    if t is None:
+        t = torch.tensor(layer_idx, dtype=torch.int64)
+        _LAYER_IDX_TENSORS[layer_idx] = t
+    return t
+
+
+class ARDiffusionPagedLayerInputs(NamedTuple):
+    """Compiled-region payload for one layer's paged self-attention.
+
+    A NamedTuple of plain tensors + ints so ``torch.compile`` treats every field
+    as a pytree graph input (no object-attribute guards, no recompiles when only
+    tensor *values* change). All layers of one branch forward share the same
+    metadata tensor objects, built once by ``prepare()``.
+
+    ``layer_idx`` is a 0-dim CPU tensor, NOT a python int: all 40 DiT blocks
+    share one compiled code object, and an int here becomes a per-layer dynamo
+    value guard (``layer_idx == k``) — 40 cache variants that blow the
+    recompile limit. A tensor input guards on shape/dtype only, so one graph
+    serves every layer.
+    """
+
+    layer_idx: torch.Tensor
+    seq_len: int
+    video_slots: torch.Tensor
+    action_slots: torch.Tensor
+    block_table: torch.Tensor
+    query_start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    max_query_len: int
+    max_seq_len: int
+
+
+@dataclass
+class ARDiffusionPagedForwardContext:
+    """Mutable branch-level state shared by all layer contexts in one forward."""
+
+    kv_cache: Any
+    adapter: Any
+    is_negative: bool
+    history_block_ids: list[int]
+    seq_len: int
+    commit_current: bool
+    max_video_tokens: int
+    current_video_block_ids: list[int] = field(default_factory=list)
+    current_video_slot_mapping: torch.Tensor | None = None
+    action_scratch_block_ids: list[int] = field(default_factory=list)
+    action_slot_mapping: torch.Tensor | None = None
+    query_len: int = 0
+    kv_len: int = 0
+    _allocated_video: bool = False
+    _committed: bool = False
+    _action_len: int = 0
+    # Set once by prepare(); shared by all layers of the branch forward.
+    block_table: torch.Tensor | None = None
+    query_start_loc: torch.Tensor | None = None
+    seq_lens: torch.Tensor | None = None
+    max_query_len: int = 0
+    max_seq_len: int = 0
+    _prepared: bool = False
+
+    @property
+    def block_size(self) -> int:
+        return int(self.kv_cache.block_size)
+
+    @property
+    def num_current_video_blocks(self) -> int:
+        if self.seq_len % self.block_size != 0:
+            raise AssertionError(
+                "AR-Diffusion paged attention expects frame-aligned seq_len "
+                f"(multiple of block_size={self.block_size}), got {self.seq_len}"
+            )
+        return self.seq_len // self.block_size
+
+    def ensure_video_slots(self, device: torch.device) -> None:
+        """Allocate/write targets for the current video tokens, once per branch."""
+        if self._allocated_video:
+            return
+
+        n_blocks = self.num_current_video_blocks
+        if self.commit_current:
+            start = int(self.adapter.num_computed_tokens)
+            self.kv_cache.allocate_token_slots(self.adapter, self.seq_len)
+            table = self.kv_cache.block_table(self.adapter)
+            start_block = start // self.block_size
+            self.current_video_block_ids = [int(b) for b in table[start_block : start_block + n_blocks]]
+            positions = torch.arange(start, start + self.seq_len, dtype=torch.long)
+            self.current_video_slot_mapping = compute_slot_mapping(table, positions, self.block_size).to(device=device)
+        else:
+            self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.is_negative, 0, n_blocks)
+            positions = torch.arange(self.seq_len, dtype=torch.long)
+            self.current_video_slot_mapping = compute_slot_mapping(
+                self.current_video_block_ids,
+                positions,
+                self.block_size,
+            ).to(device=device)
+        self._allocated_video = True
+
+    def ensure_action_slots(self, action_len: int, device: torch.device) -> None:
+        """Reserve scratch slots for action/state K/V, if present."""
+        if action_len <= 0:
+            self.action_scratch_block_ids = []
+            self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
+            self._action_len = 0
+            return
+
+        self.ensure_video_slots(device)
+        if self.action_slot_mapping is not None and self._action_len == action_len:
+            return
+
+        action_blocks = (action_len + self.block_size - 1) // self.block_size
+        scratch_offset = 0 if self.commit_current else len(self.current_video_block_ids)
+        self.action_scratch_block_ids = self.kv_cache.scratch_block_ids(
+            self.is_negative,
+            scratch_offset,
+            action_blocks,
+        )
+        positions = torch.arange(action_len, dtype=torch.long)
+        self.action_slot_mapping = compute_slot_mapping(
+            self.action_scratch_block_ids,
+            positions,
+            self.block_size,
+        ).to(device=device)
+        self._action_len = action_len
+
+    def video_block_table(self, device: torch.device) -> tuple[list[int], int]:
+        self.ensure_video_slots(device)
+        if self.max_video_tokens % self.block_size != 0:
+            raise AssertionError(
+                "AR-Diffusion paged attention requires max_video_tokens to be block-aligned, "
+                f"got max_video_tokens={self.max_video_tokens}, block_size={self.block_size}"
+            )
+        all_video_blocks = self.history_block_ids + self.current_video_block_ids
+        max_video_blocks = self.max_video_tokens // self.block_size
+        visible_video_blocks = all_video_blocks[-max_video_blocks:] if max_video_blocks else []
+        video_len = min(len(all_video_blocks) * self.block_size, self.max_video_tokens)
+        return visible_video_blocks, video_len
+
+    def build_block_table(
+        self,
+        *,
+        action_len: int,
+        query_len: int,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+        """Build FlashAttention block-table metadata for one self-attn call.
+
+        The block table is tail-padded to a fixed width and ``max_seq_len`` is a
+        constant upper bound, so across window growth only tensor *values* change
+        — tensor shapes and the int consts stay stable for ``torch.compile``
+        (and, later, CUDA-graph capture). The kernel only dereferences the first
+        ``ceil(seq_lens/block_size)`` entries, so padding is never read.
+        """
+        video_blocks, video_len = self.video_block_table(device)
+        self.ensure_action_slots(action_len, device)
+        action_blocks = self.action_scratch_block_ids if action_len > 0 else []
+        block_ids = video_blocks + action_blocks
+        if not block_ids:
+            raise RuntimeError("AR-Diffusion paged attention needs at least current video KV blocks")
+
+        # Fixed capacity: full visible video window + one action-capacity block.
+        action_capacity_blocks = max(1, (action_len + self.block_size - 1) // self.block_size)
+        width = max(self.max_video_tokens // self.block_size + action_capacity_blocks, len(block_ids))
+        padded = block_ids + [0] * (width - len(block_ids))
+
+        self.query_len = int(query_len)
+        self.kv_len = int(video_len + action_len)
+        max_seq_len = int(self.max_video_tokens + action_capacity_blocks * self.block_size)
+        block_table = torch.tensor([padded], dtype=torch.int32, device=device)
+        query_start_loc = torch.tensor([0, self.query_len], dtype=torch.int32, device=device)
+        seq_lens = torch.tensor([self.kv_len], dtype=torch.int32, device=device)
+        return block_table, query_start_loc, seq_lens, self.query_len, max_seq_len
+
+    def prepare(self, device: torch.device, action_len: int, query_len: int) -> None:
+        """Host-side, once-per-branch-forward setup (called OUTSIDE torch.compile).
+
+        Allocates the current video/action slots (still lazy: only the branch a
+        CFG-parallel rank actually runs reaches its ``_forward_blocks``), builds
+        the padded block-table metadata ONCE for all layers, and publishes the
+        pool registry for the fused custom op. The compiled per-layer code then
+        only consumes prebuilt tensors via ``ARDiffusionPagedLayerInputs``.
+        """
+        if getattr(self, "_prepared", False):
+            return
+        self.ensure_video_slots(device)
+        (
+            self.block_table,
+            self.query_start_loc,
+            self.seq_lens,
+            self.max_query_len,
+            self.max_seq_len,
+        ) = self.build_block_table(action_len=action_len, query_len=query_len, device=device)
+        if self.action_slot_mapping is None:
+            self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
+        set_current_paged_kv_cache(self.kv_cache)
+        self._prepared = True
+
+    def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs:
+        if not getattr(self, "_prepared", False):
+            raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()")
+        return ARDiffusionPagedLayerInputs(
+            layer_idx=_layer_idx_tensor(layer_idx),
+            seq_len=int(self.seq_len),
+            video_slots=self.current_video_slot_mapping,
+            action_slots=self.action_slot_mapping,
+            block_table=self.block_table,
+            query_start_loc=self.query_start_loc,
+            seq_lens=self.seq_lens,
+            max_query_len=int(self.max_query_len),
+            max_seq_len=int(self.max_seq_len),
+        )
+
+    def mark_committed(self) -> None:
+        self._committed = True
+
+
+@dataclass
+class ARDiffusionPagedLayerContext:
+    """Layer-specific handle passed through DreamZero's existing ``kv_cache`` slot."""
+
+    is_ar_diffusion_paged_context: ClassVar[bool] = True
+    layer_idx: int
+    forward_ctx: ARDiffusionPagedForwardContext
+
+    @property
+    def kv_cache(self):
+        return self.forward_ctx.kv_cache
+
+    @property
+    def adapter(self):
+        return self.forward_ctx.adapter
+
+    @property
+    def is_negative(self) -> bool:
+        return self.forward_ctx.is_negative
+
+    @property
+    def history_block_ids(self) -> list[int]:
+        return self.forward_ctx.history_block_ids
+
+    @property
+    def current_video_block_ids(self) -> list[int]:
+        return self.forward_ctx.current_video_block_ids
+
+    @property
+    def current_video_slot_mapping(self) -> torch.Tensor | None:
+        return self.forward_ctx.current_video_slot_mapping
+
+    @property
+    def action_scratch_block_ids(self) -> list[int]:
+        return self.forward_ctx.action_scratch_block_ids
+
+    @property
+    def action_slot_mapping(self) -> torch.Tensor | None:
+        return self.forward_ctx.action_slot_mapping
+
+    @property
+    def seq_len(self) -> int:
+        return self.forward_ctx.seq_len
+
+    @property
+    def query_len(self) -> int:
+        return self.forward_ctx.query_len
+
+    @property
+    def kv_len(self) -> int:
+        return self.forward_ctx.kv_len
+
+    @property
+    def commit_current(self) -> bool:
+        return self.forward_ctx.commit_current
+
+    def to_layer_inputs(self) -> ARDiffusionPagedLayerInputs:
+        """Compiled-region payload; requires ``forward_ctx.prepare()`` first."""
+        return self.forward_ctx.layer_inputs(self.layer_idx)
+
+
+def is_ar_diffusion_paged_context(value: object) -> bool:
+    return isinstance(value, ARDiffusionPagedLayerContext)
+
+
+def _reference_paged_attention(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    softmax_scale: float,
+    *,
+    causal: bool,
+) -> torch.Tensor:
+    if causal:
+        raise NotImplementedError("AR-Diffusion paged self-attention uses causal=False")
+    outs: list[torch.Tensor] = []
+    block_size = key_cache.shape[1]
+    for i in range(seq_lens.shape[0]):
+        q_start = int(query_start_loc[i].item())
+        q_end = int(query_start_loc[i + 1].item())
+        kv_len = int(seq_lens[i].item())
+        q = query[q_start:q_end]
+        positions = torch.arange(kv_len, device=query.device)
+        logical_blocks = torch.div(positions, block_size, rounding_mode="floor")
+        offsets = positions % block_size
+        physical_blocks = block_table[i, logical_blocks].long()
+        k = key_cache[physical_blocks, offsets]
+        v = value_cache[physical_blocks, offsets]
+        scores = torch.einsum("qhd,khd->hqk", q.float(), k.float()) * float(softmax_scale)
+        probs = torch.softmax(scores, dim=-1).to(v.dtype)
+        outs.append(torch.einsum("hqk,khd->qhd", probs, v))
+    return torch.cat(outs, dim=0)
+
+
+_FA_VERSION_BY_HEAD_SIZE: dict[int, int] = {}
+
+
+def _resolve_fa_version(head_size: int) -> int:
+    # get_flash_attn_version -> current_platform.get_device_capability() is not
+    # dynamo-traceable, and the answer is fixed per head size for the process.
+    version = _FA_VERSION_BY_HEAD_SIZE.get(head_size)
+    if version is None:
+        try:
+            from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
+
+            version = int(get_flash_attn_version(requires_alibi=False, head_size=head_size) or 2)
+        except Exception:
+            version = 2
+        _FA_VERSION_BY_HEAD_SIZE[head_size] = version
+    return version
+
+
+def ar_diffusion_paged_attention(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    *,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    max_query_len: int,
+    max_seq_len: int,
+    softmax_scale: float,
+    causal: bool = False,
+) -> torch.Tensor:
+    """Run non-causal paged attention over a vLLM block table.
+
+    ``query`` may be ``(B, L, H, D)`` or already flattened as ``(T, H, D)``.
+    ``key_cache`` / ``value_cache`` are ``(num_blocks, block_size, H, D)``.
+    """
+    batched = query.dim() == 4
+    if batched:
+        batch, q_len = query.shape[:2]
+        query_flat = query.reshape(batch * q_len, *query.shape[2:])
+    else:
+        query_flat = query
+
+    if not query_flat.is_cuda:
+        out = _reference_paged_attention(
+            query_flat,
+            key_cache,
+            value_cache,
+            block_table,
+            query_start_loc,
+            seq_lens,
+            softmax_scale,
+            causal=causal,
+        )
+    else:
+        from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+        fa_version = _resolve_fa_version(query_flat.shape[-1])
+
+        out = torch.empty_like(query_flat)
+        flash_attn_varlen_func(
+            q=query_flat,
+            k=key_cache,
+            v=value_cache,
+            out=out,
+            cu_seqlens_q=query_start_loc,
+            max_seqlen_q=int(max_query_len),
+            seqused_k=seq_lens,
+            max_seqlen_k=int(max_seq_len),
+            softmax_scale=float(softmax_scale),
+            causal=causal,
+            block_table=block_table,
+            fa_version=fa_version,
+        )
+
+    if batched:
+        return out.reshape(query.shape)
+    return out
+
+
+# ── Fused write+attend custom op (torch.compile-safe) ──────────────────────
+#
+# One opaque op per layer keeps the compiled DiT block fullgraph: dynamo treats
+# it as a single graph node (no eager island, no graph breaks), and the K/V slot
+# writes happen inside the op so write→read ordering with the block-table kernel
+# is internal — no hidden-mutation ordering hazards against a separate reader op
+# and no multi-GiB pool tensors as graph inputs (the flat and paged pool views
+# alias the same storage, which AOTAutograd handles poorly as inputs). Pattern
+# follows sage_attn3.py in-repo and vLLM's own unified attention ops.
+def _paged_write_attn_impl(
+    query: torch.Tensor,
+    k_curr: torch.Tensor,
+    v_curr: torch.Tensor,
+    k_act: torch.Tensor | None,
+    v_act: torch.Tensor | None,
+    layer_idx: torch.Tensor,
+    video_slots: torch.Tensor,
+    action_slots: torch.Tensor,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    max_query_len: int,
+    max_seq_len: int,
+    softmax_scale: float,
+) -> torch.Tensor:
+    kv = _CURRENT_PAGED_KV_CACHE
+    if kv is None:
+        raise RuntimeError("ar_diffusion_paged_write_attn called before prepare() set the KV pool registry")
+    layer_idx = int(layer_idx)
+    k_pool = kv._k_pools[layer_idx]
+    v_pool = kv._v_pools[layer_idx]
+    k_pool[video_slots] = k_curr.to(k_pool.dtype)
+    v_pool[video_slots] = v_curr.to(v_pool.dtype)
+    if k_act is not None and v_act is not None and k_act.shape[0] > 0:
+        k_pool[action_slots] = k_act.to(k_pool.dtype)
+        v_pool[action_slots] = v_act.to(v_pool.dtype)
+    return ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(layer_idx),
+        kv.value_cache(layer_idx),
+        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=softmax_scale,
+        causal=False,
+    )
+
+
+# hasattr guard keeps registration idempotent across test re-imports that pop
+# the module from sys.modules (same as sage_attn3.py).
+if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
+
+    @torch.library.custom_op("vllm_omni::ar_diffusion_paged_write_attn", mutates_args=())
+    def _paged_write_attn_op(
+        query: torch.Tensor,
+        k_curr: torch.Tensor,
+        v_curr: torch.Tensor,
+        k_act: torch.Tensor | None,
+        v_act: torch.Tensor | None,
+        layer_idx: torch.Tensor,
+        video_slots: torch.Tensor,
+        action_slots: torch.Tensor,
+        block_table: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        max_query_len: int,
+        max_seq_len: int,
+        softmax_scale: float,
+    ) -> torch.Tensor:
+        return _paged_write_attn_impl(
+            query,
+            k_curr,
+            v_curr,
+            k_act,
+            v_act,
+            layer_idx,
+            video_slots,
+            action_slots,
+            block_table,
+            query_start_loc,
+            seq_lens,
+            max_query_len,
+            max_seq_len,
+            softmax_scale,
+        )
+
+    @_paged_write_attn_op.register_fake
+    def _(
+        query,
+        k_curr,
+        v_curr,
+        k_act,
+        v_act,
+        layer_idx,
+        video_slots,
+        action_slots,
+        block_table,
+        query_start_loc,
+        seq_lens,
+        max_query_len,
+        max_seq_len,
+        softmax_scale,
+    ):
+        return torch.empty_like(query)
+
+
+def paged_write_attn(
+    inputs: ARDiffusionPagedLayerInputs, query, k_curr, v_curr, k_act, v_act, softmax_scale: float
+) -> torch.Tensor:
+    """Model-facing entry: routes through the custom op (traceable in fullgraph)."""
+    return torch.ops.vllm_omni.ar_diffusion_paged_write_attn(
+        query,
+        k_curr,
+        v_curr,
+        k_act,
+        v_act,
+        inputs.layer_idx,
+        inputs.video_slots,
+        inputs.action_slots,
+        inputs.block_table,
+        inputs.query_start_loc,
+        inputs.seq_lens,
+        inputs.max_query_len,
+        inputs.max_seq_len,
+        softmax_scale,
+    )

vllm_omni/experimental/ar_diffusion/kv_cache/state.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DreamZero-facing AR-Diffusion KV bridge backed by paged self-attention."""
+
+from __future__ import annotations
+
+from vllm.logger import init_logger
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+    ARDiffusionPagedForwardContext,
+    ARDiffusionPagedLayerContext,
+)
+
+_log = init_logger(__name__)
+
+
+class ARDiffusionKVState:
+    """Per-request bridge between a DreamZero rollout and the AR-Diffusion KV pool.
+
+    The self-attention path is paged-only: ``get_kv_caches`` returns lightweight
+    per-layer contexts, and each attention layer writes current K/V into the
+    pool before calling the FlashAttention block-table kernel.  No resident
+    window is gathered into a contiguous tensor.
+    """
+
+    def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
+        """One state per DreamZero session, holding BOTH CFG branches' adapters.
+
+        ``pos``/``neg`` are the classifier-free-guidance branches: ``pos`` is
+        the conditional pass (real prompt), ``neg`` the unconditional pass
+        (negative prompt). Each owns an independent request adapter — and
+        therefore independent pool blocks — because the two branches produce
+        different K/V for the same frames. Every ``is_negative`` flag in this
+        module selects between them (``False`` -> ``pos``, ``True`` -> ``neg``).
+        """
+        self.kv_cache = kv_cache
+        self.pos = pos_adapter
+        self.neg = neg_adapter
+        self.num_layers = num_layers
+        self._committed: dict[bool, int] = {False: 0, True: 0}
+        self._paged_pending: dict[bool, ARDiffusionPagedForwardContext | None] = {False: None, True: None}
+        # Cross-attn KV is populated once after text encoding (eager), then
+        # read every denoising step. Tracked per half: the text K/V depend only
+        # on the session prompt (survive window-boundary resets), while the I2V
+        # image-token K/V come from the current observation's CLIP features and
+        # must be re-projected on every window restart.
+        self._cross_text_populated: dict[bool, bool] = {False: False, True: False}
+        self._cross_img_populated: dict[bool, bool] = {False: False, True: False}
+
+    def _adapter(self, is_negative: bool):
+        return self.neg if is_negative else self.pos
+
+    def get_kv_caches(
+        self,
+        is_negative: bool,
+        seq_len: int | None = None,
+        commit_current: bool = False,
+    ) -> list[ARDiffusionPagedLayerContext]:
+        if seq_len is None:
+            raise ValueError("AR-Diffusion paged self-attention requires seq_len in get_kv_caches()")
+        return self.prepare_paged_context(is_negative, seq_len, commit_current)
+
+    def prepare_paged_context(
+        self,
+        is_negative: bool,
+        seq_len: int,
+        commit_current: bool,
+    ) -> list[ARDiffusionPagedLayerContext]:
+        """Return per-layer paged attention contexts for one branch forward.
+
+        Allocation is lazy: constructing both CFG branches' kwargs must not
+        allocate manager blocks for a branch that a CFG-parallel rank will not run.
+        The first self-attn layer that consumes the context allocates the current
+        video slots, then all layers share that branch-level state.
+        """
+        cs = self.kv_cache.spec.chunk_size
+        if int(seq_len) % cs != 0:
+            raise AssertionError(
+                f"AR-Diffusion expects frame-aligned seq_len (multiple of chunk_size={cs}), got {seq_len}"
+            )
+
+        pending = self._paged_pending.get(is_negative)
+        if pending is not None and pending.commit_current and pending._allocated_video and not pending._committed:
+            raise RuntimeError("AR-Diffusion paged context replaced before its managed current chunk was committed")
+
+        adapter = self._adapter(is_negative)
+        branch = "neg" if is_negative else "pos"
+        forward_ctx = ARDiffusionPagedForwardContext(
+            kv_cache=self.kv_cache,
+            adapter=adapter,
+            is_negative=is_negative,
+            history_block_ids=self.kv_cache.window_block_ids(adapter),
+            seq_len=int(seq_len),
+            commit_current=bool(commit_current),
+            max_video_tokens=int(self.kv_cache.spec.sliding_window),
+        )
+        self._paged_pending[is_negative] = forward_ctx
+        _log.debug(
+            "AR-Diffusion GET   [%s] source=paged-attn layers=%d history_blocks=%d seq_len=%d commit_current=%s",
+            branch,
+            self.num_layers,
+            len(forward_ctx.history_block_ids),
+            int(seq_len),
+            bool(commit_current),
+        )
+        return [ARDiffusionPagedLayerContext(layer_idx=i, forward_ctx=forward_ctx) for i in range(self.num_layers)]
+
+    def commit_paged_context(self, is_negative: bool) -> None:
+        """Commit the managed current video blocks after a successful forward."""
+        ctx = self._paged_pending.get(is_negative)
+        if ctx is None:
+            return
+        branch = "neg" if is_negative else "pos"
+        if ctx.commit_current and ctx._allocated_video:
+            n_chunks = ctx.seq_len // self.kv_cache.spec.chunk_size
+            for _ in range(n_chunks):
+                ctx.adapter.on_chunk_committed()
+            self._committed[is_negative] += ctx.seq_len
+            _log.debug(
+                "AR-Diffusion COMMIT [%s] paged-attn new_tokens=%d chunks=%d resident=%d/%d",
+                branch,
+                ctx.seq_len,
+                n_chunks,
+                len(self.kv_cache.window_block_ids(ctx.adapter)),
+                self.kv_cache.spec.window_chunks,
+            )
+        ctx.mark_committed()
+        self._paged_pending[is_negative] = None
+
+    def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
+        """Return pool-backed cross-attn cache dicts.
+
+        Under AR-Diffusion the cross-attn pool is always populated by ``_kv_populate_cross``
+        before the first read, so this must never be reached unpopulated — if it
+        were, the model would lazily project and write cross KV itself, violating
+        engine ownership. Fail loud rather than fall back to the model.
+        """
+        img_ok = not self.kv_cache._cross_k_img or self._cross_img_populated[is_negative]
+        if not (self._cross_text_populated[is_negative] and img_ok and self.kv_cache.cross_attn_length > 0):
+            raise RuntimeError(
+                f"AR-Diffusion cross-attn read before _kv_populate_cross (neg={is_negative}, "
+                f"cross_attn_length={self.kv_cache.cross_attn_length}, "
+                f"text={self._cross_text_populated[is_negative]} "
+                f"img={self._cross_img_populated[is_negative]}) — the engine must own all cross KV"
+            )
+        return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
+
+    def close(self) -> None:
+        """Final teardown: free both branches' resident pool blocks."""
+        self.kv_cache.end_request(self.pos)
+        self.kv_cache.end_request(self.neg)
+        self._paged_pending = {False: None, True: None}
+
+    def reset(self, *, keep_cross_text: bool = False) -> None:
+        """Drop this session's KV — mirrors the model-local ``state.reset()``.
+
+        DreamZero resets at the attention-window boundary (``should_reset``); the
+        model starts a fresh sliding window, so AR-Diffusion must free the resident pool
+        blocks and start a fresh adapter for each branch. The ``ARDiffusionKVState``
+        object (and thus ``pipeline._ar_diffusion_kv_state``) is preserved across the reset
+        so the runner's session mapping stays valid — only the pool-backed state
+        is recycled.
+
+        Args:
+            keep_cross_text: On window ("inference") resets the session prompt is
+                unchanged, so the text cross-attn K/V in the static pool buffers are
+                still valid — keep them and only force the image half to repopulate.
+                Session resets (prompt change / explicit reset) pass ``False``.
+        """
+        pos_id, neg_id = self.pos.request_id, self.neg.request_id
+        self.close()
+        self.pos = self.kv_cache.begin_request(pos_id)
+        self.neg = self.kv_cache.begin_request(neg_id)
+        self._committed = {False: 0, True: 0}
+        self._paged_pending = {False: None, True: None}
+        if not keep_cross_text:
+            self._cross_text_populated = {False: False, True: False}
+        self._cross_img_populated = {False: False, True: False}
+        _log.info(
+            "AR-Diffusion RESET [%s/%s] session KV cleared (window boundary; cross text %s)",
+            pos_id,
+            neg_id,
+            "kept" if keep_cross_text else "cleared",
+        )

vllm_omni/experimental/ar_diffusion/runner.py

@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ARDiffusionModelRunner — the diffusion model runner for the AR-Diffusion engine.
+
+Subclasses ``DiffusionModelRunner`` and owns the engine-level KV cache
+(:class:`ARDiffusionKVCache`). It brackets a request's rollout with the KV lifecycle
+(``begin_request`` / ``end_request``) and exposes the live ``ARDiffusionKVCache`` so the
+model pipeline (DreamZero) can drive the per-chunk allocate / slot-mapping /
+gather / commit operations during ``pipeline.forward``.
+
+Selecting the AR-Diffusion engine enables KV management (no env gate); optional KV
+parameters come from the deploy yaml's ``model_config["ar_diffusion_kv_config"]``.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import math
+import time
+from collections import OrderedDict
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import MAX_DREAMZERO_SESSIONS
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+logger = init_logger(__name__)
+
+
+def resolve_ar_diffusion_kv_config(od_config: OmniDiffusionConfig) -> ARDiffusionKVConfig:
+    """Resolve the AR-Diffusion KV config from the deploy/command arg.
+
+    Selecting the AR-Diffusion engine implies KV management, so the result is always
+    enabled. Optional overrides (``window_chunks`` / ``gpu_memory_fraction`` /
+    ``sink_chunks`` / ``reset_at_boundary``) may be supplied via the deploy yaml as
+    ``model_config["ar_diffusion_kv_config"]``. ``chunk_size`` / ``window_chunks`` are
+    finalized from the model geometry at load (see ``_preallocate_kv_cache``).
+    """
+    raw = getattr(od_config, "ar_diffusion_kv_config", None)
+    if raw is None:
+        model_config = getattr(od_config, "model_config", None)
+        if isinstance(model_config, dict):
+            raw = model_config.get("ar_diffusion_kv_config")
+    if isinstance(raw, ARDiffusionKVConfig):
+        return dataclasses.replace(raw, enable=True)
+    if isinstance(raw, dict):
+        return ARDiffusionKVConfig(**{**raw, "enable": True})
+    return ARDiffusionKVConfig(enable=True)
+
+
+class ARDiffusionModelRunner(DiffusionModelRunner):
+    def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None:
+        super().__init__(vllm_config, od_config, device)
+        self.ar_diffusion_kv_config = resolve_ar_diffusion_kv_config(od_config)
+        # Built after the model is loaded (dimensions known); stays None while KV
+        # management is disabled.
+        self.kv_cache: ARDiffusionKVCache | None = None
+        # DreamZero KV is session-scoped (persists across a session's forwards),
+        # so AR-Diffusion KV state is keyed by session_id and reused, not created per request.
+        # Bounded by the same LRU cap as DreamZeroPipeline._states: an evicted
+        # session's ARDiffusionKVState owns pool blocks (two CFG adapters), so without a
+        # bound, session-id churn in a long-running server would leak pool ownership
+        # until the KV pool is exhausted. Evicting frees those blocks (state.close()).
+        self._ar_diffusion_states: OrderedDict[str, ARDiffusionKVState] = OrderedDict()
+        self._max_ar_diffusion_states = MAX_DREAMZERO_SESSIONS
+        # Per-request server-side forward E2E times (seconds), recorded in
+        # execute_model. This is the worker-side compute time — the analog of the
+        # upstream server's per-request E2E — and excludes the engine<->worker IPC
+        # that the client's omni.generate wall time includes. Read (and cleared)
+        # by the ar_diffusion_perf_stats worker RPC for the perf-compare summary.
+        self._perf_e2e_times: list[float] = []
+
+    def build_kv_cache(
+        self,
+        *,
+        num_layers: int,
+        num_kv_heads: int,
+        head_size: int,
+        dtype,
+        block_size: int,
+        max_model_len: int,
+        available_bytes: int,
+        cross_attn_length: int = 0,
+        cross_attn_img_length: int = 0,
+        local_branches: int = 2,
+        num_frame_per_block: int = 1,
+    ) -> None:
+        """Construct the ARDiffusionKVCache (preallocating GPU pools on ``self.device``).
+
+        No-op when KV management is disabled.
+        """
+        if not self.ar_diffusion_kv_config.enable:
+            return
+        self.kv_cache = ARDiffusionKVCache(
+            self.ar_diffusion_kv_config,
+            num_layers=num_layers,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=dtype,
+            block_size=block_size,
+            max_model_len=max_model_len,
+            available_bytes=available_bytes,
+            cross_attn_length=cross_attn_length,
+            cross_attn_img_length=cross_attn_img_length,
+            local_branches=local_branches,
+            num_frame_per_block=num_frame_per_block,
+            device=self.device,
+        )
+        logger.info(
+            "AR-Diffusion KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s, "
+            "layers=%d kv_heads=%d head_dim=%d block_size=%d cross_attn_len=%d device=%s",
+            self.kv_cache.num_blocks,
+            self.ar_diffusion_kv_config.chunk_size,
+            self.ar_diffusion_kv_config.window_chunks,
+            num_layers,
+            num_kv_heads,
+            head_size,
+            block_size,
+            cross_attn_length,
+            self.device,
+        )
+
+    # -- preallocation at load -------------------------------------------------
+
+    def load_model(self, *args, **kwargs):
+        super().load_model(*args, **kwargs)
+        if self.ar_diffusion_kv_config.enable and self.pipeline is not None:
+            self._preallocate_kv_cache()
+            # Pre-capture the CUDA graphs for every window-fill shape at load time
+            # (only when compiling/cuda-graph is on) so serving is fast from chunk 0.
+            if not self.od_config.enforce_eager and self.ar_diffusion_kv_config.warmup_cudagraph:
+                self._warmup_ar_rollout()
+
+    def _infer_frame_seqlen(self) -> int:
+        """frame_seqlen = (H//8)*(W//8)//4 from the configured image_resolution."""
+        mc = getattr(self.od_config, "model_config", None) or {}
+        psc = (mc.get("policy_server_config") if isinstance(mc, dict) else None) or {}
+        res = psc.get("image_resolution", [180, 320])
+        h, w = int(res[0]), int(res[1])
+        return (h // 8) * (w // 8) // 4
+
+    def _preallocate_kv_cache(self) -> None:
+        """Build the KV pool once, from the loaded DreamZero transformer geometry."""
+        # The paged self-attention path carries one rollout per request: slot
+        # mappings/adapters are single-sequence and K/V writes address batch
+        # index 0 only. Reject a multi-sequence config at init instead of
+        # failing (or silently dropping sequences) mid-forward.
+        max_num_seqs = int(getattr(self.od_config, "max_num_seqs", 1) or 1)
+        if max_num_seqs > 1:
+            raise ValueError(
+                "AR-Diffusion paged KV supports max_num_seqs=1 (single-sequence "
+                f"rollouts); got max_num_seqs={max_num_seqs}. Per-sequence "
+                "adapters/slot maps are not implemented."
+            )
+        t = self.pipeline.transformer
+        num_layers = int(t.num_layers)
+        num_kv_heads = int(getattr(t.blocks[0].self_attn, "tp_num_heads", t.num_heads))
+        head_size = int(t.dim // t.num_heads)
+        num_frame_per_block = int(t.num_frame_per_block)
+        local_attn_size = int(t.local_attn_size)
+        # The loaded transformer is the source of truth for frame_seqlen (it derives
+        # max_attention_size from it). The deploy-config inference is only a
+        # cross-check — warn on drift rather than silently sizing chunks wrong.
+        frame_seqlen = int(getattr(t, "frame_seqlen", 0)) or self._infer_frame_seqlen()
+        inferred = self._infer_frame_seqlen()
+        if frame_seqlen != inferred:
+            logger.warning(
+                "AR-Diffusion frame_seqlen mismatch: transformer=%d deploy-config=%d; using transformer",
+                frame_seqlen,
+                inferred,
+            )
+        # The model's own attention window, in tokens (= the [-max_attention_size:]
+        # slice it applies). Read it directly rather than recomputing.
+        max_attention_size = int(t.blocks[0].self_attn.max_attention_size)
+        # Cross-attn text length — static pool for the text-encoder output.
+        # The cached k/v spans the full text sequence (the image tokens
+        # prepended in _forward_blocks are stripped inside the cross-attn
+        # forward but are not part of the cached k/v).
+        cross_attn_length = int(getattr(t, "text_len", 0))
+        # I2V image-token cross-attn length: the model splits context[:, :257] as the
+        # image tokens (the img_emb(clip_feature) output), caching k_img/v_img. T2V: 0.
+        cross_attn_img_length = 257 if getattr(t, "model_type", "t2v") == "i2v" else 0
+
+        # Frame-granular paging: 1 block = 1 frame = frame_seqlen tokens, so the
+        # resident window matches max_attention_size exactly (it need not be a whole
+        # number of num_frame_per_block causal blocks).
+        chunk_size = frame_seqlen
+        window_chunks = self.ar_diffusion_kv_config.window_chunks or (max_attention_size // frame_seqlen)
+
+        self.ar_diffusion_kv_config = dataclasses.replace(
+            self.ar_diffusion_kv_config, chunk_size=chunk_size, window_chunks=window_chunks
+        )
+        free_bytes = torch.cuda.mem_get_info(self.device)[0]
+        # Under CFG-parallel each rank executes exactly ONE branch (rank0 pos,
+        # rank1 neg; the other branch's lazy contexts never allocate on this
+        # rank), so its pool only needs one branch's capacity. A single-process
+        # run (cfg world 1) executes both branches from one pool.
+        try:
+            from vllm_omni.diffusion.distributed.parallel_state import (
+                get_classifier_free_guidance_world_size,
+            )
+
+            cfg_world = int(get_classifier_free_guidance_world_size())
+        except Exception:
+            cfg_world = 1
+        local_branches = 1 if cfg_world >= 2 else 2
+        logger.info(
+            "AR-Diffusion preallocating (paged): frame_seqlen=%d num_frame_per_block=%d "
+            "local_attn_size=%d -> chunk_size=%d window_chunks=%d (window=%d tokens)",
+            frame_seqlen,
+            num_frame_per_block,
+            local_attn_size,
+            chunk_size,
+            window_chunks,
+            window_chunks * chunk_size,
+        )
+        self.build_kv_cache(
+            num_layers=num_layers,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=self.od_config.dtype,
+            block_size=chunk_size,  # one pool block per chunk
+            max_model_len=1 << 20,
+            available_bytes=free_bytes,
+            cross_attn_length=cross_attn_length,
+            cross_attn_img_length=cross_attn_img_length,
+            local_branches=local_branches,
+            num_frame_per_block=num_frame_per_block,
+        )
+
+    def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        # KV disabled -> base behavior, unchanged.
+        if self.kv_cache is None:
+            return super().execute_model(req)
+
+        kv = self.kv_cache
+        # DreamZero KV is session-scoped (the model state persists across a
+        # session's forwards), so the AR-Diffusion KV state is keyed by session_id and
+        # reused — matching how pipeline.forward resolves the model-local state.
+        extra_args = req.sampling_params.extra_args or {}
+        session_id = str(extra_args.get("session_id") or "default")
+        state = self._ar_diffusion_states.get(session_id)
+        if state is None:
+            pos = kv.begin_request(f"bde__{session_id}")
+            neg = kv.begin_request(f"bde__{session_id}__neg")
+            state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers)
+            self._ar_diffusion_states[session_id] = state
+            # Evict the least-recently-used session(s) past the cap, freeing their
+            # pool blocks — mirrors DreamZeroPipeline._states so the AR-Diffusion pool can't
+            # outlive the model-local state map it shadows.
+            while len(self._ar_diffusion_states) > self._max_ar_diffusion_states:
+                old_id, old_state = self._ar_diffusion_states.popitem(last=False)
+                old_state.close()
+                logger.debug("AR-Diffusion evicted session=%s (LRU); freed pool blocks", old_id)
+        # Track recency on every access (hit or miss) so the LRU order is correct.
+        self._ar_diffusion_states.move_to_end(session_id)
+        logger.debug(
+            "AR-Diffusion execute_model: req=%s session=%s chunk_size=%d window_chunks=%s num_blocks=%d",
+            req.request_id,
+            session_id,
+            kv.spec.chunk_size,
+            kv.config.window_chunks,
+            kv.num_blocks,
+        )
+        self.pipeline._ar_diffusion_kv_state = state
+        # Time the worker-side forward (true per-request compute E2E). The CUDA
+        # sync below makes the stop reflect completed GPU work, not just dispatch.
+        _e2e_t0 = time.perf_counter()
+        try:
+            out = super().execute_model(req)
+        except Exception:
+            # Transactional containment: a forward that died partway may have
+            # written some layers' K/V into allocated-but-uncommitted blocks
+            # and advanced model-local state. Tear the whole session down
+            # (engine pool blocks freed via close(); model-local state dropped)
+            # so the next request with this session id starts clean instead of
+            # tripping the pending-commit guard or reading half-written KV.
+            broken = self._ar_diffusion_states.pop(session_id, None)
+            if broken is not None:
+                broken.close()
+            states = getattr(self.pipeline, "_states", None)
+            if states is not None:
+                states.pop(session_id, None)
+            logger.warning(
+                "AR-Diffusion execute_model failed for session=%s; session state torn down "
+                "(pool blocks freed) — the next request starts a fresh session",
+                session_id,
+            )
+            raise
+        finally:
+            self.pipeline._ar_diffusion_kv_state = None
+        if self.device is not None and torch.cuda.is_available():
+            torch.accelerator.synchronize(self.device)
+        self._perf_e2e_times.append(time.perf_counter() - _e2e_t0)
+        return out
+
+    # -- cuda-graph warm-up ----------------------------------------------------
+
+    #: Dedicated session id for the warm-up rollout; cleaned from both state maps.
+    _WARMUP_SID = "__ardiffusion_warmup__"
+
+    def _synth_robot_obs(self, h: int, w: int, n_frames: int) -> dict:
+        """Zero-filled roboarena observation matching the client's obs layout.
+
+        Mirrors ``examples/.../client_schedule.make_obs_from_video`` (2 ext + wrist
+        cameras, 7/6/1 state vectors) so it survives ``_transform_robot_obs`` exactly
+        like a real request. One frame -> ``(H,W,3)``; chunk -> ``(n_frames,H,W,3)``.
+        """
+        img = np.zeros((h, w, 3), dtype=np.uint8) if n_frames == 1 else np.zeros((n_frames, h, w, 3), dtype=np.uint8)
+        return {
+            "observation/exterior_image_0_left": img,
+            "observation/exterior_image_1_left": img,
+            "observation/wrist_image_left": img,
+            "observation/joint_position": np.zeros(7, dtype=np.float32),
+            "observation/cartesian_position": np.zeros(6, dtype=np.float32),
+            "observation/gripper_position": np.zeros(1, dtype=np.float32),
+            "prompt": "warmup",
+            "session_id": self._WARMUP_SID,
+        }
+
+    def _warmup_ar_rollout(self) -> None:
+        """Drive a synthetic rollout through one window-fill so every DiT graph is
+        captured at load (off the serving hot path). Frees the synthetic session and
+        leaves the KV pool exactly as found. Never raises — falls back to lazy capture.
+        """
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        kv = self.kv_cache
+        if kv is None:
+            return
+        sid = self._WARMUP_SID
+        t = self.pipeline.transformer
+        nfpb = int(t.num_frame_per_block)
+        window_chunks = int(kv.config.window_chunks or 0)
+        if window_chunks <= 0:
+            return
+        # Distinct DiT window shapes are walked by one fill: prefill + ceil((W-1)/nfpb)
+        # chunk forwards. +1 optionally to also capture the reset-cycle forward.
+        n_forwards = 1 + math.ceil(max(0, window_chunks - 1) / nfpb)
+        if self.ar_diffusion_kv_config.warmup_capture_reset:
+            n_forwards += 1
+        res = self._image_resolution()
+        h, w = int(res[0]), int(res[1])
+
+        free_before = kv.manager.block_pool.get_num_free_blocks()
+        # The step cache is left active: the DiT input shape is per-chunk (the resident
+        # window length), and the first denoise step of every chunk always computes, so
+        # every window shape is captured regardless of how many later steps are skipped.
+        logger.info(
+            "AR-Diffusion cudagraph warm-up: up to %d synthetic forwards (window_chunks=%d, nfpb=%d)",
+            n_forwards,
+            window_chunks,
+            nfpb,
+        )
+        try:
+            for i in range(n_forwards):
+                n_frames = 1 if i == 0 else 4  # client convention: 1-frame prefill, 4-frame chunks
+                obs = self._synth_robot_obs(h, w, n_frames)
+                sp = OmniDiffusionSamplingParams(extra_args={"reset": i == 0, "session_id": sid, "robot_obs": obs})
+                req = OmniDiffusionRequest(prompts=["warmup"], sampling_params=sp, request_id=f"ardiffusion-warmup-{i}")
+                self.execute_model(req)
+                # Stop once the resident window is full — the remaining shapes are
+                # already captured (the window caps/resets through the same set).
+                ar_state = self._ar_diffusion_states.get(sid)
+                if (
+                    not self.ar_diffusion_kv_config.warmup_capture_reset
+                    and ar_state is not None
+                    and len(kv.window_block_ids(ar_state.pos)) >= window_chunks
+                ):
+                    break
+        except Exception as e:  # noqa: BLE001 — warm-up must never break model load
+            logger.warning("AR-Diffusion cudagraph warm-up failed (%s); using lazy capture.", e)
+        finally:
+            # Free the synthetic session from BOTH maps (KV pool + pipeline state).
+            ar_state = self._ar_diffusion_states.pop(sid, None)
+            if ar_state is not None:
+                try:
+                    ar_state.close()  # return KV blocks to the BlockPool
+                except Exception as e:  # noqa: BLE001
+                    logger.warning("AR-Diffusion warm-up: state.close() failed (%s)", e)
+            try:
+                self.pipeline._states.pop(sid, None)
+            except Exception:  # noqa: BLE001
+                pass
+            # Drop any per-forward timings the warm-up recorded so serving stats start clean.
+            self._perf_e2e_times.clear()
+            free_after = kv.manager.block_pool.get_num_free_blocks()
+            if free_after != free_before:
+                logger.warning(
+                    "AR-Diffusion warm-up: KV pool not restored (free %d -> %d) — possible leak",
+                    free_before,
+                    free_after,
+                )
+            else:
+                logger.info(
+                    "AR-Diffusion cudagraph warm-up complete; KV pool restored (%d free blocks)",
+                    free_after,
+                )
+
+    def _image_resolution(self) -> list[int]:
+        mc = getattr(self.od_config, "model_config", None) or {}
+        psc = (mc.get("policy_server_config") if isinstance(mc, dict) else None) or {}
+        return psc.get("image_resolution", [180, 320])

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -14,7 +14,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models.
 from __future__ import annotations
 
 import dataclasses
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 
 import torch
 import torch.nn as nn
@@ -742,6 +742,27 @@ class CodePredictorWrapper(nn.Module):
     #  Forward -- re-prefill + inline sampling
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _multinomial(
+        probs: torch.Tensor,
+        generator: torch.Generator | None,
+        generators: Sequence[torch.Generator | None] | None,
+    ) -> torch.Tensor:
+        """Sample one code per row, optionally with per-row generators.
+
+        Per-row generators keep explicitly-seeded requests deterministic in a
+        multi-row batch: each row consumes draws only from its own generator,
+        so the transformer forward can stay batched (#4883).
+        """
+        if generators is None:
+            return torch.multinomial(probs, num_samples=1, generator=generator)
+        return torch.cat(
+            [
+                torch.multinomial(probs[row : row + 1], num_samples=1, generator=row_generator)
+                for row, row_generator in enumerate(generators)
+            ]
+        )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -753,9 +774,12 @@ class CodePredictorWrapper(nn.Module):
         top_k: int = 50,
         top_p: float = 1.0,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """Predict residual codebooks 1..G-1 autoregressively via re-prefill."""
         bsz = int(layer0_code.shape[0])
+        if generators is not None and len(generators) != bsz:
+            raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
         num_groups = self._num_groups
         device = layer0_code.device
 
@@ -846,7 +870,7 @@ class CodePredictorWrapper(nn.Module):
                     sorted_logits[remove_mask] = float("-inf")
                     logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
                 probs = F.softmax(logits, dim=-1, dtype=torch.float32)
-                code = torch.multinomial(probs, num_samples=1, generator=generator)
+                code = self._multinomial(probs, generator, generators)
             else:
                 # "per_call" mode: temperature-scaled + top-k
                 if use_sampling:
@@ -855,7 +879,7 @@ class CodePredictorWrapper(nn.Module):
                         topk_vals, _ = scaled.topk(top_k, dim=-1)
                         scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
                     probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
-                    code = torch.multinomial(probs, num_samples=1, generator=generator)
+                    code = self._multinomial(probs, generator, generators)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py

@@ -22,6 +22,7 @@ generation.
 
 from __future__ import annotations
 
+import json
 import os
 import re
 from typing import Any
@@ -170,14 +171,25 @@ _AUDIO_TOKENIZER_PATH_ENVS = (
 )
 
 
+def _is_higgs_audio_tokenizer_config(config_path: str) -> bool:
+    try:
+        with open(config_path, encoding="utf-8") as f:
+            config = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return False
+    return config.get("model_type") == "higgs_audio_v2_tokenizer"
+
+
 def _normalize_audio_tokenizer_dir(path: str) -> str | None:
     if not path:
         return None
     expanded = os.path.abspath(os.path.expanduser(path))
-    if os.path.isfile(os.path.join(expanded, "config.json")):
+    config_path = os.path.join(expanded, "config.json")
+    if os.path.isfile(config_path) and _is_higgs_audio_tokenizer_config(config_path):
         return expanded
     nested = os.path.join(expanded, _K2_OMNIVOICE_SUBDIR)
-    if os.path.isfile(os.path.join(nested, "config.json")):
+    nested_config_path = os.path.join(nested, "config.json")
+    if os.path.isfile(nested_config_path) and _is_higgs_audio_tokenizer_config(nested_config_path):
         return nested
     return None
 
@@ -201,7 +213,9 @@ def _resolve_audio_tokenizer_dir() -> str | None:
         filename=f"{_K2_OMNIVOICE_SUBDIR}/config.json",
     )
     if isinstance(cached_config, str) and os.path.isfile(cached_config):
-        return os.path.dirname(cached_config)
+        candidate = _normalize_audio_tokenizer_dir(os.path.dirname(cached_config))
+        if candidate is not None:
+            return candidate
 
     from huggingface_hub.constants import HF_HUB_CACHE
 
@@ -240,7 +254,9 @@ def _load_audio_tokenizer():
             _K2_OMNIVOICE_REPO,
             allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"],
         )
-        audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR)
+        audio_tokenizer_dir = _normalize_audio_tokenizer_dir(repo_path)
+        if audio_tokenizer_dir is None:
+            raise RuntimeError(f"Downloaded {_K2_OMNIVOICE_REPO} does not contain a valid Higgs audio tokenizer")
     device = current_omni_platform.get_torch_device()
     model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device)
     return model.eval()

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any
 
 import numpy as np
@@ -318,6 +318,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # OmniGPUModelRunner will store talker_mtp output under this key in
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
+        # talker_mtp samples with per-row generators, so explicitly-seeded
+        # requests stay batched instead of one scalar forward per row (#4883).
+        # Only valid while talker_mtp receives the unpadded active batch (this
+        # talker is not graph-wrapped); a padded batch would need the runner to
+        # pad the generators list as well.
+        self.talker_mtp_accepts_per_row_generators = True
 
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
@@ -1056,6 +1062,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         top_k: int | None = None,
         top_p: float | None = None,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
         **kwargs: Any,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1).
@@ -1095,6 +1102,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             top_k=top_k,
             top_p=top_p,
             generator=generator,
+            generators=generators,
         )  # [B, Q]
 
         # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes.

vllm_omni/model_executor/models/voxcpm2/runtime_config.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass(frozen=True)
+class _VoxCPM2RuntimeConfig:
+    enable_profiling: bool = False
+    enable_nvtx_profile: bool = False
+    enable_loc_dit_layer_nvtx: bool = False
+    enable_loc_dit_fused_qkv: bool = True
+    enable_loc_dit_fused_mlp: bool = True
+    enable_loc_dit_zero_dt_cache: bool = True
+    enable_loc_dit_fast_rope: bool = False
+    enable_loc_dit_skip_qkv_contig: bool = True
+    enable_loc_dit_reduce_overhead_no_cg: bool = False
+    enable_loc_dit_fullgraph_no_cg: bool = False
+    cfg_cutoff_ratio: float = 1.0
+    decode_graph_capture_policy: str = "all"
+    enable_vae_cuda_graph: bool = False
+    enable_cfm_cuda_graph: bool = False
+    enable_cfm_prealloc_output: bool = False
+    enable_batched_cfm: bool = True
+    deterministic_cfm_noise: bool = False
+    deterministic_cfm_seed: int = 20260601
+    audio_emit_every: int = 1
+    vae_decode_every: int = 3
+    enable_delayed_audio_copy: bool = False
+    delayed_audio_copy_use_events: bool = False
+    coalesce_audio_d2h: bool = True
+    enable_batched_vae_decode: bool = True
+    enable_batched_fsq_fusion: bool = False
+    batched_fsq_fusion_max_batch: int = 32
+    enable_batched_prefill_tail: bool = False
+    enable_unified_decode_graph: bool = True
+    unified_decode_graph_max_batch_size: int = 64
+
+    @classmethod
+    def from_vllm_config(cls, vllm_config: Any) -> _VoxCPM2RuntimeConfig:
+        model_config = vllm_config.model_config
+        raw = getattr(model_config, "voxcpm2_runtime_config", None)
+        if raw is None:
+            raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None)
+        if raw is None:
+            return cls()
+        if hasattr(raw, "to_dict"):
+            raw = raw.to_dict()
+        elif not isinstance(raw, dict) and hasattr(raw, "__dict__"):
+            raw = vars(raw)
+        if not isinstance(raw, dict):
+            logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw)
+            return cls()
+
+        fields = {field.name: field for field in dataclasses.fields(cls)}
+        values: dict[str, Any] = {}
+        for key, value in raw.items():
+            if key not in fields:
+                logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key)
+                continue
+            default = getattr(cls(), key)
+            values[key] = cls._coerce_value(key, value, default)
+        return cls(**values)._normalized()
+
+    def _normalized(self) -> _VoxCPM2RuntimeConfig:
+        cfg = self
+        if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph:
+            logger.warning(
+                "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; "
+                "disabling CFM CUDA Graph and keeping batched CFM."
+            )
+            cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False)
+        if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph:
+            logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.")
+            cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False)
+        if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy:
+            logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.")
+            cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False)
+        if cfg.enable_batched_vae_decode and (
+            cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph
+        ):
+            logger.warning(
+                "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, "
+                "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode."
+            )
+            cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False)
+        if cfg.enable_unified_decode_graph:
+            if cfg.enable_cfm_cuda_graph:
+                logger.info("VoxCPM2 unified decode graph supersedes independent CFM CUDA Graph; disabling it.")
+                cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False, enable_cfm_prealloc_output=False)
+            if not cfg.enable_batched_cfm:
+                logger.info("VoxCPM2 unified decode graph requires batched CFM; enabling it.")
+                cfg = dataclasses.replace(cfg, enable_batched_cfm=True)
+        return cfg
+
+    def unified_decode_graph_available(self, *, use_cuda_graph: bool) -> bool:
+        return bool(use_cuda_graph and self.enable_unified_decode_graph and not self.deterministic_cfm_noise)
+
+    @staticmethod
+    def _coerce_value(key: str, value: Any, default: Any) -> Any:
+        if isinstance(default, bool):
+            if isinstance(value, str):
+                return value.strip().lower() in ("1", "true", "yes", "on")
+            return bool(value)
+        if isinstance(default, int) and not isinstance(default, bool):
+            value = int(value)
+            if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}:
+                return max(1, value)
+            return value
+        if isinstance(default, float):
+            if key == "cfg_cutoff_ratio":
+                return min(1.0, max(0.0, float(value)))
+            return float(value)
+        if isinstance(default, str):
+            return str(value)
+        return value

vllm_omni/model_executor/models/voxcpm2/scheduler.py

@@ -6,6 +6,9 @@ from __future__ import annotations
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler
+from vllm_omni.platforms import current_omni_platform
+
+from .runtime_config import _VoxCPM2RuntimeConfig
 
 
 class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler):
@@ -19,11 +22,8 @@ class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler):
     """
 
     def _unified_decode_graph_enabled(self) -> bool:
-        hf_config = getattr(self.vllm_config.model_config, "hf_config", None)
-        runtime_config = getattr(hf_config, "voxcpm2_runtime_config", None)
-        if isinstance(runtime_config, dict):
-            return bool(runtime_config.get("enable_unified_decode_graph", False))
-        return bool(getattr(runtime_config, "enable_unified_decode_graph", False))
+        runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)
+        return runtime_config.unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())
 
     def _should_defer_waiting_for_unified_decode_graph(self) -> bool:
         if not self._unified_decode_graph_enabled():

vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py

@@ -46,6 +46,7 @@ from vllm_omni.utils.speaker_cache import (
 from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
 
 from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM
+from .runtime_config import _VoxCPM2RuntimeConfig
 from .voxcpm2_import_utils import import_voxcpm2_core
 
 logger = init_logger(__name__)
@@ -104,115 +105,6 @@ class _DecodeResidualMeta(NamedTuple):
     new_lm_hidden: torch.Tensor
 
 
-@dataclasses.dataclass(frozen=True)
-class _VoxCPM2RuntimeConfig:
-    enable_profiling: bool = False
-    enable_nvtx_profile: bool = False
-    enable_loc_dit_layer_nvtx: bool = False
-    enable_loc_dit_fused_qkv: bool = True
-    enable_loc_dit_fused_mlp: bool = True
-    enable_loc_dit_zero_dt_cache: bool = True
-    enable_loc_dit_fast_rope: bool = False
-    enable_loc_dit_skip_qkv_contig: bool = True
-    enable_loc_dit_reduce_overhead_no_cg: bool = False
-    enable_loc_dit_fullgraph_no_cg: bool = False
-    cfg_cutoff_ratio: float = 1.0
-    decode_graph_capture_policy: str = "all"
-    enable_vae_cuda_graph: bool = False
-    enable_cfm_cuda_graph: bool = False
-    enable_cfm_prealloc_output: bool = False
-    enable_batched_cfm: bool = True
-    deterministic_cfm_noise: bool = False
-    deterministic_cfm_seed: int = 20260601
-    audio_emit_every: int = 1
-    vae_decode_every: int = 3
-    enable_delayed_audio_copy: bool = False
-    delayed_audio_copy_use_events: bool = False
-    coalesce_audio_d2h: bool = True
-    enable_batched_vae_decode: bool = True
-    enable_batched_fsq_fusion: bool = False
-    batched_fsq_fusion_max_batch: int = 32
-    enable_batched_prefill_tail: bool = False
-    enable_unified_decode_graph: bool = True
-    unified_decode_graph_max_batch_size: int = 64
-
-    @classmethod
-    def from_vllm_config(cls, vllm_config: VllmConfig) -> _VoxCPM2RuntimeConfig:
-        model_config = vllm_config.model_config
-        raw = getattr(model_config, "voxcpm2_runtime_config", None)
-        if raw is None:
-            raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None)
-        if raw is None:
-            return cls()
-        if hasattr(raw, "to_dict"):
-            raw = raw.to_dict()
-        elif not isinstance(raw, dict) and hasattr(raw, "__dict__"):
-            raw = vars(raw)
-        if not isinstance(raw, dict):
-            logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw)
-            return cls()
-
-        fields = {field.name: field for field in dataclasses.fields(cls)}
-        values: dict[str, Any] = {}
-        for key, value in raw.items():
-            if key not in fields:
-                logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key)
-                continue
-            default = getattr(cls(), key)
-            values[key] = cls._coerce_value(key, value, default)
-        return cls(**values)._normalized()
-
-    def _normalized(self) -> _VoxCPM2RuntimeConfig:
-        cfg = self
-        if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph:
-            logger.warning(
-                "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; "
-                "disabling CFM CUDA Graph and keeping batched CFM."
-            )
-            cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False)
-        if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph:
-            logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.")
-            cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False)
-        if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy:
-            logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.")
-            cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False)
-        if cfg.enable_batched_vae_decode and (
-            cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph
-        ):
-            logger.warning(
-                "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, "
-                "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode."
-            )
-            cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False)
-        if cfg.enable_unified_decode_graph:
-            if cfg.enable_cfm_cuda_graph:
-                logger.info("VoxCPM2 unified decode graph supersedes independent CFM CUDA Graph; disabling it.")
-                cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False, enable_cfm_prealloc_output=False)
-            if not cfg.enable_batched_cfm:
-                logger.info("VoxCPM2 unified decode graph requires batched CFM; enabling it.")
-                cfg = dataclasses.replace(cfg, enable_batched_cfm=True)
-        return cfg
-
-    @staticmethod
-    def _coerce_value(key: str, value: Any, default: Any) -> Any:
-        if isinstance(default, bool):
-            if isinstance(value, str):
-                return value.strip().lower() in ("1", "true", "yes", "on")
-            return bool(value)
-        if isinstance(default, int) and not isinstance(default, bool):
-            value = int(value)
-            if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}:
-                return max(1, value)
-            return value
-        if isinstance(default, float):
-            if key == "cfg_cutoff_ratio":
-                return min(1.0, max(0.0, float(value)))
-            return float(value)
-        if isinstance(default, str):
-            return str(value)
-        return value
-
-
 def is_cjk_char(c: str) -> bool:
     """Check if a character is a CJK ideograph."""
     cp = ord(c)
@@ -1028,8 +920,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         self._cuda_graph_warmup_steps = 0
         self._cuda_graph_warmup_threshold = 3
         self._unified_graphs: dict[int, _CapturedUnifiedDecodeGraph] = {}
-        self._enable_unified_decode_graph = (
-            use_cuda_graph and self._runtime_config.enable_unified_decode_graph and not self._deterministic_cfm_noise
+        self._enable_unified_decode_graph = self._runtime_config.unified_decode_graph_available(
+            use_cuda_graph=use_cuda_graph
         )
         self._unified_decode_graph_max_batch_size = self._runtime_config.unified_decode_graph_max_batch_size
         self._unified_graph_bucket_sizes: frozenset[int] = self._build_unified_graph_bucket_sizes(

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -902,7 +902,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
         ) = self.execute_model_state
         # Clear ephemeral state.
         self.execute_model_state = None
-        seq_len = hidden_states.shape[0]
+        hidden_seq_len = int(hidden_states.shape[0])
+        scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
 
         # Apply structured output bitmasks if present.
         if grammar_output is not None:
@@ -1141,7 +1142,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
                                 start=start,
                                 end=end,
                                 pass_lists_through=False,
-                                seq_len=seq_len,
+                                seq_len=hidden_seq_len,
+                                scheduled_seq_len=scheduled_seq_len,
                             )
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))

vllm_omni/utils/mm_outputs.py

@@ -106,7 +106,13 @@ def _to_cpu(value):
 
 
 def to_payload_element(
-    element: object, idx: int, start: int, end: int, pass_lists_through: bool = False, seq_len: int | None = None
+    element: object,
+    idx: int,
+    start: int,
+    end: int,
+    pass_lists_through: bool = False,
+    seq_len: int | None = None,
+    scheduled_seq_len: int | None = None,
 ):
     """Build an mm payload element corresponding to one request index
     from an element containing 0 or more CPU tensors.
@@ -120,21 +126,39 @@ def to_payload_element(
             passthrough data; this should be False in normal cases, but True
             if we need to avoid splitting nonempty lists prior to calling
             postprocess, which is the case for prefix cache.
-        seq_len: Optional sequence length (i.e., dim 0 of hidden states).
-            When set, a tensor whose first dimension equals seq_len is
-            sliced per request. The prefix cache passthrough also passes
-            the total scheduled token count here so 1D (seq_len,) metadata
-            that is intentionally not cached is still split per request.
+        seq_len: Optional hidden-aligned batch length (``hidden_states.shape[0]``).
+            When a tensor's first dimension equals this value, it is sliced
+            per request using ``start:end``.
+        scheduled_seq_len: Optional scheduler-aligned batch length
+            (``scheduler_output.total_num_scheduled_tokens``). Some full-payload
+            mm tensors (e.g. batched ``codes.audio`` with tail-only hidden states)
+            are laid out by scheduled tokens instead of the hidden tail shape.
+            When omitted, ``seq_len`` is reused for backward compatibility.
     """
-    # Cached per-token tensors are merged elsewhere; here a first dim
-    # equal to seq_len means a per-request slice is required.
-    if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len:
+    if scheduled_seq_len is None:
+        scheduled_seq_len = seq_len
+
+    # Cached per-token tensors are merged elsewhere; here a first dim equal to
+    # either the hidden-aligned or scheduler-aligned batch length means a
+    # per-request slice is required.
+    if isinstance(element, torch.Tensor) and (
+        (seq_len is not None and element.shape[0] == seq_len)
+        or (scheduled_seq_len is not None and element.shape[0] == scheduled_seq_len)
+    ):
         return element[start:end].contiguous()
     # Every other case is shared between prefix cache (passthrough data)
     # and running a model without prefix caching.
     elif isinstance(element, dict):
         return {
-            sk: to_payload_element(sv, idx, start, end, pass_lists_through=pass_lists_through, seq_len=seq_len)
+            sk: to_payload_element(
+                sv,
+                idx,
+                start,
+                end,
+                pass_lists_through=pass_lists_through,
+                seq_len=seq_len,
+                scheduled_seq_len=scheduled_seq_len,
+            )
             for sk, sv in element.items()
         }
     elif isinstance(element, list):

vllm_omni/worker/gpu_ar_model_runner.py

@@ -839,7 +839,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         end: int,
         audio_sparse_output: bool,
         sparse_mm_index: dict[str, int],
-        seq_len: int,
+        hidden_seq_len: int,
+        scheduled_seq_len: int,
     ) -> dict[str, object]:
         if combined_multimodal_outputs:
             return self._build_combined_prefix_cache_mm_payload(
@@ -876,7 +877,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 start=start,
                 end=end,
                 pass_lists_through=False,
-                seq_len=seq_len,
+                seq_len=hidden_seq_len,
+                scheduled_seq_len=scheduled_seq_len,
             )
         return mm_payload
 
@@ -894,7 +896,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         mm_cpu: dict[str, object] | None,
         audio_sparse_output: bool,
         sparse_mm_index: dict[str, int],
-        seq_len: int,
+        hidden_seq_len: int,
+        scheduled_seq_len: int,
     ) -> dict[str, object]:
         payload: dict[str, object] = {}
         if not audio_sparse_output:
@@ -920,7 +923,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             end=end,
             audio_sparse_output=audio_sparse_output,
             sparse_mm_index=sparse_mm_index,
-            seq_len=seq_len,
+            hidden_seq_len=hidden_seq_len,
+            scheduled_seq_len=scheduled_seq_len,
         )
         payload.update(mm_payload)
         return payload
@@ -1676,7 +1680,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         ec_connector_output: Any,
         cudagraph_stats: Any,
         kv_extracted_req_ids: list[str] | None,
-        seq_len: int,
         num_scheduled_tokens_np: np.ndarray,
         query_start_loc_cpu: Any,
         postprocess_already_applied: bool = False,
@@ -1732,6 +1735,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         # The actual multimodal wire transport uses multimodal_outputs instead.
         pooler_output: list[dict[str, object]] | None = None
         if needs_pooler_payload:
+            hidden_seq_len = int(hidden_states.shape[0])
+            scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
             mm_cpu = None
             if self.omni_prefix_cache is not None:
                 (
@@ -1787,7 +1792,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                         mm_cpu=mm_cpu,
                         audio_sparse_output=audio_sparse_output,
                         sparse_mm_index=sparse_mm_index,
-                        seq_len=seq_len,
+                        hidden_seq_len=hidden_seq_len,
+                        scheduled_seq_len=scheduled_seq_len,
                     )
                     pooler_output.append(flatten_payload(payload))
 
@@ -1873,7 +1879,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             slot_mappings,  # OMNI: unpack slot_mappings for drafter
         ) = self.execute_model_state
         self.execute_model_state = None
-        seq_len = hidden_states.shape[0]
 
         # Apply structured output bitmasks if present.
         if grammar_output is not None:
@@ -2040,7 +2045,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     ec_connector_output=ec_connector_output,
                     cudagraph_stats=cudagraph_stats,
                     kv_extracted_req_ids=kv_extracted_req_ids,
-                    seq_len=seq_len,
                     num_scheduled_tokens_np=num_scheduled_tokens_np,
                     query_start_loc_cpu=query_start_loc_cpu,
                     postprocess_already_applied=omni_postprocess_already_applied,

vllm_omni/worker/gpu_model_runner.py

@@ -1826,7 +1826,33 @@ class OmniGPUModelRunner(GPUModelRunner):
                 seed = extra_args.get("tts_local_seed")
             return int(seed) if seed is not None else None
 
-        if decode_batch_size > 1 and any(_explicit_talker_seed(req_id) is not None for req_id in decode_req_ids):
+        def _row_generator(req_id: str) -> torch.Generator | None:
+            seed = _explicit_talker_seed(req_id)
+            if seed is None:
+                return None
+            cache = getattr(self, "_talker_mtp_generators", None)
+            if cache is None:
+                cache = {}
+                self._talker_mtp_generators = cache
+            generator = cache.get(req_id)
+            if generator is None or generator.device != req_input_ids.device:
+                generator = torch.Generator(device=req_input_ids.device)
+                generator.manual_seed(seed)
+                cache[req_id] = generator
+            return generator
+
+        row_generators = [_row_generator(req_id) for req_id in decode_req_ids]
+        cache = getattr(self, "_talker_mtp_generators", None)
+        if cache:
+            # Generators live as long as their request; drop finished ones.
+            for stale_id in [rid for rid in cache if rid not in self.requests]:
+                del cache[stale_id]
+
+        if (
+            decode_batch_size > 1
+            and any(generator is not None for generator in row_generators)
+            and not getattr(self.model, "talker_mtp_accepts_per_row_generators", False)
+        ):
             # A torch.Generator is a single stream. Using one generator for a
             # multi-row batch would make explicitly-seeded requests depend on
             # other rows in the same scheduler step, so keep that path scalar.
@@ -1849,28 +1875,17 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self.text_step.gpu[:decode_batch_size].copy_(saved_text)
             return
 
-        generator = None
-        if decode_req_ids:
-            first_req_id = decode_req_ids[0]
-            seed = _explicit_talker_seed(first_req_id)
-            if seed is not None:
-                generators = getattr(self, "_talker_mtp_generators", None)
-                if generators is None:
-                    generators = {}
-                    self._talker_mtp_generators = generators
-                generator = generators.get(first_req_id)
-                if generator is None or generator.device != req_input_ids.device:
-                    generator = torch.Generator(device=req_input_ids.device)
-                    generator.manual_seed(int(seed))
-                    generators[first_req_id] = generator
         talker_kwargs = {
             "do_sample": subtalker_params.get("do_sample"),
             "temperature": subtalker_params.get("temperature"),
             "top_k": subtalker_params.get("top_k"),
             "top_p": subtalker_params.get("top_p"),
         }
-        if generator is not None:
-            talker_kwargs["generator"] = generator
+        if decode_batch_size == 1:
+            if row_generators[0] is not None:
+                talker_kwargs["generator"] = row_generators[0]
+        elif any(generator is not None for generator in row_generators):
+            talker_kwargs["generators"] = row_generators
         if getattr(self.model, "talker_mtp_accepts_req_infos", False):
             talker_kwargs["req_ids"] = decode_req_ids
             talker_kwargs["req_infos"] = [