[https://nvbugs/6463964][fix] DSA: rebuild req_idx_per_token for MTP-Eagle draft steps - #17240
Conversation
…Eagle draft steps DSAtrtllmAttentionMetadata.req_idx_per_token maps each token of the flattened batch to its request index. It is built once per target forward in prepare_for_indices_conversion() and then deliberately reused, to keep repeat_interleave out of the CUDA-graph-captured region. That reuse is only valid while the batch layout is unchanged. The one-model MTP-Eagle draft loop breaks that premise: at the end of draft step 0 it collapses the batch to one token per request (_seq_lens.fill_(1) followed by on_update()), but AttentionMetadata.on_update() only refreshes num_tokens / num_ctx_tokens / num_generations -- it does not rebuild this mapping. Draft steps 1..k therefore consume the stale target layout ([0]*next_n + [1]*next_n + ...) where the correct mapping is arange(num_seqs), corrupting both the indexer-K slot mapping in on_update_kv_lens() and the top-k -> global index conversion in _rebuild_pool_view_cache(). Requests read and overwrite each other's indexer K-cache, and the damage compounds with draft depth. The mapping is trivially correct at batch_size == 1, so the defect only appears at batch > 1 with max_draft_len > 1, and draft step 0 stays exact. It is DSA-specific: non-DSA attention has no such buffer. Refresh the mapping whenever the layout is one token per request. num_tokens == num_seqs implies every seq_len is 1 (seq lens are >= 1 and sum to num_tokens), for which the mapping is exactly arange(num_seqs) -- an identity that also holds for ordinary single-token decode, so the refresh is unconditionally correct rather than a speculative-decoding special case. It is a copy_ from a preallocated static arange buffer, so it allocates nothing and preserves the graph-capture constraint that motivated caching the mapping in the first place. Placing it in the metadata rather than in the draft loop also covers the DeepSeek-V4 path, which builds the same mapping. Measured on GLM-5.2-NVFP4, TP=4, 80 MT-bench prompts, greedy, max_draft_len=4, mtp_eagle_one_model=True. Acceptance length at batch=8 goes from 3.283 to 3.766 (batch=1 is 3.768 and is unchanged), and the conditional acceptance rate per draft depth goes from 0.880/0.775/0.670/0.596 back to the flat 0.868/0.864/0.825/0.834 seen at batch=1. Independently reproduced with CUDA graphs enabled: 3.233 -> 3.758. max_draft_len=1 is unaffected, as expected, since the loop never reaches step 1. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
WalkthroughDSA now refreshes per-token request indices for one-token-per-sequence batches. It uses a persistent int32 arange buffer that remains safe during CUDA graph capture. The refresh runs before dependent pool views are invalidated. ChangesDSA request-index refresh
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 (2)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py (2)
1715-1746: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the required return annotation.
Declare the helper as
def _refresh_req_idx_per_token(self) -> None:.As per coding guidelines, annotate every function and use
Nonefor non-returning 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 1715 - 1746, Update the _refresh_req_idx_per_token method signature to include the required None return annotation, preserving its existing implementation unchanged.Source: Coding guidelines
1715-1746: 🗄️ Data Integrity & Integration | 🔵 Trivial | 🏗️ Heavy liftAdd an automated regression test for the MTP-Eagle transition.
Cover batch size greater than one with
max_draft_len > 1. After the draft update, assert thatreq_idx_per_tokenequalsarange(num_seqs)and that one downstream index conversion uses the refreshed mapping. Manual validation alone does not protect this cross-step contract.Based on the PR objectives, this path currently has manual validation only.
🤖 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 1715 - 1746, Add an automated regression test for the MTP-Eagle multi-request transition, using batch size greater than one and max_draft_len greater than one. Exercise the draft update that invokes _refresh_req_idx_per_token, then assert req_idx_per_token equals arange(num_seqs) and verify a downstream index conversion consumes the refreshed mapping. Use the existing attention backend test utilities and preserve coverage of the cross-step transition rather than testing the helper in isolation.
🤖 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 1715-1746: Update the _refresh_req_idx_per_token method signature
to include the required None return annotation, preserving its existing
implementation unchanged.
- Around line 1715-1746: Add an automated regression test for the MTP-Eagle
multi-request transition, using batch size greater than one and max_draft_len
greater than one. Exercise the draft update that invokes
_refresh_req_idx_per_token, then assert req_idx_per_token equals
arange(num_seqs) and verify a downstream index conversion consumes the refreshed
mapping. Use the existing attention backend test utilities and preserve coverage
of the cross-step transition rather than testing the helper in isolation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c2996b36-f366-4437-9c7b-bcd57b450cb0
📒 Files selected for processing (1)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py
|
Closing in favor of #16925, which fixes the same bug and does it more generally. I hit this independently while chasing an MTP acceptance-length gap vs vLLM on GLM-5.2 (https://nvbugs/6463964) and landed on the same root cause: #16925 is the better fix:
My one difference is that the refresh is allocation-free (a Some measurements from my investigation, in case they are useful on #16925 — the degradation is batch-dependent, which is easy to miss because it is invisible at batch=1 and invisible to accuracy tests (speculative decoding is output-lossless, so GSM8K/MMLU pass while throughput silently drops). GLM-5.2-NVFP4, TP=4, 80 MT-bench prompts, greedy,
Reproduced independently with CUDA graphs enabled: 3.233 -> 3.758. Controls: Also ruled out along the way, each with an A/B, in case it saves anyone repeating them: DSA indexer top-k sharing (#15806 — no effect, because sub-2048 sequences take the dense path and never compute a top-k to share), rejection sampling (only affects k=1), NVFP4 (the MTP layer is in the quant ignore list and runs bf16), the recycled-hidden-state norm point (TRT-LLM/vLLM/SGLang all recycle post-norm; forcing pre-norm made it worse), and KV cache dtype. @zheyuf — I also have a small CPU-only unit test for |
1 participant
DSAtrtllmAttentionMetadata.req_idx_per_token maps each token of the flattened batch to its request index. It is built once per target forward in prepare_for_indices_conversion() and then deliberately reused, to keep repeat_interleave out of the CUDA-graph-captured region. That reuse is only valid while the batch layout is unchanged.
The one-model MTP-Eagle draft loop breaks that premise: at the end of draft step 0 it collapses the batch to one token per request (seq_lens.fill(1) followed by on_update()), but AttentionMetadata.on_update() only refreshes num_tokens / num_ctx_tokens / num_generations -- it does not rebuild this mapping. Draft steps 1..k therefore consume the stale target layout ([0]*next_n + [1]*next_n + ...) where the correct mapping is arange(num_seqs), corrupting both the indexer-K slot mapping in on_update_kv_lens() and the top-k -> global index conversion in _rebuild_pool_view_cache(). Requests read and overwrite each other's indexer K-cache, and the damage compounds with draft depth.
The mapping is trivially correct at batch_size == 1, so the defect only appears at batch > 1 with max_draft_len > 1, and draft step 0 stays exact. It is DSA-specific: non-DSA attention has no such buffer.
Refresh the mapping whenever the layout is one token per request. num_tokens == num_seqs implies every seq_len is 1 (seq lens are >= 1 and sum to num_tokens), for which the mapping is exactly arange(num_seqs) -- an identity that also holds for ordinary single-token decode, so the refresh is unconditionally correct rather than a speculative-decoding special case. It is a copy_ from a preallocated static arange buffer, so it allocates nothing and preserves the graph-capture constraint that motivated caching the mapping in the first place. Placing it in the metadata rather than in the draft loop also covers the DeepSeek-V4 path, which builds the same mapping.
Measured on GLM-5.2-NVFP4, TP=4, 80 MT-bench prompts, greedy, max_draft_len=4, mtp_eagle_one_model=True. Acceptance length at batch=8 goes from 3.283 to 3.766 (batch=1 is 3.768 and is unchanged), and the conditional acceptance rate per draft depth goes from 0.880/0.775/0.670/0.596 back to the flat 0.868/0.864/0.825/0.834 seen at batch=1. Independently reproduced with CUDA graphs enabled: 3.233 -> 3.758. max_draft_len=1 is unaffected, as expected, since the loop never reaches step 1.
Dev Engineer Review
DSAtrtllmAttentionMetadata.req_idx_per_tokenwhennum_tokens == num_seqs.arangebuffer to avoid allocations and preserve CUDA graph capture support.QA Engineer Review
No test changes.
Description
For DSA models running one-model MTP-Eagle with
max_draft_len > 1, speculativeacceptance rate degrades as batch size grows: at
batch=1it matches vLLM, atbatch=8the deepest draft position loses ~half its acceptance rate (~15% acceptancelength).
DSAtrtllmAttentionMetadata.req_idx_per_tokenmaps each token of the flattened batchto its request index. It is built once per target forward in
prepare_for_indices_conversion()and then deliberately reused, to keeprepeat_interleaveout of the CUDA-graph-captured region. That reuse is only validwhile the batch layout is unchanged.
The MTP-Eagle draft loop breaks that premise. At the end of draft step 0 it collapses
the batch to one token per request (
_seq_lens.fill_(1)+on_update()), butAttentionMetadata.on_update()only refreshesnum_tokens/num_ctx_tokens/num_generations— it does not rebuild this mapping. Withnext_n = max_draft_len+1 = 5andbatch = 8, draft steps 1..k read[0,0,0,0,0,1,1,1]where the correctmapping is
[0,1,2,3,4,5,6,7]. Two consumers are then corrupted every draft step:on_update_kv_lens()writes each request's indexer-K into another request's block ata bogus offset (clobbering that sequence's real cached K), and
_rebuild_pool_view_cache()resolves each draft token's top-k indices against thewrong request's block table. Damage compounds with draft depth.
The mapping is trivially correct at
batch_size == 1([0]either way), which is whythe defect is invisible at batch 1 and why draft step 0 stays exact.
Fix: refresh the mapping whenever the layout is one token per request.
num_tokens == num_seqsimplies everyseq_lenis 1 (seq lens are ≥ 1 and sum tonum_tokens), for which the mapping is exactlyarange(num_seqs)— an identity thatalso holds for ordinary single-token decode, so the refresh is unconditionally correct
rather than a spec-dec special case. It is a
copy_from a preallocated staticarangebuffer: no allocation, no host sync, and it preserves the graph-captureconstraint that motivated caching in the first place. Placing it in the metadata
rather than in the draft loop also covers the DeepSeek-V4 path, which builds the same
mapping. Single file, +49 lines.
Test Coverage
No existing automated test guards this, and I have not added one in this PR —
flagging that explicitly rather than implying coverage exists.
Why current tests miss it:
TestGLM52::test_nvfp4runsMTPDecodingConfig(max_draft_len=1).k=1neverreaches draft step 1, so it is provably unaffected.
TestDeepSeekV32::test_fp8_blockscalecoversmtp_nextn=3only atmax_batch_size=1(latency,latency_defaultids) — batch 1 is exactly wherethe mapping is accidentally correct. The one id combining
mtp_nextn=3withmax_batch_size=24(cute_dsl_gvr_mtp3) is a GSM8K accuracy test, and acceptancerate is largely orthogonal to accuracy for speculative decoding.
Manual validation performed — GLM-5.2-NVFP4, TP=4, 80 MT-bench prompts, greedy,
max_draft_len=4,mtp_eagle_one_model=True,use_rejection_sampling=False.cond_AR[i] = AR[i]/AR[i-1]is the conditional acceptance rate at draft depth i;a healthy drafter is flat across depth.
CUDA graph off:
CUDA graph on (independently reproduced): before
batch=8AL 3.233 / cond3 0.573 →after 3.758 / cond3 0.822.
batch=1unchanged (3.739 → 3.721).PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.