[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
|
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
| @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 #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: |
|
PR_Github #63783 [ run ] completed with state |
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Checked the default path is untouched: prepare_attn_metadata_for_draft_replay and draft_kv_cache_context both still bail on not isinstance(attn_metadata, TrtllmAttentionMetadata), the new prepare_for_draft_replay/restore_after_draft_replay on TrtllmAttentionMetadata are no-ops, and _init_draft_sparse_buffers early-returns without a separate draft manager. The DSA block moved out of speculative/interface.py is byte-equivalent apart from the saved_dsa_state -> saved_backend_state["dsa"] key, which has no other reader. Confirmed the two entry points can't nest (the model_engine wrapper only surrounds cuda_graph_runner.replay, which doesn't re-run the mtp.py context manager), so save_target_sparse_state can't snapshot draft pointers as target.
Not blocking, for a follow-up: draft_kv_cache_context reaches the DSV4 sparse hooks via hasattr(attn_metadata, "save_target_sparse_state") while this same PR introduces proper base-class virtuals for the neighbouring prepare_for_draft_replay — worth making the sparse-state pair virtual too. Also _init_draft_sparse_buffers infers the draft layer's ratio from the target's compress_ratios[-1] and raises NotImplementedError from __post_init__; an explicit draft-side ratio would be less implicit.
The one open thread is CodeRabbit's List[int] -> list[int] annotation nit on an outdated diff — style only, not blocking.
| self.draft_swa_buffer_ptrs = None | ||
| self.draft_compressed_buffer_ptrs = None | ||
| draft_mgr = self.draft_kv_cache_manager | ||
| if draft_mgr is None or not hasattr(draft_mgr, "compute_sliding_block_tables"): |
There was a problem hiding this comment.
NIT: as dsv4 is only supported by DSV4 kv cache manager, I suggest remove hasattr call here and below. Or even if we want to support more types of kvcache managers, I suggest explicitly use isinstance to confirm.
There was a problem hiding this comment.
Addressed with an explicit isinstance check. The guard is still kept because DSv4 can use a non-DSv4 draft manager with other speculative-decoding methods such as EAGLE3.
|
|
||
| # DeepSeek-V4: repoint SWA/compress tables and pool pointers to the draft. | ||
| if hasattr(attn_metadata, "save_target_sparse_state"): | ||
| saved_dsv4_state = attn_metadata.save_target_sparse_state() |
There was a problem hiding this comment.
I wonder if it is possible to use our new interface, prepare_for_draft_replay and restore_after_draft_replay? If so then they are not only for replay, I suggest rename these api to prepare_for_draft_forward and restore_after_draft_forward.
There was a problem hiding this comment.
Addressed. Merged DSv4’s separate save/apply operations into one hook and renamed the pair to prepare_for_draft_forward and restore_after_draft_forward.
We retain two hook pairs because some state must be prepared during the eager (capture) path, while other state only requires refresh before graph replay. For example, DSv4’s tensor/pointer switch is needed in both paths, whereas DSA’s in-place metadata refresh is replay-specific.
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
520-533: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the draft-hook state contract.
State that each
prepare_*result is passed unchanged to its pairedrestore_*method. AddReturnsandArgssections for these public hooks.Proposed documentation update
def prepare_for_draft_replay(self) -> dict | None: - """Prepare all backend state required for draft CUDA graph replay.""" + """Prepare backend state for draft CUDA graph replay. + + Returns: + Saved backend state for `restore_after_draft_replay`, or `None`. + """ return None def restore_after_draft_replay(self, saved_state: dict | None) -> None: - """Restore backend state saved for draft CUDA graph replay.""" + """Restore backend state after draft CUDA graph replay. + + Args: + saved_state: State returned by `prepare_for_draft_replay`. + """As per coding guidelines, “Use docstrings rather than comments for externally usable interfaces, Google-style docstrings for classes and functions.”
🤖 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/trtllm.py` around lines 520 - 533, Update the docstrings for prepare_for_draft_replay, restore_after_draft_replay, prepare_for_draft_forward, and restore_after_draft_forward to use Google-style Returns and Args sections, documenting that each prepare method’s result is passed unchanged to its paired restore method and describing the saved_state argument and return values.Source: Coding guidelines
🤖 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 `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py`:
- Line 3293: Extend
TestPrepareRestoreAttnMetadataForDraftReplay.test_prepare_swaps_and_restore_recovers
beyond the None case by configuring a nested non-None draft-replay state and
asserting it is propagated during prepare and restored afterward. Preserve the
existing None coverage and add the test to the appropriate coverage listing if
required by the repository’s test tracking.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 520-533: Update the docstrings for prepare_for_draft_replay,
restore_after_draft_replay, prepare_for_draft_forward, and
restore_after_draft_forward to use Google-style Returns and Args sections,
documenting that each prepare method’s result is passed unchanged to its paired
restore method and describing the saved_state argument and return values.
🪄 Autofix
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: 04e953a5-9e7d-426e-b697-d75fc0e231bb
📒 Files selected for processing (9)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/_torch/attention_backend/trtllm.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.ymltests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
🚧 Files skipped from review as they are similar to previous changes (6)
- tensorrt_llm/_torch/pyexecutor/_util.py
- tests/integration/test_lists/test-db/l0_dgx_b200.yml
- tensorrt_llm/_torch/attention_backend/sparse/dsa.py
- tests/integration/defs/accuracy/references/gsm8k.yaml
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py
- tests/integration/defs/accuracy/test_llm_api_pytorch.py
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/dsa.py`:
- Around line 726-789: Add Google-style docstrings to prepare_for_draft_replay
and restore_after_draft_replay documenting their speculative replay contract:
explain when prepare_for_draft_replay returns None, identify the saved-state
contents, and describe the buffers it updates; document that
restore_after_draft_replay does nothing for None and restores each saved buffer
state. Replace the existing explanatory inline comments only where the new
docstrings cover the same interface behavior.
🪄 Autofix
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: 2a81b552-d1fa-4341-8baf-f578fac8ebec
📒 Files selected for processing (9)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/_torch/attention_backend/trtllm.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.ymltests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
🚧 Files skipped from review as they are similar to previous changes (8)
- tensorrt_llm/_torch/pyexecutor/_util.py
- tests/integration/defs/accuracy/test_llm_api_pytorch.py
- tests/integration/test_lists/test-db/l0_dgx_b200.yml
- tensorrt_llm/_torch/attention_backend/trtllm.py
- tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
- tensorrt_llm/_torch/speculative/interface.py
- tests/integration/defs/accuracy/references/gsm8k.yaml
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py
| def prepare_for_draft_replay(self) -> dict | None: | ||
| if (self.kv_cache_manager is None | ||
| or not hasattr(self.kv_cache_manager, "index_head_dim")): | ||
| return None | ||
|
|
||
| saved = { | ||
| "host_indexer_k_cache_block_offsets": | ||
| self.host_indexer_k_cache_block_offsets.clone(), | ||
| "indexer_k_cache_block_offsets": | ||
| self.indexer_k_cache_block_offsets.clone(), | ||
| "host_slot_mapping_fp8": | ||
| self.host_slot_mapping_fp8.clone(), | ||
| "host_slot_mapping_scale": | ||
| self.host_slot_mapping_scale.clone(), | ||
| "slot_mapping_fp8": | ||
| self.slot_mapping_fp8.clone(), | ||
| "slot_mapping_scale": | ||
| self.slot_mapping_scale.clone(), | ||
| } | ||
|
|
||
| # Derive pool indices from the draft manager's encoded block | ||
| # offsets (via _get_pool_block_indices) instead of using raw block | ||
| # IDs. With host cache offload, block IDs can exceed | ||
| # blocks_in_primary_pool after offload swaps (the block keeps its | ||
| # original high ID even though its memory now lives in the primary | ||
| # GPU pool). Using raw block IDs as pool indices causes OOB access | ||
| # in the indexer k-cache buffers. _get_pool_block_indices correctly | ||
| # decodes memPoolBlockIndex from the C++ encoded offsets. | ||
| # Note: kv_cache_manager was already swapped to draft above | ||
| # in prepare_attn_metadata_for_draft_replay() in _torch/speculative/interface.py | ||
| pool_indices = self._get_pool_block_indices() | ||
| num_blocks = pool_indices.shape[1] | ||
| self.host_indexer_k_cache_block_offsets[:self.num_seqs, : | ||
| num_blocks].copy_(pool_indices) | ||
| self.indexer_k_cache_block_offsets[:self.num_seqs].copy_( | ||
| self.host_indexer_k_cache_block_offsets[:self.num_seqs], | ||
| non_blocking=True, | ||
| ) | ||
| # Safety clamp: sanitize stale padding entries beyond num_seqs | ||
| # that may contain negative or out-of-range values, matching the | ||
| # regular DSA prepare() flow. | ||
| self.indexer_k_cache_block_offsets.clamp_(min=0) | ||
| Indexer.recompute_slot_mappings(self) | ||
|
|
||
| return saved | ||
|
|
||
| def restore_after_draft_replay(self, saved_state: dict | None) -> None: | ||
| if saved_state is None: | ||
| return | ||
|
|
||
| self.host_indexer_k_cache_block_offsets.copy_( | ||
| saved_state["host_indexer_k_cache_block_offsets"], | ||
| non_blocking=True, | ||
| ) | ||
| self.indexer_k_cache_block_offsets.copy_( | ||
| saved_state["indexer_k_cache_block_offsets"], | ||
| non_blocking=True, | ||
| ) | ||
| self.host_slot_mapping_fp8.copy_(saved_state["host_slot_mapping_fp8"]) | ||
| self.host_slot_mapping_scale.copy_( | ||
| saved_state["host_slot_mapping_scale"]) | ||
| self.slot_mapping_fp8.copy_(saved_state["slot_mapping_fp8"]) | ||
| self.slot_mapping_scale.copy_(saved_state["slot_mapping_scale"]) | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the draft replay hook contract.
prepare_for_draft_replay() and restore_after_draft_replay() are called by the speculative replay interface. Add Google-style docstrings. Define the None behavior, saved-state contents, and the buffers that each method changes or restores.
As per coding guidelines, “Use docstrings rather than comments for externally usable interfaces, Google-style docstrings for classes and functions.”
🤖 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 726 - 789,
Add Google-style docstrings to prepare_for_draft_replay and
restore_after_draft_replay documenting their speculative replay contract:
explain when prepare_for_draft_replay returns None, identify the saved-state
contents, and describe the buffers it updates; document that
restore_after_draft_replay does nothing for None and restores each saved buffer
state. Replace the existing explanatory inline comments only where the new
docstrings cover the same interface behavior.
Source: Coding guidelines
There was a problem hiding this comment.
The hook contract and None behavior are already documented on the base methods in TrtllmAttentionMetadata. These overrides follow that contract, while the existing inline comments document the DSA-specific mutation details and rationale. Repeating the contract here would be redundant, so I’d prefer to keep the documentation at the interface level.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
base_file=$(fd -a '^trtllm\.py$' tensorrt_llm/_torch/attention_backend | head -n 1)
if [[ -z "${base_file}" ]]; then
echo "Could not locate trtllm.py"
exit 1
fi
ast-grep outline "$base_file" --items all --match 'TrtllmAttentionMetadata|prepare_for_draft_replay|restore_after_draft_replay'
rg -n -A18 -B4 \
'def prepare_for_draft_replay\(|def restore_after_draft_replay\(' \
"$base_file"Length of output: 4606
@nvxuanyuc, I need to correct my previous reply.
The base methods have short docstrings, but they do not define the None behavior, saved-state contents, or DSA buffer mutations. The shell result shows only the general replay lifecycle descriptions.
Therefore, the original documentation concern remains valid. I understand your preference to keep the interface documentation centralized. If that is the intended project convention, the maintainers can choose not to add override docstrings.
You are interacting with an AI system.
|
/bot run --disable-fail-fast |
|
PR_Github #64453 [ run ] triggered by Bot. Commit: |
|
@BowenFu Thanks for the review. |
|
PR_Github #64453 [ run ] completed with state
|
…t kv cache (TEP) Signed-off-by: Xuanyu Chen <xuanyuc@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #64723 [ run ] triggered by Bot. Commit: |
7 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.
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.