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..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 @@ -931,7 +931,11 @@ 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 is read, not recomputed: super().on_update_kv_lens() + has already rebuilt it for the current seq_lens. + """ device = seq_lens.device # cu_seq_lens @@ -939,14 +943,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 a57fdadf0776..f1a45201e16e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -263,6 +263,20 @@ def _pick_dsl_expand( return factor, eff +def build_req_idx_per_token(seq_lens: torch.Tensor, + num_tokens: int) -> torch.Tensor: + """Map each flattened-batch token to its request index. + + 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, + 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,13 +843,18 @@ def on_update_kv_lens(self): # boundary so a "shared" layer never reuses a stale top-k. self.shared_topk_indices = None + # 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], + 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 - # Reuse request-per-token mapping prepared in metadata.prepare(). - # This avoids repeat_interleave in graph-capture mode. req_indices = self.req_idx_per_token[:self.num_tokens].to( dtype=torch.int64) seq_starts = torch.cumsum( 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..f4d2144bc151 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py @@ -0,0 +1,100 @@ +# 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. + +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 + +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(): + """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" + max_draft_len = 3 + num_requests = 3 + + md = object.__new__(DSAtrtllmAttentionMetadata) + md.kv_cache_manager = None + md._num_generations = 0 + # 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()-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) + 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) + + # 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) + 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]