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 7caf1c6ecc28..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 @@ -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, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index af85390cf72e..f422b1da1ec3 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,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, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6cdd95827ccf..deb78938d9c5 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/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, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0bf4298efaad..e65e59653fa8 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/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: 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)