[https://nvbugs/6513132][fix] DSA: rebuild token-to-request map inside the MTP draft loop - #16925
Conversation
79cb3be to
8169762
Compare
|
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:
WalkthroughThe DSA attention backend now builds token-to-request indices from current sequence lengths. Runtime KV-length updates refresh this mapping, and DeepSeek V4 token-position computation uses the refreshed buffer. ChangesDSA runtime mapping
Estimated code review effort: 2 (Simple) | ~15 minutes Sequence Diagram(s)sequenceDiagram
participant on_update_kv_lens
participant build_req_idx_per_token
participant _compute_token_positions
on_update_kv_lens->>build_req_idx_per_token: pass current sequence lengths and token count
build_req_idx_per_token-->>on_update_kv_lens: return request-index mapping
on_update_kv_lens->>_compute_token_positions: provide refreshed mapping
_compute_token_positions-->>on_update_kv_lens: compute token positions
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.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)
787-789: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid the redundant cast in
torch.searchsorted. Useout_int32=Trueso the result lands directly in theint32destination and skips the extra temporary/conversion in this hot path.Proposed optimization
self.req_idx_per_token[:self.num_tokens] = torch.searchsorted( cu_seq_lens, token_idx, - right=True).to(self.req_idx_per_token.dtype) + right=True, out_int32=True)🤖 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 787 - 789, Update the torch.searchsorted call in the request-index assignment to request int32 output directly with out_int32=True, then remove the trailing dtype conversion while preserving the existing cu_seq_lens, token_idx, and destination slice behavior.
🤖 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/dsa.py`:
- Around line 787-789: Update the torch.searchsorted call in the request-index
assignment to request int32 output directly with out_int32=True, then remove the
trailing dtype conversion while preserving the existing cu_seq_lens, token_idx,
and destination slice behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d777c055-8fee-478b-a7b6-ca9ed4a763c2
📒 Files selected for processing (1)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py
zheyuf
changed the title
req_idx_per_token is built once per engine step in prepare_for_indices_conversion(), from the target forward's seq_lens, which for a one-model speculative generation batch are (max_draft_len + 1) tokens per request. Eagle3OneModelWorker.forward rewrites the batch layout to one token per request inside its draft loop and calls update_for_spec_dec() -> on_update_kv_lens(), which recomputes seq_starts from the fresh all-ones seq_lens but reuses the map from prepare(). Nothing rebuilds it, so from the second draft iteration onward token j resolves to request j // (max_draft_len + 1) instead of request j; only row 0 is ever correct. convert_req_index_to_global() then resolves the sparse top-k indices through another request's block table, and Indexer._update_k_cache() scatters the draft token's indexer K through the same slot mapping into another request's pages. The damage is confined to the draft layer's own indexer state, so generated output is unaffected and the symptom is purely lost acceptance length. Because attention DP splits requests across ranks, the misrouted fraction is (B - 1) / B with B the per-rank decode batch, which is why acceptance falls monotonically with concurrency. max_draft_len = 1 is immune: the loop body runs once, so the mutation lands after the last draft forward. Affects the two architectures on the base DSA metadata class, DeepseekV32ForCausalLM and GlmMoeDsaForCausalLM. DeepseekV4TrtllmAttentionMetadata already rebuilds this map in its own on_update_kv_lens() override and is unaffected; it calls super() first, so the map is written twice with identical values. Rebuild from the current seq_lens with a device-side searchsorted, exactly equivalent to prepare()'s repeat_interleave when seq_lens is unchanged, so it is a no-op outside the draft loop. CUDA-graph safe: the write targets the existing static req_idx_per_token buffer, and the DeepSeek-V4 override already runs searchsorted on the same captured path. Measured on GLM-5.2 NVFP4, GB300 1p1d disagg, DEP8, c128, MTP k=7, T=1 with rejection sampling, 55-57k-token contexts: acceptance length 3.4916 -> 5.5222 (+58%), generation wall-clock -35%. Offline on 1xGB200 tp4/ep4 with matched prompt sets: c8 3.500 -> 3.943, c64 2.789 -> 3.909, turning a -20.3% slope into -0.9%. A max_draft_len = 1 control is unchanged (1.8433 -> 1.8402). Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
8169762 to
2cfeaab
Compare
There was a problem hiding this comment.
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/attention_backend/sparse/dsa.py (1)
790-811: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winWarm the
next_n == 1radix-filter variant.
next_nis a compile-key dimension. For compressed indexers, this guard returns when engine warmup supplies multi-token MTPnext_n, but Line 2906 routes later per-token draft iterations (next_n == 1) to the CuTe radix op. The first live draft request can therefore still incur JIT compilation. Warm1separately; for uncompressed MTP, warm both variants.Proposed fix
- if self._indexer_compress_ratio > 1 and next_n > 1: - return try: from ...custom_ops.cute_dsl_custom_ops import \ warmup_cute_dsl_radix_topk_decode except ImportError: return - warmup_cute_dsl_radix_topk_decode( - top_k=int(top_k), - num_cols=int(self.get_indexer_max_seq_len()), - next_n=next_n, - dtype=_INDEXER_LOGITS_DTYPE, - num_sms=self.num_sms, - ) + warmup_next_ns = (1,) if self._indexer_compress_ratio > 1 else ( + (next_n,) if next_n == 1 else (next_n, 1)) + for warmup_next_n in warmup_next_ns: + warmup_cute_dsl_radix_topk_decode( + top_k=int(top_k), + num_cols=int(self.get_indexer_max_seq_len()), + next_n=warmup_next_n, + dtype=_INDEXER_LOGITS_DTYPE, + num_sms=self.num_sms, + )🤖 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 790 - 811, Update the warmup logic around the guard in the radix-filter decode path so compressed indexers always warm the next_n == 1 variant, even when engine warmup receives multi-token MTP; for uncompressed MTP, warm both the supplied next_n variant and next_n == 1 when they differ. Preserve the existing ImportError handling and avoid returning before the required per-token warmup.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 790-811: Update the warmup logic around the guard in the
radix-filter decode path so compressed indexers always warm the next_n == 1
variant, even when engine warmup receives multi-token MTP; for uncompressed MTP,
warm both the supplied next_n variant and next_n == 1 when they differ. Preserve
the existing ImportError handling and avoid returning before the required
per-token warmup.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0368bd41-346a-4219-aa37-b6267c5ff79b
📒 Files selected for processing (1)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py
|
PR_Github #62727 [ run ] completed with state |
|
Hi @yunruis @pengbowang-nv, could you help review this PR to unblock GLM5.2 Agentperf submission? CI has passed. Thanks! |
|
Could I understand it is what your PR do? Before: After repair: |
|
@yunruis Yes, exactly — your walkthrough is precisely what the PR fixes. |
There was a problem hiding this comment.
Checked the regression axis: searchsorted(cumsum(seq_lens), arange(num_tokens), right=True) is exactly prepare_for_indices_conversion()'s repeat_interleave (dsa.py:1694) whenever seq_lens is unchanged — including zero-length rows — so the non-MTP DSA path is unaffected. DeepSeek-V4 is bit-identical: its override calls super().on_update_kv_lens() and then rewrites the same buffer with the identical searchsorted in _compute_token_positions() (deepseek_v4.py:913-918).
Not blocking on these, but worth a follow-up: the map is now built three ways (dsa.py:1694 CPU, dsa.py:841 device, deepseek_v4.py:913 device) — V4's copy is now redundant work every step and a shared helper would stop them drifting. And there is no unit test pinning the rebuilt map against the draft-loop layout; the benchmark deltas are convincing but they will not catch a regression here.
|
@BowenFu thanks a lot for the review. I agree that it worth a follow-up and let me do the fix in this PR. |
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 `@tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py`:
- Around line 71-90: Add a metadata-level test that constructs the relevant DSA
attention metadata, updates device seq_lens, and invokes
DSAtrtllmAttentionMetadata.on_update_kv_lens(), including the DeepSeek-V4
override path. Assert the rebuilt request indices and derived token positions,
and exercise CUDA graph replay when supported; retain
test_draft_loop_transition() as the buffer-level regression test.
🪄 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: 13da1e91-e24d-4104-a627-4a6a3cf858e7
📒 Files selected for processing (3)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/attention_backend/sparse/dsa.pytests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py
…ween DSA and DeepSeek-V4 - Extract the searchsorted rebuild into build_req_idx_per_token() and make it unconditional on num_tokens > 0 (the map depends only on seq_lens, not on the kv cache manager), so subclasses can rely on the buffer being current after super().on_update_kv_lens(). - DeepSeek-V4's _compute_token_positions() now reuses the rebuilt buffer instead of recomputing an identical searchsorted every step. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
4e51ba3 to
fae96ce
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63646 [ run ] triggered by Bot. Commit: |
|
PR_Github #63646 [ run ] completed with state
|
|
Hi @zheyuf — your PR also covers a fix for one of our NVBugs (https://nvbugs/6463964): MTP acceptance length on GLM-5.2 (DSA + one-model MTP) degrades as batch size grows, which turned out to be the same stale token→request map you're fixing here. I verified your PR resolves it locally — GLM-5.2-NVFP4, TP=4, 80 MT-bench prompts, greedy, max_draft_len=4: batch=8 acceptance length goes from 3.283 → 3.739, matching batch=1 (3.734). Same with CUDA graphs on. I had a duplicate fix open (#17240) and have closed it in favor of yours, and assigned 6463964 to you as well since this PR is what will resolve it. |
|
Thanks @zhaoyangwang-nvidia, my PR is exactly targeting the same problem that MTP acceptance length on GLM-5.2 degrades as batch size grows: here are my two NVBugs https://nvbugspro.nvidia.com/bug/6513132, https://nvbugspro.nvidia.com/bug/6513093. Sorry that I was unaware that you were also working on this problem, otherwise I should have contacted you to prevent our duplicated effort. Seems like our NVBugs are pretty duplicated 😂 |
…unit tests Two tests covering the review follow-up: - build_req_idx_per_token equals the host repeat_interleave build across layouts including zero-length rows, pinning the device and host builders against drift. - on_update_kv_lens() rebuilds the map for the draft-loop layout on a bare metadata instance; fails on pre-fix code where the stale target-forward slice misattributes every draft token to request 0. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
|
/bot run |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py (1)
41-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd complete function annotations.
Lines 41, 61, and 72 define functions without complete type annotations. Add return annotations to all functions and parameter annotations to the test parameters.
Proposed change
-def _host_reference(seq_lens: torch.Tensor) -> torch.Tensor: +def _host_reference(seq_lens: torch.Tensor) -> torch.Tensor: -def test_matches_host_repeat_interleave(seq_lens, device): +def test_matches_host_repeat_interleave( + seq_lens: list[int], device: str +) -> None: -def test_on_update_kv_lens_rebuilds_stale_map(): +def test_on_update_kv_lens_rebuilds_stale_map() -> None:As per coding guidelines, “Annotate every function.”
Also applies to: 61-61, 72-72
🤖 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_req_idx_per_token.py` around lines 41 - 42, Complete the type annotations for _host_reference and the other functions at the referenced locations, adding explicit return types and annotating all test function parameters. Use the existing tensor types where applicable and ensure every function in the file has a complete signature.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.
Nitpick comments:
In `@tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py`:
- Around line 41-42: Complete the type annotations for _host_reference and the
other functions at the referenced locations, adding explicit return types and
annotating all test function parameters. Use the existing tensor types where
applicable and ensure every function in the file has a complete signature.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2a092819-1e8a-4c9c-b551-72ea65c1dfca
📒 Files selected for processing (1)
tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py
|
PR_Github #63858 [ run ] triggered by Bot. Commit: |
Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
|
@BowenFu I have added the shared helper and the unit test pinning the rebuilt map against the draft-loop layout according to your comments. Thanks. |
|
/bot run --disable-fail-fast |
|
PR_Github #63890 [ run ] triggered by Bot. Commit: |
|
PR_Github #63858 [ run ] completed with state |
|
PR_Github #63890 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63978 [ run ] triggered by Bot. Commit: |
|
PR_Github #63978 [ run ] completed with state |
|
Still need a stamp from @NVIDIA/trt-llm-torch-attention-devs to merge this PR. I have added the shared helper and the unit test pinning the rebuilt map against the draft-loop layout according to BowenFu's comments. CI passed again. Kindly @yunruis @pengbowang-nv for review. 🙏 |
Cherry-pick of NVIDIA#16925 (d2ecec4, c7a7ab0, 79cb3be) squashed into one commit. req_idx_per_token is built once per engine step in prepare_for_indices_conversion() from the target forward's seq_lens, which for an MTP generation batch are (max_draft_len + 1) tokens per request. MTPEagleWorker.forward rewrites seq_lens to one token per request before calling update_for_spec_dec(), but on_update_kv_lens() deliberately reuses the map from prepare(). Nothing rebuilds it, so from the second draft iteration onward token j resolves to request j // (max_draft_len + 1) instead of request j; only row 0 is ever correct. Both directions are affected. convert_req_index_to_global() resolves the sparse top-k indices through the wrong request's block table, and Indexer._update_k_cache() scatters the draft token's indexer K through the same slot mapping into another request's pages. The damage is confined to the MTP layer's own indexer state, so generated output is unaffected and the symptom is purely a loss of acceptance length. Because attention DP splits requests across ranks, the misrouted fraction is (B - 1) / B with B the per-rank decode batch, which is why acceptance falls monotonically with concurrency and why max_draft_len = 1 is immune (the draft loop never mutates the layout in that case). DeepseekV4TrtllmAttentionMetadata already rebuilds this map in its on_update_kv_lens() override; models on the base DSA metadata class (deepseek_v32, glm_moe_dsa) do not. Rebuild it from the current seq_lens with a device-side searchsorted, which is CUDA-graph safe and a no-op on the target forward. Also carries the PR's regression tests (test_on_update_kv_lens_rebuilds_req_idx_in_draft_loop and test_on_update_kv_lens_matches_prepare_on_target_forward) plus the MockMetadata flag defaults they need. Branch adaptation on top of the upstream PR: this branch's on_update_kv_lens() also reads in_mtp_draft_loop for the cross-step indexer top-k reuse, which MockMetadata does not set, so the two new tests would raise AttributeError before reaching the rebuild. Default it to False in the mock alongside the topk flags. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
|
Heads-up: the unit test added here appears to break premerge CI on current main (merge skew with the Failure (seen in PR #16666's premerge, pipeline Mechanism: Suggested fix (either works):
Happy to send a one-line PR for the test-side fix if that helps. |
test_on_update_kv_lens_rebuilds_stale_map builds its metadata with object.__new__, bypassing __init__ where in_mtp_draft_loop is initialized. Since NVIDIA#16925 made on_update_kv_lens() read that flag, the test fails with AttributeError on every pre-merge run. Set the __init__ default on the stub. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
test_on_update_kv_lens_rebuilds_stale_map builds its metadata with object.__new__, bypassing __init__ where in_mtp_draft_loop is initialized. Since NVIDIA#16925 made on_update_kv_lens() read that flag, the test fails with AttributeError on every pre-merge run. Set the __init__ default on the stub. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
test_on_update_kv_lens_rebuilds_stale_map builds its metadata with object.__new__, bypassing __init__ where in_mtp_draft_loop is initialized. Since NVIDIA#16925 made on_update_kv_lens() read that flag, the test fails with AttributeError on every pre-merge run. Set the __init__ default on the stub. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
test_on_update_kv_lens_rebuilds_stale_map builds its metadata with object.__new__, bypassing __init__ where in_mtp_draft_loop is initialized. Since NVIDIA#16925 made on_update_kv_lens() read that flag, the test fails with AttributeError on every pre-merge run. Set the __init__ default on the stub. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
Thanks @longcheng-nv, seems like the test side fix is merged in this PR: #17416 and the pre-merge CI should be unblocked now. |
7 participants
The problem (for both GLM5.2 and DSV3.2):
(1) MTP acceptance rate has regression when batch size is large: https://nvbugspro.nvidia.com/bug/6513093
(2) There is a gap of AL between TRTLLM and vLLM:https://nvbugspro.nvidia.com/bug/6513132
The root cause is the stale token→request map in DSA.py. req_idx_per_token is initially built from the target forward layout, where each request contributes max_draft_len + 1 tokens. During the one-model draft loop, the layout changes to one token per request, but on_update_kv_lens() recomputes seq_starts while reusing the old map (the map is stale). This corrupts sparse-index reads through the wrong request’s block table and also corrupts Indexer-K writes into the wrong request’s KV pages.
Fix
Rebuild the map from the current
seq_lenswith a device-sidesearchsorted.Deepseek v4 doesn't have such acceptance rate regression as
DeepseekV4TrtllmAttentionMetadataalready rebuilt this map in its ownon_update_kv_lens()override, so this is a port of existing behaviour to the base class, not new logic.NVBugs fixed by this PR
https://nvbugspro.nvidia.com/bug/6513132, https://nvbugspro.nvidia.com/bug/6513093
Results — before / after
All A/B pairs are the same image and harness, differing only by this change.
GLM-5.2 NVFP4 — offline
trtllm-bench, 1×GB200 tp4/ep4 + ADP, MTP k=7, greedyDeepSeek-V3.2 NVFP4 — the second affected model, same harness, k=7 greedy
Slope c8 → c128: −6.3% → +1.5% (flat). Same qualitative signature on an independent model.
The above proves that (1) MTP acceptance rate has regression when batch size is large is fixed.
GLM-5.2 — depth sweep at c8, against a same-session vLLM reference
The above proves that (2)There is a gap of AL between TRTLLM and vLLM is fixed.
GPQA-Diamond, 198 questions — accuracy is unaffected, as the mechanism requires
GB300 1p1d disagg, c128, MTP k=7, T=1 + rejection sampling,
max_seq_len 400000:What models get fixed
DeepseekV32ForCausalLMDSAtrtllmAttentionMetadata(patched)max_draft_len ≥ 2GlmMoeDsaForCausalLMDev Engineer Review
build_req_idx_per_token()indsa.py.torch.searchsorted.on_update_kv_lens()to refresh the map before slot mapping.QA Engineer Review
tests/integration/test_lists,test-db, orqa.