[https://nvbugs/6676312][fix] Restore DSA CUDA TopK workspace lifecycle

tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py

@@ -1871,6 +1871,8 @@ def sparse_attn_indexer(
                 if gvr_prior_indices is not None
                 else None
             )
+            assert metadata.radix_aux_indices is not None
+            assert metadata.radix_aux_logits is not None
             self.top_k(
                 logits_decode,
                 topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :],
@@ -1880,6 +1882,8 @@ def sparse_attn_indexer(
                 next_n=next_n,
                 max_seq_len=indexer_max_seq_len,
                 gvr_ext_kwargs=gvr_ext_kwargs,
+                radix_aux_indices=metadata.radix_aux_indices,
+                radix_aux_logits=metadata.radix_aux_logits,
             )
 
         elif has_decode and metadata.skip_indexer_for_gen_reqs:

tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py

@@ -38,6 +38,7 @@
 # Indexer MQA-logits are currently always fp32. The dtype is part of the
 # CuTe DSL Top-K compile key, so warmup must use the runtime dtype.
 _INDEXER_LOGITS_DTYPE = torch.float32
+_MAX_RADIX_BLOCKS_PER_ROW = 10
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.speculative.interface import SpecMetadata
@@ -84,6 +85,8 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata):
     num_sparse_topk: int
     # TopK for dynamic sparse MLA
     sparse_mla_topk: int
+    radix_aux_indices: Optional[torch.Tensor] = field(default=None, init=False)
+    radix_aux_logits: Optional[torch.Tensor] = field(default=None, init=False)
     # max number of draft tokens
     max_draft_tokens: int = 0
     # Indexer head dimension
@@ -159,6 +162,7 @@ def __init__(self, *args, **kwargs):
         self._group_remap_batched = {}
         self.num_ctx_mla_kv_tokens = 0
         self.nvfp4_mla_context_fp8_scratch = None
+        self._radix_rows_per_sequence = 1
         super().__init__(*args, **kwargs)
         sparse_metadata_params = self.sparse_metadata_params
         if not isinstance(sparse_metadata_params, DSAMetadataParams):
@@ -769,6 +773,24 @@ def _create_nvfp4_mla_generation_buffers(self, capture_graph=False):
             capture_graph=capture_graph,
         )
 
+    def _create_radix_aux_buffers(self, capture_graph: bool) -> None:
+        max_num_rows = self.max_num_sequences * self._radix_rows_per_sequence
+        radix_shape = (max_num_rows, _MAX_RADIX_BLOCKS_PER_ROW, self.num_sparse_topk)
+        self.radix_aux_indices = self.get_empty(
+            self.cuda_graph_buffers,
+            radix_shape,
+            cache_name="radix_aux_indices",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
+        self.radix_aux_logits = self.get_empty(
+            self.cuda_graph_buffers,
+            radix_shape,
+            cache_name="radix_aux_logits",
+            dtype=torch.float32,
+            capture_graph=capture_graph,
+        )
+
     def create_buffers_for_indexer(self, capture_graph=False):
         sparse_metadata_params = self.sparse_metadata_params
         if not isinstance(sparse_metadata_params, DSAMetadataParams):
@@ -994,6 +1016,7 @@ def create_buffers_for_indexer(self, capture_graph=False):
                 dtype=torch.int32,
                 capture_graph=capture_graph,
             )
+        self._create_radix_aux_buffers(capture_graph)
         # Create expanded buffers for MTP support
         self.create_expanded_buffers(capture_graph=capture_graph)
 
@@ -1093,6 +1116,11 @@ def update_spec_dec_param(
             num_contexts=num_contexts,
         )
         self.max_draft_tokens = max_draft_len
+        rows_per_sequence = max(1 + max_draft_len, 1 + (max_total_draft_tokens or 0))
+        if is_spec_dec_dynamic_tree and spec_tree_manager is not None:
+            rows_per_sequence = max(rows_per_sequence, spec_tree_manager._internal_buf_dim)
+        resize_radix_workspace = rows_per_sequence != self._radix_rows_per_sequence
+        self._radix_rows_per_sequence = rows_per_sequence
         capture_graph = self.is_cuda_graph
         max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens)
         if (
@@ -1105,6 +1133,8 @@ def update_spec_dec_param(
         init_shape = self.kv_lens_expanded_host.shape[0]
         if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape:
             self.create_expanded_buffers(capture_graph=capture_graph)
+        if resize_radix_workspace:
+            self._create_radix_aux_buffers(capture_graph)
 
     def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor:
         """Refresh INDEX_KEY offsets and return their physical pool slots."""

tensorrt_llm/_torch/modules/top_k.py

@@ -11,8 +11,6 @@
 
 from tensorrt_llm.logger import logger
 
-from ..memory_buffer_utils import get_memory_buffers
-
 
 class TopKImplementation(str, Enum):
     """Top-K implementations grouped by backend and algorithm."""
@@ -36,10 +34,10 @@ class TopK(nn.Module):
 
     GVR decode state is owned by the caller so it can be shared with the
     request metadata and retain a stable address across CUDA Graph replays.
+    Native CUDA scratch is supplied by the caller so its lifetime follows the
+    caller's eager or CUDA-graph metadata.
     """
 
-    _memory_buffers = get_memory_buffers()
-
     def __init__(
         self,
         top_k: int,
@@ -89,6 +87,8 @@ def forward(
         next_n: int = 1,
         max_seq_len: int | None = None,
         gvr_ext_kwargs: dict[str, torch.Tensor | None] | None = None,
+        radix_aux_indices: torch.Tensor | None = None,
+        radix_aux_logits: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Write prefill or decode Top-K indices into ``output_indices``.
 
@@ -110,6 +110,10 @@ def forward(
                 self-sampling engine does not consume this state.
                 ``gvr_row_order`` is an optional int32 request ordering with
                 shape ``[num_requests]`` on the same device.
+            radix_aux_indices: Caller-owned int32 workspace for native CUDA
+                Radix split work.
+            radix_aux_logits: Caller-owned float32 workspace for native CUDA
+                Radix split work.
 
         Returns:
             ``output_indices`` after the selected implementation writes it.
@@ -127,6 +131,8 @@ def forward(
             next_n,
             max_seq_len,
             gvr_ext_kwargs,
+            radix_aux_indices,
+            radix_aux_logits,
         )
 
     def _forward_prefill(
@@ -176,6 +182,8 @@ def _forward_decode(
         next_n: int,
         max_seq_len: int | None,
         gvr_ext_kwargs: dict[str, torch.Tensor | None] | None,
+        radix_aux_indices: torch.Tensor | None,
+        radix_aux_logits: torch.Tensor | None,
     ) -> torch.Tensor:
         if self.decode_implementation == TopKImplementation.TORCH:
             return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n)
@@ -187,6 +195,8 @@ def _forward_decode(
                 output_indices,
                 next_n,
                 max_seq_len=max_seq_len,
+                radix_aux_indices=radix_aux_indices,
+                radix_aux_logits=radix_aux_logits,
                 **(gvr_ext_kwargs or {}),
             )
 
@@ -196,6 +206,8 @@ def _forward_decode(
             scan_lengths,
             output_indices,
             next_n,
+            radix_aux_indices,
+            radix_aux_logits,
         )
 
     def _forward_decode_radix(
@@ -205,6 +217,8 @@ def _forward_decode_radix(
         scan_lengths: torch.Tensor,
         output_indices: torch.Tensor,
         next_n: int,
+        radix_aux_indices: torch.Tensor | None,
+        radix_aux_logits: torch.Tensor | None,
     ) -> torch.Tensor:
         use_cute_dsl = self.decode_implementation == TopKImplementation.CUTE_DSL_RADIX and not (
             self.compress_ratio > 1 and next_n > 1
@@ -219,7 +233,9 @@ def _forward_decode_radix(
             )
             return output_indices
 
-        radix_indices, radix_values = self._get_radix_workspace(scores)
+        radix_indices, radix_values = self._get_radix_workspace(
+            scores, radix_aux_indices, radix_aux_logits
+        )
         torch.ops.trtllm.indexer_topk_decode(
             scores,
             sequence_lengths,
@@ -232,52 +248,26 @@ def _forward_decode_radix(
         )
         return output_indices
 
-    def _get_workspace(
+    def _get_radix_workspace(
         self,
         scores: torch.Tensor,
-        shape: tuple[int, ...],
-        dtype: torch.dtype,
-        buffer_name: str,
-    ) -> torch.Tensor:
-        device_buffer_name = f"{buffer_name}_{scores.device}"
-        if scores.is_cuda:
-            with torch.cuda.device(scores.device):
-                capture_graph = torch.cuda.is_current_stream_capturing()
-                return self._memory_buffers.get_buffer(
-                    shape,
-                    dtype=dtype,
-                    buffer_name=device_buffer_name,
-                    reserve_buffer=capture_graph,
-                )
-        return self._memory_buffers.get_buffer(
-            shape,
-            dtype=dtype,
-            buffer_name=device_buffer_name,
-            reserve_buffer=False,
-        )
-
-    def _get_radix_workspace(
-        self, scores: torch.Tensor
+        radix_aux_indices: torch.Tensor | None,
+        radix_aux_logits: torch.Tensor | None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         if scores.dtype != torch.float32:
             # The C++ bf16/fp16 entry has no split-work tier or aux-buffer
             # arguments and rejects widths that would require split work.
             return None, None
 
-        shape = (scores.shape[0], _MAX_RADIX_BLOCKS_PER_ROW, self.top_k)
-        radix_indices = self._get_workspace(
-            scores,
-            shape,
-            torch.int32,
-            "top_k_radix_indices_workspace",
-        )
-        radix_values = self._get_workspace(
-            scores,
-            shape,
-            torch.float32,
-            "top_k_radix_values_workspace",
+        if radix_aux_indices is None or radix_aux_logits is None:
+            raise ValueError(
+                "Native CUDA Radix TopK requires radix_aux_indices and radix_aux_logits"
+            )
+        num_rows = scores.shape[0]
+        return (
+            radix_aux_indices[:num_rows, :_MAX_RADIX_BLOCKS_PER_ROW, : self.top_k],
+            radix_aux_logits[:num_rows, :_MAX_RADIX_BLOCKS_PER_ROW, : self.top_k],
         )
-        return radix_indices, radix_values
 
     def _forward_decode_gvr(
         self,
@@ -286,6 +276,8 @@ def _forward_decode_gvr(
         output_indices: torch.Tensor,
         next_n: int,
         max_seq_len: int | None,
+        radix_aux_indices: torch.Tensor | None,
+        radix_aux_logits: torch.Tensor | None,
         gvr_prior_indices: torch.Tensor | None = None,
         gvr_row_order: torch.Tensor | None = None,
     ) -> torch.Tensor:
@@ -339,7 +331,9 @@ def _forward_decode_gvr(
                 "falling back to the CUDA insertion/radix Top-K path.",
                 key="selfsampling_topk_fallthrough",
             )
-            radix_indices, radix_values = self._get_radix_workspace(scores)
+            radix_indices, radix_values = self._get_radix_workspace(
+                scores, radix_aux_indices, radix_aux_logits
+            )
             torch.ops.trtllm.indexer_topk_decode(
                 scores,
                 sequence_lengths,

tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py

@@ -410,6 +410,7 @@ def test_shared_topk_lifecycle(monkeypatch):
     metadata.max_num_sequences = 2
     metadata.max_num_tokens = 4
     metadata.num_sparse_topk = 3
+    metadata._radix_rows_per_sequence = 1
     metadata.num_sms = 1
     metadata.cuda_graph_buffers = None
     metadata.kv_cache_manager = SimpleNamespace(max_blocks_per_seq=2)
@@ -640,6 +641,60 @@ def test_indexer_projection_dtype_follows_bf16_flag(monkeypatch, flag_value, exp
     assert indexer.weights_proj.dtype == expected_dtype
 
 
+@pytest.mark.parametrize(
+    "is_tree,is_dynamic,max_draft_len,max_total_draft_tokens,spec_tree_manager,rows_per_sequence",
+    [
+        (False, False, 7, 7, None, 8),
+        (True, False, 3, 7, None, 8),
+        (
+            True,
+            True,
+            6,
+            31,
+            SimpleNamespace(_internal_buf_dim=60, dynamic_tree_max_topK=10),
+            60,
+        ),
+    ],
+)
+def test_metadata_spec_update_resizes_radix_workspace(
+    is_tree,
+    is_dynamic,
+    max_draft_len,
+    max_total_draft_tokens,
+    spec_tree_manager,
+    rows_per_sequence,
+):
+    metadata = object.__new__(DSAtrtllmAttentionMetadata)
+    metadata.max_num_sequences = 2
+    metadata.max_draft_tokens = 0
+    metadata._radix_rows_per_sequence = 1
+    metadata.num_sparse_topk = 2
+    metadata.is_cuda_graph = True
+    metadata.nvfp4_mla_fp8_scratch = None
+    metadata.kv_lens_cuda_2d = torch.empty(2, 1)
+    metadata.kv_lens_expanded_host = torch.empty(2)
+    metadata._create_kv_lens_2d_buffer = Mock()
+    metadata.create_expanded_buffers = Mock()
+    metadata._create_radix_aux_buffers = Mock()
+
+    with patch(
+        "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata."
+        "TrtllmAttentionMetadata.update_spec_dec_param"
+    ):
+        metadata.update_spec_dec_param(
+            batch_size=2,
+            is_spec_decoding_enabled=True,
+            is_spec_dec_tree=is_tree,
+            is_spec_dec_dynamic_tree=is_dynamic,
+            max_draft_len=max_draft_len,
+            max_total_draft_tokens=max_total_draft_tokens,
+            spec_tree_manager=spec_tree_manager,
+        )
+
+    assert metadata._radix_rows_per_sequence == rows_per_sequence
+    metadata._create_radix_aux_buffers.assert_called_once_with(True)
+
+
 def _ceil_to_ue8m0(x: torch.Tensor):
     """Round tensor values up to the nearest power of two (UE8M0 format)."""
     return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))