diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 2f93f0fb028f..67a8411bf98f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/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, ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5c6c89e88b3c..44dcc8ddecf7 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/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", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 5d7c50e676cb..a50f2af0252f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/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: ... diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py index 4a0e9a5f7064..26715e1c8515 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py +++ b/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. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 191f2398e92a..a0eeb9c0ca01 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/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) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 1362f671ccfb..7f40e8efd637 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -623,6 +623,13 @@ "kind": "value", "path": "kv_cache_config.enable_partial_reuse" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_config.enable_swa_ring_reuse" + }, { "allowed_values": [], "annotation": "", @@ -3043,6 +3050,13 @@ "kind": "value", "path": "kv_cache_config.enable_partial_reuse" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_config.enable_swa_ring_reuse" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 8c8daa74b115..2cf3733dd882 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/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