From 7f78f393939974724d61c92de7f7518ac1b6adaf Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 3 Aug 2026 20:26:02 -0700 Subject: [PATCH] [https://nvbugs/6463964][fix] DSA: rebuild req_idx_per_token for MTP-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 --- .../_torch/attention_backend/sparse/dsa.py | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 26d8a96e11ce..ec22dc20cbdd 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -825,6 +825,7 @@ def on_update_kv_lens(self): # caches so they are recomputed (and captured) on every _forward_step". Invalidate the # pool_view cache here so it is recomputed on the next # transform_local_topk_and_prepare_pool_view() call. + self._refresh_req_idx_per_token() self._invalidate_pool_view_cache() # Clear per-step cross-layer top-k, but keep it inside the MTP draft # loop so the step-0 stash survives for the reuse branch. @@ -1108,6 +1109,22 @@ def create_buffers_for_indexer(self, capture_graph=False): ) self.host_req_idx_per_token = torch.empty_like( self.req_idx_per_token, device='cpu', pin_memory=prefer_pinned()) + # Static arange used to refresh ``req_idx_per_token`` when the batch + # layout becomes one token per request (see _refresh_req_idx_per_token). + # Content is constant, so it is filled once here and only ever read; + # refreshing via copy_ from it allocates nothing and is therefore safe + # under CUDA graph capture (which is why prepare_for_indices_conversion + # avoids repeat_interleave in the first place). + self._arange_req_idx = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens, ), + cache_name="arange_req_idx", + dtype=torch.int32, + capture_graph=capture_graph, + ) + torch.arange(self.max_num_tokens, + dtype=torch.int32, + out=self._arange_req_idx) # Block table for topk_indices conversion (shared for context and generation) self.block_table = self.get_empty( self.cuda_graph_buffers, @@ -1695,6 +1712,38 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, else: self.max_gen_seq_len = 0 + def _refresh_req_idx_per_token(self): + """Keep ``req_idx_per_token`` consistent with the *current* batch layout. + + ``prepare_for_indices_conversion`` builds this mapping once per target + forward and it is deliberately reused afterwards to keep + ``repeat_interleave`` out of the CUDA-graph-captured region. That reuse + is only valid while the layout is unchanged. + + The MTP-Eagle draft loop breaks that assumption: after draft step 0 it + rewrites the batch to one token per request + (``eagle3.py::_forward_linear_draft_loop`` -> ``_seq_lens.fill_(1)`` + + ``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 then read the stale *target* layout + ([0]*next_n + [1]*next_n + ...) where the correct mapping is + arange(num_seqs), which mis-addresses both the indexer-K slot mapping + (``on_update_kv_lens``) and the top-k -> global index conversion + (``_rebuild_pool_view_cache`` -> ``convert_req_index_to_global``). + The result is a monotonic acceptance-rate collapse with draft depth at + batch>1. It is a no-op at batch==1, where the mapping is [0] either way. + + ``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)``. That identity holds for ordinary single-token + decode too, so this refresh is unconditionally correct -- it is not + specific to speculative decoding. + """ + n = self.num_tokens + if n > 0 and n == self.num_seqs: + self.req_idx_per_token[:n].copy_(self._arange_req_idx[:n], + non_blocking=True) + def prepare_for_indices_conversion(self): # Build req_idx_per_token for topk_indices conversion # Use pinned staging buffer to avoid pageable H2D memcpy