[None][feat] Integrate KVCMv2 commit_min_snapshot in runtime - #16185
Conversation
jiaganc
changed the title
commit_min_snapshot in runtime
|
This PR depends on #15752 and should wait for it to be merged first. |
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
52fe3a7 to
2a525b0
Compare
📝 WalkthroughWalkthroughChangesKV cache reuse behavior
Estimated code review effort: 3 (Moderate) | ~20 minutes 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}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)
2739-2762: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueTODO is a legitimate, self-scoped perf follow-up — not blocking.
The
is_end=TrueTODO matches thecommit()/_KVCachesemantics shown in context (partial-block "move vs copy" on final chunk), so this is purely a performance optimization for disaggregated-prefill, not a correctness gap. Fine to defer.Happy to draft the
is_end=Truewiring for the last context chunk (usingrequest.context_remaining_length == 0) if you'd like to pick this up now rather than later — let me know.🤖 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/kv_cache_manager_v2.py` around lines 2739 - 2762, No code change is required; retain the existing TODO in try_commit_blocks. Treat is_end=True wiring for the final disaggregated-prefill context chunk as a deferred performance follow-up, preserving the current kv_cache.commit(tokens) behavior.tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)
1524-1586: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCoverage gap: no test exercises
KVCacheManagerV2._build_cache_config'scommit_min_snapshotformula against a manager that actually has SSM layers.
TestSSMSupport._make_ssm_configconstructsKVCacheManagerConfigby hand and passescommit_min_snapshotexplicitly, so it never goes through the real_build_cache_configcomputation added in this PR. Given the concern raised ontensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py(lines 1545-1559) aboutenable_block_reuse=False+ SSM layers potentially violating thecommit_min_snapshotinvariant, this file would be the natural place to add a regression test (e.g. build a manager subclass with SSM layers via_build_cache_config, setenable_block_reuse=False, and assert construction either raises clearly or is handled). Coverage here is currently insufficient for that interaction and should be added either in this PR or as an explicit follow-up intests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py.🤖 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 1524 - 1586, Add a regression test in TestSSMSupport that exercises KVCacheManagerV2._build_cache_config with SSM layers and enable_block_reuse=False, rather than relying on _make_ssm_config’s explicitly supplied commit_min_snapshot. Assert the construction behavior required by the invariant—either a clear error is raised or the configuration is handled correctly—and preserve existing SSM test setup and teardown.Source: Path instructions
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py (1)
939-955: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueExtract the shared
commit_min_snapshotpredicate
This is duplicated inpyexecutor/kv_cache_manager_v2.pyand here; a small shared helper onKVCacheManagerV2would keep future subclasses consistent.🤖 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/attention_backend/sparse/deepseek_v4/cache_manager.py` around lines 939 - 955, Extract the duplicated commit_min_snapshot predicate into a shared helper on KVCacheManagerV2, then update the configuration construction here and in pyexecutor/kv_cache_manager_v2.py to call that helper. Preserve the existing condition based on enable_block_reuse and block_reuse_policy.
🤖 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/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py`:
- Around line 38-49: Update the BlockReusePolicy import in the DeepSeek V4 cache
manager test to import it from
tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 instead of resource_manager,
leaving the other imports unchanged.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py`:
- Around line 939-955: Extract the duplicated commit_min_snapshot predicate into
a shared helper on KVCacheManagerV2, then update the configuration construction
here and in pyexecutor/kv_cache_manager_v2.py to call that helper. Preserve the
existing condition based on enable_block_reuse and block_reuse_policy.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 2739-2762: No code change is required; retain the existing TODO in
try_commit_blocks. Treat is_end=True wiring for the final disaggregated-prefill
context chunk as a deferred performance follow-up, preserving the current
kv_cache.commit(tokens) behavior.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 1524-1586: Add a regression test in TestSSMSupport that exercises
KVCacheManagerV2._build_cache_config with SSM layers and
enable_block_reuse=False, rather than relying on _make_ssm_config’s explicitly
supplied commit_min_snapshot. Assert the construction behavior required by the
invariant—either a clear error is raised or the configuration is handled
correctly—and preserve existing SSM test setup and teardown.
🪄 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: 504a6fa7-9b10-4b22-9c23-24ef87a47e3a
📒 Files selected for processing (6)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/executor/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
|
/bot run |
|
PR_Github #58925 [ run ] triggered by Bot. Commit: |
|
PR_Github #58925 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58993 [ run ] triggered by Bot. Commit: |
|
PR_Github #58993 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59084 [ run ] triggered by Bot. Commit: |
|
PR_Github #59084 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59113 [ run ] triggered by Bot. Commit: |
|
PR_Github #59113 [ run ] completed with state |
5 participants
Summary by CodeRabbit
Description
Integrate
commit_min_snapshotsupport for SSM reuse from #15752 and connect it to the KV cache manager V2 block-reuse policies.The cache manager now commits a minimal SSM snapshot when block reuse is enabled for policies other than
all_reusable, propagates partial-reuse configuration, and handles the final context commit path. The change includes coverage for SSM/SWA reuse, partial-page commits, cache statistics, DeepSeek V4 sparse cache configuration, and V2 reuse-policy selection.Test Coverage
python3 -m py_compilepassed for the directly modified Python implementation and test files.pytestinstalled.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.