diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index d6655133dbf6..6b82f25919cf 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/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: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index f74c6cacbbc7..6e6956218f53 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/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.""" diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 8093a96a016d..482c01a2687f 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/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, diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 6ac79f345954..ee847e3611a9 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/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()))) diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index a9d319eac9f5..af7a3f0d09d2 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -4,7 +4,7 @@ import sys from types import SimpleNamespace -from unittest.mock import Mock, call +from unittest.mock import Mock import pytest import torch @@ -96,9 +96,6 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: output = torch.empty(2, 2, dtype=torch.int32) radix_indices = torch.empty(2, 10, 2, dtype=torch.int32) radix_values = torch.empty(2, 10, 2) - buffers = Mock() - buffers.get_buffer.side_effect = [radix_indices, radix_values] - monkeypatch.setattr(TopK, "_memory_buffers", buffers) top_k( scores, output, @@ -106,32 +103,19 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: sequence_lengths=logical_lengths, scan_lengths=scan_lengths, next_n=2, - ) - cute_dsl.assert_not_called() - assert buffers.get_buffer.call_args_list == [ - call( - (2, 10, 2), - dtype=torch.int32, - buffer_name="top_k_radix_indices_workspace_cpu", - reserve_buffer=False, - ), - call( - (2, 10, 2), - dtype=torch.float32, - buffer_name="top_k_radix_values_workspace_cpu", - reserve_buffer=False, - ), - ] - trtllm.assert_called_once_with( - scores, - logical_lengths, - output, - 2, - 2, - compress_ratio=4, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, ) + cute_dsl.assert_not_called() + trtllm.assert_called_once() + runtime_call = trtllm.call_args + assert runtime_call.args[0] is scores + assert runtime_call.args[1] is logical_lengths + assert runtime_call.args[2] is output + assert runtime_call.args[3:] == (2, 2) + assert runtime_call.kwargs["compress_ratio"] == 4 + assert runtime_call.kwargs["radix_aux_indices"].data_ptr() == radix_indices.data_ptr() + assert runtime_call.kwargs["radix_aux_logits"].data_ptr() == radix_values.data_ptr() def test_gvr_uses_caller_prior_state(monkeypatch) -> None: @@ -362,33 +346,77 @@ def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: output = torch.empty((1, 1), dtype=torch.int32) radix_indices = torch.empty(1, 10, 1, dtype=torch.int32) radix_values = torch.empty(1, 10, 1) - buffers = Mock() - buffers.get_buffer.side_effect = [radix_indices, radix_values] - monkeypatch.setattr(TopK, "_memory_buffers", buffers) - result = top_k( scores, output, is_prefill=False, sequence_lengths=lengths, scan_lengths=lengths, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, ) assert top_k.prefill_implementation == TopKImplementation.CUDA_RADIX assert top_k.decode_implementation == TopKImplementation.CUDA_RADIX assert result is output - assert buffers.get_buffer.call_count == 2 - decode.assert_called_once_with( - scores, - lengths, - output, - 1, - 1, - compress_ratio=1, - radix_aux_indices=radix_indices, - radix_aux_logits=radix_values, + decode.assert_called_once() + runtime_call = decode.call_args + assert runtime_call.args[0] is scores + assert runtime_call.args[1] is lengths + assert runtime_call.args[2] is output + assert runtime_call.args[3:] == (1, 1) + assert runtime_call.kwargs["compress_ratio"] == 1 + assert runtime_call.kwargs["radix_aux_indices"].data_ptr() == radix_indices.data_ptr() + assert runtime_call.kwargs["radix_aux_logits"].data_ptr() == radix_values.data_ptr() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_graph_workspace_isolated_from_eager_growth(monkeypatch) -> None: + """Caller-owned graph scratch must survive unrelated eager execution.""" + device = torch.device("cuda", torch.cuda.current_device()) + num_columns = 32_768 + top_k = 8 + top_k_module = TopK(top_k, decode_implementation=TopKImplementation.CUDA_RADIX) + static_scores = torch.randn(1, num_columns, device=device) + static_lengths = torch.full((1,), num_columns, dtype=torch.int32, device=device) + static_output = torch.empty(1, top_k, dtype=torch.int32, device=device) + + graph_indices = torch.empty(1, 10, top_k, dtype=torch.int32, device=device) + graph_values = torch.empty(1, 10, top_k, device=device) + graph_pointer = graph_indices.data_ptr() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + top_k_module( + static_scores, + static_output, + is_prefill=False, + sequence_lengths=static_lengths, + scan_lengths=static_lengths, + radix_aux_indices=graph_indices, + radix_aux_logits=graph_values, + ) + + eager_scores = torch.randn(2, num_columns, device=device) + eager_lengths = torch.full((2,), num_columns, dtype=torch.int32, device=device) + eager_output = torch.empty(2, top_k, dtype=torch.int32, device=device) + top_k_module( + eager_scores, + eager_output, + is_prefill=False, + sequence_lengths=eager_lengths, + scan_lengths=eager_lengths, + radix_aux_indices=torch.empty(2, 10, top_k, dtype=torch.int32, device=device), + radix_aux_logits=torch.empty(2, 10, top_k, device=device), ) + assert graph_indices.data_ptr() == graph_pointer + graph.replay() + torch.cuda.synchronize() + expected = torch.topk(static_scores, top_k, dim=-1).values.sort(dim=-1).values + actual = static_scores.gather(1, static_output.long()).sort(dim=-1).values + torch.testing.assert_close(actual, expected) + def test_cute_dsl_prefill_dispatches_to_blackwell_kernel(monkeypatch) -> None: prefill = Mock()