[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,
@@ -776,6 +777,11 @@ def prepare(self):
             self._compress_ratios_sorted,
         )
 
+        # LJF row-reorder for the GVR DSL top-k path. V4 does not chain the base
+        # DSAtrtllmAttentionMetadata.prepare(), so call the shared helper here so
+        # the reorder engages for DeepSeek-V4 too (not just V3.2).
+        self._compute_kv_lens_row_reorder()
+
     def prepare_compressed_kv_metadata(
         self,
         kv_lens: torch.Tensor,

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,8 @@ 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
+        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 +859,30 @@ 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 unconditionally from every ``prepare()`` (base and DeepSeek-V4)
+        so the GVR op sees a fresh valid permutation each step and never a
+        stale one left over from a prior decode step. Copies into the stable
+        buffer (not a fresh tensor) so the CUDA-Graph-captured op reads a valid
+        permutation on every replay.
+        """
+        if (self.enable_heuristic_topk and self.use_cute_dsl_topk
+                and self.num_generations >= 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 +1130,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 +1250,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 +1806,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 +2792,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, 1, 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_mtp1"
         ])
     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: