From 7c92a3bcbe761c8ae048b19a650e22c3e2820fb8 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:29:23 +0000 Subject: [PATCH 1/6] [None][feat] Two-level GVR decode top-K dispatch from the sparse-attention config enable_heuristic_topk keeps selecting the GVR family over the exact radix path; a new use_self_sampling_topk config field (default True) selects the hint-free self-sampling engine over the temporal-hint engines. The CUTE_DSL_GVR_V2 enum folds into CUTE_DSL_GVR behind a gvr_self_sampling module flag, TopK.needs_gvr_prior follows the two-level decision, and the retired TRTLLM_GVR_SELF_SAMPLING env only warns. The field threads llm_args -> model_config -> DSAParams/DSAMetadataParams -> indexer and the warmup mirror (whose top_k source also moves off a dead index_topk getattr to sparse_mla_topk). Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 51 ++++++++++++++----- .../attention_backend/sparse/dsa/metadata.py | 16 +++--- .../attention_backend/sparse/dsa/params.py | 5 ++ tensorrt_llm/_torch/model_config.py | 6 +++ tensorrt_llm/_torch/modules/top_k.py | 37 +++++++------- tensorrt_llm/llmapi/llm_args.py | 23 +++++++-- .../attention/sparse/dsa/test_dsa_indexer.py | 50 ++++++++++++++++++ tests/unittest/_torch/modules/test_top_k.py | 32 +++++++++--- 8 files changed, 172 insertions(+), 48 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 972f1160a53f..0884e831b327 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -690,16 +690,17 @@ def __init__( self._enable_heuristic_topk = ( sparse_params.enable_heuristic_topk and get_sm_version() >= 100 ) - # Opt-in self-sampling GVR top-K decode (CuTeDSL, env-gated: - # TRTLLM_GVR_SELF_SAMPLING=1). Same operator contract as the tiered - # heuristic path (per-request device kv_lens, raw prev-top-K hints, - # per-row MTP window, in-kernel n <= topK short path); tuning is - # frozen from indexer_max_seq_len at capture time, so the launch is - # CUDA-graph-replay safe. The TopK module's hardware-format gate - # falls through to the CUDA GVR path with a one-time warning; - # contract violations inside the engine raise. + # Two-level GVR dispatch: enable_heuristic_topk selects the GVR + # family over the exact radix path; use_self_sampling_topk (default + # True) selects the hint-free self-sampling engine over the + # temporal-hint engines. Tuning is frozen from indexer_max_seq_len + # at capture time, so the launch is CUDA-graph-replay safe. The + # TopK module's hardware-format gate falls back to the exact + # insertion/radix path with a one-time warning; contract violations + # inside the engine raise. self._use_self_sampling_topk = ( - os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") == "1" + sparse_params.use_self_sampling_topk + and self._enable_heuristic_topk and IS_CUTLASS_DSL_AVAILABLE # datacenter Blackwell only; consumer Blackwell (sm_120/121) # lacks thread-block clusters @@ -707,6 +708,29 @@ def __init__( and sparse_params.index_topk in (512, 1024, 2048) and compress_ratio in (1, 4) ) + if os.environ.get("TRTLLM_GVR_SELF_SAMPLING") is not None: + logger.warning_once( + "TRTLLM_GVR_SELF_SAMPLING is retired and ignored: the " + "self-sampling GVR engine is selected by the " + "use_self_sampling_topk sparse-attention config field " + "(default True) when enable_heuristic_topk is set.", + key="gvr_self_sampling_env_retired", + ) + if ( + self._enable_heuristic_topk + and sparse_params.use_self_sampling_topk + and not self._use_self_sampling_topk + ): + logger.warning_once( + "use_self_sampling_topk=True but the self-sampling GVR " + "prerequisites are not met " + f"(cutlass_dsl={IS_CUTLASS_DSL_AVAILABLE}, " + f"sm={get_sm_version()}, " + f"index_topk={sparse_params.index_topk}, " + f"compress_ratio={compress_ratio}); using the temporal GVR " + "path instead.", + key="gvr_self_sampling_prereq_fallback", + ) self.mtp_index_share = sparse_params.mtp_index_share if self.use_cute_dsl_topk: @@ -719,15 +743,16 @@ def __init__( decode_top_k_implementation = TopKImplementation.CUDA_GVR else: decode_top_k_implementation = TopKImplementation.CUDA_RADIX - if self._use_self_sampling_topk and self._enable_heuristic_topk: - # env opt-in overrides the decode implementation; the GVR prior - # contract is identical - decode_top_k_implementation = TopKImplementation.CUTE_DSL_GVR_V2 + if self._use_self_sampling_topk: + # The self-sampling engine overrides the temporal decode + # implementation regardless of use_cute_dsl_topk. + decode_top_k_implementation = TopKImplementation.CUTE_DSL_GVR self.top_k = TopK( self.index_topk, prefill_implementation=TopKImplementation.CUDA_RADIX, decode_implementation=decode_top_k_implementation, compress_ratio=self.compress_ratio, + gvr_self_sampling=self._use_self_sampling_topk, ) # GVR emission-assisted decode (opt-in, experimental): the FP4/FP8 # indexer epilogue emits candidates the GVR Top-K consumes (see diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index c93fe8c021dd..2b1c0521ebe6 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -188,6 +188,7 @@ def __post_init__(self): self.enable_gvr_topk = ( sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 ) + self.use_self_sampling_topk = sparse_metadata_params.use_self_sampling_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 @@ -382,18 +383,19 @@ def warmup_selfsampling_topk( warmed keys are the ones dispatch actually looks up. Batches outside this set still compile lazily on first touch. The helper enumerates one representative row per distinct engine compile key, so large - batch lists warm in bounded time and memory. No-op unless the opt-in - gate (TRTLLM_GVR_SELF_SAMPLING=1) selects the engine. + batch lists warm in bounded time and memory. No-op unless the + two-level dispatch (enable_heuristic_topk + use_self_sampling_topk) + selects the self-sampling engine. """ - if os.environ.get("TRTLLM_GVR_SELF_SAMPLING", "0") != "1": - return - # same hardware gates as the dispatch flag (indexer __init__): never - # compile these kernels on unsupported stacks during warmup + # same two-level dispatch and hardware gates as the indexer __init__: + # never compile these kernels on unsupported stacks during warmup if not IS_CUTLASS_DSL_AVAILABLE or get_sm_version() not in (100, 103): return if not self.enable_gvr_topk or self.kv_cache_manager is None: return - top_k = getattr(self.sparse_metadata_params, "index_topk", None) + if not self.use_self_sampling_topk: + return + top_k = self.sparse_mla_topk if not top_k or int(top_k) not in (512, 1024, 2048): return cr = int(self._indexer_compress_ratio) if self._indexer_compress_ratio else 1 diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py index d72f15c45b2d..78957f38542a 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py @@ -41,6 +41,7 @@ class DSAMetadataParams(SparseMetadataParams): q_split_threshold: int has_shared_indexer_layers: bool = False mtp_index_share: bool = False + use_self_sampling_topk: bool = True @dataclass(frozen=True) @@ -58,6 +59,10 @@ class DSAParams(SparseParams): q_split_threshold: int = 8192 indexer_rope_interleave: bool = False enable_heuristic_topk: bool = False + # Second-level GVR dispatch: hint-free self-sampling engine (True) vs + # temporal previous-step-hint engines (False). Only meaningful when + # enable_heuristic_topk is set. + use_self_sampling_topk: bool = True indexer_k_dtype: Literal["fp8", "fp4"] = "fp8" # Shared layers reuse the preceding full layer's top-k. is_full_indexer_layer: bool = True diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 74b6757628bd..2daf38752f3b 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -997,6 +997,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): q_split_threshold = sparse_attention_config.q_split_threshold indexer_rope_interleave = sparse_attention_config.indexer_rope_interleave enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk + use_self_sampling_topk = sparse_attention_config.use_self_sampling_topk indexer_k_dtype = sparse_attention_config.indexer_k_dtype else: index_n_heads = pretrained_config.index_n_heads @@ -1010,6 +1011,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): q_split_threshold = 8192 indexer_rope_interleave = False enable_heuristic_topk = False + use_self_sampling_topk = True default_sparse_attention_config = DeepSeekV4SparseAttentionConfig( ) indexer_k_dtype = default_sparse_attention_config.indexer_k_dtype @@ -1026,6 +1028,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): indexer_config['q_split_threshold'] = q_split_threshold indexer_config['indexer_rope_interleave'] = indexer_rope_interleave indexer_config['enable_heuristic_topk'] = enable_heuristic_topk + indexer_config['use_self_sampling_topk'] = use_self_sampling_topk indexer_config['indexer_k_dtype'] = indexer_k_dtype return indexer_config @@ -1063,6 +1066,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): use_cute_dsl_paged_mqa_logits = sparse_attention_config.use_cute_dsl_paged_mqa_logits q_split_threshold = sparse_attention_config.q_split_threshold enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk + use_self_sampling_topk = sparse_attention_config.use_self_sampling_topk indexer_k_dtype = sparse_attention_config.indexer_k_dtype index_share_for_mtp_iteration = sparse_attention_config.index_share_for_mtp_iteration else: @@ -1075,6 +1079,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): use_cute_dsl_paged_mqa_logits = False q_split_threshold = 8192 enable_heuristic_topk = False + use_self_sampling_topk = True indexer_k_dtype = "fp8" index_share_for_mtp_iteration = None kwargs[ @@ -1091,6 +1096,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): q_split_threshold=q_split_threshold, indexer_rope_interleave=indexer_rope_interleave, enable_heuristic_topk=enable_heuristic_topk, + use_self_sampling_topk=use_self_sampling_topk, indexer_k_dtype=indexer_k_dtype, index_share_for_mtp_iteration= index_share_for_mtp_iteration) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 6f323d839bf9..9a140d7fe35f 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -22,17 +22,11 @@ class TopKImplementation(str, Enum): CUTE_DSL_RADIX = "cute_dsl_radix" CUDA_GVR = "cuda_gvr" CUTE_DSL_GVR = "cute_dsl_gvr" - CUTE_DSL_GVR_V2 = "cute_dsl_gvr_v2" _GVR_IMPLEMENTATIONS = { TopKImplementation.CUDA_GVR, TopKImplementation.CUTE_DSL_GVR, - TopKImplementation.CUTE_DSL_GVR_V2, -} -_TEMPORAL_GVR_IMPLEMENTATIONS = { - TopKImplementation.CUDA_GVR, - TopKImplementation.CUTE_DSL_GVR, } _MAX_RADIX_BLOCKS_PER_ROW = 10 @@ -53,6 +47,7 @@ def __init__( prefill_implementation: TopKImplementation | None = None, decode_implementation: TopKImplementation | None = None, compress_ratio: int = 1, + gvr_self_sampling: bool = True, ) -> None: super().__init__() self.top_k = top_k @@ -63,10 +58,12 @@ def __init__( decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio - # emission-assisted GVR (opt-in via prepare_gvr_emission): the - # module owns the closed-loop emission state; the caller passes - # the returned kwargs to the scoring op, and the consume side is - # injected into the GVR Top-K call while the step stays armed + # Second-level GVR dispatch for CUTE_DSL_GVR: True selects the + # hint-free self-sampling engine, False the temporal-hint engine. + self.gvr_self_sampling = gvr_self_sampling + # emission-assisted GVR (opt-in via prepare_gvr_emission): the module + # owns the closed-loop emission state; only reachable on the temporal + # (gvr_self_sampling=False) V1 path. self._gvr_emission_state = None self._gvr_emission_route = None self._gvr_emission_armed = False @@ -74,7 +71,12 @@ def __init__( @property def needs_gvr_prior(self) -> bool: """Return whether decode consumes previous-step Top-K indices.""" - return self.decode_implementation in _TEMPORAL_GVR_IMPLEMENTATIONS + if self.decode_implementation == TopKImplementation.CUDA_GVR: + return True + return ( + self.decode_implementation == TopKImplementation.CUTE_DSL_GVR + and not self.gvr_self_sampling + ) def forward( self, @@ -103,10 +105,11 @@ def forward( next_n: Number of decode rows per request. max_seq_len: Maximum decode score width used for GVR kernel tuning. gvr_ext_kwargs: GVR-only keyword arguments. ``gvr_prior_indices`` - is required by the temporal CUDA and CuTe DSL GVR paths. It is + is required by the temporal GVR paths (``CUDA_GVR``, or + ``CUTE_DSL_GVR`` with ``gvr_self_sampling=False``). It is caller-owned int32 previous selection with shape - ``[num_requests, top_k]`` on ``scores.device``. GVR V2 does - not consume this state. + ``[num_requests, top_k]`` on ``scores.device``. The + self-sampling engine does not consume this state. ``gvr_row_order`` is an optional int32 request ordering with shape ``[num_requests]`` on the same device. @@ -279,7 +282,7 @@ def _forward_decode_gvr( gvr_prior_indices: torch.Tensor | None = None, gvr_row_order: torch.Tensor | None = None, ) -> torch.Tensor: - if self.decode_implementation == TopKImplementation.CUTE_DSL_GVR_V2: + if self.decode_implementation == TopKImplementation.CUTE_DSL_GVR and self.gvr_self_sampling: assert max_seq_len is not None if ( # engine hardware-format gate (falls through otherwise): @@ -306,7 +309,7 @@ def _forward_decode_gvr( f"next_n={next_n}, hint-free).", key="selfsampling_topk_engaged", ) - # Self-sampling GVR varlen engine (TRTLLM_GVR_SELF_SAMPLING=1): + # Self-sampling GVR varlen engine: # one launch for the batch; per-row n from device kv_lens, # capture-stable tuning from the max-seq-len engine constant # (no host reads — CUDA-graph safe). The module receives @@ -323,7 +326,7 @@ def _forward_decode_gvr( ) return output_indices logger.warning_once( - "TRTLLM_GVR_SELF_SAMPLING=1 but the decode scores do not " + "self-sampling GVR is selected but the decode scores do not " "satisfy the engine's hardware-format gate " f"(dtype={scores.dtype}, strides={tuple(scores.stride())}); " "falling back to the CUDA insertion/radix Top-K path.", diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 343cef0fbd63..b68cf55f0974 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -979,11 +979,20 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): default=False, description= "Whether to enable Guess-Verify-Refine (GVR) Top-K for the DSA decode " - "indexer. GVR reuses previous-step Top-K indices as hints to reduce " - "threshold search iterations. Currently supported for index_topk ∈ " - "{512, 1024, 2048} on Blackwell (SM100+), with compress_ratio ∈ {1, 4} " - "(DSv3.2 + DSv4 indexers). Falls back to the production insertion/" - "radix Top-K path when prerequisites are not met.") + "indexer instead of the exact insertion/radix Top-K path. Currently " + "supported for index_topk ∈ {512, 1024, 2048} on Blackwell (SM100+), " + "with compress_ratio ∈ {1, 4} (DSv3.2 + DSv4 indexers). Falls back to " + "the production insertion/radix Top-K path when prerequisites are not " + "met. `use_self_sampling_topk` selects the GVR engine generation.") + use_self_sampling_topk: bool = Field( + default=True, + description= + "Select the GVR engine generation when enable_heuristic_topk is set: " + "True (default) runs the hint-free self-sampling engine, which derives " + "its search bracket from the current row and keeps no cross-step " + "state; False runs the temporal-hint engines, which reuse the " + "previous decode step's Top-K indices as hints. Ignored when " + "enable_heuristic_topk is False.") indexer_k_dtype: Literal["fp8", "fp4"] = Field( default="fp8", description= @@ -1125,6 +1134,7 @@ def _value(name: str, default=None): q_split_threshold=self.q_split_threshold, indexer_rope_interleave=self.indexer_rope_interleave, enable_heuristic_topk=self.enable_heuristic_topk, + use_self_sampling_topk=self.use_self_sampling_topk, indexer_k_dtype=self.indexer_k_dtype, is_full_indexer_layer=self._is_full_indexer_layer( pretrained_config, kwargs.get("layer_idx")), @@ -1157,6 +1167,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_self_sampling_topk=self.use_self_sampling_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, @@ -1244,6 +1255,7 @@ def _value(name: str, default=None): q_split_threshold=self.q_split_threshold, indexer_rope_interleave=self.indexer_rope_interleave, enable_heuristic_topk=self.enable_heuristic_topk, + use_self_sampling_topk=self.use_self_sampling_topk, indexer_k_dtype=self.indexer_k_dtype, compress_ratios=self.compress_ratios, window_size=self.window_size, @@ -1269,6 +1281,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_self_sampling_topk=self.use_self_sampling_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, 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 a5bb5545f785..4de20744923a 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -476,6 +476,56 @@ def test_indexer_configures_one_top_k_module( assert isinstance(indexer.top_k, TopK) assert indexer.top_k.prefill_implementation == TopKImplementation.CUDA_RADIX assert indexer.top_k.decode_implementation == expected_decode + if enable_heuristic: + # index_topk=128 misses the self-sampling prerequisites, so the + # default use_self_sampling_topk=True falls back to the temporal path. + assert not indexer.top_k.gvr_self_sampling + assert indexer.top_k.needs_gvr_prior + + +@skip_pre_hopper +@pytest.mark.parametrize( + "use_self_sampling,use_cute_dsl,expected_decode", + [ + (True, False, TopKImplementation.CUTE_DSL_GVR), + (True, True, TopKImplementation.CUTE_DSL_GVR), + (False, True, TopKImplementation.CUTE_DSL_GVR), + (False, False, TopKImplementation.CUDA_GVR), + ], +) +def test_indexer_two_level_gvr_dispatch( + monkeypatch, + use_self_sampling, + use_cute_dsl, + expected_decode, +): + # The retired TRTLLM_GVR_SELF_SAMPLING env must be ignored: with + # use_self_sampling_topk=False the temporal path must win regardless. + monkeypatch.setenv("TRTLLM_GVR_SELF_SAMPLING", "1") + sparse_config = DeepSeekSparseAttentionConfig( + index_head_dim=128, + index_n_heads=32, + index_topk=512, + use_cute_dsl_topk=use_cute_dsl, + enable_heuristic_topk=True, + use_self_sampling_topk=use_self_sampling, + ) + + with ( + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.IS_CUTLASS_DSL_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.get_sm_version", + return_value=100, + ), + ): + indexer = create_indexer(sparse_config) + + assert indexer.top_k.decode_implementation == expected_decode + assert indexer.top_k.gvr_self_sampling == use_self_sampling + assert indexer.top_k.needs_gvr_prior == (not use_self_sampling) @skip_pre_hopper diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index daed3f24c9ba..28a95ff092d5 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -144,6 +144,7 @@ def test_gvr_uses_caller_prior_state(monkeypatch) -> None: 2, decode_implementation=TopKImplementation.CUTE_DSL_GVR, compress_ratio=4, + gvr_self_sampling=False, ) scores = torch.randn(1, 8) logical_lengths = torch.tensor([32], dtype=torch.int32) @@ -181,7 +182,11 @@ def test_gvr_uses_caller_prior_state(monkeypatch) -> None: def test_gvr_uses_caller_prepared_row_order(monkeypatch) -> None: gvr = Mock(side_effect=lambda *args, **kwargs: args[3].zero_()) monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) - top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + ) next_n = 2 lengths = torch.tensor([4, 1, 8, 2], dtype=torch.int32) row_order = torch.tensor([2, 0, 3, 1], dtype=torch.int32) @@ -231,7 +236,7 @@ def test_gvr_v2_decode_is_hint_free(monkeypatch) -> None: runner = _install_fake_selfsampling_runner(monkeypatch) top_k = TopK( 2, - decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, compress_ratio=4, ) @@ -252,7 +257,7 @@ def test_gvr_v2_hardware_gate_falls_back_without_prior(monkeypatch) -> None: monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) top_k = TopK( 2, - decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, compress_ratio=4, ) scores = torch.randn(1, 8, dtype=torch.bfloat16) @@ -289,7 +294,7 @@ def test_gvr_v2_decode_rejects_output_width_mismatch(monkeypatch) -> None: runner = _install_fake_selfsampling_runner(monkeypatch) top_k = TopK( 2, - decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, compress_ratio=4, ) with pytest.raises(AssertionError): @@ -309,7 +314,11 @@ def test_gvr_v2_decode_rejects_output_width_mismatch(monkeypatch) -> None: ], ) def test_update_gvr_prior_from_prefill_uses_last_request_rows(device) -> None: - top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + ) prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32, device=device) prior_indices = torch.zeros(3, 2, dtype=torch.int32, device=device) @@ -326,7 +335,7 @@ def test_update_gvr_prior_from_prefill_uses_last_request_rows(device) -> None: def test_gvr_v2_does_not_update_prior_from_prefill() -> None: - top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR_V2) + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) prior_indices = torch.zeros(1, 2, dtype=torch.int32) top_k.update_gvr_prior_from_prefill( @@ -339,6 +348,17 @@ def test_gvr_v2_does_not_update_prior_from_prefill() -> None: assert not top_k.needs_gvr_prior +def test_needs_gvr_prior_follows_two_level_dispatch() -> None: + assert TopK(2, decode_implementation=TopKImplementation.CUDA_GVR).needs_gvr_prior + assert not TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR).needs_gvr_prior + assert TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + ).needs_gvr_prior + assert not TopK(2).needs_gvr_prior + + def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: decode = Mock() monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) From d690470a44202fbc91e3d3fc5802f9da9831a02d Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:41:51 +0000 Subject: [PATCH 2/6] [None][refactor] Remove the CUDA GVR heuristic top-K decode Nothing selects the CUDA heuristic once the unified DSL GVR router is in: delete heuristicTopKDecode.{cu,h} / heuristic_topk.cuh, the canUseHeuristic dispatch and the GVR SchemeX bounds in indexerTopK.cu (radix keeps a cached SM-count helper), shrink the indexer_topk_decode thop schema and its register_fake (pre_idx / heuristic_scratch gone), drop the CUDA_GVR enum plus module branch, and retire the heuristic-only distribution / hostile-hint / tie-plateau test arms. The radix insertion / histogram / split-work tiers are untouched. C++ changes are not compiled yet: build + CI plus the 886x11 CUDA-v1 vs DSL-v1 paired A/B sign-off gate this draft. Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/IndexerTopK.h | 32 +- .../kernels/heuristicTopKDecode.cu | 285 --- .../kernels/heuristicTopKDecode.h | 61 - cpp/tensorrt_llm/kernels/heuristic_topk.cuh | 2006 ----------------- cpp/tensorrt_llm/kernels/indexerTopK.cu | 256 +-- cpp/tensorrt_llm/thop/IndexerTopKOp.cpp | 61 +- .../_torch/custom_ops/cpp_custom_ops.py | 2 - tensorrt_llm/_torch/modules/top_k.py | 82 +- tests/unittest/_torch/modules/test_top_k.py | 70 - .../_torch/thop/parallel/test_indexer_topk.py | 873 +------ 10 files changed, 61 insertions(+), 3667 deletions(-) delete mode 100644 cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu delete mode 100644 cpp/tensorrt_llm/kernels/heuristicTopKDecode.h delete mode 100644 cpp/tensorrt_llm/kernels/heuristic_topk.cuh diff --git a/cpp/tensorrt_llm/kernels/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 6e6e7b29b059..656675e641ad 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -34,20 +34,16 @@ namespace kernels // (a value <= 0 selects the internal default). int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitWorkThreshold = 0); -/// fp32 indexer TopK decode — L2-aware BS-threshold dispatcher with four -/// fallback tiers: -/// - GVR Heuristic (preIdx provided, kSeqSmall ≤ N < splitWork, BS < kBsLarge, K ∈ {512,1024,2048}) +/// fp32 indexer TopK decode — three dispatch tiers: /// - Insertion sort (N < kSortingAlgorithmThreshold) /// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork) /// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux) void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, - int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, - int const preIdxCount = 0, float* heuristicScratch = nullptr, int const compressRatio = 1, + int const stride1, int const next_n, int const topK = 2048, int const compressRatio = 1, cudaStream_t const stream = 0); -/// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except -/// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and +/// bf16 indexer TopK decode — same dispatch tiers as the fp32 entry, except /// the split-work tier is unsupported (the bf16/fp16 entry does not expose /// the float aux buffers required for split-work). Insertion + radix tiers /// share topKPerRowDecode with fp32 — histogram and sort run on float keys @@ -57,35 +53,17 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic /// that regime must use the fp32 entry. void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, - int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, - int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1, - cudaStream_t const stream = 0); + int const next_n, int const topK = 2048, int const compressRatio = 1, cudaStream_t const stream = 0); /// fp16 indexer TopK decode — see bf16 overload for dispatcher contract. void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, - int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - __half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0); + int const topK = 2048, int const compressRatio = 1, cudaStream_t const stream = 0); void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, int const numRows, int const numColumns, int const stride0, int const stride1, int const topK = 2048, cudaStream_t const stream = 0); -/// Returns true iff invokeIndexerTopKDecode would route to the GVR Heuristic -/// kernel for this (numRows, numColumns, topK) triple, assuming valid preIdx -/// is provided and stride1 == 1. Useful for callers that need to provision a -/// preIdx tensor or heuristicScratch buffer only when GVR will be selected. -/// -/// Mirrors the gating logic of the dispatcher: K ∈ {512, 1024, 2048}, -/// numColumns ∈ [kSeqSmall, splitWorkThreshold), numRows < kBsLarge, where -/// kBsLarge = min(kBsWave, kBsL2) and kBsL2 scales with bytesPerElem. -/// -/// @param numRows logits rows (batch · next_n) -/// @param numColumns logits columns (max sequence length) -/// @param topK requested output size -/// @param bytesPerElem element size of logits (4 for fp32, 2 for bf16/fp16) -bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytesPerElem = 4); - } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu deleted file mode 100644 index 839ae5483c05..000000000000 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu +++ /dev/null @@ -1,285 +0,0 @@ -/* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "tensorrt_llm/kernels/heuristicTopKDecode.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/envUtils.h" - -// Import gvrTopKJob (__device__ __noinline__, the GVR micro-kernel) and -// all helpers. gvrTopKJob is independently optimized by ptxas, matching standalone -// SASS quality regardless of the caller's prologue code. -#include "tensorrt_llm/kernels/heuristic_topk.cuh" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace -{ - -using heuristic_topk::BLOCK_SIZE; -using heuristic_topk::GvrDtypeTraits; -using heuristic_topk::GvrParams; -using heuristic_topk::gvrTopKJob; -using heuristic_topk::gvrTopKJobDtype; -using heuristic_topk::KernelSmemTplK; - -// Templated on TopK so the launcher can dispatch K=512/1024/2048 to the -// same kernel template. Smem layout is derived from GvrParams -// at compile time. -template -__global__ void __launch_bounds__(BLOCK_SIZE) - heuristicTopKMultiRowKernel(float const* __restrict__ logits, int const* __restrict__ seqLens, - int const* __restrict__ preIdx, float* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0, - int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio) -{ - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - - int const rowIdx = blockIdx.x; - int const seq_len = seqLens[rowIdx / next_n]; - // seqLens is in uncompressed token space; the logits/preIdx live in - // compressed-index space when compressRatio > 1 (DSv4 indexer). - int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; - int const N = actual_kv_len / compressRatio; - - float const* __restrict__ input = logits + static_cast(rowIdx) * stride0; - int const* __restrict__ rowPreIdx = preIdx + static_cast(rowIdx / next_n) * preIdxStride; - float* __restrict__ outputValues = scratchValues + static_cast(rowIdx) * topK; - int* __restrict__ outputIndices = outIndices + static_cast(rowIdx) * topK; - - extern __shared__ unsigned char smem_raw[]; - auto* smem = reinterpret_cast(smem_raw); - - if (N <= topK) - { - int const tid = threadIdx.x; - for (int i = tid; i < N; i += BLOCK_SIZE) - { - outputValues[i] = input[i]; - outputIndices[i] = i; - } - for (int i = N + tid; i < topK; i += BLOCK_SIZE) - { - outputValues[i] = -FLT_MAX; - outputIndices[i] = -1; - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - return; - } - - // Temporal-shift offset to map prev-step's top-K indices into this step's - // KV index space. - // compressRatio == 1 (DSv3.2): +1 — KV grew by exactly 1 token per - // decode step; prev indices were at seq_len-1 so a uniform +1 maps - // them to the equivalent positions under the indexer's "newest-first" - // layout. The (rowIdx % next_n) addend extends this to MTP windows. - // compressRatio == 4 (DSv4): 0 — in compressed-index space new - // compressed entries are appended at the end; prev indices in - // [0, c_prev-1] remain valid as-is. Per-row Δc varies (0 or 1) with - // prev kv_len mod 4 alignment, but a uniform offset of 0 stays - // within-bounds for all rows and preserves the temporal-correlation - // hint (vertical top-K consistency validated offline). - int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0; - gvrTopKJob(input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -// ============================================================================ -// Multi-dtype path (bf16 / fp16) -// ============================================================================ -// Mirrors heuristicTopKMultiRowKernel for bf16/fp16 inputs. The kernel body -// is structurally identical; only the input/output dtype, the smem-key -// dtype, and the GVR job (gvrTopKJobDtype) differ. - -// Templated on (InputT, TopK). Smem layout is derived from -// GvrParams. -template -__global__ void __launch_bounds__(BLOCK_SIZE) - heuristicTopKMultiRowKernelDtype(InputT const* __restrict__ logits, int const* __restrict__ seqLens, - int const* __restrict__ preIdx, InputT* __restrict__ scratchValues, int* __restrict__ outIndices, int stride0, - int next_n, int topK, int preIdxStride, int preIdxCount, int compressRatio) -{ - // dtype path uses fp32 keys[] in smem (down-conversion deferred to writeback). - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - - int const rowIdx = blockIdx.x; - int const seq_len = seqLens[rowIdx / next_n]; - int const actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; - int const N = actual_kv_len / compressRatio; - - InputT const* __restrict__ input = logits + static_cast(rowIdx) * stride0; - int const* __restrict__ rowPreIdx = preIdx + static_cast(rowIdx / next_n) * preIdxStride; - InputT* __restrict__ outputValues = scratchValues + static_cast(rowIdx) * topK; - int* __restrict__ outputIndices = outIndices + static_cast(rowIdx) * topK; - - extern __shared__ unsigned char smem_raw[]; - auto* smem = reinterpret_cast(smem_raw); - - if (N <= topK) - { - int const tid = threadIdx.x; - for (int i = tid; i < N; i += BLOCK_SIZE) - { - outputValues[i] = input[i]; - outputIndices[i] = i; - } - InputT const neg_max = GvrDtypeTraits::from_fp32(-FLT_MAX); - for (int i = N + tid; i < topK; i += BLOCK_SIZE) - { - outputValues[i] = neg_max; - outputIndices[i] = -1; - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - return; - } - - // See fp32 path: cr==1 → (rowIdx % next_n)+1; cr!=1 (DSv4) → 0. - int const preIdxOffset = (compressRatio == 1) ? ((rowIdx % next_n) + 1) : 0; - gvrTopKJobDtype( - input, N, rowPreIdx, preIdxCount, topK, outputValues, outputIndices, smem, preIdxOffset); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -// Explicit instantiations — 6 (dtype × K) combos. Launchers dispatch on -// runtime topK via switch, so all 6 must be available at link time. -// Trailing `int` is the compressRatio parameter (1 = V3.2, 4 = V4 indexer). -template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 512>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); -template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 1024>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); -template __global__ void heuristicTopKMultiRowKernelDtype<__nv_bfloat16, 2048>( - __nv_bfloat16 const*, int const*, int const*, __nv_bfloat16*, int*, int, int, int, int, int, int); -template __global__ void heuristicTopKMultiRowKernelDtype<__half, 512>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); -template __global__ void heuristicTopKMultiRowKernelDtype<__half, 1024>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); -template __global__ void heuristicTopKMultiRowKernelDtype<__half, 2048>( - __half const*, int const*, int const*, __half*, int*, int, int, int, int, int, int); -template __global__ void heuristicTopKMultiRowKernel<512>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); -template __global__ void heuristicTopKMultiRowKernel<1024>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); -template __global__ void heuristicTopKMultiRowKernel<2048>( - float const*, int const*, int const*, float*, int*, int, int, int, int, int, int); - -// Dispatch on topK at runtime — each TopK-instantiation gets its own smem -// size (driven by GvrParams::kC/kNumBins) and own kfn pointer -// (cudaFuncSetAttribute / cudaLaunchKernelEx target the right kernel). -// -// fp32 routes to heuristicTopKMultiRowKernel; bf16/fp16 route to -// heuristicTopKMultiRowKernelDtype. Vector-load alignment -// requirement is 4 elements for fp32 (float4) and 8 elements for bf16/fp16 -// (int4 of 16-bit). In TRT-LLM the logits stride is always a multiple of -// tokens_per_block (≥64), so the alignment check is never hit at runtime -// — it's an assert against caller misuse. -template -void launchHeuristicTopKDecodeImpl(InputT const* logits, int const* seqLens, int const* preIdx, int* outIndices, - InputT* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO( - topK == 512 || topK == 1024 || topK == 2048, "heuristicTopKDecode requires topK ∈ {512, 1024, 2048}"); - - constexpr int kAlign = std::is_same_v ? 4 : 8; - TLLM_CHECK_WITH_INFO(stride0 % kAlign == 0 || numRows <= 1, - "heuristicTopKDecode requires logits stride0 divisible by %d for multi-row launch", kAlign); - - auto launchOne = [&]() - { - // bf16/fp16 path also uses fp32 keys[] in smem (down-conversion deferred). - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - size_t const smemSize = sizeof(SmemT); - - auto kfn = []() - { - if constexpr (std::is_same_v) - return heuristicTopKMultiRowKernel; - else - return heuristicTopKMultiRowKernelDtype; - }(); - - if (smemSize > 48u * 1024u) - { - cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemSize)); - } - - cudaLaunchConfig_t config; - config.gridDim = numRows; - config.blockDim = BLOCK_SIZE; - config.dynamicSmemBytes = smemSize; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); - config.numAttrs = 1; - config.attrs = attrs; - - cudaLaunchKernelEx(&config, kfn, logits, seqLens, preIdx, scratchValues, outIndices, stride0, next_n, topK, - preIdxStride, preIdxCount, compressRatio); - }; - - switch (topK) - { - case 512: launchOne.template operator()<512>(); break; - case 1024: launchOne.template operator()<1024>(); break; - case 2048: launchOne.template operator()<2048>(); break; - default: TLLM_THROW("heuristicTopKDecode: topK validated above; unreachable"); - } -} - -} // anonymous namespace - -void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices, - float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream) -{ - launchHeuristicTopKDecodeImpl(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); -} - -void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices, - __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream) -{ - launchHeuristicTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, - topK, preIdxStride, preIdxCount, numRows, compressRatio, stream); -} - -void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices, - __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream) -{ - launchHeuristicTopKDecodeImpl<__half>(logits, seqLens, preIdx, outIndices, scratchValues, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h b/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h deleted file mode 100644 index 0d2330f76545..000000000000 --- a/cpp/tensorrt_llm/kernels/heuristicTopKDecode.h +++ /dev/null @@ -1,61 +0,0 @@ -/* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -inline constexpr int kHeuristicTopK = 2048; -inline constexpr int kHeuristicSize = 2048; - -/// Launch heuristic TopK decode kernel — fp32 input. -/// @param scratchValues Caller-owned buffer of size [numRows * topK] floats. -/// Required for CUDA Graph compatibility — must have a stable device address. -/// @param compressRatio KV compression ratio (1 = V3.2 indexer; 4 = V4 indexer -/// whose logits/preIdx live in compressed-token-index space). For -/// compressRatio != 1, preIdxOffset is forced to 0 (append-at-end in -/// compressed space → prev-step indices remain valid as-is); the -/// existing (rowIdx % next_n)+1 shift is used only when compressRatio==1. -void launchHeuristicTopKDecode(float const* logits, int const* seqLens, int const* preIdx, int* outIndices, - float* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream); - -/// Launch heuristic TopK decode kernel — bf16 input. -/// scratchValues is [numRows * topK] of bf16 (matches input dtype). -/// @param compressRatio See fp32 overload. -void launchHeuristicTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int const* preIdx, int* outIndices, - __nv_bfloat16* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream); - -/// Launch heuristic TopK decode kernel — fp16 input. -/// scratchValues is [numRows * topK] of fp16 (matches input dtype). -/// @param compressRatio See fp32 overload. -void launchHeuristicTopKDecode(__half const* logits, int const* seqLens, int const* preIdx, int* outIndices, - __half* scratchValues, int stride0, int next_n, int topK, int preIdxStride, int preIdxCount, int numRows, - int compressRatio, cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh b/cpp/tensorrt_llm/kernels/heuristic_topk.cuh deleted file mode 100644 index d75933bd5d38..000000000000 --- a/cpp/tensorrt_llm/kernels/heuristic_topk.cuh +++ /dev/null @@ -1,2006 +0,0 @@ -/* - * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -// ============================================================================ -// heuristic_topk.cuh — Heuristic-Guided Top-K (Sort-Free, Histogram-Based) -// -// Outer name: "heuristic" (algorithm family + public dispatcher / launchers). -// Inner name: "gvr" (Guess-Verify-Refine) — the single-CTA single-row -// micro-kernel implementing the algorithm of: -// "Guess-Verify-Refine: Data-Aware Top-K for Sparse-Attention Decoding -// on Blackwell via Temporal Correlation" -// -// Optimised for NVIDIA B200 (Blackwell, sm_100), single thread-block kernel. -// -// GVR phase mapping: -// P1 (preIdx stats) ┐ Guess: estimate the K-th-value -// P2 (secant threshold search) ┘ threshold from previous-step top-K -// indices, then refine the guess via -// count-only secant iterations. -// P3 (collect) — Verify: scatter the elements that -// pass the guessed threshold into -// shared memory and confirm the -// candidate count is in the safe band. -// P4 (histogram snap + partition) — Refine: 2048-bin histogram snap to -// the exact K-th value, then partition -// the candidates into the output set. -// ============================================================================ - -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace heuristic_topk -{ - -// ============================================================================ -// Multi-dtype Trait Layer -// ============================================================================ -// Encapsulates dtype-specific cvt intrinsics + vector load width so the -// kernel body can be templated cleanly. For fp32 the trait is identity. -// -// Arithmetic (threshold, accumulators, bin index) is always fp32; only -// the HBM input container, smem keys, and output values follow InputT. -// GVR is HBM-bandwidth-bound, so fp32 ALU has no measurable cost. - -template -struct GvrDtypeTraits; - -template <> -struct GvrDtypeTraits -{ - using SmemKey = float; - static constexpr int VEC_W = 4; // int4 = 4 × fp32 - static constexpr int SMEM_KEY_BYTES = 4; - - __device__ static __forceinline__ float to_fp32(float v) - { - return v; - } - - __device__ static __forceinline__ float from_fp32(float v) - { - return v; - } - - __device__ static __forceinline__ void unpack4(int4 raw, float* out) - { - out[0] = __int_as_float(raw.x); - out[1] = __int_as_float(raw.y); - out[2] = __int_as_float(raw.z); - out[3] = __int_as_float(raw.w); - } -}; - -template <> -struct GvrDtypeTraits<__nv_bfloat16> -{ - using SmemKey = __nv_bfloat16; - static constexpr int VEC_W = 8; // int4 = 8 × bf16 = 4 × bf162 - static constexpr int SMEM_KEY_BYTES = 2; - - __device__ static __forceinline__ float to_fp32(__nv_bfloat16 v) - { - return __bfloat162float(v); - } - - __device__ static __forceinline__ __nv_bfloat16 from_fp32(float v) - { - return __float2bfloat16_rn(v); - } - - __device__ static __forceinline__ void unpack8(int4 raw, float* out) - { - auto* p = reinterpret_cast<__nv_bfloat162*>(&raw); -#pragma unroll - for (int j = 0; j < 4; j++) - { - out[2 * j] = __low2float(p[j]); - out[2 * j + 1] = __high2float(p[j]); - } - } -}; - -template <> -struct GvrDtypeTraits<__half> -{ - using SmemKey = __half; - static constexpr int VEC_W = 8; // int4 = 8 × fp16 = 4 × half2 - static constexpr int SMEM_KEY_BYTES = 2; - - __device__ static __forceinline__ float to_fp32(__half v) - { - return __half2float(v); - } - - __device__ static __forceinline__ __half from_fp32(float v) - { - return __float2half_rn(v); - } - - __device__ static __forceinline__ void unpack8(int4 raw, float* out) - { - auto* p = reinterpret_cast<__half2*>(&raw); -#pragma unroll - for (int j = 0; j < 4; j++) - { - out[2 * j] = __low2float(p[j]); - out[2 * j + 1] = __high2float(p[j]); - } - } -}; - -// ============================================================================ -// Configuration Constants -// ============================================================================ - -constexpr int BLOCK_SIZE = 512; -constexpr int WARP_SIZE = 32; -constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; - -constexpr int TOP_K = 2048; -constexpr int HEURISTIC_SIZE = 2048; -constexpr int SAFETY_MARGIN = 2048; -constexpr int MAX_CANDIDATES = TOP_K + SAFETY_MARGIN * 2; // 6144 - -constexpr int MAX_REFINE_ITERS = 15; -// Phase-3 repair budget: bisecting on the uint32 key image collapses any -// bracket to adjacent floats in <= 32 steps; 40 adds slack. -constexpr int MAX_REPAIR_ITERS = 40; -constexpr int NUM_BINS = 2048; - -static_assert(TOP_K % BLOCK_SIZE == 0); -static_assert(MAX_CANDIDATES % BLOCK_SIZE == 0); - -// ============================================================================ -// Multi-K Trait Layer -// ============================================================================ -// Per-(InputT, TopK) compile-time trait encoding the secant-search target -// `kFTarget`, candidate-buffer cap `kC`, and Phase-4 histogram bin count -// `kNumBins`. -// -// kFTarget under V3.2-decode preIdx semantics (preIdx = top-K of prev row): -// Phase-1 pmean lands near the right tail of the prev-row top-K, so the -// Phase-2 secant initial bracket is biased high. A tighter K-proportional -// target converges faster than the M=2048 era's flat target. -// -// `kFTarget` is the secant solver's **soft steering target**, not the -// convergence condition. The Phase-2 loop converges whenever the -// candidate count falls within `[kK, kCC]` (see the `done = 1` check at -// the end of `gvrTopKJob`'s P1+P2 scope). A `kFTarget` below `kK` is -// intentional and useful for small K — with preIdx-seeded P1 landing -// the initial threshold near the right tail of the prev-row top-K, the -// secant's first interpolation often overshoots; biasing the target -// below kK pulls the next iteration's threshold down more aggressively -// and reaches the legal `[kK, kCC]` band in fewer secant steps. The -// concrete multipliers below were tuned empirically over V4 M=K cells -// (commit 8 in this PR): -// K=512: 0.75K = 384 -// K=1024: 2.5K = 2560 -// K=2048: kept at 1.5K (3072) for fp32 to preserve SASS byte-identity -// with the V2e production hot path. bf16/fp16 K=2048 use 2K -// (4096), where there is no prior production baseline to honor. -// -// kC = 5120 for K=512/1024 across all dtypes — drops smem footprint enough -// to leave headroom for 4-5 CTA/SM theoretical occupancy when register -// allocation permits. fp32 K=2048 keeps kC=6144 to preserve V2e SASS -// byte-identity (production fp32 K=2048 hot path is the correctness -// floor; smem layout change would alter SASS). -// -// kNumBins varies per (T, K) by atomic-contention vs Phase-4 setup-cost -// trade-off: -// - Below ~1024 bins, atomicAdd contention on the bin-counter array -// dominates Phase-4 cost. -// - Above 1024 bins, the Phase-4 histogram clear+scan setup cost -// dominates over the contention savings. -// The optimum varies because (a) candidate-buffer size kC scales the -// atomic-contention denominator, and (b) bf16/fp16 paths read smem keys -// through Trait::to_fp32, shifting the Phase-3 ↔ Phase-4 ratio. Hence -// per-(T, K) tabulation rather than a closed-form rule. -// -// Primary template intentionally left undefined: any unsupported (T, K) -// combination triggers a compile-time error rather than a runtime fall- -// through. -template -struct GvrParams; // primary undefined → compile-time error for bad combos - -template <> -struct GvrParams -{ - // kFTarget=kK aligns the secant's soft steering target with the band's - // lower edge; eliminates the upper-clamp saturation on tight-σ + high-A2 - // layers (L36/L42/L28). Cross-prompt simulator validation on swe-bench - // 32k/64k/100k showed 2.19× / 1.77× / 1.51× total P2-iter reduction with - // zero cap-hits, zero per-layer regression vs the prior kFTarget=384. - static constexpr int kFTarget = 512; - static constexpr int kC = 5120; - static constexpr int kNumBins = 1024; -}; - -template <> -struct GvrParams -{ - // kFTarget = kK (see GvrParams rationale). Q9k Pro 32k - // K=1024 native sweep (M=K=1024) finds kFT=1024 reduces sum_mean - // P2 iters from 35.33 (kFT=2560) → 30.21 (1.17× speedup) with zero - // per-layer regression and zero cap-hits. The prior kFT=2560 setting - // was tuned with M=512 K=1024 (sparse_attention_config default - // index_topk=512 inherited Flash's K), which does not represent - // production Pro behavior (production: M = K). - static constexpr int kFTarget = 1024; - static constexpr int kC = 5120; - static constexpr int kNumBins = 1024; -}; - -// fp32 K=2048 preserves V2e SASS byte-identity with the production hot path. -// Changing kFTarget or kC would alter ptxas output for the existing -// `gvrTopKJob<2048>` / `heuristicTopKMultiRowKernel<2048>` instantiations. -template <> -struct GvrParams -{ - static constexpr int kFTarget = 3072; - static constexpr int kC = 6144; - static constexpr int kNumBins = NUM_BINS; -}; - -template <> -struct GvrParams<__nv_bfloat16, 512> -{ - // kFTarget aligned to kK — see GvrParams rationale. - static constexpr int kFTarget = 512; - static constexpr int kC = 5120; - static constexpr int kNumBins = 512; -}; - -template <> -struct GvrParams<__nv_bfloat16, 1024> -{ - // kFTarget = kK — see GvrParams rationale. - static constexpr int kFTarget = 1024; - static constexpr int kC = 5120; - static constexpr int kNumBins = 512; -}; - -template <> -struct GvrParams<__nv_bfloat16, 2048> -{ - static constexpr int kFTarget = 4096; - static constexpr int kC = 5120; - static constexpr int kNumBins = NUM_BINS; -}; - -template <> -struct GvrParams<__half, 512> -{ - // kFTarget aligned to kK — see GvrParams rationale. - static constexpr int kFTarget = 512; - static constexpr int kC = 5120; - static constexpr int kNumBins = 512; -}; - -template <> -struct GvrParams<__half, 1024> -{ - // kFTarget = kK — see GvrParams rationale. - static constexpr int kFTarget = 1024; - static constexpr int kC = 5120; - static constexpr int kNumBins = 1024; -}; - -template <> -struct GvrParams<__half, 2048> -{ - static constexpr int kFTarget = 4096; - static constexpr int kC = 5120; - static constexpr int kNumBins = NUM_BINS; -}; - -// kC must remain divisible by BLOCK_SIZE (vector loads). -static_assert(GvrParams::kC % BLOCK_SIZE == 0); -static_assert(GvrParams::kC % BLOCK_SIZE == 0); -static_assert(GvrParams::kC % BLOCK_SIZE == 0); - -// ============================================================================ -// Shared Memory Layout -// ============================================================================ -// Templated on (SmemKey, candidate-cap, num-bins). Default -// (MAX_CANDIDATES=6144, NUM_BINS=2048) sizes: -// fp32 : ~59 KB -// bf16/fp16 : ~47 KB -// K=512/1024 instantiations cap candidates at kC=5120 (~51 KB fp32 / -// ~41 KB bf16/fp16) per GvrParams::kC. - -template -struct KernelSmemTplK -{ - alignas(16) SmemKey keys[CCap]; // CCap × sizeof(SmemKey) (4B fp32 / 2B bf16/fp16) - alignas(16) int vals[CCap]; // CCap × 4B - - int warp_counts[NUM_WARPS]; // 64 B - int histogram[NumBinsT]; // NumBinsT × 4B (default 2048 → 8 KB) - int per_thread_counts[BLOCK_SIZE]; // cached from the most recent blockCountGE call (Phase-3 reuse) - - float threshold; - int cand_count; - int done; - - float val_lo, val_hi; - int cnt_lo, cnt_hi; - - float pmax_saved; - int out_count; -}; - -// Convenience alias for the default-cap layout (kC=6144, kNumBins=2048), -// used by the K=2048 instantiations. -template -using KernelSmemTpl = KernelSmemTplK; - -using KernelSmem = KernelSmemTpl; - -// ============================================================================ -// Warp-Level Reduction Primitives -// ============================================================================ - -#if __CUDA_ARCH__ >= 800 - -__device__ __forceinline__ int warpReduceSum(int val) -{ - return __reduce_add_sync(0xffffffffu, val); -} - -__device__ __forceinline__ unsigned floatToOrderedUint(float f) -{ - unsigned u = __float_as_uint(f); - return (u & 0x80000000u) ? ~u : (u | 0x80000000u); -} - -__device__ __forceinline__ float orderedUintToFloat(unsigned u) -{ - return __uint_as_float((u & 0x80000000u) ? (u & ~0x80000000u) : ~u); -} - -__device__ __forceinline__ float warpReduceMin(float val) -{ - unsigned u = floatToOrderedUint(val); - u = __reduce_min_sync(0xffffffffu, u); - return orderedUintToFloat(u); -} - -__device__ __forceinline__ float warpReduceMax(float val) -{ - unsigned u = floatToOrderedUint(val); - u = __reduce_max_sync(0xffffffffu, u); - return orderedUintToFloat(u); -} - -#else - -__device__ __forceinline__ int warpReduceSum(int val) -{ -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - val += __shfl_down_sync(0xffffffffu, val, off); - return val; -} - -__device__ __forceinline__ float warpReduceMin(float val) -{ -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - val = fminf(val, __shfl_xor_sync(0xffffffffu, val, off)); - return val; -} - -__device__ __forceinline__ float warpReduceMax(float val) -{ -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - val = fmaxf(val, __shfl_xor_sync(0xffffffffu, val, off)); - return val; -} - -#endif - -// ============================================================================ -// Order-preserving float <-> uint32 map (arch-independent) -// ============================================================================ -// Same bijection as floatToOrderedUint above but defined for every -// __CUDA_ARCH__; the Phase-3 repair bisects on this key image so the -// bracket provably collapses (a float-average midpoint has no such bound). -__device__ __forceinline__ unsigned gvrOrderKey(float f) -{ - unsigned u = __float_as_uint(f); - return (u & 0x80000000u) ? ~u : (u | 0x80000000u); -} - -__device__ __forceinline__ float gvrOrderKeyToFloat(unsigned u) -{ - return __uint_as_float((u & 0x80000000u) ? (u & ~0x80000000u) : ~u); -} - -// ============================================================================ -// Device: Block count ≥ threshold in GLOBAL memory (1-sync pattern) -// ============================================================================ - -// Templated on SmemT so K=512/1024 paths can pass a kC=5120 layout. -// Default (KernelSmem) targets the K=2048 kC=6144 layout. -template -__device__ __forceinline__ void blockCountGE( - float const* __restrict__ input, int N, float threshold, SmemT* smem, int tid, int warp_id, int lane) -{ - int c = 0; - for (int i = tid * 4; i + 3 < N; i += BLOCK_SIZE * 4) - { - float4 v4 = __ldg(reinterpret_cast(input + i)); - c += (v4.x >= threshold) + (v4.y >= threshold) + (v4.z >= threshold) + (v4.w >= threshold); - } - for (int i = (N & ~3) + tid; i < N; i += BLOCK_SIZE) - c += (__ldg(&input[i]) >= threshold); - - // cache per-thread count for Phase 3 sub-pass 1 reuse - smem->per_thread_counts[tid] = c; - - c = warpReduceSum(c); - - if (lane == 0) - smem->warp_counts[warp_id] = c; - __syncthreads(); - - if (tid == 0) - { - int t = 0; - for (int w = 0; w < NUM_WARPS; w++) - t += smem->warp_counts[w]; - smem->cand_count = t; - } -} - -// ============================================================================ -// Fused snap iteration (2 syncs per call) -// ============================================================================ - -// Templated on (TopK, SmemT) so K=512/1024 paths reuse the same helper. -template -__device__ __forceinline__ void blockFusedSnapIter(SmemT* smem, int count, int tid, int warp_id, int lane) -{ - float const thr = smem->threshold; - - int lge = 0, lgt = 0; - float s_up = FLT_MAX, s_down = -FLT_MAX; - - for (int i = tid; i < count; i += BLOCK_SIZE) - { - float v = smem->keys[i]; - lge += (v >= thr); - lgt += (v > thr); - if (v > thr) - s_up = fminf(s_up, v); - if (v < thr) - s_down = fmaxf(s_down, v); - } - - int packed = (lge << 16) | lgt; - packed = warpReduceSum(packed); - s_up = warpReduceMin(s_up); - s_down = warpReduceMax(s_down); - - if (lane == 0) - { - smem->warp_counts[warp_id] = packed; - smem->histogram[warp_id] = __float_as_int(s_up); - smem->histogram[NUM_WARPS + warp_id] = __float_as_int(s_down); - } - __syncthreads(); - - if (tid == 0) - { - int tp = 0; - float total_up = FLT_MAX, total_down = -FLT_MAX; - for (int w = 0; w < NUM_WARPS; w++) - { - tp += smem->warp_counts[w]; - total_up = fminf(total_up, __int_as_float(smem->histogram[w])); - total_down = fmaxf(total_down, __int_as_float(smem->histogram[NUM_WARPS + w])); - } - smem->cnt_lo = tp >> 16; - smem->cnt_hi = tp & 0xFFFF; - - int cge = smem->cnt_lo; - int cgt = smem->cnt_hi; - - if (cgt >= TopK) - { - if (total_up < FLT_MAX) - smem->threshold = total_up; - } - else if (cge < TopK) - { - if (total_down > -FLT_MAX) - smem->threshold = total_down; - } - } - __syncthreads(); -} - -// ============================================================================ -// Dtype-templated helpers (bf16 / fp16) -// ============================================================================ -// Mirror of `blockCountGE` for bf16/fp16 inputs (8-wide vector load via -// `Trait::unpack8` + fp32 up-cast before threshold compare). The Phase-4 -// snap iter is NOT mirrored: `gvrTopKJobDtype` stores smem `keys[]` as -// fp32 even on bf16/fp16 paths (deferred-conversion optimization, see -// the `gvrTopKJobDtype` comment block below), so it reuses the fp32 -// `blockFusedSnapIter` helper directly. - -// Templated on SmemT so K=512/1024 dtype paths can pass a kC=5120 layout. -// Default targets the K=2048 kC=6144 layout. -template ::SmemKey>> -__device__ __forceinline__ void blockCountGEDtype( - InputT const* __restrict__ input, int N, float threshold, SmemT* smem, int tid, int warp_id, int lane) -{ - using Trait = GvrDtypeTraits; - static_assert(Trait::VEC_W == 8, "blockCountGEDtype is for bf16/fp16 (8-wide vector); use blockCountGE for fp32"); - - int c = 0; - for (int i = tid * 8; i + 7 < N; i += BLOCK_SIZE * 8) - { - int4 raw = __ldg(reinterpret_cast(input + i)); - float v[8]; - Trait::unpack8(raw, v); -#pragma unroll - for (int j = 0; j < 8; j++) - c += (v[j] >= threshold); - } - for (int i = (N & ~7) + tid; i < N; i += BLOCK_SIZE) - c += (Trait::to_fp32(__ldg(&input[i])) >= threshold); - - smem->per_thread_counts[tid] = c; - - c = warpReduceSum(c); - - if (lane == 0) - smem->warp_counts[warp_id] = c; - __syncthreads(); - - if (tid == 0) - { - int t = 0; - for (int w = 0; w < NUM_WARPS; w++) - t += smem->warp_counts[w]; - smem->cand_count = t; - } -} - -// ============================================================================ -// Device function: algorithm body (independently optimized by ptxas) -// __noinline__ ensures ptxas allocates registers and schedules instructions -// for this function independently from the caller, matching standalone SASS. -// ============================================================================ - -// Templated on TopK so K=512/1024/2048 fp32 paths share this body. The -// runtime `topK` parameter is kept for header-API compatibility with -// callers that pass it; the kernel asserts at entry that it matches -// the template instantiation. -template -__device__ __noinline__ void gvrTopKJob(float const* __restrict__ input, int const N, int const* __restrict__ preIdx, - int const M, int const topK, float* __restrict__ outputValues, int* __restrict__ outputIndices, - KernelSmemTplK::kC, GvrParams::kNumBins>* smem, - int const preIdxOffset = 0) -{ - using Params = GvrParams; - constexpr int kK = TopK; - constexpr int kCC = Params::kC; - constexpr int kBins = Params::kNumBins; - constexpr int kFTarget = Params::kFTarget; - - int const tid = threadIdx.x; - int const warp_id = tid / WARP_SIZE; - int const lane = tid & (WARP_SIZE - 1); - unsigned const full_mask = 0xffffffffu; - - { - // ================================================================ - // Phase 1 (GVR Guess, part 1) — Min/Max/Mean of pre-indexed values - // ================================================================ - - float local_min = FLT_MAX; - float local_max = -FLT_MAX; - float local_sum = 0.0f; - int local_cnt = 0; - for (int i = tid; i < M; i += BLOCK_SIZE) - { - int idx = __ldg(&preIdx[i]) + preIdxOffset; - if (idx >= 0 && idx < N) - { - float v = __ldg(&input[idx]); - local_min = fminf(local_min, v); - local_max = fmaxf(local_max, v); - local_sum += v; - local_cnt++; - } - } - - float wmin = warpReduceMin(local_min); - float wmax = warpReduceMax(local_max); - float wsum = local_sum; -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - wsum += __shfl_down_sync(0xffffffffu, wsum, off); - int wcnt = warpReduceSum(local_cnt); - - if (lane == 0) - { - smem->histogram[warp_id] = __float_as_int(wmin); - smem->histogram[NUM_WARPS + warp_id] = __float_as_int(wmax); - smem->histogram[NUM_WARPS * 2 + warp_id] = __float_as_int(wsum); - smem->histogram[NUM_WARPS * 3 + warp_id] = wcnt; - } - __syncthreads(); - - if (tid == 0) - { - float pmin = FLT_MAX, pmax = -FLT_MAX, psum = 0.0f; - int pcnt = 0; - for (int w = 0; w < NUM_WARPS; w++) - { - pmin = fminf(pmin, __int_as_float(smem->histogram[w])); - pmax = fmaxf(pmax, __int_as_float(smem->histogram[NUM_WARPS + w])); - psum += __int_as_float(smem->histogram[NUM_WARPS * 2 + w]); - pcnt += smem->histogram[NUM_WARPS * 3 + w]; - } - float pmean = (pcnt > 0) ? psum / (float) pcnt : (pmin + pmax) * 0.5f; - - smem->pmax_saved = pmax; - smem->threshold = pmean; - smem->val_lo = pmin; - smem->val_hi = pmax; - smem->cnt_lo = M + M / 4; - smem->cnt_hi = 1; - smem->done = 0; - } - __syncthreads(); - - // Degenerate hint (all gathered values identical or out of range): - // reset to a trusted bracket instead of emitting row[0:K]; done = 2 - // skips the secant (it cannot converge on a full-range bracket) and - // hands the row to the Phase-3 repair. The hint only affects speed. - if (smem->val_hi <= -FLT_MAX || smem->val_lo >= smem->val_hi) - { - if (tid == 0) - { - float const seed = (smem->val_hi <= -FLT_MAX) ? 0.0f : smem->pmax_saved; - smem->val_lo = -FLT_MAX; - smem->val_hi = FLT_MAX; - smem->cnt_lo = N; - smem->cnt_hi = 0; - smem->threshold = seed; - smem->done = 2; - } - __syncthreads(); - } - - // ================================================================ - // Phase 2 (GVR Guess, part 2) — Secant-interpolation threshold search - // ================================================================ - - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - - if (tid == 0) - { - int c = smem->cand_count; - if (c >= kK && c <= kCC) - smem->done = 1; - else if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->cnt_lo = c; - } - else - { - smem->val_hi = smem->threshold; - smem->cnt_hi = c; - } - } - __syncthreads(); - - for (int iter = 0; iter < MAX_REFINE_ITERS; iter++) - { - if (smem->done) - break; - if (tid == 0) - { - float vlo = smem->val_lo, vhi = smem->val_hi; - int clo = smem->cnt_lo, chi = smem->cnt_hi; - constexpr int target = kFTarget; - float range = vhi - vlo; - float nv; - if (clo > chi && range > 1e-10f) - { - float f = (float) (clo - target) / (float) (clo - chi); - f = fmaxf(0.05f, fminf(0.95f, f)); - if (iter == 0) - f = fminf(f, 0.50f); - nv = vlo + range * f; - } - else - nv = (vlo + vhi) * 0.5f; - if (nv <= vlo) - nv = vlo + range * 0.05f; - if (nv >= vhi) - nv = vhi - range * 0.05f; - if (nv == vlo || nv == vhi) - { - nv = (vlo + vhi) * 0.5f; - if (nv == vlo || nv == vhi) - { - smem->threshold = vlo; - smem->done = 2; - } - else - smem->threshold = nv; - } - else - smem->threshold = nv; - } - __syncthreads(); - if (smem->done) - break; - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0) - { - int c = smem->cand_count; - if (c >= kK && c <= kCC) - smem->done = 1; - else if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->cnt_lo = c; - } - else - { - smem->val_hi = smem->threshold; - smem->cnt_hi = c; - } - } - __syncthreads(); - } - - if (tid == 0 && !smem->done) - { - if (smem->cnt_lo <= kCC * 2) - smem->threshold = smem->val_lo; - else - smem->threshold = smem->val_hi; - smem->done = 2; - } - __syncthreads(); - } // end of P1+P2 scope - - // ================================================================ - // Phase 3 (GVR Verify) — Ballot-free candidate collect - // ================================================================ - - // done==1: Phase 2 verified cand_count in [kK, kCC]; skip the re-check. - // Otherwise the secant did not converge and `threshold` carries no - // guarantee: repair BOTH sides (the old loop only handled overflow, so - // an undershooting threshold shipped a -1-padded, silently wrong top-K). - if (smem->done != 1) - { - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - // Anchor the untested bracket end at a float extreme: Phase 1 seeds - // both ends from HINTED values with invented counts, so they can sit - // on the same side of the K-th value. count(-FLT_MAX) >= kK, - // count(FLT_MAX) = 0. - if (tid == 0) - { - int c = smem->cand_count; - if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->val_hi = FLT_MAX; - } - else if (c < kK) - { - smem->val_hi = smem->threshold; - smem->val_lo = -FLT_MAX; - } - } - __syncthreads(); - - // Invariant maintained below: count(val_lo) >= kK. - for (int retry = 0; retry < MAX_REPAIR_ITERS && (smem->cand_count > kCC || smem->cand_count < kK); retry++) - { - unsigned const klo = gvrOrderKey(smem->val_lo); - unsigned const khi = gvrOrderKey(smem->val_hi); - if (khi <= klo + 1u) - break; // bracket collapsed to adjacent representable values - if (tid == 0) - smem->threshold = gvrOrderKeyToFloat(klo + ((khi - klo) >> 1)); - __syncthreads(); - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0) - { - int c = smem->cand_count; - if (c > kCC) - smem->val_lo = smem->threshold; - else if (c < kK) - smem->val_hi = smem->threshold; - } - __syncthreads(); - } - - // Still short of kK: the bracket collapsed; val_lo admits >= kK by - // the anchor invariant (or the row has < kK finite entries and the - // -1 tail pad is the correct answer). - if (smem->cand_count < kK) - { - if (tid == 0) - smem->threshold = smem->val_lo; - __syncthreads(); - blockCountGE(input, N, smem->threshold, smem, tid, warp_id, lane); - } - // blockCountGE publishes cand_count from tid 0 only; the branch below - // must be uniform across the block. - __syncthreads(); - - // Collapsed bracket still over kCC = a tie plateau wider than the - // candidate buffer: emit everything strictly above val_lo (< kK by - // construction) plus arbitrary ties — a valid tie-aware top-K. The - // adjacency guard keeps the emit sound if the loop ever ran dry. - if (smem->cand_count > kCC && gvrOrderKey(smem->val_hi) <= gvrOrderKey(smem->val_lo) + 1u) - { - float const thr = smem->threshold; - if (tid == 0) - smem->out_count = 0; - __syncthreads(); - for (int i = tid; i < N; i += BLOCK_SIZE) - { - float const v = __ldg(&input[i]); - if (v > thr) - { - int const p = atomicAdd(&smem->out_count, 1); - if (p < kK) - { - outputValues[p] = v; - outputIndices[p] = i; - } - } - } - __syncthreads(); - int const n_gt = min(smem->out_count, kK); - if (tid == 0) - smem->out_count = n_gt; - __syncthreads(); - for (int i = tid; i < N && smem->out_count < kK; i += BLOCK_SIZE) - { - float const v = __ldg(&input[i]); - if (v == thr) - { - int const p = atomicAdd(&smem->out_count, 1); - if (p < kK) - { - outputValues[p] = v; - outputIndices[p] = i; - } - } - } - __syncthreads(); - for (int i = min(smem->out_count, kK) + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = -FLT_MAX; - outputIndices[i] = -1; - } - return; - } - } - - // Reuse per-thread counts cached by the last blockCountGE call (saves - // one full N-scan; blockCountGE's __syncthreads guarantees visibility). - int my_total_qual = smem->per_thread_counts[tid]; - - int thread_prefix = my_total_qual; -#pragma unroll - for (int off = 1; off < WARP_SIZE; off *= 2) - { - int other = __shfl_up_sync(full_mask, thread_prefix, off); - if (lane >= off) - thread_prefix += other; - } - int my_excl_offset = thread_prefix - my_total_qual; - int warp_total_qual = __shfl_sync(full_mask, thread_prefix, WARP_SIZE - 1); - - if (lane == 0) - smem->warp_counts[warp_id] = warp_total_qual; - __syncthreads(); - - if (tid == 0) - { - int total = 0; - for (int w = 0; w < NUM_WARPS; w++) - { - int cnt = smem->warp_counts[w]; - smem->warp_counts[w] = total; - total += cnt; - } - smem->cand_count = total; - } - __syncthreads(); - - int my_write_pos = smem->warp_counts[warp_id] + my_excl_offset; - - { - float const thr = smem->threshold; - for (int i = tid * 4; i + 3 < N; i += BLOCK_SIZE * 4) - { - float4 v4 = __ldg(reinterpret_cast(input + i)); -#pragma unroll - for (int j = 0; j < 4; j++) - { - float val = (&v4.x)[j]; - if (val >= thr && my_write_pos < kCC) - { - smem->keys[my_write_pos] = val; - smem->vals[my_write_pos] = i + j; - my_write_pos++; - } - } - } - for (int i = (N & ~3) + tid; i < N; i += BLOCK_SIZE) - { - float val = __ldg(&input[i]); - if (val >= thr && my_write_pos < kCC) - { - smem->keys[my_write_pos] = val; - smem->vals[my_write_pos] = i; - my_write_pos++; - } - } - } - __syncthreads(); - - // ================================================================ - // Phase 4 (GVR Refine) — Histogram-based selection + partition - // ================================================================ - - int const cand_count = min(smem->cand_count, kCC); - - if (cand_count == kK) - { - for (int i = tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = smem->keys[i]; - outputIndices[i] = smem->vals[i]; - } - return; - } - - if (cand_count > kK) - { - float cmin = FLT_MAX, cmax = -FLT_MAX; - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - float v = smem->keys[i]; - cmin = fminf(cmin, v); - cmax = fmaxf(cmax, v); - } - cmin = warpReduceMin(cmin); - cmax = warpReduceMax(cmax); - if (lane == 0) - { - smem->warp_counts[warp_id] = __float_as_int(cmin); - smem->histogram[warp_id] = __float_as_int(cmax); - } - __syncthreads(); - - float block_min = FLT_MAX, block_max = -FLT_MAX; - for (int w = 0; w < NUM_WARPS; w++) - { - block_min = fminf(block_min, __int_as_float(smem->warp_counts[w])); - block_max = fmaxf(block_max, __int_as_float(smem->histogram[w])); - } - if (block_max <= block_min) - block_max = block_min + 1e-6f; - - for (int i = tid; i < kBins; i += BLOCK_SIZE) - smem->histogram[i] = 0; - __syncthreads(); - - float range1 = block_max - block_min; - float inv1 = (range1 > 0.0f) ? ((float) (kBins - 1) + 0.99f) / range1 : 0.0f; - - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - int bin = (int) ((smem->keys[i] - block_min) * inv1); - bin = min(max(bin, 0), kBins - 1); - atomicAdd(&smem->histogram[bin], 1); - } - __syncthreads(); - - // Parallel K-th bin search (3-step). - // Step 1: each warp sums BINS_PER_WARP consecutive bins (high→low). - // Step 2: tid=0 locates the target warp in NUM_WARPS steps. - // Step 3: one thread in that warp scans its BINS_PER_WARP bins. - // Total serial depth: NUM_WARPS + BINS_PER_WARP steps vs full kBins. - { - constexpr int BINS_PER_WARP = kBins / NUM_WARPS; - static_assert(kBins % NUM_WARPS == 0, "kBins must be divisible by NUM_WARPS"); - // Step 1: each warp accumulates its slice of bins (high→low) - int warp_bin_sum = 0; - for (int j = 0; j < BINS_PER_WARP; j++) - warp_bin_sum += smem->histogram[kBins - 1 - warp_id * BINS_PER_WARP - j]; - if (lane == 0) - smem->warp_counts[warp_id] = warp_bin_sum; - } - __syncthreads(); // S-4b3a - - // Step 2: tid=0 finds which warp contains the K-th element - if (tid == 0) - { - int cum = 0, tw = NUM_WARPS - 1; - for (int w = 0; w < NUM_WARPS; w++) - { - cum += smem->warp_counts[w]; - if (cum >= kK) - { - tw = w; - break; - } - } - // Recompute prefix before target warp for step 3 - cum = 0; - for (int w = 0; w < tw; w++) - cum += smem->warp_counts[w]; - smem->cnt_lo = cum; // prefix count before target warp - smem->cnt_hi = tw; // target warp index - } - __syncthreads(); // S-4b3b - - // Step 3: one thread in target warp scans its BINS_PER_WARP bins - if (warp_id == smem->cnt_hi && lane == 0) - { - constexpr int BINS_PER_WARP = kBins / NUM_WARPS; - int base_cum = smem->cnt_lo; - float thr = block_min; - for (int j = 0; j < BINS_PER_WARP; j++) - { - int b = kBins - 1 - smem->cnt_hi * BINS_PER_WARP - j; - base_cum += smem->histogram[b]; - if (base_cum >= kK) - { - thr = block_min + (float) b * range1 / (float) kBins; - break; - } - } - smem->threshold = thr; - } - __syncthreads(); // S-4b3c - - // snap_limit must equal cand_count to guarantee convergence: each - // iteration either strictly decreases cgt (raise thr to next - // distinct value above) or strictly increases cge (lower thr to - // next distinct value below) by >= 1, so worst-case convergence - // takes cand_count - kK + 1 iters. The older bound `cand_count/4` - // silently accepted a non-converged threshold; Pass 1 then picked - // K elements in scan order from `cgt > kK` candidates, missing - // some true top-K members (~0.09 % intermittent at small-kNumBins - // mean-zero distributions). Common path still converges in 1-3 - // iters; the higher upper bound only affects the long-tail cells. - bool snap_converged = false; - int snap_limit = cand_count; - for (int si = 0; si < snap_limit; si++) - { - blockFusedSnapIter(smem, cand_count, tid, warp_id, lane); - int cge = smem->cnt_lo; - int cgt = smem->cnt_hi; - if (cgt < kK && cge >= kK) - { - snap_converged = true; - break; - } - } - (void) snap_converged; - - float sel_thr = smem->threshold; - if (tid == 0) - smem->out_count = 0; - __syncthreads(); - - // Two-pass selection: pass 1 emits strictly-greater-than-threshold - // candidates, pass 2 fills remaining slots with tie-values. An - // interleaved single-pass implementation would be unstable across - // rows with many ties at the K-th rank — equal-valued candidates - // could displace strictly-greater ones depending on their relative - // order in the candidate buffer. Splitting the passes makes the - // selection deterministic regardless of buffer ordering. - - // Pass 1: strictly greater than sel_thr - for (int base = warp_id * WARP_SIZE; base < cand_count; base += BLOCK_SIZE) - { - int i = base + lane; - float v = (i < cand_count) ? smem->keys[i] : -FLT_MAX; - - bool emit_gt = (i < cand_count) && (v > sel_thr); - unsigned mask_gt = __ballot_sync(full_mask, emit_gt); - if (mask_gt) - { - int cnt = __popc(mask_gt); - int moff = __popc(mask_gt & ((1u << lane) - 1u)); - int bp = 0; - if (lane == 0) - bp = atomicAdd(&smem->out_count, cnt); - bp = __shfl_sync(full_mask, bp, 0); - if (emit_gt && bp + moff < kK) - { - outputValues[bp + moff] = v; - outputIndices[bp + moff] = smem->vals[i]; - } - } - } - __syncthreads(); - - // Pass 2: equal to sel_thr (fills remaining slots) - for (int base = warp_id * WARP_SIZE; base < cand_count; base += BLOCK_SIZE) - { - int i = base + lane; - float v = (i < cand_count) ? smem->keys[i] : -FLT_MAX; - - bool emit_eq = (i < cand_count) && (v == sel_thr); - unsigned mask_eq = __ballot_sync(full_mask, emit_eq); - if (mask_eq) - { - int cnt = __popc(mask_eq); - int moff = __popc(mask_eq & ((1u << lane) - 1u)); - int bp = 0; - if (lane == 0) - bp = atomicAdd(&smem->out_count, cnt); - bp = __shfl_sync(full_mask, bp, 0); - if (emit_eq && bp + moff < kK) - { - outputValues[bp + moff] = v; - outputIndices[bp + moff] = smem->vals[i]; - } - } - } - __syncthreads(); - - int filled = min(smem->out_count, kK); - for (int i = filled + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = -FLT_MAX; - outputIndices[i] = -1; - } - return; - } - - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - outputValues[i] = smem->keys[i]; - outputIndices[i] = smem->vals[i]; - } - for (int i = cand_count + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = -FLT_MAX; - outputIndices[i] = -1; - } -} - -// ============================================================================ -// gvrTopKKernel — single-row global wrapper (1 CTA, 1 row). -// Calls gvrTopKJob (independently-optimized device function). -// For multi-row decode launches, see heuristicTopKMultiRowKernel in -// heuristicTopKDecode.cu — both share the same micro-kernel job. -// ============================================================================ - -// Templated on TopK so the launcher can dispatch K=512/1024/2048 to the -// same kernel template. -// -// __launch_bounds__ uses the single-arg form (no minBlocksPerSM hint) so -// nvcc applies the same register heuristic as `heuristicTopKMultiRowKernel` -// in heuristicTopKDecode.cu. Adding `, 1` would lower theoretical occupancy -// from 75% (REG=40) to 50% (REG=64) for the K=2048 fp32 path. -template -__global__ void __launch_bounds__(BLOCK_SIZE) - gvrTopKKernel(float const* __restrict__ input, int const N, int const* __restrict__ preIdx, int const M, - int const topK, float* __restrict__ outputValues, int* __restrict__ outputIndices) -{ - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - extern __shared__ unsigned char smem_raw[]; - auto* smem = reinterpret_cast(smem_raw); - - gvrTopKJob(input, N, preIdx, M, topK, outputValues, outputIndices, smem, /*preIdxOffset=*/0); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -// ============================================================================ -// gvrTopKJobDtype — bf16/fp16 device function -// ============================================================================ -// Mirror of gvrTopKJob with trait-driven dtype substitutions: -// - HBM input read : Trait::to_fp32(__ldg(&input[i])) -// - outputValues : InputT (Trait::from_fp32 at writeback) -// - vector load : 8-wide via Trait::unpack8 -// - blockCountGE : blockCountGEDtype -// `blockFusedSnapIter` is reused directly from the fp32 path — -// smem `keys[]` are stored as fp32 here (see deferred-conversion note -// below) so no dtype-specialized snap helper is needed. -// All arithmetic (threshold, accumulators, bin index) stays fp32. The fp32 -// dtype path uses gvrTopKJob (above), not this template; instantiated only -// for bf16 and fp16. -// -// smem `keys[]` are stored as fp32 even on the bf16/fp16 paths: deferring -// the down-conversion out of the Phase-3 collect loop saves more than the -// extra ~10 KB of smem costs. The fp32 keys move conversion to the output -// writeback (one cvt per surviving candidate) instead of every smem store. -template -__device__ __noinline__ void gvrTopKJobDtype(InputT const* __restrict__ input, int const N, - int const* __restrict__ preIdx, int const M, int const topK, InputT* __restrict__ outputValues, - int* __restrict__ outputIndices, - KernelSmemTplK::kC, GvrParams::kNumBins>* smem, - int const preIdxOffset = 0) -{ - using Trait = GvrDtypeTraits; - using SmemKey = float; // keys stay fp32; conversion deferred to output writeback - using Params = GvrParams; - constexpr int kK = TopK; - constexpr int kCC = Params::kC; - constexpr int kBins = Params::kNumBins; - constexpr int kFTarget = Params::kFTarget; - static_assert(Trait::VEC_W == 8, "gvrTopKJobDtype is for bf16/fp16 (8-wide); fp32 uses gvrTopKJob"); - - int const tid = threadIdx.x; - int const warp_id = tid / WARP_SIZE; - int const lane = tid & (WARP_SIZE - 1); - unsigned const full_mask = 0xffffffffu; - - { - // ================================================================ - // Phase 1 — Min/Max/Mean of pre-indexed values - // ================================================================ - - float local_min = FLT_MAX; - float local_max = -FLT_MAX; - float local_sum = 0.0f; - int local_cnt = 0; - for (int i = tid; i < M; i += BLOCK_SIZE) - { - int idx = __ldg(&preIdx[i]) + preIdxOffset; - if (idx >= 0 && idx < N) - { - float v = Trait::to_fp32(__ldg(&input[idx])); - local_min = fminf(local_min, v); - local_max = fmaxf(local_max, v); - local_sum += v; - local_cnt++; - } - } - - float wmin = warpReduceMin(local_min); - float wmax = warpReduceMax(local_max); - float wsum = local_sum; -#pragma unroll - for (int off = WARP_SIZE / 2; off > 0; off >>= 1) - wsum += __shfl_down_sync(0xffffffffu, wsum, off); - int wcnt = warpReduceSum(local_cnt); - - if (lane == 0) - { - smem->histogram[warp_id] = __float_as_int(wmin); - smem->histogram[NUM_WARPS + warp_id] = __float_as_int(wmax); - smem->histogram[NUM_WARPS * 2 + warp_id] = __float_as_int(wsum); - smem->histogram[NUM_WARPS * 3 + warp_id] = wcnt; - } - __syncthreads(); - - if (tid == 0) - { - float pmin = FLT_MAX, pmax = -FLT_MAX, psum = 0.0f; - int pcnt = 0; - for (int w = 0; w < NUM_WARPS; w++) - { - pmin = fminf(pmin, __int_as_float(smem->histogram[w])); - pmax = fmaxf(pmax, __int_as_float(smem->histogram[NUM_WARPS + w])); - psum += __int_as_float(smem->histogram[NUM_WARPS * 2 + w]); - pcnt += smem->histogram[NUM_WARPS * 3 + w]; - } - float pmean = (pcnt > 0) ? psum / (float) pcnt : (pmin + pmax) * 0.5f; - - smem->pmax_saved = pmax; - smem->threshold = pmean; - smem->val_lo = pmin; - smem->val_hi = pmax; - smem->cnt_lo = M + M / 4; - smem->cnt_hi = 1; - smem->done = 0; - } - __syncthreads(); - - // Degenerate hint (all gathered values identical or out of range): - // reset to a trusted bracket instead of emitting row[0:K]; done = 2 - // skips the secant (it cannot converge on a full-range bracket) and - // hands the row to the Phase-3 repair. The hint only affects speed. - if (smem->val_hi <= -FLT_MAX || smem->val_lo >= smem->val_hi) - { - if (tid == 0) - { - float const seed = (smem->val_hi <= -FLT_MAX) ? 0.0f : smem->pmax_saved; - smem->val_lo = -FLT_MAX; - smem->val_hi = FLT_MAX; - smem->cnt_lo = N; - smem->cnt_hi = 0; - smem->threshold = seed; - smem->done = 2; - } - __syncthreads(); - } - - // ================================================================ - // Phase 2 — Secant-interpolation threshold search - // ================================================================ - - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - - if (tid == 0) - { - int c = smem->cand_count; - if (c >= kK && c <= kCC) - smem->done = 1; - else if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->cnt_lo = c; - } - else - { - smem->val_hi = smem->threshold; - smem->cnt_hi = c; - } - } - __syncthreads(); - - for (int iter = 0; iter < MAX_REFINE_ITERS; iter++) - { - if (smem->done) - break; - if (tid == 0) - { - float vlo = smem->val_lo, vhi = smem->val_hi; - int clo = smem->cnt_lo, chi = smem->cnt_hi; - constexpr int target = kFTarget; - float range = vhi - vlo; - float nv; - if (clo > chi && range > 1e-10f) - { - float f = (float) (clo - target) / (float) (clo - chi); - f = fmaxf(0.05f, fminf(0.95f, f)); - if (iter == 0) - f = fminf(f, 0.50f); - nv = vlo + range * f; - } - else - nv = (vlo + vhi) * 0.5f; - if (nv <= vlo) - nv = vlo + range * 0.05f; - if (nv >= vhi) - nv = vhi - range * 0.05f; - if (nv == vlo || nv == vhi) - { - nv = (vlo + vhi) * 0.5f; - if (nv == vlo || nv == vhi) - { - smem->threshold = vlo; - smem->done = 2; - } - else - smem->threshold = nv; - } - else - smem->threshold = nv; - } - __syncthreads(); - if (smem->done) - break; - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0) - { - int c = smem->cand_count; - if (c >= kK && c <= kCC) - smem->done = 1; - else if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->cnt_lo = c; - } - else - { - smem->val_hi = smem->threshold; - smem->cnt_hi = c; - } - } - __syncthreads(); - } - - if (tid == 0 && !smem->done) - { - if (smem->cnt_lo <= kCC * 2) - smem->threshold = smem->val_lo; - else - smem->threshold = smem->val_hi; - smem->done = 2; - } - __syncthreads(); - } // end of P1+P2 scope - - // ================================================================ - // Phase 3 — Ballot-free candidate collect - // ================================================================ - - // Mirror of the fp32 Phase-3 repair in gvrTopKJob (see comments there). - if (smem->done != 1) - { - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - // See the fp32 path: anchor the untested bracket end at a float extreme - // so count(val_lo) >= kK > count(val_hi) holds by construction. - if (tid == 0) - { - int c = smem->cand_count; - if (c > kCC) - { - smem->val_lo = smem->threshold; - smem->val_hi = FLT_MAX; - } - else if (c < kK) - { - smem->val_hi = smem->threshold; - smem->val_lo = -FLT_MAX; - } - } - __syncthreads(); - - // Invariant maintained below: count(val_lo) >= kK. - for (int retry = 0; retry < MAX_REPAIR_ITERS && (smem->cand_count > kCC || smem->cand_count < kK); retry++) - { - unsigned const klo = gvrOrderKey(smem->val_lo); - unsigned const khi = gvrOrderKey(smem->val_hi); - if (khi <= klo + 1u) - break; // bracket collapsed to adjacent representable values - if (tid == 0) - smem->threshold = gvrOrderKeyToFloat(klo + ((khi - klo) >> 1)); - __syncthreads(); - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - if (tid == 0) - { - int c = smem->cand_count; - if (c > kCC) - smem->val_lo = smem->threshold; - else if (c < kK) - smem->val_hi = smem->threshold; - } - __syncthreads(); - } - - if (smem->cand_count < kK) - { - if (tid == 0) - smem->threshold = smem->val_lo; - __syncthreads(); - blockCountGEDtype(input, N, smem->threshold, smem, tid, warp_id, lane); - } - // blockCountGEDtype publishes cand_count from tid 0 only; the branch - // below must be uniform across the block. - __syncthreads(); - - // Collapsed bracket with > kCC elements at the threshold: emit the - // strictly-greater set plus arbitrary ties directly (see fp32 path). - // The direct emit below is only valid once the bracket has collapsed: - // it assumes count(> thr) < kK, which is exactly "val_hi is the next - // representable value above val_lo and count(val_hi) < kK". If the - // loop ran out of iterations without collapsing (it cannot, given - // MAX_REPAIR_ITERS >= 32, but the guard keeps that an invariant rather - // than an assumption) fall through to the ordinary collect. - if (smem->cand_count > kCC && gvrOrderKey(smem->val_hi) <= gvrOrderKey(smem->val_lo) + 1u) - { - float const thr = smem->threshold; - if (tid == 0) - smem->out_count = 0; - __syncthreads(); - for (int i = tid; i < N; i += BLOCK_SIZE) - { - float const v = Trait::to_fp32(__ldg(&input[i])); - if (v > thr) - { - int const p = atomicAdd(&smem->out_count, 1); - if (p < kK) - { - outputValues[p] = Trait::from_fp32(v); - outputIndices[p] = i; - } - } - } - __syncthreads(); - int const n_gt = min(smem->out_count, kK); - if (tid == 0) - smem->out_count = n_gt; - __syncthreads(); - for (int i = tid; i < N && smem->out_count < kK; i += BLOCK_SIZE) - { - float const v = Trait::to_fp32(__ldg(&input[i])); - if (v == thr) - { - int const p = atomicAdd(&smem->out_count, 1); - if (p < kK) - { - outputValues[p] = Trait::from_fp32(v); - outputIndices[p] = i; - } - } - } - __syncthreads(); - InputT const neg_max = Trait::from_fp32(-FLT_MAX); - for (int i = min(smem->out_count, kK) + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = neg_max; - outputIndices[i] = -1; - } - return; - } - } - - int my_total_qual = smem->per_thread_counts[tid]; - - int thread_prefix = my_total_qual; -#pragma unroll - for (int off = 1; off < WARP_SIZE; off *= 2) - { - int other = __shfl_up_sync(full_mask, thread_prefix, off); - if (lane >= off) - thread_prefix += other; - } - int my_excl_offset = thread_prefix - my_total_qual; - int warp_total_qual = __shfl_sync(full_mask, thread_prefix, WARP_SIZE - 1); - - if (lane == 0) - smem->warp_counts[warp_id] = warp_total_qual; - __syncthreads(); - - if (tid == 0) - { - int total = 0; - for (int w = 0; w < NUM_WARPS; w++) - { - int cnt = smem->warp_counts[w]; - smem->warp_counts[w] = total; - total += cnt; - } - smem->cand_count = total; - } - __syncthreads(); - - int my_write_pos = smem->warp_counts[warp_id] + my_excl_offset; - - { - float const thr = smem->threshold; - // 8-wide vector load (int4 = 8 × bf16/fp16) - for (int i = tid * 8; i + 7 < N; i += BLOCK_SIZE * 8) - { - int4 raw = __ldg(reinterpret_cast(input + i)); - float v[8]; - Trait::unpack8(raw, v); -#pragma unroll - for (int j = 0; j < 8; j++) - { - float val = v[j]; - if (val >= thr && my_write_pos < kCC) - { - smem->keys[my_write_pos] = val; // P0: defer convert to output - smem->vals[my_write_pos] = i + j; - my_write_pos++; - } - } - } - // Tail loop (N % 8) - for (int i = (N & ~7) + tid; i < N; i += BLOCK_SIZE) - { - float val = Trait::to_fp32(__ldg(&input[i])); - if (val >= thr && my_write_pos < kCC) - { - smem->keys[my_write_pos] = val; // P0: defer convert to output - smem->vals[my_write_pos] = i; - my_write_pos++; - } - } - } - __syncthreads(); - - // ================================================================ - // Phase 4 — Histogram-based selection + partition - // ================================================================ - - int const cand_count = min(smem->cand_count, kCC); - - if (cand_count == kK) - { - for (int i = tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = Trait::from_fp32(smem->keys[i]); // P0: convert at output - outputIndices[i] = smem->vals[i]; - } - return; - } - - if (cand_count > kK) - { - float cmin = FLT_MAX, cmax = -FLT_MAX; - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - float v = smem->keys[i]; // P0: keys already fp32 - cmin = fminf(cmin, v); - cmax = fmaxf(cmax, v); - } - cmin = warpReduceMin(cmin); - cmax = warpReduceMax(cmax); - if (lane == 0) - { - smem->warp_counts[warp_id] = __float_as_int(cmin); - smem->histogram[warp_id] = __float_as_int(cmax); - } - __syncthreads(); - - float block_min = FLT_MAX, block_max = -FLT_MAX; - for (int w = 0; w < NUM_WARPS; w++) - { - block_min = fminf(block_min, __int_as_float(smem->warp_counts[w])); - block_max = fmaxf(block_max, __int_as_float(smem->histogram[w])); - } - if (block_max <= block_min) - block_max = block_min + 1e-6f; - - for (int i = tid; i < kBins; i += BLOCK_SIZE) - smem->histogram[i] = 0; - __syncthreads(); - - float range1 = block_max - block_min; - float inv1 = (range1 > 0.0f) ? ((float) (kBins - 1) + 0.99f) / range1 : 0.0f; - - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - int bin = (int) ((smem->keys[i] - block_min) * inv1); // P0: keys fp32 - bin = min(max(bin, 0), kBins - 1); - atomicAdd(&smem->histogram[bin], 1); - } - __syncthreads(); - - // Parallel K-th bin search (2-step) - { - constexpr int BINS_PER_WARP = kBins / NUM_WARPS; - static_assert(kBins % NUM_WARPS == 0, "kBins must be divisible by NUM_WARPS"); - int warp_bin_sum = 0; - for (int j = 0; j < BINS_PER_WARP; j++) - warp_bin_sum += smem->histogram[kBins - 1 - warp_id * BINS_PER_WARP - j]; - if (lane == 0) - smem->warp_counts[warp_id] = warp_bin_sum; - } - __syncthreads(); - - if (tid == 0) - { - int cum = 0, tw = NUM_WARPS - 1; - for (int w = 0; w < NUM_WARPS; w++) - { - cum += smem->warp_counts[w]; - if (cum >= kK) - { - tw = w; - break; - } - } - cum = 0; - for (int w = 0; w < tw; w++) - cum += smem->warp_counts[w]; - smem->cnt_lo = cum; - smem->cnt_hi = tw; - } - __syncthreads(); - - if (warp_id == smem->cnt_hi && lane == 0) - { - constexpr int BINS_PER_WARP = kBins / NUM_WARPS; - int base_cum = smem->cnt_lo; - float thr = block_min; - for (int j = 0; j < BINS_PER_WARP; j++) - { - int b = kBins - 1 - smem->cnt_hi * BINS_PER_WARP - j; - base_cum += smem->histogram[b]; - if (base_cum >= kK) - { - thr = block_min + (float) b * range1 / (float) kBins; - break; - } - } - smem->threshold = thr; - } - __syncthreads(); - - // snap_limit must equal cand_count to guarantee convergence; see - // detailed rationale in the fp32 `gvrTopKJob` path. - bool snap_converged = false; - int snap_limit = cand_count; - for (int si = 0; si < snap_limit; si++) - { - blockFusedSnapIter(smem, cand_count, tid, warp_id, lane); // P0: keys fp32 → use fp32 snap helper - int cge = smem->cnt_lo; - int cgt = smem->cnt_hi; - if (cgt < kK && cge >= kK) - { - snap_converged = true; - break; - } - } - (void) snap_converged; - - float sel_thr = smem->threshold; - if (tid == 0) - smem->out_count = 0; - __syncthreads(); - - // Pass 1: strictly greater than sel_thr - for (int base = warp_id * WARP_SIZE; base < cand_count; base += BLOCK_SIZE) - { - int i = base + lane; - float v = (i < cand_count) ? smem->keys[i] : -FLT_MAX; // P0: keys fp32 - - bool emit_gt = (i < cand_count) && (v > sel_thr); - unsigned mask_gt = __ballot_sync(full_mask, emit_gt); - if (mask_gt) - { - int cnt = __popc(mask_gt); - int moff = __popc(mask_gt & ((1u << lane) - 1u)); - int bp = 0; - if (lane == 0) - bp = atomicAdd(&smem->out_count, cnt); - bp = __shfl_sync(full_mask, bp, 0); - if (emit_gt && bp + moff < kK) - { - outputValues[bp + moff] = Trait::from_fp32(v); - outputIndices[bp + moff] = smem->vals[i]; - } - } - } - __syncthreads(); - - // Pass 2: equal to sel_thr (fills remaining slots) - for (int base = warp_id * WARP_SIZE; base < cand_count; base += BLOCK_SIZE) - { - int i = base + lane; - float v = (i < cand_count) ? smem->keys[i] : -FLT_MAX; // P0: keys fp32 - - bool emit_eq = (i < cand_count) && (v == sel_thr); - unsigned mask_eq = __ballot_sync(full_mask, emit_eq); - if (mask_eq) - { - int cnt = __popc(mask_eq); - int moff = __popc(mask_eq & ((1u << lane) - 1u)); - int bp = 0; - if (lane == 0) - bp = atomicAdd(&smem->out_count, cnt); - bp = __shfl_sync(full_mask, bp, 0); - if (emit_eq && bp + moff < kK) - { - outputValues[bp + moff] = Trait::from_fp32(v); - outputIndices[bp + moff] = smem->vals[i]; - } - } - } - __syncthreads(); - - int filled = min(smem->out_count, kK); - InputT const neg_max = Trait::from_fp32(-FLT_MAX); - for (int i = filled + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = neg_max; - outputIndices[i] = -1; - } - return; - } - - // cand_count < kK fallback - for (int i = tid; i < cand_count; i += BLOCK_SIZE) - { - outputValues[i] = Trait::from_fp32(smem->keys[i]); // P0: convert at output - outputIndices[i] = smem->vals[i]; - } - InputT const neg_max = Trait::from_fp32(-FLT_MAX); - for (int i = cand_count + tid; i < kK; i += BLOCK_SIZE) - { - outputValues[i] = neg_max; - outputIndices[i] = -1; - } -} - -// ============================================================================ -// gvrTopKKernelDtype — bf16/fp16 single-row global wrapper -// ============================================================================ -// Templated on (InputT, TopK). __launch_bounds__ uses the single-arg form -// so nvcc applies the same register heuristic as the multi-row dtype kernel -// in heuristicTopKDecode.cu. See `gvrTopKKernel` note above. - -template -__global__ void __launch_bounds__(BLOCK_SIZE) - gvrTopKKernelDtype(InputT const* __restrict__ input, int const N, int const* __restrict__ preIdx, int const M, - int const topK, InputT* __restrict__ outputValues, int* __restrict__ outputIndices) -{ - using SmemKey = typename GvrDtypeTraits::SmemKey; - using SmemT - = KernelSmemTplK::kC, GvrParams::kNumBins>; // dtype keys fp32 - extern __shared__ unsigned char smem_raw[]; - auto* smem = reinterpret_cast(smem_raw); - - gvrTopKJobDtype(input, N, preIdx, M, topK, outputValues, outputIndices, smem, - /*preIdxOffset=*/0); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -// ============================================================================ -// Explicit kernel instantiations — 9 (T × K) combos. -// ============================================================================ -// Mirrors the same pattern used in heuristicTopKDecode.cu for the multi-row -// kernels. Forces nvcc to emit each `gvrTopKKernel` / `gvrTopKKernelDtype -// ` host-side wrapper stub *before* `launchHeuristicTopK` takes their -// address via `cudaLaunchKernelEx`. Without these declarations, certain nvcc -// stubgen versions (CI containers on sm_89 / sm_120f) emit the implicit stub -// inside the kernel body and then conflict with their own subsequent -// "explicit specialization of __wrapper__device_stub_*" pass — surfacing -// as a "specialization after instantiation" build error against -// cudafe1.stub.c. Header is included only by heuristicTopKDecode.cu (one TU) -// so no ODR concern. -template __global__ void gvrTopKKernel<512>(float const*, int, int const*, int, int, float*, int*); -template __global__ void gvrTopKKernel<1024>(float const*, int, int const*, int, int, float*, int*); -template __global__ void gvrTopKKernel<2048>(float const*, int, int const*, int, int, float*, int*); -template __global__ void gvrTopKKernelDtype<__nv_bfloat16, 512>( - __nv_bfloat16 const*, int, int const*, int, int, __nv_bfloat16*, int*); -template __global__ void gvrTopKKernelDtype<__nv_bfloat16, 1024>( - __nv_bfloat16 const*, int, int const*, int, int, __nv_bfloat16*, int*); -template __global__ void gvrTopKKernelDtype<__nv_bfloat16, 2048>( - __nv_bfloat16 const*, int, int const*, int, int, __nv_bfloat16*, int*); -template __global__ void gvrTopKKernelDtype<__half, 512>(__half const*, int, int const*, int, int, __half*, int*); -template __global__ void gvrTopKKernelDtype<__half, 1024>(__half const*, int, int const*, int, int, __half*, int*); -template __global__ void gvrTopKKernelDtype<__half, 2048>(__half const*, int, int const*, int, int, __half*, int*); - -// ============================================================================ -// Launch Wrapper -// ============================================================================ - -namespace detail -{ -// Per-(T, TopK) launcher implementation. Hoisted out of `launchHeuristicTopK` -// (was a C++20 templated lambda `[&]()`) because that pattern -// confuses nvcc's cudafe1 stub generator: taking the address of -// `gvrTopKKernel` / `gvrTopKKernelDtype` from inside a -// templated capturing lambda triggers an "explicit specialization of -// `__wrapper__device_stub_gvrTopKKernel` after instantiation" error -// against the auto-generated host wrapper stub. A regular function template -// avoids the quirk and stays in C++17 (no templated-lambda extension warning). -// -// Kernel body, GvrParams traits, kfn selection, opt-in smem, PDL attr, and -// cudaLaunchKernelEx call are byte-identical to the previous lambda body — -// SASS is unchanged for all 9 (T, K) instantiations. -template -cudaError_t launchHeuristicTopKImpl(T const* input, int N, int const* preIdx, int M, int topK, T* outputValues, - int* outputIndices, cudaStream_t stream, bool enablePDL) -{ - // dtype path uses fp32 smem keys (deferred convert). Launcher - // allocates smem with float keys regardless of input dtype. - using SmemT = KernelSmemTplK::kC, GvrParams::kNumBins>; - size_t const smemSize = sizeof(SmemT); - - // Resolve target kernel function pointer at compile time. - auto kfn = []() - { - if constexpr (std::is_same_v) - return gvrTopKKernel; - else - return gvrTopKKernelDtype; - }(); - - if (smemSize > 48u * 1024u) - { - int device; - cudaGetDevice(&device); - int maxSmem; - cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); - if (smemSize > static_cast(maxSmem)) - return cudaErrorInvalidConfiguration; - cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemSize)); - } - - cudaLaunchConfig_t config{}; - config.gridDim = dim3(1); - config.blockDim = dim3(BLOCK_SIZE); - config.dynamicSmemBytes = smemSize; - config.stream = stream; - - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = enablePDL ? 1 : 0; - config.attrs = attrs; - config.numAttrs = 1; - - cudaLaunchKernelEx(&config, kfn, input, N, preIdx, M, topK, outputValues, outputIndices); - return cudaGetLastError(); -} -} // namespace detail - -template -cudaError_t launchHeuristicTopK(T const* input, int N, IdxT const* preIdx, int M, int topK, T* outputValues, - IdxT* outputIndices, cudaStream_t stream = 0) -{ - static_assert(sizeof(IdxT) == sizeof(int), "launchHeuristicTopK only supports 32-bit indices"); - static_assert(std::is_same_v || std::is_same_v || std::is_same_v, - "launchHeuristicTopK supports only fp32 / bf16 / fp16"); - - // GvrParams specializations cover K ∈ {512, 1024, 2048}; reject others. - if (topK != 512 && topK != 1024 && topK != 2048) - return cudaErrorInvalidValue; - - // Dispatch on (T, topK) → 9 distinct kernel-pointer paths. Each - // instantiation captures its own (kFTarget, kC, kNumBins) tuple via - // GvrParams so all values are compile-time constants inside - // the kernel body. Opt-in smem + cudaLaunchKernelEx + PDL handling is - // shared across all 9 paths via `detail::launchHeuristicTopKImpl`. - - // Honor the standard TRTLLM_ENABLE_PDL env var (default on; set "0" to - // disable). - bool enablePDL = true; - if (char const* env = std::getenv("TRTLLM_ENABLE_PDL")) - { - if (env[0] == '0' && env[1] == '\0') - enablePDL = false; - } - - switch (topK) - { - case 512: - return detail::launchHeuristicTopKImpl( - input, N, preIdx, M, topK, outputValues, outputIndices, stream, enablePDL); - case 1024: - return detail::launchHeuristicTopKImpl( - input, N, preIdx, M, topK, outputValues, outputIndices, stream, enablePDL); - case 2048: - return detail::launchHeuristicTopKImpl( - input, N, preIdx, M, topK, outputValues, outputIndices, stream, enablePDL); - default: return cudaErrorInvalidValue; - } -} - -// Explicit instantiations — fp32 + bf16/fp16 -template cudaError_t launchHeuristicTopK( - float const*, int, int const*, int, int, float*, int*, cudaStream_t); -template cudaError_t launchHeuristicTopK<__nv_bfloat16, int>( - __nv_bfloat16 const*, int, int const*, int, int, __nv_bfloat16*, int*, cudaStream_t); -template cudaError_t launchHeuristicTopK<__half, int>( - __half const*, int, int const*, int, int, __half*, int*, cudaStream_t); - -} // namespace heuristic_topk diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index d63f203a048c..b1bce6f2bf67 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -19,7 +19,6 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaTypeUtils.cuh" #include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/heuristicTopKDecode.h" #include "tensorrt_llm/kernels/noAuxTcKernels.h" #include #include @@ -718,105 +717,19 @@ constexpr int kMaxBlocksPerRowDecode = 10; // one full histogram pass worth of columns. constexpr int kDecodeMinColsPerSubBlock = kNumBins; -// Scheme X bound calculator — shared between fp32 and bf16/fp16 dispatchers. -// Caches hardware attrs (SM count, L2 capacity) and the small-N threshold -// once per process via std::call_once. Per-call cost is just two reads -// from cached static variables plus a small arithmetic block, no syscalls. -struct SchemeXBounds -{ - int smCount; - int l2Bytes; - int kBsWave; - int kBsL2; - int kBsLarge; - int kSeqSmall; -}; - -// Uniform small-N lower bound for the Heuristic GVR path across all K. -// Aligns the GVR routing boundary with the Radix multi-CTA split-work -// threshold (maxByCols = N / kDecodeMinColsPerSubBlock(=2048) ≥ 2 at -// N ≥ 4096), so the dispatcher's algorithmic-handoff point is consistent: -// below 4096 the Radix path resolves to single-CTA insertion-sort and GVR -// is not attempted; at or above 4096 GVR may be considered. -// DSv4 swe-bench synth sweeps on B200/B300 (V3.2-Q19c protocol, May 2026): -// N=4K cells across K ∈ {512, 1024, 2048} all win — GVR R/H bf16 = 3.07× -// (K=512) / 2.57× (K=1024) / 1.34× (K=2048). -// N=2K cells across the same 9 (K × dtype) combos all show GVR R/H < 1 -// (0.55× – 0.84×), justifying 4K as the floor. -inline int kSeqSmallDefaultForK(int /*topK*/) -{ - return 4096; -} - -inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem, int topK) +// Cached device SM count for the wave-aware blocks-per-row dispatch below. +inline int getDeviceSmCount() { static std::once_flag sOnce; static int sSm = 0; - static int sL2 = 0; - // ----------------------------------------------------------------------- - // Diagnostic / tuning escape-hatch env overrides. Both are OFF by default - // and the K-aware / hardware-derived defaults below are expected to be - // optimal for production. Use only for microbenchmarks, regression - // bisection, or workload-specific tuning where the defaults are clearly - // suboptimal. - // - // TRTLLM_HEURISTIC_NMIN (valid range [1024, 200000]) - // Overrides `kSeqSmall` (Heuristic small-N threshold) for ALL K. - // Lower risk: only shifts a perf threshold; the kernel still - // produces an exact top-K either way. Setting it too low routes - // more N → Heuristic and may be slower than the fallback for - // small N; correctness is preserved. - // - // TRTLLM_HEURISTIC_BSMAX (valid range [1, 65536]) - // Overrides `kBsLarge` (BS upper bound for Heuristic) past the - // hardware-derived min(kBsWave, kBsL2). Higher risk: bypasses - // L2/occupancy safety bounds, so heuristic may run in working-set - // ranges where it has not been tuned (L2 thrash, suboptimal grid - // configs). Primary use is indexer microbenchmarks that need a - // BS-scaling comparison against the Radix path on identical inputs. - // ----------------------------------------------------------------------- - // sNMinEnv > 0 iff TRTLLM_HEURISTIC_NMIN is set to a valid value. When set, - // it overrides the per-K default for ALL K. - static int sNMinEnv = 0; - static int sBsMax = 0; std::call_once(sOnce, []() { int dev = 0; cudaGetDevice(&dev); cudaDeviceGetAttribute(&sSm, cudaDevAttrMultiProcessorCount, dev); - cudaDeviceGetAttribute(&sL2, cudaDevAttrL2CacheSize, dev); - char const* env = std::getenv("TRTLLM_HEURISTIC_NMIN"); - if (env != nullptr) - { - int const v = std::atoi(env); - sNMinEnv = (v >= 1024 && v <= 200000) ? v : 0; - } - char const* env_bsmax = std::getenv("TRTLLM_HEURISTIC_BSMAX"); - if (env_bsmax != nullptr) - { - int const v = std::atoi(env_bsmax); - sBsMax = (v >= 1 && v <= 65536) ? v : 0; - } }); - - SchemeXBounds b; - b.smCount = sSm; - b.l2Bytes = sL2; - b.kBsWave = (sSm > 0) ? (sSm * 3 - sSm / 8) : 426; - b.kBsL2 = (sL2 > 0 && numColumns > 0) - ? static_cast(static_cast(sL2) * 9 / 10 / (static_cast(numColumns) * bytesPerElem)) - : b.kBsWave; - b.kBsLarge = std::min(b.kBsWave, b.kBsL2 > 0 ? b.kBsL2 : b.kBsWave); - if (sBsMax > 0) - { - // BSMAX env override bypasses the hardware-derived L2/occupancy bound - // (see the BSMAX section in the call_once block above for risk notes). - b.kBsLarge = sBsMax; - } - // NMIN env override (if set) wins over the per-K default for ALL K. - b.kSeqSmall = (sNMinEnv > 0) ? sNMinEnv : kSeqSmallDefaultForK(topK); - return b; + return sSm; } } // namespace @@ -838,12 +751,8 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW // Query the actual SM count from the driver so the dispatch tracks the // hardware rather than a baked-in target (H100=132, B200=148, …). - // topK=0: blocks-per-row computation is K-agnostic; kSeqSmall is uniform - // 4096 across K, so the topK arg is unused for the kSeqSmall lookup as - // well, and only smCount/kBsWave/kBsL2 are consumed here. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4, /*topK=*/0); - TLLM_CHECK_WITH_INFO(bounds.smCount > 0, "indexerTopK: failed to query device SM count"); - int const smCount = bounds.smCount; + int const smCount = getDeviceSmCount(); + TLLM_CHECK_WITH_INFO(smCount > 0, "indexerTopK: failed to query device SM count"); int const maxByCols = std::max(1, numColumns / kDecodeMinColsPerSubBlock); int const maxBp = std::min(maxByCols, kMaxBlocksPerRowDecode); @@ -889,110 +798,9 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, - int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, - int const preIdxCount, float* heuristicScratch, int const compressRatio, cudaStream_t const stream) + int const stride1, int const next_n, int const topK, int const compressRatio, cudaStream_t const stream) { constexpr int kNumThreadsPerBlock = 512; - int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; - - // ======================================================================== - // Small-N dispatch axis. - // - // GVR Heuristic Top-K has a *fixed* per-launch overhead from Phase-1 - // (preIdx stats reduction over M=2048) and Phase-4 (2048-bin histogram - // snap), totaling ~11 µs regardless of N. For small N (≤16K), this - // fixed cost dominates and the kernel loses to the existing - // insertion-sort/radix path. Empirically (random data, B200 BS=1): - // N=8192 : Heuristic 16.5 µs vs Radix 11.2 µs (radix 1.47× faster) - // N=16384 : Heuristic 21.9 µs vs Radix 22.0 µs (parity) - // N=32768 : Heuristic 26.1 µs vs Radix 32.9 µs (heuristic 1.26× faster) - // N=131072 : Heuristic 43.4 µs vs Radix 76.1 µs (heuristic 1.75× faster) - // - // Route N < kSeqSmall to the existing Radix/Insertion path (which itself - // splits at kSortingAlgorithmThreshold=12288). kSeqSmall is set at the - // empirical crossover point. - // - // ======================================================================== - // Architecture-derived BS-threshold dispatch — jointly bounded by - // occupancy AND L2 cache capacity. - // - // Two physical constraints bound when the per-row heuristic kernel - // remains faster than a radix streaming kernel: - // - // (A) Occupancy bound — 3·SM − SM/8 (wave geometry + setup margin) - // Each CTA uses ~58 KB SMEM (fixed, independent of N), so B200's - // 228 KB dynamic SMEM allows max 3 CTA/SM. Above 3·SM rows per - // launch, tail-wave imbalance causes stragglers. The -SM/8 margin - // (~1/8 wave) covers CTA setup + L2 ingestion overhead. - // On B200(148 SM): 3×148 − 18 = 426. - // - // (B) L2 cache bound — 0.9·L2 / (4·N) per-CTA logits fit - // Each CTA streams its row (N×4B) through L2 per Phase-2 iter. - // With num_concurrent_CTAs × N × 4B > L2, eviction dominates. - // On B200(126 MB L2) with N=70K: 0.9·126MB/(4·70690) ≈ 440, - // which is ~ equal to (A)=426 — the two constraints cross over - // near the SWE-Bench data point. - // For N > 73K the L2 bound tightens below (A) and must take - // over; e.g. N=128K → kBsL2=238, N=196K → kBsL2=155. - // - // Dispatch threshold = min(kBsWave, kBsL2), still data-agnostic (only - // queries hardware attrs). At N≈70K both bounds produce ~426, so the - // L2 axis is a no-op there; for larger N it auto-tightens the threshold. - // - // Small-N lower bound `kSeqSmall` is uniform 4096 across all K (see - // kSeqSmallDefaultForK). 4K is the dispatcher's algorithmic-handoff - // point: below 4096 the Radix path resolves to single-CTA insertion-sort - // (maxByCols = N/2048 = 1 → bp=1; useRadixSort = N≥12288 = false), and - // GVR is empirically slower than insertion-sort below 4K across all - // K ∈ {512, 1024, 2048} × dtype ∈ {fp32, bf16, fp16} (R/H ∈ [0.55, 0.84] - // at N=2K; DSv4 V3.2-Q19c synth sweeps May 2026). Configurable via - // TRTLLM_HEURISTIC_NMIN env (>=1024), which overrides the default. - // ======================================================================== - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4, topK); - int const kBsWave = bounds.kBsWave; - int const kBsL2 = bounds.kBsL2; - int const kBsLarge = bounds.kBsLarge; - int const kSeqSmall = bounds.kSeqSmall; - - bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - // compressRatio == 1: DSv3.2 indexer (no compressor). - // compressRatio == 4: DSv4 indexer (overlap compressor); logits/preIdx in - // compressed-token-index space. Kernel handles N = actual_kv_len/cr and - // forces preIdxOffset=0 internally for cr != 1. - bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4); - bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK - && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold - && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; - - // Optional env-gated dispatch trace (set TRTLLM_SCHEMEX_DEBUG=1 to enable) - { - static std::once_flag sDebugOnceFlag; - static bool sDebug = false; - std::call_once(sDebugOnceFlag, - []() - { - char const* env = std::getenv("TRTLLM_SCHEMEX_DEBUG"); - sDebug = (env != nullptr && env[0] == '1'); - }); - if (sDebug) - { - fprintf(stderr, - "[Scheme X] numRows=%d numColumns=%d kBsWave=%d kBsL2=%d kBsLarge=%d kSeqSmall=%d smCount=%d " - "L2=%dMB -> %s path%s\n", - numRows, numColumns, kBsWave, kBsL2, kBsLarge, kSeqSmall, bounds.smCount, - bounds.l2Bytes / (1024 * 1024), canUseHeuristic ? "Heuristic" : "Radix", - (numColumns < kSeqSmall) ? " (small-N route)" : ""); - } - } - - if (canUseHeuristic) - { - launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); - sync_check_cuda_error(stream); - return; - } - int const blocksPerRow = computeIndexerTopKDecodeBlocksPerRow(numRows, numColumns, splitWorkThreshold); cudaLaunchAttribute attrs[1]; @@ -1052,13 +860,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic // ============================================================================ // bf16 / fp16 dispatcher overloads // ============================================================================ -// Reuses the BS-threshold + small-N dispatch axes (kBsLarge, kSeqSmall) from -// the fp32 dispatcher, except kBsL2 uses sizeof(InputT) bytes/element instead -// of 4 — L2 footprint is half, so bf16/fp16 path remains valid for larger BS -// than fp32 at the same N. -// -// Fallback chain when GVR-Heuristic preconditions are not met (preIdx -// missing, BS too large, or numColumns < kSeqSmall): +// Dispatch chain: // numColumns < kSortingAlgorithmThreshold (12288) → insertion sort // kSortingAlgorithmThreshold ≤ numColumns < splitWorkThreshold → radix sort // numColumns ≥ splitWorkThreshold (200K default) → unsupported @@ -1078,8 +880,7 @@ namespace template void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, int const compressRatio, - cudaStream_t const stream) + int const compressRatio, cudaStream_t const stream) { static_assert(std::is_same_v || std::is_same_v, "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); @@ -1087,25 +888,7 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* constexpr int kNumThreadsPerBlock = 512; int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; - // bf16/fp16: bytes_per_element = sizeof(InputT) = 2 → kBsL2 doubles vs fp32. - // K-aware kSeqSmall — see fp32 dispatcher for rationale. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT)), topK); - int const kBsLarge = bounds.kBsLarge; - int const kSeqSmall = bounds.kSeqSmall; - - bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - // See fp32 path: cr==1 (V3.2) and cr==4 (V4 indexer) are both supported. - bool const compressRatioOk = (compressRatio == 1 || compressRatio == 4); - bool const canUseHeuristic = compressRatioOk && preIdx != nullptr && stride1 == 1 && isSupportedTopK - && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold - && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; - - if (canUseHeuristic) - { - launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, - preIdxStride, preIdxCount, numRows, compressRatio, stream); - } - else if (numColumns < kSortingAlgorithmThreshold) + if (numColumns < kSortingAlgorithmThreshold) { // Insertion sort path — InputT propagated; histogram/sort run on float keys. auto* kernel_instance = &topKPerRowDecode(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + stride0, stride1, next_n, topK, compressRatio, stream); } void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, int const compressRatio, - cudaStream_t const stream) + int const compressRatio, cudaStream_t const stream) { invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + stride1, next_n, topK, compressRatio, stream); } void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, @@ -1199,17 +980,6 @@ void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int con sync_check_cuda_error(stream); } -bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytesPerElem) -{ - bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - if (!isSupportedTopK) - { - return false; - } - auto const bounds = getSchemeXBounds(numColumns, bytesPerElem, topK); - return numColumns >= bounds.kSeqSmall && numColumns < kDefaultSplitWorkThreshold && numRows < bounds.kBsLarge; -} - } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index 7e9f3bd7070d..2352f82def77 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -36,9 +36,8 @@ namespace torch_ext { void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, th::Tensor const& indices, - int64_t next_n, int64_t index_topk, std::optional const& pre_idx, - std::optional const& heuristic_scratch, int64_t compress_ratio, - std::optional const& radix_aux_indices, std::optional const& radix_aux_logits) + int64_t next_n, int64_t index_topk, int64_t compress_ratio, std::optional const& radix_aux_indices, + std::optional const& radix_aux_logits) { TORCH_CHECK(compress_ratio > 0, "compress_ratio must be greater than 0"); @@ -70,53 +69,18 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t TORCH_CHECK(logits_stride_0 >= 0, "logits_stride_0 must be greater than or equal to 0"); TORCH_CHECK(logits_stride_1 >= 0, "logits_stride_1 must be greater than or equal to 0"); - int32_t const* preIdxPtr = nullptr; - int32_t preIdxStride = 0; - int32_t preIdxCount = 0; - if (pre_idx.has_value()) - { - auto const& preIdxTensor = pre_idx.value(); - TORCH_CHECK(preIdxTensor.is_cuda(), "pre_idx must be a CUDA tensor"); - TORCH_CHECK(preIdxTensor.device() == logits.device(), "pre_idx must be on the same device as logits"); - TORCH_CHECK(preIdxTensor.is_contiguous(), "pre_idx must be contiguous"); - TORCH_CHECK(preIdxTensor.dim() == 2, "pre_idx must be a 2D Tensor"); - TORCH_CHECK(preIdxTensor.size(0) * next_n == numRows64, - "pre_idx first dimension must equal logits.size(0)/next_n (one hint row per batch element)"); - preIdxPtr = preIdxTensor.data_ptr(); - preIdxStride = static_cast(preIdxTensor.stride(0)); - preIdxCount = static_cast(preIdxTensor.size(1)); - } - - // Caller-owned scratch buffer for heuristic TopK output values. - // Must be pre-allocated with stable address for CUDA Graph compatibility. - // scratch dtype must match input dtype. auto const logits_dtype = logits.scalar_type(); TORCH_CHECK(logits_dtype == at::ScalarType::Float || logits_dtype == at::ScalarType::BFloat16 || logits_dtype == at::ScalarType::Half, "indexer_topk_decode: logits dtype must be float32, bfloat16, or float16; got ", logits_dtype); - void* heuristicScratchPtr = nullptr; - if (heuristic_scratch.has_value()) - { - auto const& scratchTensor = heuristic_scratch.value(); - TORCH_CHECK(scratchTensor.is_cuda(), "heuristic_scratch must be a CUDA tensor"); - TORCH_CHECK( - scratchTensor.device() == logits.device(), "heuristic_scratch must be on the same device as logits"); - TORCH_CHECK(scratchTensor.is_contiguous(), "heuristic_scratch must be contiguous"); - TORCH_CHECK(scratchTensor.numel() >= static_cast(num_rows) * index_topk, - "heuristic_scratch must have at least numRows * index_topk elements"); - TORCH_CHECK(scratchTensor.scalar_type() == logits_dtype, - "heuristic_scratch dtype must match logits dtype (got scratch=", scratchTensor.scalar_type(), - ", logits=", logits_dtype, ")"); - heuristicScratchPtr = scratchTensor.data_ptr(); - } int32_t splitWorkThreshold = 200 * 1000; auto stream = at::cuda::getCurrentCUDAStream(logits.get_device()); if (logits_dtype == at::ScalarType::Float) { - // fp32 path — full Scheme X v1.2 dispatcher (GVR / Insertion / Radix / - // Radix-split-work). Caller-owned radix_aux_{indices,logits} are the + // fp32 path — Insertion / Radix / Radix-split-work dispatcher. + // Caller-owned radix_aux_{indices,logits} are the // split-work scratch buffers and are dereferenced only when // blocksPerRow > 1; for blocksPerRow == 1 the dispatcher passes // nullptr through to the kernel and never touches them (see @@ -126,8 +90,7 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t float* aux_logits_ptr = nullptr; 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). The + // Caller-owned scratch with stable address (CUDA Graph safe). 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(); @@ -160,23 +123,22 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t } tk::invokeIndexerTopKDecode(logits.data_ptr(), seq_lens.data_ptr(), indices.data_ptr(), aux_logits_ptr, aux_indices_ptr, splitWorkThreshold, num_rows, num_columns, logits_stride_0, - logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, - preIdxCount, static_cast(heuristicScratchPtr), static_cast(compress_ratio), stream); + logits_stride_1, static_cast(next_n), static_cast(index_topk), + static_cast(compress_ratio), stream); } else if (logits_dtype == at::ScalarType::BFloat16) { tk::invokeIndexerTopKDecode(reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()), seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, - logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, - preIdxStride, preIdxCount, static_cast<__nv_bfloat16*>(heuristicScratchPtr), + logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), static_cast(compress_ratio), stream); } else // Half { tk::invokeIndexerTopKDecode(reinterpret_cast<__half const*>(logits.data_ptr()), seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, - static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, - static_cast<__half*>(heuristicScratchPtr), static_cast(compress_ratio), stream); + static_cast(next_n), static_cast(index_topk), static_cast(compress_ratio), + stream); } } @@ -226,8 +188,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( "indexer_topk_decode(Tensor logits, Tensor seq_lens, Tensor indices, int next_n, int index_topk=2048, " - "Tensor? pre_idx=None, Tensor? heuristic_scratch=None, int compress_ratio=1, " - "Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> ()"); + "int compress_ratio=1, Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index f5086beecbf5..c35412398a2c 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -339,8 +339,6 @@ def _(logits, indices, next_n, index_topk, - pre_idx=None, - heuristic_scratch=None, compress_ratio=1, radix_aux_indices=None, radix_aux_logits=None): diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 9a140d7fe35f..2d7565fe8906 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -20,12 +20,10 @@ class TopKImplementation(str, Enum): TORCH = "torch" CUDA_RADIX = "cuda_radix" CUTE_DSL_RADIX = "cute_dsl_radix" - CUDA_GVR = "cuda_gvr" CUTE_DSL_GVR = "cute_dsl_gvr" _GVR_IMPLEMENTATIONS = { - TopKImplementation.CUDA_GVR, TopKImplementation.CUTE_DSL_GVR, } _MAX_RADIX_BLOCKS_PER_ROW = 10 @@ -71,8 +69,6 @@ def __init__( @property def needs_gvr_prior(self) -> bool: """Return whether decode consumes previous-step Top-K indices.""" - if self.decode_implementation == TopKImplementation.CUDA_GVR: - return True return ( self.decode_implementation == TopKImplementation.CUTE_DSL_GVR and not self.gvr_self_sampling @@ -105,8 +101,8 @@ def forward( next_n: Number of decode rows per request. max_seq_len: Maximum decode score width used for GVR kernel tuning. gvr_ext_kwargs: GVR-only keyword arguments. ``gvr_prior_indices`` - is required by the temporal GVR paths (``CUDA_GVR``, or - ``CUTE_DSL_GVR`` with ``gvr_self_sampling=False``). It is + is required by the temporal GVR path (``CUTE_DSL_GVR`` with + ``gvr_self_sampling=False``). It is caller-owned int32 previous selection with shape ``[num_requests, top_k]`` on ``scores.device``. The self-sampling engine does not consume this state. @@ -217,8 +213,6 @@ def _forward_decode_radix( output_indices, next_n, self.top_k, - pre_idx=None, - heuristic_scratch=None, compress_ratio=self.compress_ratio, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, @@ -339,8 +333,6 @@ def _forward_decode_gvr( output_indices, next_n, self.top_k, - pre_idx=None, - heuristic_scratch=None, compress_ratio=self.compress_ratio, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, @@ -348,52 +340,32 @@ def _forward_decode_gvr( return output_indices assert gvr_prior_indices is not None - if self.decode_implementation == TopKImplementation.CUDA_GVR: - workspace = self._get_workspace( - scores, - (scores.shape[0], self.top_k), - scores.dtype, - "top_k_cuda_gvr_workspace", - ) - radix_indices, radix_values = self._get_radix_workspace(scores) - torch.ops.trtllm.indexer_topk_decode( - scores, - sequence_lengths, - output_indices, - next_n, - self.top_k, - pre_idx=gvr_prior_indices, - heuristic_scratch=workspace, - compress_ratio=self.compress_ratio, - radix_aux_indices=radix_indices, - radix_aux_logits=radix_values, + assert max_seq_len is not None + # V1 temporal (DSL). Emission-assisted candidates (opt-in) are only + # armed on this hint-first path; the self-sampling V2 path above never + # arms them. + emission_kwargs: dict = {} + if self._gvr_emission_armed: + state = self._gvr_emission_state + num_rows = scores.shape[0] + emission_kwargs = state.topk_ext_kwargs( + self._gvr_emission_route, + num_rows, + state.block_max[:num_rows] if state.block_max is not None else None, ) - elif self.decode_implementation == TopKImplementation.CUTE_DSL_GVR: - assert max_seq_len is not None - emission_kwargs: dict = {} - if self._gvr_emission_armed: - state = self._gvr_emission_state - num_rows = scores.shape[0] - emission_kwargs = state.topk_ext_kwargs( - self._gvr_emission_route, - num_rows, - state.block_max[:num_rows] if state.block_max is not None else None, - ) - self._gvr_emission_armed = False - torch.ops.trtllm.cute_dsl_gvr_topk_decode( - scores, - gvr_prior_indices, - sequence_lengths, - output_indices, - self.top_k, - next_n=next_n, - compress_ratio=self.compress_ratio, - max_seq_len=max_seq_len, - order_row=gvr_row_order, - **emission_kwargs, - ) - else: - raise AssertionError(f"Unexpected GVR implementation: {self.decode_implementation}") + self._gvr_emission_armed = False + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + scores, + gvr_prior_indices, + sequence_lengths, + output_indices, + self.top_k, + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=max_seq_len, + order_row=gvr_row_order, + **emission_kwargs, + ) return output_indices def prepare_gvr_emission( diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 28a95ff092d5..366903b0be09 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -3,7 +3,6 @@ """Tests for the reusable sparse index-selection Top-K module.""" import sys -from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import Mock, call @@ -129,8 +128,6 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: output, 2, 2, - pre_idx=None, - heuristic_scratch=None, compress_ratio=4, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, @@ -280,8 +277,6 @@ def test_gvr_v2_hardware_gate_falls_back_without_prior(monkeypatch) -> None: output, 1, 2, - pre_idx=None, - heuristic_scratch=None, compress_ratio=4, radix_aux_indices=None, radix_aux_logits=None, @@ -349,7 +344,6 @@ def test_gvr_v2_does_not_update_prior_from_prefill() -> None: def test_needs_gvr_prior_follows_two_level_dispatch() -> None: - assert TopK(2, decode_implementation=TopKImplementation.CUDA_GVR).needs_gvr_prior assert not TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR).needs_gvr_prior assert TopK( 2, @@ -390,76 +384,12 @@ def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: output, 1, 1, - pre_idx=None, - heuristic_scratch=None, compress_ratio=1, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, ) -def test_cuda_gvr_reserves_workspace_during_capture(monkeypatch) -> None: - decode = Mock(side_effect=lambda *args, **kwargs: args[2].copy_(torch.tensor([[3, 1]]))) - monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) - device_context = Mock(side_effect=lambda _: nullcontext()) - monkeypatch.setattr(torch.cuda, "device", device_context) - - top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) - scores = Mock( - shape=(1, 8), - dtype=torch.float32, - is_cuda=True, - device=torch.device("cuda", 3), - ) - lengths = torch.tensor([8], dtype=torch.int32) - output = torch.empty(1, 2, dtype=torch.int32) - radix_indices = torch.empty(1, 10, 2, dtype=torch.int32) - radix_values = torch.empty(1, 10, 2) - workspace = torch.empty(1, 2) - prior_indices = torch.zeros(1, 2, dtype=torch.int32) - buffers = Mock() - buffers.get_buffer.side_effect = [workspace, radix_indices, radix_values] - monkeypatch.setattr(TopK, "_memory_buffers", buffers) - - top_k( - scores, - output, - is_prefill=False, - sequence_lengths=lengths, - scan_lengths=lengths, - gvr_ext_kwargs={"gvr_prior_indices": prior_indices}, - ) - - assert buffers.get_buffer.call_args_list == [ - call( - (scores.shape[0], 2), - dtype=scores.dtype, - buffer_name="top_k_cuda_gvr_workspace_cuda:3", - reserve_buffer=True, - ), - call( - (scores.shape[0], 10, 2), - dtype=torch.int32, - buffer_name="top_k_radix_indices_workspace_cuda:3", - reserve_buffer=True, - ), - call( - (scores.shape[0], 10, 2), - dtype=torch.float32, - buffer_name="top_k_radix_values_workspace_cuda:3", - reserve_buffer=True, - ), - ] - assert device_context.call_args_list == [call(scores.device)] * 3 - runtime_call = decode.call_args_list[-1] - assert runtime_call.kwargs["pre_idx"] is prior_indices - assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == workspace.data_ptr() - assert runtime_call.kwargs["radix_aux_indices"] is radix_indices - assert runtime_call.kwargs["radix_aux_logits"] is radix_values - assert prior_indices.tolist() == [[0, 0]] - - def test_unsupported_prefill_implementation_raises() -> None: top_k = TopK(1, prefill_implementation=TopKImplementation.CUTE_DSL_RADIX) diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index 33ebdb2c82ca..e065db98ac59 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -7,38 +7,12 @@ # http://www.apache.org/licenses/LICENSE-2.0 """ -Distribution-parameterized correctness tests for the heuristic indexer_topk_decode. - -Logits are sampled from four distribution families that characterise -negative-shifted decode-phase logit spaces (means −0.5 to −4.5): - - beta — bounded, bell-shaped; near-zero / moderate / deep negative mean - logistic — heavy-tailed symmetric (leptokurtic) - lognorm — positively skewed, wide support - weibull_min — right-skewed extreme-value; narrow and wide spread variants - -Two extensions beyond the baseline test_indexer_topk.py: - - pre_idx (heuristic candidates) - Shape [batch_size, index_topk]. For each batch element b, pre_idx[b] is - built from the base row's actual top-K: - - pre_idx[b, 0] = argmax (kernel invariant) - - floor(index_topk * success_ratio) slots drawn from actual top-K indices - - remaining slots filled with random valid indices - success_ratio is a pytest parameter (>= 0.4). - - MTP structure (next_n > 1) - When next_n > 1, consecutive rows within each batch element share most of - their logit values. For batch element b with valid_base = row_ends[b*next_n] - and MTP offset nni = 1…next_n-1: - logits[b*next_n + nni, nni : nni+valid_base] = logits[b*next_n, 0 : valid_base] - Positions 0..nni-1 are independently sampled (new token positions); - positions >= nni+valid_base remain -inf. - -Logit shapes (batch_size, next_n, num_tokens) match test_indexer_topk.py. +Correctness tests for the indexer Top-K custom ops: the CUDA +`indexer_topk_decode` / `indexer_topk_prefill` dispatchers +(insertion / radix / radix-split-work tiers) and the CuTe DSL +radix / filtered Top-K kernels. """ -import numpy as np import pytest import torch from utils.util import getSMVersion, skip_pre_blackwell, skip_pre_hopper @@ -54,15 +28,6 @@ if not torch.cuda.is_available(): pytest.skip("CUDA is required for indexer_topk tests", allow_module_level=True) -try: - import scipy.stats as _scipy_stats - from scipy.special import gamma as _gamma - - _HAS_SCIPY = True -except ImportError: - _HAS_SCIPY = False - - # --------------------------------------------------------------------------- # Prefill parameter helpers (unchanged from test_indexer_topk.py) # --------------------------------------------------------------------------- @@ -380,7 +345,7 @@ def test_indexer_topk_decode_launch_policy_transitions( # # The fix added two optional kwargs `radix_aux_indices` and # `radix_aux_logits` so the caller can supply persistent stable-address -# buffers (matching the existing `heuristic_scratch` convention). +# buffers with stable addresses (CUDA-graph safe). # # These tests verify: # (a) caller-owned-aux output matches the default (th::empty) path, @@ -870,460 +835,6 @@ def test_filtered_topk_varlen_odd_k(top_k, dtype_name): ) -# --------------------------------------------------------------------------- -# Distribution configs for heuristic decode correctness tests -# -# Each entry is a dict with keys: -# dist — distribution family -# mean — target mean (negative; typical decode logit range −0.5 to −4.5) -# std — target standard deviation -# full_range — support width (high − low), used as the bounding interval -# c — Weibull shape parameter (weibull_min only; c≈14 for moderate skew) -# -# Parameter derivation (all analytical, no external data dependencies): -# -# beta: -# low = mean − full_range/2, high = mean + full_range/2 -# mu01 = (mean − low) / full_range -# conc = mu01*(1−mu01) / (std/full_range)² − 1 -# α = conc*mu01, β = conc*(1−mu01) -# -# logistic: -# scale = std * √3 / π [std(logistic) = scale*π/√3] -# CDF inversion: x = mean + scale * ln(u/(1−u)), u ~ U(0,1) -# -# lognorm (left-shifted to loc = mean − full_range/2): -# pos_mean = full_range / 2 -# σ = √(log(1 + (std/pos_mean)²)) → matches target std exactly -# scale = exp(log(pos_mean) − σ²/2) → matches target mean exactly -# -# weibull_min (shape c fixed; loc/scale solved from mean/std): -# scale = std / √(Γ(1+2/c) − Γ²(1+1/c)) -# loc = mean − scale * Γ(1+1/c) -# --------------------------------------------------------------------------- - -_DECODE_DIST_CONFIGS = [ - # --- Beta: bounded, bell-shaped --- - # shallow negative mean, wide spread - dict(dist="beta", mean=-0.75, std=1.90, full_range=13.60), - # moderate negative mean - dict(dist="beta", mean=-2.96, std=1.68, full_range=12.85), - # deep negative mean, narrow spread - dict(dist="beta", mean=-4.51, std=1.75, full_range=11.24), - # --- Logistic: heavy-tailed symmetric (leptokurtic) --- - dict(dist="logistic", mean=-0.47, std=1.46, full_range=12.32), - # --- Lognorm: positively skewed, wide support --- - dict(dist="lognorm", mean=-4.12, std=2.55, full_range=17.28), - # --- Weibull minimum: right-skewed extreme-value --- - # wider spread - dict(dist="weibull_min", mean=-3.04, std=1.57, full_range=12.30, c=14.0), - # narrower spread - dict(dist="weibull_min", mean=-2.26, std=1.28, full_range=9.71, c=14.0), -] - -# Human-readable pytest IDs: dist_mean_std -_DECODE_DIST_IDS = [ - f"{c['dist']}_m{abs(c['mean']):.2f}_s{c['std']:.2f}" for c in _DECODE_DIST_CONFIGS -] - - -# --------------------------------------------------------------------------- -# Distribution-aware logit generator -# --------------------------------------------------------------------------- - - -def _fit_beta_params(mean: float, std: float, low: float, high: float): - """Fit Beta(α, β) on [low, high] to match target mean and std.""" - r = high - low - mu = (mean - low) / r - var = min((std / r) ** 2, mu * (1 - mu) * 0.99) - conc = mu * (1 - mu) / var - 1 - return conc * mu, conc * (1 - mu) - - -def _fit_weibull_params(mean: float, std: float, c: float): - """Fit Weibull_min(c, loc, scale) to match target mean and std.""" - g1 = _gamma(1 + 1 / c) - g2 = _gamma(1 + 2 / c) - scale = std / np.sqrt(g2 - g1**2) - return c, mean - scale * g1, scale - - -def create_distributed_logits( - cfg: dict, - row_starts: torch.Tensor, - row_ends: torch.Tensor, - dtype: torch.dtype, - seed: int, -) -> torch.Tensor: - """ - Generate a logits tensor sampled from the distribution specified by *cfg*. - - Values outside [row_start, row_end) are set to -inf. All distribution - parameters are derived analytically from (mean, std, full_range). - - Args: - cfg: One entry from _DECODE_DIST_CONFIGS - row_starts: (num_rows,) inclusive start column per row - row_ends: (num_rows,) exclusive end column per row - dtype: Target torch dtype - seed: NumPy RNG seed - - Returns: - Tensor (num_rows, max_len) with sampled values and -inf padding. - """ - rng = np.random.default_rng(seed) - num_rows = int(row_starts.shape[0]) - max_len = int(row_ends.cpu().max().item()) - # Pad to multiple of 8 so stride0 satisfies the alignment requirement of - # launchHeuristicTopKDecode for both fp32 (float4 = 4 elements) and - # bf16/fp16 (int4 = 16 B = 8 elements) in multi-row mode (matches TRT-LLM - # runtime where strides are always multiples of tokens_per_block >= 64). - max_len = (max_len + 7) & ~7 - size = (num_rows, max_len) - - dist = cfg["dist"] - mean, std, full_range = cfg["mean"], cfg["std"], cfg["full_range"] - low = mean - full_range / 2 - high = mean + full_range / 2 - - if dist == "beta": - alpha, beta_p = _fit_beta_params(mean, std, low, high) - samples = (rng.beta(alpha, beta_p, size=size) * (high - low) + low).astype(np.float32) - - elif dist == "logistic": - s = std * np.sqrt(3) / np.pi - u = rng.uniform(1e-7, 1 - 1e-7, size=size) - samples = (mean + s * np.log(u / (1 - u))).astype(np.float32) - - elif dist == "lognorm": - loc = low - pos_mean = max(mean - loc, 1e-6) # = full_range / 2 - sigma = float(np.sqrt(np.log(1 + (std / pos_mean) ** 2))) - scale = np.exp(np.log(pos_mean) - sigma**2 / 2) - samples = _scipy_stats.lognorm.rvs( - s=sigma, - loc=loc, - scale=scale, - size=size, - random_state=int(rng.integers(2**31 - 1)), - ).astype(np.float32) - - elif dist == "weibull_min": - c, loc, scale = _fit_weibull_params(mean, std, cfg.get("c", 14.0)) - samples = _scipy_stats.weibull_min.rvs( - c, - loc=loc, - scale=scale, - size=size, - random_state=int(rng.integers(2**31 - 1)), - ).astype(np.float32) - - else: - raise ValueError(f"Unknown distribution: {dist!r}") - - # Clip to [low, high] to bound the effective value range to exactly full_range. - # Unbounded distributions (lognorm, logistic, weibull_min) can produce outliers - # above `high` that inflate the histogram bin width in the kernel's 256-bin - # threshold search, causing boundary-element misidentification. - samples = np.clip(samples, low, high).astype(np.float32) - - logits = torch.from_numpy(samples).to(dtype=dtype, device="cuda") - col_idx = torch.arange(max_len, device="cuda").unsqueeze(0) - mask = (col_idx < row_starts.unsqueeze(1)) | (col_idx >= row_ends.unsqueeze(1)) - logits[mask] = float("-inf") - return logits - - -# --------------------------------------------------------------------------- -# MTP structure: make consecutive rows within a batch correlated -# --------------------------------------------------------------------------- - - -def apply_mtp_structure( - logits: torch.Tensor, - batch_size: int, - next_n: int, - row_ends: torch.Tensor, -) -> torch.Tensor: - """ - Enforce MTP (Multi-Token Prediction) logit correlation within each batch. - - For batch element b with base row valid length valid_base = row_ends[b*next_n], - each MTP offset nni = 1…next_n-1 satisfies: - - logits[b*next_n + nni, nni : nni+valid_base] = logits[b*next_n, 0 : valid_base] - - Positions 0..nni-1 of each MTP row remain independently sampled (new token - positions). Positions nni+valid_base.. are already -inf from - create_distributed_logits (since row_ends[b*next_n+nni] = valid_base+nni), - so no additional masking is required after this function. - - Args: - logits: (batch_size*next_n, max_len) float tensor; modified in-place - batch_size: number of batch elements - next_n: MTP factor; returns logits unchanged when next_n == 1 - row_ends: (batch_size*next_n,) exclusive end column per row - - Returns: - Same logits tensor with MTP segments overwritten. - """ - if next_n == 1: - return logits - - for b in range(batch_size): - base = b * next_n - valid_base = int(row_ends[base].item()) # valid length of base row - for nni in range(1, next_n): - # Copy base row positions [0:valid_base] → MTP row positions [nni:nni+valid_base] - # Positions 0..nni-1 stay independently sampled; positions >=nni+valid_base stay -inf. - logits[base + nni, nni : nni + valid_base] = logits[base, :valid_base] - - return logits - - -# --------------------------------------------------------------------------- -# pre_idx generator: heuristic candidate indices for the indexer kernel -# --------------------------------------------------------------------------- - - -def generate_pre_idx( - logits: torch.Tensor, - row_ends: torch.Tensor, - batch_size: int, - next_n: int, - index_topk: int, - success_ratio: float = 0.6, - seed: int = 0, -) -> torch.Tensor: - """ - Build the heuristic pre-prediction index tensor for each batch element. - - The V3.2 multi-row kernel (`heuristicTopKMultiRowKernel{,Dtype}` in - cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu) internally adds - `preIdxOffset = (rowIdx % next_n) + 1` to every pre_idx slot during its - Phase-1 stats reduction (heuristic_topk.cuh:654/1209). Production V3.2 - callers therefore pass pre_idx in PREVIOUS-step coordinates so the - kernel's +1 / +2 / +3 shift maps prev positions to current-step - positions correctly. - - This test builds pre_idx from `current_logits.topk()`, then applies a - `-1` shift before returning, so the kernel's internal `+1` brings every - hint back to its intended current-step position (preserves the kernel's - argmax invariant: kernel reads `input[(argmax_pos - 1) + 1] = input[argmax_pos]`). - - For batch element b (base row = b*next_n): - - pre_idx[b, 0] = argmax of the base row (kernel invariant) - - n_hit slots = floor(index_topk * success_ratio) indices drawn - WITHOUT replacement from the actual top-K - - n_fill = index_topk - n_hit slots - = indices drawn WITHOUT replacement from the - non-top-K pool (all valid indices except top-K) - - No element appears more than once in pre_idx[b]. The hit and fill pools - are disjoint by construction, so cross-pool duplicates are impossible. - - Edge case: when valid_len < index_topk (short sequences, ~10% of batches), - the non-top-K pool may be smaller than n_fill. In that case all available - non-top-K indices are used first; any remaining slots are filled from the - unused top-K tail (topk_idx[n_hit:]) to preserve the no-duplicate guarantee - as far as possible. - - Args: - logits: (batch_size*next_n, max_len) logits tensor - row_ends: (batch_size*next_n,) valid lengths per row - batch_size: number of batch elements - next_n: MTP factor; base row index is b*next_n - index_topk: number of pre-predicted candidates (K) - success_ratio: fraction of pre_idx drawn from actual top-K (>= 0.4) - seed: torch manual seed for reproducibility - - Returns: - pre_idx: int32 tensor of shape (batch_size, index_topk), no duplicates - """ - torch.manual_seed(seed) - pre_idx = torch.zeros(batch_size, index_topk, dtype=torch.int32, device=logits.device) - - for b in range(batch_size): - base = b * next_n - valid_len = int(row_ends[base].item()) - k = min(index_topk, valid_len) - - # Actual top-K of the base row (no duplicates); index 0 = argmax (kernel invariant) - _, topk_idx = logits[base, :valid_len].topk(k) - - # --- Hit slots: sample n_hit from top-K without replacement --- - # Always include argmax at position 0. - n_hit = max(1, int(k * success_ratio)) - n_hit = min(n_hit, k) - - if n_hit > 1: - perm = torch.randperm(k - 1, device=logits.device)[: n_hit - 1] - hits = torch.cat([topk_idx[:1], topk_idx[1:][perm]]) - else: - hits = topk_idx[:1] - - pre_idx[b, :n_hit] = hits.int() - - # --- Fill slots: sample n_fill from non-top-K pool without replacement --- - # The non-top-K pool is disjoint from topk_idx, so no cross-pool duplicates. - n_fill = index_topk - n_hit - if n_fill > 0: - # Build non-top-K pool: all valid indices that are NOT in topk_idx - topk_mask = torch.zeros(valid_len, dtype=torch.bool, device=logits.device) - topk_mask[topk_idx] = True - non_topk = torch.where(~topk_mask)[0] # shape: (valid_len - k,) - - if len(non_topk) >= n_fill: - # Normal case: enough non-top-K candidates - perm = torch.randperm(len(non_topk), device=logits.device)[:n_fill] - pre_idx[b, n_hit:] = non_topk[perm].int() - else: - # Edge case (valid_len ≈ index_topk): use all non-top-K first, - # then fill remaining from the unused top-K tail (topk_idx[n_hit:]) - pre_idx[b, n_hit : n_hit + len(non_topk)] = non_topk.int() - leftover = n_fill - len(non_topk) - topk_tail = topk_idx[n_hit:] # not yet in pre_idx[b] - take = min(leftover, len(topk_tail)) - if take > 0: - pre_idx[b, n_hit + len(non_topk) : n_hit + len(non_topk) + take] = topk_tail[ - :take - ].int() - - # V3.2 compensation: kernel adds `(rowIdx % next_n) + 1` to every pre_idx - # entry during P1 stats reduction. Shifting by -1 here means that for the - # base row (rowIdx % next_n == 0, offset = +1) the kernel reads the exact - # current-step positions our `topk()` selected. Negative entries are - # silently dropped by the kernel's `idx >= 0 && idx < N` range check. - pre_idx -= 1 - return pre_idx - - -def apply_mtp_structure_compressed( - logits: torch.Tensor, - batch_size: int, - next_n: int, - row_ends: torch.Tensor, -) -> torch.Tensor: - """ - cr=4-safe variant of apply_mtp_structure. - - apply_mtp_structure assumes ``row_ends[b*next_n + nni] == row_ends[b*next_n] - + nni`` (the V3.2 cr=1 invariant where each MTP draft adds exactly one KV - token). Under cr=4, ``row_ends = floor(actual_kv_len / 4)`` and that - invariant breaks: when ``actual_kv_len[base] mod 4`` lies in {1, 2, 3} - (75% of seq_lens) we get ``row_ends[base+nni] == row_ends[base]``, so the - copy ``[nni : nni+valid_base]`` overruns ``row_ends[base+nni]`` and writes - finite values into what create_distributed_logits left as -inf. The - polluted positions then leak into torch.topk's reference (which doesn't - know about the row's true compressed N), producing off-by-one counts vs. - the kernel. - - This variant clips the per-row copy length to fit within row b*next_n+nni's - valid compressed range, preserving MTP correlation where it fits and - leaving -inf positions untouched. - """ - if next_n == 1: - return logits - - for b in range(batch_size): - base = b * next_n - valid_base = int(row_ends[base].item()) - for nni in range(1, next_n): - row = base + nni - valid_row = int(row_ends[row].item()) - # Largest copy_len such that [nni, nni+copy_len) ⊆ [0, valid_row). - copy_len = max(0, min(valid_base, valid_row - nni)) - if copy_len > 0: - logits[row, nni : nni + copy_len] = logits[base, :copy_len] - - return logits - - -def generate_pre_idx_v4( - logits: torch.Tensor, - row_ends: torch.Tensor, - batch_size: int, - next_n: int, - index_topk: int, - success_ratio: float = 0.6, - seed: int = 0, -) -> torch.Tensor: - """ - DSv4 (compress_ratio=4) variant of generate_pre_idx — no `-1` shift. - - Unlike V3.2 where the kernel applies preIdxOffset = (rowIdx % next_n) + 1 - to every preIdx entry (KV grew by 1 per decode step in uncompressed space), - the V4 indexer operates in compressed-token-index space where consecutive - decode steps may add 0 or 1 compressed entries (each compressed entry - fuses 4 real tokens). Per-row Δc varies with prev kv_len mod 4 alignment, - but new compressed entries are always appended at the end so prev-step - indices in [0, c_prev-1] remain valid as-is in [0, c_curr-1]. The kernel - therefore forces preIdxOffset = 0 when compressRatio != 1, and tests must - pass preIdx in CURRENT-step coordinates (no -1 shift). - - Structure of the returned pre_idx[b]: - - pre_idx[b, 0] = argmax of the base row (kernel invariant) - - floor(K * success_ratio) slots from the actual top-K (without replace) - - remaining slots from non-top-K pool (without replace) - - Edge case (valid_len < K) handled identically to generate_pre_idx. - - Args: - logits, row_ends, batch_size, next_n, index_topk, success_ratio, seed: - See generate_pre_idx — the V4 helper mirrors its sampling logic. - - Returns: - pre_idx: int32 tensor of shape (batch_size, index_topk), entries in - the compressed current-step index space (no negative entries since - the kernel uses offset = 0). - """ - torch.manual_seed(seed) - pre_idx = torch.zeros(batch_size, index_topk, dtype=torch.int32, device=logits.device) - - for b in range(batch_size): - base = b * next_n - valid_len = int(row_ends[base].item()) - k = min(index_topk, valid_len) - - # Actual top-K of the base row; index 0 = argmax (kernel invariant). - _, topk_idx = logits[base, :valid_len].topk(k) - - n_hit = max(1, int(k * success_ratio)) - n_hit = min(n_hit, k) - - if n_hit > 1: - perm = torch.randperm(k - 1, device=logits.device)[: n_hit - 1] - hits = torch.cat([topk_idx[:1], topk_idx[1:][perm]]) - else: - hits = topk_idx[:1] - - pre_idx[b, :n_hit] = hits.int() - - n_fill = index_topk - n_hit - if n_fill > 0: - topk_mask = torch.zeros(valid_len, dtype=torch.bool, device=logits.device) - topk_mask[topk_idx] = True - non_topk = torch.where(~topk_mask)[0] - - if len(non_topk) >= n_fill: - perm = torch.randperm(len(non_topk), device=logits.device)[:n_fill] - pre_idx[b, n_hit:] = non_topk[perm].int() - else: - pre_idx[b, n_hit : n_hit + len(non_topk)] = non_topk.int() - leftover = n_fill - len(non_topk) - topk_tail = topk_idx[n_hit:] - take = min(leftover, len(topk_tail)) - if take > 0: - pre_idx[b, n_hit + len(non_topk) : n_hit + len(non_topk) + take] = topk_tail[ - :take - ].int() - - # No shift: kernel reads input[preIdx[i] + 0] = input[preIdx[i]] directly - # in compressed current-step coordinates. - return pre_idx - - # radix filter single-cta test. @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") @skip_pre_blackwell @@ -1643,283 +1154,6 @@ def run_fn(logits, seq_lens): ) -# ============================================================================ -# Heuristic Decode Distribution-Parameterised Tests -# ============================================================================ - - -@skip_pre_blackwell -@pytest.mark.skipif(not _HAS_SCIPY, reason="scipy required for distribution tests") -@pytest.mark.parametrize("success_ratio", [0.5, 0.9]) -@pytest.mark.parametrize("dist_cfg", _DECODE_DIST_CONFIGS, ids=_DECODE_DIST_IDS) -@pytest.mark.parametrize("batch_size", [1, 64, 128]) -@pytest.mark.parametrize("next_n", [1, 2, 3]) -@pytest.mark.parametrize("index_topk", [512, 1024, 2048]) -# num_tokens=4096 added to cover the new uniform kSeqSmall=4096 boundary -# across all K (indexerTopK.cu kSeqSmallDefaultForK). num_tokens=4096 sits -# right at the GVR routing threshold for every K ∈ {512, 1024, 2048} so the -# assertion validates the just-inside-GVR path correctness. -@pytest.mark.parametrize("num_tokens", [4096, 8192, 16384]) -@pytest.mark.parametrize( - "dtype", - [torch.float32, torch.bfloat16, torch.float16], - ids=["fp32", "bf16", "fp16"], -) -def test_indexer_topk_decode_dist( - dist_cfg, batch_size, next_n, index_topk, num_tokens, success_ratio, dtype -): - """ - Correctness test for the heuristic indexer_topk_decode across realistic - logit distributions, MTP correlation structures, pre_idx accuracy levels, - GVR-supported K values, and supported logit dtypes. - """ - torch.manual_seed(24) - torch.cuda.manual_seed(24) - - num_gen_tokens = batch_size * next_n - row_starts = torch.zeros(num_gen_tokens, dtype=torch.int32, device="cuda") - row_indices = torch.arange(num_gen_tokens, device="cuda") // next_n - next_n_offset = torch.arange(num_gen_tokens, device="cuda") % next_n - - seq_lens = generate_seq_lens(batch_size, index_topk, num_tokens) - # Clamp so that every base row has valid_len >= 1 (i.e., seq_len >= next_n). - # Without this, seq_len < next_n produces non-positive row_ends for offset 0. - seq_lens = seq_lens.clamp(min=next_n) - row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1 - - # 1. Sample logits from the target distribution - logits = create_distributed_logits(dist_cfg, row_starts, row_ends, dtype, seed=42) - - # 2. Apply MTP correlation: consecutive rows share their tail logits - if next_n > 1: - logits = apply_mtp_structure(logits, batch_size, next_n, row_ends) - - # 3. Build heuristic pre-prediction indices - pre_idx = generate_pre_idx( - logits, - row_ends, - batch_size, - next_n, - index_topk, - success_ratio=success_ratio, - seed=7, - ) - - # 4. Run heuristic CUDA kernel — heuristic_scratch dtype must match logits. - indices = torch.empty((num_gen_tokens, index_topk), dtype=torch.int32, device="cuda") - heuristic_scratch = torch.empty(num_gen_tokens * index_topk, dtype=dtype, device="cuda") - # Supply Radix split-work aux scratch. For dtype=fp32 with num_columns - # below kSeqSmall the dispatcher falls through GVR to the Radix path and - # the cpp op rejects blocks_per_row > 1 without caller-owned scratch; for - # bf16/fp16 these kwargs are simply ignored. - radix_aux_indices, radix_aux_logits = _build_radix_aux_buffers(num_gen_tokens, index_topk) - torch.ops.trtllm.indexer_topk_decode( - logits, - seq_lens, - indices, - next_n, - index_topk, - pre_idx, - heuristic_scratch, - radix_aux_indices=radix_aux_indices, - radix_aux_logits=radix_aux_logits, - ) - torch.cuda.synchronize() - - # 5. Reference: exact torch.topk masked to valid range - max_row_len = int(row_ends.max().item()) - torch_indices = logits.topk(min(index_topk, max_row_len), dim=-1)[1] - mask_lo = torch_indices >= 0 - mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0 - torch_indices = torch_indices.masked_fill(~(mask_lo & mask_hi), -1) - - # GVR Top-K is an exact algorithm: with same-dtype `logits.topk` as the - # reference, the sorted output values must be bit-identical (bf16 -> fp32 - # promotion inside the kernel is lossless and order-preserving, so the - # K-th cutoff is identical in both comparison spaces). Any value gap is - # a real kernel bug, not algorithmic noise — keep the default 1e-5 - # tolerance and let CI surface regressions. - assert compare_top_k_results( - logits, - indices, - torch_indices, - row_starts, - row_ends, - index_topk, - ), ( - f"heuristic indexer_topk_decode mismatch: dist={dist_cfg['dist']}, " - f"mean={dist_cfg['mean']}, std={dist_cfg['std']}, " - f"next_n={next_n}, success_ratio={success_ratio}, dtype={dtype}" - ) - - -# ============================================================================ -# DSv4 Heuristic Decode Test (compress_ratio = 4) -# ============================================================================ -# -# Exercises the V4 indexer GVR Top-K path enabled by the -# `compressRatio == 1 || compressRatio == 4` relaxation in -# canUseHeuristic (cpp/tensorrt_llm/kernels/indexerTopK.cu). For -# compressRatio != 1 the kernel: -# 1. Computes N = (seq_len - next_n + (rowIdx % next_n) + 1) / compressRatio, -# i.e. the row's compressed-KV length (vs. uncompressed N in the V3.2 -# path). -# 2. Forces preIdxOffset = 0 (vs. (rowIdx % next_n) + 1 in V3.2), since -# compressed entries are appended at the end of the compressed KV and -# prev-step indices remain valid as-is. -# -# To reach the GVR (Heuristic) path with cr=4 we need the *compressed* -# numColumns ≥ kSeqSmall (≈12288), so the test uses num_tokens ∈ -# {65536, 131072} which gives compressed range ≈ {16K, 32K}. Smaller cr=4 -# cases (where compressed N falls below kSeqSmall) are already covered by -# test_indexer_topk_decode parametrized on compress_ratio ∈ [1, 4] — those -# exercise the Radix/Insertion fallback for the same gate. - - -def _run_indexer_topk_decode_v4_gvr_check( - batch_size: int, - next_n: int, - index_topk: int, - num_tokens: int, - dtype: torch.dtype, - dist_cfg: dict, - success_ratio: float, -): - """Run the V4 (compress_ratio=4) heuristic indexer_topk_decode check.""" - torch.manual_seed(24) - torch.cuda.manual_seed(24) - - compress_ratio = 4 - num_gen_tokens = batch_size * next_n - row_starts = torch.zeros(num_gen_tokens, dtype=torch.int32, device="cuda") - row_indices = torch.arange(num_gen_tokens, device="cuda") // next_n - next_n_offset = torch.arange(num_gen_tokens, device="cuda") % next_n - - # Uncompressed seq_lens are what the kernel receives in `seq_lens`. - # Clamp so that compressed_actual_kv_len > kSeqSmall for every row; the - # kernel will divide actual_kv_len by compress_ratio internally, so a - # floor of (kSeqSmall + 1) * compress_ratio + next_n on the uncompressed - # seq_len guarantees compressed N stays in the GVR window. - # kSeqSmall is uniform 4096 across K (matches indexerTopK.cu - # kSeqSmallDefaultForK). - ksmall = 4096 - min_uncompressed = (ksmall + 1) * compress_ratio + next_n - if min_uncompressed >= num_tokens: - pytest.skip( - f"num_tokens={num_tokens} too small to clamp into the GVR window for " - f"K={index_topk} (needs uncompressed > {min_uncompressed})" - ) - seq_lens = generate_seq_lens(batch_size, min_uncompressed, num_tokens) - seq_lens = seq_lens.clamp(min=min_uncompressed) - - # row_ends is the compressed-KV length per row (= what logits' columns - # represent in V4 — the indexer operates in compressed-token-index space). - actual_kv_lens = seq_lens[row_indices] - next_n + next_n_offset + 1 - row_ends = actual_kv_lens // compress_ratio - - # 1. Sample logits over the compressed shape. - logits = create_distributed_logits(dist_cfg, row_starts, row_ends, dtype, seed=42) - - # 2. Apply MTP correlation between rows within each batch element. - # Use the compressed-aware variant: cr=4 breaks the cr=1 invariant - # row_ends[base+nni] = row_ends[base]+nni, so the copy length must be - # clipped per-row to avoid overrunning the row's valid range. - if next_n > 1: - logits = apply_mtp_structure_compressed(logits, batch_size, next_n, row_ends) - - # 3. Build heuristic pre-prediction indices — V4 variant (no -1 shift). - pre_idx = generate_pre_idx_v4( - logits, - row_ends, - batch_size, - next_n, - index_topk, - success_ratio=success_ratio, - seed=7, - ) - - # 4. Run heuristic CUDA kernel with compress_ratio=4. The kernel: - # - reads logits in compressed-index space (numColumns = logits.shape[1]) - # - divides seq_lens by compress_ratio to derive per-row N - # - uses preIdxOffset = 0 (preIdx already in current-step coords) - indices = torch.empty((num_gen_tokens, index_topk), dtype=torch.int32, device="cuda") - heuristic_scratch = torch.empty(num_gen_tokens * index_topk, dtype=dtype, device="cuda") - # Supply Radix split-work aux scratch — same rationale as the V3.2 helper: - # required by the cpp op when blocks_per_row > 1, harmless otherwise. - radix_aux_indices, radix_aux_logits = _build_radix_aux_buffers(num_gen_tokens, index_topk) - torch.ops.trtllm.indexer_topk_decode( - logits, - seq_lens, - indices, - next_n, - index_topk, - pre_idx, - heuristic_scratch, - compress_ratio=compress_ratio, - radix_aux_indices=radix_aux_indices, - radix_aux_logits=radix_aux_logits, - ) - torch.cuda.synchronize() - - # 5. Reference: torch.topk masked to the compressed row_ends. - max_row_len = int(row_ends.max().item()) - torch_indices = logits.topk(min(index_topk, max_row_len), dim=-1)[1] - mask = (torch_indices >= 0) & ((torch_indices - (row_ends - row_starts)[:, None]) < 0) - torch_indices = torch_indices.masked_fill(~mask, -1) - - assert compare_top_k_results( - logits, indices, torch_indices, row_starts, row_ends, index_topk - ), ( - f"V4 heuristic indexer_topk_decode (cr=4) mismatch: dist={dist_cfg['dist']}, " - f"mean={dist_cfg['mean']}, std={dist_cfg['std']}, batch_size={batch_size}, " - f"next_n={next_n}, index_topk={index_topk}, num_tokens={num_tokens}, " - f"success_ratio={success_ratio}, dtype={dtype}" - ) - - -# Param matrix is intentionally tighter than test_indexer_topk_decode_dist: -# only one logit distribution and one success_ratio because the GVR algorithm -# is dist-/hint-quality-invariant for correctness (an exact algorithm). The -# axes that *do* differ in V4 vs V3.2 are exercised in full: -# compress_ratio = 4 (fixed — sole purpose of this test) -# next_n in {1, 2, 3} (decode + MTP windows) -# index_topk in {512, 1024, 2048} (all GVR-supported K) -# num_tokens in {65536, 131072} (compressed N ≈ 16K and 32K) -# dtype: fp32 / bf16 / fp16 (both kernel templates) -# batch_size: 1 (single-row), 64 (multi-row) -@skip_pre_blackwell -@pytest.mark.skipif(not _HAS_SCIPY, reason="scipy required for distribution tests") -@pytest.mark.parametrize("success_ratio", [0.7]) -@pytest.mark.parametrize("batch_size", [1, 64]) -@pytest.mark.parametrize("next_n", [1, 2, 3]) -@pytest.mark.parametrize("index_topk", [512, 1024, 2048]) -# num_tokens=32768 added so that all K ∈ {512, 1024, 2048} hit the uniform -# kSeqSmall=4096 boundary at compress_ratio=4 (helper's clamp floor = -# (4096+1)*4+next_n ≈ 16389 < 32768, so no skip is triggered for any K). -@pytest.mark.parametrize("num_tokens", [32768, 65536, 131072]) -@pytest.mark.parametrize( - "dtype", - [torch.float32, torch.bfloat16, torch.float16], - ids=["fp32", "bf16", "fp16"], -) -def test_indexer_topk_decode_dist_v4_cr4( - batch_size, next_n, index_topk, num_tokens, success_ratio, dtype -): - """ - Correctness test for the DSv4 heuristic indexer_topk_decode with - compress_ratio=4 across MTP windows, all GVR-supported K, and all - supported logit dtypes. Uses one representative distribution; broader - distribution coverage is left to test_indexer_topk_decode_dist (cr=1). - """ - # Logistic chosen as the single representative distribution — its - # heavy-tailed symmetric shape produces the wide K-th-value spread that - # stresses GVR's secant threshold search most. - dist_cfg = dict(dist="logistic", mean=-0.47, std=1.46, full_range=12.32) - _run_indexer_topk_decode_v4_gvr_check( - batch_size, next_n, index_topk, num_tokens, dtype, dist_cfg, success_ratio - ) - - # ============================================================================ # CuTE DSL Prefill Top-K Tests # ============================================================================ @@ -2350,100 +1584,3 @@ def test_prefill_overflow_policy_overflow( dtype, row_start_offset=row_start_offset, ) - - -# ============================================================================ -# GVR Phase-3 threshold-repair regressions: hints that defeat the threshold -# search (undershoot / degenerate hint / tie plateau wider than kC) used to -# produce a silently wrong top-K (-1 pads or row[0:K]). -# ============================================================================ - - -def _gvr_decode_exact_check(logits_row, pre_idx_row, index_topk, tag): - """Run indexer_topk_decode (cr=4, BS=1) and assert a tie-aware exact top-K.""" - n = logits_row.shape[-1] - dtype = logits_row.dtype - logits = logits_row.view(1, n).contiguous() - pre_idx = pre_idx_row.view(1, index_topk).to(torch.int32).contiguous() - seq_lens = torch.full((1,), n * 4, dtype=torch.int32, device="cuda") - # -1 sentinel so unwritten slots trip the assertions below. - indices = torch.full((1, index_topk), -1, dtype=torch.int32, device="cuda") - scratch = torch.empty(index_topk, dtype=dtype, device="cuda") - aux_indices, aux_logits = _build_radix_aux_buffers(1, index_topk) - torch.ops.trtllm.indexer_topk_decode( - logits, - seq_lens, - indices, - 1, - index_topk, - pre_idx, - scratch, - compress_ratio=4, - radix_aux_indices=aux_indices, - radix_aux_logits=aux_logits, - ) - torch.cuda.synchronize() - - assert int((indices < 0).sum()) == 0, ( - f"{tag}: {int((indices < 0).sum())} of {index_topk} output slots are -1" - ) - # Distinctness: a duplicate+omission pair on a tie plateau would leave - # the sorted value multiset below unchanged. - n_unique = int(torch.unique(indices[0]).numel()) - assert n_unique == index_topk, ( - f"{tag}: only {n_unique} of {index_topk} output indices are distinct" - ) - flat = logits[0].float() - got = flat[indices[0].long()].sort().values - ref = flat.topk(index_topk).values.sort().values - assert torch.equal(got, ref), f"{tag}: selected values differ from torch.topk" - - -@skip_pre_blackwell -@pytest.mark.parametrize("index_topk", [512, 1024, 2048]) -@pytest.mark.parametrize("num_tokens", [65536, 131072]) -@pytest.mark.parametrize( - "dtype", [torch.float32, torch.bfloat16, torch.float16], ids=["fp32", "bf16", "fp16"] -) -@pytest.mark.parametrize("hint", ["bottom_k", "uniform_max", "random"]) -def test_indexer_topk_decode_gvr_hostile_hint(index_topk, num_tokens, dtype, hint): - """A hint that points away from the top-K must not change the result. - - ``uniform_max`` (every slot = argmax) additionally collapses Phase 1's - min/max bracket to a point, which used to short-circuit the kernel into - emitting row[0:K]. - """ - torch.manual_seed(1234) - logits = torch.randn(num_tokens, dtype=torch.float32, device="cuda").to(dtype) - flat = logits.float() - if hint == "bottom_k": - pre = flat.topk(index_topk, largest=False).indices - elif hint == "uniform_max": - pre = flat.argmax().repeat(index_topk) - else: - pre = torch.randint(0, num_tokens, (index_topk,), device="cuda") - _gvr_decode_exact_check(logits, pre, index_topk, f"hint={hint}") - - -@skip_pre_blackwell -@pytest.mark.parametrize("index_topk", [512, 1024, 2048]) -@pytest.mark.parametrize("n_tie", [6000, 20000, 100000]) -@pytest.mark.parametrize( - "dtype", [torch.float32, torch.bfloat16, torch.float16], ids=["fp32", "bf16", "fp16"] -) -def test_indexer_topk_decode_gvr_tie_plateau(index_topk, n_tie, dtype): - """More ties at the K-th value than the candidate buffer can hold: no - threshold lands in [K, kC], so the repair must emit the strictly-greater - set plus arbitrary ties. bf16/fp16 cover the reduced-precision driver's - separate direct-emit block.""" - torch.manual_seed(1234) - num_tokens = 131072 - n_above = index_topk // 2 - logits = torch.full((num_tokens,), -1.0, dtype=torch.float32, device="cuda") - logits[:n_above] = torch.linspace(2.0, 3.0, n_above, device="cuda") - logits[n_above : n_above + n_tie] = 1.0 - # Plateau (1.0) and floor (-1.0) are exact in every dtype; casting can - # only merge strictly-greater values with each other, which is tolerated. - logits = logits[torch.randperm(num_tokens, device="cuda")].contiguous().to(dtype) - pre = torch.randint(0, num_tokens, (index_topk,), device="cuda") - _gvr_decode_exact_check(logits, pre, index_topk, f"n_tie={n_tie}") From aaba6b0753954fa9f9ad926c6f4a4c0262b6440e Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:40:42 +0000 Subject: [PATCH 3/6] [None][refactor] Skip GVR prior state for the self-sampling engine The self-sampling engine keeps no cross-step state, so the framework no longer allocates it any: the per-layer gvr_prior_indices arena, the LJF row-reorder buffer, prefill seeding, the aux-stream write-back, and the indexer-side prior slice all key on needs_gvr_prior = two-level dispatch selecting the temporal engine. A shared use_self_sampling_gvr() predicate in dsa/params.py keeps the indexer's per-layer TopK construction and the metadata's allocation decision in agreement (live indexers only exist on cr in {1, 4} layers, matching the metadata's representative ratio). With the CUDA heuristic gone, the temporal engine requires the CuTe DSL on SM100/103; enable_heuristic_topk without it falls back to exact radix with a one-time warning. Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 58 +++++++++-------- .../attention_backend/sparse/dsa/metadata.py | 48 +++++++++----- .../attention_backend/sparse/dsa/params.py | 25 ++++++++ .../attention/sparse/dsa/test_dsa_indexer.py | 64 ++++++++++++++++--- 4 files changed, 145 insertions(+), 50 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 0884e831b327..e60dfc34a746 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -38,7 +38,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -from .params import DSAParams +from .params import DSAParams, use_self_sampling_gvr ModelConfig = tensorrt_llm.bindings.ModelConfig @@ -698,15 +698,13 @@ def __init__( # TopK module's hardware-format gate falls back to the exact # insertion/radix path with a one-time warning; contract violations # inside the engine raise. - self._use_self_sampling_topk = ( - sparse_params.use_self_sampling_topk - and self._enable_heuristic_topk - and IS_CUTLASS_DSL_AVAILABLE - # datacenter Blackwell only; consumer Blackwell (sm_120/121) - # lacks thread-block clusters - and get_sm_version() in (100, 103) - and sparse_params.index_topk in (512, 1024, 2048) - and compress_ratio in (1, 4) + self._use_self_sampling_topk = use_self_sampling_gvr( + enable_heuristic_topk=self._enable_heuristic_topk, + use_self_sampling_topk=sparse_params.use_self_sampling_topk, + index_topk=sparse_params.index_topk, + compress_ratio=compress_ratio, + is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, + sm_version=get_sm_version(), ) if os.environ.get("TRTLLM_GVR_SELF_SAMPLING") is not None: logger.warning_once( @@ -727,26 +725,35 @@ def __init__( f"(cutlass_dsl={IS_CUTLASS_DSL_AVAILABLE}, " f"sm={get_sm_version()}, " f"index_topk={sparse_params.index_topk}, " - f"compress_ratio={compress_ratio}); using the temporal GVR " - "path instead.", + f"compress_ratio={compress_ratio}); falling back to the " + "temporal GVR path (exact radix when the DSL engine is " + "unavailable).", key="gvr_self_sampling_prereq_fallback", ) self.mtp_index_share = sparse_params.mtp_index_share - if self.use_cute_dsl_topk: + if ( + self._enable_heuristic_topk + and IS_CUTLASS_DSL_AVAILABLE + # datacenter Blackwell only; consumer Blackwell (sm_120/121) + # lacks the thread-block clusters both GVR engines use + and get_sm_version() in (100, 103) + ): + decode_top_k_implementation = TopKImplementation.CUTE_DSL_GVR + else: + if self._enable_heuristic_topk: + logger.warning_once( + "enable_heuristic_topk=True but the DSL GVR engine is " + f"unavailable (cutlass_dsl={IS_CUTLASS_DSL_AVAILABLE}, " + f"sm={get_sm_version()}); using the exact radix decode " + "top-K instead.", + key="gvr_prereq_radix_fallback", + ) decode_top_k_implementation = ( - TopKImplementation.CUTE_DSL_GVR - if self._enable_heuristic_topk - else TopKImplementation.CUTE_DSL_RADIX + TopKImplementation.CUTE_DSL_RADIX + if self.use_cute_dsl_topk + else TopKImplementation.CUDA_RADIX ) - elif self._enable_heuristic_topk: - decode_top_k_implementation = TopKImplementation.CUDA_GVR - else: - decode_top_k_implementation = TopKImplementation.CUDA_RADIX - if self._use_self_sampling_topk: - # The self-sampling engine overrides the temporal decode - # implementation regardless of use_cute_dsl_topk. - decode_top_k_implementation = TopKImplementation.CUTE_DSL_GVR self.top_k = TopK( self.index_topk, prefill_implementation=TopKImplementation.CUDA_RADIX, @@ -1466,7 +1473,8 @@ def sparse_attn_indexer( num_gen_tokens = num_tokens - num_ctx_tokens gvr_prior_indices = None - if self._enable_heuristic_topk: + if self.top_k.needs_gvr_prior: + assert metadata.gvr_prior_indices is not None local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] gvr_prior_indices = metadata.gvr_prior_indices[local_layer] if is_generation is None: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 2b1c0521ebe6..91f6fc167139 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -31,7 +31,7 @@ _pick_dsl_expand, _select_indexer_compress_ratio, ) -from .params import DSAMetadataParams +from .params import DSAMetadataParams, use_self_sampling_gvr ModelConfig = tensorrt_llm.bindings.ModelConfig @@ -120,6 +120,11 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # Number of compressed KV tokens for context requests num_ctx_kv_tokens: int = 0 gen_indexer_kv_lens_cuda_runtime: Optional[torch.Tensor] = None + # Temporal-GVR prior state: allocated only when the two-level dispatch + # selects the temporal engine (the self-sampling engine keeps no + # cross-step state). + needs_gvr_prior: bool = field(default=False, init=False) + gvr_prior_indices: Optional[torch.Tensor] = field(default=None, init=False) def __init__(self, *args, **kwargs): """Initialize DSA metadata with SM count and indexer chunk size.""" @@ -188,7 +193,6 @@ def __post_init__(self): self.enable_gvr_topk = ( sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 ) - self.use_self_sampling_topk = sparse_metadata_params.use_self_sampling_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 @@ -204,6 +208,23 @@ def __post_init__(self): if hasattr(self.kv_cache_manager, "compressed_block_sizes"): tpb = tpb // _effective_compress_ratio_divisor(self._indexer_compress_ratio) self._tokens_per_block = tpb + # Mirror the indexer's two-level GVR decision. The representative + # compress ratio matches every live indexer: the DeepSeek-V4 backend + # only builds indexers on cr=4 layers and plain DSA uses cr=1. + self.use_self_sampling_topk = use_self_sampling_gvr( + enable_heuristic_topk=self.enable_gvr_topk, + use_self_sampling_topk=sparse_metadata_params.use_self_sampling_topk, + index_topk=self.num_sparse_topk, + compress_ratio=self._indexer_compress_ratio, + is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, + sm_version=get_sm_version(), + ) + self.needs_gvr_prior = ( + self.enable_gvr_topk + and IS_CUTLASS_DSL_AVAILABLE + and get_sm_version() in (100, 103) + and not self.use_self_sampling_topk + ) self.create_buffers_for_mla_rope_append(capture_graph=capture_graph) self.create_buffers_for_indexer(capture_graph=capture_graph) @@ -643,11 +664,7 @@ def _run_fused_dsa_decode_metadata(self): 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_gvr_topk - and self.use_cute_dsl_topk - and self.num_generations * next_n >= 2 * self.num_sms - ): + if self.needs_gvr_prior 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) @@ -957,7 +974,7 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - if self.enable_gvr_topk: + if self.needs_gvr_prior: self.gvr_prior_indices = self.get_empty( self.cuda_graph_buffers, ( @@ -970,14 +987,13 @@ def create_buffers_for_indexer(self, capture_graph=False): capture_graph=capture_graph, ) self.gvr_prior_indices.zero_() - 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, - ) + 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, + ) # Create expanded buffers for MTP support self.create_expanded_buffers(capture_graph=capture_graph) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py index 78957f38542a..fbea670a5978 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py @@ -20,6 +20,31 @@ pass +def use_self_sampling_gvr( + *, + enable_heuristic_topk: bool, + use_self_sampling_topk: bool, + index_topk: int | None, + compress_ratio: int, + is_cute_dsl_available: bool, + sm_version: int, +) -> bool: + """Return whether the two-level dispatch picks the self-sampling engine. + + Shared by the indexer (per-layer TopK construction) and the attention + metadata (prior-state allocation and warmup) so both sides of the + dispatch agree. + """ + return ( + enable_heuristic_topk + and use_self_sampling_topk + and is_cute_dsl_available + and sm_version in (100, 103) + and index_topk in (512, 1024, 2048) + and compress_ratio in (1, 4) + ) + + @dataclass(kw_only=True, slots=True) class DSABackendForwardArgs(SparseBackendForwardArgs): """DSA inputs passed from the MLA module to its backend.""" 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 4de20744923a..0c82c5a5911b 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -60,6 +60,7 @@ from tensorrt_llm._torch.attention_backend.sparse.dsa.indexer import ( transform_local_topk_and_prepare_pool_view_grouped, ) +from tensorrt_llm._torch.attention_backend.sparse.dsa.params import use_self_sampling_gvr from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation @@ -125,11 +126,15 @@ def _set_torch_top_k(indexer: Indexer) -> None: ) -def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): +@pytest.mark.parametrize("use_self_sampling", [True, False]) +def test_metadata_cache_geometry_comes_from_sparse_metadata_params(use_self_sampling): sparse_config = DeepSeekV4SparseAttentionConfig( compress_ratios=[1, 4, 128], index_head_dim=96, + index_topk=512, indexer_k_dtype="fp8", + enable_heuristic_topk=True, + use_self_sampling_topk=use_self_sampling, ) sparse_metadata_params = sparse_config.to_sparse_metadata_params() metadata = object.__new__(DSAtrtllmAttentionMetadata) @@ -143,8 +148,18 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): metadata.create_buffers_for_mla_rope_append = Mock() metadata.create_buffers_for_indexer = Mock() - with patch( - "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.TrtllmAttentionMetadata.__post_init__" + with ( + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.TrtllmAttentionMetadata.__post_init__" + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.IS_CUTLASS_DSL_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.get_sm_version", + return_value=100, + ), ): DSAtrtllmAttentionMetadata.__post_init__(metadata) @@ -152,6 +167,35 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): assert metadata.compress_ratios == [1, 4, 128] assert metadata._indexer_compress_ratio == 4 assert metadata._tokens_per_block == 64 + # The metadata mirror of the two-level dispatch drives prior allocation. + assert metadata.use_self_sampling_topk == use_self_sampling + assert metadata.needs_gvr_prior == (not use_self_sampling) + + +@pytest.mark.parametrize( + "kwargs,expected", + [ + (dict(), True), + (dict(sm_version=103, index_topk=2048, compress_ratio=4), True), + (dict(enable_heuristic_topk=False), False), + (dict(use_self_sampling_topk=False), False), + (dict(is_cute_dsl_available=False), False), + (dict(sm_version=120), False), + (dict(index_topk=256), False), + (dict(compress_ratio=2), False), + ], +) +def test_use_self_sampling_gvr(kwargs, expected): + base = dict( + enable_heuristic_topk=True, + use_self_sampling_topk=True, + index_topk=512, + compress_ratio=1, + is_cute_dsl_available=True, + sm_version=100, + ) + base.update(kwargs) + assert use_self_sampling_gvr(**base) is expected @pytest.mark.parametrize( @@ -217,8 +261,7 @@ def make_mock(num_generations, kv_lens_list): kv_lens_cuda = torch.tensor(kv_lens_list, dtype=torch.int32, device="cuda") row_order_buffer = torch.zeros(64, dtype=torch.int32, device="cuda") return SimpleNamespace( - enable_gvr_topk=True, - use_cute_dsl_topk=True, + needs_gvr_prior=True, num_generations=num_generations, num_sms=num_sms, max_draft_tokens=next_n - 1, @@ -277,7 +320,9 @@ def test_gvr_prior_writeback_uses_aux_stream(): enable_indexer_skip=True, ) indexer = create_indexer(sparse_config) - indexer._enable_heuristic_topk = True + # temporal GVR consumes the prior; force it independent of hardware + indexer.top_k.decode_implementation = TopKImplementation.CUTE_DSL_GVR + indexer.top_k.gvr_self_sampling = False indexer.aux_stream = torch.cuda.Stream() metadata.gvr_prior_indices = torch.zeros( (cache_manager.num_local_layers, batch_size, index_topk), @@ -342,6 +387,7 @@ def test_shared_topk_lifecycle(monkeypatch): metadata.enable_context_mla_with_cached_kv = False metadata.enable_indexer_skip = False metadata.enable_gvr_topk = False + metadata.needs_gvr_prior = False metadata.get_empty = Mock( side_effect=lambda _, shape, **kwargs: torch.empty(tuple(shape), dtype=kwargs["dtype"]) ) @@ -444,7 +490,7 @@ def test_indexer_post_load_weights_caches_fused_weight(): [ (False, False, TopKImplementation.CUDA_RADIX), (True, False, TopKImplementation.CUTE_DSL_RADIX), - (False, True, TopKImplementation.CUDA_GVR), + (False, True, TopKImplementation.CUTE_DSL_GVR), (True, True, TopKImplementation.CUTE_DSL_GVR), ], ) @@ -490,7 +536,7 @@ def test_indexer_configures_one_top_k_module( (True, False, TopKImplementation.CUTE_DSL_GVR), (True, True, TopKImplementation.CUTE_DSL_GVR), (False, True, TopKImplementation.CUTE_DSL_GVR), - (False, False, TopKImplementation.CUDA_GVR), + (False, False, TopKImplementation.CUTE_DSL_GVR), ], ) def test_indexer_two_level_gvr_dispatch( @@ -4021,8 +4067,8 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, index_topk, prefill_implementation=TopKImplementation.CUDA_RADIX, decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, ) - indexer._enable_heuristic_topk = True metadata_skip.gvr_prior_indices = torch.zeros( (cache_manager.num_local_layers, batch_size, index_topk), device="cuda", From 9aad420c7eb4c229281562e77852878a710d0b38 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:11:11 +0000 Subject: [PATCH 4/6] [None][feat] Emission block-skip as a third GVR decode top-K dispatch param Promote the emission-assisted block-skip optimization from the TRTLLM_GVR_EMISSION env var to a `use_gvr_emission` sparse-attention config field (default False). It only takes effect on the temporal-hint (V1) GVR path with FP4 paged-MQA logits; the self-sampling (V2) engine derives its bracket from the current row and never uses emission. Threads the field through llm_args -> model_config (V4 + V3.2 rebuilds) -> DSAParams / DSAMetadataParams -> the indexer gate, and adds config-threading unit tests. Made-with: Claude Code (Fable 5) Co-Authored-By: Claude Fable 5 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 12 ++++---- .../attention_backend/sparse/dsa/params.py | 4 +++ .../_torch/custom_ops/cute_dsl_custom_ops.py | 6 ++-- tensorrt_llm/_torch/model_config.py | 6 ++++ tensorrt_llm/llmapi/llm_args.py | 14 ++++++++++ .../attention/sparse/dsa/test_dsa_indexer.py | 28 +++++++++++++++++++ 6 files changed, 62 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index e60dfc34a746..c6dd1915d773 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -761,12 +761,14 @@ def __init__( compress_ratio=self.compress_ratio, gvr_self_sampling=self._use_self_sampling_topk, ) - # GVR emission-assisted decode (opt-in, experimental): the FP4/FP8 - # indexer epilogue emits candidates the GVR Top-K consumes (see - # gvr_emission / gvr_routing; state lives on the TopK module) - # only the FP4 scoring op accepts emission kwargs + # Emission block-skip is a temporal-hint (V1) optimization: the FP4 + # indexer epilogue emits per-block max logits the GVR Top-K consumes to + # skip whole blocks (see gvr_emission / gvr_routing; state lives on the + # TopK module). Off on the self-sampling path (no cross-step state to + # assist) and off non-FP4 / non-paged-MQA layers. self.use_gvr_emission = ( - os.environ.get("TRTLLM_GVR_EMISSION", "0") == "1" + sparse_params.use_gvr_emission + and not self._use_self_sampling_topk and decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR and self.use_cute_dsl_paged_mqa_logits and self.use_fp4 diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py index fbea670a5978..f9ed77bb04c1 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py @@ -67,6 +67,7 @@ class DSAMetadataParams(SparseMetadataParams): has_shared_indexer_layers: bool = False mtp_index_share: bool = False use_self_sampling_topk: bool = True + use_gvr_emission: bool = False @dataclass(frozen=True) @@ -88,6 +89,9 @@ class DSAParams(SparseParams): # temporal previous-step-hint engines (False). Only meaningful when # enable_heuristic_topk is set. use_self_sampling_topk: bool = True + # Emission block-skip for the temporal-hint engine; only meaningful with + # enable_heuristic_topk=True and use_self_sampling_topk=False on FP4. + use_gvr_emission: bool = False indexer_k_dtype: Literal["fp8", "fp4"] = "fp8" # Shared layers reuse the preceding full layer's top-k. is_full_indexer_layer: bool = True diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 22e2c7c901a6..b76ed27aee9d 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7840,9 +7840,9 @@ def cute_dsl_gvr_topk_decode( arg is None at call time (re-verified on the pinned torch), and most calls pass no hints. Under torch.compile/functionalization the undeclared write is invisible, so the hint path is eager / - CUDA-graph only. ``TRTLLM_GVR_EMISSION=1`` gates the - emission-assisted wiring that feeds these tensors (opt-in, - experimental). + CUDA-graph only. The ``use_gvr_emission`` sparse-attention config + field gates the emission-assisted wiring that feeds these tensors + (opt-in; temporal-hint path only). """ if not is_sm_100f(): raise ValueError( diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 2daf38752f3b..ed6f225c221a 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -998,6 +998,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): indexer_rope_interleave = sparse_attention_config.indexer_rope_interleave enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk use_self_sampling_topk = sparse_attention_config.use_self_sampling_topk + use_gvr_emission = sparse_attention_config.use_gvr_emission indexer_k_dtype = sparse_attention_config.indexer_k_dtype else: index_n_heads = pretrained_config.index_n_heads @@ -1012,6 +1013,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): indexer_rope_interleave = False enable_heuristic_topk = False use_self_sampling_topk = True + use_gvr_emission = False default_sparse_attention_config = DeepSeekV4SparseAttentionConfig( ) indexer_k_dtype = default_sparse_attention_config.indexer_k_dtype @@ -1029,6 +1031,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): indexer_config['indexer_rope_interleave'] = indexer_rope_interleave indexer_config['enable_heuristic_topk'] = enable_heuristic_topk indexer_config['use_self_sampling_topk'] = use_self_sampling_topk + indexer_config['use_gvr_emission'] = use_gvr_emission indexer_config['indexer_k_dtype'] = indexer_k_dtype return indexer_config @@ -1067,6 +1070,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): q_split_threshold = sparse_attention_config.q_split_threshold enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk use_self_sampling_topk = sparse_attention_config.use_self_sampling_topk + use_gvr_emission = sparse_attention_config.use_gvr_emission indexer_k_dtype = sparse_attention_config.indexer_k_dtype index_share_for_mtp_iteration = sparse_attention_config.index_share_for_mtp_iteration else: @@ -1080,6 +1084,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): q_split_threshold = 8192 enable_heuristic_topk = False use_self_sampling_topk = True + use_gvr_emission = False indexer_k_dtype = "fp8" index_share_for_mtp_iteration = None kwargs[ @@ -1097,6 +1102,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): indexer_rope_interleave=indexer_rope_interleave, enable_heuristic_topk=enable_heuristic_topk, use_self_sampling_topk=use_self_sampling_topk, + use_gvr_emission=use_gvr_emission, indexer_k_dtype=indexer_k_dtype, index_share_for_mtp_iteration= index_share_for_mtp_iteration) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index b68cf55f0974..c062fef23ca8 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -993,6 +993,16 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): "state; False runs the temporal-hint engines, which reuse the " "previous decode step's Top-K indices as hints. Ignored when " "enable_heuristic_topk is False.") + use_gvr_emission: bool = Field( + default=False, + description= + "Enable the emission-assisted block-skip optimization for the " + "temporal-hint GVR engine. When set, the FP4 indexer epilogue emits " + "per-block max logits so the GVR Top-K can skip whole blocks. Only " + "takes effect with enable_heuristic_topk=True, use_self_sampling_topk=" + "False, and the FP4 paged-MQA-logits path; ignored otherwise. The " + "self-sampling engine derives its bracket from the current row and " + "does not use emission.") indexer_k_dtype: Literal["fp8", "fp4"] = Field( default="fp8", description= @@ -1135,6 +1145,7 @@ def _value(name: str, default=None): indexer_rope_interleave=self.indexer_rope_interleave, enable_heuristic_topk=self.enable_heuristic_topk, use_self_sampling_topk=self.use_self_sampling_topk, + use_gvr_emission=self.use_gvr_emission, indexer_k_dtype=self.indexer_k_dtype, is_full_indexer_layer=self._is_full_indexer_layer( pretrained_config, kwargs.get("layer_idx")), @@ -1168,6 +1179,7 @@ def _value(name: str, default=None): enable_indexer_skip=self.skip_indexer_for_short_seqs, enable_heuristic_topk=self.enable_heuristic_topk, use_self_sampling_topk=self.use_self_sampling_topk, + use_gvr_emission=self.use_gvr_emission, 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, @@ -1256,6 +1268,7 @@ def _value(name: str, default=None): indexer_rope_interleave=self.indexer_rope_interleave, enable_heuristic_topk=self.enable_heuristic_topk, use_self_sampling_topk=self.use_self_sampling_topk, + use_gvr_emission=self.use_gvr_emission, indexer_k_dtype=self.indexer_k_dtype, compress_ratios=self.compress_ratios, window_size=self.window_size, @@ -1282,6 +1295,7 @@ def _value(name: str, default=None): enable_indexer_skip=self.skip_indexer_for_short_seqs, enable_heuristic_topk=self.enable_heuristic_topk, use_self_sampling_topk=self.use_self_sampling_topk, + use_gvr_emission=self.use_gvr_emission, 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, 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 0c82c5a5911b..88c840bf6515 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -198,6 +198,34 @@ def test_use_self_sampling_gvr(kwargs, expected): assert use_self_sampling_gvr(**base) is expected +@pytest.mark.parametrize("use_self_sampling_topk", [True, False]) +@pytest.mark.parametrize("use_gvr_emission", [False, True]) +def test_use_gvr_emission_threads_to_params(use_gvr_emission, use_self_sampling_topk): + """The emission block-skip flag (third GVR dispatch param) threads from the + sparse-attention config into both DSAParams and DSAMetadataParams, + independently of the V1/V2 selection. The runtime indexer gate additionally + requires the temporal-hint (V1) path + FP4 + paged-MQA to take effect.""" + sparse_config = DeepSeekV4SparseAttentionConfig( + compress_ratios=[1, 4, 128], + index_head_dim=96, + index_topk=512, + indexer_k_dtype="fp8", + enable_heuristic_topk=True, + use_self_sampling_topk=use_self_sampling_topk, + use_gvr_emission=use_gvr_emission, + ) + assert sparse_config.to_sparse_params().use_gvr_emission is use_gvr_emission + assert sparse_config.to_sparse_metadata_params().use_gvr_emission is use_gvr_emission + + +def test_use_gvr_emission_defaults_off(): + """Default keeps the emission block-skip optimization disabled end to end.""" + sparse_config = DeepSeekV4SparseAttentionConfig(index_topk=512) + assert sparse_config.use_gvr_emission is False + assert sparse_config.to_sparse_params().use_gvr_emission is False + assert sparse_config.to_sparse_metadata_params().use_gvr_emission is False + + @pytest.mark.parametrize( "enable_heuristic,use_cute_dsl,sm_version,compress_ratio,next_n,should_warmup", [ From 59d051ab0a8025560137998042b4ce1ab0173bc2 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 09:45:23 +0000 Subject: [PATCH 5/6] [None][test] Regenerate the LLM args telemetry golden manifest for the GVR dispatch fields Adds the two new sparse-attention config fields (use_self_sampling_topk, use_gvr_emission; both bool, captured by value) via scripts/generate_llm_args_golden_manifest.py so test_build_capture_manifest_matches_committed_golden passes again. Made-with: Claude Code (Fable 5.1) Co-Authored-By: Claude Fable 5.1 Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- tensorrt_llm/usage/llm_args_golden_manifest.json | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index d4a688885151..4a66a2fd2f37 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1828,6 +1828,20 @@ "kind": "value", "path": "sparse_attention_config.use_cute_dsl_topk" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.use_gvr_emission" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.use_self_sampling_topk" + }, { "allowed_values": [], "annotation": "", From b3a4fb2b4d3cc300612e528fd67d48dbb9a6d8ad Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sun, 6 Sep 2026 04:03:32 +0000 Subject: [PATCH 6/6] [None][infra] Restore blossom-ci allowlist to match main The merge-from-main left the branch behind on the blossom-ci authorized-user list (a stale-base artifact, not a GVR change); restore it so the PR does not drop 5 authorized users. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .github/workflows/blossom-ci.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/blossom-ci.yml b/.github/workflows/blossom-ci.yml index 8038f0a2fb0d..a4f22e88854b 100644 --- a/.github/workflows/blossom-ci.yml +++ b/.github/workflows/blossom-ci.yml @@ -168,6 +168,7 @@ jobs: "JadoTu", "jaedeok-nvidia", "janbernloehr", + "jasxu-nvidia", "jdebache", "jdemouth-nvidia", "JennyLiu-nv", @@ -191,6 +192,7 @@ jobs: "jthomson04", "juney-nvidia", "JunyiXu-nv", + "jupiterepoch", "JyChang012", "kaiyux", "Kambili", @@ -298,8 +300,10 @@ jobs: "RayenTian", "raymochen", "reasonsolo", + "richardc-nv", "richardhuo-nv", "rmccorm4", + "rmeghwal-nv", "roborluo", "RoeyAzran1992", "roikoren755", @@ -390,6 +394,7 @@ jobs: "xwang233", "xxi-nv", "yali-arch", + "yanxinzhangcs", "yechank-nvidia", "yibinl-nvidia", "yifeizhang-c",