[https://nvbugs/6013562][fix] fix kv cache allocation is double the budget for vswa + eagle#12855
Conversation
|
/bot run |
📝 WalkthroughWalkthroughThe Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 783-790: Resolve the draft manager class before mutating shared
budget and only override max_gpu_total_bytes when the draft manager is
V1-compatible; wrap the temporary override in try/finally to always restore the
original budget. Concretely: in the logic that decides draft_kv_cache_config
(the branch using self._kv_cache_manager_cls, is_vswa_enabled, and calling
_split_kv_cache_budget_for_draft()), first compute the class that would be used
for the draft (by invoking the same resolution used in
_create_kv_cache_manager() or otherwise deriving the draft manager class from
self._draft_model_engine), check whether that class is NOT a subclass of
KVCacheManagerV2, and only then set the shared
self._kv_cache_config.max_gpu_total_bytes override; perform the override inside
a try/finally so the original max_gpu_total_bytes is restored even if creating
the draft manager or computing draft_kv_cache_config raises. Apply the same
pattern to the other identical block handling draft budget (the second
occurrence referenced in the review).
🪄 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: Pro
Run ID: 087c7dc6-7041-4d8a-aced-1b58d18e9cbc
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/_util.py
|
PR_Github #42387 [ run ] triggered by Bot. Commit: |
|
PR_Github #42387 [ run ] completed with state
|
dongfengy
force-pushed
the
user/dongfengy/kvcachealloc
branch
from
5a5bfe2 to
ab9dafd
Compare
|
/bot run |
|
PR_Github #42417 [ run ] triggered by Bot. Commit: |
|
PR_Github #42417 [ run ] completed with state
|
dongfengy
force-pushed
the
user/dongfengy/kvcachealloc
branch
from
ab9dafd to
bc673b7
Compare
|
/bot run |
|
PR_Github #42474 [ run ] triggered by Bot. Commit: |
|
PR_Github #42474 [ run ] completed with state
|
dongfengy
force-pushed
the
user/dongfengy/kvcachealloc
branch
2 times, most recently
from
5eadb9e to
240420f
Compare
|
/bot run |
|
PR_Github #42567 [ run ] triggered by Bot. Commit: |
|
PR_Github #42567 [ run ] completed with state
|
|
/bot run |
|
PR_Github #42587 [ run ] triggered by Bot. Commit: |
|
PR_Github #42587 [ run ] completed with state
|
dongfengy
force-pushed
the
user/dongfengy/kvcachealloc
branch
from
240420f to
a52ea30
Compare
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #42734 [ run ] triggered by Bot. Commit: |
|
PR_Github #43594 [ run ] completed with state
|
dongfengy
force-pushed
the
user/dongfengy/kvcachealloc
branch
from
153ca92 to
8fcdfda
Compare
|
/bot run |
|
PR_Github #43659 [ run ] triggered by Bot. Commit: |
|
PR_Github #43659 [ run ] completed with state
|
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #43670 [ run ] triggered by Bot. Commit: |
|
PR_Github #43670 [ run ] completed with state
|
dongfengy
force-pushed
the
user/dongfengy/kvcachealloc
branch
from
8fcdfda to
fb6184c
Compare
|
/bot run |
|
PR_Github #43788 [ run ] triggered by Bot. Commit: |
|
PR_Github #43788 [ run ] completed with state
|
|
/bot run |
|
PR_Github #43834 [ run ] triggered by Bot. Commit: |
|
PR_Github #43834 [ run ] completed with state
|
dongfengy
force-pushed
the
user/dongfengy/kvcachealloc
branch
from
fb6184c to
e40f8ef
Compare
|
/bot run |
|
PR_Github #43936 [ run ] triggered by Bot. Commit: |
|
PR_Github #43936 [ run ] completed with state
|
…udget for vswa Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
dongfengy
force-pushed
the
user/dongfengy/kvcachealloc
branch
from
e40f8ef to
6f7432c
Compare
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #44045 [ run ] triggered by Bot. Commit: |
|
PR_Github #44045 [ run ] completed with state
|
|
/bot run |
|
PR_Github #44134 [ run ] triggered by Bot. Commit: |
|
PR_Github #44134 [ run ] completed with state |
3 participants
…udget for vswa
Summary by CodeRabbit
Description
Test Coverage
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)
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.