[None][feat] Reuse freshly out-of-window SWA pages for new decode blocks in KV cache manager v2#16288
Open
lishicheng1996-nv wants to merge 1 commit into
Open
Conversation
…cks in KV cache manager v2 During decode with a sliding window, every tokens_per_block tokens retire one out-of-window block and allocate one new block. Today the staled page takes a full free-list round-trip (release with a ready event, later re-allocation with cross-stream event bookkeeping) while the new block draws a different slot from the allocator. This change adds an opt-in per-sequence reuse pocket (kv_cache_config.enable_swa_decode_slot_reuse, default off): a resize() that orphans a freshly staled page parks it, and a later resize() that needs a new block consumes the parked page's slot directly. On the same CUDA stream the recycled slot is immediately safe to overwrite, so for a 131-token window the sequence's SWA footprint is served by the same ~3 physical blocks in rotation. Only fully-orphaned pages are parked: uncommitted, with committing disallowed for the cache (so the block radix tree cannot retain them) and GPU-resident. Committed pages kept for prefix reuse are never touched. The pocket is bounded (2 pages per life cycle), consumed and refilled only on resize() success paths so the OutOfPagesError revert is unaffected, and drained on suspend()/close(). Adds TestSwaDecodeSlotReuse: a disaggregated-decode-style uncommitted generation with FakeEngine refcheck at every step plus deterministic allocator traffic assertions (flag off: one allocation per boundary crossing per life cycle; flag on: the windowed life cycle allocates only until its first page is parked). Signed-off-by: Shicheng Li <shicli@nvidia.com>
Contributor
📝 WalkthroughWalkthroughAdds an experimental SWA decode slot-reuse configuration, wires it through KV cache manager construction, recycles eligible stale pages during resize, clears parked pages on lifecycle transitions, and tests allocator traffic with the option enabled and disabled. ChangesSWA decode slot reuse
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant KVCacheManagerV2
participant _KVCache
participant StorageManager
participant SWAReusePocket
KVCacheManagerV2->>_KVCache: resize with reuse enabled
_KVCache->>SWAReusePocket: consume parked page slots
_KVCache->>StorageManager: request remaining GPU slots
_KVCache->>SWAReusePocket: park eligible stale orphaned pages
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)
658-712: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSolid steady-state coverage; consider adding pocket-lifecycle edge cases in
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py.Coverage for the primary decode steady-state (disabled vs. enabled allocator traffic + refcheck correctness) is sufficient and the math checks out. However, the PR's own safety claims aren't directly exercised here:
suspend()/close()draining_swa_reuse_pocket(e.g., suspend mid-decode with a warmed-up pocket, then resume and verify the next crossing allocates fresh rather than reusing a stale/invalid pocket entry).- The 2-entries-per-life-cycle pocket cap (e.g., a resize() that stales two blocks in one call, such as a larger history_length jump).
OutOfPagesErrorrollback interaction with the pocket (PR description states rollback is unchanged, but no test asserts pocket state survives a failedresize()untouched).These are follow-up-worthy rather than blocking, given the core mechanism is already well-covered.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py` around lines 658 - 712, Extend the KV cache manager tests around _run_uncommitted_decode to cover pocket lifecycle safety: suspend and close must drain _swa_reuse_pocket, and after resume the next block crossing must allocate fresh storage rather than reuse stale entries. Add a resize scenario that stales two blocks in one call to verify the two-entry-per-lifecycle cap. Add an OutOfPagesError resize test asserting rollback leaves the pocket state unchanged.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`:
- Around line 888-905: Revalidate each pocketed holder’s page immediately before
reuse in the `_swa_reuse_pocket` pop path, including `cache_level` and
`has_valid_slot` (and existing commitment checks), before calling
`move_to_new_slot()`. Remove stale or no-longer-eligible holders instead of
reusing them, or otherwise pin them against migration so pocketed pages remain
valid until reuse.
---
Nitpick comments:
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 658-712: Extend the KV cache manager tests around
_run_uncommitted_decode to cover pocket lifecycle safety: suspend and close must
drain _swa_reuse_pocket, and after resume the next block crossing must allocate
fresh storage rather than reuse stale entries. Add a resize scenario that stales
two blocks in one call to verify the two-entry-per-lifecycle cap. Add an
OutOfPagesError resize test asserting rollback leaves the pocket state
unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7c06be1f-ab1e-49b8-b52b-8a2d1ccc7082
📒 Files selected for processing (6)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
Comment on lines
+888
to
+905
| # 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_slot_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 | ||
| self._swa_reuse_pocket[holder_lc].append(holder) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect how held (unlocked) uncommitted pages interact with eviction scheduling.
rg -n "scheduled_for_eviction|schedule_for_eviction|exclude_from_eviction" tensorrt_llm/runtime/kv_cache_manager_v2/_page.py tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 4142
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== _core/_kv_cache.py around the pocket consumer =="
sed -n '790,835p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
echo
echo "== _page.py around page state, eviction, and _PageHolder =="
sed -n '1,220p' tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
echo
sed -n '220,380p' tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
echo
echo "== _storage_manager.py around schedule/exclude and migration paths =="
sed -n '520,720p' tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 23938
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== relevant definitions and references =="
rg -n "class _PageHolder|def __del__|has_valid_slot|scheduled_for_eviction|lock|unlock|move_to_new_slot|is_evictable|exclude_from_eviction|schedule_for_eviction" \
tensorrt_llm/runtime/kv_cache_manager_v2/_page.py \
tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py \
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== storage is_evictable and eviction controller hooks =="
sed -n '1,220p' tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller.py
echo
sed -n '220,360p' tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller.py
echo
echo "== storage_manager migrations around last-level pages =="
sed -n '1020,1085p' tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py
echo
echo "== page unlock/destructor paths =="
sed -n '240,320p' tensorrt_llm/runtime/kv_cache_manager_v2/_page.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 322
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== kv_cache_manager_v2 file tree =="
fd -a . tensorrt_llm/runtime/kv_cache_manager_v2
echo
echo "== likely eviction-controller references =="
rg -n "def is_evictable|class .*Controller|schedule_for_eviction|exclude_from_eviction|node_ref" tensorrt_llm/runtime/kv_cache_manager_v2Repository: NVIDIA/TensorRT-LLM
Length of output: 8203
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== storage_manager is_evictable =="
sed -n '380,430p' tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py
echo
echo "== eviction controller =="
sed -n '156,230p' tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py
echo
echo "== page holder / page destructor paths =="
sed -n '60,110p' tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
echo
sed -n '260,310p' tensorrt_llm/runtime/kv_cache_manager_v2/_page.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 9569
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Slot methods and move_to_new_slot semantics =="
sed -n '1,220p' tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py
echo
echo "== any direct callers of move_to_new_slot =="
rg -n "move_to_new_slot\(" tensorrt_llm/runtime/kv_cache_manager_v2Repository: NVIDIA/TensorRT-LLM
Length of output: 6893
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Slot.move_to_new_slot =="
sed -n '250,310p' tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py
echo
echo "== page slot-transfer helpers =="
sed -n '480,520p' tensorrt_llm/runtime/kv_cache_manager_v2/_page.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3460
Re-check pocketed pages before reuse
Held pages at GPU_LEVEL are still evictable when lower cache levels exist, so _swa_reuse_pocket can’t assume cache_level and has_valid_slot stay unchanged between park and pop. Re-validate before move_to_new_slot() or pin pocketed pages against migration.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py` around lines 888
- 905, Revalidate each pocketed holder’s page immediately before reuse in the
`_swa_reuse_pocket` pop path, including `cache_level` and `has_valid_slot` (and
existing commitment checks), before calling `move_to_new_slot()`. Remove stale
or no-longer-eligible holders instead of reusing them, or otherwise pin them
against migration so pocketed pages remain valid until reuse.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
During decode with a sliding-window-attention (SWA) layer, every
tokens_per_blocktokens the window slides across a block boundary: one out-of-window block is retired and one new block is allocated. Today the staled page takes a full free-list round-trip (release with a ready event, later re-allocation with cross-stream event bookkeeping) while the new block draws a different slot from the global allocator.This PR adds an opt-in per-sequence reuse pocket to KV cache manager v2 (
kv_cache_config.enable_swa_decode_slot_reuse, default off): aresize()that orphans a freshly staled page parks it; a laterresize()that needs a new block consumes the parked page's slot directly. On the same CUDA stream the recycled slot is immediately safe to overwrite, so e.g. for a 131-token window a sequence's SWA footprint is served by the same ~3 physical blocks in rotation instead of continuous allocate-then-free.Safety gating (behavior is unchanged except for which physical slot backs the new block):
resize()success paths, so theOutOfPagesErrorrevert path (_lock_held_blocks) is unaffected;suspend()/close().Plumbed through
KvCacheConfig.enable_swa_decode_slot_reuse(prototype) -> pyexecutorKVCacheManagerV2->KVCacheManagerConfig.Test Coverage
New
TestSwaDecodeSlotReuseintests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py: drives a disaggregated-decode-style uncommitted generation across 8 window advances with FakeEngine refcheck at every step, plus deterministic allocator-traffic assertions:2 x crossings);crossings + 1).Existing
TestNoBatchingbattery passes unchanged under the patched module (flag off is a no-op: two cheap branches).Perf A/B on DeepSeek-V4 disaggregated GEN (GB200, dep32/EP32, 256K ISL, per-rank batch 84, MTP3, MoE backend MEGAMOE_DEEPGEMM, 2688/2688 requests completed on both arms):
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Summary by CodeRabbit
New Features
Bug Fixes
Tests