[TRTLLM-15216][fix] Kimi K3 on KVCacheManagerV2: conv-state layout and SSM iteration stats - #17447
Open
brnguyen2 wants to merge 2 commits into
Open
[TRTLLM-15216][fix] Kimi K3 on KVCacheManagerV2: conv-state layout and SSM iteration stats#17447brnguyen2 wants to merge 2 commits into
brnguyen2 wants to merge 2 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
WalkthroughChangesHybrid KV-cache management
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CacheManager
participant BlockRadixTree
participant KvCache
participant PythonBinding
CacheManager->>BlockRadixTree: match lookup path
BlockRadixTree->>BlockRadixTree: compute attention-only prefix
BlockRadixTree->>BlockRadixTree: apply SSM-aware pruning
BlockRadixTree-->>CacheManager: return ReuseMatch with both counts
CacheManager->>KvCache: initialize diagnostic count
PythonBinding->>KvCache: request pre-hybrid token count
KvCache-->>PythonBinding: return numTokensBeforeHybridPruning
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
Comment on lines
+2397
to
+2434
| 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() |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether the new test is registered in CI test lists.
set -euo pipefail
rg -n 'test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation' tests/integration/test_lists/ || echo "Not found in test_lists"Repository: NVIDIA/TensorRT-LLM
Length of output: 181
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- changed-file summary ---'
git diff --stat
echo '--- target test occurrences and context ---'
rg -n -C 8 'test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation|class TestSSMSupport' \
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
echo '--- test-list roots and matching entries ---'
git ls-files 'tests/integration/test_lists/**' | sed -n '1,120p'
rg -n -i -C 2 'kv_cache_manager_v2|kv.cache.manager.v2|hybrid_pruning|recurrent_truncation|TestSSMSupport' \
tests/integration/test_lists/ || true
echo '--- relevant changed target-file diff ---'
git diff -- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py | sed -n '1,260p'
echo '--- implementation/statistics symbols ---'
rg -n -C 5 'iter_intra_device_copy|KVCacheIterationStatsDelta|num_tokens_before_hybrid_pruning|deferred|resume\(' \
--glob '_kv_cache.py' --glob 'kvCache.cpp' --glob '*.py' --glob '*.cpp' --glob '*.h' --glob '*.hpp' . | sed -n '1,500p'Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- repository state ---'
git status --short
echo '--- target lines ---'
sed -n '2360,2460p' tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
echo '--- exact test-name search ---'
rg -n 'test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation|_get_num_tokens_before_hybrid_pruning' \
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py || true
echo '--- tracked implementation candidates ---'
git ls-files | rg '(^|/)(_kv_cache\.py|kvCache\.(cpp|h|hpp)|.*iteration.*stats.*|.*kv.*cache.*)' | sed -n '1,240p'
echo '--- exact statistics symbols in likely source/test files ---'
rg -l 'KVCacheIterationStatsDelta|iter_intra_device_copy_blocks|iter_intra_device_copy_bytes' \
--glob '*.py' --glob '*.cpp' --glob '*.h' --glob '*.hpp' . | sed -n '1,160p'
echo '--- test-list conventions around kv_cache_manager_v2_tests ---'
for f in tests/integration/test_lists/test-db/l0_a10.yml \
tests/integration/test_lists/test-db/l0_h100.yml \
tests/integration/test_lists/test-db/l0_b200.yml; do
echo "### $f"
grep -n -B 4 -A 5 'kv_cache_manager_v2_tests' "$f"
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 17668
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- Python KV cache methods and stats references ---'
rg -n -C 12 '^[[:space:]]*def (resume|suspend|close)|iter_intra_device_copy|pending_stats|copy_engine|copy' \
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py | sed -n '1,420p'
echo '--- C++ KV cache methods and stats references ---'
rg -n -C 12 'KVCache::(resume|suspend|close)|iterIntraDeviceCopy|intraDeviceCopy|intra_device_copy|PendingStats|resume\\(' \
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/batch_manager/kv_cache_manager_v2/pendingStats.h \
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h | sed -n '1,520p'
echo '--- existing SSM and iteration-stats tests ---'
rg -n -C 8 'SSM|ssm|iter_intra_device_copy|intra_device_copy|resume' \
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py \
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py \
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py \
tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py | sed -n '1,520p'
echo '--- source-level verifier: list resume methods and nearby stats/copy calls ---'
python3 - <<'PY'
from pathlib import Path
import ast
paths = [
Path("tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py"),
]
for path in paths:
tree = ast.parse(path.read_text())
print(f"### {path}")
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "resume":
print(f"resume:{node.lineno}-{node.end_lineno}")
for child in ast.walk(node):
if isinstance(child, ast.Call):
fn = ast.unparse(child.func)
if any(word in fn.lower() for word in ("copy", "stat", "resume", "pending")):
print(f" call line {child.lineno}: {fn}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 19976
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- complete Python resume path ---'
sed -n '1140,1375p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
echo '--- Python copy/statistics methods ---'
sed -n '490,550p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
sed -n '1375,1515p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
echo '--- C++ resume and copy/statistics symbols ---'
rg -n -F -C 10 'KVCache::resume' cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp
rg -n -F -C 10 'iterIntraDeviceCopy' 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/batch_manager/kv_cache_manager_v2/pendingStats.h \
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/stats.h || true
rg -n -F -C 10 'intraDeviceCopy' cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp \
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h || true
echo '--- existing copy/statistics assertions ---'
rg -n -C 6 'iter_intra_device_copy|intra_device_copy|SSM|ssm' \
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py \
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py \
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py \
tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py | sed -n '1,420p'
echo '--- source-level verifier for resume calls and assignments ---'
python3 - <<'PY'
from pathlib import Path
import ast
path = Path("tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == "resume":
print(f"resume:{node.lineno}-{node.end_lineno}")
for child in sorted(ast.walk(node), key=lambda n: (getattr(n, "lineno", 0), getattr(n, "col_offset", 0))):
if isinstance(child, ast.Call):
fn = ast.unparse(child.func)
if any(word in fn.lower() for word in ("copy", "stat", "pending", "commit")):
print(f"call:{child.lineno}:{fn}")
elif isinstance(child, ast.Assign):
value = ast.unparse(child.value)
if "IterationStats" in value or "intra_device" in value:
print(f"assign:{child.lineno}:{ast.unparse(child)}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 20497
🏁 Script executed:
#!/bin/bash
set -u
echo '--- exact deferred-copy and stat lines ---'
rg -n -F 'batched_copy(' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
rg -n -F '_record_direct_iteration_stats(' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
rg -n -F 'iter_intra_device_copy_blocks' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py \
tests/unittest/kv_cache_manager_v2_tests tests/integration/defs/kv_cache || true
echo '--- all suite registrations ---'
rg -n -F 'unittest/kv_cache_manager_v2_tests' tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
echo '--- AST verifier for the deferred-copy loop ---'
python3 - <<'PY'
from pathlib import Path
import ast
path = Path("tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py")
tree = ast.parse(path.read_text())
resume = next(
node for node in ast.walk(tree)
if isinstance(node, ast.FunctionDef) and node.name == "resume"
)
print(f"resume_range={resume.lineno}-{resume.end_lineno}")
for node in ast.walk(resume):
if isinstance(node, ast.For):
calls = [
ast.unparse(child.func)
for child in ast.walk(node)
if isinstance(child, ast.Call)
]
if "batched_copy" in calls:
print(f"copy_loop_range={node.lineno}-{node.end_lineno}")
print(f"copy_loop_calls={calls}")
for child in node.body:
if isinstance(child, ast.If):
print(f"loop_if_line={child.lineno}")
print(ast.dump(child.test, include_attributes=False))
branch_calls = [
ast.unparse(desc.func)
for desc in ast.walk(child)
if isinstance(desc, ast.Call)
]
print(f"branch_calls={branch_calls}")
PY
echo '--- existing test assertions around iteration stats ---'
rg -n -C 5 'iter_intra_device_copy|iteration_stats|stats_enabled|commit_pending_stats' \
tests/unittest/kv_cache_manager_v2_tests tests/integration/defs/kv_cache | sed -n '1,360p'Repository: NVIDIA/TensorRT-LLM
Length of output: 32509
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
source = Path("tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py").read_text()
start = source.index(" # Deferred copy: for partial blocks and SSM")
end = source.index(" self._never_resumed = False", start)
section = source[start:end]
copy_pos = section.index("batched_copy(")
ssm_guard_pos = section.index("if lc_idx == ssm_lc_id:")
stats_guard_pos = section.index("if lc_idx != ssm_lc_id:")
stats_call_pos = section.index("self._record_direct_iteration_stats(")
assert copy_pos < stats_guard_pos < stats_call_pos
assert ssm_guard_pos < stats_guard_pos
assert "iter_intra_device_copy_blocks=1" in section[stats_call_pos:]
print("deferred_copy_has_ssm_branch_before_non_ssm_stats_branch: true")
print("ssm_deferred_copy_is_outside_direct_stats_branch: true")
lists = [
Path("tests/integration/test_lists/test-db/l0_a10.yml"),
Path("tests/integration/test_lists/test-db/l0_b200.yml"),
Path("tests/integration/test_lists/test-db/l0_h100.yml"),
]
assert all("unittest/kv_cache_manager_v2_tests" in path.read_text() for path in lists)
print("kv_cache_manager_v2_suite_registered_in_test_db: true")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 329
Add SSM deferred-copy iteration-stat coverage.
- Added test:
test_num_tokens_before_hybrid_pruning_isolates_recurrent_truncation. - The suite is registered in
l0_a10.yml,l0_b200.yml, andl0_h100.yml. - Coverage verdict: insufficient.
resume()performs the SSM deferred copy, butiter_intra_device_copy_*is recorded only for non-SSM pages. Add a regression assertion after resuming a reused SSM snapshot and record the SSM delta.
🤖 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_manager_v2.py` around
lines 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.
Source: Path instructions
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
|
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
2 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] Summary