[None][feat] Add per-conversation KV cache block reuse - #16115
Conversation
b5a8913 to
8260269
Compare
8260269 to
2ef8d7b
Compare
📝 WalkthroughWalkthroughAdds ChangesConversation-scoped KV reuse
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Request
participant KVCacheManagerV2
participant ConversationManager
participant _KVCache
Request->>KVCacheManagerV2: prepare first context chunk
KVCacheManagerV2->>ConversationManager: update_conversation(request)
Request->>KVCacheManagerV2: finish context
KVCacheManagerV2->>ConversationManager: save_drop_plan(request, kv_cache)
ConversationManager->>_KVCache: plan_committed_block_drop()
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: 3
🤖 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/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 157-208: Update ConversationManager and the free_resources cleanup
path to release per-conversation state when a tracked request terminates before
save_drop_plan, clearing current_request_id only when it matches the freed
request and preserving any planned drop handle as appropriate. Ensure all
early-termination paths use this cleanup, and evict conversation entries that no
longer contain active request or drop-plan state so _conversation_states cannot
grow without bound.
In `@tests/unittest/_torch/executor/test_kv_cache_manager_v2.py`:
- Around line 36-63: Add a cache_salt field to the _ContextRequest test helper,
separate from cache_salt_id, so KVCacheManagerV2._prepare_context_impl() and
prepare_resources() can access req.cache_salt for first-chunk requests without
AttributeError.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 613-647: Update the nested plan_drop helper in
test_planned_drop_handle to consume the stream holder’s pending finish event via
take_finish_event() after the cache operations complete and before closing the
cache. Preserve the existing drop-handle creation and return behavior while
matching neighboring test cleanup.
🪄 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: d942f280-a57f-448c-91c3-55b6c7f0c639
📒 Files selected for processing (11)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/_torch/executor/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/llmapi/test_llm_args.py
jiaganc
added
the
api-compatible
|
/bot run |
|
PR_Github #59372 [ run ] triggered by Bot. Commit: |
|
PR_Github #59372 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59399 [ run ] triggered by Bot. Commit: |
|
PR_Github #59399 [ run ] completed with state
|
|
[by Codex] @VALLIS-NERIA Friendly reminder to review this PR as the primary KV-cache-manager reviewer when you have a chance. Thanks! |
|
PR_Github #59855 [ run ] completed with state |
|
PR_Github #59860 [ run ] completed with state
|
JunyiXu-nv
left a comment
JunyiXu-nv
left a comment
There was a problem hiding this comment.
LGTM for LLM api changes.
|
/bot run |
|
PR_Github #59933 [ run ] triggered by Bot. Commit: |
|
PR_Github #59933 [ run ] completed with state
|
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
|
/bot run |
|
PR_Github #59984 [ run ] triggered by Bot. Commit: |
|
PR_Github #59984 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59992 [ run ] triggered by Bot. Commit: |
|
PR_Github #59992 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60249 [ run ] triggered by Bot. Commit: |
|
PR_Github #60249 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60316 [ run ] triggered by Bot. Commit: |
|
PR_Github #60316 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60376 [ run ] triggered by Bot. Commit: |
|
PR_Github #60376 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60390 [ run ] triggered by Bot. Commit: |
|
PR_Github #60390 [ run ] completed with state |
9 participants
Summary by CodeRabbit
New Features
per_conversationas a KV-cache block reuse policy.Tests
Description
Add a
per_conversationKV cache block reuse policy for KV cache manager V2.The core runtime can now create a
PlannedDropHandlefor committed SWA-window pages needed only by the next conversation turn. Dropping the handle removes already-droppable pages when no other plan covers them. Full-attention and attention-sink pages remain managed by the existing global reuse policy, while SSM state support remains a TODO.The PyTorch executor tracks the current request and latest planned drop handle for each conversation. Completing a context saves its drop plan and applies the preceding plan only when all required pages are available. Finishing or cancelling a request clears its active conversation state. Concurrent requests for the same conversation are ignored by conversation tracking, and requests without conversation parameters retain
per_requestbehavior.Test Coverage
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py.tests/unittest/_torch/executor/test_kv_cache_manager_v2.py.KvCacheConfigvalidation coverage intests/unittest/llmapi/test_llm_args.py.git diff --check, and relevant pre-commit hooks passed locally.PR Checklist
Please review the following before submitting this 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.