[TRTLLM-15218][chore] KVCacheManagerV2: report the prefix attention alone supports - #17448
Open
brnguyen2 wants to merge 2 commits into
Open
[TRTLLM-15218][chore] KVCacheManagerV2: report the prefix attention alone supports#17448brnguyen2 wants to merge 2 commits into
brnguyen2 wants to merge 2 commits into
Conversation
…KV cache manager V2 When Kimi K3 prefix reuse underperforms, the useful number is the one that separates "attention pages matched N tokens" from "recurrent-state snapshot availability cut it to M". V2 exposed only the final M, so the two causes were indistinguishable from the outside. BlockRadixTree::pruneMatch now takes the SSM life cycle as a parameter; passing std::nullopt skips the recurrent-snapshot constraint and yields the attention-only prefix. match() records that value on ReuseMatch, and KvCache carries it to a _get_num_tokens_before_hybrid_pruning() accessor (C++, nanobind and the Python runtime mirror). Models without an SSM life cycle skip the extra prune pass entirely and report the final match length, so only hybrid models pay for the diagnostic. Diagnostic only, no behavior change, and reachable only under use_kv_cache_manager_v2=True. Test: test_kv_cache_manager_v2.py::test_ssm_reuse_keeps_snapshots_from_multiple_commits asserts the diagnostic reports 48 where the committed reuse is 32, i.e. that recurrent pruning rather than a short attention match caused the truncation. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
… differ The assertion added to test_ssm_reuse_keeps_snapshots_from_multiple_commits could not hold. That test runs without partial reuse, so a match is block-aligned: with tokens_per_block=32 a 48-token lookup matches only the one complete block, the attention-only prefix is 32, and the diagnostic is indistinguishable from num_committed_tokens. Restore that test to its original assertions and cover the diagnostic in a test that configures enable_partial_reuse=True, where attention partially covers 48 tokens while the latest reusable SSM snapshot sits at 32. That is the case the counter exists to explain. The second half asserts the diagnostic collapses onto num_committed_tokens when the snapshot and the attention match agree, so the test fails if it ever reports the lookup length instead. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (8)
WalkthroughHybrid KV-cache matching now records the attention prefix before SSM pruning. C++ and Python cache APIs expose this count, and an SSM partial-reuse test validates the diagnostic values. ChangesHybrid pruning diagnostics
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant BlockRadixTree
participant KvCache
participant KVCacheBinding
BlockRadixTree->>BlockRadixTree: compute attention-only prefix
BlockRadixTree->>BlockRadixTree: apply SSM snapshot pruning
BlockRadixTree-->>KvCache: return ReuseMatch with both counts
KvCache->>KVCacheBinding: expose diagnostic accessor
KVCacheBinding-->>KvCache: return pre-hybrid-pruning count
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
Author
|
/bot run |
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
Split out of #17447 so the two fixes there are not held up by a design discussion about a
diagnostic.
When Kimi K3 prefix reuse underperforms, the useful number is the one that separates
"attention pages matched N tokens" from "recurrent-snapshot availability cut it to M".
KVCacheManagerV2 exposed only the final M, so the two causes were indistinguishable from
the outside.
BlockRadixTree::pruneMatchnow takes the SSM life cycle as a parameter; passingstd::nulloptskips the recurrent-snapshot constraint and yields the attention-onlyprefix.
match()records that onReuseMatchandKvCachecarries it to a_get_num_tokens_before_hybrid_pruning()accessor (C++, nanobind, and the Python runtimemirror). Models without an SSM life cycle skip the extra prune pass entirely and report
the final match length, so only hybrid models pay for the diagnostic.
Diagnostic only, no behavior change, and reachable only under
use_kv_cache_manager_v2=True.Reviewer note — this is a redefinition, not a port. An equivalent counter exists
elsewhere with a different meaning, computed at a snapshot point in a staged prune. This
implementation cannot reproduce that:
pruneMatchhere is a single fixed-point loop withpageCoverage, so that snapshot point has no equivalent. The counter is defined insteadas the prefix the attention pages alone support, which is slightly narrower — it does
not separate out SWA-window pruning. If a different definition is wanted, this is the
commit to say so on.
Test Coverage
test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation: attentionpartially covers a 48-token lookup while the latest reusable SSM snapshot sits at 32, so
the diagnostic reports 48 where
num_committed_tokensis 32. The second half asserts itcollapses onto
num_committed_tokenswhen the snapshot and the attention match agree, sothe test fails if it ever reports the lookup length instead.
Partial reuse must be enabled for the two numbers to differ at all: without it a match is
block-aligned, both are cut at the same block boundary, and the diagnostic is
indistinguishable from
num_committed_tokens.Verified on this branch standalone (not just as part of the combined branch it was split
from): C++ builds clean and
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pyis 113 passed / 13 skipped / 0 failed.
PR Checklist
[JIRA/NVBUG/None][type] SummaryDev Engineer Review
KvCacheand nanobind.QA Engineer Review
TestSSMSupport.test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation.tests/integration/test_lists/entry is reported for this test.