[None][refactor] Modularize sparse Top-K selection

cpp/tensorrt_llm/thop/IndexerTopKOp.cpp

@@ -127,9 +127,8 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t
         if (radix_aux_indices.has_value() && radix_aux_logits.has_value())
         {
             // Caller-owned scratch with stable address (CUDA Graph safe;
-            // matches the heuristic_scratch convention noted above). All
-            // in-tree callers under CUDA Graph capture go through dsa.py's
-            // DSAtrtllmAttentionMetadata which always pre-allocates these.
+            // matches the heuristic_scratch convention noted above). The
+            // Python TopK module supplies these from its reusable buffer arena.
             auto const& ai = radix_aux_indices.value();
             auto const& al = radix_aux_logits.value();
             TORCH_CHECK(ai.is_cuda() && al.is_cuda(), "radix_aux_{indices,logits} must be CUDA tensors");

tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py

@@ -20,7 +20,6 @@
     rotate_activation,
     split_prefill_chunks,
     transform_local_topk_and_prepare_pool_view,
-    warmup_heuristic_topk_decode,
 )
 from .metadata import DSAtrtllmAttentionMetadata, build_req_idx_per_token
 from .params import DSABackendForwardArgs, DSAMetadataParams, DSAParams
@@ -49,5 +48,4 @@
     "rotate_activation",
     "split_prefill_chunks",
     "transform_local_topk_and_prepare_pool_view",
-    "warmup_heuristic_topk_decode",
 ]

tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py

@@ -31,14 +31,8 @@
 
 ModelConfig = tensorrt_llm.bindings.ModelConfig
 
-# dtype of the indexer MQA-logits that feed the top-k. All paged_mqa_logits
-# paths produce fp32 today (DSL fp8/fp4 default output_dtype=fp32; DeepGEMM
-# fp8 hardcodes kFloat; DeepGEMM fp4 defaults logits_dtype=kFloat32 and is not
-# overridden here), and the decode forward feeds logits to the top-k without a
-# cast. dtype is a top-k compile-key dimension, so the warmup pre-compiles for
-# exactly this value. If a paged_mqa_logits caller ever emits a different dtype
-# (e.g. overriding the DeepGEMM fp4 logits_dtype to bf16), update this constant
-# or the warmup silently compiles the wrong variant.
+# 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
 
 if TYPE_CHECKING:
@@ -58,8 +52,6 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata):
     """Attention metadata for DSA (Dense Sparse Attention) with indexer state."""
 
     sparse_metadata_params: Optional[DSAMetadataParams] = None
-    # Store reference to indexer for preparation stage
-    indexer: Optional["Indexer"] = None
     # Chunked prefill metadata for indexer (prefill-only, no CUDA graph needed)
     indexer_prefill_chunks: Optional[List[IndexerPrefillChunkMetadata]] = None
     # Max chunk size for two-level chunking:
@@ -160,6 +152,9 @@ def __post_init__(self):
         self.use_cute_dsl_topk = (
             sparse_metadata_params.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE
         )
+        self.enable_gvr_topk = (
+            sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100
+        )
         self.kv_lens_row_reorder = None
         capture_graph = self.is_cuda_graph
         # Plain DSA has no compression and uses the default [1]. DeepSeek-V4's
@@ -313,38 +308,17 @@ def get_indexer_max_seq_len(self) -> int:
         return max(1, self.kv_cache_manager.max_seq_len // self._indexer_compress_ratio)
 
     def warmup_cute_dsl_radix_topk(self, next_n: int) -> None:
-        """Pre-compile the radix-filter CuTe DSL decode top-k during warmup.
-
-        Eager decode iters (mixed prefill+decode batch, or ``cuda_graph``
-        disabled) whose ``num_rows`` lands in a ``cluster_size`` band that
-        graph capture did not exercise otherwise pay a first-touch JIT stall
-        on a live request. ``num_cols`` is fixed at ``indexer_max_seq_len``,
-        so only the ``cluster_size`` dimension needs sweeping; delegate to the
-        custom-op warmup helper, which owns the band enumeration.
-
-        ``next_n`` (a compile-key dimension) is supplied by the caller from
-        the engine's static spec-decode config.
-
-        No-op unless decode actually routes to
-        ``cute_dsl_indexer_topk_decode``: heuristic top-k uses the GVR kernel
-        and plain (no cute_dsl_topk) decode uses the C++ op. Called once from
-        ``ModelEngine.warmup``.
-        """
-        if not self.use_cute_dsl_topk or self.enable_heuristic_topk:
+        """Pre-compile CuTe DSL radix variants not covered by engine warmup."""
+        sparse_params = self.sparse_metadata_params
+        if not self.use_cute_dsl_topk or (
+            sparse_params.enable_heuristic_topk and get_sm_version() >= 100
+        ):
             return
         if self.kv_cache_manager is None:
             return
-        top_k = getattr(self.sparse_metadata_params, "index_topk", None)
+        top_k = self.sparse_mla_topk
         if not top_k:
             return
-        # The radix-filter DSL kernel does not support a compressed indexer
-        # combined with multi-row MTP: decode dispatches to it only when
-        # compress_ratio == 1 or next_n == 1. The compress_ratio > 1 &&
-        # next_n > 1 case routes to the C++ op (or GVR when heuristic top-k is
-        # on), so there is nothing to pre-compile here.
-        # TODO: extending the radix-filter path to compress_ratio > 1 &&
-        # next_n > 1 is straightforward; once the dispatch above is relaxed to
-        # use it there, drop this guard so the case is pre-compiled too.
         if self._indexer_compress_ratio > 1 and next_n > 1:
             return
         try:
@@ -353,6 +327,7 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None:
             )
         except ImportError:
             return
+
         warmup_cute_dsl_radix_topk_decode(
             top_k=int(top_k),
             num_cols=int(self.get_indexer_max_seq_len()),
@@ -480,11 +455,11 @@ def on_update_kv_lens(self):
         self._compute_kv_lens_row_reorder()
         self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True)
 
-    def _compute_kv_lens_row_reorder(self):
-        """Prepare longest-job-first row order for GVR top-k."""
+    def _compute_kv_lens_row_reorder(self) -> None:
+        """Prepare the longest-job-first GVR row order once per forward step."""
         next_n = 1 + self.max_draft_tokens
         if (
-            self.enable_heuristic_topk
+            self.enable_gvr_topk
             and self.use_cute_dsl_topk
             and self.num_generations * next_n >= 2 * self.num_sms
         ):
@@ -560,37 +535,6 @@ def create_buffers_for_mla_rope_append(self, capture_graph=False):
             pin_memory=prefer_pinned(),
         )
 
-    def _create_radix_aux_buffers(self, capture_graph=False):
-        # Persistent scratch for Radix-split-work indexer path (blocks_per_row > 1).
-        # Mirrors the fix the Heuristic path applied: per-call th::empty inside
-        # indexer_topk_decode produces stale pointers under CUDA Graph replay when
-        # the caching allocator is perturbed by chunked prefill at high CONC.
-        # Sized to the worst case kMaxBlocksPerRowDecode=10 from
-        # cpp/tensorrt_llm/kernels/indexerTopK.cu, times the max number of
-        # generation rows (num_seqs * (1 + max_draft_tokens)); the cpp op aborts
-        # if this is smaller than num_rows*blocks_per_row*index_topk. Allocated
-        # unconditionally: even with enable_heuristic_topk=True the dispatcher can
-        # fall back to Radix when canUseHeuristic returns False (small numColumns,
-        # etc.). MUST be re-created whenever max_draft_tokens changes (see
-        # update_spec_dec_param) or it is left too small once MTP raises the
-        # generation-row count.
-        _radix_max_blocks_per_row = 10
-        _radix_max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens)
-        self.radix_aux_indices = self.get_empty(
-            self.cuda_graph_buffers,
-            (_radix_max_gen_tokens, _radix_max_blocks_per_row, self.num_sparse_topk),
-            cache_name="radix_aux_indices",
-            dtype=torch.int32,
-            capture_graph=capture_graph,
-        )
-        self.radix_aux_logits = self.get_empty(
-            self.cuda_graph_buffers,
-            (_radix_max_gen_tokens, _radix_max_blocks_per_row, self.num_sparse_topk),
-            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):
@@ -695,7 +639,8 @@ def create_buffers_for_indexer(self, capture_graph=False):
                 pin_memory=prefer_pinned(),
             )
         # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation.
-        # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory.
+        # Allocate these buffers dynamically in Indexer.prepare()
+        # based on the actual total_kv_len to save memory.
         if self.enable_context_mla_with_cached_kv:
             self.slot_mapping_fp8_fullkv = None
             self.slot_mapping_scale_fullkv = None
@@ -786,37 +731,19 @@ def create_buffers_for_indexer(self, capture_graph=False):
                 device="cpu",
                 pin_memory=prefer_pinned(),
             )
-        # Per-layer persistent buffers for heuristic TopK pre_idx.
-        # Indexed by [local_layer_idx, generation_position, :].
-        # The graph captures reads/writes on these stable-address buffers;
-        # each replay's write becomes the next replay's read (feedback loop).
-        self.enable_heuristic_topk = (
-            sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100
-        )
-        if self.enable_heuristic_topk:
-            num_local_layers = self.kv_cache_manager.num_local_layers
-            self.heuristic_prev_topk = self.get_empty(
+        if self.enable_gvr_topk:
+            self.gvr_prior_indices = self.get_empty(
                 self.cuda_graph_buffers,
-                (num_local_layers, self.max_num_sequences, self.num_sparse_topk),
-                cache_name="heuristic_prev_topk",
+                (
+                    self.kv_cache_manager.num_local_layers,
+                    self.max_num_sequences,
+                    self.num_sparse_topk,
+                ),
+                cache_name="gvr_prior_indices",
                 dtype=torch.int32,
                 capture_graph=capture_graph,
             )
-            # Zero-initialize so the first decode step's pre_idx (kernel
-            # adds +1 offset) points to index 1 — a valid but benign candidate.
-            # Without this, uninitialized memory produces random hint indices.
-            self.heuristic_prev_topk.zero_()
-            # The C++ top-k path needs a stable scratch address.
-            if not self.use_cute_dsl_topk:
-                max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens)
-                self.heuristic_scratch_values = self.get_empty(
-                    self.cuda_graph_buffers,
-                    (max_gen_tokens, self.num_sparse_topk),
-                    cache_name="heuristic_scratch_values",
-                    dtype=torch.float32,
-                    capture_graph=capture_graph,
-                )
-            # GVR row order also needs a stable address for CUDA graphs.
+            self.gvr_prior_indices.zero_()
             if self.use_cute_dsl_topk:
                 self.kv_lens_row_reorder_buffer = self.get_empty(
                     self.cuda_graph_buffers,
@@ -825,12 +752,6 @@ def create_buffers_for_indexer(self, capture_graph=False):
                     dtype=torch.int32,
                     capture_graph=capture_graph,
                 )
-
-        # Persistent scratch for the Radix-split-work indexer path. Re-created
-        # in update_spec_dec_param when max_draft_tokens changes so it stays
-        # large enough for the MTP generation-row count.
-        self._create_radix_aux_buffers(capture_graph=capture_graph)
-
         # Create expanded buffers for MTP support
         self.create_expanded_buffers(capture_graph=capture_graph)
 
@@ -938,23 +859,6 @@ 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)
-            # Resize heuristic scratch buffer for new max_draft_tokens.
-            if self.enable_heuristic_topk and not self.use_cute_dsl_topk:
-                max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens)
-                self.heuristic_scratch_values = self.get_empty(
-                    self.cuda_graph_buffers,
-                    (max_gen_tokens, self.num_sparse_topk),
-                    cache_name="heuristic_scratch_values",
-                    dtype=torch.float32,
-                    capture_graph=capture_graph,
-                )
-            # The Radix-split-work scratch (radix_aux_*) is sized the same way
-            # (num_seqs * (1 + max_draft_tokens) rows) and is allocated
-            # unconditionally, so it must be resized here too -- otherwise the
-            # cpp indexer_topk_decode op aborts once MTP raises max_draft_tokens
-            # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk
-            # elements").
-            self._create_radix_aux_buffers(capture_graph=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/kv_cache_compression/triattention/triattention.py

@@ -36,6 +36,7 @@
 from tensorrt_llm.logger import logger
 
 from ...distributed import allgather
+from ...modules.top_k import TopK, TopKImplementation
 from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
 from ...pyexecutor.llm_request import LlmRequestState
 from ...pyexecutor.resource_manager import KVCacheCompressionManager
@@ -632,12 +633,13 @@ def _select_kept_ordinals(self, request_count: int) -> None:
         """Select top-k tokens and settle score ties into kept-ordinal rows."""
         rows = request_count * self._selection_rows_per_request
         # The trailing 1 is next_n: decode scores one query token per request.
-        torch.ops.trtllm.cute_dsl_indexer_topk_decode(
+        self._selection_top_k(
             self._selection_scores_rows[:rows],
-            self._selection_row_lengths[:rows],
             self._provisional_rows[:rows],
-            self.budget,
-            1,
+            is_prefill=False,
+            sequence_lengths=self._selection_row_lengths[:rows],
+            scan_lengths=self._selection_row_lengths[:rows],
+            next_n=1,
         )
         settle_ties(
             self._selection_scores_rows,
@@ -799,6 +801,10 @@ def _allocate_metadata_buffers(
 
     def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> None:
         """Allocate fixed manager-lifetime TopK inputs and outputs."""
+        self._selection_top_k = TopK(
+            self.budget,
+            decode_implementation=TopKImplementation.CUTE_DSL_RADIX,
+        )
         request_capacity = self._request_capacity
         selection_width = self._selection_width_capacity
         union = self.eviction_mode == "union"