[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