[None][feat] Support DSA MTP indexer top-k sharing#15806
Conversation
📝 WalkthroughWalkthroughAdds an ChangesMTP Top-K Reuse for DSA Indexer
Sequence Diagram(s)sequenceDiagram
participant Eagle3Worker as Eagle3OneModelWorker
participant Metadata as DSAtrtllmAttentionMetadata
participant Indexer as Indexer.sparse_attn_indexer
Eagle3Worker->>Metadata: set_in_mtp_draft_loop(true)
loop draft iterations
Eagle3Worker->>Metadata: set_skip_topk(iteration > 0)
Eagle3Worker->>Indexer: sparse_attn_indexer(metadata)
alt indexer_skip_topk and shared_topk_indices available
Indexer->>Metadata: read shared_topk_indices
Indexer-->>Eagle3Worker: reuse topk_indices_buffer
else compute new Top-K
Indexer->>Indexer: compute Top-K
Indexer->>Metadata: write shared_topk_indices
end
end
Eagle3Worker->>Metadata: set_skip_topk(false), set_in_mtp_draft_loop(false)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~30 minutes 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: 3
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)
1236-1240: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate and document the new metadata setter API.
These setters are called outside this file from Eagle3, so add
-> Noneand short Google-style docstrings.As per coding guidelines, “Always annotate functions with return types” and “For interfaces used outside a file, prefer docstrings over comments.”
Proposed cleanup
- def set_skip_topk(self, skip: bool): + def set_skip_topk(self, skip: bool) -> None: + """Set whether the DSA indexer should reuse shared MTP Top-K.""" self.indexer_skip_topk = skip - def set_in_mtp_draft_loop(self, active: bool): + def set_in_mtp_draft_loop(self, active: bool) -> None: + """Set whether the current forward is inside an MTP draft loop.""" self.in_mtp_draft_loop = active🤖 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 1236 - 1240, The new metadata setter APIs in dsa.py are missing return annotations and documentation; update set_skip_topk and set_in_mtp_draft_loop to explicitly return None and add short Google-style docstrings. Keep the change localized to these setter methods so external callers from Eagle3 can rely on the documented interface and annotated signatures.Source: Coding guidelines
tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py (1)
2628-2630: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for the public config wiring.
Coverage is sufficient for the raw stash/reuse path, but insufficient for
DeepSeekSparseAttentionConfig.index_share_for_mtp_iteration -> DSAParams.mtp_index_share -> Indexer.mtp_index_share. Prefer setting the config field beforecreate_indexer()instead of mutatingindexer.mtp_index_sharedirectly, or add a small conversion test in this file.As per path instructions, test feedback should state whether coverage is sufficient or insufficient and suggest concrete test locations.
🤖 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/dsa/test_dsa_indexer.py` around lines 2628 - 2630, Coverage for the public config wiring is insufficient: the test currently mutates Indexer.mtp_index_share directly instead of verifying the DeepSeekSparseAttentionConfig.index_share_for_mtp_iteration -> DSAParams.mtp_index_share -> Indexer.mtp_index_share path. Update the existing test around create_indexer() to set index_share_for_mtp_iteration on sparse_attn_config before creating the indexer, or add a small conversion-focused test in this file that asserts the config field is propagated into Indexer.mtp_index_share.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/dsa.py`:
- Around line 1071-1072: The Top-K reuse state in DSA is currently stored in a
single shared slot on metadata, so one layer can overwrite another layer’s
indices when multiple layer-specific Indexer instances call
sparse_attn_indexer(). Update the shared_topk_indices handling in
sparse_attn_indexer and the related MTP reuse path to scope the stash by
self.layer_idx (or the local layer identity) so each layer reuses its own Top-K,
or add an explicit guard/assertion that only one DSA indexer is supported. Make
the change consistently wherever shared_topk_indices is read or written.
In `@tensorrt_llm/_torch/speculative/eagle3.py`:
- Around line 779-787: Reset the MTP Top-K state even when the draft loop fails
by wrapping the `draft_kv_cache_context` loop in a try/finally and clearing the
flags in the finally block. In `eagle3.py`, update the `self.is_mtp_eagle` /
`attn_metadata.set_skip_topk` handling inside the draft-loop section so
`set_in_mtp_draft_loop(False)` and `set_skip_topk(False)` are always restored
after `runtime_draft_len` iteration, regardless of exceptions; apply the same
cleanup pattern in the other draft-loop block referenced by the same
`reuse_mtp_topk` logic.
In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py`:
- Around line 2701-2704: The sparse_attn_indexer custom Top-K test is swallowing
all failures by catching a broad Exception, which can hide real regressions.
Update the test around sparse_attn_indexer in test_dsa_indexer so it only skips
for the specific “custom topk not available” condition, or remove the try/except
entirely and let unexpected errors fail the test. Keep the logic centered on the
sparse_attn_indexer call and the use_custom_topk=True path.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 1236-1240: The new metadata setter APIs in dsa.py are missing
return annotations and documentation; update set_skip_topk and
set_in_mtp_draft_loop to explicitly return None and add short Google-style
docstrings. Keep the change localized to these setter methods so external
callers from Eagle3 can rely on the documented interface and annotated
signatures.
In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py`:
- Around line 2628-2630: Coverage for the public config wiring is insufficient:
the test currently mutates Indexer.mtp_index_share directly instead of verifying
the DeepSeekSparseAttentionConfig.index_share_for_mtp_iteration ->
DSAParams.mtp_index_share -> Indexer.mtp_index_share path. Update the existing
test around create_indexer() to set index_share_for_mtp_iteration on
sparse_attn_config before creating the indexer, or add a small
conversion-focused test in this file that asserts the config field is propagated
into Indexer.mtp_index_share.
🪄 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: 874ef455-3105-4c9b-bb7e-8025098a3bd4
📒 Files selected for processing (4)
tensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/_torch/speculative/eagle3.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
| # MTP cross-step indexer Top-K reuse state. | ||
| self.shared_topk_indices: Optional[torch.Tensor] = None |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Scope shared Top-K state per indexer layer.
metadata.shared_topk_indices is a single slot, but sparse_attn_indexer() is called from layer-specific Indexer instances. With more than one DSA indexer in the MTP path, step 0’s last layer overwrites the stash and every layer in later draft steps reuses that layer’s Top-K. Key this state by self.layer_idx/local layer, or assert this path only supports one DSA indexer.
Also applies to: 2544-2552, 2798-2813
🤖 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 1071 -
1072, The Top-K reuse state in DSA is currently stored in a single shared slot
on metadata, so one layer can overwrite another layer’s indices when multiple
layer-specific Indexer instances call sparse_attn_indexer(). Update the
shared_topk_indices handling in sparse_attn_indexer and the related MTP reuse
path to scope the stash by self.layer_idx (or the local layer identity) so each
layer reuses its own Top-K, or add an explicit guard/assertion that only one DSA
indexer is supported. Make the change consistently wherever shared_topk_indices
is read or written.
| reuse_mtp_topk = (self.is_mtp_eagle | ||
| and hasattr(attn_metadata, "set_skip_topk")) | ||
| if reuse_mtp_topk: | ||
| attn_metadata.set_in_mtp_draft_loop(True) | ||
|
|
||
| with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): | ||
| for i in range(runtime_draft_len): | ||
| if reuse_mtp_topk: | ||
| attn_metadata.set_skip_topk(i > 0) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reset MTP Top-K flags in a finally block.
If any draft-loop operation raises, the metadata can remain in skip/reuse mode and poison the next forward with stale Top-K state.
Proposed structure
if reuse_mtp_topk:
attn_metadata.set_in_mtp_draft_loop(True)
- with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager):
- for i in range(runtime_draft_len):
+ try:
+ with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager):
+ for i in range(runtime_draft_len):
if reuse_mtp_topk:
attn_metadata.set_skip_topk(i > 0)
...
- next_draft_tokens = torch.stack(next_draft_tokens, dim=1)
-
- if reuse_mtp_topk:
- attn_metadata.set_skip_topk(False)
- attn_metadata.set_in_mtp_draft_loop(False)
+ next_draft_tokens = torch.stack(next_draft_tokens, dim=1)
+ finally:
+ if reuse_mtp_topk:
+ attn_metadata.set_skip_topk(False)
+ attn_metadata.set_in_mtp_draft_loop(False)Also applies to: 945-947
🤖 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/speculative/eagle3.py` around lines 779 - 787, Reset the
MTP Top-K state even when the draft loop fails by wrapping the
`draft_kv_cache_context` loop in a try/finally and clearing the flags in the
finally block. In `eagle3.py`, update the `self.is_mtp_eagle` /
`attn_metadata.set_skip_topk` handling inside the draft-loop section so
`set_in_mtp_draft_loop(False)` and `set_skip_topk(False)` are always restored
after `runtime_draft_len` iteration, regardless of exceptions; apply the same
cleanup pattern in the other draft-loop block referenced by the same
`reuse_mtp_topk` logic.
| try: | ||
| topk0 = indexer.sparse_attn_indexer(meta0, h, q, k_fp8, k_scale, w, use_custom_topk=True) | ||
| except Exception as e: | ||
| pytest.skip(f"Custom topk not available: {e}") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Don’t hide Top-K failures behind a broad skip.
Catching Exception here will skip on shape bugs, stale metadata, or kernel regressions. Let the test fail, or catch only the specific unavailable-custom-op error.
As per coding guidelines, “catch specific exceptions, not bare except.”
Preferred simplification
- try:
- topk0 = indexer.sparse_attn_indexer(meta0, h, q, k_fp8, k_scale, w, use_custom_topk=True)
- except Exception as e:
- pytest.skip(f"Custom topk not available: {e}")
+ topk0 = indexer.sparse_attn_indexer(meta0, h, q, k_fp8, k_scale, w, use_custom_topk=True)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| topk0 = indexer.sparse_attn_indexer(meta0, h, q, k_fp8, k_scale, w, use_custom_topk=True) | |
| except Exception as e: | |
| pytest.skip(f"Custom topk not available: {e}") | |
| topk0 = indexer.sparse_attn_indexer(meta0, h, q, k_fp8, k_scale, w, use_custom_topk=True) |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 2703-2703: Do not catch blind exception: Exception
(BLE001)
🤖 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/dsa/test_dsa_indexer.py` around lines
2701 - 2704, The sparse_attn_indexer custom Top-K test is swallowing all
failures by catching a broad Exception, which can hide real regressions. Update
the test around sparse_attn_indexer in test_dsa_indexer so it only skips for the
specific “custom topk not available” condition, or remove the try/except
entirely and let unexpected errors fail the test. Keep the logic centered on the
sparse_attn_indexer call and the use_custom_topk=True path.
Sources: Coding guidelines, Linters/SAST tools
ff5300d to
cdedd71
Compare
cdedd71 to
3a21336
Compare
nvxuanyuc
added
the
api-compatible
|
/bot run --disable-fail-fast |
|
PR_Github #58604 [ run ] triggered by Bot. Commit: |
|
PR_Github #58604 [ run ] completed with state
|
3a21336 to
9cadc26
Compare
9cadc26 to
2a518bb
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58845 [ run ] triggered by Bot. Commit: |
|
PR_Github #58845 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58953 [ run ] triggered by Bot. Commit: |
|
PR_Github #58953 [ run ] completed with state
|
|
PR_Github #60455 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60508 [ run ] triggered by Bot. Commit: |
|
PR_Github #60508 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60607 [ run ] triggered by Bot. Commit: |
|
PR_Github #60607 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60733 [ run ] triggered by Bot. Commit: |
| if num_generations > 0: | ||
| next_n = num_gen_tokens // num_generations | ||
| rows = topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + | ||
| num_gen_tokens][next_n - 1::next_n] |
There was a problem hiding this comment.
This stashes topk_indices_buffer[...][next_n - 1::next_n], i.e. each gen request's last padded row. Every gen request is padded to next_n = draft_len+1 rows for cuda-graph, so any row past num_accepted sits on the target's rejected-branch continuation. Under partial acceptance (the common case at k≥7), the reuse steps at i≥1 then attend a top-k that was selected by a rejected-branch query. After the kv-lens rewind those indices also point at stale rejected-branch KV, so every reuse step of a partially accepted round is corrupted.
vLLM handles this by compacting to each request's last accepted token before the reuse steps: after set_skip_topk(True) it calls compact_topk_indices(token_indices_to_sample), where token_indices_to_sample are the last-accepted rows.
- vllm/v1/spec_decode/llm_base_proposer.py#L598-L603
- vllm/model_executor/models/deepseek_mtp.py#L196-L209
The fix here is to plumb num_accepted_tokens into the metadata and stash base + (num_accepted-1).clamp(0, next_n-1) per gen request. The context stash stays as is, and this is cuda-graph safe.
The same wrong-row pattern also shows up in the new _enable_heuristic_topk seed at line 2846. Lower severity there since it only warms up the next top-k rather than freezing it, but it wants the same num_accepted-1 selection.
See an example patch: http://nv/jvo
There was a problem hiding this comment.
Thanks! Have added the fix for partial acceptance indexing and factored into a helper that heuristic topk path can reuse.
|
PR_Github #60733 [ run ] completed with state
|
Squashed from NVIDIA#15806 (single commit, PR head 25f9efb). Gated off by default (sparse_attention_config.index_share_for_mtp_iteration). Signed-off-by: Iman Tabrizian <10105175+Tabrizian@users.noreply.github.com>
… row PR NVIDIA#15806's cross-step DSA indexer top-k stash saved topk_indices_buffer[next_n-1::next_n] (each gen request's last padded row). TRT pads every gen request to next_n=(draft_len+1) rows for cuda-graph capture; rows beyond num_accepted lie on the target's rejected-branch continuation. Under partial acceptance (the majority of rounds at k>=7) the frozen indices are then a rejected-branch query's top-k, and after the kv-lens rewind the reuse steps attend in-loop draft KV -- a KVShare violation that corrupts every reuse step of every partially accepted round. vLLM/SGLang compact the reuse buffer at each request's last ACCEPTED token. Mirror that: pass num_accepted_tokens into the metadata and stash the row at start + (num_accepted-1) per gen request (context stash unchanged). ~15 lines, cuda-graph safe. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> (cherry picked from commit f074a2f98b3d60f60448857d2f0a1670a0eca6cb) Signed-off-by: Iman Tabrizian <10105175+Tabrizian@users.noreply.github.com>
Signed-off-by: Xuanyu Chen <xuanyuc@nvidia.com>
25f9efb to
e829211
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61104 [ run ] triggered by Bot. Commit: |
|
PR_Github #61104 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61411 [ run ] triggered by Bot. Commit: |
|
PR_Github #61411 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61435 [ run ] triggered by Bot. Commit: |
|
PR_Github #61435 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61496 [ run ] triggered by Bot. Commit: |
|
PR_Github #61496 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61524 [ run ] triggered by Bot. Commit: |
|
PR_Github #61524 [ run ] completed with state
|
6 participants
Summary by CodeRabbit
Description
Implement GLM-5.2 IndexShare for the one-model MTP path: the indexer top-k indices are computed on the first draft step and reused for all following draft steps.
Based on #15574 for full GLM-5.2 IndexShare support - this PR covers MTP side.
A new
index_share_for_mtp_iterationfield is added toDeepSeekSparseAttentionConfig. It is resolved from the HF checkpoint config when left unset and threaded into the DSA indexer asmtp_index_share. On draft step 0 the per-request last-token top-k is stashed on the attention metadata; for draft steps > 0 the indexer skips the logits computation and the top-k selection, reusing the stashed indices via a copy instead.Test Coverage
Unit test —
tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py::test_indexer_decode_mtp_topk_reuse: verifies that draft step 0 stores each request's last-token top-k and that later draft steps reuse it bit-exactly (no recompute). Covered for both the steady-state generation case and the first generation step right after prefill, at BS 1 and 2.E2E tests — in
TestGLM52:test_nvfp4_mtp_index_share— GSM8K accuracy (GLM-5.2-NVFP4, CuteDSL MoE, MTP 3, TEP8). Measured 93.86.test_nvfp4_mtp_index_share_mtp_ar— MTP acceptance-rate guard. Measured AR 62.75%.Accuracy A/B validation (index-share in MTP is lossless).
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.