From 2cfeaaba9359d2e8740aa7787cfab6651eda5dd8 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 28 Jul 2026 15:48:23 -0700 Subject: [PATCH 1/4] [None][fix] DSA: rebuild token-to-request map inside the draft loop 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 --- tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index a57fdadf0776..1e314f700c83 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -834,8 +834,17 @@ def on_update_kv_lens(self): # Runtime cached lengths after overlap/spec-dec correction. start_positions = self.kv_lens_cuda[:self.num_seqs] - seq_lens - # Reuse request-per-token mapping prepared in metadata.prepare(). - # This avoids repeat_interleave in graph-capture mode. + # Rebuild the token->request map from the current seq_lens: the map + # built in prepare() describes the target forward's layout, but the + # draft loop rewrites seq_lens to one token per request. Equivalent to + # prepare()'s repeat_interleave whenever seq_lens is unchanged. + cu_seq_lens = torch.cumsum(seq_lens, dim=0, dtype=torch.int32) + token_idx = torch.arange(self.num_tokens, + device=seq_lens.device, + dtype=torch.int32) + self.req_idx_per_token[:self.num_tokens] = torch.searchsorted( + cu_seq_lens, token_idx, + right=True).to(self.req_idx_per_token.dtype) req_indices = self.req_idx_per_token[:self.num_tokens].to( dtype=torch.int64) seq_starts = torch.cumsum( From fae96ced767d5e4246388846f529b6d2ecdfc29d Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 4 Aug 2026 01:16:51 +0000 Subject: [PATCH 2/4] [https://nvbugs/6513132][fix] Share the req_idx_per_token rebuild between 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 --- .../sparse/deepseek_v4/deepseek_v4.py | 16 +++---- .../_torch/attention_backend/sparse/dsa.py | 42 ++++++++++++++----- 2 files changed, 40 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 1f76e5a73e51..d2ed2da01142 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -931,7 +931,13 @@ def _compute_token_positions( cu_seq_lens_buf: torch.Tensor, req_idx_per_token_buf: torch.Tensor, ) -> torch.Tensor: - """Compute cu_seq_lens, req_idx_per_token, and token_positions (eager).""" + """Compute cu_seq_lens and token_positions (eager). + + ``req_idx_per_token_buf`` must already describe the current seq_lens + layout; ``super().on_update_kv_lens()`` rebuilds it via + ``build_req_idx_per_token`` before this runs, so it is read here + rather than recomputed. + """ device = seq_lens.device # cu_seq_lens @@ -939,14 +945,10 @@ def _compute_token_positions( torch.cumsum(seq_lens.to(torch.int), dim=0), (1, 0) ) - # req_idx_per_token via searchsorted - token_idx = torch.arange(num_tokens, dtype=torch.int32, device=device) - req_idx = torch.searchsorted( - cu_seq_lens_buf[1 : batch_size + 1].to(torch.int32), token_idx, right=True - ) - req_idx_per_token_buf[:num_tokens] = req_idx + req_idx = req_idx_per_token_buf[:num_tokens].to(torch.int64) # token positions + token_idx = torch.arange(num_tokens, dtype=torch.int32, device=device) base_pos = cached_tokens[req_idx].to(torch.int32) offsets = token_idx - cu_seq_lens_buf[req_idx].to(torch.int32) return base_pos + offsets diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 1e314f700c83..ec1cfba59857 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -263,6 +263,26 @@ def _pick_dsl_expand( return factor, eff +def build_req_idx_per_token(seq_lens: torch.Tensor, + num_tokens: int) -> torch.Tensor: + """Map each token of the flattened batch to the request it belongs to. + + Token ``t`` belongs to the first request whose cumulative token end + exceeds ``t``; ``right=True`` makes duplicated cumsum boundaries skip + zero-length rows, so the result matches ``torch.repeat_interleave( + arange(num_seqs), seq_lens)`` for every layout (see the unit test). + + Device-side, fixed-shape counterpart of the host build in + ``prepare_for_indices_conversion()`` — usable between draft-loop + iterations and inside CUDA-graph capture, where a host round-trip is not. + """ + cu_seq_lens = torch.cumsum(seq_lens, dim=0, dtype=torch.int32) + token_idx = torch.arange(num_tokens, + device=seq_lens.device, + dtype=torch.int32) + return torch.searchsorted(cu_seq_lens, token_idx, right=True) + + def _compute_slot_mappings( global_positions: torch.Tensor, block_offsets: torch.Tensor, @@ -829,22 +849,22 @@ def on_update_kv_lens(self): # boundary so a "shared" layer never reuses a stale top-k. self.shared_topk_indices = None + # Rebuild the token->request map from the current seq_lens: the map + # built in prepare() describes the target forward's layout, but the + # draft loop rewrites seq_lens to one token per request. Equivalent to + # prepare()'s repeat_interleave whenever seq_lens is unchanged. + # Unconditional (no kv_cache_manager guard) so subclasses such as + # DeepSeek-V4 can reuse the rebuilt buffer instead of recomputing it. + if self.num_tokens > 0: + self.req_idx_per_token[:self.num_tokens] = build_req_idx_per_token( + self.seq_lens_cuda[:self.num_seqs], + self.num_tokens).to(self.req_idx_per_token.dtype) + if self.kv_cache_manager is not None and self.num_tokens > 0: seq_lens = self.seq_lens_cuda[:self.num_seqs] # Runtime cached lengths after overlap/spec-dec correction. start_positions = self.kv_lens_cuda[:self.num_seqs] - seq_lens - # Rebuild the token->request map from the current seq_lens: the map - # built in prepare() describes the target forward's layout, but the - # draft loop rewrites seq_lens to one token per request. Equivalent to - # prepare()'s repeat_interleave whenever seq_lens is unchanged. - cu_seq_lens = torch.cumsum(seq_lens, dim=0, dtype=torch.int32) - token_idx = torch.arange(self.num_tokens, - device=seq_lens.device, - dtype=torch.int32) - self.req_idx_per_token[:self.num_tokens] = torch.searchsorted( - cu_seq_lens, token_idx, - right=True).to(self.req_idx_per_token.dtype) req_indices = self.req_idx_per_token[:self.num_tokens].to( dtype=torch.int64) seq_starts = torch.cumsum( From 8d4c26fcf88a41d3ea67bfe12f65f6e36d1a882a Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 4 Aug 2026 23:21:27 +0000 Subject: [PATCH 3/4] [https://nvbugs/6513132][fix] Pin the req_idx_per_token rebuild with 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 --- .../sparse/dsa/test_req_idx_per_token.py | 114 ++++++++++++++++++ 1 file changed, 114 insertions(+) create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py new file mode 100644 index 000000000000..7e344e61c73e --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Tests for the token->request map used by DSAtrtllmAttentionMetadata. + +Two invariants, one test each: + +1. build_req_idx_per_token (device searchsorted) must equal + prepare_for_indices_conversion()'s host repeat_interleave for every batch + layout, so the two builders cannot drift. +2. on_update_kv_lens() must rebuild the map for the CURRENT seq_lens. The MTP + draft loop rewrites seq_lens to one token per request without re-running + prepare(); reusing prepare()'s map misattributes every draft token to + request 0, corrupting indexer K-writes and top-k reads through the wrong + block table (https://nvbugs/6513132, https://nvbugs/6513093). +""" + +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.dsa import ( + DSAtrtllmAttentionMetadata, + build_req_idx_per_token, +) + + +def _host_reference(seq_lens: torch.Tensor) -> torch.Tensor: + """The host build from prepare_for_indices_conversion().""" + return torch.repeat_interleave( + torch.arange(len(seq_lens), dtype=torch.int32, device=seq_lens.device), + seq_lens, + dim=0, + ) + + +@pytest.mark.parametrize( + "seq_lens", + [ + pytest.param([4, 4, 4], id="target_forward_mtp3"), + pytest.param([1, 1, 1], id="draft_loop"), + pytest.param([37, 5, 1, 1], id="mixed_ctx_gen"), + pytest.param([2, 0, 3], id="zero_length_row"), + pytest.param([0, 4], id="leading_zero_row"), + ], +) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_matches_host_repeat_interleave(seq_lens, device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) + num_tokens = int(seq_lens.sum()) + + result = build_req_idx_per_token(seq_lens, num_tokens) + + assert result.to(torch.int32).tolist() == _host_reference(seq_lens).tolist() + + +def test_on_update_kv_lens_rebuilds_stale_map(): + """The draft-loop transition: the regression this PR fixes. + + Bare-instance construction (object.__new__ + backing fields) mirrors + test_dsa_indexer.py; kv_cache_manager=None and num_generations=0 confine + on_update_kv_lens() to the map rebuild under test. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + device = "cuda" + max_draft_len = 3 + num_requests = 3 + + md = object.__new__(DSAtrtllmAttentionMetadata) + md.kv_cache_manager = None + md._num_generations = 0 + # Collaborators invoked at the end of on_update_kv_lens(); unrelated to + # the map rebuild under test (same stubbing style as test_dsa_indexer.py). + md.kv_lens_cuda = torch.tensor([100, 200, 300], dtype=torch.int32, device=device) + md._compute_kv_lens_row_reorder = Mock() + md.prepare_dense_topk_indices = Mock() + + # prepare() state for the target forward: 1 + max_draft_len tokens per + # request, map built host-side via repeat_interleave. + target_seq_lens = torch.full((num_requests,), 1 + max_draft_len, dtype=torch.int32) + md._seq_lens = target_seq_lens + md._seq_lens_cuda = target_seq_lens.to(device) + md._num_tokens = int(target_seq_lens.sum()) + md.req_idx_per_token = torch.empty(md._num_tokens, dtype=torch.int32, device=device) + md.req_idx_per_token[:] = _host_reference(md._seq_lens_cuda) + + # The draft loop rewrites seq_lens to one token per request (what + # _preprocess_inputs does between draft iterations) without re-running + # prepare(). The stale prefix misattributes every token to request 0. + draft_seq_lens = torch.ones(num_requests, dtype=torch.int32) + md._seq_lens = draft_seq_lens + md._seq_lens_cuda = draft_seq_lens.to(device) + md._num_tokens = num_requests + assert md.req_idx_per_token[:num_requests].tolist() == [0, 0, 0] + + md.on_update_kv_lens() + + assert md.req_idx_per_token[:num_requests].tolist() == [0, 1, 2] From 1a9fb816aa4802fe63ae58500df400754d17c7cd Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 5 Aug 2026 01:27:56 +0000 Subject: [PATCH 4/4] [https://nvbugs/6513132][fix] Trim comments to the essentials Signed-off-by: Zheyu Fu --- .../sparse/deepseek_v4/deepseek_v4.py | 6 ++-- .../_torch/attention_backend/sparse/dsa.py | 20 +++--------- .../sparse/dsa/test_req_idx_per_token.py | 32 ++++++------------- 3 files changed, 16 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index d2ed2da01142..c07d82cf9f6c 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -933,10 +933,8 @@ def _compute_token_positions( ) -> torch.Tensor: """Compute cu_seq_lens and token_positions (eager). - ``req_idx_per_token_buf`` must already describe the current seq_lens - layout; ``super().on_update_kv_lens()`` rebuilds it via - ``build_req_idx_per_token`` before this runs, so it is read here - rather than recomputed. + req_idx_per_token_buf is read, not recomputed: super().on_update_kv_lens() + has already rebuilt it for the current seq_lens. """ device = seq_lens.device diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index ec1cfba59857..f1a45201e16e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -265,16 +265,10 @@ def _pick_dsl_expand( def build_req_idx_per_token(seq_lens: torch.Tensor, num_tokens: int) -> torch.Tensor: - """Map each token of the flattened batch to the request it belongs to. + """Map each flattened-batch token to its request index. - Token ``t`` belongs to the first request whose cumulative token end - exceeds ``t``; ``right=True`` makes duplicated cumsum boundaries skip - zero-length rows, so the result matches ``torch.repeat_interleave( - arange(num_seqs), seq_lens)`` for every layout (see the unit test). - - Device-side, fixed-shape counterpart of the host build in - ``prepare_for_indices_conversion()`` — usable between draft-loop - iterations and inside CUDA-graph capture, where a host round-trip is not. + Capture-safe device counterpart of prepare_for_indices_conversion()'s + host repeat_interleave; right=True keeps zero-length rows equivalent. """ cu_seq_lens = torch.cumsum(seq_lens, dim=0, dtype=torch.int32) token_idx = torch.arange(num_tokens, @@ -849,12 +843,8 @@ def on_update_kv_lens(self): # boundary so a "shared" layer never reuses a stale top-k. self.shared_topk_indices = None - # Rebuild the token->request map from the current seq_lens: the map - # built in prepare() describes the target forward's layout, but the - # draft loop rewrites seq_lens to one token per request. Equivalent to - # prepare()'s repeat_interleave whenever seq_lens is unchanged. - # Unconditional (no kv_cache_manager guard) so subclasses such as - # DeepSeek-V4 can reuse the rebuilt buffer instead of recomputing it. + # prepare()'s map is stale once the draft loop rewrites seq_lens. + # Unconditional so subclasses (DeepSeek-V4) can reuse the buffer. if self.num_tokens > 0: self.req_idx_per_token[:self.num_tokens] = build_req_idx_per_token( self.seq_lens_cuda[:self.num_seqs], diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py index 7e344e61c73e..f4d2144bc151 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py @@ -15,16 +15,11 @@ """ Tests for the token->request map used by DSAtrtllmAttentionMetadata. -Two invariants, one test each: - -1. build_req_idx_per_token (device searchsorted) must equal - prepare_for_indices_conversion()'s host repeat_interleave for every batch - layout, so the two builders cannot drift. -2. on_update_kv_lens() must rebuild the map for the CURRENT seq_lens. The MTP - draft loop rewrites seq_lens to one token per request without re-running - prepare(); reusing prepare()'s map misattributes every draft token to - request 0, corrupting indexer K-writes and top-k reads through the wrong - block table (https://nvbugs/6513132, https://nvbugs/6513093). +1. build_req_idx_per_token must match the host repeat_interleave build for + every layout, so the device and host builders cannot drift. +2. on_update_kv_lens() must rebuild the map after the MTP draft loop rewrites + seq_lens, or every draft token is misattributed to request 0 + (https://nvbugs/6513132, https://nvbugs/6513093). """ from unittest.mock import Mock @@ -70,12 +65,7 @@ def test_matches_host_repeat_interleave(seq_lens, device): def test_on_update_kv_lens_rebuilds_stale_map(): - """The draft-loop transition: the regression this PR fixes. - - Bare-instance construction (object.__new__ + backing fields) mirrors - test_dsa_indexer.py; kv_cache_manager=None and num_generations=0 confine - on_update_kv_lens() to the map rebuild under test. - """ + """on_update_kv_lens() must replace prepare()'s stale map (fails pre-fix).""" if not torch.cuda.is_available(): pytest.skip("CUDA not available") device = "cuda" @@ -85,14 +75,12 @@ def test_on_update_kv_lens_rebuilds_stale_map(): md = object.__new__(DSAtrtllmAttentionMetadata) md.kv_cache_manager = None md._num_generations = 0 - # Collaborators invoked at the end of on_update_kv_lens(); unrelated to - # the map rebuild under test (same stubbing style as test_dsa_indexer.py). + # Stub collaborators unrelated to the map rebuild (test_dsa_indexer.py style). md.kv_lens_cuda = torch.tensor([100, 200, 300], dtype=torch.int32, device=device) md._compute_kv_lens_row_reorder = Mock() md.prepare_dense_topk_indices = Mock() - # prepare() state for the target forward: 1 + max_draft_len tokens per - # request, map built host-side via repeat_interleave. + # prepare()-time state: target forward, 1 + max_draft_len tokens/request. target_seq_lens = torch.full((num_requests,), 1 + max_draft_len, dtype=torch.int32) md._seq_lens = target_seq_lens md._seq_lens_cuda = target_seq_lens.to(device) @@ -100,9 +88,7 @@ def test_on_update_kv_lens_rebuilds_stale_map(): md.req_idx_per_token = torch.empty(md._num_tokens, dtype=torch.int32, device=device) md.req_idx_per_token[:] = _host_reference(md._seq_lens_cuda) - # The draft loop rewrites seq_lens to one token per request (what - # _preprocess_inputs does between draft iterations) without re-running - # prepare(). The stale prefix misattributes every token to request 0. + # Draft loop: one token per request; the stale prefix reads [0, 0, 0]. draft_seq_lens = torch.ones(num_requests, dtype=torch.int32) md._seq_lens = draft_seq_lens md._seq_lens_cuda = draft_seq_lens.to(device)