From 96fad10732d7d23a1d275bd52d95884ee376a1ad 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 5d0d4bc314c0..af1f0f3e4798 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 @@ -780,7 +780,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 @@ -788,7 +789,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 @@ -876,10 +876,21 @@ std::vector BlockRadixTree::pruneMatch(std::vector< BlockRadixTree::ReuseMatch BlockRadixTree::match( ReuseScope const& reuseScope, std::vector const& tokens, bool enablePartialMatch) const { - auto const matched = pruneMatch(matchTokenPath(reuseScope, tokens, enablePartialMatch)); + auto rawMatched = matchTokenPath(reuseScope, tokens, 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 c4a1debcb093..964dec138350 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 @@ -307,6 +307,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; }; ReuseMatch match( @@ -348,7 +354,10 @@ class BlockRadixTree private: std::vector matchTokenPath( ReuseScope const& reuseScope, std::vector const& tokens, 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 6beb2bd4f9eb..74e428563523 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 @@ -70,6 +70,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 2fca81c8d17e..ae76baba5399 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 @@ -321,6 +321,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; @@ -609,6 +616,7 @@ class KvCache : public std::enable_shared_from_this TypedVec mBlocks; std::vector mCommittedTokens; + 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 62faf997cf1b..def45d345e3b 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1484,6 +1484,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 66e225ee63ba..d01e572e4b67 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 @@ -50,6 +50,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") @@ -560,13 +564,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 @@ -628,13 +637,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 d71bf3db6acb..cebf57a6733e 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", @@ -272,6 +273,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. @@ -326,6 +329,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 @@ -1043,6 +1049,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 48c10a006fc0..c88d7ab64866 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 @@ -2386,6 +2386,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 59f1ccc40a86d0ec9c331c947e355f19749f33fb 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 c88d7ab64866..36deb40aeb79 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 @@ -2386,9 +2386,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() @@ -2397,6 +2394,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)