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..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 @@ -2394,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)