From f74da7e1c56a3a2f2b50dd1730b97750d4ece61f Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:15:08 +0000 Subject: [PATCH 1/7] [None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e Wire the existing CuTe DSL GVR (Guess-Verify-Refine) top-k decode op into the DSA indexer decode dispatch. The flag combo use_cute_dsl_topk=True AND enable_heuristic_topk=True now routes to cute_dsl_gvr_topk_decode, covering all compress_ratio and next_n values. - DSAMetadataParams: add use_cute_dsl_topk; propagate through all three to_sparse_metadata_params / make_deepseek_v4_sparse_metadata_params sites. - Metadata prepare(): compute the LJF row-reorder (argsort of gen kv_lens, descending) only when num_generations >= num_sms, writing into a pre-allocated stable buffer so CUDA Graph replay reads a valid permutation instead of a freed fresh allocation. - Indexer decode dispatch: add GVR DSL as the first branch; keep the radix DSL and C++ scheduler branches as fallbacks. - Skip the redundant radix warmup when the GVR path is active. - Integration test: add cute_dsl_gvr_mtp1 case to test_fp8_blockscale. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../sparse/deepseek_v4/deepseek_v4.py | 6 + .../_torch/attention_backend/sparse/dsa.py | 103 ++++++++++++++---- tensorrt_llm/llmapi/llm_args.py | 2 + .../defs/accuracy/test_llm_api_pytorch.py | 45 +++++--- 4 files changed, 118 insertions(+), 38 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 52f16c84ffe9..89eef4ba4f2d 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/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, @@ -772,6 +773,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, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index af85390cf72e..358dddc03e82 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/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, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index cc6a66e53a08..f4b498aa8d34 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -954,6 +954,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, ) @@ -1062,6 +1063,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, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 2305c0852664..91be418775a5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3318,24 +3318,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) @@ -3362,16 +3372,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: From d8df5a25a87437b70c908086213bd11d39155b9d Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:22:56 +0000 Subject: [PATCH 2/7] [None][fix] fix GVR row-reorder threshold and add unit test Gate _compute_kv_lens_row_reorder on num_generations * next_n >= 2 * num_sms instead of num_generations >= num_sms, aligning with the kernel-side tuning note that records the win region starting at num_rows >= 2 * num_sms. The old threshold was only correct for next_n == 2; for next_n == 1 it engaged inside the measured regression band, and for next_n == 4 it missed the win region between 2*num_sms and 4*num_sms rows. Add test_kv_lens_row_reorder_threshold to cover the threshold boundary and verify the descending argsort output. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 9 +++- .../cute_dsl_kernels/blackwell/utils.py | 1 + .../attention/sparse/dsa/test_dsa_indexer.py | 52 +++++++++++++++++++ 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 358dddc03e82..ccac5567e292 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -873,8 +873,15 @@ def _compute_kv_lens_row_reorder(self): 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 >= self.num_sms): + 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) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py index 795934f776f1..1d969a1d417f 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py @@ -208,6 +208,7 @@ def fmin(a: Union[float, cutlass.Float32], ip=None) -> cutlass.Float32: return cutlass.Float32( nvvm.fmin( + T.f32(), cutlass.Float32(a).ir_value(loc=loc, ip=ip), cutlass.Float32(b).ir_value(loc=loc, ip=ip), nan=nan, 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 98247a7f1c79..cf4e9af27de4 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/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) From 8a277cc6c558eaa36121f9b4dda7208f64035041 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:27:27 +0000 Subject: [PATCH 3/7] [None][fix] remove redundant _compute_kv_lens_row_reorder call in V4 prepare() The reorder is already triggered via on_update_kv_lens() -> super() -> _compute_kv_lens_row_reorder() every forward step since _preprocess_inputs() calls on_update_kv_lens() unconditionally before prepare(). The explicit call in DeepseekV4TrtllmAttentionMetadata.prepare() was redundant and caused the argsort + copy to run twice per step when active. Also fix the docstring: _compute_kv_lens_row_reorder is called from on_update_kv_lens(), not from prepare(). Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../sparse/deepseek_v4/deepseek_v4.py | 5 ----- tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 10 +++++----- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 8282660c7af1..e003d44ff3c8 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -777,11 +777,6 @@ 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, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index ccac5567e292..e41690de1d0a 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -867,11 +867,11 @@ def _compute_kv_lens_row_reorder(self): 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. + 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 From a44d9acba6955bd8c381010316567825321e4320 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:29:19 +0000 Subject: [PATCH 4/7] [None][fix] gate metadata use_cute_dsl_topk with IS_CUTLASS_DSL_AVAILABLE Indexer.use_cute_dsl_topk is gated by IS_CUTLASS_DSL_AVAILABLE, but the metadata copy was the raw flag. When DSL is unavailable, this mismatch caused metadata to skip heuristic_scratch_values allocation while Indexer took the C++ indexer_topk_decode path, which silently falls back to radix when heuristicScratch is nullptr. Derive the metadata flag from the same gated value to keep them in sync. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index e41690de1d0a..f422b1da1ec3 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -674,7 +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.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 From 1b67d1507f61160f61c7d40efb15481db366441f Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:33:35 +0000 Subject: [PATCH 5/7] [None][test] add cute_dsl_gvr_mtp1 and cute_dsl_gvr_mtp3 test cases Add two integration test cases for the GVR DSL top-k path: - cute_dsl_gvr_mtp1 (next_n=2): basic GVR coverage - cute_dsl_gvr_mtp3 (next_n=4): covers next_n > 2, backing the PR claim that GVR supports all next_n values (unlike the radix DSL path) Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 68ab07c7c55f..dc67121c97b1 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3361,11 +3361,13 @@ class TestDeepSeekV32(LlmapiAccuracyTestHarness): False), (8, 1, 8, 1, False, True, True, True, 24, "_DEFAULT", False, True, True), + (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", - "cute_dsl_gvr_mtp1" + "cute_dsl_gvr_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, From b782847c66ae25e1cb45912f04a17158edc44aea Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:35:34 +0000 Subject: [PATCH 6/7] [None][test] replace cute_dsl_gvr_mtp1 with cute_dsl_gvr_mtp3 mtp3 (next_n=4) covers next_n > 2, directly backing the PR claim that GVR supports all next_n values. mtp1 (next_n=2) is redundant given mtp3 already exercises a strictly larger kernel variant. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index dc67121c97b1..e65e59653fa8 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3359,15 +3359,13 @@ class TestDeepSeekV32(LlmapiAccuracyTestHarness): 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), (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", - "cute_dsl_gvr_mtp1", "cute_dsl_gvr_mtp3" + "cute_dsl_gvr_mtp3" ]) def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, attention_dp, cuda_graph, overlap_scheduler, From 329dd6da165def69714c41dd8a5831a3ae6eddad Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:30:31 +0000 Subject: [PATCH 7/7] [None][fix] revert accidental T.f32() addition to nvvm.fmin in utils.py The T.f32() argument was accidentally bundled into the threshold-fix commit. Under the repo-pinned nvidia-cutlass-dsl==4.5.0 the nvvm.fmin signature is fmin(a, b, *, ...) with no leading result-type parameter, so T.f32() raises TypeError at JIT time. Revert to the original 4.5.0-compatible call. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py index 1d969a1d417f..795934f776f1 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py @@ -208,7 +208,6 @@ def fmin(a: Union[float, cutlass.Float32], ip=None) -> cutlass.Float32: return cutlass.Float32( nvvm.fmin( - T.f32(), cutlass.Float32(a).ir_value(loc=loc, ip=ip), cutlass.Float32(b).ir_value(loc=loc, ip=ip), nan=nan,