[TRTLLM-15218][chore] KVCacheManagerV2: report the prefix attention alone supports by brnguyen2 · Pull Request #17448 · NVIDIA/TensorRT-LLM

brnguyen2 · 2026-08-09T17:26:09Z

Description

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::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 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.

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: pruneMatch here is a single fixed-point loop with
pageCoverage, so that snapshot point has no equivalent. The counter is defined instead
as 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: attention
partially covers a 48-token lookup while the latest reusable SSM snapshot sits at 32, so
the diagnostic reports 48 where num_committed_tokens is 32. The second half asserts it
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.

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.

Verified on this branch standalone (not just as part of the combined branch it was split
from): C++ builds clean and tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
is 113 passed / 13 skipped / 0 failed.

PR Checklist

PR title is [JIRA/NVBUG/None][type] Summary
Commits are signed off (DCO)
New test added and passing
No new dependencies

Dev Engineer Review

Adds diagnostic plumbing for the attention-only prefix length before hybrid SSM pruning.
Updates the C++ implementation and Python runtime mirror consistently.
Exposes the value through KvCache and nanobind.
Preserves behavior for non-hybrid models.
Reports the final match length when no SSM lifecycle exists.
Adds no configuration changes.
No test-list files were modified.
The change is diagnostic-only. It adds an additional diagnostic match.
Standalone C++ validation reportedly passed 113 tests, skipped 13 tests, and failed 0 tests.

QA Engineer Review

Added TestSSMSupport.test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation.
The test covers partial reuse where attention supports 48 tokens and recurrent snapshots limit committed reuse to 32 tokens.
The test also covers matching attention and recurrent limits at 64 tokens.
No corresponding tests/integration/test_lists/ entry is reported for this test.
Verdict: needs follow-up.

coderabbitai · 2026-08-10T02:54:46Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 68725e7c-d255-41d4-b494-68ce08859fdc

📥 Commits

Reviewing files that changed from the base of the PR and between 5e09668 and 2e63075.

📒 Files selected for processing (8)

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h
cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp
tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

🚧 Files skipped from review as they are similar to previous changes (8)

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp
tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py

Included review availability: Your plan includes up to 12 reviews per rolling hour; 8 remain after this review.

Walkthrough

Hybrid KV-cache matching now records the attention prefix before SSM pruning. C++ and Python cache APIs expose this count, and an SSM partial-reuse test validates the diagnostic values.

Changes

Hybrid pruning diagnostics

Layer / File(s)	Summary
Separate attention and SSM pruning `cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.*`, `tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py`	Matching computes the attention-only prefix before applying SSM snapshot constraints and stores both token counts in `ReuseMatch`.
Propagate the diagnostic count `cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.*`, `cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp`, `tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`	KV-cache objects initialize and expose the pre-hybrid-pruning token count through C++ and Python bindings.
Validate partial SSM reuse `tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`	The regression test checks 32 committed tokens versus 48 pre-hybrid-pruning tokens for partial reuse, and 64 for an exact snapshot match.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 2e630

This change adds a diagnostic for prefix reuse without changing runtime behavior; no actionable merge-blocking risk remains beyond normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant BlockRadixTree
  participant KvCache
  participant KVCacheBinding
  BlockRadixTree->>BlockRadixTree: compute attention-only prefix
  BlockRadixTree->>BlockRadixTree: apply SSM snapshot pruning
  BlockRadixTree-->>KvCache: return ReuseMatch with both counts
  KvCache->>KVCacheBinding: expose diagnostic accessor
  KVCacheBinding-->>KvCache: return pre-hybrid-pruning count

Suggested reviewers: yizhang-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 54.55% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the diagnostic change for KVCacheManagerV2 and follows the required ticket and type format.
Description check	✅ Passed	The description explains the problem, solution, diagnostic scope, test coverage, results, and checklist items with sufficient detail.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

brnguyen2 · 2026-08-10T03:05:41Z

/bot run

tensorrt-cicd · 2026-08-10T03:11:48Z

PR_Github #64935 [ run ] triggered by Bot. Commit: 59f1ccc Link to invocation

tensorrt-cicd · 2026-08-10T05:01:04Z

PR_Github #64935 [ run ] completed with state SUCCESS. Commit: 59f1ccc
/LLM/main/L0_MergeRequest_PR pipeline #52770 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-10T05:12:28Z

/bot run

tensorrt-cicd · 2026-08-10T05:18:41Z

PR_Github #64956 [ run ] triggered by Bot. Commit: 59f1ccc Link to invocation

tensorrt-cicd · 2026-08-10T05:32:12Z

PR_Github #64956 [ run ] completed with state FAILURE. Commit: 59f1ccc
/LLM/main/L0_MergeRequest_PR pipeline #52790 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-10T05:32:50Z

/bot run

tensorrt-cicd · 2026-08-10T05:38:45Z

PR_Github #64961 [ run ] triggered by Bot. Commit: 59f1ccc Link to invocation

tensorrt-cicd · 2026-08-10T06:47:48Z

PR_Github #64961 [ run ] completed with state SUCCESS. Commit: 59f1ccc
/LLM/main/L0_MergeRequest_PR pipeline #52795 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-10T23:24:35Z

/bot run

tensorrt-cicd · 2026-08-10T23:29:59Z

PR_Github #65149 [ run ] triggered by Bot. Commit: 59f1ccc Link to invocation

tensorrt-cicd · 2026-08-11T01:38:40Z

PR_Github #65149 [ run ] completed with state FAILURE. Commit: 59f1ccc
/LLM/main/L0_MergeRequest_PR pipeline #52943 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-11T01:52:29Z

/bot run

tensorrt-cicd · 2026-08-11T01:58:52Z

PR_Github #65164 [ run ] triggered by Bot. Commit: 59f1ccc Link to invocation

tensorrt-cicd · 2026-08-11T06:13:28Z

PR_Github #65164 [ run ] completed with state FAILURE. Commit: 59f1ccc
/LLM/main/L0_MergeRequest_PR pipeline #52958 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

brnguyen2 · 2026-08-12T15:09:38Z

/bot run

brnguyen2 · 2026-08-12T17:16:35Z

/bot run

brnguyen2 · 2026-08-12T19:23:52Z

/bot run

brnguyen2 · 2026-08-12T21:35:18Z

/bot run

brnguyen2 · 2026-08-12T23:36:13Z

/bot run

brnguyen2 · 2026-08-13T01:42:08Z

/bot run

brnguyen2 · 2026-08-13T03:50:29Z

/bot run

brnguyen2 · 2026-08-15T05:34:08Z

/bot run

brnguyen2 · 2026-08-15T07:37:59Z

/bot run

brnguyen2 · 2026-08-15T09:40:55Z

/bot run

brnguyen2 · 2026-08-15T19:19:50Z

/bot run

brnguyen2 · 2026-08-15T21:23:34Z

/bot run

brnguyen2 · 2026-08-15T23:27:12Z

/bot run

brnguyen2 · 2026-08-16T01:30:45Z

/bot run

brnguyen2 · 2026-08-16T03:35:29Z

/bot run

brnguyen2 · 2026-08-16T05:39:05Z

/bot run

brnguyen2 · 2026-08-16T07:42:44Z

/bot run

brnguyen2 · 2026-08-16T09:46:10Z

/bot run

brnguyen2 · 2026-08-16T19:31:30Z

/bot run

brnguyen2 · 2026-08-16T21:35:43Z

/bot run

brnguyen2 · 2026-08-16T23:37:17Z

/bot run

brnguyen2 · 2026-08-17T01:43:15Z

/bot run

…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>

brnguyen2 · 2026-08-17T03:43:19Z

/bot run

coderabbitai · 2026-08-17T04:49:50Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

brnguyen2 · 2026-08-17T04:50:17Z

/bot run

tensorrt-cicd · 2026-08-17T04:56:19Z

PR_Github #66635 [ run ] triggered by Bot. Commit: 2e63075 Link to invocation

tensorrt-cicd · 2026-08-17T05:09:43Z

PR_Github #66635 [ run ] completed with state FAILURE. Commit: 2e63075
/LLM/main/L0_MergeRequest_PR pipeline #54255 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-17T05:49:55Z

/bot run

tensorrt-cicd · 2026-08-17T05:56:09Z

PR_Github #66658 [ run ] triggered by Bot. Commit: 2e63075 Link to invocation

tensorrt-cicd · 2026-08-17T08:26:01Z

PR_Github #66658 [ run ] completed with state SUCCESS. Commit: 2e63075
/LLM/main/L0_MergeRequest_PR pipeline #54275 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

brnguyen2 · 2026-08-17T08:35:32Z

/bot run

tensorrt-cicd · 2026-08-17T08:41:03Z

PR_Github #66706 [ run ] triggered by Bot. Commit: 2e63075 Link to invocation

tensorrt-cicd · 2026-08-17T21:04:46Z

PR_Github #66706 [ run ] completed with state SUCCESS. Commit: 2e63075
/LLM/main/L0_MergeRequest_PR pipeline #54317 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

github-actions Bot assigned brnguyen2 Aug 9, 2026

brnguyen2 mentioned this pull request Aug 9, 2026

[TRTLLM-15216][fix] Kimi K3 on KVCacheManagerV2: conv-state layout and SSM iteration stats #17447

Merged

4 tasks

brnguyen2 marked this pull request as ready for review August 10, 2026 02:47

brnguyen2 requested a review from a team as a code owner August 10, 2026 02:47

brnguyen2 requested review from liji-nv and yizhang-nv August 10, 2026 02:47

brnguyen2 added the ci: full pre-merge approved label Aug 12, 2026

brnguyen2 added 2 commits August 16, 2026 20:15

brnguyen2 force-pushed the k3/kvcm-v2-pruning-diagnostic branch from 59f1ccc to 2e63075 Compare August 17, 2026 04:49

brnguyen2 merged commit c763b04 into NVIDIA:main Aug 17, 2026
10 checks passed

Uh oh!

Conversation

brnguyen2 commented Aug 9, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

Dev Engineer Review

QA Engineer Review

Uh oh!

coderabbitai Bot commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

brnguyen2 commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

brnguyen2 commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

brnguyen2 commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

brnguyen2 commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

brnguyen2 commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

brnguyen2 commented Aug 12, 2026

Uh oh!

brnguyen2 commented Aug 12, 2026

Uh oh!

brnguyen2 commented Aug 12, 2026

Uh oh!

brnguyen2 commented Aug 12, 2026

Uh oh!

brnguyen2 commented Aug 12, 2026

Uh oh!

brnguyen2 commented Aug 13, 2026

Uh oh!

brnguyen2 commented Aug 13, 2026

Uh oh!

brnguyen2 commented Aug 15, 2026

Uh oh!

brnguyen2 commented Aug 15, 2026

Uh oh!

brnguyen2 commented Aug 15, 2026

Uh oh!

brnguyen2 commented Aug 15, 2026

Uh oh!

brnguyen2 commented Aug 15, 2026

Uh oh!

brnguyen2 commented Aug 15, 2026

Uh oh!

brnguyen2 commented Aug 16, 2026

Uh oh!

brnguyen2 commented Aug 16, 2026

Uh oh!

brnguyen2 commented Aug 16, 2026

Uh oh!

brnguyen2 commented Aug 16, 2026

brnguyen2 commented Aug 9, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 10, 2026 •

edited

Loading