[https://nvbugs/6463967][fix] DeepSeek-V4 one-model MTP separate draft kv cache (TEP) - #16887
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughDeepSeek-V4 sparse attention now supports separate draft KV-cache buffers for one-model MTP. Backend hooks manage draft forward and replay state. Cache selection and integration coverage are updated. ChangesDeepSeek-V4 draft KV-cache support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant SpeculativeInterface
participant DeepseekV4TrtllmAttentionMetadata
participant DSAtrtllmAttentionMetadata
SpeculativeInterface->>DeepseekV4TrtllmAttentionMetadata: prepare draft forward or replay
DeepseekV4TrtllmAttentionMetadata->>DSAtrtllmAttentionMetadata: prepare DSA replay state
SpeculativeInterface->>DeepseekV4TrtllmAttentionMetadata: restore backend state
DeepseekV4TrtllmAttentionMetadata->>DSAtrtllmAttentionMetadata: restore DSA replay state
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: 2
🤖 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 603-606: Modernize the annotations at all three sites: in
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py:603-606,
change the per_layer_compress_ratios property return type to list[int]; in
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py:653-761,
annotate helper parameters and returns, using None for mutating helpers and the
precise snapshot type for snapshot-returning helpers; in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:3880, add -> None to the
test method.
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 196-204: Make the DeepseekV4TrtllmAttentionMetadata mutation in
prepare_attn_metadata_for_draft_replay exception-safe by wrapping
apply_draft_sparse_state() with restoration of the saved sparse state when it
raises. Restore the saved metadata before re-raising the original exception,
while preserving the existing saved-state behavior on success.
🪄 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: 78af6d9d-93bc-409f-919e-60efd609be8d
📒 Files selected for processing (7)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/speculative/interface.pytests/integration/defs/accuracy/references/gsm8k.yamltests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_dgx_b200.yml
|
/bot run --disable-fail-fast |
|
PR_Github #62069 [ run ] triggered by Bot. Commit: |
|
PR_Github #62069 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62256 [ run ] triggered by Bot. Commit: |
|
PR_Github #62256 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #66876 [ run ] triggered by Bot. Commit: |
|
PR_Github #66876 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66933 [ run ] triggered by Bot. Commit: |
|
PR_Github #66933 [ run ] completed with state |
…t kv cache (TEP) Signed-off-by: Xuanyu Chen <xuanyuc@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #67134 [ run ] triggered by Bot. Commit: |
|
PR_Github #67134 [ run ] completed with state
|
|
/bot run |
|
PR_Github #67183 [ run ] triggered by Bot. Commit: |
|
PR_Github #67183 [ run ] completed with state
|
|
/bot run |
|
PR_Github #67265 [ run ] triggered by Bot. Commit: |
|
PR_Github #67265 [ run ] completed with state
|
|
/bot run |
|
PR_Github #67346 [ run ] triggered by Bot. Commit: |
|
PR_Github #67346 [ run ] completed with state
|
|
/bot run |
|
PR_Github #67504 [ run ] triggered by Bot. Commit: |
|
PR_Github #67504 [ run ] completed with state |
10 participants
Summary
Adds separate draft KV-cache support for DeepSeek-V4 sparse attention during one-model MTP under TEP.
PP > 1.l0_dgx_b200.yml.PP=1, andbeam_width=1.Dev Engineer Review
NotImplementedError.PP > 1fallback prevents unsupported cache allocation.797c4a5and the fourth run fore56cedb.Verdict: sufficient
QA Engineer Review
Added test function:
TestDeepSeekV4Flash::test_tep_mtp_separate_draft_kv_cacheCoverage:
tests/integration/test_lists/test-db/l0_dgx_b200.yml.Verdict: sufficient
Description
Wires DeepSeek-V4's sparse attention forward to the existing separate draft KV cache mechanism for one-model MTP under TEP.
Fix issue #17024
Design
DeepseekV4CacheManagerholds the MTP draft layers; the metadata carriesdraft-sized sliding block tables + draft pool base pointers, allocated separately.
_num_tablesand precomputed sliding tables are populated inprepare()viadraft_mgr.compute_sliding_block_tables(...), mirroring the target manager.save/apply/restore_target_sparse_stateat both theeager
draft_kv_cache_contextand the CUDA-graphprepare_attn_metadata_for_draft_replay,gated by
isinstance(attn_metadata, DeepseekV4TrtllmAttentionMetadata)(mirrors the DSAbranch). The swap is a pointer rebind plus per-batch copy.
compress_ratio1, asserted);PP=1only (folds otherwise);beam_width=1(pre-existing DSv4 + MTP constraint).Validation Results
[DeepSeek-V4-Flash | TEP4 | fp8kv]
Separate vs unified draft KV
Perf [nextn=1 | ISL/OSL 2k/2k | max conc 64]
Test Coverage
TestDeepSeekV4Flash::test_tep_mtp_separate_draft_kv_cache(GSM8K), registered inl0_dgx_b200.yml. Local run: 95.53.GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.