[None][feat] add KV cache reuse probe

[lowsfer]

Summary:

• Move full KV cache reuse matching and pruning into BlockRadixTree.
• Have KVCacheManager create a ReuseMatch before constructing _KVCache.
• Add KVCacheManager.probe_reuse() for advisory prefix-match length without holding pages.
• Update reuse salting and reuse-scope tests.

Tests:

• ruff check _block_radix_tree.py _core/_kv_cache.py _core/_kv_cache_manager.py init.pyi tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
• ruff format --check _block_radix_tree.py _core/_kv_cache.py _core/_kv_cache_manager.py init.pyi tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
• PYTHONPYCACHEPREFIX=/tmp/kvcm_pycache python -m py_compile _block_radix_tree.py _core/_kv_cache.py _core/_kv_cache_manager.py tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
• PYTHONPATH=/home/yaoy/tekit/tensorrt_llm/runtime PYTHONPYCACHEPREFIX=/tmp/kvcm_pycache python tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py -v
• PYTHONPATH=/home/yaoy/tekit/tensorrt_llm/runtime PYTHONPYCACHEPREFIX=/tmp/kvcm_pycache python tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py TestNoBatching.test_reuse_scope_isolates_reuse -v
• CUDA_VISIBLE_DEVICES=4 PYTHONPATH=/home/yaoy/tekit/tensorrt_llm/runtime PYTHONPYCACHEPREFIX=/tmp/kvcm_pycache python tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py -v: passed, 61 tests run, 12 skipped

[coderabbitai]

📝 Walkthrough

Walkthrough

Refactored KV cache reuse matching to return a structured ReuseMatch data object from BlockRadixTree.match() instead of a generator. Updated _KVCache and KVCacheManager to use the new contract, added probe_reuse() API to report reusable prefix length without page acquisition, and updated tests.

Changes

KV Cache Reuse Matching Contract and Integration

Layer / File(s)	Summary
ReuseMatch contract and BlockRadixTree matching refactor tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py, tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi	ReuseMatch NamedTuple packages match results as (blocks, num_tokens). BlockRadixTree.match() refactored to return ReuseMatch instead of a generator. Internal helpers _match_token_path() and _prune_match() apply lifecycle constraints; static helpers _has_pages, _has_page, _num_matched_tokens support matching logic. Public type stubs updated.
_KVCache reuse initialization with ReuseMatch tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py	_KVCache.__init__() now accepts `reuse_match: ReuseMatch
KVCacheManager reuse matching and probe API tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py, tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi	Imports ReuseMatch. Private _match_reuse() helper centralizes radix-tree matching. create_kv_cache() computes optional reuse_match and passes it to _KVCache. Public probe_reuse() method returns advisory prefix length (num_tokens) without acquiring/holding pages.
Test updates for ReuseMatch and probe_reuse APIs tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py, tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py	_EmptyLifeCycles stub implements lifecycle interface with ssm_life_cycle_id property and attention_life_cycles() method. Radix-tree match assertions updated to use structured ReuseMatch fields (match.blocks, match.num_tokens). Test helper validates probe_reuse() behavior.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#13793: Both PRs modify BlockRadixTree matching logic in the KV cache v2 reuse path; #13793 adds cache salt and task id parameters, while this PR refactors the match return contract to ReuseMatch.
• NVIDIA/TensorRT-LLM#14140: Both PRs refactor KVCacheManagerV2 reuse plumbing via ReuseScope and BlockRadixTree.match() signatures; this PR builds on that foundation by changing match() to return ReuseMatch and updating cache initialization accordingly.

Suggested reviewers

• joyang-nv
• eopXD
• yizhang-nv

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.71% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The PR description lacks several required template sections including PR title format, detailed explanation of the issue/solution, and explicit test coverage listing.	Add a properly formatted title following [JIRA/GitHub/None][type] format, provide a detailed description section explaining the issue and solution, and list relevant test coverage in a dedicated section.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly indicates the main feature addition: a KV cache reuse probe capability for the KVCacheManager.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)

517-524: QA list updates look unnecessary here.

This PR still looks unit-scope only, so I don’t see a new entry needed under tests/integration/test_lists/qa/ or test-db for this change. The new assertion covers the happy path for probe_reuse(), but I’d still want direct unit regressions for the two pruning edge cases in _block_radix_tree.py before relying on this refactor.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py`:
- Around line 510-517: The current SSM truncation can keep a partially matched
block that has an SSM snapshot, causing _KVCache._setup_for_reuse() to seed
_ssm_blocks from a future state; instead, when scanning matched for the last
block with a snapshot (ssm_lc_id), require that the block is fully matched
before accepting it: in the reversed loop over matched, only set ssm_trunc = i +
1 if matched[i][0].storage[ssm_lc_id] is not None AND the matched entry
indicates a full match (e.g., matched[i][1] equals the block's full token count
such as matched[i][0].num_tokens or matched[i][0].length); otherwise skip that
index so any partially matched tail is dropped, then slice matched =
matched[:ssm_trunc] as before.
- Around line 538-540: The current slice keeps the first missing SWA window
block when n==0, causing probe_reuse() to over-report and later
_setup_for_reuse() to hit unwrap_optional(None); adjust the truncation to drop
that missing block by slicing matched to exclude the found index itself (use
matched[: len(matched) - n - 1] instead of matched[: len(matched) - n]) after
computing n via find_index(reversed(matched[stale_end:]), check_no_page_stale),
ensuring the newest missing block is removed from the reusable prefix so
probe_reuse() and _setup_for_reuse() no longer see an invalid block.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 094c3dea-a403-4c95-ba3e-b900bf4da3a3

📥 Commits

Reviewing files that changed from the base of the PR and between fb06a2f and b26561f.

📒 Files selected for processing (6)

• tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi
• tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
• tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
• tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py

[lowsfer]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49350 [ run ] triggered by Bot. Commit: f9968c6 Link to invocation

[lancelly]

LGTM～ Thanks!

[tensorrt-cicd]

PR_Github #49350 [ run ] completed with state SUCCESS. Commit: f9968c6
/LLM/main/L0_MergeRequest_PR pipeline #39006 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

[lancelly]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49628 [ run ] triggered by Bot. Commit: f9968c6 Link to invocation

[tensorrt-cicd]

PR_Github #49628 [ run ] completed with state SUCCESS. Commit: f9968c6
/LLM/main/L0_MergeRequest_PR pipeline #39247 completed with status: 'SUCCESS'

CI Report

Link to invocation