[None][fix] Fix DeepSeek V4 KV cache warmup handling and serveral other issues - #16466
Conversation
liji-nv
changed the title
📝 WalkthroughWalkthroughDeepSeek V4 cache sizing now includes extra KV tokens. CUDA graph capture preserves speculative-decode KV-length metadata across warmups and capture. KV-transfer tests now parse indexer-compressed buffers using explicit manager-provided sizes and scale dtypes. ChangesDeepSeek V4 cache sizing
Speculative decode CUDA graph capture
KV-transfer buffer parsing
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant CUDAGraphRunner
participant attn_metadata
participant CUDA_Graph
CUDAGraphRunner->>attn_metadata: Snapshot kv_lens_cuda
CUDAGraphRunner->>attn_metadata: Restore state before warmup
CUDAGraphRunner->>CUDA_Graph: Capture warmup and graph operations
CUDAGraphRunner->>attn_metadata: Restore state after capture
attn_metadata->>attn_metadata: Call on_update_kv_lens()
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}
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)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (1)
450-464: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRestore the snapshot on exceptional exits.
If a warmup forward,
postprocess_fn, or CUDA graph capture raises, the restore call is skipped andattn_metadata.kv_lens_cudacan remain mutated. Wrap the capture sequence intry/finally, retaining the per-warmup restore while guaranteeing cleanup on failure.🤖 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/cuda_graph_runner.py` around lines 450 - 464, Wrap the warmup loop and CUDA graph capture sequence in a try/finally so _restore_spec_decode_capture_state(attn_metadata, saved_kv_lens_cuda) always runs when _setup_spec_decoding_and_forward, postprocess_fn, or graph capture raises. Keep the existing restore after each warmup iteration, and use the finally block for the overall cleanup path.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (1)
42-55: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a typed attention-metadata protocol instead of
Any.The new helpers cannot statically verify
kv_lens_cuda,num_seqs, oron_update_kv_lens. Define a smallProtocolor shared metadata interface for this contract.As per coding guidelines, avoid
Anyand use precise type annotations.🤖 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/cuda_graph_runner.py` around lines 42 - 55, Replace the Any annotation in _save_spec_decode_capture_state and define a typed attention-metadata Protocol or shared interface exposing kv_lens_cuda, num_seqs, and on_update_kv_lens. Use that type for both helper signatures, preserving the existing optional-state behavior and restore logic.Source: Coding guidelines
tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)
152-169: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd capture-level coverage for the integration.
Coverage is currently insufficient: this test exercises only the helpers and would pass even if
CUDAGraphRunner.captureomitted restoration or restored at the wrong point. Add a focused test intests/unittest/_torch/executor/test_pytorch_model_engine.pycovering both warmups and the final post-capture restore, or track that integration coverage as follow-up.As per path instructions, test feedback should identify whether coverage is sufficient and provide a concrete file for follow-up.
🤖 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/_torch/executor/test_pytorch_model_engine.py` around lines 152 - 169, Add an integration-focused test for CUDAGraphRunner.capture in tests/unittest/_torch/executor/test_pytorch_model_engine.py that verifies speculative-decode KV lengths are restored before each of the two warmup forwards and again after capture completes. Exercise the capture flow rather than only _save_spec_decode_capture_state and _restore_spec_decode_capture_state, and assert the original metadata remains intact at every required point.Source: Path instructions
tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py (1)
87-113: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover the
max_num_tokens is Nonebranch.This test only validates
max_num_tokens + num_extra_kv_tokens; the fallbacktypical_seq_len + num_extra_kv_tokensbranch remains uncovered. Add a second case with_max_num_tokens=Noneand an explicitavg_seq_len, then assert the resulting capacity.🤖 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/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py` around lines 87 - 113, Extend test_mtp_extra_tokens_are_in_context_capacity with a second setup where _max_num_tokens is None and avg_seq_len is explicitly provided, then build the cache configuration and assert the capacity uses avg_seq_len plus num_extra_kv_tokens. Keep the existing max_num_tokens case unchanged.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 `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py`:
- Around line 891-893: Update the generation descriptor sizing and decode warmup
constraint near the context-capacity calculation to include
self.num_extra_kv_tokens consistently with _get_generation_bytes(). Ensure
generation warmups and pool sizing reserve the same extra-token capacity, rather
than using bare typical_seq_len or max_seq_len.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 450-464: Wrap the warmup loop and CUDA graph capture sequence in a
try/finally so _restore_spec_decode_capture_state(attn_metadata,
saved_kv_lens_cuda) always runs when _setup_spec_decoding_and_forward,
postprocess_fn, or graph capture raises. Keep the existing restore after each
warmup iteration, and use the finally block for the overall cleanup path.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 42-55: Replace the Any annotation in
_save_spec_decode_capture_state and define a typed attention-metadata Protocol
or shared interface exposing kv_lens_cuda, num_seqs, and on_update_kv_lens. Use
that type for both helper signatures, preserving the existing optional-state
behavior and restore logic.
In
`@tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py`:
- Around line 87-113: Extend test_mtp_extra_tokens_are_in_context_capacity with
a second setup where _max_num_tokens is None and avg_seq_len is explicitly
provided, then build the cache configuration and assert the capacity uses
avg_seq_len plus num_extra_kv_tokens. Keep the existing max_num_tokens case
unchanged.
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 152-169: Add an integration-focused test for
CUDAGraphRunner.capture in
tests/unittest/_torch/executor/test_pytorch_model_engine.py that verifies
speculative-decode KV lengths are restored before each of the two warmup
forwards and again after capture completes. Exercise the capture flow rather
than only _save_spec_decode_capture_state and
_restore_spec_decode_capture_state, and assert the original metadata remains
intact at every required point.
🪄 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: ff92d7bd-790b-4ed5-84bb-e8400e305b1e
📒 Files selected for processing (6)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
5f73620 to
5330b1f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59632 [ run ] triggered by Bot. Commit: |
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
5330b1f to
904a0c9
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59708 [ run ] triggered by Bot. Commit: |
|
PR_Github #59632 [ run ] completed with state |
|
PR_Github #59708 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59817 [ run ] triggered by Bot. Commit: |
|
PR_Github #59817 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59898 [ run ] completed with state |
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #59918 [ run ] triggered by Bot. Commit: |
|
PR_Github #59918 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59972 [ run ] triggered by Bot. Commit: |
|
PR_Github #59972 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59989 [ run ] triggered by Bot. Commit: |
|
PR_Github #59989 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60065 [ run ] triggered by Bot. Commit: |
|
PR_Github #60065 [ run ] completed with state |
9 participants
Summary by CodeRabbit
Bug Fixes
Tests
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)
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.