[TRTLLM-15218][chore] KVCacheManagerV2: report the prefix attention alone supports

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp

@@ -780,15 +780,15 @@ std::vector<BlockRadixTree::MatchResult> BlockRadixTree::matchTokenPath(
     return results;
 }
 
-std::vector<BlockRadixTree::MatchResult> BlockRadixTree::pruneMatch(std::vector<MatchResult> matched) const
+std::vector<BlockRadixTree::MatchResult> BlockRadixTree::pruneMatch(
+    std::vector<MatchResult> matched, std::optional<LifeCycleId> ssmLcId) const
 {
     // All blocks except the last must be fully matched (mirrors Python: matched[:-1]).
     TLLM_CHECK_DEBUG(matched.size() <= 1
         || std::all_of(matched.begin(), matched.end() - 1,
             [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::MatchResult> BlockRadixTree::pruneMatch(std::vector<
 BlockRadixTree::ReuseMatch BlockRadixTree::match(
     ReuseScope const& reuseScope, std::vector<TokenIdExt> 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<int> 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<int>(tokens.size());
+    result.numTokensBeforeHybridPruning = attnOnlyTokens.value_or(result.numTokens);
     result.blocks.reserve(BlockOrdinal{static_cast<int>(matched.size())});
     for (auto const& match : matched)
     {

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<MatchResult> matchTokenPath(
         ReuseScope const& reuseScope, std::vector<TokenIdExt> const& tokens, bool enablePartialMatch) const;
-    std::vector<MatchResult> pruneMatch(std::vector<MatchResult> 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<MatchResult> pruneMatch(std::vector<MatchResult> matched, std::optional<LifeCycleId> ssmLcId) const;
 
     // Erase any pending empty root blocks from mRoots.
     // Const-qualified: deferred cleanup is not a logical mutation.

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp

@@ -70,6 +70,7 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional<B
     , mHistoryLength(0)
     , mExpectedPromptLength(
           expectedPromptLength.has_value() ? std::optional<int>{std::max(*expectedPromptLength, 0)} : std::nullopt)
+    , mNumTokensBeforeHybridPruning(reuseMatch.has_value() ? reuseMatch->numTokensBeforeHybridPruning : 0)
     , mNumCommittedBlocks(0)
     , mTokensPerBlock(manager.tokensPerBlock())
 {

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h

@@ -321,6 +321,13 @@ class KvCache : public std::enable_shared_from_this<KvCache>
         return static_cast<int>(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<TokenIdExt> const& committedTokens() const noexcept
     {
         return mCommittedTokens;
@@ -609,6 +616,7 @@ class KvCache : public std::enable_shared_from_this<KvCache>
     TypedVec<BlockOrdinal, SeqBlock> mBlocks;
 
     std::vector<TokenIdExt> mCommittedTokens;
+    int mNumTokensBeforeHybridPruning;
     int mNumCommittedBlocks;
     std::optional<CachedCudaEvent> mFinishEvent;
     int mTokensPerBlock;

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); })

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:

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)

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)