[None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py

@@ -250,6 +250,7 @@ def make_deepseek_v4_sparse_metadata_params(
         ),
         enable_indexer_skip=sparse_attention_config.skip_indexer_for_short_seqs,
         enable_heuristic_topk=sparse_attention_config.enable_heuristic_topk,
+        use_cute_dsl_topk=sparse_attention_config.use_cute_dsl_topk,
         use_cute_dsl_paged_mqa_logits=(sparse_attention_config.use_cute_dsl_paged_mqa_logits),
         q_split_threshold=sparse_attention_config.q_split_threshold,
         compress_ratios=sparse_attention_config.compress_ratios,

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -90,6 +90,7 @@ class DSAMetadataParams(SparseMetadataParams):
     index_head_dim: int
     enable_indexer_skip: bool
     enable_heuristic_topk: bool
+    use_cute_dsl_topk: bool
     use_cute_dsl_paged_mqa_logits: bool
     q_split_threshold: int
 
@@ -673,6 +674,9 @@ def __post_init__(self):
         self.indexer_head_dim = sparse_metadata_params.index_head_dim
         self.indexer_quant_block_size = 128
         self.enable_indexer_skip = (sparse_metadata_params.enable_indexer_skip)
+        self.use_cute_dsl_topk = (sparse_metadata_params.use_cute_dsl_topk
+                                  and IS_CUTLASS_DSL_AVAILABLE)
+        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
         # metadata params carry the model-specific compression ratios.
@@ -856,8 +860,37 @@ def on_update_kv_lens(self):
                     _DG_SCHEDULE_BLOCK_KV, self.num_sms)
                 self.scheduler_metadata_buffer_expanded.copy_(
                     scheduler_metadata_buffer_expanded, non_blocking=True)
+        self._compute_kv_lens_row_reorder()
         self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True)
 
+    def _compute_kv_lens_row_reorder(self):
+        """LJF (longest-job-first) row-reorder for the GVR DSL top-k path.
+
+        Writes ``argsort(gen_kv_lens, descending)`` into the stable buffer when
+        the multi-wave threshold is met, otherwise leaves ``order_row`` None.
+        Called from ``on_update_kv_lens()`` (both base and DeepSeek-V4 via
+        super()) unconditionally every forward step so the GVR op sees a fresh
+        valid permutation and never a stale one from a prior step.  Copies into
+        the stable buffer (not a fresh tensor) so the CUDA-Graph-captured op
+        reads a valid permutation on every replay.
+        """
+        # Gate on row count (num_generations * next_n) rather than request count
+        # so the threshold aligns with the kernel-side tuning note that records
+        # the win region starting at num_rows >= 2 * num_sms.  Using
+        # num_generations alone is only correct for next_n == 2; for next_n == 1
+        # it engages inside the measured regression band, and for next_n == 4 it
+        # misses the win region between 2*num_sms and 4*num_sms rows.
+        next_n = 1 + self.max_draft_tokens
+        if (self.enable_heuristic_topk and self.use_cute_dsl_topk
+                and self.num_generations * next_n >= 2 * self.num_sms):
+            gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.num_seqs]
+            order = torch.argsort(gen_kv_lens, descending=True).to(torch.int32)
+            self.kv_lens_row_reorder_buffer[:self.num_generations].copy_(order)
+            self.kv_lens_row_reorder = \
+                self.kv_lens_row_reorder_buffer[:self.num_generations]
+        else:
+            self.kv_lens_row_reorder = None
+
     def update_for_spec_dec(self):
         super().update_for_spec_dec()
         # host
@@ -1105,15 +1138,30 @@ def create_buffers_for_indexer(self, capture_graph=False):
             # Pre-allocated with stable address for CUDA Graph compatibility
             # (replaces cudaMallocAsync/cudaFreeAsync inside the kernel launcher).
             # Shape: [max_gen_tokens, topK] where max_gen_tokens = max_batch * (1 + max_draft).
-            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,
-            )
+            # Only the C++ indexer_topk_decode path consumes it; the GVR DSL
+            # path does not, so skip the allocation when use_cute_dsl_topk.
+            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,
+                )
+            # Stable-address buffer for the GVR DSL LJF row-reorder
+            # (order_row = argsort(gen_kv_lens, descending)). Must not be
+            # fresh-allocated per step: under CUDA Graph the captured op reads
+            # a frozen address, so prepare() copies into this buffer instead.
+            if self.use_cute_dsl_topk:
+                self.kv_lens_row_reorder_buffer = self.get_empty(
+                    self.cuda_graph_buffers,
+                    (self.max_num_sequences, ),
+                    cache_name="kv_lens_row_reorder_buffer",
+                    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
@@ -1210,7 +1258,9 @@ def update_spec_dec_param(
         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:
+            # Skip when use_cute_dsl_topk (GVR path never consumes it), matching
+            # the allocation guard in create_buffers_for_indexer.
+            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(
@@ -1764,7 +1814,7 @@ def __init__(self,
                 or self.use_cute_dsl_paged_mqa_logits) and layer_idx == 0:
             from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops
 
-            if self.use_cute_dsl_topk:
+            if self.use_cute_dsl_topk and not self._enable_heuristic_topk:
                 # the dtype of topk input tensor, which is float32 now.
                 # Note, need to update it if the dtype of topk input tensor is changed.
                 cute_dsl_custom_ops.warmup_cute_dsl_indexer_topk(
@@ -2750,17 +2800,34 @@ def sparse_attn_indexer(
                     # handled inside the C++ kernel (preIdxOffset += 1).
                     pre_idx = metadata.heuristic_prev_topk[
                         local_layer, :num_generations]
-                    heuristic_scratch = \
-                        metadata.heuristic_scratch_values[
-                            :num_gen_tokens]
-
-                # CuTE DSL top-k allocates O(num_gen_tokens * kv_len) global
-                # memory. Beyond 256 tokens the extra memory becomes significant,
-                # so we cap it at 256 for now and fall back to the CUDA C++
-                # indexer_topk_decode. This limit can be removed if GPU memory
-                # is not a bottleneck.
-                if (self.use_cute_dsl_topk and num_gen_tokens <= 256
-                        and (self.compress_ratio == 1 or next_n == 1)):
+                    # heuristic_scratch is only consumed by the C++
+                    # indexer_topk_decode path; the GVR DSL op does not take it.
+                    # Guard on the metadata flag so this stays consistent with
+                    # the buffer allocation (also gated on the same flag).
+                    if not metadata.use_cute_dsl_topk:
+                        heuristic_scratch = \
+                            metadata.heuristic_scratch_values[
+                                :num_gen_tokens]
+
+                if self.use_cute_dsl_topk and self._enable_heuristic_topk:
+                    # GVR DSL: supports all compress_ratio and next_n values.
+                    torch.ops.trtllm.cute_dsl_gvr_topk_decode(
+                        logits_decode,
+                        pre_idx,
+                        gen_kv_lens_cuda,
+                        topk_indices_buffer[num_ctx_tokens:num_ctx_tokens +
+                                            num_gen_tokens, :],
+                        self.index_topk,
+                        next_n=next_n,
+                        compress_ratio=self.compress_ratio,
+                        max_seq_len=indexer_max_seq_len,
+                        order_row=metadata.kv_lens_row_reorder,
+                    )
+                # CuTE DSL radix top-k allocates O(num_gen_tokens * kv_len)
+                # global memory. Beyond 256 tokens the extra memory becomes
+                # significant, so we cap it at 256 and fall back to C++.
+                elif (self.use_cute_dsl_topk and num_gen_tokens <= 256
+                      and (self.compress_ratio == 1 or next_n == 1)):
                     torch.ops.trtllm.cute_dsl_indexer_topk_decode(
                         logits_decode, context_lens
                         if self.compress_ratio > 1 else gen_kv_lens_cuda,

tensorrt_llm/llmapi/llm_args.py

@@ -1014,6 +1014,7 @@ def _value(name: str, default=None):
             index_head_dim=_value("index_head_dim", 128),
             enable_indexer_skip=self.skip_indexer_for_short_seqs,
             enable_heuristic_topk=self.enable_heuristic_topk,
+            use_cute_dsl_topk=self.use_cute_dsl_topk,
             use_cute_dsl_paged_mqa_logits=(self.use_cute_dsl_paged_mqa_logits),
             q_split_threshold=self.q_split_threshold,
         )
@@ -1122,6 +1123,7 @@ def _value(name: str, default=None):
             index_head_dim=_value("index_head_dim", 128),
             enable_indexer_skip=self.skip_indexer_for_short_seqs,
             enable_heuristic_topk=self.enable_heuristic_topk,
+            use_cute_dsl_topk=self.use_cute_dsl_topk,
             use_cute_dsl_paged_mqa_logits=(self.use_cute_dsl_paged_mqa_logits),
             q_split_threshold=self.q_split_threshold,
             compress_ratios=self.compress_ratios,

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -3343,24 +3343,34 @@ class TestDeepSeekV32(LlmapiAccuracyTestHarness):
     @skip_pre_hopper
     @pytest.mark.skip_less_device_memory(140000)
     @pytest.mark.parametrize(
-        "tp_size,pp_size,ep_size,mtp_nextn,fp8kv,attention_dp,cuda_graph,overlap_scheduler,max_batch_size,moe_backend,disable_skip_indexer,enable_heuristic_topk",
+        "tp_size,pp_size,ep_size,mtp_nextn,fp8kv,attention_dp,cuda_graph,overlap_scheduler,max_batch_size,moe_backend,disable_skip_indexer,enable_heuristic_topk,use_cute_dsl_topk",
         [
-            (8, 1, 8, 0, False, True, True, True, 24, "_DEFAULT", False, False),
-            (8, 1, 8, 1, False, True, True, True, 24, "_DEFAULT", False, False),
-            (8, 1, 8, 0, True, True, True, True, 24, "_DEFAULT", False, False),
-            (8, 1, 8, 3, False, False, True, True, 1, "TRTLLM", False, False),
-            (8, 1, 8, 3, False, False, True, True, 1, "_DEFAULT", False, False),
-            (8, 1, 8, 1, False, True, True, True, 24, "_DEFAULT", True, False),
-            (8, 1, 8, 1, False, True, True, True, 24, "_DEFAULT", False, True),
+            (8, 1, 8, 0, False, True, True, True, 24, "_DEFAULT", False, False,
+             False),
+            (8, 1, 8, 1, False, True, True, True, 24, "_DEFAULT", False, False,
+             False),
+            (8, 1, 8, 0, True, True, True, True, 24, "_DEFAULT", False, False,
+             False),
+            (8, 1, 8, 3, False, False, True, True, 1, "TRTLLM", False, False,
+             False),
+            (8, 1, 8, 3, False, False, True, True, 1, "_DEFAULT", False, False,
+             False),
+            (8, 1, 8, 1, False, True, True, True, 24, "_DEFAULT", True, False,
+             False),
+            (8, 1, 8, 1, False, True, True, True, 24, "_DEFAULT", False, True,
+             False),
+            (8, 1, 8, 3, False, True, True, True, 24, "_DEFAULT", False, True,
+             True),
         ],
         ids=[
             "baseline", "baseline_mtp1", "baseline_fp8kv", "latency",
-            "latency_default", "disable_skip_indexer", "heuristic_topk_mtp1"
+            "latency_default", "disable_skip_indexer", "heuristic_topk_mtp1",
+            "cute_dsl_gvr_mtp3"
         ])
     def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv,
                             attention_dp, cuda_graph, overlap_scheduler,
                             max_batch_size, moe_backend, disable_skip_indexer,
-                            enable_heuristic_topk):
+                            enable_heuristic_topk, use_cute_dsl_topk):
         if get_sm_version() == 100 or get_sm_version() == 103:
             moe_backend = "DEEPGEMM" if moe_backend == "_DEFAULT" else moe_backend
             moe_config = MoeConfig(backend=moe_backend, max_num_tokens=16384)
@@ -3387,16 +3397,19 @@ def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv,
                 )
             kv_cache_config.dtype = "fp8"
 
-        if enable_heuristic_topk and get_sm_version() < 100:
+        if (enable_heuristic_topk
+                or use_cute_dsl_topk) and get_sm_version() < 100:
             pytest.skip("Heuristic TopK requires Blackwell (SM >= 100)")
 
-        dsa_config = None
+        dsa_kwargs = {}
         if disable_skip_indexer:
-            dsa_config = DeepSeekSparseAttentionConfig(
-                skip_indexer_for_short_seqs=False)
+            dsa_kwargs["skip_indexer_for_short_seqs"] = False
         if enable_heuristic_topk:
-            dsa_config = DeepSeekSparseAttentionConfig(
-                enable_heuristic_topk=True)
+            dsa_kwargs["enable_heuristic_topk"] = enable_heuristic_topk
+        if use_cute_dsl_topk:
+            dsa_kwargs["use_cute_dsl_topk"] = use_cute_dsl_topk
+        dsa_config = DeepSeekSparseAttentionConfig(
+            **dsa_kwargs) if dsa_kwargs else None
 
         mtp_config = None
         if mtp_nextn > 0:

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

@@ -3330,3 +3330,55 @@ def _run_indexer():
     assert "indexer_topk_out_buffer" in metadata.cuda_graph_buffers.buffers, (
         "indexer topk-output buffer must be drawn from the cuda_graph_buffers arena"
     )
+
+
+def test_kv_lens_row_reorder_threshold():
+    """_compute_kv_lens_row_reorder engages iff num_generations * next_n >= 2 * num_sms,
+    and produces a descending argsort of gen_kv_lens when active."""
+    num_sms = 16  # small synthetic value; threshold = 2 * 16 = 32 rows
+    next_n = 2  # max_draft_tokens=1 → next_n = 1 + 1 = 2
+
+    def make_mock(num_generations, kv_lens_list):
+        kv_cuda = torch.tensor(kv_lens_list, dtype=torch.int32, device="cuda")
+        buf = torch.zeros(64, dtype=torch.int32, device="cuda")
+        ns = SimpleNamespace(
+            enable_heuristic_topk=True,
+            use_cute_dsl_topk=True,
+            num_generations=num_generations,
+            num_sms=num_sms,
+            max_draft_tokens=next_n - 1,
+            num_contexts=0,
+            num_seqs=num_generations,
+            kv_lens_cuda=kv_cuda,
+            kv_lens_row_reorder_buffer=buf,
+            kv_lens_row_reorder=None,
+        )
+        ns._compute_kv_lens_row_reorder = (
+            lambda: DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(ns)
+        )
+        return ns
+
+    # Fixed unsorted sequence for deterministic sort verification (len == num_sms)
+    kv_vals = [4, 1, 8, 2, 16, 3, 12, 6, 7, 9, 5, 11, 13, 10, 14, 15]
+
+    # Below threshold: 1 * 2 = 2 < 32 → None
+    md_below = make_mock(1, [1000])
+    md_below._compute_kv_lens_row_reorder()
+    assert md_below.kv_lens_row_reorder is None
+
+    # At threshold: num_sms * 2 = 32 → engages, verify descending argsort
+    md_at = make_mock(num_sms, kv_vals)
+    md_at._compute_kv_lens_row_reorder()
+    assert md_at.kv_lens_row_reorder is not None
+    reorder = md_at.kv_lens_row_reorder.cpu().tolist()
+    assert [kv_vals[i] for i in reorder] == sorted(kv_vals, reverse=True), (
+        "order_row must be a descending argsort of gen_kv_lens"
+    )
+
+    # Above threshold: (num_sms + 1) * 2 = 34 > 32 → also engages with correct sort
+    kv_vals2 = kv_vals + [100]
+    md_above = make_mock(num_sms + 1, kv_vals2)
+    md_above._compute_kv_lens_row_reorder()
+    assert md_above.kv_lens_row_reorder is not None
+    reorder2 = md_above.kv_lens_row_reorder.cpu().tolist()
+    assert [kv_vals2[i] for i in reorder2] == sorted(kv_vals2, reverse=True)