From fc80deeff6f5fa812395da79d50d20093d23ce16 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 07:04:04 -0700 Subject: [PATCH 1/2] [TRTLLM-15218][chore] plumb num_tokens_before_hybrid_pruning through 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 --- .../kv_cache_manager_v2/blockRadixTree.cpp | 17 ++++++++-- .../kv_cache_manager_v2/blockRadixTree.h | 11 ++++++- .../kv_cache_manager_v2/kvCache.cpp | 1 + .../kv_cache_manager_v2/kvCache.h | 8 +++++ .../batch_manager/kvCacheManagerV2.cpp | 1 + .../kv_cache_manager_v2/_block_radix_tree.py | 33 +++++++++++++++---- .../kv_cache_manager_v2/_core/_kv_cache.py | 10 ++++++ .../test_kv_cache_manager_v2.py | 3 ++ 8 files changed, 73 insertions(+), 11 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index a6db17976f30..81ae502ed754 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -767,7 +767,8 @@ std::vector BlockRadixTree::matchTokenPath( return results; } -std::vector BlockRadixTree::pruneMatch(std::vector matched) const +std::vector BlockRadixTree::pruneMatch( + std::vector matched, std::optional ssmLcId) const { // All blocks except the last must be fully matched (mirrors Python: matched[:-1]). TLLM_CHECK_DEBUG(matched.size() <= 1 @@ -775,7 +776,6 @@ std::vector BlockRadixTree::pruneMatch(std::vector< [this](auto const& m) { return m.numMatchedTokens == mTokensPerBlock; })); auto attnLcs = mLifeCycles.attentionLifeCycles(); - auto ssmLcId = mLifeCycles.ssmLifeCycleId(); // Fixed-point loop: SSM may select an earlier exact snapshot, while attention may // shorten the match to the coverage of a required page. Every retry strictly @@ -863,10 +863,21 @@ std::vector BlockRadixTree::pruneMatch(std::vector< BlockRadixTree::ReuseMatch BlockRadixTree::match( ReuseScope const& reuseScope, TokenSpan tokens, bool knownNoDigest, bool enablePartialMatch) const { - auto const matched = pruneMatch(matchTokenPath(reuseScope, tokens, knownNoDigest, enablePartialMatch)); + auto rawMatched = matchTokenPath(reuseScope, tokens, knownNoDigest, enablePartialMatch); + auto const ssmLcId = mLifeCycles.ssmLifeCycleId(); + // Diagnostic only: re-prune ignoring recurrent-snapshot availability to get + // the prefix the attention pages alone support. Only hybrid models pay for + // the second pass; without an SSM life cycle the two results are identical. + std::optional attnOnlyTokens; + if (ssmLcId.has_value()) + { + attnOnlyTokens = numMatchedTokens(pruneMatch(rawMatched, std::nullopt), mTokensPerBlock); + } + auto const matched = pruneMatch(std::move(rawMatched), ssmLcId); ReuseMatch result{}; result.numTokens = numMatchedTokens(matched, mTokensPerBlock); result.numLookupTokens = static_cast(tokens.size()); + result.numTokensBeforeHybridPruning = attnOnlyTokens.value_or(result.numTokens); result.blocks.reserve(BlockOrdinal{static_cast(matched.size())}); for (auto const& match : matched) { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h index aff76497dc40..0120d411ee36 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h @@ -368,6 +368,12 @@ class BlockRadixTree int numTokens; // Total query length passed to match() (== len(tokens)). int numLookupTokens; + // Internal diagnostic: the prefix the attention pages alone would + // support, i.e. before recurrent-state (SSM) snapshot availability + // shortens it. Equal to numTokens when the model has no SSM life + // cycle. Separates "attention prefix matched N tokens" from + // "recurrent-snapshot pruning cut it to M". + int numTokensBeforeHybridPruning; }; // knownNoDigest: from external text_only knowledge, never a scan (see Hasher::update). @@ -413,7 +419,10 @@ class BlockRadixTree // knownNoDigest: from external text_only knowledge, never a scan (see Hasher::update). std::vector matchTokenPath( ReuseScope const& reuseScope, TokenSpan tokens, bool knownNoDigest, bool enablePartialMatch) const; - std::vector pruneMatch(std::vector matched) const; + // Shorten `matched` to the prefix that is actually reusable. Passing + // std::nullopt for `ssmLcId` skips the recurrent-snapshot constraint and + // yields the attention-only prefix (used for numTokensBeforeHybridPruning). + std::vector pruneMatch(std::vector matched, std::optional ssmLcId) const; // Erase any pending empty root blocks from mRoots. // Const-qualified: deferred cleanup is not a logical mutation. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index a1595db74c65..6709af5a5ce6 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -71,6 +71,7 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional{std::max(*expectedPromptLength, 0)} : std::nullopt) + , mNumTokensBeforeHybridPruning(reuseMatch.has_value() ? reuseMatch->numTokensBeforeHybridPruning : 0) , mNumCommittedBlocks(0) , mTokensPerBlock(manager.tokensPerBlock()) { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index afb6e3cce152..8688223b6e12 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -281,6 +281,13 @@ class KvCache : public std::enable_shared_from_this return static_cast(mCommittedTokens.size()); } + // Internal diagnostic: prefix supported by the attention pages alone, + // before recurrent-state (SSM) snapshot pruning shortened the reuse. + int numTokensBeforeHybridPruning() const noexcept + { + return mNumTokensBeforeHybridPruning; + } + std::vector const& committedTokens() const noexcept { return mCommittedTokens; @@ -581,6 +588,7 @@ class KvCache : public std::enable_shared_from_this std::vector mCommittedTokens; // Resolved per-sequence text-only state after applying the manager default. bool mTextOnly = false; + int mNumTokensBeforeHybridPruning; int mNumCommittedBlocks; std::optional mFinishEvent; int mTokensPerBlock; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index dcebcc4e6446..59d5c6cbf19f 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1618,6 +1618,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_prop_ro("num_blocks", [](kv::KvCache const& self) { return self.numBlocks().value(); }) .def_prop_ro("num_committed_blocks", &kv::KvCache::numCommittedBlocks) .def_prop_ro("num_committed_tokens", &kv::KvCache::numCommittedTokens) + .def("_get_num_tokens_before_hybrid_pruning", &kv::KvCache::numTokensBeforeHybridPruning) .def_prop_rw("history_length", &kv::KvCache::historyLength, [](kv::KvCache& self, int hist) { self.setHistoryLength(hist); }) .def_prop_rw("capacity", &kv::KvCache::capacity, [](kv::KvCache& self, int cap) { self.setCapacity(cap); }) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index 0adbb847e783..0c78ac996298 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -202,6 +202,10 @@ class ReuseMatch(NamedTuple): blocks: list["Block"] num_tokens: int num_lookup_tokens: int + # Internal diagnostic: the prefix the attention pages alone would support, + # i.e. before recurrent-state (SSM) snapshot availability shortened it. + # Equal to num_tokens when the model has no SSM life cycle. + num_tokens_before_hybrid_pruning: int = 0 Child = TypeVar("Child", bound="Block | RootBlock") @@ -705,13 +709,18 @@ def _match_token_path( block = partial_block yield block, match_len - def _prune_match(self, matched: list[tuple[Block, int]]) -> list[tuple[Block, int]]: + def _prune_match( + self, matched: list[tuple[Block, int]], ssm_lc_id: LifeCycleId | None + ) -> list[tuple[Block, int]]: + """Shorten `matched` to the prefix that is actually reusable. + + Passing ssm_lc_id=None skips the recurrent-snapshot constraint and yields + the attention-only prefix (used for num_tokens_before_hybrid_pruning). + """ tokens_per_block = self._tokens_per_block assert all(b[1] == tokens_per_block for b in matched[:-1]) - life_cycles = self._life_cycles - attn_life_cycles = list(life_cycles.attention_life_cycles()) - ssm_lc_id = life_cycles.ssm_life_cycle_id + attn_life_cycles = list(self._life_cycles.attention_life_cycles()) # Fixed-point loop: SSM may select an earlier exact snapshot, while attention may # shorten the match to the coverage of a required page. Every retry strictly @@ -773,13 +782,23 @@ def match( The result is volatile: callers that need to reuse the returned blocks must acquire ownership of the pages before depending on them. """ - matched = self._prune_match( - list(self._match_token_path(reuse_scope, tokens, enable_partial_match)) + raw_matched = list(self._match_token_path(reuse_scope, tokens, enable_partial_match)) + ssm_lc_id = self._life_cycles.ssm_life_cycle_id + # Diagnostic only: re-prune ignoring recurrent-snapshot availability to get + # the prefix the attention pages alone support. Only hybrid models pay for + # the second pass; without an SSM life cycle the two results are identical. + attn_only_tokens = ( + self._num_matched_tokens(self._prune_match(list(raw_matched), None)) + if ssm_lc_id is not None + else None ) + matched = self._prune_match(raw_matched, ssm_lc_id) + num_tokens = self._num_matched_tokens(matched) return ReuseMatch( [block for block, _ in matched], - self._num_matched_tokens(matched), + num_tokens, len(tokens), + num_tokens if attn_only_tokens is None else attn_only_tokens, ) def _check_sanity(self) -> bool: 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 a3e3e2779cf2..db31d740c9e0 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 @@ -238,6 +238,7 @@ class _KVCache: "_blocks", "_base_page_indices", "_committed_tokens", + "_num_tokens_before_hybrid_pruning", "_num_committed_blocks", "_finish_event", "_tokens_per_block", @@ -273,6 +274,8 @@ class _KVCache: # be computed on the fly, but that would be slow due to python. _base_page_indices: TypedIndexList[BeamIndex, TypedIndexList[LifeCycleId, IndexSeq]] _committed_tokens: list[TokenIdExt] + # Internal diagnostic captured from the reuse match: see ReuseMatch. + _num_tokens_before_hybrid_pruning: int # Sometimes we can't commit a block because all its tokens are already covered by another block in # the radix tree. But it's unsafe to just use the other block because: 1. the data may have numeric # difference, 2. if our block is a partial block, we can't write to memory of the other blocks. @@ -330,6 +333,9 @@ def __init__( self.beam_width, ) self._committed_tokens = [] + self._num_tokens_before_hybrid_pruning = ( + reuse_match.num_tokens_before_hybrid_pruning if reuse_match is not None else 0 + ) self._num_committed_blocks = BlockOrdinal(0) self._finish_event = None self._tokens_per_block = manager.tokens_per_block @@ -1072,6 +1078,10 @@ def commit( def num_committed_tokens(self) -> int: return len(self._committed_tokens) + def _get_num_tokens_before_hybrid_pruning(self) -> int: + """Return the pre-hybrid-pruning prefix for internal diagnostics.""" + return self._num_tokens_before_hybrid_pruning + @property def committed_tokens(self) -> list[TokenIdExt]: return list(self._committed_tokens) 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 67868935fd83..09823ac0ad2a 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 @@ -2483,6 +2483,9 @@ def test_ssm_reuse_keeps_snapshots_from_multiple_commits(self) -> None: kv3 = self.manager.create_kv_cache(input_tokens=prompt[:48]) self.assertEqual(kv3.num_committed_tokens, 32) + # Attention pages cover all 48 tokens; the latest reusable SSM snapshot + # is at 32, so recurrent pruning is what cut the reuse (TRTLLM-15218). + self.assertEqual(kv3._get_num_tokens_before_hybrid_pruning(), 48) kv3.resume(stream) kv3.close() From 2e63075e6ebca1c07217ff25830ce90228a96984 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 9 Aug 2026 09:08:33 -0700 Subject: [PATCH 2/2] [TRTLLM-15218][test] cover the hybrid-pruning diagnostic where it can 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 --- .../test_kv_cache_manager_v2.py | 42 +++++++++++++++++-- 1 file changed, 39 insertions(+), 3 deletions(-) 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 09823ac0ad2a..bcc02b3feba7 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 @@ -2483,9 +2483,6 @@ def test_ssm_reuse_keeps_snapshots_from_multiple_commits(self) -> None: kv3 = self.manager.create_kv_cache(input_tokens=prompt[:48]) self.assertEqual(kv3.num_committed_tokens, 32) - # Attention pages cover all 48 tokens; the latest reusable SSM snapshot - # is at 32, so recurrent pruning is what cut the reuse (TRTLLM-15218). - self.assertEqual(kv3._get_num_tokens_before_hybrid_pruning(), 48) kv3.resume(stream) kv3.close() @@ -2494,6 +2491,45 @@ def test_ssm_reuse_keeps_snapshots_from_multiple_commits(self) -> None: kv4.resume(stream) kv4.close() + def test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation(self) -> None: + """The diagnostic separates a short attention match from recurrent pruning. + + Partial reuse is required for the two numbers to differ at all: without + it a match is block-aligned, so the attention-only prefix and the final + committed prefix are cut at the same block boundary and the diagnostic + is indistinguishable from num_committed_tokens. + """ + cfg = self._make_ssm_config(tokens_per_block=32, enable_partial_reuse=True) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + + prompt = [self.next_token() for _ in range(96)] + kv1 = self.manager.create_kv_cache() + kv1.resume(stream) + kv1.capacity = 32 + kv1.commit(prompt[:32]) + kv1.capacity = 64 + kv1.commit(prompt[32:64]) + kv1.close() + + # Attention pages partially cover all 48 lookup tokens, but the latest + # reusable SSM snapshot sits at 32 — so recurrent pruning, not a short + # attention match, is what cut the reuse. + kv = self.manager.create_kv_cache(input_tokens=prompt[:48]) + self.assertEqual(kv.num_committed_tokens, 32) + self.assertEqual(kv._get_num_tokens_before_hybrid_pruning(), 48) + kv.resume(stream) + kv.close() + + # When the snapshot and the attention match agree, the diagnostic must + # collapse onto num_committed_tokens rather than reporting the lookup. + kv = self.manager.create_kv_cache(input_tokens=prompt[:64]) + self.assertEqual(kv.num_committed_tokens, 64) + self.assertEqual(kv._get_num_tokens_before_hybrid_pruning(), 64) + kv.resume(stream) + kv.close() + def test_ssm_planned_drop_targets_latest_snapshot_with_shared_plans(self) -> None: """Shared plans drop only their conversation endpoint snapshot.""" cfg = self._make_ssm_config(tokens_per_block=32)