[None][feat] Ring-buffer reuse of out-of-window SWA pages during decode in KV cache manager v2

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -1549,6 +1549,7 @@ def _build_cache_config(
             max_util_for_resume=kv_cache_config.max_util_for_resume,
             enable_stats=self.enable_stats,
             swa_scratch_reuse=scratch_reuse_config,
+            enable_swa_ring_reuse=kv_cache_config.enable_swa_ring_reuse,
             initial_pool_ratio=kv_cache_config.pool_ratio,
             layers=layer_configs,
         )

tensorrt_llm/llmapi/llm_args.py

@@ -3481,6 +3481,16 @@ class KvCacheConfig(StrictBaseModel, PybindMirror):
         "The maximum utilization of the KV cache for resume. Default is 95%. Only used when using KV cache manager v2 (experimental)."
     )
 
+    enable_swa_ring_reuse: bool = Field(
+        default=False,
+        status="prototype",
+        description=
+        "Recycle a sequence's freshly out-of-window sliding-window-attention pages in "
+        "place to back the new blocks its decode allocates, instead of a free-list "
+        "round-trip. Only fully-orphaned pages (no prefix-reuse retention) are recycled, "
+        "so behavior is unchanged except for which physical slot backs the new block. "
+        "Only used when using KV cache manager v2 (experimental).")
+
     enable_kv_pool_rebalance: bool = Field(
         default=False,
         status="prototype",

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi

@@ -175,6 +175,7 @@ class KVCacheManagerConfig:
     initial_pool_ratio: list[float] | None = None
     swa_scratch_reuse: SwaScratchReuseConfig | None = None
     commit_min_snapshot: bool = False
+    enable_swa_ring_reuse: bool = False
     enable_stats: bool = True
     @property
     def enable_swa_scratch_reuse(self) -> bool: ...

tensorrt_llm/runtime/kv_cache_manager_v2/_config.py

@@ -242,6 +242,24 @@ class KVCacheManagerConfig:
     Required when SSM layers are present.
     """
 
+    enable_swa_ring_reuse: bool = False
+    """
+    When True, a sequence's freshly out-of-window (stale) SWA pages are parked in a small
+    per-sequence pocket and recycled in place to back the new blocks allocated by upcoming
+    resize() calls -- the steady state of decode with a sliding window, where each block
+    boundary crossing retires one block and adds one. This makes the window's storage behave
+    as a ring buffer (the same few physical blocks in rotation) and skips the allocator
+    free-list round-trip; on the same CUDA stream the recycled slot is immediately safe to
+    overwrite.
+
+    Only fully-orphaned stale pages are parked: uncommitted, with committing disallowed for
+    the cache (so the radix tree cannot retain them) and resident on GPU. Committed pages
+    kept for prefix reuse are never touched, so behavior is identical to the default
+    allocate-then-free path except for which physical slot backs the new block. Parked
+    pages are pinned against inter-level eviction (migrating them would waste bandwidth
+    on dead data) until consumed or drained; the pocket is drained on suspend()/close().
+    """
+
     enable_stats: bool = True
     """
     Collect V2 KV cache allocation, reuse, and transfer statistics.

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py

@@ -197,6 +197,7 @@ class _KVCache:
         "_never_resumed",
         "_enable_swa_scratch_reuse",
         "_scratch_slots",
+        "_swa_reuse_pocket",
         "_pending_stats",
         "__rawref__",
     )
@@ -245,6 +246,9 @@ class _KVCache:
     # Managed via delta in resize(): existing slots are reused across resize calls,
     # only the additional needed slots are allocated. Freed on teardown/suspend.
     _scratch_slots: TypedIndexList[LifeCycleId, list[ScratchSlotLock]]
+    # Freshly out-of-window, fully-orphaned pages parked for in-place reuse by upcoming
+    # decode blocks (enable_swa_ring_reuse). Drained on suspend()/close().
+    _swa_reuse_pocket: TypedIndexList[LifeCycleId, list[_PageHolder]]
     _pending_stats: _PendingStats
 
     def __init__(
@@ -288,6 +292,9 @@ def __init__(
         self._scratch_slots = make_typed(
             lambda _: list[ScratchSlotLock](), manager._storage.num_life_cycles
         )
+        self._swa_reuse_pocket = make_typed(
+            lambda _: list[_PageHolder](), manager._storage.num_life_cycles
+        )
         self._pending_stats = _PendingStats()
         self.__rawref__ = rawref.NULL
         if reuse_match is not None:
@@ -519,6 +526,10 @@ def close(self) -> None:
             return
         self.discard_pending_stats()
         self.stop_committing()
+        if any(self._swa_reuse_pocket):
+            with self._record_event():
+                for pocket in self._swa_reuse_pocket:
+                    pocket.clear()  # release parked pages back to the allocator
         assert NDEBUG or self._check_sanity()
         manager = self.manager
         if self.capacity > 0:
@@ -752,8 +763,32 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
                         num_new_blocks_to_add = new_num_blocks - max(stale_end, old_num_blocks)
                     num_new_slots[lc] = num_new_blocks_to_add * beam_width
 
+            # SWA ring reuse: consume pages parked in the reuse pocket by earlier
+            # resize() calls (freshly out-of-window, fully-orphaned pages -- see the refill
+            # below) to back new blocks, instead of requesting slots from the allocator.
+            # Scratch mode has its own slot machinery, so the pocket is bypassed there.
+            enable_ring_reuse = manager._init_config.enable_swa_ring_reuse and not enable_scratch
+            num_pocket_slots = filled_list(0, num_life_cycles)
+            if enable_ring_reuse:
+                for lc in typed_range(num_life_cycles):
+                    # Re-validate parked pages before use. Parked pages are excluded from
+                    # eviction (see the refill below), so with the current eviction policy
+                    # this drops nothing; it keeps the consume path locally correct even if
+                    # a future policy migrates or invalidates held pages.
+                    pocket = self._swa_reuse_pocket[lc]
+                    if pocket:
+                        pocket[:] = [
+                            h
+                            for h in pocket
+                            if not h.page.is_committed()
+                            and h.page.has_valid_slot
+                            and h.page.cache_level == GPU_LEVEL
+                        ]
+                    num_pocket_slots[lc] = min(len(pocket), num_new_slots[lc])
+
             net_alloc_counts = make_typed(
-                lambda lc: num_new_slots[lc] + delta_scratch_slots[lc], num_life_cycles
+                lambda lc: num_new_slots[lc] - num_pocket_slots[lc] + delta_scratch_slots[lc],
+                num_life_cycles,
             )
             storage = self._storage
             if any(c > 0 for c in net_alloc_counts):
@@ -778,9 +813,18 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
             # Combine slots and distribute
             slots = make_typed(lambda _: list[Slot](), num_life_cycles)
             for lc in typed_range(num_life_cycles):
-                slots[lc] = new_slots[lc] + [
-                    lock.detach_slot() for lock in excess_scratch_slots[lc]
-                ]
+                # Pocket pages transfer their slots directly: moving the slot out
+                # invalidates the page, making its destruction a no-op. Their last use
+                # was issued on this cache's stream, so in-order execution makes them
+                # immediately safe for the new block without an event wait.
+                slots[lc] = (
+                    new_slots[lc]
+                    + [lock.detach_slot() for lock in excess_scratch_slots[lc]]
+                    + [
+                        self._swa_reuse_pocket[lc].pop().page.move_to_new_slot()
+                        for _ in range(num_pocket_slots[lc])
+                    ]
+                )
                 new_slots[lc].clear()
                 excess_scratch_slots[lc].clear()
 
@@ -851,6 +895,30 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
                         ).lock(self, beam_index, ordinal, lc, skip_wait=True)
                 self._blocks.append(SeqBlock(block, None))
             assert all(len(slots[lc]) == 0 for lc in typed_range(num_life_cycles))
+            # Park this call's freshly staled, fully-orphaned pages for reuse by upcoming
+            # resize() calls. Qualifying pages are uncommitted with committing disallowed
+            # (the radix tree cannot retain them; backup_holders holds the sole reference)
+            # and GPU-resident. Holding the _PageHolder keeps the slot alive. Bounded per
+            # life cycle to the ring working set (one block retired per boundary crossing,
+            # +1 for straddle).
+            if enable_ring_reuse and self._commit_state != self.CommitState.ALLOWED:
+                for _, _, holder_lc, holder in backup_holders:
+                    if len(self._swa_reuse_pocket[holder_lc]) >= 2:
+                        continue
+                    page = holder.page
+                    if (
+                        page.is_committed()
+                        or not page.has_valid_slot
+                        or page.cache_level != GPU_LEVEL
+                    ):
+                        continue
+                    # Pin against eviction: unlocking (in _unlock_stale_blocks) scheduled
+                    # this held page for eviction when lower cache levels exist. Migrating
+                    # it would waste bandwidth on dead data and move it off GPU_LEVEL,
+                    # invalidating the park-time checks above before the page is consumed.
+                    if page.scheduled_for_eviction:
+                        storage.exclude_from_eviction(page)
+                    self._swa_reuse_pocket[holder_lc].append(holder)
         self._capacity = capacity
         self._history_length = history_length
         self._refresh_generation_alloc_ready()
@@ -1020,6 +1088,8 @@ def suspend(self) -> None:
                     self.set_base_page_index_buf(beam_idx, lc, None)
         ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id
         with self._record_event():  # used by _SharedPageLock.__del__
+            for pocket in self._swa_reuse_pocket:
+                pocket.clear()  # release parked pages back to the allocator
             for ordinal, beam_idx, lc_idx in self._active_pages():
                 beam_block = (
                     self._block(ordinal, beam_idx)

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -623,6 +623,13 @@
       "kind": "value",
       "path": "kv_cache_config.enable_partial_reuse"
     },
+    {
+      "allowed_values": [],
+      "annotation": "<class 'bool'>",
+      "converter": "",
+      "kind": "value",
+      "path": "kv_cache_config.enable_swa_ring_reuse"
+    },
     {
       "allowed_values": [],
       "annotation": "<class 'bool'>",
@@ -3043,6 +3050,13 @@
       "kind": "value",
       "path": "kv_cache_config.enable_partial_reuse"
     },
+    {
+      "allowed_values": [],
+      "annotation": "<class 'bool'>",
+      "converter": "",
+      "kind": "value",
+      "path": "kv_cache_config.enable_swa_ring_reuse"
+    },
     {
       "allowed_values": [],
       "annotation": "<class 'bool'>",

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

@@ -697,6 +697,88 @@ def test_naive_perf(self, interval, profile: bool) -> None:
             profiler.dump_stats("profiler.prof")
 
 
+class TestSwaRingReuse(TestNoBatching):
+    """KVCacheManagerConfig.enable_swa_ring_reuse.
+
+    Drives a disaggregated-decode-style uncommitted generation (committing disallowed,
+    history advanced directly via resize) across many sliding-window advances and checks:
+      * data correctness via FakeEngine refcheck at every step, and
+      * allocator traffic: with the flag ON, the windowed layer's freshly staled pages are
+        parked in the reuse pocket and recycled in place, so after the pocket warms up the
+        global allocator receives no further requests for the windowed life cycle (only the
+        full-attention life cycle still allocates one block per boundary crossing); with the
+        flag OFF, every boundary crossing requests one block for each life cycle.
+    """
+
+    def _run_uncommitted_decode(self, enable: bool) -> "tuple[int, int]":
+        tokens_per_block = 32
+        window = 2 * tokens_per_block
+        self.cfg = create_config(
+            tokens_per_block,
+            32 << 20,
+            0,
+            0,
+            num_layers=2,  # layer 0: sliding window; layer 1: full attention
+            window_size=window,
+            sink_tokens=0,
+        )
+        self.cfg.enable_swa_ring_reuse = enable
+        self.engine = FakeEngine(self.cfg)
+        self.manager = KVCacheManager(self.cfg)
+
+        prompt_len = window  # no prompt block is stale at generation start
+        decode_len = 8 * tokens_per_block
+        req = self.new_request(0, None, prompt_len, decode_len)
+        kv_cache = req.kv_cache
+
+        alloc_requests: "list[int]" = []
+        orig_new_gpu_slots = StorageManager.new_gpu_slots
+
+        def spying_new_gpu_slots(mgr, counts, *args, **kwargs):
+            alloc_requests.append(sum(int(c) for c in counts))
+            return orig_new_gpu_slots(mgr, counts, *args, **kwargs)
+
+        with TemporaryCudaStream([]) as s:
+            stream = cast(CudaStream, s.handle)
+            self.assertTrue(kv_cache.resume(stream))
+            # Disaggregated decode style: KV arrives from elsewhere; nothing is committed,
+            # so freshly staled blocks are fully orphaned.
+            kv_cache.stop_committing()
+            self.assertTrue(kv_cache.resize(prompt_len, prompt_len))
+            self.engine.execute([Step(kv_cache, req.prompt, [])], stream)
+            history = list(req.prompt)
+            StorageManager.new_gpu_slots = spying_new_gpu_slots  # type: ignore[method-assign]
+            try:
+                for _ in range(decode_len):
+                    token = self.next_token()
+                    self.assertTrue(kv_cache.resize(len(history) + 1, len(history)))
+                    self.engine.execute([Step(kv_cache, [token], history)], stream)
+                    history.append(token)
+            finally:
+                StorageManager.new_gpu_slots = orig_new_gpu_slots  # type: ignore[method-assign]
+            self.engine.execute([Step(kv_cache, [], history)], stream)
+        s.take_finish_event().synchronize()
+        kv_cache.close()
+        total_requested = sum(alloc_requests)
+        num_crossings = div_up(prompt_len + decode_len, tokens_per_block) - div_up(
+            prompt_len, tokens_per_block
+        )
+        return total_requested, num_crossings
+
+    def test_disabled_baseline_traffic(self) -> None:
+        total, crossings = self._run_uncommitted_decode(enable=False)
+        # One new block per boundary crossing for the windowed AND the full-attention
+        # life cycle.
+        self.assertEqual(total, 2 * crossings)
+
+    def test_enabled_reuses_stale_window_pages(self) -> None:
+        total, crossings = self._run_uncommitted_decode(enable=True)
+        # The windowed life cycle allocates only until its first page goes stale (one
+        # block); every later boundary crossing is served from the reuse pocket. The
+        # full-attention life cycle still allocates one block per crossing.
+        self.assertEqual(total, crossings + 1)
+
+
 class TestBatching(TestKVCacheManagerV2):
     num_requests: int
     avg_length: int