[None][fix] KVCacheManagerV2: keep partial rewind endpoints reusable via per-page token coverage - #17122
[None][fix] KVCacheManagerV2: keep partial rewind endpoints reusable via per-page token coverage#17122lowsfer wants to merge 5 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe KV cache now records token coverage on committed pages. Block insertion, snapshotting, rebasing, pruning, event generation, introspection, and reuse use this coverage to preserve valid partial pages and reject insufficient pages. ChangesCoverage-aware KV cache lifecycle
Estimated code review effort: 4 (Complex) | ~60 minutes 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}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h (1)
113-114: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign the constructor parameter name with the definition.
The declaration names the parameter
numTokensInBlock, which is also the member name. The definition inpage.cpp(lines 99-106) names itnumTokensInBlock_. Use one name in both places.As per coding guidelines: "keep declaration and definition parameter names consistent".
♻️ Proposed rename in the declaration
- CommittedPage(StorageManager* mgr, SharedPtr<Block> blk, LifeCycleId lc, CacheLevel level, int numTokensInBlock, - Priority prio); + CommittedPage(StorageManager* mgr, SharedPtr<Block> blk, LifeCycleId lc, CacheLevel level, int numTokensInBlock_, + Priority prio);🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h` around lines 113 - 114, Rename the CommittedPage constructor declaration parameter numTokensInBlock to numTokensInBlock_ so it matches the definition in page.cpp, without changing behavior or the member initialization.Source: Coding guidelines
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp (1)
221-229: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace
dynamic_castwithstatic_cast.Line 216 already tests
page->isCommitted(), so the pointer type is known inside this branch. Thedynamic_castand the following null test add runtime type information lookup with no benefit.As per coding guidelines: "Use the least forceful cast; ... reserve
reinterpret_castas a last resort, and avoiddynamic_cast."♻️ Proposed cast simplification
- auto* cp = dynamic_cast<CommittedPage*>(page.get()); - if (cp) - { - if (cp->block == nullptr || cp->block->isOrphan() || !cp->block->holdsPage(*cp)) - manager->excludeFromEviction(*page); - } + auto* cp = static_cast<CommittedPage*>(page.get()); + if (cp->block == nullptr || cp->block->isOrphan() || !cp->block->holdsPage(*cp)) + { + manager->excludeFromEviction(*page); + }🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp` around lines 221 - 229, In the committed-page handling branch guarded by page->isCommitted(), replace the dynamic_cast to CommittedPage with static_cast and remove the now-unnecessary null check, while preserving the existing block/orphan/holdsPage eviction exclusion logic.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp`:
- Around line 221-229: In the committed-page handling branch guarded by
page->isCommitted(), replace the dynamic_cast to CommittedPage with static_cast
and remove the now-unnecessary null check, while preserving the existing
block/orphan/holdsPage eviction exclusion logic.
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h`:
- Around line 113-114: Rename the CommittedPage constructor declaration
parameter numTokensInBlock to numTokensInBlock_ so it matches the definition in
page.cpp, without changing behavior or the member initialization.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bce1479c-4a06-4289-8c34-37efae466b24
📒 Files selected for processing (15)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.hcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpptensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_introspection.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
|
/bot run --disable-fail-fast |
|
PR_Github #63022 [ run ] triggered by Bot. Commit: |
30c9f9f to
e13eda5
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63027 [ run ] triggered by Bot. Commit: |
|
PR_Github #63022 [ run ] completed with state |
e13eda5 to
31a499d
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63041 [ run ] triggered by Bot. Commit: |
|
PR_Github #63027 [ run ] completed with state |
|
PR_Github #63041 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63092 [ run ] triggered by Bot. Commit: |
|
PR_Github #63092 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63122 [ run ] triggered by Bot. Commit: |
|
PR_Github #63122 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63171 [ run ] triggered by Bot. Commit: |
|
PR_Github #63171 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63201 [ run ] triggered by Bot. Commit: |
|
PR_Github #63201 [ run ] completed with state |
|
PR_Github #63315 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63359 [ run ] triggered by Bot. Commit: |
|
PR_Github #63359 [ run ] completed with state
|
e5718fb to
425ca49
Compare
|
/bot run |
|
PR_Github #63450 [ run ] triggered by Bot. Commit: |
|
PR_Github #63450 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63458 [ run ] triggered by Bot. Commit: |
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
A partial trailing tree block (a rewind endpoint, e.g. 16 tokens of a 32-token block) was destroyed when a longer sibling covering the same tokens was created, and refused when it was created second. The covering block does not necessarily hold a page for every life cycle at that token boundary -- typically a SWA life cycle whose page was never allocated because the block was already outside the sliding window when it was committed -- so the reusable endpoint was lost either way. Generalize the mechanism SSM snapshots already used. num_tokens_in_block moves from SsmCommittedPage up to CommittedPage, meaning the number of leading tokens of the owning block that the page's data is valid for. Attention reads it as a prefix-valid length; SSM keeps its exact-endpoint meaning. Each (block, life cycle) slot keeps only the widest page. Block.__init__ now moves a covered sibling's pages into the new block instead of dropping them, and _snapshot_partial_block_to_tree attaches partial pages to a longer sibling. Prefix matching honours the recorded count for attention as well as SSM, and _commit_block's rebase path no longer adopts a page that covers fewer tokens than the block spans, which would have fed uninitialized KV into a live request. Block.storage reads now go through Block.get_page(); the slot never holds a dangling ref because CommittedPage.__del__ unlinks before invalidating. KVCacheEventManager drops _resolve_page_ref and does not announce a life cycle whose page covers less than the whole block, since the event payload carries the block's full token list. SsmCommittedPage is deleted; it no longer adds any field. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Mirrors the Python changes on this branch so both backends agree. numTokensInBlock moves from SsmCommittedPage up to CommittedPage, and SsmCommittedPage is deleted -- it no longer adds any field. The count means the number of leading tokens of the owning block the page's data is valid for: a prefix length for attention, an exact checkpoint for SSM. Block gains pageCoverage(), holdsPage(), reservePageSlot() and adoptPagesFrom(). addOrGetExistingBlock() now inserts the new block before detaching the covered sibling and moves that sibling's pages over, so a rewind endpoint survives the longer block that replaces it. pruneMatch() becomes the same fixed-point loop as Python: SSM truncates to an exact snapshot, then every attention life cycle clamps the match to the coverage of each page it still needs, skipping the stale range in between. That subsumes the old full-attention and SWA-sink passes. commitBlock()'s rebase path no longer adopts a tree page covering fewer tokens than the block spans, which would have fed uninitialized KV into a live request. snapshotPartialBlockToTree() drops its exact-span guard and reports the life cycles it attached. eventManager's pageCoversBlock() now checks every page instead of only SSM ones -- with numTokensInBlock on the base class the old dynamic_cast would have silently stopped filtering attention pages. _checkSanity() gains the scratch carve-out from the sibling Python commit, and _getTreeBlock() no longer asserts a page still points at that block, since a page may be moved to a longer sibling or replaced by a wider one. reuse_match_pages() reports the recorded count for all pages, not just SSM. The new coverage tests now go through _introspection so they exercise both backends; test_planned_drop_handle_rejects_partial_coverage stays Python-only because it has to write a page field directly. kvCacheManagerV2StatsTest.cpp is updated for the widened CommittedPage constructor. Without it the google-tests target fails to compile, which fails every *-CPP-* CI stage at fixture setup -- that target is built inside those stages, not by the Build-x86_64/Build-SBSA jobs. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
_copyPageToTreeBlock() built its CommittedPage by routing through a temporary UncommittedPage. That temporary claims the (kvCache, ordinal, beam, lifecycle) identity which the live sequence's own uncommitted page already holds, so ~UncommittedPage's slot-ownership check failed -- a throwing destructor, hence std::terminate under TLLM_DEBUG_MODE=1. Build the CommittedPage straight from the new slot, as _copy_page_to_tree_block() does in Python. The end state is the same page with the same slot, post-copy ready event, priority and manager, installed through the same replacePage(); the temporary is simply not created. The failure it removes is debug-only -- in release the temporary registered nothing and its destructor released nothing, since the slot had already been transferred -- but the fix is on a non-debug path, so it belongs with the per-page-coverage work that made this path hot rather than in the debug-check cleanup. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
425ca49 to
53419c6
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63503 [ run ] triggered by Bot. Commit: |
|
PR_Github #63458 [ run ] completed with state |
|
PR_Github #63503 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63568 [ run ] triggered by Bot. Commit: |
tongyuantongyu
added
the
Release Blocker
|
/bot run --disable-fail-fast |
|
PR_Github #63640 [ run ] triggered by Bot. Commit: |
|
PR_Github #63568 [ run ] completed with state |
6 participants
Dev Engineer Review
CommittedPage.numTokensInBlock.SsmCommittedPagetype and conversion path.QA Engineer Review
tests/integration/test_lists/,test-db/, orqa/.Description
A partial trailing tree block — a rewind endpoint, e.g. 16 tokens of a 32-token block — was lost in both orderings:
Either way the reusable endpoint disappeared. The covering block does not necessarily hold a page for every life cycle at that token boundary — typically a SWA life cycle whose page was never allocated, because the block was already outside the sliding window when it was committed.
This is the issue #16713 addresses. Rather than preventing block replacement, this PR generalizes the mechanism SSM snapshots already used.
num_tokens_in_blockmoves fromSsmCommittedPageup toCommittedPage, where it means the number of leading tokens of the owning block that this page's data is valid for. The two life-cycle families read it differently:>=);Each
(block, life cycle)slot keeps only the page with the largest recorded count. That deliberately allows at most one checkpoint per block: two conversation turns rarely end inside the same block, and if they do, a reuse miss is acceptable.SsmCommittedPageis deleted — it no longer adds any field.Consequences:
Block.__init__/addOrGetExistingBlockmoves a covered sibling's pages into the new block instead of dropping them._snapshot_partial_block_to_treeattaches partial pages to a longer sibling._commit_block's rebase path no longer adopts a tree page covering fewer tokens than the block spans. That was a latent bug this feature exposed: adopting such a page would have fed uninitialized KV for the uncovered tail into a live request.KVCacheEventManagerdoes not announce a life cycle whose page covers less than the whole block, since the event payload carries the block's full token list and cannot express a shorter valid prefix.A slot's page is installed through
Block.replace_page()/replacePage()and detached onlythrough
Block.unlink_page()/unlinkPage(), which are the only two places the block-page link ismutated.
can_replace_page()is a pure predicate, so a caller can test whether a wider page may takeover before doing work that can fail —
convert_to_committed()and_copy_page_to_tree_block()bothrely on that, so an
OutOfPagesErrormid-way cannot destroy a still-usable shorter snapshot.replace_page()clears the superseded page'sblockback-pointer, which is required for memorysafety in C++. The displaced page can outlive the call while a live request still holds it, and it
is no longer reachable from
storage, so_release_pages()/releasePages()— which walkstorage— would never clear it. Once the block died,~PageHolder(block->isOrphan(),block->holdsPage()) and~CommittedPage(blk->unlinkPage()) would read freed memory. Python issafe either way, since
blockis a rawref thatBlock.__del__invalidates, but both backends nowclear it so the two agree and the invariant is uniform: a committed page's
blockis set exactlywhile it occupies that block's slot.
This is reachable through ordinary batched serving, with no eviction pressure: a request ends inside
a block and stays open holding the pages it committed; a second turn grows that block and commits
over them, displacing the held pages; a third turn replaces the block with a longer sibling and
destroys it.
Two smaller fixes ride along:
_check_sanity's scratch carve-out kept using the current scratch range, which goes empty exactly whencommit()raiseshistory_lengthbefore committing blocks. It now useshistory_length=0, the widest possible range, since the prior history length and capacity are not retained.Block.storagereads go throughBlock.get_page(). The slot never holds a dangling ref, becauseCommittedPage.__del__unlinks from its block before invalidating its rawref._resolve_page_refis dropped — its non-callable branch was dead, and a runtimecallable()check has no C++ equivalent.The change is implemented in both the Python and C++ backends so the two stay in agreement. Notably,
eventManager.cpp'spageCoversBlockfiltered viadynamic_cast<SsmCommittedPage const*>; once the field moved to the base class that code would still compile but silently stop filtering attention pages, so deleting the class was load-bearing.Test Coverage
New
TestPartialCoverageReuse(tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py):test_rewind_endpoint_survives_longer_sibling_created_aftertest_rewind_endpoint_attaches_to_longer_sibling_created_beforetest_reused_partial_coverage_kv_is_correct— runsFakeEngine, which validates every reused history token against expected values for both layer groups, so a clean run proves the salvaged endpoint holds real KV rather than uninitialized memorytest_exact_boundary_ignores_stale_last_block_partial_coveragetest_page_coverage_only_growstest_replaced_page_does_not_outlive_its_block— the displaced-page lifetime above, reached through committest_replaced_reused_page_does_not_outlive_its_block— same defect reached through reusePlus
TestSSMSupport::test_ssm_snapshot_moves_to_covering_block,TestNoBatching::test_planned_drop_handle_rejects_partial_coverage, andtest_v2_kv_cache_event_manager_omits_partial_life_cycle_coverage.These go through
_introspectionrather than the Python object graph, so they run against bothbackends. The one exception is
test_planned_drop_handle_rejects_partial_coverage, which has towrite a page field directly and is gated to the Python backend.
The two lifetime tests assert the invariant via a new
_introspection.committed_page_is_linked()hook rather than waiting for a crash. That is deliberate: the underlying fault is a use-after-free
whose overwhelmingly likely symptom is silent — freed-but-mapped memory reads back plausibly, so
the pages merely leak an eviction slot. Both tests were confirmed to fail on both backends with the
back-pointer clear removed, and to pass with it restored. Catching the fault itself still needs a
sanitizer build, which is out of scope here.
Full suite, both backends, no failures:
Skips are pre-existing perf tests plus backend-specific tests; this PR adds exactly one skip (the
white-box test above).
TLLM_DEBUG_MODE=1was also checked. It has pre-existing failures onmainfor the C++ backend,verified by A/B rebuild with the change stashed against the branch base at the time; this PR's
debug-mode failure set was a strict subset of the baseline's and added none.
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.
🤖 Generated with Claude Code