[AR-Diffusion] Bound KV metadata for long rollouts (#7498)

Signed-off-by: wtz2333 <2955110911@qq.com>

tests/diffusion/ar_diffusion/test_capability_runner.py

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

tests/diffusion/ar_diffusion/test_kv_orchestration.py

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

vllm_omni/experimental/ar_diffusion/capability.py

@@ -81,6 +81,7 @@ class ARDiffusionKVCacheSpec:
     sink_frames: int = 0
     reset_at_boundary: bool = False
     cross_attention: tuple[ARDiffusionCrossAttentionKVSpec, ...] = ()
+    # Initial vLLM storage horizon; extended to fit an in-flight span.
     max_model_len: int = 1 << 20
     max_scratch_tokens_per_branch: int = 0
     model_owned_state_bytes_per_session: int = 0

vllm_omni/experimental/ar_diffusion/kv_cache/manager.py

@@ -15,6 +15,7 @@ from __future__ import annotations
 import inspect
 import os
 from collections.abc import Collection, Iterable, Sequence
+from typing import cast
 
 import torch
 from vllm.logger import init_logger
@@ -30,6 +31,7 @@ from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor
 from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
+    ChunkWindowManager,
     ChunkWindowSpec,
     allocate_kv_pool_with_views,
     chunk_slot_mapping,
@@ -49,10 +51,10 @@ class ARDiffusionRequestAdapter:
     real ``KVCacheManager`` against this adapter so the surface cannot silently
     drift across vLLM versions.
 
-    An 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.
+    ``completed_chunks`` and ``absolute_num_computed_tokens`` track the model
+    timeline. vLLM sees compacted storage positions through ``num_computed_tokens``
+    and ``num_tokens``. A chunk's denoise steps reuse the same slots; only a
+    successful commit advances the timeline and compacts evicted metadata.
     """
 
     def __init__(
@@ -66,6 +68,7 @@ class ARDiffusionRequestAdapter:
         self._chunk_size = chunk_size
         self._prefill = prefill_prefix_tokens
         self._completed_chunks = 0
+        self.compacted_tokens = 0
         # Filled only when cross-request prefix reuse is enabled (Phase 3).
         self.block_hashes: list = []
         self.skip_reading_prefix_cache = True
@@ -79,14 +82,19 @@ class ARDiffusionRequestAdapter:
         self.num_in_flight_tokens = 0
 
     @property
-    def num_computed_tokens(self) -> int:
-        """Persistent KV already materialized (committed chunks + prefill)."""
+    def absolute_num_computed_tokens(self) -> int:
+        """Cumulative model position, unaffected by storage compaction."""
         return self._prefill + self._completed_chunks * self._chunk_size
 
+    @property
+    def num_computed_tokens(self) -> int:
+        """Committed position in the compact vLLM block table."""
+        return self.absolute_num_computed_tokens - self.compacted_tokens
+
     @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
+        """Storage position once the in-flight chunk is committed."""
+        return self.num_computed_tokens + self._chunk_size
 
     @property
     def num_prompt_tokens(self) -> int:
@@ -540,12 +548,9 @@ class ARDiffusionKVCache:
     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).
+        Returns the request's compact storage block table.
         """
-        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)
+        table = self.allocate_token_slots(adapter, self.spec.chunk_size)
         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",
@@ -561,6 +566,10 @@ class ARDiffusionKVCache:
         """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}")
+        # vLLM clips slot allocation at max_model_len. Admit the entire
+        # in-flight span in compact storage coordinates, independent of the
+        # model's absolute frame positions.
+        self.manager.max_model_len = max(self.manager.max_model_len, adapter.num_computed_tokens + 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")
@@ -618,6 +627,8 @@ class ARDiffusionKVCache:
             adapter.num_computed_tokens,
             num_prompt_tokens=adapter.num_prompt_tokens,
         )
+        window_manager = cast(ChunkWindowManager, self.manager.coordinator.single_type_managers[0])
+        adapter.compacted_tokens += window_manager.compact_block_table(adapter.request_id)
         _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
 
     # -- pool-backed K/V access --------------------------------------------

vllm_omni/experimental/ar_diffusion/kv_cache/paged.py

@@ -197,6 +197,28 @@ class ChunkWindowManager(SlidingWindowManager):
             reset_at_boundary=spec.reset_at_boundary,
         )
 
+    def compact_block_table(self, request_id: str) -> int:
+        """Remove the evicted gap after the sink; return its size in tokens.
+
+        Call only after commit/eviction, then shift the request's storage
+        position by the returned amount. Physical pages and model positions
+        are unchanged, including any allocated but not yet committed tail.
+        """
+        if self.enable_caching:
+            raise RuntimeError("AR block-table compaction requires prefix caching to be disabled")
+        blocks = self.req_to_blocks.get(request_id, [])
+        start = self.kv_cache_spec.sink_chunks
+        end = start
+        while end < len(blocks) and blocks[end] == self._null_block:
+            end += 1
+        if end == start:
+            return 0
+        del blocks[start:end]
+        if request_id in self.num_cached_block:
+            cached = self.num_cached_block[request_id]
+            self.num_cached_block[request_id] = min(cached, start) + max(0, cached - end)
+        return (end - start) * self.block_size
+
     def remove_skipped_blocks(
         self,
         request_id: str,