[https://nvbugs/6463967][fix] DeepSeek-V4 one-model MTP separate draft kv cache (TEP) - #16887
[https://nvbugs/6463967][fix] DeepSeek-V4 one-model MTP separate draft kv cache (TEP)#16887nvxuanyuc wants to merge 1 commit into
Conversation
WalkthroughDeepSeek-V4 sparse attention now supports separate draft KV-cache buffers for one-model MTP. Draft replay snapshots and restores sparse metadata, DSA indexer state, and target mappings. Cache selection and B200 accuracy 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: snapshot target sparse state
SpeculativeInterface->>DeepseekV4TrtllmAttentionMetadata: apply draft sparse state
DeepseekV4TrtllmAttentionMetadata->>DSAtrtllmAttentionMetadata: prepare DSA replay state
SpeculativeInterface->>DeepseekV4TrtllmAttentionMetadata: restore target sparse state
DeepseekV4TrtllmAttentionMetadata->>DSAtrtllmAttentionMetadata: restore DSA replay state
Possibly related PRs
Suggested labels: 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
| @property | ||
| def per_layer_compress_ratios(self) -> List[int]: | ||
| """Per-layer compress ratios (MTP-extended), indexed by global layer.""" | ||
| return self._compress_ratios |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the required modern type annotations.
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py#L603-L606: uselist[int]instead ofList[int].tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py#L653-L761: annotate helper arguments and return types; useNonefor mutating helpers and a precise snapshot type.tests/integration/defs/accuracy/test_llm_api_pytorch.py#L3880-L3880: annotate the test method with-> None.
📍 Affects 3 files
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py#L603-L606(this comment)tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py#L653-L761tests/integration/defs/accuracy/test_llm_api_pytorch.py#L3880-L3880
🤖 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 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.
Source: Coding guidelines
|
/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 |
|
PR_Github #62353 [ run ] triggered by Bot. Commit: |
|
PR_Github #62353 [ run ] completed with state |
| # Only SWA-only (compress_ratio 1) draft layers are supported; compress | ||
| # (128) / indexer (4) draft layers are intentionally unsupported. | ||
| draft_ratios = { | ||
| draft_mgr.per_layer_compress_ratios[layer_idx] for layer_idx in draft_mgr.pp_layers |
There was a problem hiding this comment.
It seems that you can directly use self.compress_ratios to get the list and don't need this per_layer_compress_ratios.
There was a problem hiding this comment.
Addressed. Given the MTP layers duplicate the final configured ratio, the guard now checks self.compress_ratios[-1] directly. And the metadata now locally extends self.compress_ratios for pointer construction, per_layer_compress_ratios helper has been removed.
| # draft manager (the DSA block above only handles the indexer k-cache). | ||
| from ..attention_backend.sparse.deepseek_v4.deepseek_v4 import \ | ||
| DeepseekV4TrtllmAttentionMetadata | ||
| if isinstance(attn_metadata, DeepseekV4TrtllmAttentionMetadata): |
There was a problem hiding this comment.
It'd be better to add a prepare_for_draft_replay in deepseek_v4 metadata. And here we can get saved = attn_metadata.prepare_for_draft_replay.
Same for the above DSA metadata.
There was a problem hiding this comment.
Addressed. Added replay hooks to TrtllmAttentionMetadata, with the existing indexer preparation and restoration moved into the DSA metadata implementation. The DSv4 metadata implementation extends those hooks with its sparse-state apply and restore.
| if isinstance(attn_metadata, DeepseekV4TrtllmAttentionMetadata): | ||
| saved_dsv4_state = attn_metadata.save_target_sparse_state() | ||
| try: | ||
| if saved_dsv4_state is not None: |
There was a problem hiding this comment.
I think this part can be simplified to
# DeepSeek-V4: xxxxxxxxx
if hasattr(attn_metadata, "save_target_sparse_state"):
saved_dsv4_state = attn_metadata.save_target_sparse_state()
attn_metadata.apply_draft_sparse_state()
There was a problem hiding this comment.
Addressed. Replaced the type check with hasattr. The None guard remains because this context also serves non-MTP spec dec paths.
|
[by Codex] @liji-nv Friendly review reminder: this PR is awaiting your review. Thanks! |
…t kv cache (TEP) Signed-off-by: Xuanyu Chen <xuanyuc@nvidia.com>
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py (1)
791-799: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a single cached flag for the draft-support check.
The
hasattr(draft_mgr, "compute_sliding_block_tables")check appears here and in_init_draft_sparse_buffers. The two gates must stay in agreement, becauseapply_draft_sparse_statecopies tables that only this branch computes. Set one boolean in_init_draft_sparse_buffers(for exampleself._has_draft_sparse_state) and test it in both places.🤖 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/deepseek_v4.py` around lines 791 - 799, Introduce a cached draft-support boolean in _init_draft_sparse_buffers based on whether draft_kv_cache_manager provides compute_sliding_block_tables, then reuse that flag in this preparation block instead of repeating hasattr. Ensure apply_draft_sparse_state and the table-computation branch remain gated by the same flag.tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)
717-730: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider persistent snapshot buffers instead of
clone()per draft step.
prepare_for_draft_replayruns on every draft replay step. Each call allocates six new tensors, including twomax_num_tokensint64 device buffers and their host counterparts. Allocate the snapshot buffers once and usecopy_into them. This removes per-step allocator traffic on a hot path and keeps addresses stable.🤖 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/dsa.py` around lines 717 - 730, Update the snapshot handling in prepare_for_draft_replay to use six persistent buffers allocated once during initialization, rather than cloning the tensors on every draft replay step. Copy each current buffer into its corresponding persistent snapshot with copy_, preserving the existing saved keys and tensor contents while keeping device and host buffer addresses stable.
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py`:
- Around line 791-799: Introduce a cached draft-support boolean in
_init_draft_sparse_buffers based on whether draft_kv_cache_manager provides
compute_sliding_block_tables, then reuse that flag in this preparation block
instead of repeating hasattr. Ensure apply_draft_sparse_state and the
table-computation branch remain gated by the same flag.
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 717-730: Update the snapshot handling in prepare_for_draft_replay
to use six persistent buffers allocated once during initialization, rather than
cloning the tensors on every draft replay step. Copy each current buffer into
its corresponding persistent snapshot with copy_, preserving the existing saved
keys and tensor contents while keeping device and host buffer addresses stable.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b1125ddd-fd86-4f7f-9530-0a7811ee864f
📒 Files selected for processing (3)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/_torch/attention_backend/trtllm.py
|
PR_Github #63542 [ run ] triggered by Bot. Commit: |
|
PR_Github #63542 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63580 [ run ] triggered by Bot. Commit: |
|
PR_Github #63580 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63597 [ run ] triggered by Bot. Commit: |
|
PR_Github #63597 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63783 [ run ] triggered by Bot. Commit: |
5 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
try/finallyhandling.NotImplementedError.PP > 1configurations.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.ymlwith a 60-minute timeout.Verdict: sufficient
Description
Wires DeepSeek-V4's sparse attention forward to the existing separate draft KV cache mechanism
for one-model MTP under TEP.
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] ~1-2% decode-rate overhead
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.