[TRTLLM-15216][fix] Kimi K3 on KVCacheManagerV2: conv-state layout and SSM iteration stats - #17447
Open
brnguyen2 wants to merge 6 commits into
Open
[TRTLLM-15216][fix] Kimi K3 on KVCacheManagerV2: conv-state layout and SSM iteration stats#17447brnguyen2 wants to merge 6 commits into
brnguyen2 wants to merge 6 commits into
Conversation
…nager V2 The kimi_linear branch of _create_kv_cache_manager passed model_type="qwen3_next" unconditionally. MambaHybridCacheManagerV2 has no model_type parameter: it absorbs it into **kwargs and selects the KDA convolution-state sectioning from conv_state_layout, which defaults to "x_b_c". An explicit use_kv_cache_manager_v2=True opt-in therefore built the KDA conv state with the wrong section layout, silently, with no error and no warning. Select the kwarg from the manager class, matching what the qwen3_hybrid branch in the same function already does: conv_state_layout="q_k_v" for MambaHybridCacheManagerV2, model_type="qwen3_next" for the V1 managers. This path is opt-in only today (use_kv_cache_manager_v2 must be set to True; "auto" does not select V2 for kimi_linear), so no default configuration changes behavior. Test: test_kimi_explicit_v2_manager_uses_qkv_convolution_layout already existed as a strict xfail describing this bug; the marker is dropped. Adds test_kimi_v1_manager_still_selects_qwen3_next_model_type to guard the V1 branch. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…teration stats The V2 page-movement recorders (_recordDirectIterationStats, _recordMigratedSlots, _recordDroppedPages, and their Python mirrors) returned early or skipped the page whenever the life cycle was not an AttnLifeCycle. Offload, onboard, intra-device copy and host-tier drop of recurrent (SSM/KDA) state were therefore never recorded, and any recurrent-cache iteration statistics read back as zeros. Iteration statistics are already keyed by life cycle, so recurrent movement stays distinguishable from attention movement without the filter. Keep the filter only where it is semantically required: the global cache-hit counters (allocTotalBlocks / allocNewBlocks) and the block-reuse hit/miss range accounting stay attention-only. Observability only, no change to allocation or reuse behavior, and reachable only under use_kv_cache_manager_v2=True. Test: tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py drives the Python recorders directly (no GPU needed) and asserts offload and host-drop are reported for both life-cycle kinds while the global cache-hit counters remain attention-only. The C++ recorders have no equivalent CPU-only seam. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
brnguyen2
marked this pull request as ready for review
August 9, 2026 17:22
brnguyen2
requested review from
liji-nv,
mikeiovine,
yizhang-nv and
zhaoyangwang-nvidia
August 9, 2026 17:22
4 tasks
brnguyen2
force-pushed
the
k3/kvcm-v2-fixes
branch
from
August 9, 2026 17:28
6db9d2e to
3f6c398
Compare
brnguyen2
changed the title
Contributor
WalkthroughThe change initializes hybrid-pruning state, records KV-cache statistics for attention and SSM lifecycles, and centralizes manager-specific layout selection for hybrid cache managers. ChangesHybrid KV-cache management
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp (1)
617-629: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winRecord SSM deferred-copy iteration stats in Python
Move
_record_direct_iteration_stats()outside theif lc_idx != ssm_lc_id:block intensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py. Keep onlyrecord_allocation_range()inside the block so Python reports SSM copy blocks and bytes like C++.🤖 Prompt for 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. In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp` around lines 617 - 629, The C++ method _recordDirectIterationStats records statistics for every lifecycle, including SSM, while the Python flow currently skips SSM stats. In tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py lines 505-513, move _record_direct_stats() outside the if lc_idx != ssm_lc_id block, leaving only record_allocation_range() inside it so SSM copy blocks and bytes are reported consistently.
🧹 Nitpick comments (1)
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py (1)
42-125: 🗄️ Data Integrity & Integration | 🔵 Trivial | 🏗️ Heavy liftAdd coverage for
KvCache.resume()deferred-copy path
- Added tests:
test_offload_is_recorded_for_every_life_cycle,test_host_drop_is_recorded_for_every_life_cycle, andtest_onboard_counts_globally_only_for_attention. No modified or removed tests.- The file is registered through
unittest/kv_cache_manager_v2_testsinl0_b200.ymlandl0_h100.yml.- Coverage is insufficient. The tests call recorder methods directly and do not exercise
KvCache.resume()for SSM deferred copies. Add this case and assert the per-life-cycleiter_intra_device_copy_*statistics.🤖 Prompt for 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. In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py` around lines 42 - 125, Add a test that exercises KvCache.resume() with an SSM deferred-copy scenario rather than invoking recorder methods directly. Configure the deferred copy and resume flow using the existing test fixtures/helpers, then assert the committed SSM statistics include the expected iter_intra_device_copy_blocks and iter_intra_device_copy_bytes values.Source: Path instructions
🤖 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 `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 2397-2434: Extend
test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation to capture
the relevant iter_intra_device_copy_* counter before and after kv.resume(stream)
for the reused SSM snapshot, then assert the expected delta from the deferred
copy. Ensure the assertion specifically covers the SSM resume path rather than
relying only on token-count diagnostics.
---
Outside diff comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp`:
- Around line 617-629: The C++ method _recordDirectIterationStats records
statistics for every lifecycle, including SSM, while the Python flow currently
skips SSM stats. In tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
lines 505-513, move _record_direct_stats() outside the if lc_idx != ssm_lc_id
block, leaving only record_allocation_range() inside it so SSM copy blocks and
bytes are reported consistently.
---
Nitpick comments:
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py`:
- Around line 42-125: Add a test that exercises KvCache.resume() with an SSM
deferred-copy scenario rather than invoking recorder methods directly. Configure
the deferred copy and resume flow using the existing test fixtures/helpers, then
assert the committed SSM statistics include the expected
iter_intra_device_copy_blocks and iter_intra_device_copy_bytes values.
🪄 Autofix
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: 128cc974-fab3-4f1b-8330-d1ed9fa1fe99
📒 Files selected for processing (11)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.hcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpptensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytests/unittest/_torch/executor/test_mamba_cache_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64869 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64869 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64877 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64877 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64887 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64887 [ run ] completed with state
|
zhaoyangwang-nvidia
requested changes
Aug 10, 2026
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…backends
Review follow-up. The Python backend used to record
iter_intra_device_copy_{blocks,bytes} only for non-SSM life cycles in the
resume() deferred-copy loop, unlike the default C++ backend; the previous
commit aligned the call sites. Add the coverage that was missing:
- test_kv_cache_stats_life_cycles.py now drives
_record_direct_iteration_stats (it was bound into the stand-in but
never exercised) for both attention and SSM life cycles.
- TestSSMSupport.test_ssm_resume_records_intra_device_copy asserts the
SSM intra-device-copy delta end to end against the selected backend,
so the default C++ implementation is covered and a Python-backend
parity break fails the same test (verified: it fails against the
pre-fix Python backend).
Also comment why the recorder call sits outside the ssm_lc_id guard.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…d dedupe layout dispatch Review follow-up, two hardenings against the bug class fixed earlier in this PR: - MambaHybridCacheManagerV2.__init__ now raises TypeError when handed the V1 managers' model_type kwarg instead of conv_state_layout. Silently absorbing it into **kwargs is exactly how the wrong-layout bug went unnoticed. - The three copies of the "V2 takes conv_state_layout, V1 takes model_type" dispatch in _create_kv_cache_manager (kimi, nemotron, qwen3_hybrid branches) collapse into one helper with a single model_type-to-layout mapping, so a future hybrid branch cannot pick one convention and drop the other. Test: test_v2_manager_rejects_model_type_kwarg; existing kimi/qwen3 layout-selection tests cover the helper refactor. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
The module imported KvCache but the class is named _KVCache, so the file failed pytest collection and none of its tests ever ran. Review follow-up discovered while adding the _record_direct_iteration_stats case. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
2025-2026: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse precise annotations for the new helper.
kv_cache_manager_cls: typeand-> dictare too broad for this new function. Use a concrete manager-class type and a parameterized mapping, such astype[MambaHybridCacheManager]anddict[str, str].As per coding guidelines, every function must be annotated and imprecise
dicttypes must use precise generic types.Suggested annotation update
-def _mamba_conv_layout_kwargs(kv_cache_manager_cls: type, - model_type: str) -> dict: +def _mamba_conv_layout_kwargs( + kv_cache_manager_cls: type[MambaHybridCacheManager], + model_type: str) -> dict[str, str]:🤖 Prompt for 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. In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 2025 - 2026, Update _mamba_conv_layout_kwargs annotations to use type[MambaHybridCacheManager] for kv_cache_manager_cls and dict[str, str] for the return value, ensuring the referenced manager type is available in scope.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 2025-2026: Update _mamba_conv_layout_kwargs annotations to use
type[MambaHybridCacheManager] for kv_cache_manager_cls and dict[str, str] for
the return value, ensuring the referenced manager type is available in scope.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4af1b507-0961-4d9f-affd-80ae1ac73890
📒 Files selected for processing (6)
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytests/unittest/_torch/executor/test_mamba_cache_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_life_cycles.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #65017 [ run ] triggered by Bot. Commit: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Two fixes to the Kimi K3 path on KVCacheManagerV2, one commit per ticket.
TRTLLM-15216 — wrong conv-state layout (fix). The kimi_linear branch of the manager
selection hardcoded
model_type="qwen3_next"and never passedconv_state_layout. TheV2 manager defaults to
x_b_cand absorbsmodel_typethrough**kwargs, so an explicitV2 opt-in silently built the wrong layout with no error. Selects the kwarg from the
manager class, following the idiom already used by the qwen3_hybrid branch. A
strict=Truexfail test describing exactly this bug already existed; the marker isdropped and a companion test covers the V1 branch.
TRTLLM-15217 — SSM life cycles missing from iteration stats (fix). The page-movement
recorders dropped every non-attention life cycle, so KDA recurrent-state offload, onboard
and drop were invisible in iteration statistics. Filters removed at the recording sites
and in the Python mirror. The filter is kept for the global
allocTotalBlocks/allocNewBlockscounters and for reuse hit/miss range accounting,which are attention-only by definition.
Both are reachable only with an explicit
use_kv_cache_manager_v2=True;"auto"does notselect V2 for this model today, so default runs are unaffected.
Not in scope:
not held up by a design discussion about what that counter should mean.
Test Coverage
use_kv_cache_manager_v2=True), plus a newtest for the non-V2 branch.
the production filtering logic is what is exercised.
PR Checklist
[JIRA/NVBUG/None][type] SummaryDev Engineer Review
conv_state_layout; V1 managers retainmodel_type.model_typearguments."auto"is unchanged.QA Engineer Review
test_kimi_k3_v2_cache_manager_layout.model_typecoverage.test_offload_is_recorded_for_every_life_cycle.test_host_drop_is_recorded_for_every_life_cycle.test_direct_iteration_stats_are_recorded_for_every_life_cycle.test_onboard_counts_globally_only_for_attention.test_ssm_resume_records_intra_device_copy.test-db/,qa/, orwaives.txtchanges are reported.