[TRTLLM-15218][chore] KVCacheManagerV2: report the prefix attention alone supports - #17448
Draft
brnguyen2 wants to merge 2 commits into
Draft
[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>
Open
4 tasks
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
Draft. 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.Validation status: the C++ was compiled and the suite run (239 passed) as part of the
combined branch before the split; this branch has not yet been independently rebuilt.
Kept in draft until it has.
PR Checklist
[JIRA/NVBUG/None][type] Summary