[LingBot World] Eleven bit-identical removals of repeated work in the realtime path (#7648)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

tests/diffusion/ar_diffusion/test_kv_cache.py

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

tests/diffusion/ar_diffusion/test_paged_attention.py

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

tests/diffusion/distributed/test_comm.py

@@ -461,3 +461,46 @@ def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
         device_kind="cuda",
         master_port=29501,
     )
+
+
+def _run_fused_qkv_matches_three_4d(
+    local_rank: int,
+    world_size: int,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    """LingBot's fused (B, S/N, 3, H, D) exchange must equal three per-tensor 4D exchanges."""
+    from vllm_omni.diffusion.distributed.comm import all_to_all_5D
+
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
+    _init_worker(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(ulysses_degree=world_size)
+    sp_group = get_sp_group().ulysses_group
+    try:
+        batch, seq_per_rank, heads, head_size = 1, 6, 4 * world_size, 8
+        # Distinguishable values: rank, q/k/v index, position, head, feature.
+        base = torch.arange(seq_per_rank * heads * head_size, dtype=torch.float32, device=device).reshape(
+            batch, seq_per_rank, heads, head_size
+        )
+        tensors = [base + 1000.0 * local_rank + 100000.0 * which for which in range(3)]
+        separate = [SeqAllToAll4D.apply(sp_group, t, 2, 1, False) for t in tensors]
+        fused = all_to_all_5D(torch.stack(tensors, dim=2), scatter_idx=3, gather_idx=1, group=sp_group)
+        assert fused.shape == (batch, seq_per_rank * world_size, 3, heads // world_size, head_size)
+        for which, expected in enumerate(separate):
+            assert torch.equal(fused.unbind(2)[which], expected), f"q/k/v index {which} differs on rank {local_rank}"
+    finally:
+        destroy_distributed_env()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
+    torch.multiprocessing.spawn(
+        _run_fused_qkv_matches_three_4d,
+        args=(world_size, "cpu", 29613),
+        nprocs=world_size,
+    )

tests/diffusion/distributed/test_wan_spatial_shard.py

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

tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py

@@ -74,6 +74,17 @@ def _install_vllm_stubs() -> None:
         _SeqAllToAll4D,
     )
 
+    def _all_to_all_5D(value, scatter_idx=3, gather_idx=1, group=None, use_sync=False):
+        # Single-rank stub: the fused q/k/v exchange is the identity when N == 1.
+        del scatter_idx, gather_idx, group, use_sync
+        return value
+
+    setattr(
+        sys.modules["vllm_omni.diffusion.distributed.comm"],
+        "all_to_all_5D",
+        _all_to_all_5D,
+    )
+
     def get_sp_group():
         return SimpleNamespace(
             ulysses_world_size=1,
@@ -588,26 +599,7 @@ def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyP
     torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
 
 
-def test_single_token_text_kv_shard_owns_compact_storage(monkeypatch):
-    module = _load_module()
-    monkeypatch.setattr(
-        module,
-        "get_sp_group",
-        lambda: SimpleNamespace(
-            ulysses_world_size=2,
-            ulysses_rank=1,
-            ulysses_group=None,
-        ),
-    )
-    attention = module.LingBotCrossAttention(dim=8, num_heads=4)
-    full = torch.arange(8, dtype=torch.float32).reshape(1, 1, 4, 2)
-    shard = attention.shard_kv_heads(full)
-    torch.testing.assert_close(shard, full[:, :, 2:], rtol=0, atol=0)
-    assert shard.is_contiguous()
-    assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()
-
-
-@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention", "LingBotCrossAttention"])
+@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention"])
 def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class):
     module = _load_module()
     monkeypatch.setattr(
@@ -621,3 +613,35 @@ def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class):
     )
     with pytest.raises(ValueError, match="heads must be divisible"):
         getattr(module, attention_class)(dim=8, num_heads=4)
+
+
+def test_supplied_camera_modulation_matches_computing_it_inline() -> None:
+    """Passing a prebuilt modulation must be arithmetically identical to letting the block build its own."""
+    module = _load_module()
+    torch.manual_seed(0)
+    block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+    camera = torch.randn(1, 3, 4)
+
+    scale, shift = block.camera_modulation(camera)
+    inline_scale, inline_shift = block.camera_modulation(camera)
+
+    torch.testing.assert_close(scale, inline_scale, rtol=0, atol=0)
+    torch.testing.assert_close(shift, inline_shift, rtol=0, atol=0)
+
+
+def test_the_camera_modulation_reads_only_the_camera_tokens() -> None:
+    """The reuse window is only sound because the modulation depends on nothing else that varies per forward."""
+    module = _load_module()
+    torch.manual_seed(0)
+    block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+    camera = torch.randn(1, 3, 4)
+
+    first = block.camera_modulation(camera)
+    # Same camera, different everything else a forward would carry: the modulation cannot move.
+    second = block.camera_modulation(camera.clone())
+
+    torch.testing.assert_close(first[0], second[0], rtol=0, atol=0)
+    torch.testing.assert_close(first[1], second[1], rtol=0, atol=0)
+    # A different camera must move it, or the cache key would be hiding a real dependency.
+    other = block.camera_modulation(camera + 1)
+    assert not torch.equal(first[0], other[0])

tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py

@@ -837,6 +837,9 @@ def _rollout(model, mode, dtype, batch):
             frames_per_block=_FRAMES,
             max_scratch_tokens_per_branch=_FRAMES * _TOKENS_PER_FRAME,
             cross_attention_lengths={"text": 5},
+            # Cross-attention keeps every local head on every SP rank, so its pool is wider than the
+            # self-attention share above.
+            cross_attention_kv_heads={"text": model.blocks[0].cross_attn.num_local_heads},
             device=device,
         )
         state = ARDiffusionKVState(kv, "numeric", {"main": kv.begin_request("numeric")}, num_layers=_LAYERS)
@@ -896,7 +899,6 @@ def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous):
     from vllm_omni.diffusion.distributed.parallel_state import (
         destroy_distributed_env,
         destroy_model_parallel,
-        get_sp_group,
         init_distributed_environment,
         initialize_model_parallel,
     )
@@ -948,11 +950,10 @@ def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous):
                         bound = 1e-5 if dtype == torch.float32 else 1e-2
                         error = (result - expected).double().norm() / expected.double().norm()
                         assert error <= bound, f"rank={rank}, TP{tp} SP{sp}: relative L2 {error.item():.3g} > {bound}"
-                        local_heads = _HEADS // tp // sp
-                        first = (
-                            get_tensor_model_parallel_rank() * (_HEADS // tp)
-                            + get_sp_group().ulysses_rank * local_heads
-                        )
+                        # Cross-attention K/V is replicated across the Ulysses group, not split over it, so
+                        # every SP rank holds its whole TP shard of heads. Only TP slices the cache now.
+                        local_heads = _HEADS // tp
+                        first = get_tensor_model_parallel_rank() * local_heads
                         for (k, v), (ek, ev) in zip(cross, expected_cross, strict=True):
                             tolerance = 1e-5 if dtype == torch.float32 else 2e-2
                             torch.testing.assert_close(
@@ -1033,3 +1034,56 @@ def test_temporal_rope_extension_keeps_spatial_limit():
         model._rotary_embedding(
             frames=3, height=17, width=1, start_frame=16, dtype=torch.float32, device=torch.device("cpu")
         )
+
+
+@pytest.mark.cpu
+def test_the_camera_injector_runs_once_per_held_cache_and_gives_the_same_output() -> None:
+    """Two forwards sharing one caller-held cache: the injector runs once per block, outputs unchanged."""
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=2).eval()
+    hidden_states = torch.randn(1, 36, 1, 4, 4)
+    timestep = torch.tensor([1.0])
+    encoder_hidden_states = torch.randn(1, 3, 6)
+    camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+
+    def run(cache, camera_cache=None):
+        return model(
+            hidden_states,
+            timestep,
+            encoder_hidden_states,
+            camera_hidden_states,
+            cache=cache,
+            start_frame=0,
+            update_cache=False,
+            camera_modulation_cache=camera_cache,
+        )
+
+    # Baseline: no cache, so every forward rebuilds the injector in every block.
+    uncached_first = run(_cache(module, model))
+    uncached_second = run(_cache(module, model))
+    injector_calls = [attention_tests._record_outputs(block.cam_injector_layer1) for block in model.blocks]
+    held = module.CameraModulationCache()
+    cached_first = run(_cache(module, model), held)
+    cached_second = run(_cache(module, model), held)
+
+    # One build per block for both forwards, instead of one per block per forward.
+    assert [len(calls) for calls in injector_calls] == [1, 1]
+    torch.testing.assert_close(cached_first, uncached_first, rtol=0, atol=0)
+    torch.testing.assert_close(cached_second, uncached_second, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_timestep_projection_is_staged_into_one_static_buffer_per_shape() -> None:
+    """Two forwards' projections land in the same buffer; the blocks never see the fresh per-forward tensor."""
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1).eval()
+    first = torch.randn(1, 8, 6, model.dim)
+    second = torch.randn(1, 8, 6, model.dim)
+
+    staged_first = model._stage_timestep_projection(first)
+    assert staged_first is not first and torch.equal(staged_first, first)
+    staged_second = model._stage_timestep_projection(second)
+    assert staged_second is staged_first and torch.equal(staged_second, second)
+    # A different shape gets its own buffer; the first one is kept.
+    other = model._stage_timestep_projection(torch.randn(1, 4, 6, model.dim))
+    assert other is not staged_first and len(model._timestep_projection_buffers) == 2

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -4,11 +4,13 @@
 from __future__ import annotations
 
 import gc
+import os
 import weakref
 from pathlib import Path
 from threading import Lock
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
@@ -29,6 +31,7 @@ from vllm_omni.diffusion.models.lingbot_world.camera import (
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.utils import StepRequestState
+from vllm_omni.experimental.ar_diffusion.capability import supports_chunk_step_grouping
 from vllm_omni.experimental.ar_diffusion.tick_protocol import (
     ARDiffusionControlInput,
     ARDiffusionTickRequest,
@@ -127,6 +130,10 @@ class _RecordingTransformer(nn.Module):
             sink_size=3,
         )
         self.blocks = nn.ModuleList([nn.Identity(), nn.Identity()])
+        for block in self.blocks:
+            # The spec reads head geometry from both attentions; cross-attention keeps every local head.
+            block.self_attn = SimpleNamespace(num_sp_heads=2)
+            block.cross_attn = SimpleNamespace(num_local_heads=2)
         self.calls: list[dict] = []
         self.cache_allocations: list[dict] = []
         self.raise_on_call = raise_on_call
@@ -146,6 +153,7 @@ class _RecordingTransformer(nn.Module):
             "cache_id": id(kwargs["cache"]),
             "start_frame": kwargs["start_frame"],
             "update_cache": kwargs["update_cache"],
+            "camera_cache": kwargs.get("camera_modulation_cache"),
         }
         self.calls.append(call)
         if self.raise_on_call == len(self.calls):
@@ -534,14 +542,15 @@ def test_ar_diffusion_capability_uses_transformer_local_head_geometry() -> None:
     assert spec.sink_frames == 3
     assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)]
     assert spec.cross_attention_lengths == {"text": 512}
-    # One condition block, committed + pending encoder caches, and decoder state.
+    # One condition block, committed + pending encoder caches, the session's
+    # cached prompt embedding (512 tokens x text_dim 8 x fp32), and decoder state.
     # The stub encoder retains 2 RGB frames at 16x16 and 2 feature frames at 2x2.
     encoder_bytes = (2 * 3 * 16 * 16 + 2 * 16 * 2 * 2) * 4
-    assert spec.model_owned_state_bytes_per_session == 960 + 2 * encoder_bytes + 2_421_248
+    assert spec.model_owned_state_bytes_per_session == 960 + 384 + 2 * encoder_bytes + 2_421_248
 
 
 def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None:
-    """Admission includes encoder state even when the decoder cannot stream."""
+    """Admission includes encoder state and the prompt embedding even when the decoder cannot stream."""
     module = _load_pipeline_module()
 
     def spec_pipeline():
@@ -560,13 +569,29 @@ def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None:
     # than the image condition alone could account for.
     assert widened - small > 2_000_000
 
-    # Removing decoder support leaves the condition and both encoder histories.
+    # Removing decoder support leaves the condition, both encoder histories and
+    # the prompt-embedding cache.
     bare = spec_pipeline()
     for attribute in ("decoder", "post_quant_conv"):
         if hasattr(bare.vae, attribute):
             delattr(bare.vae, attribute)
     bare.vae._cached_conv_counts.pop("decoder")
-    assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 960 + 2 * 6_656
+    assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 960 + 384 + 2 * 6_656
+
+
+@pytest.mark.parametrize("sp_size", [1, 4])
+def test_session_admission_budgets_camera_modulation_per_rank(sp_size: int) -> None:
+    pipeline = _pipeline(_load_pipeline_module())
+    pipeline._ar_height, pipeline._ar_width = 480, 832
+    pipeline.od_config.parallel_config.ulysses_degree = sp_size
+    pipeline.transformer.config.num_layers = 40
+    pipeline.transformer.config.num_attention_heads = 40
+    pipeline.transformer.config.attention_head_dim = 128
+    pipeline.transformer._dtype = torch.bfloat16
+    forty_layers = pipeline.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session
+    pipeline.transformer.config.num_layers = 1
+    one_layer = pipeline.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session
+    assert forty_layers - one_layer == 2 * 39 * (4680 // sp_size) * 5120 * 2
 
 
 def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None:
@@ -1672,7 +1697,6 @@ def test_typed_ticks_generate_one_global_block_and_return_standard_metadata(
         "applied_event_ids": [1],
     }
     assert pipeline._ar_sessions["world-1"].next_chunk_index == 2
-    assert pipeline._ar_sessions["world-1"].prompt == "enter the snowy valley"
     assert pipeline._ar_sessions["world-1"].camera_tail is not None
     assert pipeline._ar_sessions["world-1"].camera_tail.poses.shape == (1, 4, 4)
     expected_generator = torch.Generator(device="cpu").manual_seed(17)
@@ -1846,8 +1870,11 @@ def test_typed_ticks_keep_bounded_encoder_state_beyond_ten_chunks(
         pipeline(_request(sampling=sampling))
 
     assert pipeline.vae.encode_inputs == []
-    assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == [1] + [4] * 32
+    # The stub encoder's history settles after its second block; the nine
+    # blocks after that reuse the settled condition without an encode.
+    assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == [1] + [4] * 5
     assert all(torch.count_nonzero(video) == 0 for video in pipeline.vae.encoder.inputs[1:])
+    assert pipeline._ar_sessions["world-1"].condition_fixed_point is not None
     encoder_cache = pipeline._ar_sessions["world-1"].encoder_cache
     assert sum(t.numel() * t.element_size() for t in encoder_cache) == pipeline._condition_encoder_cache_bytes()
     assert pipeline.vae._enc_feat_map == ["module-owned"]
@@ -1943,6 +1970,76 @@ def test_condition_continues_real_causal_vae_beyond_initial_horizon(
     assert pipeline._ar_sessions == {}
 
 
+def test_condition_encoder_stops_at_its_fixed_point_and_restarts_with_the_session(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Once a block returns the history it was given, later blocks reuse the condition without encoding.
+
+    The whole-clip encode is the reference for every block, so reusing the
+    stored condition is only allowed to be a speedup, never a change.
+    """
+    pipeline = _tiny_condition_pipeline(monkeypatch)
+    pipeline._ar_diffusion_kv_state = object()
+    blocks = 12
+    inputs = pipeline._parse_request(_request(sampling=_SamplingParams(num_frames=(blocks * 3 - 1) * 4 + 1)))
+    with torch.inference_mode():
+        expected = pipeline._prepare_condition(inputs, dtype=torch.float32)
+    encoder_calls = []
+    original_encoder = pipeline.vae.encoder.forward
+
+    def encoder(*args, **kwargs):
+        encoder_calls.append(len(seen))
+        return original_encoder(*args, **kwargs)
+
+    seen: list[torch.Tensor] = []
+
+    def generate_block(**kwargs):
+        index = len(seen) % blocks
+        torch.testing.assert_close(
+            kwargs["condition"], expected[:, :, index * 3 : (index + 1) * 3], rtol=1e-6, atol=1e-6
+        )
+        seen.append(kwargs["condition"])
+        return torch.zeros_like(kwargs["condition"][:, :16])
+
+    monkeypatch.setattr(pipeline.vae.encoder, "forward", encoder)
+    monkeypatch.setattr(pipeline, "_generate_block", generate_block)
+
+    def tick(index):
+        pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=index))))
+
+    settled_at = None
+    for index in range(blocks):
+        tick(index)
+        state = pipeline._ar_sessions["world-1"]
+        if settled_at is None and state.condition_fixed_point is not None:
+            settled_at = index
+        assert state.pending_encoder_cache is None and state.encoder_cache is not None
+    assert settled_at is not None and 0 < settled_at < blocks - 1, settled_at
+    # Three encoder passes per encoded block, none once the fixed point is stored.
+    assert encoder_calls == [index for index in range(settled_at + 1) for _ in range(3)]
+    # The stored condition is the session's own, not shared with the callers.
+    assert all(seen[settled_at] is not later for later in seen[settled_at + 1 :])
+
+    # A session restarted from chunk 0 encodes again from the source image.
+    pipeline.reset_ar_diffusion_session("world-1")
+    tick(0)
+    state = pipeline._ar_sessions["world-1"]
+    assert state.condition_fixed_point is None
+    assert encoder_calls[-3:] == [blocks] * 3
+
+
+def test_condition_encoder_fixed_point_is_released_with_the_session(monkeypatch: pytest.MonkeyPatch) -> None:
+    pipeline = _tiny_condition_pipeline(monkeypatch)
+    pipeline._ar_diffusion_kv_state = object()
+    monkeypatch.setattr(pipeline, "_generate_block", lambda **kwargs: torch.zeros_like(kwargs["condition"][:, :16]))
+    for index in range(8):
+        pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=index))))
+    state = pipeline._ar_sessions["world-1"]
+    assert state.condition_fixed_point is not None
+    pipeline.close_ar_diffusion_session("world-1")
+    assert state.condition_fixed_point is None and state.encoder_cache is None
+
+
 def test_condition_encoder_interleaving_retry_and_reset(monkeypatch: pytest.MonkeyPatch) -> None:
     pipeline = _tiny_condition_pipeline(monkeypatch)
     pipeline._ar_diffusion_kv_state = object()
@@ -2209,6 +2306,40 @@ def _stepwise_chunks(pipeline, state, ar_state):
             yield output
 
 
+def test_stepwise_block_shares_one_camera_cache_across_its_five_forwards() -> None:
+    """The four probes and the commit of one stepwise block reuse one camera cache; the next block gets its own.
+
+    Request mode creates one cache per loop; the stepwise path's forwards are
+    separate scheduler steps, so the pipeline holds the block's cache and
+    passes it to each one.
+    """
+    module = _load_pipeline_module()
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "latent"
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        outputs = _run_stepwise(pipeline, state)
+
+    assert len(outputs) == 2 and len(transformer.calls) == 10
+    caches = [call["camera_cache"] for call in transformer.calls]
+    assert all(c is not None for c in caches)
+    assert len({id(c) for c in caches[:5]}) == 1 and len({id(c) for c in caches[5:]}) == 1
+    assert caches[0] is not caches[5]
+    # The block's cache does not outlive its commit.
+    assert "camera_cache" not in state.extra
+
+
+def test_pipeline_declares_chunk_step_grouping() -> None:
+    """A realtime block's probes and commit may run without a scheduler cycle between them."""
+    module = _load_pipeline_module()
+    assert module.LingBotWorldCausalDMDPipeline.supports_chunk_step_grouping is True
+    assert supports_chunk_step_grouping(_pipeline(module))
+
+
 def test_pipeline_declares_step_execution_support() -> None:
     module = _load_pipeline_module()
     pipeline = _pipeline(module)
@@ -2337,7 +2468,10 @@ def test_stepwise_matches_tick_transformer_trace_and_latents(num_chunks) -> None
         torch.testing.assert_close(step_call["camera_hidden_states"], tick_call["camera_hidden_states"])
     assert stepwise_state.extra["condition"].shape[2] == 3
     assert pipeline.vae.encode_inputs == []
-    assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == ([1] + [4] * (num_chunks * 3 - 1)) * 2
+    # The stub encoder's history settles after its second block, so both
+    # modes encode two blocks and reuse the settled condition after that.
+    encoded_blocks = min(num_chunks, 2)
+    assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == ([1] + [4] * (encoded_blocks * 3 - 1)) * 2
     assert stepwise_fake.commits == ["main"] * num_chunks
     assert len(stepwise_outputs) == num_chunks and stepwise_outputs[-1].finished
     assert pipeline.vae.decode_inputs == []
@@ -2406,21 +2540,16 @@ def test_stepwise_emits_latents_without_touching_the_vae() -> None:
     assert pipeline.vae.decode_inputs == []
 
 
-def _capture_video_processor(monkeypatch) -> list[tuple[tuple[int, ...], str]]:
-    """Record what reaches postprocess_video, and name each chunk's output."""
-    import diffusers.video_processor as video_processor_module
+def _capture_pixel_conversion(monkeypatch) -> list[tuple[int, ...]]:
+    """Record each decoded block reaching the uint8 conversion, and name its output."""
+    module = _load_pipeline_module()
+    processed: list[tuple[int, ...]] = []
 
-    processed: list[tuple[tuple[int, ...], str]] = []
+    def uint8_frames(video):
+        processed.append(tuple(video.shape))
+        return f"frames-{len(processed)}"
 
-    class VideoProcessor:
-        def __init__(self, *, vae_scale_factor):
-            assert vae_scale_factor == 8
-
-        def postprocess_video(self, video, *, output_type):
-            processed.append((tuple(video.shape), output_type))
-            return f"frames-{len(processed)}"
-
-    monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
+    monkeypatch.setattr(module, "_uint8_frames", uint8_frames)
     return processed
 
 
@@ -2434,7 +2563,7 @@ def _streaming_pipeline(module):
 def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
     """A non-latent request must stream pixels under the primary "video" key."""
     module = _load_pipeline_module()
-    processed = _capture_video_processor(monkeypatch)
+    processed = _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
@@ -2443,7 +2572,7 @@ def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> Non
     with pipeline.bind_ar_diffusion_state(state.request_id, fake):
         outputs = _run_stepwise(pipeline, state)
 
-    assert [output_type for _, output_type in processed] == ["np", "np"]
+    assert processed == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
     assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
     # Identity metadata survives the switch to pixel output.
     assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1]
@@ -2457,7 +2586,7 @@ def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> N
     the same latents would: 9 frames, then 12.
     """
     module = _load_pipeline_module()
-    processed = _capture_video_processor(monkeypatch)
+    processed = _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
@@ -2466,7 +2595,7 @@ def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> N
     with pipeline.bind_ar_diffusion_state(state.request_id, fake):
         _run_stepwise(pipeline, state)
 
-    assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
+    assert processed == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
     # One opening frame for the session, not one per chunk.
     assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5
     # The whole-clip decode path is not involved, and the module-owned cache
@@ -2511,7 +2640,7 @@ def test_prepare_next_chunk_rejects_live_camera_with_scripted_paths() -> None:
 def test_peek_chunk_media_matches_streaming_decoder_frame_counts(monkeypatch) -> None:
     """Camera timelines must track the streaming decoder's 9-then-12 media counts."""
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
@@ -2536,7 +2665,7 @@ def test_peek_chunk_media_matches_streaming_decoder_frame_counts(monkeypatch) ->
 def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None:
     """Decoder state is keyed by request id and released with the AR session."""
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(request_id="req-stream", num_frames=21)
     state.sampling.output_type = "np"
@@ -2560,7 +2689,7 @@ def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None:
 def test_streaming_decode_state_is_dropped_on_session_reset(monkeypatch) -> None:
     """A reset session must not resume the temporal context it just abandoned."""
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(request_id="req-reset", num_frames=21)
     state.sampling.output_type = "np"
@@ -2584,7 +2713,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
     interleaving the scheduler actually produces.
     """
     module = _load_pipeline_module()
-    processed = _capture_video_processor(monkeypatch)
+    processed = _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     states = []
     for request_id in ("req-a", "req-b"):
@@ -2617,7 +2746,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
     assert pipeline.vae.decoder.first_chunk_flags == (
         [True, False, False] + [True, False, False] + [False] * 3 + [False] * 3
     )
-    assert [shape for shape, _ in processed] == [
+    assert processed == [
         (1, 3, 9, 16, 16),
         (1, 3, 9, 16, 16),
         (1, 3, 12, 16, 16),
@@ -2628,7 +2757,7 @@ def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> No
 def test_streaming_decode_failure_drops_the_half_advanced_cache(monkeypatch) -> None:
     """A failed chunk must not leave a cache a later chunk would continue from."""
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(request_id="req-boom", num_frames=21)
     state.sampling.output_type = "np"
@@ -2662,7 +2791,7 @@ def test_streaming_decode_receives_rescaled_latents(monkeypatch) -> None:
     never reaches, so this asserts on what ``post_quant_conv`` received.
     """
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     state = _stepwise_state(num_frames=21)
     state.sampling.output_type = "np"
@@ -2695,7 +2824,7 @@ def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None:
     perf regression in the streaming decoder would be too.
     """
     module = _load_pipeline_module()
-    _capture_video_processor(monkeypatch)
+    _capture_pixel_conversion(monkeypatch)
     pipeline = _pipeline(
         module,
         transformer=_RecordingTransformer(),
@@ -2723,21 +2852,18 @@ def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None:
 def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> None:
     """Tiled decode drives its own cache, so a tiling shape keeps per-chunk decode."""
     module = _load_pipeline_module()
-    processed = _capture_video_processor(monkeypatch)
+    processed = _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     # A latent of 2x2 exceeds a one-latent-cell tile, so vae.decode would tile.
     _enable_vae_tiling(pipeline, tile_sample_min=8)
-    outputs = [
-        pipeline._decode_chunk_to_pixels(torch.zeros(1, 16, 3, 2, 2), output_type="np", session_id="tiled")
-        for _ in range(2)
-    ]
+    outputs = [pipeline._decode_chunk_to_pixels(torch.zeros(1, 16, 3, 2, 2), session_id="tiled") for _ in range(2)]
 
     # One whole-clip decode per AR block, each seeing only that block's frames.
     assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
     assert pipeline.vae.decoder.first_chunk_flags == []
     assert pipeline._streaming_decode_states == {}
     # Nine frames per chunk instead of 9 then 12: the frame loss being removed.
-    assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)]
+    assert processed == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)]
     assert outputs == ["frames-1", "frames-2"]
 
 
@@ -2749,7 +2875,7 @@ def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monke
     on a configuration that would never have tiled.
     """
     module = _load_pipeline_module()
-    processed = _capture_video_processor(monkeypatch)
+    processed = _capture_pixel_conversion(monkeypatch)
     pipeline = _streaming_pipeline(module)
     # A 2x2 latent is within a tile this size, so vae.decode would not tile.
     _enable_vae_tiling(pipeline, tile_sample_min=256)
@@ -2761,7 +2887,7 @@ def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monke
 
     assert pipeline.vae.decode_inputs == []
     assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5
-    assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
+    assert processed == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)]
 
 
 def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
@@ -2904,3 +3030,193 @@ def test_stepwise_requests_with_the_same_prompt_share_one_encode() -> None:
 
     assert encoded == ["move through the room"]
     assert states[0].prompt_embeds is states[1].prompt_embeds
+
+
+def test_session_text_caches_publish_key_only_after_a_successful_build() -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    state = module._LingBotARSessionState()
+    invalidations: list[bool] = []
+    fail_next = {"raise": False}
+
+    def fake_text_caches(prompt_embeds, *, invalidate):
+        invalidations.append(invalidate)
+        if fail_next["raise"]:
+            fail_next["raise"] = False
+            raise RuntimeError("injected failure inside _ar_text_caches")
+        return ["caches"]
+
+    pipeline._ar_text_caches = fake_text_caches
+    embeds = torch.ones(1, 512, 8)
+    caches = lambda prompt, length=512, e=embeds: pipeline._session_text_caches(  # noqa: E731
+        e, prompt=prompt, max_sequence_length=length, session_state=state
+    )
+
+    # First build of a session: runner state is empty -> no invalidation, key published.
+    caches("A")
+    assert invalidations == [False] and state.text_cache_key == ("A", 512, str(embeds.dtype))
+    assert state.text_cache_dirty is False
+    # Same inputs: reuse without invalidation.
+    caches("A")
+    assert invalidations == [False, False]
+    # Regression (Codex review): B's build fails after invalidating -> no key is
+    # left behind, so retrying A must invalidate again instead of consuming
+    # partial/mixed K/V.
+    fail_next["raise"] = True
+    with pytest.raises(RuntimeError, match="injected failure"):
+        caches("B")
+    assert state.text_cache_key is None and state.text_cache_dirty is True
+    caches("A")
+    assert invalidations[-2:] == [True, True] and state.text_cache_key == ("A", 512, str(embeds.dtype))
+    assert state.text_cache_dirty is False
+    # The key covers every embedding-defining input: same text, other length or dtype -> rebuild.
+    caches("A", length=256)
+    caches("A", e=embeds.to(torch.bfloat16))
+    assert invalidations[-2:] == [True, True]
+
+
+# ---------------------------------------------------------------------------
+# a source image is decoded once while the file behind it is unchanged
+# ---------------------------------------------------------------------------
+
+
+def _write_png(path, colour=(10, 20, 30)):
+    Image.new("RGB", (64, 64), colour).save(path)
+    return path
+
+
+def _counting_decode(module):
+    """Replace the file decode with a counter, so the tests measure decodes rather than pixels."""
+    decoded: list[str] = []
+    original = module._decode_source_image_file
+
+    def decode(path):
+        decoded.append(os.fspath(path))
+        return Image.new("RGB", (64, 64), (1, 2, 3))
+
+    module._decode_source_image_file = decode
+    return decoded, original
+
+
+def test_the_same_source_image_is_decoded_once(tmp_path: Path) -> None:
+    module = _load_pipeline_module()
+    module._decode_source_image_fingerprinted.cache_clear()
+    path = _write_png(tmp_path / "frame.png")
+    decoded, original = _counting_decode(module)
+    try:
+        first = module._load_source_image(path)
+        second = module._load_source_image(path)
+        assert decoded == [os.fspath(os.path.realpath(path))] or len(decoded) == 1
+        # Every caller owns its copy: mutating one must not reach the next caller or the cache.
+        assert first is not second
+        first.putpixel((0, 0), (255, 255, 255))
+        assert module._load_source_image(path).getpixel((0, 0)) == (1, 2, 3)
+        assert len(decoded) == 1
+    finally:
+        module._decode_source_image_file = original
+        module._decode_source_image_fingerprinted.cache_clear()
+
+
+def test_a_rewritten_source_image_is_decoded_again(tmp_path: Path) -> None:
+    module = _load_pipeline_module()
+    module._decode_source_image_fingerprinted.cache_clear()
+    path = _write_png(tmp_path / "frame.png")
+    decoded, original = _counting_decode(module)
+    try:
+        module._load_source_image(path)
+        module._load_source_image(path)
+        assert len(decoded) == 1
+        # Same path, different bytes: the mtime/size in the key must force a fresh decode.
+        os.utime(path, ns=(0, 0))
+        _write_png(path, colour=(200, 100, 50))
+        module._load_source_image(path)
+        assert len(decoded) == 2
+    finally:
+        module._decode_source_image_file = original
+        module._decode_source_image_fingerprinted.cache_clear()
+
+
+def test_an_image_rewritten_during_the_decode_is_never_cached(tmp_path: Path) -> None:
+    module = _load_pipeline_module()
+    module._decode_source_image_fingerprinted.cache_clear()
+    path = _write_png(tmp_path / "frame.png")
+    decoded: list[str] = []
+    original = module._decode_source_image_file
+
+    def decode_then_rewrite(p):
+        decoded.append(os.fspath(p))
+        _write_png(path, colour=(len(decoded), 0, 0))
+        os.utime(path, ns=(len(decoded) * 1_000_000, len(decoded) * 1_000_000))
+        return Image.new("RGB", (64, 64), (1, 2, 3))
+
+    module._decode_source_image_file = decode_then_rewrite
+    try:
+        module._load_source_image(path)
+        # The entry is keyed by the fingerprint taken before the decode, which the rewrite invalidated: the
+        # next call must decode again rather than serve the image that was decoded mid-rewrite.
+        module._load_source_image(path)
+        assert len(decoded) == 2
+    finally:
+        module._decode_source_image_file = original
+        module._decode_source_image_fingerprinted.cache_clear()
+
+
+def test_an_unstattable_path_keeps_upstream_validation(tmp_path: Path) -> None:
+    module = _load_pipeline_module()
+    module._decode_source_image_fingerprinted.cache_clear()
+    missing = tmp_path / "does-not-exist.png"
+    # The uncached path owns the error contract; caching must not change which exception a caller sees.
+    with pytest.raises(ValueError):
+        module._load_source_image(missing)
+    assert module._decode_source_image_fingerprinted.cache_info().currsize == 0
+
+
+def test_the_text_cache_is_sized_by_cross_attention_heads_not_the_self_attention_share() -> None:
+    """Cross-attention replicates heads per rank, so its pool must not inherit the self-attention head count."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    pipeline.transformer.blocks[0].self_attn = SimpleNamespace(num_sp_heads=1)
+    pipeline.transformer.blocks[0].cross_attn = SimpleNamespace(num_local_heads=4)
+
+    spec = pipeline.ar_diffusion_kv_cache_spec()
+
+    assert spec.num_kv_heads == 1
+    assert spec.cross_attention_kv_heads == {"text": 4}
+    assert spec.cross_attention_lengths == {"text": 512}
+
+
+@pytest.mark.parametrize("output_type", ["np", "pil"])
+def test_streamed_chunk_is_uint8_frames_whatever_pixel_output_type_was_asked(monkeypatch, output_type) -> None:
+    """A realtime tick delivers (F, H, W, 3) uint8 for every pixel output type; only latent differs."""
+    module = _load_pipeline_module()
+    pipeline = _streaming_pipeline(module)
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = output_type
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        outputs = _run_stepwise(pipeline, state)
+
+    videos = [output.output["payload"]["video"] for output in outputs]
+    assert [type(video) for video in videos] == [np.ndarray, np.ndarray]
+    assert [video.dtype for video in videos] == [np.uint8, np.uint8]
+    assert [video.shape for video in videos] == [(9, 16, 16, 3), (12, 16, 16, 3)]
+
+
+def test_uint8_frames_match_the_pil_conversion_byte_for_byte() -> None:
+    """The device-side conversion is the PIL path's arithmetic, so the bytes are the same."""
+    from diffusers.video_processor import VideoProcessor
+
+    module = _load_pipeline_module()
+    generator = torch.Generator().manual_seed(7)
+    # Cover the clamp on both sides and values that sit exactly on rounding edges.
+    video = torch.rand(1, 3, 5, 8, 8, generator=generator) * 2.4 - 1.2
+    video[0, :, 0, 0, :4] = torch.tensor([-1.0, 1.0, 0.0, 1 / 255 - 1.0])
+
+    reference = VideoProcessor(vae_scale_factor=8).postprocess_video(video, output_type="pil")[0]
+    expected = np.stack([np.asarray(frame) for frame in reference])
+
+    frames = module._uint8_frames(video)
+
+    assert frames.dtype == np.uint8 and frames.shape == (5, 8, 8, 3)
+    assert frames.flags["C_CONTIGUOUS"]
+    np.testing.assert_array_equal(frames, expected)

tests/diffusion/test_diffusion_step_pipeline.py

@@ -103,6 +103,78 @@ class _StepPipeline:
         return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
 
 
+class _ChunkedStepPipeline(_StepPipeline):
+    """Streaming stub: two chunks of two denoise steps each, one decode per chunk."""
+
+    supports_chunk_step_grouping = False
+
+    def prepare_encode(self, state, **kwargs):
+        del kwargs
+        self.prepare_calls += 1
+        state.timesteps = [torch.tensor(10), torch.tensor(5), torch.tensor(10), torch.tensor(5)]
+        state.latents = torch.tensor([0.0])
+        state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
+        state.chunk_num_steps = 2
+        state.total_chunks = 2
+        state.chunk_index = 0
+        state.step_in_chunk = 0
+        return state
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del noise_pred, kwargs
+        self.scheduler_calls += 1
+        state.step_in_chunk += 1
+        state.step_index = state.step_in_chunk
+
+    def post_decode(self, state, **kwargs):
+        del kwargs
+        self.decode_calls += 1
+        result = DiffusionOutput(output=torch.tensor([float(state.chunk_index)]))
+        state.chunk_index += 1
+        state.step_in_chunk = 0
+        return result
+
+
+class _BatchDependentChunkPipeline(_ChunkedStepPipeline):
+    """Streaming stub whose noise depends on the batch's latents and timesteps, so a stale batch shows."""
+
+    def __init__(self):
+        super().__init__()
+        self.seen: list[tuple[float, float]] = []
+
+    def prepare_encode(self, state, **kwargs):
+        super().prepare_encode(state, **kwargs)
+        state.latents = torch.tensor([1.0])
+        return state
+
+    def denoise_step(self, input_batch, **kwargs):
+        self.denoise_calls += 1
+        latent = float(input_batch.latents.reshape(-1)[0])
+        timestep = float(input_batch.timesteps.reshape(-1)[0])
+        self.seen.append((latent, timestep))
+        return torch.full((input_batch.latents.shape[0], 1), latent + timestep)
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del kwargs
+        self.scheduler_calls += 1
+        state.latents = state.latents + noise_pred.reshape(-1)[0]
+        state.step_in_chunk += 1
+        state.step_index = state.step_in_chunk
+
+    def post_decode(self, state, **kwargs):
+        self.final_latent = state.latents.reshape(-1)[0].clone()
+        return super().post_decode(state, **kwargs)
+
+
+class _FailsMidChunkPipeline(_ChunkedStepPipeline):
+    """Streaming stub whose second denoise step of the first chunk raises."""
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        super().step_scheduler(state, noise_pred, **kwargs)
+        if state.chunk_index == 0 and state.step_in_chunk == 2:
+            raise RuntimeError("boom mid chunk")
+
+
 class _PerRequestErrorStepPipeline(_StepPipeline):
     def prepare_encode(self, state, **kwargs):
         if state.prompt == "fail-prepare":
@@ -327,8 +399,9 @@ def _make_vllm_config():
 
 def _make_runner(
     diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY,
+    runner_cls: type[DiffusionModelRunner] = DiffusionModelRunner,
 ):
-    runner = object.__new__(DiffusionModelRunner)
+    runner = object.__new__(runner_cls)
     runner.vllm_config = _make_vllm_config()
     runner.od_config = SimpleNamespace(
         cache_backend=None,
@@ -511,6 +584,37 @@ def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_compositi
     torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3))
 
 
+@pytest.mark.cpu
+def test_make_batch_identity_mapping_never_reads_the_device_tensor_back(monkeypatch):
+    """The default index mapping is built on the host; reading it back from the device would sync every step."""
+
+    def forbidden(self, *args, **kwargs):
+        raise AssertionError("make_batch read a tensor back to the host")
+
+    monkeypatch.setattr(torch.Tensor, "tolist", forbidden)
+    monkeypatch.setattr(torch.Tensor, "cpu", forbidden)
+    states = [_make_input_batch_state("req-1", 1.0), _make_input_batch_state("req-2", 2.0)]
+
+    batch = InputBatch.make_batch(states)
+    repacked = InputBatch.make_batch(states, cached_batch=batch)
+
+    assert repacked is batch
+    assert list(batch.idx_mapping_np) == [0, 1]
+    assert batch.idx_mapping.dtype is torch.int32 and batch.idx_mapping.shape == (2,)
+    assert batch.request_ids == ["req-1", "req-2"]
+
+
+@pytest.mark.cpu
+def test_make_batch_explicit_mapping_still_selects_and_orders_states():
+    states = [_make_input_batch_state("req-1", 1.0), _make_input_batch_state("req-2", 2.0)]
+
+    batch = InputBatch.make_batch(states, idx_mapping=torch.tensor([1, 0]))
+
+    assert batch.request_ids == ["req-2", "req-1"]
+    assert list(batch.idx_mapping_np) == [1, 0]
+    torch.testing.assert_close(batch.latents, torch.tensor([[2.0], [1.0]]))
+
+
 @pytest.mark.cpu
 def test_step_profiler_reports_denoise_step_as_diffuse():
     pipeline = _AutoDenoiseProfilerPipeline()
@@ -628,6 +732,146 @@ class TestRunner:
         assert runner.pipeline.scheduler_calls == 2
         assert runner.pipeline.decode_calls == 1
 
+    @staticmethod
+    def _make_grouping_runner(pipeline, *, grouping: bool):
+        """An AR runner without KV (so it inherits the stepwise entry) over a streaming stub pipeline."""
+        from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+        runner = _make_runner(runner_cls=ARDiffusionModelRunner)
+        runner.od_config.streaming_output = True
+        runner.kv_cache = None
+        runner.pipeline = pipeline
+        runner.pipeline.supports_chunk_step_grouping = grouping
+        return runner
+
+    @staticmethod
+    def _drive_to_completion(runner, num_steps=4):
+        outputs = [runner.execute_stepwise(_make_scheduler_output(_make_step_request(num_steps)))]
+        step = 1
+        while not outputs[-1].get_request_output("req-1").finished:
+            outputs.append(runner.execute_stepwise(_make_cached_scheduler_output(step_id=step)))
+            step += 1
+        return [o.get_request_output("req-1") for o in outputs]
+
+    @pytest.mark.parametrize("grouping", [False, True])
+    def test_ar_runner_runs_every_step_of_a_chunk_in_one_call_when_the_pipeline_declares_it(
+        self, monkeypatch, grouping
+    ):
+        """With the capability declared, one AR runner call drives a whole chunk; the outputs are the same either way."""
+        runner = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=grouping)
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        denoise_calls_after_each_call = []
+        original = runner.execute_stepwise
+
+        def counting(scheduler_output):
+            output = original(scheduler_output)
+            denoise_calls_after_each_call.append(runner.pipeline.denoise_calls)
+            return output
+
+        runner.execute_stepwise = counting
+
+        per_call = self._drive_to_completion(runner)
+
+        # Chunk outputs, and the steps they were emitted at, are identical.
+        emitted = [(o.step_index, float(o.result.output[0])) for o in per_call if o.result is not None]
+        assert emitted == [(2, 0.0), (2, 1.0)]
+        assert per_call[-1].finished is True
+        assert runner.pipeline.denoise_calls == 4 and runner.pipeline.scheduler_calls == 4
+        assert runner.pipeline.decode_calls == 2
+        assert runner.pipeline.prepare_calls == 1
+        # One call per chunk with the capability, one per step without; a
+        # grouped call never crosses into the next chunk.
+        assert denoise_calls_after_each_call == ([2, 4] if grouping else [1, 2, 3, 4])
+        assert "req-1" not in runner.state_cache
+
+    def test_ar_runner_continues_a_grouped_chunk_as_a_cached_request(self, monkeypatch):
+        """The first step admits the new request; every further step of the chunk is a cached-request continuation."""
+        runner = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=True)
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        seen = []
+        original_core = DiffusionModelRunner._execute_stepwise_core
+
+        def recording_core(self_, scheduler_output, **kwargs):
+            seen.append(scheduler_output)
+            return original_core(self_, scheduler_output, **kwargs)
+
+        monkeypatch.setattr(DiffusionModelRunner, "_execute_stepwise_core", recording_core)
+
+        first = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4), finished_req_ids={"old"}))
+
+        assert first.get_request_output("req-1").result is not None
+        assert runner.pipeline.prepare_calls == 1
+        assert len(seen) == 2
+        assert [n.request_id for n in seen[0].scheduled_new_reqs] == ["req-1"]
+        continuation = seen[1]
+        assert continuation.scheduled_new_reqs == []
+        assert continuation.scheduled_cached_reqs.request_ids == ["req-1"]
+        assert continuation.finished_req_ids == set()
+        assert continuation.kv_prefetch_job is None and continuation.kv_connector_metadata is None
+        assert continuation.step_id == seen[0].step_id
+
+    def test_ar_runner_grouped_steps_see_the_same_inputs_as_per_step_calls(self, monkeypatch):
+        """Every grouped step sees the latents and timestep step_scheduler just advanced, exactly as per-step calls do."""
+        finals = {}
+        seen = {}
+        for grouping in (False, True):
+            runner = self._make_grouping_runner(_BatchDependentChunkPipeline(), grouping=grouping)
+            monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+            self._drive_to_completion(runner)
+            finals[grouping] = float(runner.pipeline.final_latent)
+            seen[grouping] = runner.pipeline.seen
+
+        # 1 -> +(1+10)=12 -> +(12+5)=29 -> +(29+10)=68 -> +(68+5)=141 in both modes.
+        assert seen[False] == [(1.0, 10.0), (12.0, 5.0), (29.0, 10.0), (68.0, 5.0)]
+        assert seen[True] == seen[False]
+        assert finals[True] == finals[False] == 141.0
+
+    def test_ar_runner_grouped_chunk_stops_at_a_failing_step_and_drops_the_state(self, monkeypatch):
+        """A failure mid-chunk ends the grouped call with that step's error; nothing runs after it."""
+        runner = self._make_grouping_runner(_FailsMidChunkPipeline(), grouping=True)
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+        output = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4)))
+
+        failed = output.get_request_output("req-1")
+        assert failed.finished is True
+        assert failed.result is not None and "boom mid chunk" in failed.result.error
+        assert runner.pipeline.denoise_calls == 2 and runner.pipeline.decode_calls == 0
+        assert "req-1" not in runner.state_cache
+
+    def test_ar_runner_fails_a_chunk_that_never_advances(self, monkeypatch):
+        pipeline = _ChunkedStepPipeline()
+        monkeypatch.setattr(pipeline, "step_scheduler", lambda *args, **kwargs: None)
+        runner = self._make_grouping_runner(pipeline, grouping=True)
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        output = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4)))
+        failed = output.get_request_output("req-1")
+        assert failed.finished
+        assert "within 2 denoise steps" in failed.result.error
+        assert pipeline.denoise_calls == 2
+        assert "req-1" not in runner.state_cache
+
+    def test_base_runner_never_groups_and_the_ar_policy_needs_capability_streaming_and_one_request(self, monkeypatch):
+        """The shared runner is one step per call; the AR runner groups only a lone streaming request that declares it."""
+        base = _make_runner()
+        base.od_config.streaming_output = True
+        base.pipeline = _ChunkedStepPipeline()
+        base.pipeline.supports_chunk_step_grouping = True
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        first = DiffusionModelRunner.execute_stepwise(base, _make_scheduler_output(_make_step_request(4)))
+        assert first.get_request_output("req-1").result is None and base.pipeline.denoise_calls == 1
+
+        ar = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=False)
+        one = _make_cached_scheduler_output()
+        two = _make_batch_scheduler_output([_make_step_request(4), _make_step_request(4)])
+        two.scheduled_new_reqs[1].request_id = two.scheduled_new_reqs[1].req.request_id = "req-2"
+        assert ar._groups_chunk_steps(one) is False  # no capability declared
+        ar.pipeline.supports_chunk_step_grouping = True
+        assert ar._groups_chunk_steps(one) is True
+        assert ar._groups_chunk_steps(two) is False  # not for a batch
+        ar.od_config.streaming_output = False
+        assert ar._groups_chunk_steps(one) is False  # not without streaming output
+
     def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch):
         runner = _make_runner()
         runner.pipeline = _ProfilingStepPipeline()