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 01/10] [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 02/10] [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 03/10] [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 04/10] [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 05/10] [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 602a652c1048f21dce873847c332b45f42e49092 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:12:44 +0000 Subject: [PATCH 06/10] [None][feat] add Rubin topology support to GVR V2 decode Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 29 +- .../attention_backend/sparse/dsa/metadata.py | 26 +- .../attention_backend/sparse/dsa/params.py | 23 +- .../top_k/gvr_topk_decode_self_sampling.py | 76 ++-- .../gvr_topk_decode_self_sampling_host.py | 422 ++++++++++++++---- .../attention/sparse/dsa/test_dsa_indexer.py | 51 ++- .../parallel/test_gvr_selfsampling_topk.py | 41 +- .../test_gvr_selfsampling_topk_host.py | 328 ++++++++++++++ 8 files changed, 857 insertions(+), 139 deletions(-) create mode 100644 tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk_host.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index c6dd1915d773..56ac5459712b 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -16,7 +16,10 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.interface import MLAParams, PositionalEmbeddingParams -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE, +) from tensorrt_llm._torch.distributed.ops import allgather from tensorrt_llm._torch.modules.layer_norm import LayerNorm from tensorrt_llm._torch.modules.linear import Linear @@ -38,7 +41,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -from .params import DSAParams, use_self_sampling_gvr +from .params import DSAParams, is_gvr_cute_dsl_supported, use_self_sampling_gvr ModelConfig = tensorrt_llm.bindings.ModelConfig @@ -705,6 +708,13 @@ def __init__( compress_ratio=compress_ratio, is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, sm_version=get_sm_version(), + is_cute_dsl_rubin_available=IS_CUTLASS_DSL_RUBIN_AVAILABLE, + ) + gvr_cute_dsl_supported = is_gvr_cute_dsl_supported( + is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, + is_cute_dsl_rubin_available=IS_CUTLASS_DSL_RUBIN_AVAILABLE, + sm_version=get_sm_version(), + use_self_sampling_topk=self._use_self_sampling_topk, ) if os.environ.get("TRTLLM_GVR_SELF_SAMPLING") is not None: logger.warning_once( @@ -723,28 +733,23 @@ def __init__( "use_self_sampling_topk=True but the self-sampling GVR " "prerequisites are not met " f"(cutlass_dsl={IS_CUTLASS_DSL_AVAILABLE}, " + f"cutlass_dsl_rubin={IS_CUTLASS_DSL_RUBIN_AVAILABLE}, " f"sm={get_sm_version()}, " f"index_topk={sparse_params.index_topk}, " - f"compress_ratio={compress_ratio}); falling back to the " - "temporal GVR path (exact radix when the DSL engine is " - "unavailable).", + f"compress_ratio={compress_ratio}); falling back to temporal " + "GVR on SM100/103 or exact radix otherwise.", key="gvr_self_sampling_prereq_fallback", ) self.mtp_index_share = sparse_params.mtp_index_share - 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) - ): + if self._enable_heuristic_topk and gvr_cute_dsl_supported: 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"cutlass_dsl_rubin={IS_CUTLASS_DSL_RUBIN_AVAILABLE}, " f"sm={get_sm_version()}); using the exact radix decode " "top-K instead.", key="gvr_prereq_radix_fallback", diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 91f6fc167139..16af5f4c26a1 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -14,7 +14,10 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE, +) from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -31,7 +34,7 @@ _pick_dsl_expand, _select_indexer_compress_ratio, ) -from .params import DSAMetadataParams, use_self_sampling_gvr +from .params import DSAMetadataParams, is_gvr_cute_dsl_supported, use_self_sampling_gvr ModelConfig = tensorrt_llm.bindings.ModelConfig @@ -218,12 +221,16 @@ def __post_init__(self): compress_ratio=self._indexer_compress_ratio, is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, sm_version=get_sm_version(), + is_cute_dsl_rubin_available=IS_CUTLASS_DSL_RUBIN_AVAILABLE, + ) + gvr_cute_dsl_supported = is_gvr_cute_dsl_supported( + is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, + is_cute_dsl_rubin_available=IS_CUTLASS_DSL_RUBIN_AVAILABLE, + sm_version=get_sm_version(), + use_self_sampling_topk=self.use_self_sampling_topk, ) 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.enable_gvr_topk and gvr_cute_dsl_supported and not self.use_self_sampling_topk ) self.create_buffers_for_mla_rope_append(capture_graph=capture_graph) @@ -410,7 +417,12 @@ def warmup_selfsampling_topk( """ # 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): + if not is_gvr_cute_dsl_supported( + is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, + is_cute_dsl_rubin_available=IS_CUTLASS_DSL_RUBIN_AVAILABLE, + sm_version=get_sm_version(), + use_self_sampling_topk=True, + ): return if not self.enable_gvr_topk or self.kv_cache_manager is None: return diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py index f9ed77bb04c1..b16b530661d0 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py @@ -20,6 +20,20 @@ pass +def is_gvr_cute_dsl_supported( + *, + is_cute_dsl_available: bool, + is_cute_dsl_rubin_available: bool, + sm_version: int, + use_self_sampling_topk: bool, +) -> bool: + """Return whether the CuTe DSL stack can run GVR on this architecture.""" + return is_cute_dsl_available and ( + sm_version in (100, 103) + or (sm_version == 107 and is_cute_dsl_rubin_available and use_self_sampling_topk) + ) + + def use_self_sampling_gvr( *, enable_heuristic_topk: bool, @@ -28,6 +42,7 @@ def use_self_sampling_gvr( compress_ratio: int, is_cute_dsl_available: bool, sm_version: int, + is_cute_dsl_rubin_available: bool = False, ) -> bool: """Return whether the two-level dispatch picks the self-sampling engine. @@ -38,8 +53,12 @@ def use_self_sampling_gvr( return ( enable_heuristic_topk and use_self_sampling_topk - and is_cute_dsl_available - and sm_version in (100, 103) + and is_gvr_cute_dsl_supported( + is_cute_dsl_available=is_cute_dsl_available, + is_cute_dsl_rubin_available=is_cute_dsl_rubin_available, + sm_version=sm_version, + use_self_sampling_topk=True, + ) and index_topk in (512, 1024, 2048) and compress_ratio in (1, 4) ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index a07ba529eb6d..c1a856b35025 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Self-sampling GVR top-K decode kernels (CuTe DSL, Blackwell sm_100a). +"""Self-sampling GVR top-K decode kernels (CuTe DSL, SM100/103 and SM107). + +SM107/Rubin is enabled only when the installed CuTe DSL exposes its Rubin +helpers. This is source support, not an R200 performance-tuning claim. Sample-calibrated threshold ladders for exact single-pass top-K: the kernel derives its selection threshold from an in-kernel sample of the row itself @@ -3038,18 +3041,29 @@ def workspace_bytes() -> int: return WS_BYTES +def _host_dispatch_for(logits): + """Load the host router and resolve the tensor's execution-domain SMs. + + Standalone/debug entries share the production topology resolver so SM107 + cannot silently inherit the historical B200 148-SM plan. The import stays + lazy to avoid a device/host cycle during module initialization. + """ + try: + from . import gvr_topk_decode_self_sampling_host as ct_dispatch + except ImportError: + import gvr_topk_decode_self_sampling_host as ct_dispatch + return ct_dispatch, ct_dispatch._available_num_sms(logits.get_device()) + + def run(logits, pre_idx, n: int, out, ws): """torch-facing single-call entry: routes (b, n, k) through ct_dispatch, asserts the shape lands on gvr_main, launches the matching variant. ws: zero-initialised >=20,973,568-B CUDA buffer (reused across launches; the kernel restores the zeros it consumes).""" - try: - from . import gvr_topk_decode_self_sampling_host as ct_dispatch - except ImportError: - import gvr_topk_decode_self_sampling_host as ct_dispatch + ct_dispatch, num_sms = _host_dispatch_for(logits) b, npad = logits.shape k = pre_idx.shape[1] - r = ct_dispatch.route(b, int(n), npad, k) + r = ct_dispatch.route(b, int(n), npad, k, num_sms=num_sms) assert r["kernel"] == "main", f"shape routes to {r['kernel']}, not gvr_main" assert ws.numel() * ws.element_size() >= WS_BYTES rt = r["rt"] @@ -4369,12 +4383,14 @@ def reg_topk(logits, pre_idx, n, out, rd=None): logits [b, npad] f32, pre_idx [b, k] i32, out [b, >=k] i32, n = valid len. rd: optional pre-computed ct_dispatch.route() dict (must be reg/regimg). """ + ct_dispatch, num_sms = _host_dispatch_for(logits) + expected = ct_dispatch.route( + logits.shape[0], int(n), logits.shape[1], pre_idx.shape[1], num_sms=num_sms + ) if rd is None: - try: - from .gvr_topk_decode_self_sampling_host import route - except ImportError: - from gvr_topk_decode_self_sampling_host import route - rd = route(logits.shape[0], int(n), logits.shape[1], pre_idx.shape[1]) + rd = expected + else: + assert rd == expected, "pre-computed route does not match the current execution domain" assert rd["kernel"] in ("reg", "regimg"), rd["kernel"] tpl = rd["tpl"] rt = rd["rt"] @@ -4473,6 +4489,8 @@ def __init__( next_n: int = 1, cr_shift: int = 0, hint_free: bool = False, + *, + num_sms: int, ) -> None: assert blk == 1024, "gvr_clus is always BLK=1024" assert minb == 1, "gvr_clus is __launch_bounds__(BLK, 1)" @@ -4489,6 +4507,8 @@ def __init__( if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) self.hint_free = bool(hint_free) # hint-free: gather_hint sites compiled out + self.num_sms = int(num_sms) + assert self.num_sms >= 1 self.lcs = cs.bit_length() - 1 # log2(CS) for the per-row Q shift self.blk = blk self.u = u @@ -4739,7 +4759,7 @@ def kern( if x6 - ri * ri > ri: r6 = ri + cutlass.Int32(1) # aim_base: R = CS > 1 always for this family; bigf is the - # launch-computed occupancy flag (num_rows * CS <= 148). + # launch-computed occupancy flag (num_rows * CS <= num_sms). aim = k << cutlass.Int32(1) if bigf == cutlass.Int32(0): if k >= cutlass.Int32(1024): @@ -5631,7 +5651,7 @@ def __call__( # of (rows, CS) so it is launch-computed, not an ABI scalar. b = out.shape[0] bigf = cutlass.Int32(0) - if b * cutlass.Int32(self.cs) <= cutlass.Int32(148): + if b * cutlass.Int32(self.cs) <= cutlass.Int32(self.num_sms): bigf = cutlass.Int32(1) self.kern( logits, pre_idx, kv_lens, out, n, npad, k, SCAP, CMP, SMP, TGT, Q, SS2, TGT2, bigf @@ -5659,6 +5679,8 @@ def get_compiled__clus( next_n: int = 1, cr_shift: int = 0, hint_free: bool = False, + *, + num_sms: int, ) -> Any: """Compile (or fetch) the gvr_clus variant for constexpr tuple tpl = (BLK, U, MINB, NBS, CS); scap/cmp are smem-extent keys (every @@ -5672,6 +5694,7 @@ def get_compiled__clus( int(next_n), int(cr_shift), bool(hint_free), + int(num_sms), ) hit = _COMPILE_CACHE__clus.get(key) if hit is not None: @@ -5689,6 +5712,7 @@ def get_compiled__clus( next_n=next_n, cr_shift=cr_shift, hint_free=hint_free, + num_sms=num_sms, ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -5727,18 +5751,15 @@ def run__clus(logits, pre_idx, n: int, out): gvr_clus takes NO workspace.""" import torch # debug-entry only: module stays torch-free at import - try: - from . import gvr_topk_decode_self_sampling_host as ct_dispatch - except ImportError: - import gvr_topk_decode_self_sampling_host as ct_dispatch + ct_dispatch, num_sms = _host_dispatch_for(logits) b, npad = logits.shape k = pre_idx.shape[1] - r = ct_dispatch.route(b, int(n), npad, k) + r = ct_dispatch.route(b, int(n), npad, k, num_sms=num_sms) assert r["kernel"] == "clus", f"shape routes to {r['kernel']}, not gvr_clus" rt = r["rt"] - kobj = GvrClusKernel(*r["tpl"], scap=rt["SCAP"], cmp_=rt["CMP"]) + kobj = GvrClusKernel(*r["tpl"], scap=rt["SCAP"], cmp_=rt["CMP"], num_sms=num_sms) assert r["smem"] == kobj.dyn_bytes, (r["smem"], kobj.dyn_bytes) - fn = get_compiled__clus(tuple(r["tpl"]), scap=rt["SCAP"], cmp_=rt["CMP"]) + fn = get_compiled__clus(tuple(r["tpl"]), scap=rt["SCAP"], cmp_=rt["CMP"], num_sms=num_sms) dkv = torch.zeros(1, dtype=torch.int32, device=logits.device) # dead varlen slot fn( logits, @@ -5765,7 +5786,8 @@ def run_manual(logits, pre_idx, n: int, out, tpl, rt): the same CS; U only changes the chunk geometry).""" import torch # debug-entry only: module stays torch-free at import - fn = get_compiled__clus(tuple(tpl), scap=rt["SCAP"], cmp_=rt["CMP"]) + _, num_sms = _host_dispatch_for(logits) + fn = get_compiled__clus(tuple(tpl), scap=rt["SCAP"], cmp_=rt["CMP"], num_sms=num_sms) dkv = torch.zeros(1, dtype=torch.int32, device=logits.device) # dead varlen slot fn( logits, @@ -6662,12 +6684,14 @@ def regclus_topk(logits, pre_idx, n, out, rd=None): logits [b, npad] f32, pre_idx [b, k] i32, out [b, >=k] i32, n = valid len. rd: optional pre-computed ct_dispatch.route() dict (must be reg_clus). """ + ct_dispatch, num_sms = _host_dispatch_for(logits) + expected = ct_dispatch.route( + logits.shape[0], int(n), logits.shape[1], pre_idx.shape[1], num_sms=num_sms + ) if rd is None: - try: - from .gvr_topk_decode_self_sampling_host import route - except ImportError: - from gvr_topk_decode_self_sampling_host import route - rd = route(logits.shape[0], int(n), logits.shape[1], pre_idx.shape[1]) + rd = expected + else: + assert rd == expected, "pre-computed route does not match the current execution domain" assert rd["kernel"] == "reg_clus", rd["kernel"] tpl = tuple(rd["tpl"]) assert pre_idx.shape[1] <= tpl[0], "k <= BLK enforced by dispatch" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index cea15314e1f4..51d837997789 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -15,15 +15,18 @@ """Self-sampling GVR top-K decode — host side (dispatch, workspace, entry). Companion to ``gvr_topk_decode_self_sampling.py`` (the device module). +Supports SM100/103 and capability-gated SM107/Rubin; the Rubin path is not +yet performance-tuned or R200-validated. Three sections: 1. dispatch — the CUDA host dispatch as a pure function - ``route(b, n, npad, k)``; -2. workspace — one zero-initialised per-device slab (20,973,568 B) via the - torch caching allocator, with keep-alive + double-checked locking; + ``route(b, n, npad, k, num_sms=148)``; +2. workspace — one zero-initialised slab per device/execution domain + (20,973,568 B) via the matching torch caching allocator or localized + mempool, with keep-alive + double-checked locking; 3. operator entry — ``run(logits, pre_idx, n_valid, indices)`` / ``run_ws(..., workspace)`` DPS forms with input hardening and a - bind-once launch cache keyed on ``(b, n, npad, k)``. + bind-once launch cache keyed on shape plus execution topology. OPERATOR CONTRACT (batch-uniform entries): ``n_valid`` is one host python int for the whole batch — every row shares the same valid prefix, in @@ -72,8 +75,8 @@ def _device(): # =========================================================================== """Pure-Python mirror of the GVR CUDA host dispatch (gvr_topk_launch). -route(b, n, npad, k) is a PURE function of its four ints -- no env knobs, no -GPU, stdlib only. It returns the kernel family, its compile-time template +route(b, n, npad, k, num_sms=148) is a PURE function of its integer inputs -- +no env knobs or GPU queries. It returns the kernel family, its compile-time template tuple, the runtime scalar pack `rt`, grid/cluster/block geometry, smem size, and whether the family needs the workspace. @@ -105,19 +108,40 @@ def _device(): SNB = 256 # streaming-path bin count CMPC = 4096 # crossing-bin slots per CTA, clustered register path BLKC = 1024 # CTA size of the clustered register path +DEFAULT_NUM_SMS = 148 # B200 default; preserves the original pure-function contract +_CLUSTER8_MAX_BATCH = 15 # B200-validated GPC packing limit +MAX_SPLIT_ROWS = 160 # workspace MAXC; must match the device kernel -def route(b: int, n: int, npad: int, k: int) -> dict[str, object]: +def _cluster8_is_supported(batch_size: int, num_sms: int) -> bool: + """Apply the validated cluster-8 packing limit conservatively. + + The 15-row limit is a GPC-packing constraint, not an SM-wave ratio. Keep + it unchanged for Rubin until R200 cluster residency is characterized; + the ``num_sms`` argument makes the architecture-policy seam explicit. + """ + return batch_size <= min(_CLUSTER8_MAX_BATCH, num_sms) + + +def route( + b: int, + n: int, + npad: int, + k: int, + num_sms: int = DEFAULT_NUM_SMS, +) -> dict[str, object]: """Mirror of the CUDA gvr_topk_launch dispatch. Pure. See module doc.""" if b < 1: raise RuntimeError(f"route requires b >= 1, got {b}") - wide = b <= 148 + if num_sms < 1: + raise RuntimeError(f"route requires num_sms >= 1, got {num_sms}") + wide = b <= num_sms # ======================= register-resident block ======================== n4 = n >> 2 CMP = n if n < 2560 else 2560 - QC = 1024 if b > 148 else QUADC - CURE = not (n < 2 * k and b > 148) + QC = 1024 if b > num_sms else QUADC + CURE = not (n < 2 * k and b > num_sms) DEGE = (n <= 3 * k) or (n <= 4 * k + 64) if DEGE and CMP < n: CMP = n @@ -184,7 +208,7 @@ def _reg(BLK, VPT, MINB, NBH): # ---- clustered register-resident path ---- if n4 > 4096 and n4 <= 8 * BLKC * 4 and k <= BLKC: - av = 148 // (b if b > 0 else 1) # truncating + av = num_sms // (b if b > 0 else 1) # truncating amax = 1 while (amax << 1) <= av and amax < 8: amax <<= 1 @@ -197,7 +221,7 @@ def _reg(BLK, VPT, MINB, NBH): c = 1 # 64-bit product in C while c * BLKC * v < n4: c <<= 1 - if c == 8 and b > 15: # the veto + if c == 8 and not _cluster8_is_supported(b, num_sms): continue if c <= amax: vsel = v @@ -222,7 +246,7 @@ def _reg(BLK, VPT, MINB, NBH): # ====================== streaming / collect path ======================== R = 1 if b <= 32: - r1 = 148 // b + r1 = num_sms // b if r1 < 1: r1 = 1 r2 = ((n >> 2) + 1023) // 1024 @@ -231,7 +255,8 @@ def _reg(BLK, VPT, MINB, NBH): R = r1 if r1 < r2 else r2 if R < 1: R = 1 - elif b <= 74 and (n >> 2) >= 16384 and k <= 1024: # shallow R=2 split + elif b <= min(num_sms // 2, MAX_SPLIT_ROWS) and (n >> 2) >= 16384 and k <= 1024: + # shallow R=2 split; MAXC limits split rows, not CTAs per row R = 2 useclus = False @@ -241,12 +266,12 @@ def _reg(BLK, VPT, MINB, NBH): p2 <<= 1 # gvr_clus cs=8 hits the same GPC packing wall as the clustered # register path; same veto, same b > 15 threshold. - if p2 == 8 and b > 15: + if p2 == 8 and not _cluster8_is_supported(b, num_sms): p2 = 4 R = p2 useclus = True - big = b * R <= 148 + big = b * R <= num_sms SCAP = (16384 if R == 1 else 8192) if big else (8192 if k > 1024 else 4096) CMP = (4096 if k > 1024 else 2048) if big else 1024 @@ -352,7 +377,7 @@ def _reg(BLK, VPT, MINB, NBH): "ws": False, } - smem_main = (SCAP + 4) * (8 if (R > 1 or b <= 296) else 4) + (CMP + 1) * 8 + smem_main = (SCAP + 4) * (8 if (R > 1 or b <= 2 * num_sms) else 4) + (CMP + 1) * 8 def _main(BLK, MINB, U, SPLIT): # KPT ladder 1/2/4/8; grid = (R, b). @@ -391,7 +416,7 @@ def _main(BLK, MINB, U, SPLIT): per = Q >> 10 U = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) return _main(1024, 1, U, R > 1) # SPLIT iff R>1 - if b <= 296: + if b <= 2 * num_sms: return _main(512, 2, 8, False) return _main(256, 4, 8, False) @@ -425,13 +450,13 @@ def _main(BLK, MINB, U, SPLIT): # route_static(b, n, npad, k) — everything that must be frozen per launch: # family, compile tuple, grid, cluster, block, and the rt scalars that # change only at discrete n-thresholds; -# route_dynamic(static, n) — the n-continuous scalars a per-row kernel +# route_dynamic(static, n, num_sms=...) — the n-continuous scalars a per-row kernel # recomputes from its own row length (the device code will mirror these # formulas): n, CMP (reg families), the sampling ladder # SMP/TGT/SS2/TGT2/Q (streaming families), and the reg-family smem # footprint. # INVARIANT: merging route_dynamic back into route_static reproduces -# route() EXACTLY for every n. The policy of which n to freeze the static +# route() EXACTLY for every n and execution topology. The policy of which n to freeze the static # half at (e.g. max_seq_len) is a perf-only choice — the factorization # itself is lossless. @@ -445,11 +470,17 @@ def _main(BLK, MINB, U, SPLIT): _DYN_SMEM = ("reg", "regimg") # smem depends on CMP/IMGW -> recomputed per n -def route_static(b: int, n: int, npad: int, k: int) -> dict[str, object]: +def route_static( + b: int, + n: int, + npad: int, + k: int, + num_sms: int = DEFAULT_NUM_SMS, +) -> dict[str, object]: """route() with the n-continuous fields redacted (see _DYN_RT/_DYN_SMEM). Constant on maximal n-intervals ("bands"); every redacted field is - reconstructible from (static, n) by route_dynamic.""" - plan = route(b, n, npad, k) + reconstructible from (static, n, num_sms) by route_dynamic.""" + plan = route(b, n, npad, k, num_sms=num_sms) st = {key: (dict(val) if isinstance(val, dict) else val) for key, val in plan.items()} for f in _DYN_RT[st["kernel"]]: st["rt"].pop(f) @@ -458,8 +489,17 @@ def route_static(b: int, n: int, npad: int, k: int) -> dict[str, object]: return st -def route_dynamic(static: dict[str, object], n: int) -> tuple[dict[str, object], int]: - """Recompute the redacted n-continuous scalars from (static, n). +def route_dynamic( + static: dict[str, object], + n: int, + *, + num_sms: int, +) -> tuple[dict[str, object], int]: + """Recompute redacted scalars from ``(static, n, num_sms)``. + + ``num_sms`` is required because the static plan cannot in general reveal + which execution-domain topology produced it. + Returns (rt_updates, smem). Must stay equivalent to route(); the device-side per-row engine mirrors exactly these formulas.""" fam = static["kernel"] @@ -485,7 +525,7 @@ def route_dynamic(static: dict[str, object], n: int) -> tuple[dict[str, object], else: R = static["rt"]["R"] scap = static["rt"]["SCAP_"] - big = b * R <= 148 + big = b * R <= num_sms aim = ( ((4 * k if k >= 1024 else 2 * k) if R == 1 else 2 * k) if big @@ -535,11 +575,17 @@ def route_dynamic(static: dict[str, object], n: int) -> tuple[dict[str, object], ) -def route_split(b: int, n: int, npad: int, k: int) -> dict[str, object]: +def route_split( + b: int, + n: int, + npad: int, + k: int, + num_sms: int = DEFAULT_NUM_SMS, +) -> dict[str, object]: """route_static + route_dynamic recombined — must equal route() exactly (the factorization fuzz in the unit tests asserts this).""" - st = route_static(b, n, npad, k) - dyn, smem = route_dynamic(st, n) + st = route_static(b, n, npad, k, num_sms=num_sms) + dyn, smem = route_dynamic(st, n, num_sms=num_sms) plan = {key: (dict(val) if isinstance(val, dict) else val) for key, val in st.items()} plan["rt"].update(dyn) plan["smem"] = smem @@ -547,7 +593,12 @@ def route_split(b: int, n: int, npad: int, k: int) -> dict[str, object]: def route_streaming( - b: int, n: int, npad: int, k: int, force_main: bool = False + b: int, + n: int, + npad: int, + k: int, + force_main: bool = False, + num_sms: int = DEFAULT_NUM_SMS, ) -> dict[str, object]: """route() restricted to its STREAMING half (main / clus) — the varlen capture policy: per-row kernels must be picked from the families that are @@ -558,23 +609,25 @@ def route_streaming( min(r1, r2) R matches the CUDA else-branch exactly.""" if b < 1: raise RuntimeError(f"route_streaming requires b >= 1, got {b}") + if num_sms < 1: + raise RuntimeError(f"route_streaming requires num_sms >= 1, got {num_sms}") R = 1 if b <= 32: - r1 = max(148 // b, 1) + r1 = max(num_sms // b, 1) r2 = max(((n >> 2) + 1023) // 1024, 1) R = max(min(r1, r2), 1) - elif b <= 74 and (n >> 2) >= 16384 and k <= 1024: + elif b <= min(num_sms // 2, MAX_SPLIT_ROWS) and (n >> 2) >= 16384 and k <= 1024: R = 2 useclus = False if not force_main and 2 <= R <= 8 and k <= 1024: p2 = 1 while (p2 << 1) <= R: p2 <<= 1 - if p2 == 8 and b > 15: + if p2 == 8 and not _cluster8_is_supported(b, num_sms): p2 = 4 R = p2 useclus = True - big = b * R <= 148 + big = b * R <= num_sms scap = (16384 if R == 1 else 8192) if big else (8192 if k > 1024 else 4096) cmp_ = (4096 if k > 1024 else 2048) if big else 1024 aim = ( @@ -639,7 +692,7 @@ def route_streaming( "smem": smc, "ws": False, } - smem_main = (scap + 4) * (8 if (R > 1 or b <= 296) else 4) + (cmp_ + 1) * 8 + smem_main = (scap + 4) * (8 if (R > 1 or b <= 2 * num_sms) else 4) + (cmp_ + 1) * 8 def _main(blk_, minb_, u_, split_): kpt = 1 if k <= blk_ else (2 if k <= 2 * blk_ else (4 if k <= 4 * blk_ else 8)) @@ -671,12 +724,83 @@ def _main(blk_, minb_, u_, split_): per = q_ >> 10 u_ = 8 if per >= 8 else (4 if per >= 4 else (2 if per >= 2 else 1)) return _main(1024, 1, u_, R > 1) - if b <= 296: + if b <= 2 * num_sms: return _main(512, 2, 8, False) return _main(256, 4, 8, False) _VARLEN_CACHE = {} +_DEVICE_COMPUTE_INFO = {} + + +def _current_locality_domain() -> int | None: + """Return the thread-local locality domain without importing it eagerly.""" + from tensorrt_llm._torch.locality_domain_utils import get_current_locality_domain + + return get_current_locality_domain() + + +def _locality_domain_topology() -> tuple[tuple[int, int], ...]: + """Return the initialized public CUDA locality-domain compute split.""" + from tensorrt_llm._torch.locality_domain.runtime import LocalityDomainRuntime + + return LocalityDomainRuntime().topology_identity() + + +def _device_ordinal(device: torch.device | int) -> int: + """Normalize an explicit/current CUDA device to its integer ordinal.""" + if isinstance(device, int): + return device + cuda_device = torch.device(device) + if cuda_device.type != "cuda": + raise RuntimeError(f"expected a CUDA device, got {device}") + return cuda_device.index if cuda_device.index is not None else torch.cuda.current_device() + + +def _execution_domain(device: torch.device | int) -> tuple[int, int | None]: + """Return available SMs and the current locality-domain cache identity.""" + device_index = _device_ordinal(device) + compute_info = _DEVICE_COMPUTE_INFO.get(device_index) + if compute_info is None: + properties = torch.cuda.get_device_properties(device_index) + sm_version = int(properties.major) * 10 + int(properties.minor) + compute_info = (sm_version, int(properties.multi_processor_count)) + _DEVICE_COMPUTE_INFO[device_index] = compute_info + sm_version, full_device_num_sms = compute_info + + # B200/GB200 has no locality-domain execution. Avoid importing or + # querying that runtime on its hot path. + if sm_version != 107: + return full_device_num_sms, None + + locality_domain_id = _current_locality_domain() + if locality_domain_id is not None: + # The low-level topology cache is keyed by torch's current device. + # Query it under the tensor's device rather than the caller's ambient + # device, which may differ in multi-GPU serving processes. + with torch.cuda.device(device_index): + topology = _locality_domain_topology() + if not 0 <= locality_domain_id < len(topology): + raise RuntimeError(f"invalid current locality domain {locality_domain_id}") + partition_num_sms, total_num_sms = topology[locality_domain_id] + if not 0 < partition_num_sms <= total_num_sms: + raise RuntimeError( + "locality-domain compute topology is unavailable or invalid: " + f"domain={locality_domain_id}, counts={(partition_num_sms, total_num_sms)}" + ) + if total_num_sms != full_device_num_sms: + raise RuntimeError( + "locality-domain topology does not match the target device: " + f"device={device_index}, topology_total={total_num_sms}, " + f"device_total={full_device_num_sms}" + ) + return partition_num_sms, locality_domain_id + return full_device_num_sms, None + + +def _available_num_sms(device: torch.device | int) -> int: + """Return SMs available to launches in the current execution domain.""" + return _execution_domain(device)[0] def _varlen_launcher( @@ -686,12 +810,14 @@ def _varlen_launcher( n_env: int, next_n: int, cr: int, + num_sms: int = DEFAULT_NUM_SMS, + locality_domain_id: int | None = None, ) -> tuple: """Capture-time varlen plan + compiled launcher. The gvr_main port is the universally correct fallback; specialist family tiers below. Every choice here is a function of capture-stable quantities only — mirroring the in-tree runner's pick_tuning(graph_capture=...) discipline.""" - key = (num_rows, npad, k, n_env, next_n, cr) + key = (num_rows, npad, k, n_env, next_n, cr, num_sms, locality_domain_id) hit = _VARLEN_CACHE.get(key) if hit is not None: return hit @@ -703,7 +829,7 @@ def _varlen_launcher( # admission window (n4 <= 32768) fits capture-frozen envelopes. The # choice is a pure function of this cache key, so CUDA-graph replay # safety is unchanged; per-row n / short-row handling lives in-kernel. - plan_free = route(num_rows, n_eff, npad, k) + plan_free = route(num_rows, n_eff, npad, k, num_sms=num_sms) if plan_free["kernel"] == "reg_clus": fn = dev.get_compiled__regclus( tuple(plan_free["tpl"]), @@ -755,6 +881,7 @@ def _varlen_launcher( next_n=next_n, cr_shift=cr_shift, hint_free=True, + num_sms=num_sms, ) lc = ( "clus", @@ -763,7 +890,14 @@ def _varlen_launcher( ) _VARLEN_CACHE[key] = lc return lc - plan = route_streaming(num_rows, n_eff, npad, k, force_main=True) + plan = route_streaming( + num_rows, + n_eff, + npad, + k, + force_main=True, + num_sms=num_sms, + ) tpl = tuple(plan["tpl"]) # (BLK, U, MINB, SNB, KPT, SPLIT, TSHG) rt = plan["rt"] r_const = rt["R"] @@ -771,7 +905,7 @@ def _varlen_launcher( # machinery in whenever SPLIT); normalize it out of the compile key so # row counts differing only in that slot share one engine fn = dev.get_compiled(tpl[:6] + (False,) + (next_n, cr_shift, r_const), hint_free=True) - big = num_rows * r_const <= 148 + big = num_rows * r_const <= num_sms aim_base = ( ((4 * k if k >= 1024 else 2 * k) if r_const == 1 else 2 * k) if big @@ -797,7 +931,12 @@ def _varlen_launcher( def route_bands( - b: int, npad: int, k: int, n_lo: int | None = None, n_hi: int | None = None + b: int, + npad: int, + k: int, + n_lo: int | None = None, + n_hi: int | None = None, + num_sms: int = DEFAULT_NUM_SMS, ) -> list[tuple[int, int, dict[str, object]]]: """Enumerate maximal n-intervals on which route_static is constant. Dense O(n_hi - n_lo) scan of the pure host dispatch — an offline / @@ -808,7 +947,7 @@ def route_bands( bands = [] cur_key, cur_lo, cur_plan = None, lo, None for n in range(lo, hi + 1): - st = route_static(b, n, npad, k) + st = route_static(b, n, npad, k, num_sms=num_sms) key = repr(st) if key != cur_key: if cur_key is not None: @@ -822,11 +961,11 @@ def route_bands( # =========================================================================== # ==== workspace ============================================================ # =========================================================================== -"""Per-device workspace slab for the multi-CTA SPLIT path. +"""Per-execution-domain workspace slab for the multi-CTA SPLIT path. Semantics: - * ONE zero-initialised slab workspace per device, lazily allocated through - the torch caching allocator; + * ONE zero-initialised slab workspace per full device or locality domain, + lazily allocated through the matching torch caching allocator/mempool; * keep-alive store: module dict `_ws_keep` (tensor refcount = keep-alive); * double-checked locking: lock-free hot-path load (a GIL-atomic dict get plays an acquire load), slow path re-checks under a mutex; @@ -835,8 +974,9 @@ def route_bands( input checks, so a CPU logits tensor dies here with "device index out of range: -1"). -Concurrent STREAMS on one device that may both take the multi-CTA SPLIT path -must pass their own workspace via run_ws(). +Concurrent STREAMS on one device or in the same locality domain that may +both take the multi-CTA SPLIT path must pass distinct, zero-initialised +workspaces via run_ws() / run_varlen(workspace=...). Size: workspace_bytes() = GVR_WS_BUF_OFF + MAXC*GCAP*sizeof(int2) = 2048 + 160*16384*8 = 20,973,568 B. @@ -852,14 +992,30 @@ def route_bands( # workspace geometry constants -- must match the device kernels GVR_MAX_DEV = 64 -_MAXC = 160 +_MAXC = MAX_SPLIT_ROWS _GCAP = 16384 _GVR_WS_BUF_OFF = 2048 WS_BYTES = _GVR_WS_BUF_OFF + _MAXC * _GCAP * 8 # 20,973,568 assert WS_BYTES == 20_973_568 _mu = threading.Lock() # slow-path mutex -_ws_keep = {} # device index -> keep-alive int32 view +_ws_keep = {} # device or (device, locality domain) -> keep-alive int32 view + + +def _workspace_cache_key( + device_index: int, locality_domain_id: int | None +) -> int | tuple[int, int]: + """Return the workspace identity for the current execution domain.""" + if locality_domain_id is None: + return device_index + return device_index, locality_domain_id + + +def _optional_locality_domain_mem_pool(): + """Return the current locality-domain allocation context lazily.""" + from tensorrt_llm._torch.locality_domain_utils import optional_locality_domain_mem_pool + + return optional_locality_domain_mem_pool() def workspace_bytes() -> int: @@ -867,7 +1023,10 @@ def workspace_bytes() -> int: return WS_BYTES -def default_workspace(ref: torch.Tensor) -> torch.Tensor: +def _default_workspace( + ref: torch.Tensor, + locality_domain_id: int | None, +) -> torch.Tensor: """Per-device cached workspace slab. Returns the kernel-facing 1-D int32 view (zero-initialised on first use; @@ -876,21 +1035,32 @@ def default_workspace(ref: torch.Tensor) -> torch.Tensor: d = ref.get_device() if not (0 <= d < GVR_MAX_DEV): raise RuntimeError(f"device index out of range: {d}") - ws = _ws_keep.get(d) # hot path: one (GIL-atomic) load + workspace_key = _workspace_cache_key(d, locality_domain_id) + ws = _ws_keep.get(workspace_key) # hot path: one (GIL-atomic) load if ws is not None: return ws with _mu: # slow path: double-checked - ws = _ws_keep.get(d) + ws = _ws_keep.get(workspace_key) if ws is not None: return ws - # lazy zeros via the torch caching allocator, viewed int32 for the - # DSL launch signature. - buf = torch.zeros(WS_BYTES, dtype=torch.uint8, device=ref.device) + if locality_domain_id is None: + buf = torch.zeros(WS_BYTES, dtype=torch.uint8, device=ref.device) + else: + # Route first touch to the current domain's localized mempool. + with torch.cuda.device(d): + with _optional_locality_domain_mem_pool(): + buf = torch.zeros(WS_BYTES, dtype=torch.uint8, device=ref.device) ws = buf.view(torch.int32) - _ws_keep[d] = ws # keep-alive (ws_keep[d] = tensor) + _ws_keep[workspace_key] = ws return ws +def default_workspace(ref: torch.Tensor) -> torch.Tensor: + """Return the default workspace for the current execution domain.""" + _, locality_domain_id = _execution_domain(ref.get_device()) + return _default_workspace(ref, locality_domain_id) + + def validate_run_ws(workspace: torch.Tensor, logits: torch.Tensor) -> None: """run_ws() workspace hardening, in a fixed predicate order: CUDA + same device as logits; numel*element_size >= workspace_bytes(); @@ -930,10 +1100,40 @@ def kernel_view(workspace: torch.Tensor) -> torch.Tensor: return t +def _workspace_for_varlen_launch( + logits: torch.Tensor, + workspace: torch.Tensor | None, + locality_domain_id: int | None, +) -> torch.Tensor: + """Resolve a varlen workspace after its launcher is capture-ready. + + The caller must resolve the domain-specific launcher first. During CUDA + graph capture a cold default workspace cannot be allocated safely: an + aborted capture could otherwise publish a slab whose one-time zeroing + only existed in the discarded graph. + """ + if workspace is not None: + validate_run_ws(workspace, logits) + return kernel_view(workspace) + + device_index = logits.get_device() + workspace_key = _workspace_cache_key(device_index, locality_domain_id) + ws = _ws_hot.get(workspace_key) + if ws is not None: + return ws + if _is_capturing(): + raise RuntimeError( + "default workspace is not initialized for this execution domain " + "— warm up before CUDA graph capture" + ) + return _default_workspace(logits, locality_domain_id) + + def _reset_for_tests() -> None: """Drop cached slabs (tests only; NOT part of the C contract).""" with _mu: _ws_keep.clear() + _DEVICE_COMPUTE_INFO.clear() # =========================================================================== @@ -1004,8 +1204,8 @@ def _dummy_kv(dev_index, device): # --------------------------------------------------------------------------- # per-family launcher builders (cold path: once per distinct shape key) # --------------------------------------------------------------------------- -def _build_launcher(b, n, npad, k): - rd = route(b, n, npad, k) +def _build_launcher(b, n, npad, k, num_sms=DEFAULT_NUM_SMS): + rd = route(b, n, npad, k, num_sms=num_sms) fam = rd["kernel"] tpl = tuple(rd["tpl"]) rt = rd["rt"] @@ -1053,7 +1253,12 @@ def fn(lg, pi, o, w, *a, _raw=raw): # ABI: (logits, pre_idx, kv_lens, out, n, npad, k, SCAP, CMP, SMP, # TGT, Q, SS2, TGT2) -- NO workspace; kv_lens is the dead # varlen slot in batch-uniform mode (cached dummy tensor) - fn = dev.get_compiled__clus(tpl, scap=rt["SCAP"], cmp_=rt["CMP"]) + fn = dev.get_compiled__clus( + tpl, + scap=rt["SCAP"], + cmp_=rt["CMP"], + num_sms=num_sms, + ) args = ( rt["n"], rt["npad"], @@ -1090,7 +1295,16 @@ def _call(lg, pi, idx, _fn=fn, _n=n_arg): # --------------------------------------------------------------------------- # shared implementation of the batch-uniform entries # --------------------------------------------------------------------------- -def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): +def _run_impl( + logits, + pre_idx, + n_valid, + indices, + ws, + values=None, + num_sms=DEFAULT_NUM_SMS, + locality_domain_id=None, +): if not (logits.is_cuda and pre_idx.is_cuda and indices.is_cuda): raise RuntimeError("all tensors must be CUDA") if logits.dtype is not _F32: @@ -1183,10 +1397,10 @@ def _run_impl(logits, pre_idx, n_valid, indices, ws, values=None): values[:, n:] = torch.finfo(_F32).min # -FLT_MAX pad return - key = (b, n, npad, k) + key = (b, n, npad, k, num_sms, locality_domain_id) lc = _LAUNCH_CACHE.get(key) if lc is None: - lc = _build_launcher(b, n, npad, k) + lc = _build_launcher(b, n, npad, k, num_sms=num_sms) _LAUNCH_CACHE[key] = lc fn, args, needs_ws = lc try: @@ -1227,10 +1441,21 @@ def run( d = logits.get_device() if not 0 <= d < _GVR_MAX_DEV: # checked on EVERY call raise RuntimeError(f"device index out of range: {d}") - ws = _ws_hot.get(d) + num_sms, locality_domain_id = _execution_domain(d) + workspace_key = _workspace_cache_key(d, locality_domain_id) + ws = _ws_hot.get(workspace_key) if ws is None: - ws = default_workspace(logits) - _run_impl(logits, pre_idx, n_valid, indices, ws, values) + ws = _default_workspace(logits, locality_domain_id) + _run_impl( + logits, + pre_idx, + n_valid, + indices, + ws, + values, + num_sms=num_sms, + locality_domain_id=locality_domain_id, + ) def run_ws( @@ -1245,7 +1470,17 @@ def run_ws( Explicit-workspace form for multi-stream callers.""" validate_run_ws(workspace, logits) - _run_impl(logits, pre_idx, n_valid, indices, kernel_view(workspace), values) + num_sms, locality_domain_id = _execution_domain(logits.get_device()) + _run_impl( + logits, + pre_idx, + n_valid, + indices, + kernel_view(workspace), + values, + num_sms=num_sms, + locality_domain_id=locality_domain_id, + ) def run_varlen( @@ -1287,6 +1522,10 @@ def run_varlen( implementation-specifically). Finite inputs — including +/-inf and denormals — are tie-aware exact. + ``workspace``, when supplied, must be zero-initialised before its first + launch and owned by this in-flight invocation; concurrent launches in the + same locality domain must not share it. + CONTRACT: correct and dispatched for any ``num_rows`` (BS 1..1024+ x next_n) and any envelope up to 1M kv tokens. Family selection (streaming main / clustered register-resident) is a pure @@ -1322,15 +1561,7 @@ def run_varlen( d = logits.get_device() if not 0 <= d < _GVR_MAX_DEV: raise RuntimeError(f"device index out of range: {d}") - if workspace is not None: - # multi-stream escape hatch (run_ws parity): concurrent varlen - # launches on one device must not share the SPLIT publish slab - validate_run_ws(workspace, logits) - ws = kernel_view(workspace) - else: - ws = _ws_hot.get(d) - if ws is None: - ws = default_workspace(logits) + num_sms, locality_domain_id = _execution_domain(d) # ---- per-row in-kernel engine (gvr_main varlen port) ---------------- # Full validation battery (the engine bypasses _run_impl — every @@ -1395,14 +1626,27 @@ def run_varlen( # R increment (bounded plans, bounded _VARLEN_CACHE) n_env = 1 << max(n_env - 1, 1).bit_length() n_env = min(max(n_env, 1), npad) - key = (num_rows, npad, k, n_env, nn, cr) + key = (num_rows, npad, k, n_env, nn, cr, num_sms, locality_domain_id) lc = _VARLEN_CACHE.get(key) if lc is None: if _is_capturing(): raise RuntimeError( "varlen launcher not compiled for this shape — warm up before CUDA graph capture" ) - lc = _varlen_launcher(num_rows, npad, k, n_env, nn, cr) + lc = _varlen_launcher( + num_rows, + npad, + k, + n_env, + nn, + cr, + num_sms=num_sms, + locality_domain_id=locality_domain_id, + ) + # Resolve the launcher before touching a cold domain workspace. If a + # capture reaches this point, both the launcher and the default slab + # must already have been warmed in this exact execution domain. + ws = _workspace_for_varlen_launch(logits, workspace, locality_domain_id) idx = indices if idx.shape[1] != k: idx = idx.reshape(-1)[: num_rows * k].view(num_rows, k) @@ -1484,6 +1728,7 @@ def warmup_varlen( """ dev = torch.cuda.current_device() + num_sms, locality_domain_id = _execution_domain(dev) nn = max(1, int(next_n)) # round each request down to a next_n multiple (min next_n) and dedup req_rows = sorted({max(int(r) - int(r) % nn, nn) for r in num_rows_list}) @@ -1502,7 +1747,13 @@ def warmup_varlen( r = nn r_max = req_rows[-1] while r <= r_max: - plan_free = route(r, max(min(n_env_c, npad_c), int(top_k) + 1), npad_c, int(top_k)) + plan_free = route( + r, + max(min(n_env_c, npad_c), int(top_k) + 1), + npad_c, + int(top_k), + num_sms=num_sms, + ) if plan_free["kernel"] == "reg_clus": ekey = ("reg_clus", tuple(plan_free["tpl"])) elif plan_free["kernel"] in ("reg", "regimg"): @@ -1514,6 +1765,7 @@ def warmup_varlen( npad_c, int(top_k), force_main=True, + num_sms=num_sms, ) ekey = ("main", tuple(p["tpl"][:6]), p["rt"]["R"]) if ekey not in seen_keys: @@ -1537,8 +1789,11 @@ def warmup_varlen( int(max_seq_len), int(compress_ratio), nn, + tuple(req_rows), tuple(rows_list), npad, + num_sms, + locality_domain_id, ) with _VARLEN_WARMUP_LOCK: if key in _VARLEN_WARMUP_DONE: @@ -1567,6 +1822,15 @@ def warmup_varlen( # CUDA-graph capture at any requested geometry finds its key immediately. n_env_l = min(max(int(max_seq_len) >> (0 if int(compress_ratio) == 1 else 2), 1), npad) for r in req_rows: - _varlen_launcher(r, npad, int(top_k), n_env_l, nn, int(compress_ratio)) + _varlen_launcher( + r, + npad, + int(top_k), + n_env_l, + nn, + int(compress_ratio), + num_sms=num_sms, + locality_domain_id=locality_domain_id, + ) with _VARLEN_WARMUP_LOCK: _VARLEN_WARMUP_DONE.add(key) 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 2eca4ab764e6..95800a39ad25 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -57,7 +57,10 @@ 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.sparse.dsa.params import ( + is_gvr_cute_dsl_supported, + 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 @@ -181,6 +184,8 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(use_self_samp (dict(enable_heuristic_topk=False), False), (dict(use_self_sampling_topk=False), False), (dict(is_cute_dsl_available=False), False), + (dict(sm_version=107), False), + (dict(sm_version=107, is_cute_dsl_rubin_available=True), True), (dict(sm_version=120), False), (dict(index_topk=256), False), (dict(compress_ratio=2), False), @@ -193,12 +198,24 @@ def test_use_self_sampling_gvr(kwargs, expected): index_topk=512, compress_ratio=1, is_cute_dsl_available=True, + is_cute_dsl_rubin_available=False, sm_version=100, ) base.update(kwargs) assert use_self_sampling_gvr(**base) is expected +def test_sm107_gvr_admission_is_self_sampling_only(): + """Temporal GVR remains restricted to its validated SM100/103 path.""" + common = dict( + is_cute_dsl_available=True, + is_cute_dsl_rubin_available=True, + sm_version=107, + ) + assert is_gvr_cute_dsl_supported(**common, use_self_sampling_topk=True) + assert not is_gvr_cute_dsl_supported(**common, use_self_sampling_topk=False) + + @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): @@ -603,6 +620,38 @@ def test_indexer_two_level_gvr_dispatch( assert indexer.top_k.needs_gvr_prior == (not use_self_sampling) +@skip_pre_hopper +def test_indexer_sm107_temporal_gvr_falls_back_to_radix(): + """SM107 support is V2-only; requesting temporal V1 must not select GVR.""" + sparse_config = DeepSeekSparseAttentionConfig( + index_head_dim=128, + index_n_heads=32, + index_topk=512, + use_cute_dsl_topk=False, + enable_heuristic_topk=True, + use_self_sampling_topk=False, + ) + with ( + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.IS_CUTLASS_DSL_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.IS_CUTLASS_DSL_RUBIN_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.get_sm_version", + return_value=107, + ), + ): + indexer = create_indexer(sparse_config) + + assert indexer.top_k.decode_implementation == TopKImplementation.CUDA_RADIX + assert not indexer.top_k.gvr_self_sampling + assert not indexer.top_k.needs_gvr_prior + + @skip_pre_hopper @pytest.mark.parametrize( "flag_value,expected_dtype", diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 3408dd325fd6..3decf9b74303 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -32,7 +32,10 @@ from utils.util import getSMVersion import tensorrt_llm # noqa: F401 -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE, +) if not torch.cuda.is_available(): pytest.skip("CUDA is required for gvr_selfsampling_topk tests", allow_module_level=True) @@ -40,21 +43,35 @@ if not IS_CUTLASS_DSL_AVAILABLE: pytest.skip("cutlass DSL is required for gvr_selfsampling_topk tests", allow_module_level=True) -if getSMVersion() not in (100, 103): +_SM_VERSION = getSMVersion() +if _SM_VERSION not in (100, 103) and not (_SM_VERSION == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE): pytest.skip( - "self-sampling GVR kernels target datacenter Blackwell (sm_100/103) " - "— same gate as the production dispatch; consumer Blackwell " - "(sm_120/121) lacks thread-block clusters", + "self-sampling GVR kernels require SM100/103 or SM107 with Rubin " + "CuTe DSL helpers — same gate as the production dispatch", allow_module_level=True, ) -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( # noqa: E402 gvr_topk_decode_self_sampling_host as ss_host, ) _DEV = "cuda" +def _varlen_cache_key(rows, npad, top_k, n_env, next_n, compress_ratio): + num_sms, locality_domain_id = ss_host._execution_domain(torch.device(_DEV)) + return ( + rows, + npad, + top_k, + n_env, + next_n, + compress_ratio, + num_sms, + locality_domain_id, + ) + + def _make_case(batch_size, n_valid, top_k, seed, hit_ratio=0.6): """Decode-like fp32 logits + prev-step hint. The padded tail is poisoned with +3e38 so any read past n_valid corrupts the top-K values.""" @@ -705,7 +722,7 @@ def test_selfsampling_varlen_regclus_parity_and_oracle(): kv = torch.tensor([msl_c * cr, 900, nn - 1], dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) - key = (rows, npad, k, msl_c, nn, cr) + key = _varlen_cache_key(rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg_clus", ss_host._VARLEN_CACHE[key][0] torch.cuda.synchronize() ref = _reference_varlen_indices(lg, kv, nn, cr, k) @@ -770,7 +787,7 @@ def test_selfsampling_varlen_reg_parity_and_oracle(): kv = torch.tensor([msl, max((k - 3) * cr, nn), nn - 1], dtype=torch.int32, device=_DEV) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl) - key = (rows, npad, k, msl_c, nn, cr) + key = _varlen_cache_key(rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] torch.cuda.synchronize() ref = _reference_varlen_indices(lg, kv, nn, cr, k) @@ -810,7 +827,7 @@ def test_selfsampling_varlen_clus_parity_and_oracle(): ) out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, kv, out, next_n=nn, compress_ratio=cr, max_seq_len=msl_c * cr) - key = (rows, npad, k, msl_c, nn, cr) + key = _varlen_cache_key(rows, npad, k, msl_c, nn, cr) assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] torch.cuda.synchronize() ref = _reference_varlen_indices(lg, kv, nn, cr, k) @@ -837,7 +854,7 @@ def test_selfsampling_varlen_clus_cuda_graph(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr) + key = _varlen_cache_key(rows, msl_c, k, msl_c, 1, cr) assert ss_host._VARLEN_CACHE[key][0] == "clus", ss_host._VARLEN_CACHE[key][0] g = torch.cuda.CUDAGraph() out.fill_(-7) @@ -863,7 +880,7 @@ def test_selfsampling_varlen_reg_cuda_graph(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr) + key = _varlen_cache_key(rows, msl_c, k, msl_c, 1, cr) assert ss_host._VARLEN_CACHE[key][0] == "reg", ss_host._VARLEN_CACHE[key][0] g = torch.cuda.CUDAGraph() out.fill_(-7) @@ -891,7 +908,7 @@ def test_selfsampling_varlen_full_row_range(): out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) ss_host.run_varlen(lg, kv, out, next_n=1, compress_ratio=cr, max_seq_len=msl_c * cr) torch.cuda.synchronize() - key = (rows, msl_c, k, msl_c, 1, cr) + key = _varlen_cache_key(rows, msl_c, k, msl_c, 1, cr) assert key in ss_host._VARLEN_CACHE, "row count must dispatch in-engine" ref_v = torch.topk(lg.float(), k, dim=1).values.sort(dim=1).values got = lg.float().gather(1, out.long().clamp_min(0)).sort(dim=1).values diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk_host.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk_host.py new file mode 100644 index 000000000000..0bbec4ad1d71 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk_host.py @@ -0,0 +1,328 @@ +# 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. + +"""CPU-only tests for self-sampling GVR host routing and locality caches.""" + +import importlib.util +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_HOST_PATH = ( + Path(__file__).parents[5] + / "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py" +) +_SPEC = importlib.util.spec_from_file_location("gvr_self_sampling_host_cpu_test", _HOST_PATH) +assert _SPEC is not None and _SPEC.loader is not None +ss_host = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(ss_host) + + +def test_default_route_preserves_b200_plans(): + """The new topology argument must preserve every default B200 decision.""" + for batch_size in (1, 8, 16, 64, 148, 296, 1024): + for n_valid in (4096, 65536, 131072, 262144): + for top_k in (512, 1024, 2048): + default = ss_host.route(batch_size, n_valid, n_valid, top_k) + explicit = ss_host.route( + batch_size, + n_valid, + n_valid, + top_k, + num_sms=148, + ) + assert default == explicit + assert ( + ss_host.route_split( + batch_size, + n_valid, + n_valid, + top_k, + num_sms=148, + ) + == explicit + ) + + +def test_rubin_full_device_route_uses_available_sms(): + """Single-row splits may use all SMs; MAXC limits rows, not per-row CTAs.""" + b200 = ss_host.route(1, 1 << 20, 1 << 20, 1024) + rubin_212 = ss_host.route(1, 1 << 20, 1 << 20, 1024, num_sms=212) + rubin_224 = ss_host.route(1, 1 << 20, 1 << 20, 1024, num_sms=224) + assert b200["rt"]["R"] == 148 + assert rubin_212["kernel"] == rubin_224["kernel"] == "main" + assert rubin_212["rt"]["R"] == rubin_212["grid"][0] == 212 + assert rubin_224["rt"]["R"] == rubin_224["grid"][0] == 224 + + # The 74-row B200 half-wave threshold scales to 106 rows on R200. + assert ss_host.route(80, 262144, 262144, 1024, num_sms=148)["kernel"] == "main" + assert ss_host.route(80, 262144, 262144, 1024, num_sms=212)["kernel"] == "clus" + + # A synthetic topology above 320 SMs must not select SPLIT for more + # than the workspace's 160 row slabs. + assert ss_host.route(160, 262144, 262144, 1024, num_sms=400)["grid"][0] == 2 + assert ss_host.route(161, 262144, 262144, 1024, num_sms=400)["grid"][0] == 1 + + +def test_route_dynamic_requires_matching_execution_topology(): + """A non-default static plan must not silently use the B200 SM count.""" + args = (80, 8192, 8192, 512) + static = ss_host.route_static(*args, num_sms=64) + dynamic, smem = ss_host.route_dynamic(static, args[1], num_sms=64) + recombined = { + key: (dict(value) if isinstance(value, dict) else value) for key, value in static.items() + } + recombined["rt"].update(dynamic) + recombined["smem"] = smem + assert recombined == ss_host.route(*args, num_sms=64) + with pytest.raises(TypeError, match="num_sms"): + ss_host.route_dynamic(static, args[1]) + + +def test_execution_domain_uses_full_and_partition_sm_counts(monkeypatch): + """Full-device properties are cached; a current domain uses its partition.""" + calls = [] + device_contexts = [] + + @contextmanager + def fake_device(device): + device_contexts.append(device) + yield + + monkeypatch.setattr(ss_host, "_current_locality_domain", lambda: None) + monkeypatch.setattr(ss_host.torch.cuda, "device", fake_device) + monkeypatch.setattr( + ss_host.torch.cuda, + "get_device_properties", + lambda device: calls.append(device) + or SimpleNamespace(major=10, minor=7, multi_processor_count=212), + ) + ss_host._DEVICE_COMPUTE_INFO.clear() + assert ss_host._execution_domain(0) == (212, None) + assert ss_host._execution_domain(ss_host.torch.device("cuda:0")) == (212, None) + assert calls == [0] + + monkeypatch.setattr(ss_host, "_current_locality_domain", lambda: 1) + monkeypatch.setattr(ss_host, "_locality_domain_topology", lambda: ((104, 212), (108, 212))) + assert ss_host._execution_domain(0) == (108, 1) + assert device_contexts == [0] + + monkeypatch.setattr(ss_host, "_locality_domain_topology", lambda: ((104, 224), (108, 224))) + with pytest.raises(RuntimeError, match="topology.*target device"): + ss_host._execution_domain(ss_host.torch.device("cuda:0")) + + +def test_b200_execution_domain_does_not_query_locality_runtime(monkeypatch): + """B200 keeps the pre-Rubin dependency and hot-path behavior.""" + monkeypatch.setattr( + ss_host.torch.cuda, + "get_device_properties", + lambda device: SimpleNamespace(major=10, minor=0, multi_processor_count=148), + ) + monkeypatch.setattr( + ss_host, + "_current_locality_domain", + lambda: (_ for _ in ()).throw(AssertionError("locality runtime queried on B200")), + ) + ss_host._DEVICE_COMPUTE_INFO.clear() + assert ss_host._execution_domain(0) == (148, None) + + +class _FakeDeviceModule: + STATIC_BYTES = 0 + + def __init__(self): + self.calls = [] + + def _compiled(self, family, *args, **kwargs): + marker = object() + self.calls.append((family, args, kwargs, marker)) + return marker + + def get_compiled(self, *args, **kwargs): + return self._compiled("main", *args, **kwargs) + + def get_compiled__reg(self, *args, **kwargs): + return self._compiled("reg", *args, **kwargs) + + def get_compiled__regclus(self, *args, **kwargs): + return self._compiled("reg_clus", *args, **kwargs) + + def get_compiled__clus(self, *args, **kwargs): + return self._compiled("clus", *args, **kwargs) + + +def test_varlen_cache_separates_full_device_and_locality_domain(monkeypatch): + """Equal SM counts in different execution domains must not alias launchers.""" + fake_device = _FakeDeviceModule() + monkeypatch.setattr(ss_host, "_device", lambda: fake_device) + ss_host._VARLEN_CACHE.clear() + args = (64, 262144, 1024, 262144, 1, 4) + full = ss_host._varlen_launcher(*args, num_sms=212, locality_domain_id=None) + local = ss_host._varlen_launcher(*args, num_sms=212, locality_domain_id=0) + assert full is not local + assert (*args, 212, None) in ss_host._VARLEN_CACHE + assert (*args, 212, 0) in ss_host._VARLEN_CACHE + clus_calls = [call for call in fake_device.calls if call[0] == "clus"] + assert clus_calls + assert all(call[2]["num_sms"] == 212 for call in clus_calls) + + +def test_default_workspace_is_locality_domain_scoped(monkeypatch): + """Concurrent locality streams receive distinct, domain-local slabs.""" + allocations = [] + pool_entries = [] + device_contexts = [] + + class FakeBuffer: + def __init__(self, allocation_id): + self.allocation_id = allocation_id + + def view(self, dtype): + return self, dtype + + @contextmanager + def fake_pool(): + pool_entries.append("enter") + yield + + @contextmanager + def fake_device(device): + device_contexts.append(device) + yield + + ref = SimpleNamespace(get_device=lambda: 0, device="cuda:0") + monkeypatch.setattr(ss_host, "_optional_locality_domain_mem_pool", fake_pool) + monkeypatch.setattr(ss_host.torch.cuda, "device", fake_device) + monkeypatch.setattr( + ss_host.torch, + "zeros", + lambda *args, **kwargs: allocations.append((args, kwargs)) or FakeBuffer(len(allocations)), + ) + ss_host._ws_keep.clear() + full = ss_host._default_workspace(ref, None) + domain0 = ss_host._default_workspace(ref, 0) + domain1 = ss_host._default_workspace(ref, 1) + assert ss_host._default_workspace(ref, 0) is domain0 + assert len({id(full), id(domain0), id(domain1)}) == 3 + assert set(ss_host._ws_keep) == {0, (0, 0), (0, 1)} + assert len(allocations) == 3 + assert len(pool_entries) == 2 + assert device_contexts == [0, 0] + + +def test_varlen_capture_workspace_miss_does_not_allocate(monkeypatch): + """A cold domain slab is never allocated or published inside capture.""" + allocations = [] + ref = SimpleNamespace(get_device=lambda: 0) + monkeypatch.setattr(ss_host, "_is_capturing", lambda: True) + monkeypatch.setattr( + ss_host, + "_default_workspace", + lambda *args: allocations.append(args), + ) + ss_host._ws_keep.clear() + + with pytest.raises(RuntimeError, match="workspace.*warm up"): + ss_host._workspace_for_varlen_launch(ref, None, 1) + assert allocations == [] + assert ss_host._ws_keep == {} + + +def test_varlen_capture_checks_launcher_before_workspace(monkeypatch): + """A launcher miss aborts capture before workspace resolution begins.""" + + class FakeTensor: + is_cuda = True + + def __init__(self, shape, dtype): + self.shape = shape + self.dtype = dtype + + def dim(self): + return len(self.shape) + + def get_device(self): + return 0 + + def stride(self, dim): + return self.shape[1] if dim == 0 else 1 + + def is_contiguous(self): + return True + + def data_ptr(self): + return 16 + + logits = FakeTensor((1, 4096), ss_host.torch.float32) + kv_lens = FakeTensor((1,), ss_host.torch.int32) + indices = FakeTensor((1, 512), ss_host.torch.int32) + workspace_calls = [] + monkeypatch.setattr(ss_host, "_TENSOR", FakeTensor) + monkeypatch.setattr(ss_host, "_execution_domain", lambda device: (108, 1)) + monkeypatch.setattr(ss_host, "_is_capturing", lambda: True) + monkeypatch.setattr( + ss_host, + "_workspace_for_varlen_launch", + lambda *args: workspace_calls.append(args), + ) + ss_host._VARLEN_CACHE.clear() + + with pytest.raises(RuntimeError, match="launcher.*warm up"): + ss_host.run_varlen(logits, kv_lens, indices, max_seq_len=4096) + assert workspace_calls == [] + + +def test_warmup_done_key_includes_exact_requested_rows(monkeypatch): + """Equal band representatives must not hide a new capture row count.""" + + class FakeTensor: + def __getitem__(self, key): + return self + + exact_launchers = [] + monkeypatch.setattr(ss_host.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(ss_host.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(ss_host, "_execution_domain", lambda device: (212, 0)) + monkeypatch.setattr(ss_host.torch, "zeros", lambda *args, **kwargs: FakeTensor()) + monkeypatch.setattr(ss_host.torch, "full", lambda *args, **kwargs: FakeTensor()) + monkeypatch.setattr(ss_host.torch, "empty", lambda *args, **kwargs: FakeTensor()) + monkeypatch.setattr(ss_host, "run_varlen", lambda *args, **kwargs: None) + monkeypatch.setattr( + ss_host, + "_varlen_launcher", + lambda rows, *args, **kwargs: exact_launchers.append(rows), + ) + ss_host._VARLEN_WARMUP_DONE.clear() + + common = dict( + top_k=1024, + max_seq_len=262144, + compress_ratio=4, + row_stride=65536, + ) + ss_host.warmup_varlen(**common, num_rows_list=(32, 128)) + first_call_count = len(exact_launchers) + ss_host.warmup_varlen(**common, num_rows_list=(64, 128)) + + assert first_call_count > 0 + assert len(exact_launchers) > first_call_count + assert exact_launchers[-2:] == [64, 128] + + +def test_route_rejects_invalid_sm_count(): + with pytest.raises(RuntimeError, match="num_sms >= 1"): + ss_host.route(1, 4096, 4096, 512, num_sms=0) From 3c8be431c80d2b9949732223f5799d5238a0752b Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:26:09 +0000 Subject: [PATCH 07/10] [None][feat] add Rubin uGPU row sharding to GVR V2 decode Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 1 + .../attention_backend/sparse/dsa/metadata.py | 40 ++ .../attention_backend/sparse/dsa/params.py | 5 + .../_torch/locality_domain/gvr_topk.py | 177 +++++++ .../_torch/locality_domain/runtime.py | 7 + tensorrt_llm/_torch/locality_domain_utils.py | 19 +- tensorrt_llm/_torch/model_config.py | 7 + tensorrt_llm/_torch/modules/top_k.py | 256 +++++++++- tensorrt_llm/llmapi/llm_args.py | 21 + .../usage/llm_args_golden_manifest.json | 7 + .../attention/sparse/dsa/test_dsa_indexer.py | 14 + .../modeling/test_modeling_deepseekv4.py | 2 + .../_torch/modules/test_gvr_topk_locality.py | 476 ++++++++++++++++++ tests/unittest/_torch/test_model_config.py | 43 ++ .../parallel/test_locality_domain_utils.py | 48 +- 15 files changed, 1106 insertions(+), 17 deletions(-) create mode 100644 tensorrt_llm/_torch/locality_domain/gvr_topk.py create mode 100644 tests/unittest/_torch/modules/test_gvr_topk_locality.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 56ac5459712b..6b5a158336c5 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -765,6 +765,7 @@ def __init__( decode_implementation=decode_top_k_implementation, compress_ratio=self.compress_ratio, gvr_self_sampling=self._use_self_sampling_topk, + use_gvr_locality_domain=sparse_params.use_gvr_locality_domain, ) # Emission block-skip is a temporal-hint (V1) optimization: the FP4 # indexer epilogue emits per-block max logits the GVR Top-K consumes to diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 16af5f4c26a1..cebfbafc0571 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -196,6 +196,7 @@ def __post_init__(self): self.enable_gvr_topk = ( sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 ) + self.use_gvr_locality_domain = sparse_metadata_params.use_gvr_locality_domain 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 @@ -472,6 +473,45 @@ def warmup_selfsampling_topk( num_rows_list=tuple(sorted(rows)), row_stride=row_stride, ) + if self.use_gvr_locality_domain and IS_CUTLASS_DSL_RUBIN_AVAILABLE: + from ....locality_domain.gvr_topk import plan_gvr_topk_row_shards + from ....locality_domain.runtime import LocalityDomainRuntime + from ....locality_domain_utils import is_locality_domain_enabled + + if is_locality_domain_enabled(torch.cuda.current_device()): + runtime = LocalityDomainRuntime(num_partitions=2) + topology = runtime.topology_identity() + shard_rows: list[set[int]] = [set(), set()] + for num_rows in rows: + plan = plan_gvr_topk_row_shards( + num_rows=num_rows, + next_n=nn, + score_width=msl_c, + top_k=int(top_k), + topology=topology, + ) + if plan is None: + continue + if msl_c % 4 and any(shard.num_rows == 1 for shard in plan.shards): + continue + for shard in plan.shards: + shard_rows[shard.partition_id].add(shard.num_rows) + + if all(shard_rows): + runtime.fork() + try: + for partition_id, partition_rows in enumerate(shard_rows): + with runtime.partition_context(partition_id): + _ss_host.warmup_varlen( + int(top_k), + msl_c * cr, + compress_ratio=cr, + next_n=nn, + num_rows_list=tuple(sorted(partition_rows)), + row_stride=row_stride, + ) + finally: + runtime.join() except torch.cuda.OutOfMemoryError: # warmup is best-effort: the dispatch works without it (engines # JIT lazily outside capture), so do not fail engine init diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py index b16b530661d0..0b214b54a437 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py @@ -87,6 +87,7 @@ class DSAMetadataParams(SparseMetadataParams): mtp_index_share: bool = False use_self_sampling_topk: bool = True use_gvr_emission: bool = False + use_gvr_locality_domain: bool = False @dataclass(frozen=True) @@ -111,6 +112,10 @@ class DSAParams(SparseParams): # 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 + # Prototype Rubin-only row sharding for self-sampling GVR V2. The logits + # producer remains full-device; only non-overlapping Top-K row slices are + # submitted to the two locality-domain streams. + use_gvr_locality_domain: 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/locality_domain/gvr_topk.py b/tensorrt_llm/_torch/locality_domain/gvr_topk.py new file mode 100644 index 000000000000..3fc3a37cd514 --- /dev/null +++ b/tensorrt_llm/_torch/locality_domain/gvr_topk.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Pure row-sharding policy for Rubin locality-domain GVR Top-K.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +# These are conservative prototype thresholds, not Rubin architectural +# constants. They only reject workloads where two extra launches plus the +# fork/join events are clearly unlikely to amortize. Retune them from an R200 +# BS/ISL sweep before enabling locality-domain GVR by default. +GVR_LOCALITY_MIN_TOTAL_SCORE_ELEMENTS = 1 << 20 +GVR_LOCALITY_MIN_SCORE_ELEMENTS_PER_SM = 1 << 12 + + +@dataclass(frozen=True, slots=True) +class GvrTopKRowShard: + """One contiguous request-aligned row slice assigned to a locality domain.""" + + partition_id: int + request_start: int + request_end: int + row_start: int + row_end: int + num_sms: int + + @property + def num_requests(self) -> int: + return self.request_end - self.request_start + + @property + def num_rows(self) -> int: + return self.row_end - self.row_start + + +@dataclass(frozen=True, slots=True) +class GvrTopKRowShardPlan: + """A two-domain decode plan whose row slices are disjoint and exhaustive.""" + + shards: tuple[GvrTopKRowShard, GvrTopKRowShard] + topology: tuple[tuple[int, int], tuple[int, int]] + + +def _validate_topology( + topology: Sequence[tuple[int, int]], +) -> tuple[tuple[int, int], tuple[int, int]]: + if len(topology) != 2: + raise ValueError(f"GVR locality sharding requires 2 partitions, got {len(topology)}") + + normalized = tuple( + (int(partition_sms), int(total_sms)) for partition_sms, total_sms in topology + ) + if any(partition_sms <= 0 for partition_sms, _ in normalized): + raise ValueError(f"partition SM counts must be positive, got {normalized}") + if any(total_sms <= 0 or partition_sms > total_sms for partition_sms, total_sms in normalized): + raise ValueError(f"invalid locality-domain topology {normalized}") + if normalized[0][1] != normalized[1][1]: + raise ValueError(f"locality domains disagree on the full-device SM count: {normalized}") + if normalized[0][0] + normalized[1][0] > normalized[0][1]: + raise ValueError( + f"locality-domain partitions overlap the full-device topology: {normalized}" + ) + return normalized[0], normalized[1] + + +def is_gvr_topk_locality_workload_large_enough( + *, + num_rows: int, + next_n: int, + score_width: int, + top_k: int, + min_total_score_elements: int = GVR_LOCALITY_MIN_TOTAL_SCORE_ELEMENTS, +) -> bool: + """Apply the topology-independent part of the provisional gain gate.""" + return ( + next_n > 0 + and num_rows >= 2 * next_n + and num_rows % next_n == 0 + and score_width > top_k + and num_rows * score_width >= min_total_score_elements + ) + + +def plan_gvr_topk_row_shards( + *, + num_rows: int, + next_n: int, + score_width: int, + top_k: int, + topology: Sequence[tuple[int, int]], + min_total_score_elements: int = GVR_LOCALITY_MIN_TOTAL_SCORE_ELEMENTS, + min_score_elements_per_sm: int = GVR_LOCALITY_MIN_SCORE_ELEMENTS_PER_SM, +) -> GvrTopKRowShardPlan | None: + """Plan two proportional, request-aligned decode row slices. + + ``num_rows`` contains ``next_n`` consecutive rows for every request. The + split is therefore made in request space and then converted back to rows; + this preserves the leaf GVR mapping ``request = local_row // next_n``. + + The workload checks are deliberately capture-stable: they use tensor + geometry and the engine score-width envelope, never device KV lengths. + They are a provisional overhead guard rather than a claim of measured + Rubin speedup. + """ + if num_rows < 0: + raise ValueError(f"num_rows must be non-negative, got {num_rows}") + if next_n < 1: + raise ValueError(f"next_n must be positive, got {next_n}") + if num_rows % next_n: + raise ValueError(f"num_rows {num_rows} must be divisible by next_n {next_n}") + if score_width < 0: + raise ValueError(f"score_width must be non-negative, got {score_width}") + if top_k < 1: + raise ValueError(f"top_k must be positive, got {top_k}") + if min_total_score_elements < 0 or min_score_elements_per_sm < 0: + raise ValueError("GVR locality workload thresholds must be non-negative") + + normalized_topology = _validate_topology(topology) + if not is_gvr_topk_locality_workload_large_enough( + num_rows=num_rows, + next_n=next_n, + score_width=score_width, + top_k=top_k, + min_total_score_elements=min_total_score_elements, + ): + return None + num_requests = num_rows // next_n + + sms_0 = normalized_topology[0][0] + sms_1 = normalized_topology[1][0] + partition_sms = sms_0 + sms_1 + # Nearest-integer proportional split, clamped so both domains receive at + # least one complete request. This also handles asymmetric public splits. + split_request = (num_requests * sms_0 + partition_sms // 2) // partition_sms + split_request = min(max(split_request, 1), num_requests - 1) + split_row = split_request * next_n + + shards = ( + GvrTopKRowShard(0, 0, split_request, 0, split_row, sms_0), + GvrTopKRowShard( + 1, + split_request, + num_requests, + split_row, + num_rows, + sms_1, + ), + ) + if any( + shard.num_rows * score_width < min_score_elements_per_sm * shard.num_sms for shard in shards + ): + return None + return GvrTopKRowShardPlan(shards=shards, topology=normalized_topology) + + +__all__ = [ + "GVR_LOCALITY_MIN_SCORE_ELEMENTS_PER_SM", + "GVR_LOCALITY_MIN_TOTAL_SCORE_ELEMENTS", + "GvrTopKRowShard", + "GvrTopKRowShardPlan", + "is_gvr_topk_locality_workload_large_enough", + "plan_gvr_topk_row_shards", +] diff --git a/tensorrt_llm/_torch/locality_domain/runtime.py b/tensorrt_llm/_torch/locality_domain/runtime.py index 7ab54df17f70..fd84a120fc2d 100644 --- a/tensorrt_llm/_torch/locality_domain/runtime.py +++ b/tensorrt_llm/_torch/locality_domain/runtime.py @@ -29,6 +29,7 @@ from tensorrt_llm._torch.locality_domain.policy import PartitionPlan from tensorrt_llm._torch.locality_domain_utils import ( end_for_all_locality_domain, + get_current_locality_domain, get_locality_domain_compute_sm_counts, get_locality_domain_mempool, get_locality_domain_stream, @@ -84,6 +85,12 @@ def partition_context(self, partition_id: int) -> Iterator[None]: forward execution. Kernel runners read it via get_current_locality_domain(). """ with locality_domain_device(partition_id): + current_partition = get_current_locality_domain() + if current_partition != partition_id: + raise RuntimeError( + "failed to enter the requested locality domain: " + f"requested={partition_id}, current={current_partition}" + ) with torch.cuda.stream(self.partition_stream(partition_id)): yield diff --git a/tensorrt_llm/_torch/locality_domain_utils.py b/tensorrt_llm/_torch/locality_domain_utils.py index ce7127ca1141..b1deafc524fe 100644 --- a/tensorrt_llm/_torch/locality_domain_utils.py +++ b/tensorrt_llm/_torch/locality_domain_utils.py @@ -35,7 +35,6 @@ import tensorrt_llm as trtllm import tensorrt_llm.bindings.internal.runtime as _tbr -from tensorrt_llm._utils import get_sm_version __all__ = [ "get_locality_domain_stream", @@ -226,18 +225,24 @@ def is_locality_domain_supported(device: int | None = None) -> bool: return False -@lru_cache(maxsize=1) -def is_locality_domain_enabled() -> bool: +@lru_cache(maxsize=None) +def is_locality_domain_enabled(device: int | None = None) -> bool: """ - Check if LOCALITY_DOMAIN localization is enabled on this system. + Check whether LOCALITY_DOMAIN localization is enabled on a CUDA device. + + Callers that may switch CUDA devices should pass an explicit ordinal so + the cached result cannot be inherited from a different device. """ if os.getenv("DISABLE_LOCALITY_DOMAINS", "0") == "1": return False if not torch.cuda.is_available(): return False - if get_sm_version() != 107: + if device is None: + device = torch.cuda.current_device() + properties = torch.cuda.get_device_properties(device) + if int(properties.major) * 10 + int(properties.minor) != 107: return False - return is_locality_domain_supported() + return is_locality_domain_supported(device) def get_current_locality_domain() -> int | None: @@ -334,7 +339,7 @@ def locality_domain_device(locality_domain_id: int | None): raise ValueError(f"locality_domain_id must be 0, 1, or None, got {locality_domain_id}") # If LOCALITY_DOMAIN is not enabled, do nothing and keep current LOCALITY_DOMAIN as None - if not is_locality_domain_enabled(): + if not is_locality_domain_enabled(torch.cuda.current_device()): yield return diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index e4daf4ffb167..e7141c9e068d 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -1015,6 +1015,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): 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 + use_gvr_locality_domain = sparse_attention_config.use_gvr_locality_domain indexer_k_dtype = sparse_attention_config.indexer_k_dtype else: index_n_heads = pretrained_config.index_n_heads @@ -1030,6 +1031,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): enable_heuristic_topk = False use_self_sampling_topk = True use_gvr_emission = False + use_gvr_locality_domain = False default_sparse_attention_config = DeepSeekV4SparseAttentionConfig( ) indexer_k_dtype = default_sparse_attention_config.indexer_k_dtype @@ -1048,6 +1050,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): 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['use_gvr_locality_domain'] = use_gvr_locality_domain indexer_config['indexer_k_dtype'] = indexer_k_dtype return indexer_config @@ -1087,6 +1090,8 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): 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 + use_gvr_locality_domain = ( + sparse_attention_config.use_gvr_locality_domain) indexer_k_dtype = sparse_attention_config.indexer_k_dtype index_share_for_mtp_iteration = sparse_attention_config.index_share_for_mtp_iteration else: @@ -1101,6 +1106,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): enable_heuristic_topk = False use_self_sampling_topk = True use_gvr_emission = False + use_gvr_locality_domain = False indexer_k_dtype = "fp8" index_share_for_mtp_iteration = None kwargs[ @@ -1119,6 +1125,7 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): enable_heuristic_topk=enable_heuristic_topk, use_self_sampling_topk=use_self_sampling_topk, use_gvr_emission=use_gvr_emission, + use_gvr_locality_domain=use_gvr_locality_domain, 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 8093a96a016d..939756fbdbd1 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -4,15 +4,26 @@ from __future__ import annotations +import os +from contextlib import nullcontext from enum import Enum +from typing import TYPE_CHECKING, Callable import torch import torch.nn as nn from tensorrt_llm.logger import logger +from ..locality_domain.gvr_topk import ( + GvrTopKRowShardPlan, + is_gvr_topk_locality_workload_large_enough, + plan_gvr_topk_row_shards, +) from ..memory_buffer_utils import get_memory_buffers +if TYPE_CHECKING: + from ..locality_domain.runtime import LocalityDomainRuntime + class TopKImplementation(str, Enum): """Top-K implementations grouped by backend and algorithm.""" @@ -48,6 +59,7 @@ def __init__( decode_implementation: TopKImplementation | None = None, compress_ratio: int = 1, gvr_self_sampling: bool = True, + use_gvr_locality_domain: bool = False, ) -> None: super().__init__() self.top_k = top_k @@ -61,6 +73,14 @@ def __init__( # 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 + # Rubin locality-domain row sharding is an explicit prototype opt-in. + # Runtime resources remain lazy so default and unsupported paths keep + # the original full-device GVR lifecycle. + self.use_gvr_locality_domain = use_gvr_locality_domain + self._gvr_locality_runtime: LocalityDomainRuntime | None = None + self._gvr_locality_capability: dict[int, bool] = {} + self._gvr_locality_topologies: dict[int, tuple[tuple[int, int], ...]] = {} + self._gvr_locality_ready_launches: set[tuple] = set() # 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. @@ -279,6 +299,215 @@ def _get_radix_workspace( ) return radix_indices, radix_values + def _build_gvr_locality_launch( + self, + scores: torch.Tensor, + next_n: int, + max_seq_len: int, + ) -> tuple["LocalityDomainRuntime", GvrTopKRowShardPlan, tuple] | None: + """Build a capture-stable Rubin row-sharding launch, if eligible.""" + if not self.use_gvr_locality_domain or not scores.is_cuda: + return None + + score_width = min(int(max_seq_len), int(scores.shape[1])) + if not is_gvr_topk_locality_workload_large_enough( + num_rows=int(scores.shape[0]), + next_n=int(next_n), + score_width=score_width, + top_k=self.top_k, + ): + # In particular, BS=1 and small score envelopes never initialize + # locality-domain resources and retain the full-device launch. + return None + + device_index = scores.get_device() + with torch.cuda.device(device_index): + from ..locality_domain_utils import ( + get_current_locality_domain, + is_locality_domain_supported, + ) + + # Avoid recursively partitioning a Top-K already submitted under + # another locality-domain composite. + if get_current_locality_domain() is not None: + return None + + capturing = torch.cuda.is_current_stream_capturing() + capable = self._gvr_locality_capability.get(device_index) + if capable is None: + if capturing: + raise RuntimeError( + "GVR locality-domain capability is cold during CUDA " + "Graph capture; run this shape once eagerly before capture" + ) + from ..cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + properties = torch.cuda.get_device_properties(device_index) + sm_version = int(properties.major) * 10 + int(properties.minor) + capable = ( + sm_version == 107 + and IS_CUTLASS_DSL_RUBIN_AVAILABLE + and os.environ.get("DISABLE_LOCALITY_DOMAINS", "0") != "1" + and is_locality_domain_supported(device_index) + ) + self._gvr_locality_capability[device_index] = capable + if not capable: + logger.warning_once( + "use_gvr_locality_domain=True but Rubin locality-domain " + "execution is unavailable; keeping self-sampling GVR on " + "the full-device stream.", + key="gvr_locality_domain_unavailable", + ) + return None + + runtime = self._gvr_locality_runtime + if runtime is None: + if capturing: + raise RuntimeError( + "GVR locality-domain resources are cold during CUDA " + "Graph capture; run this shape once eagerly before capture" + ) + from ..locality_domain.runtime import LocalityDomainRuntime + + runtime = LocalityDomainRuntime(num_partitions=2) + self._gvr_locality_runtime = runtime + + topology = self._gvr_locality_topologies.get(device_index) + if topology is None: + if capturing: + raise RuntimeError( + "GVR locality-domain topology is cold during CUDA Graph " + "capture; run this shape once eagerly before capture" + ) + topology = runtime.topology_identity() + self._gvr_locality_topologies[device_index] = topology + + properties = torch.cuda.get_device_properties(device_index) + device_num_sms = int(properties.multi_processor_count) + if any(total_sms != device_num_sms for _, total_sms in topology): + raise RuntimeError( + "GVR locality-domain topology does not match the score " + f"device: topology={topology}, device_num_sms={device_num_sms}" + ) + try: + plan = plan_gvr_topk_row_shards( + num_rows=int(scores.shape[0]), + next_n=int(next_n), + score_width=score_width, + top_k=self.top_k, + topology=topology, + ) + except ValueError as error: + raise RuntimeError(f"invalid GVR locality-domain topology: {error}") from error + if plan is None: + return None + + # run_varlen derives npad from shape[1] for a one-row launch. Do + # not turn a legal multi-row strided view into an illegal shard. + if any(shard.num_rows == 1 for shard in plan.shards) and scores.shape[1] % 4: + return None + + shard_geometry = tuple( + ( + shard.num_rows, + int(scores.shape[1]) if shard.num_rows == 1 else int(scores.stride(0)), + shard.num_sms, + ) + for shard in plan.shards + ) + launch_key = ( + device_index, + shard_geometry, + self.top_k, + int(max_seq_len), + int(next_n), + self.compress_ratio, + plan.topology, + ) + if capturing and launch_key not in self._gvr_locality_ready_launches: + raise RuntimeError( + "GVR locality-domain launchers or workspaces are cold " + "during CUDA Graph capture; run this shape once eagerly " + "before capture" + ) + return runtime, plan, launch_key + + def _run_gvr_locality_domain( + self, + runner: Callable[..., None], + scores: torch.Tensor, + sequence_lengths: torch.Tensor, + output_indices: torch.Tensor, + next_n: int, + max_seq_len: int, + ) -> bool: + """Run two non-overlapping GVR row slices on Rubin locality domains.""" + launch = self._build_gvr_locality_launch(scores, next_n, max_seq_len) + if launch is None: + return False + runtime, plan, launch_key = launch + + num_rows = int(scores.shape[0]) + num_requests = num_rows // int(next_n) + if tuple(output_indices.shape) != (num_rows, self.top_k): + raise RuntimeError( + "GVR locality-domain output shape must match the unsplit " + f"launch: expected {(num_rows, self.top_k)}, got " + f"{tuple(output_indices.shape)}" + ) + if sequence_lengths.dim() != 1 or int(sequence_lengths.shape[0]) != num_requests: + raise RuntimeError( + "GVR locality-domain sequence_lengths must have one entry " + f"per request: expected {(num_requests,)}, got " + f"{tuple(sequence_lengths.shape)}" + ) + if sequence_lengths.device != scores.device or output_indices.device != scores.device: + raise RuntimeError( + "GVR locality-domain scores, sequence_lengths, and output_indices " + "must be on the same device" + ) + + # CPU tensors are useful for orchestration tests with a mocked launch + # plan/runtime. Production locality launches always enter the CUDA + # device guard above. + device_context = torch.cuda.device(scores.device) if scores.is_cuda else nullcontext() + with device_context: + runtime.fork() + try: + for shard in plan.shards: + with runtime.partition_context(shard.partition_id): + runner( + scores[shard.row_start : shard.row_end], + sequence_lengths[shard.request_start : shard.request_end], + output_indices[shard.row_start : shard.row_end], + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=max_seq_len * self.compress_ratio, + ) + except Exception: + # Preserve the launch failure if cleanup also fails; the + # original exception is the actionable cause. + try: + runtime.join() + except Exception: + logger.exception( + "failed to join Rubin locality-domain streams while " + "handling a GVR Top-K launch failure" + ) + raise + runtime.join() + + # Mark ready only after both launches and the join complete. The key + # covers every leaf cache/workspace dimension needed during capture. + self._gvr_locality_ready_launches.add(launch_key) + logger.info_once( + "Rubin locality-domain self-sampling GVR Top-K engaged; only " + "the already-produced Top-K rows are sharded (the logits " + "producer remains full-device).", + key="gvr_locality_domain_engaged", + ) + return True + def _forward_decode_gvr( self, scores: torch.Tensor, @@ -323,14 +552,25 @@ def _forward_decode_gvr( # max_seq_len in compressed index space; run_varlen's value # is in KV-token space like sequence_lengths, so multiply it # back by the compression ratio. - selfsampling_topk_run_varlen( - scores, - sequence_lengths, - output_indices, - next_n=next_n, - compress_ratio=self.compress_ratio, - max_seq_len=max_seq_len * self.compress_ratio, - ) + if not ( + self.use_gvr_locality_domain + and self._run_gvr_locality_domain( + selfsampling_topk_run_varlen, + scores, + sequence_lengths, + output_indices, + next_n, + max_seq_len, + ) + ): + selfsampling_topk_run_varlen( + scores, + sequence_lengths, + output_indices, + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=max_seq_len * self.compress_ratio, + ) return output_indices logger.warning_once( "self-sampling GVR is selected but the decode scores do not " diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1b4ba4426e87..03424bfbeb82 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1105,6 +1105,23 @@ class DeepSeekSparseAttentionConfig(SeqLenAwareSparseAttentionConfig): "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.") + use_gvr_locality_domain: bool = Field( + default=False, + status="prototype", + description="Enable prototype Rubin locality-domain row sharding for the " + "self-sampling GVR V2 decode Top-K. Only the already-produced logits " + "rows are split; the indexer logits producer remains full-device. " + "Only takes effect with enable_heuristic_topk=True, " + "use_self_sampling_topk=True, K in {512, 1024, 2048}, compression " + "ratio in {1, 4}, a supported GVR input format, SM107, Rubin CuTe " + "DSL support, at least two requests, and a sufficiently large " + "capture-stable workload envelope. Eligible smaller workloads keep " + "the full-device GVR path; other unsupported GVR configurations keep " + "the existing radix fallback. Disabled by default while the " + "provisional gain thresholds are tuned on R200. This prototype " + "supports one in-flight model execution per CUDA device; concurrent " + "GraphExec replay or multiple engines sharing a device are not yet " + "supported.") indexer_k_dtype: Literal["fp8", "fp4"] = Field( default="fp8", description= @@ -1248,6 +1265,7 @@ def _value(name: str, default=None): enable_heuristic_topk=self.enable_heuristic_topk, use_self_sampling_topk=self.use_self_sampling_topk, use_gvr_emission=self.use_gvr_emission, + use_gvr_locality_domain=self.use_gvr_locality_domain, indexer_k_dtype=self.indexer_k_dtype, is_full_indexer_layer=self._is_full_indexer_layer( pretrained_config, kwargs.get("layer_idx")), @@ -1282,6 +1300,7 @@ def _value(name: str, default=None): enable_heuristic_topk=self.enable_heuristic_topk, use_self_sampling_topk=self.use_self_sampling_topk, use_gvr_emission=self.use_gvr_emission, + use_gvr_locality_domain=self.use_gvr_locality_domain, 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, @@ -1371,6 +1390,7 @@ def _value(name: str, default=None): enable_heuristic_topk=self.enable_heuristic_topk, use_self_sampling_topk=self.use_self_sampling_topk, use_gvr_emission=self.use_gvr_emission, + use_gvr_locality_domain=self.use_gvr_locality_domain, indexer_k_dtype=self.indexer_k_dtype, compress_ratios=self.compress_ratios, window_size=self.window_size, @@ -1398,6 +1418,7 @@ def _value(name: str, default=None): enable_heuristic_topk=self.enable_heuristic_topk, use_self_sampling_topk=self.use_self_sampling_topk, use_gvr_emission=self.use_gvr_emission, + use_gvr_locality_domain=self.use_gvr_locality_domain, 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/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index c1d613bb408a..2e225c3e1965 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1836,6 +1836,13 @@ "kind": "value", "path": "sparse_attention_config.use_gvr_emission" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.use_gvr_locality_domain" + }, { "allowed_values": [], "annotation": "", 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 95800a39ad25..da1b9b5250ed 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -244,6 +244,18 @@ def test_use_gvr_emission_defaults_off(): assert sparse_config.to_sparse_metadata_params().use_gvr_emission is False +@pytest.mark.parametrize("enabled", [False, True]) +def test_use_gvr_locality_domain_threads_to_params(enabled): + """The default-off prototype flag reaches both execution and warmup params.""" + sparse_config = DeepSeekV4SparseAttentionConfig( + index_topk=512, + use_gvr_locality_domain=enabled, + ) + + assert sparse_config.to_sparse_params().use_gvr_locality_domain is enabled + assert sparse_config.to_sparse_metadata_params().use_gvr_locality_domain is enabled + + @pytest.mark.parametrize( "enable_heuristic,use_cute_dsl,sm_version,compress_ratio,next_n,should_warmup", [ @@ -601,6 +613,7 @@ def test_indexer_two_level_gvr_dispatch( use_cute_dsl_topk=use_cute_dsl, enable_heuristic_topk=True, use_self_sampling_topk=use_self_sampling, + use_gvr_locality_domain=True, ) with ( @@ -617,6 +630,7 @@ def test_indexer_two_level_gvr_dispatch( assert indexer.top_k.decode_implementation == expected_decode assert indexer.top_k.gvr_self_sampling == use_self_sampling + assert indexer.top_k.use_gvr_locality_domain assert indexer.top_k.needs_gvr_prior == (not use_self_sampling) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index d37194578e63..fff3c5f1762b 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -789,6 +789,7 @@ def test_deepseek_v4_sparse_ratios_prefer_checkpoint_defaults(tmp_path, monkeypa compress_ratios=[1, 1, 4, 128, 4, 128, 4], q_split_threshold=2048, skip_indexer_for_short_seqs=False, + use_gvr_locality_domain=True, ) model_config = ModelConfig.from_pretrained( @@ -799,6 +800,7 @@ def test_deepseek_v4_sparse_ratios_prefer_checkpoint_defaults(tmp_path, monkeypa ) assert model_config.sparse_attention_config.compress_ratios == [128, 128, 4, 128, 4, 128, 1, 4] + assert model_config.sparse_attention_config.use_gvr_locality_domain is True # V4 sparse MLA hardcodes window_size==128 (FMHA kernel TileSizeKV; see # the runtime assertion in deepseek_v4.py:DeepseekV4TrtllmAttentionMetadata # __post_init__), so this is the only legal value here. diff --git a/tests/unittest/_torch/modules/test_gvr_topk_locality.py b/tests/unittest/_torch/modules/test_gvr_topk_locality.py new file mode 100644 index 000000000000..baffbd7d9231 --- /dev/null +++ b/tests/unittest/_torch/modules/test_gvr_topk_locality.py @@ -0,0 +1,476 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""CPU tests for Rubin locality-domain GVR row sharding.""" + +import sys +from contextlib import contextmanager, nullcontext +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.locality_domain.gvr_topk import plan_gvr_topk_row_shards +from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation + + +def test_plan_keeps_every_next_n_request_group_together() -> None: + plan = plan_gvr_topk_row_shards( + num_rows=15, + next_n=3, + score_width=4096, + top_k=512, + topology=((104, 208), (104, 208)), + min_total_score_elements=0, + min_score_elements_per_sm=0, + ) + + assert plan is not None + assert [(shard.request_start, shard.request_end) for shard in plan.shards] == [ + (0, 3), + (3, 5), + ] + assert [(shard.row_start, shard.row_end) for shard in plan.shards] == [ + (0, 9), + (9, 15), + ] + assert all(shard.row_start % 3 == 0 and shard.row_end % 3 == 0 for shard in plan.shards) + + +def test_plan_uses_real_asymmetric_partition_sm_counts() -> None: + plan = plan_gvr_topk_row_shards( + num_rows=10, + next_n=1, + score_width=4096, + top_k=512, + topology=((64, 208), (128, 208)), + min_total_score_elements=0, + min_score_elements_per_sm=0, + ) + + assert plan is not None + assert [shard.num_requests for shard in plan.shards] == [3, 7] + assert [shard.num_sms for shard in plan.shards] == [64, 128] + + +@pytest.mark.parametrize( + "kwargs", + [ + # A single request cannot be split without breaking the local-row ABI. + {"num_rows": 4, "next_n": 4, "score_width": 1 << 18}, + # The total launch is too small to amortize two launches and events. + {"num_rows": 2, "next_n": 1, "score_width": 4096}, + # Total work passes, but the smaller shard has too little work per SM. + {"num_rows": 4, "next_n": 1, "score_width": 1 << 18}, + ], +) +def test_provisional_gain_gate_keeps_unprofitable_shapes_full_device(kwargs) -> None: + topology = ((8, 208), (200, 208)) if kwargs["num_rows"] == 4 else ((104, 208), (104, 208)) + assert ( + plan_gvr_topk_row_shards( + top_k=512, + topology=topology, + **kwargs, + ) + is None + ) + + +@pytest.mark.parametrize( + "topology,error", + [ + (((104, 208), (104, 210)), "disagree"), + (((120, 208), (120, 208)), "overlap"), + ], +) +def test_invalid_topology_is_rejected_before_launch(topology, error) -> None: + with pytest.raises(ValueError, match=error): + plan_gvr_topk_row_shards( + num_rows=8, + next_n=1, + score_width=1 << 18, + top_k=512, + topology=topology, + ) + + +class _FakeRuntime: + def __init__(self) -> None: + self.events: list[object] = [] + self.current_partition: int | None = None + + def fork(self) -> None: + self.events.append("fork") + + @contextmanager + def partition_context(self, partition_id: int): + self.events.append(("enter", partition_id)) + self.current_partition = partition_id + try: + yield + finally: + self.current_partition = None + self.events.append(("exit", partition_id)) + + def join(self) -> None: + self.events.append("join") + + +def _install_fake_selfsampling_runner(monkeypatch) -> Mock: + runner = Mock() + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k", + SimpleNamespace(selfsampling_topk_run_varlen=runner), + ) + return runner + + +def test_topk_dispatches_disjoint_request_slices_and_joins(monkeypatch) -> None: + runtime = _FakeRuntime() + plan = plan_gvr_topk_row_shards( + num_rows=6, + next_n=2, + score_width=8, + top_k=2, + topology=((1, 3), (2, 3)), + min_total_score_elements=0, + min_score_elements_per_sm=0, + ) + assert plan is not None + launch_key = ("cpu-mock",) + monkeypatch.setattr( + TopK, + "_build_gvr_locality_launch", + lambda self, scores, next_n, max_seq_len: (runtime, plan, launch_key), + ) + runner = _install_fake_selfsampling_runner(monkeypatch) + + def run_slice(scores, lengths, output, **kwargs) -> None: + partition_id = runtime.current_partition + runtime.events.append(("launch", partition_id, scores.shape[0], lengths.tolist())) + output.fill_(10 + int(partition_id)) + + runner.side_effect = run_slice + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + use_gvr_locality_domain=True, + ) + scores = torch.randn(6, 8) + lengths = torch.tensor([32, 48, 64], dtype=torch.int32) + output = torch.empty(6, 2, dtype=torch.int32) + + result = top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + next_n=2, + max_seq_len=16, + ) + + assert result is output + assert runtime.events == [ + "fork", + ("enter", 0), + ("launch", 0, 2, [32]), + ("exit", 0), + ("enter", 1), + ("launch", 1, 4, [48, 64]), + ("exit", 1), + "join", + ] + assert output[:2].eq(10).all() + assert output[2:].eq(11).all() + assert top_k._gvr_locality_ready_launches == {launch_key} + assert runner.call_count == 2 + for call in runner.call_args_list: + assert call.kwargs == {"next_n": 2, "compress_ratio": 4, "max_seq_len": 64} + + +def test_topk_joins_and_does_not_mark_ready_after_shard_failure(monkeypatch) -> None: + runtime = _FakeRuntime() + plan = plan_gvr_topk_row_shards( + num_rows=2, + next_n=1, + score_width=8, + top_k=2, + topology=((1, 2), (1, 2)), + min_total_score_elements=0, + min_score_elements_per_sm=0, + ) + assert plan is not None + monkeypatch.setattr( + TopK, + "_build_gvr_locality_launch", + lambda self, scores, next_n, max_seq_len: (runtime, plan, ("failed",)), + ) + runner = _install_fake_selfsampling_runner(monkeypatch) + runner.side_effect = RuntimeError("synthetic launch failure") + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + use_gvr_locality_domain=True, + ) + + with pytest.raises(RuntimeError, match="synthetic"): + top_k( + torch.randn(2, 8), + torch.empty(2, 2, dtype=torch.int32), + is_prefill=False, + sequence_lengths=torch.tensor([8, 8], dtype=torch.int32), + scan_lengths=torch.tensor([8, 8], dtype=torch.int32), + max_seq_len=8, + ) + + assert runtime.events[-1] == "join" + assert not top_k._gvr_locality_ready_launches + + +def test_topk_preserves_launch_failure_if_join_also_fails(monkeypatch) -> None: + runtime = _FakeRuntime() + plan = plan_gvr_topk_row_shards( + num_rows=2, + next_n=1, + score_width=8, + top_k=2, + topology=((1, 2), (1, 2)), + min_total_score_elements=0, + min_score_elements_per_sm=0, + ) + assert plan is not None + monkeypatch.setattr( + TopK, + "_build_gvr_locality_launch", + lambda self, scores, next_n, max_seq_len: (runtime, plan, ("failed",)), + ) + runner = _install_fake_selfsampling_runner(monkeypatch) + runner.side_effect = RuntimeError("synthetic launch failure") + runtime.join = Mock(side_effect=RuntimeError("synthetic join failure")) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + use_gvr_locality_domain=True, + ) + + with pytest.raises(RuntimeError, match="synthetic launch failure"): + top_k( + torch.randn(2, 8), + torch.empty(2, 2, dtype=torch.int32), + is_prefill=False, + sequence_lengths=torch.tensor([8, 8], dtype=torch.int32), + scan_lengths=torch.tensor([8, 8], dtype=torch.int32), + max_seq_len=8, + ) + + runtime.join.assert_called_once_with() + assert not top_k._gvr_locality_ready_launches + + +@pytest.mark.parametrize( + "lengths_shape,output_shape,error", + [ + ((3,), (2, 2), "sequence_lengths"), + ((2,), (3, 2), "output shape"), + ], +) +def test_topk_rejects_global_shapes_that_sharding_could_mask( + monkeypatch, lengths_shape, output_shape, error +) -> None: + runtime = _FakeRuntime() + plan = plan_gvr_topk_row_shards( + num_rows=2, + next_n=1, + score_width=8, + top_k=2, + topology=((1, 2), (1, 2)), + min_total_score_elements=0, + min_score_elements_per_sm=0, + ) + assert plan is not None + monkeypatch.setattr( + TopK, + "_build_gvr_locality_launch", + lambda self, scores, next_n, max_seq_len: (runtime, plan, ("invalid",)), + ) + _install_fake_selfsampling_runner(monkeypatch) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + use_gvr_locality_domain=True, + ) + + with pytest.raises(RuntimeError, match=error): + top_k( + torch.randn(2, 8), + torch.empty(output_shape, dtype=torch.int32), + is_prefill=False, + sequence_lengths=torch.full(lengths_shape, 8, dtype=torch.int32), + scan_lengths=torch.full(lengths_shape, 8, dtype=torch.int32), + max_seq_len=8, + ) + + assert runtime.events == [] + + +def test_cold_capture_fails_before_locality_runtime_initialization(monkeypatch) -> None: + from tensorrt_llm._torch import locality_domain_utils + + class _FakeCudaScores: + is_cuda = True + shape = (4, 1 << 18) + + @staticmethod + def get_device() -> int: + return 0 + + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + use_gvr_locality_domain=True, + ) + # Capability discovery is an eager, non-allocating step. Pretend it has + # already succeeded so this test isolates the capture lifecycle guard. + top_k._gvr_locality_capability[0] = True + monkeypatch.setattr(torch.cuda, "device", lambda device: nullcontext()) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + monkeypatch.setattr(locality_domain_utils, "get_current_locality_domain", lambda: None) + + with pytest.raises(RuntimeError, match="resources are cold"): + top_k._build_gvr_locality_launch(_FakeCudaScores(), next_n=1, max_seq_len=1 << 18) + + assert top_k._gvr_locality_runtime is None + + +def test_cold_capture_does_not_run_capability_discovery(monkeypatch) -> None: + from tensorrt_llm._torch import locality_domain_utils + + class _FakeCudaScores: + is_cuda = True + shape = (4, 1 << 18) + + @staticmethod + def get_device() -> int: + return 0 + + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + use_gvr_locality_domain=True, + ) + properties = Mock(side_effect=AssertionError("must not query properties during capture")) + supported = Mock(side_effect=AssertionError("must not query driver support during capture")) + monkeypatch.setattr(torch.cuda, "device", lambda device: nullcontext()) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_properties", properties) + monkeypatch.setattr(locality_domain_utils, "get_current_locality_domain", lambda: None) + monkeypatch.setattr(locality_domain_utils, "is_locality_domain_supported", supported) + + with pytest.raises(RuntimeError, match="capability is cold"): + top_k._build_gvr_locality_launch(_FakeCudaScores(), next_n=1, max_seq_len=1 << 18) + + properties.assert_not_called() + supported.assert_not_called() + + +def test_explicit_opt_in_on_unsupported_cpu_uses_one_full_device_call(monkeypatch) -> None: + runner = _install_fake_selfsampling_runner(monkeypatch) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + use_gvr_locality_domain=True, + ) + scores = torch.randn(4, 8) + lengths = torch.full((4,), 8, dtype=torch.int32) + output = torch.empty(4, 2, dtype=torch.int32) + + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + max_seq_len=1 << 18, + ) + + runner.assert_called_once_with( + scores, + lengths, + output, + next_n=1, + compress_ratio=1, + max_seq_len=1 << 18, + ) + + +def test_default_off_does_not_call_locality_helper(monkeypatch) -> None: + runner = _install_fake_selfsampling_runner(monkeypatch) + locality_helper = Mock(side_effect=AssertionError("default path must not enter helper")) + monkeypatch.setattr(TopK, "_run_gvr_locality_domain", locality_helper) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + ) + scores = torch.randn(4, 8) + lengths = torch.full((4,), 8, dtype=torch.int32) + output = torch.empty(4, 2, dtype=torch.int32) + + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + max_seq_len=1 << 18, + ) + + locality_helper.assert_not_called() + runner.assert_called_once() + + +def test_temporal_gvr_never_enters_locality_row_sharding(monkeypatch) -> None: + build_launch = Mock(side_effect=AssertionError("V1 must not shard")) + monkeypatch.setattr(TopK, "_build_gvr_locality_launch", build_launch) + temporal_runner = Mock() + monkeypatch.setattr( + torch.ops.trtllm, + "cute_dsl_gvr_topk_decode", + temporal_runner, + raising=False, + ) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + use_gvr_locality_domain=True, + ) + scores = torch.randn(2, 8) + lengths = torch.full((2,), 8, dtype=torch.int32) + + top_k( + scores, + torch.empty(2, 2, dtype=torch.int32), + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + max_seq_len=1 << 18, + gvr_ext_kwargs={"gvr_prior_indices": torch.zeros(2, 2, dtype=torch.int32)}, + ) + + build_launch.assert_not_called() + temporal_runner.assert_called_once() diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 6e3daea5e695..980d3ad93869 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + import json import struct import types @@ -230,6 +245,34 @@ def test_deepseek_v4_missing_compress_ratios_raises(tmp_path, monkeypatch): ModelConfig.from_pretrained(str(tmp_path)) +def test_deepseek_v32_preserves_gvr_locality_domain_config(tmp_path, monkeypatch): + from tensorrt_llm._torch import model_config as model_config_module + from tensorrt_llm._torch.configs.deepseek_v3 import DeepseekV3Config + from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig + + pretrained_config = DeepseekV3Config( + architectures=["DeepseekV32ForCausalLM"], + index_n_heads=64, + index_head_dim=128, + index_topk=2048, + indexer_rope_interleave=False, + ) + monkeypatch.setattr( + model_config_module, "load_pretrained_config", lambda *args, **kwargs: pretrained_config + ) + + model_config = ModelConfig.from_pretrained( + str(tmp_path), + sparse_attention_config=DeepSeekSparseAttentionConfig( + use_gvr_locality_domain=True, + ), + attn_backend="TRTLLM", + moe_backend="TRTLLM", + ) + + assert model_config.sparse_attention_config.use_gvr_locality_domain is True + + def test_model_config_sets_is_encoder_decoder_from_pretrained_config(): model_config = ModelConfig( pretrained_config=make_pretrained_config( diff --git a/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py b/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py index a5b7b5480104..77e057ef1cbc 100644 --- a/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py +++ b/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py @@ -80,7 +80,10 @@ def test_is_locality_domain_enabled_requires_rubin(self): is_locality_domain_enabled.cache_clear() with ( patch("torch.cuda.is_available", return_value=True), - patch("tensorrt_llm._torch.locality_domain_utils.get_sm_version", return_value=100), + patch( + "torch.cuda.get_device_properties", + return_value=SimpleNamespace(major=10, minor=0), + ), patch( "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_supported", return_value=True, @@ -93,7 +96,10 @@ def test_is_locality_domain_enabled_allows_rubin_when_supported(self): is_locality_domain_enabled.cache_clear() with ( patch("torch.cuda.is_available", return_value=True), - patch("tensorrt_llm._torch.locality_domain_utils.get_sm_version", return_value=107), + patch( + "torch.cuda.get_device_properties", + return_value=SimpleNamespace(major=10, minor=7), + ), patch( "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_supported", return_value=True, @@ -102,6 +108,24 @@ def test_is_locality_domain_enabled_allows_rubin_when_supported(self): assert is_locality_domain_enabled() is_locality_domain_enabled.cache_clear() + def test_is_locality_domain_enabled_caches_per_explicit_device(self): + is_locality_domain_enabled.cache_clear() + + def properties(device): + return SimpleNamespace(major=10, minor=7 if device == 1 else 0) + + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_properties", side_effect=properties), + patch( + "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_supported", + return_value=True, + ), + ): + assert not is_locality_domain_enabled(0) + assert is_locality_domain_enabled(1) + is_locality_domain_enabled.cache_clear() + class TestLocalityDomainComputeTopology: """Pure mocked tests for compute topology and grid sizing.""" @@ -243,6 +267,26 @@ def test_runtime_rejects_unsupported_num_partitions(self): with pytest.raises(ValueError, match="num_partitions"): LocalityDomainRuntime(num_partitions=bad) + def test_runtime_rejects_silent_partition_context_noop(self, monkeypatch): + @contextmanager + def no_op_partition_context(partition_id): + yield + + monkeypatch.setattr( + locality_domain_runtime, + "locality_domain_device", + no_op_partition_context, + ) + monkeypatch.setattr( + locality_domain_runtime, + "get_current_locality_domain", + lambda: None, + ) + + with pytest.raises(RuntimeError, match="failed to enter"): + with LocalityDomainRuntime().partition_context(0): + pass + @pytest.mark.parametrize("locality_domain_id", [-1, 2]) def test_compute_sm_counts_reject_invalid_partition(self, locality_domain_id): with pytest.raises(ValueError, match="locality_domain_id must be 0 or 1"): From 7ffe80bca0afa8b9677967a3516249baaf71efc4 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Fri, 4 Sep 2026 04:49:33 +0000 Subject: [PATCH 08/10] [None][feat] Self-sampling GVR V2 prefill indexer top-K Extend the hint-free self-sampling GVR top-K from decode to the DSA prefill phase, replacing the CUDA radix prefill on the layers that already select self-sampling for decode (same enable_heuristic_topk x use_self_sampling_topk dispatch; no new config/API). Kernel: a const_expr prefill mode on GvrMainKernel for per-row [ks, ke) windows. ks/ke ride the unused pre_idx/kv_lens ABI slots (byte-identical signature; distinct compile key). The base rounds down to a 16B boundary and the <=3 lead lanes are masked positionally (no materialized -inf); the clamps tighten to the last in-window float4 so the reads are exact. Output is the local (column - ks) frame with a trailing -1 pad; nv <= k emits identity. All edits are const_expr-gated, so decode/legacy codegen is unchanged. Host: run_prefill (stride(0) for all row counts, no device reads, <=32768-row slabs for gridDim.y, R=1) + warmup_prefill (<=6 engines/k). Module: a CUTE_DSL_GVR prefill branch with an all-short short-circuit and a format-gate radix fallthrough. Indexer selects the engine for prefill iff it does for decode; metadata warms the prefill leg. Compression-ratio agnostic, so V3.2, V4 Flash and V4 Pro share one prefill path. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 15 +- .../attention_backend/sparse/dsa/metadata.py | 13 + .../blackwell/top_k/__init__.py | 2 + .../top_k/gvr_topk_decode_self_sampling.py | 258 ++++++++++++--- .../gvr_topk_decode_self_sampling_host.py | 244 ++++++++++++++ tensorrt_llm/_torch/modules/top_k.py | 54 ++- .../attention/sparse/dsa/test_dsa_indexer.py | 54 +++ tests/unittest/_torch/modules/test_top_k.py | 118 +++++++ .../parallel/test_gvr_selfsampling_topk.py | 313 ++++++++++++++++++ 9 files changed, 1014 insertions(+), 57 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 6b5a158336c5..9281b5242e8e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -759,9 +759,22 @@ def __init__( if self.use_cute_dsl_topk else TopKImplementation.CUDA_RADIX ) + # The self-sampling engine has a prefill form (per-row [ks, ke) + # windows); select it for prefill on exactly the layers where the + # two-level dispatch picks self-sampling for decode, so both phases + # share one config and one warmup. The temporal-hint engine has no + # prefill form, so those layers keep the exact radix prefill. + prefill_top_k_implementation = ( + TopKImplementation.CUTE_DSL_GVR + if ( + decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR + and self._use_self_sampling_topk + ) + else TopKImplementation.CUDA_RADIX + ) self.top_k = TopK( self.index_topk, - prefill_implementation=TopKImplementation.CUDA_RADIX, + prefill_implementation=prefill_top_k_implementation, decode_implementation=decode_top_k_implementation, compress_ratio=self.compress_ratio, gvr_self_sampling=self._use_self_sampling_topk, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index cebfbafc0571..98e21cb0dcb8 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -450,6 +450,19 @@ def warmup_selfsampling_topk( for bs in batch_sizes or (): rows.add(int(bs) * nn) msl_c = int(self.get_indexer_max_seq_len()) + # Prefill leg: the self-sampling engine also serves prefill (per-row + # [ks, ke) windows). It is placed BEFORE the DeepGEMM decode-stride + # guard below (which would return early for an odd msl_c) because the + # DeepGEMM prefill stride is always a 256-multiple. Bounded to the six + # tier x U engines per k; best-effort under the same OOM guard as the + # decode leg. + try: + _ss_host.warmup_prefill(int(top_k), max(msl_c, 32768)) + except torch.cuda.OutOfMemoryError: + logger.warning( + "self-sampling GVR prefill warmup ran out of memory; prefill " + "engines will JIT-compile lazily on first touch instead." + ) if self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits: # mirror the DSL paged-MQA arena stride (rows round up to 256 # elements). A drift here only degrades warmup to unused keys — diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index d5cab389489b..25d06bff8bd0 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py @@ -20,6 +20,7 @@ from .gvr_topk_decode_direct import DirectTopKKernel from .gvr_topk_decode_dispatch import is_tiered_topk_supported, tiered_topk from .gvr_topk_decode_reg import GvrRegKernel +from .gvr_topk_decode_self_sampling_host import run_prefill as selfsampling_topk_run_prefill from .gvr_topk_decode_self_sampling_host import run_varlen as selfsampling_topk_run_varlen from .gvr_topk_decode_tp import GvrTpKernel from .single_pass_multi_cta_radix_topk import SinglePassMultiCTARadixTopKKernel @@ -36,4 +37,5 @@ "tiered_topk", "is_tiered_topk_supported", "selfsampling_topk_run_varlen", + "selfsampling_topk_run_prefill", ] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index c1a856b35025..cee8b6bde066 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -1331,6 +1331,7 @@ def __init__( cr_shift: int = 0, r_const: int = 1, hint_free: bool = False, + prefill: bool = False, ) -> None: assert nbs == 256, "SNB must stay 256" assert blk in (256, 512, 1024) and u in (1, 2, 4, 8) @@ -1353,6 +1354,21 @@ def __init__( self.r_const = int(r_const) # hint-free: gather_hint sites compiled out (sentinel pass-through) self.hint_free = bool(hint_free) + # prefill: per-row window [ks, ke) from row_starts/row_ends (rides the + # kv_lens / pre_idx ABI slots); base rounds down to a 16B boundary and + # the <=3 lead lanes are positionally masked. Single-CTA-per-row only + # (no SPLIT/workspace/TSH); next_n==1, cr_shift==0 (ks/ke are already + # in compressed column units). All prefill edits are const_expr-gated + # so legacy/varlen codegen stays byte-identical. + self.prefill = bool(prefill) + if self.prefill: + assert ( + self.varlen + and self.hint_free + and self.next_n == 1 + and self.cr_shift == 0 + and not self.split + ) if self.varlen: assert self.next_n >= 1 and self.cr_shift in (0, 2) and self.r_const >= 1 # TSH-floor staging arm. SPLIT-only compile-time key; the CUDA form @@ -1509,17 +1525,41 @@ def kern( short = cutlass.Int32(0) n_row = cutlass.Int32(0) tsh_run = cutlass.Int32(1) + # prefill window offset: lead = ks & 3 low lanes masked, col0 = ks + # rounded down to a float4 boundary (declared before the dynamic ifs + # per the scoping rule; stay 0 in every non-prefill compile). + lead = cutlass.Int32(0) + col0 = cutlass.Int32(0) if cutlass.const_expr(self.varlen): - req = row // cutlass.Int32(self.next_n) - rr = row % cutlass.Int32(self.next_n) - kvl = kv_lens[req] - nv = (kvl - cutlass.Int32(self.next_n) + rr + cutlass.Int32(1)) >> cutlass.Int32( - self.cr_shift - ) - if nv < cutlass.Int32(0): - nv = cutlass.Int32(0) - if nv > npad: - nv = npad + if cutlass.const_expr(self.prefill): + # per-row window [ks, ke) already in compressed column units + # (kv_lens slot = row_starts, pre_idx slot = row_ends); no + # next_n / cr_shift math. Clamp only for memory safety — the + # indexer guarantees 0 <= ks <= ke <= logits.shape[1]. + ks = kv_lens[row] + ke = pre_idx[row] + if ks < cutlass.Int32(0): + ks = cutlass.Int32(0) + if ke > npad: + ke = npad + if ks > ke: + ks = ke + nv = ke - ks + lead = ks & cutlass.Int32(3) + col0 = ks - lead + if col0 > npad - cutlass.Int32(4): + col0 = npad - cutlass.Int32(4) + else: + req = row // cutlass.Int32(self.next_n) + rr = row % cutlass.Int32(self.next_n) + kvl = kv_lens[req] + nv = (kvl - cutlass.Int32(self.next_n) + rr + cutlass.Int32(1)) >> cutlass.Int32( + self.cr_shift + ) + if nv < cutlass.Int32(0): + nv = cutlass.Int32(0) + if nv > npad: + nv = npad n_row = nv if nv <= k: short = cutlass.Int32(1) @@ -1533,7 +1573,12 @@ def kern( TGT2 = cutlass.Int32(0x3FFFFFFF) Q = cutlass.Int32(0) if short == cutlass.Int32(0): - n = nv + if cutlass.const_expr(self.prefill): + # scan extent from the rounded-down base col0 spans the + # lead pad plus the real window: [col0, ke) = nv + lead. + n = nv + lead + else: + n = nv n4v = n >> cutlass.Int32(2) # Ladder-scalar baselines only: the real SMP/SS2/TGT/TGT2 are # derived by warp0 alone in the block below (bit-identical @@ -1598,6 +1643,13 @@ def kern( s_lad = smem.allocate_tensor( cutlass.Int32, cute.make_ordered_layout((4,), order=(0,)), byte_alignment=16 ) + if cutlass.const_expr(self.prefill): + # per-row lead (0..3) broadcast slot: warp0 publishes it under the + # existing s_lad barrier; every masked-lane site reloads it from + # smem so no live register is carried on the 64-register arms. + s_lead = smem.allocate_tensor( + cutlass.Int32, cute.make_ordered_layout((1,), order=(0,)), byte_alignment=4 + ) blob = smem.allocate_tensor( # dynamic-equivalent region cutlass.Int8, cute.make_ordered_layout((self.dyn_bytes,), order=(0,)), byte_alignment=16 ) @@ -1643,14 +1695,27 @@ def kern( row64 = cutlass.Int64(row) # _pin_i64: keep the row base a REGISTER across the attempt/tile scf # regions (NVVM otherwise re-derives ld.param+%ctaid.y+mul per region) - x_addr = _pin_i64(logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4)) + if cutlass.const_expr(self.prefill): + # base rounded down to the col0 float4 boundary (16B aligned since + # the row base is 16B aligned and col0 is a multiple of 4). + x_addr = _pin_i64( + logits.iterator.toint() + + (row64 * cutlass.Int64(npad) + cutlass.Int64(col0)) * cutlass.Int64(4) + ) + else: + x_addr = _pin_i64( + logits.iterator.toint() + row64 * cutlass.Int64(npad) * cutlass.Int64(4) + ) # varlen: pre_idx is REQUEST-level [num_rows/next_n, k] — a request's # next_n rows share one hint row (production contract); legacy mode - # keeps the per-row mapping (next_n == 1 makes them identical). - prow64 = row64 - if cutlass.const_expr(self.varlen): - prow64 = cutlass.Int64(row // cutlass.Int32(self.next_n)) - p_addr = pre_idx.iterator.toint() + prow64 * cutlass.Int64(k) * cutlass.Int64(4) + # keeps the per-row mapping (next_n == 1 makes them identical). In + # prefill the pre_idx slot is 1-D row_ends (already consumed in the + # prologue), so this dead hint pointer is compiled out. + if cutlass.const_expr(not self.prefill): + prow64 = row64 + if cutlass.const_expr(self.varlen): + prow64 = cutlass.Int64(row // cutlass.Int32(self.next_n)) + p_addr = pre_idx.iterator.toint() + prow64 * cutlass.Int64(k) * cutlass.Int64(4) out_row = out[row, None] ws_addr = ws.iterator.toint() gdon_addr = ws_addr # slab views @@ -1792,11 +1857,22 @@ def kern( s_lad[1] = SS2 s_lad[2] = TGT s_lad[3] = TGT2 + if cutlass.const_expr(self.prefill): + s_lead[0] = lead # Register-free L2 hints for the first U-batch of this CTA's own # P3 slice (clamped in-row): the data P3 touches first starts # flowing while warp0 walks the chain. Short rows clamp every # hint to the row's last line — harmless. - plim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) + # prefill: the base is shifted to col0, so the clamp must stay in + # the row's own window [col0, ke) — an npad-based clamp would + # over-read col0 columns past the last row's allocation. n4-1 is + # the last full in-window float4 (>=0 even for the n=0 short pass). + if cutlass.const_expr(self.prefill): + plim4 = n4 - cutlass.Int32(1) + if plim4 < cutlass.Int32(0): + plim4 = cutlass.Int32(0) + else: + plim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) for uu in cutlass.range_constexpr(U): # NOTE: names must not collide with the PRIME-LATE block's # i_/ic — the DSL kills inner-scope names at region exit and @@ -1823,6 +1899,17 @@ def kern( p4 = tidx * SS2 * cutlass.Int32(2) C.ld_g_f32x4(atom128, x_addr, p4, fsa) C.ld_g_f32x4(atom128, x_addr, p4 + cutlass.Int32(1), fsb) + if cutlass.const_expr(self.prefill): + # only thread 0's fsa (float4 index 0) can hold the <=3 masked + # lead lanes; substitute the always-valid lane 3 so the sample + # min/max fold and histogram stay finite and count-invariant + # (a materialized -inf would drive f2s_rz to INT_MIN and write + # out of bounds in the sample histogram at :1930). + if tidx == cutlass.Int32(0): + ld_ = s_lead[0] + for q in cutlass.range_constexpr(3): + if cutlass.Int32(q) < ld_: + fsa[q] = fsa[3] # ============ P2: quantile rung from the sample ====================== smn = cutlass.Float32(float("inf")) @@ -1856,7 +1943,13 @@ def kern( cute.arch.barrier() # ---- barrier (sample redux publish) ---- # PRIME-LATE prefetch block: strictly after the barrier. - lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) + # prefill clamps to the last in-window float4 (see plim4 note above). + if cutlass.const_expr(self.prefill): + lim4 = n4 - cutlass.Int32(1) + if lim4 < cutlass.Int32(0): + lim4 = cutlass.Int32(0) + else: + lim4 = (npad >> cutlass.Int32(2)) - cutlass.Int32(1) pf = [cute.make_rmem_tensor((4,), cutlass.Float32) for _ in range(max(PFD, 1))] if cutlass.const_expr(self.pf): fullsl = cutlass.Int32(0) @@ -2145,6 +2238,12 @@ def kern( if okq != cutlass.Int32(0): # ok-gated (+inf-pad escape) for q in cutlass.range_constexpr(4): M = M | (cutlass.Int32(vv[q] >= TF) << cutlass.Int32(uu * 4 + q)) + if cutlass.const_expr(self.prefill): + # the <=lead lead lanes live only in bits 0..lead-1 of the + # i0==0 tile (float4 0, thread 0, part 0); clear them so the + # reservation, survivor walk and re-reads never see them. + if i0 == cutlass.Int32(0): + M = M & (cutlass.Int32(-1) << s_lead[0]) # prefetch roll-forward BEFORE reservation/walk if cutlass.const_expr(self.pf): hasnext = cutlass.Int32(0) @@ -2453,7 +2552,12 @@ def kern( if bq >= B: p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) if p < lim1: - out_row[p] = idv + # prefill: staged idx are in the col0 frame; the + # local output frame is relative to ks = col0+lead. + if cutlass.const_expr(self.prefill): + out_row[p] = idv - s_lead[0] + else: + out_row[p] = idv else: if whole == cutlass.Int32(0): q2 = p - above @@ -2476,22 +2580,32 @@ def kern( i_ = i0_ if i0_ >= hi2: i_ = tail0 + (i0_ - hi2) - x = C.ldg_f32(x_addr, i_) - if x >= TF: - bq = C.f2s_rz((x - TF) * SC) - if bq > cutlass.Int32(NBS - 1): - bq = cutlass.Int32(NBS - 1) - if bq >= B: - p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) - if p < lim1: - out_row[p] = i_ - else: - if whole == cutlass.Int32(0): - q2 = p - above - if q2 < cutlass.Int32(CMPB): - s_ck64[q2] = ( - cutlass.Uint64(C.fkey(x)) << cutlass.Uint64(32) - ) | cutlass.Uint64(cutlass.Uint32(i_)) + masked = cutlass.Int32(0) + if cutlass.const_expr(self.prefill): + # skip the <=lead lead lanes (col0-frame positions + # 0..lead-1 hold the previous request's finite logits) + if i_ < s_lead[0]: + masked = cutlass.Int32(1) + if masked == cutlass.Int32(0): + x = C.ldg_f32(x_addr, i_) + if x >= TF: + bq = C.f2s_rz((x - TF) * SC) + if bq > cutlass.Int32(NBS - 1): + bq = cutlass.Int32(NBS - 1) + if bq >= B: + p = C.atomic_add_cta(s_hist.iterator + bq, cutlass.Int32(1)) + if p < lim1: + if cutlass.const_expr(self.prefill): + out_row[p] = i_ - s_lead[0] + else: + out_row[p] = i_ + else: + if whole == cutlass.Int32(0): + q2 = p - above + if q2 < cutlass.Int32(CMPB): + s_ck64[q2] = ( + cutlass.Uint64(C.fkey(x)) << cutlass.Uint64(32) + ) | cutlass.Uint64(cutlass.Uint32(i_)) i0_ = i0_ + cutlass.Int32(BLK) # ---- P6 refine ---- @@ -2520,11 +2634,14 @@ def kern( cutlass.Uint64(s_ck64[mc2]) > cutlass.Uint64(u64v) ) if r_ < need: - out_row[above + r_] = cutlass.Int32( + idv6 = cutlass.Int32( cutlass.Uint32( cutlass.Uint64(u64v) & cutlass.Uint64(0xFFFFFFFF) ) ) + if cutlass.const_expr(self.prefill): + idv6 = idv6 - s_lead[0] + out_row[above + r_] = idv6 i = i + cutlass.Int32(BLK) else: # key-space narrowing over ck64 @@ -2626,6 +2743,11 @@ def kern( p1 = cutlass.Int32(1) if iu == ethr: p2 = cutlass.Int32(1) + # staged idx are col0-frame; shift to the ks-relative + # local output frame (p1=p2=0 for i>=mc, so the -lead + # on the idv=0 default is never emitted). + if cutlass.const_expr(self.prefill): + idv = idv - s_lead[0] self._ballot_pair_emit( p1, p2, @@ -2789,6 +2911,10 @@ def kern( if tie_m != cutlass.Int32(0): if iu == ethr: p2 = cutlass.Int32(1) + # staged idx (already >= lead via the P3 M-mask) -> local + # frame; the x_addr re-read above stays in the col0 frame. + if cutlass.const_expr(self.prefill): + idv = idv - s_lead[0] self._ballot_pair_emit( p1, p2, idv, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane ) @@ -2799,7 +2925,13 @@ def kern( rhi = cutlass.Uint32(0xFFFFFFFF) above2 = cutlass.Int32(0) need2 = k - m2 = n + # prefill: the genuine window is [lead, n); the <=lead lead + # lanes are excluded from the histogram, the emit and the + # candidate count so they never join a tie class. + lead_db = cutlass.Int32(0) + if cutlass.const_expr(self.prefill): + lead_db = s_lead[0] + m2 = n - lead_db ethr = cutlass.Int64(0) tie_m = cutlass.Int32(1) if tidx < cutlass.Int32(NBS): @@ -2829,8 +2961,8 @@ def kern( if sh2 < cutlass.Int32(0): sh2 = cutlass.Int32(0) sh2u = cutlass.Uint32(sh2) - i = tidx - while i < n: # whole row + i = tidx + lead_db # prefill: skip lead lanes + while i < n: # whole row (window [lead, n)) uq = C.fkey(C.ldg_f32(x_addr, i)) if uq >= cutlass.Uint32(rlo): if uq <= cutlass.Uint32(rhi): @@ -2879,15 +3011,16 @@ def kern( p1 = cutlass.Int32(0) p2 = cutlass.Int32(0) if i < n: - uq = C.fkey(C.ldg_f32(x_addr, i)) - iu = cutlass.Int64(uq) - if iu > ethr: - p1 = cutlass.Int32(1) - if tie_m != cutlass.Int32(0): - if iu == ethr: - p2 = cutlass.Int32(1) + if i >= lead_db: # prefill: exclude lead lanes + uq = C.fkey(C.ldg_f32(x_addr, i)) + iu = cutlass.Int64(uq) + if iu > ethr: + p1 = cutlass.Int32(1) + if tie_m != cutlass.Int32(0): + if iu == ethr: + p2 = cutlass.Int32(1) self._ballot_pair_emit( - p1, p2, i, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane + p1, p2, i - lead_db, cutlass.Int32(0), nA, nA, nT, out_row, s_scal, lane ) it = it + cutlass.Int32(1) @@ -2969,16 +3102,26 @@ def __call__( _COMPILE_CACHE = {} -def get_compiled(tpl: tuple, options_extra: str = "", hint_free: bool = False) -> Any: +def get_compiled( + tpl: tuple, options_extra: str = "", hint_free: bool = False, prefill: bool = False +) -> Any: """Compile (or fetch) the gvr_main variant for constexpr tuple tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG) — legacy, or tpl = (BLK, U, MINB, NBS, KPT, SPLIT, TSHG, NEXT_N, CR_SHIFT, R_CONST) — per-row varlen mode (TSHG slot is ignored: varlen compiles the TSH - machinery in whenever SPLIT and gates it per row at runtime).""" - key = (tuple(tpl), options_extra, bool(hint_free)) + machinery in whenever SPLIT and gates it per row at runtime). + + ``prefill`` selects the per-row window mode. It shares the varlen tuple + (next_n=1, cr_shift=0) but has a distinct prologue, so it MUST be part of + the cache key — otherwise a DSv3.2 decode varlen engine and the prefill + engine collide on the same tuple. The prefill compile also retypes the + pre_idx ABI slot to a 1-D align-4 fake (it carries 4B-aligned row_ends).""" + key = (tuple(tpl), options_extra, bool(hint_free), bool(prefill)) hit = _COMPILE_CACHE.get(key) if hit is not None: return hit + if prefill: + assert len(tpl) == 10, "prefill compile requires the varlen tuple" if len(tpl) == 7: blk, u, minb, nbs, kpt, split, tshg = tpl kern = GvrMainKernel( @@ -2999,6 +3142,7 @@ def get_compiled(tpl: tuple, options_extra: str = "", hint_free: bool = False) - cr_shift=cr_shift, r_const=r_const, hint_free=bool(hint_free), + prefill=bool(prefill), ) r0, c0 = cute.sym_int(), cute.sym_int() r1, c1 = cute.sym_int(), cute.sym_int() @@ -3008,9 +3152,15 @@ def get_compiled(tpl: tuple, options_extra: str = "", hint_free: bool = False) - logits_fake = _crt.make_fake_compact_tensor( cutlass.Float32, (r0, c0), stride_order=(1, 0), assumed_align=16 ) - pre_fake = _crt.make_fake_compact_tensor( - cutlass.Int32, (r1, c1), stride_order=(1, 0), assumed_align=16 - ) + if prefill: + # pre_idx slot carries row_ends [rows] int32 (4B-aligned slices). + pre_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (r1,), stride_order=(0,), assumed_align=4 + ) + else: + pre_fake = _crt.make_fake_compact_tensor( + cutlass.Int32, (r1, c1), stride_order=(1, 0), assumed_align=16 + ) out_fake = _crt.make_fake_compact_tensor( cutlass.Int32, (r2, c2), stride_order=(1, 0), assumed_align=16 ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index 51d837997789..e13070fda535 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -802,6 +802,69 @@ def _available_num_sms(device: torch.device | int) -> int: """Return SMs available to launches in the current execution domain.""" return _execution_domain(device)[0] +# ---- prefill launcher cache ------------------------------------------------ +# Prefill routes always force R==1 (single CTA per row): route_streaming gives +# R>1 only for b<=74, so the representative row counts below (first row of each +# route band) pin R=1 and reduce the engine set to <=6 per k. The launcher +# compiled function depends only on the row TIER, k and the envelope bucket +# (which selects U on the tier-0 1024-thread arm; tiers 1/2 fix U), never on +# the exact row count (arbitrary q-tile / q-split remainders) or npad (a +# runtime scalar), so the cache stays bounded over a long-running server. +_PREFILL_CACHE = {} +_PREFILL_ROW_SLAB = 32768 # gridDim.y <= 65535; slab so keys stay bounded +_PREFILL_TIER_ROWS = (75, 149, 297) # (rows<=148, 149..296, >296) band reps + + +def _prefill_tier(rows: int) -> int: + return 0 if rows <= 148 else 1 if rows <= 296 else 2 + + +def _prefill_bucket(n_env: int) -> int: + # pow2-quantize the envelope so a growing envelope reuses one plan; cap at + # 32768 because U=8 for every n>=32768 on the tier-0 arm. + return min(1 << max(int(n_env) - 1, 1).bit_length(), 32768) + + +def _prefill_cache_key(tier: int, k: int, n_bucket: int): + # tiers 1/2 fix U, so the bucket does not change their engine — collapse it + # to one key so warmup covers them with a single launch. + return (tier, k, n_bucket if tier == 0 else 0) + + +def _prefill_launcher(tier: int, k: int, n_bucket: int) -> tuple: + """Capture-time prefill plan + compiled launcher (main family, R=1). + + Mirrors ``_varlen_launcher``'s main branch but with r_const=1, split=False + (so tsh_en=0) and the prefill compile flag. SCAP_/CMP_/aim are envelope + upper bounds; npad is filled per call in ``run_prefill``.""" + key = _prefill_cache_key(tier, k, n_bucket) + hit = _PREFILL_CACHE.get(key) + if hit is not None: + return hit + b_route = _PREFILL_TIER_ROWS[tier] + n_route = max(n_bucket, k + 1) + plan = route_streaming(b_route, n_route, n_route, k, force_main=True) + if plan["kernel"] != "main": + raise RuntimeError(f"prefill route did not land on gvr_main: {plan['kernel']}") + rt = plan["rt"] + if rt["R"] != 1: + raise RuntimeError(f"prefill requires R==1 (got {rt['R']})") + tpl = tuple(plan["tpl"]) + dev = _device() + fn = dev.get_compiled(tpl[:6] + (False,) + (1, 0, 1), hint_free=True, prefill=True) + big = tier == 0 + # r_const==1 branch of the _varlen_launcher tuning scalars + aim_base = ( + (4 * k if k >= 1024 else 2 * k) if big else ((11 * k) // 8 if k >= 1024 else (3 * k) // 2) + ) + sfac = 64 if k >= 1024 else 32 + amin = (7 * k) // 2 + sd_en = 1 if (k > 1024 and not big) else 0 + tail = (aim_base, sfac, amin, sd_en, 0) # tsh_en=0 (split=False) + lc = ("main", fn, (rt["SCAP_"], rt["CMP_"]), tail) + _PREFILL_CACHE[key] = lc + return lc + def _varlen_launcher( num_rows: int, @@ -1675,6 +1738,127 @@ def run_varlen( return +def run_prefill( + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + max_row_len: int | None = None, + workspace: torch.Tensor | None = None, +) -> None: + """Hint-free self-sampling Top-K for the prefill phase, per-row windows. + + Row semantics (mirror of ``topKPerRowPrefill`` / ``indexer_topk_prefill``): + row ``r`` selects the Top-K of ``logits[r, ks:ke]`` where + ``ks = row_starts[r]``, ``ke = row_ends[r]`` (both int32, in the SAME + compressed column units the DeepGEMM prefill producer emits — no + ``next_n`` / ``compress_ratio`` math). ``k`` comes from + ``indices.shape[1]``. The output is written in the LOCAL frame (column + minus ``ks``) with a trailing ``-1`` pad; rows with ``nv = ke - ks <= k`` + get the identity ``0..nv-1`` (matching the radix short-row contract). The + engine reads exactly ``[r*npad + (ks & ~3), r*npad + ke)`` — no dependence + on any producer slack. + + Envelope: ``max_row_len`` (a capture-stable engine constant) or, when + omitted, ``logits.shape[1]`` — a host int, so the call performs NO device + reads and is CUDA-graph-replay safe (it refuses to compile a new plan + under capture). Launches in ``<=65535``-row slabs so ``gridDim.y`` never + overflows. + + KNOWN LIMITATION: rows containing NaN inside the window are out of + contract (as for the radix reference — both order NaN implementation- + specifically). DeepGEMM prefill logits are finite in-window. Trusted + invariant: ``0 <= ks <= ke <= logits.shape[1]`` (the indexer guarantees + it); the kernel clamps ``ke <= npad`` for memory safety only. + """ + if logits.dtype is not _F32: + raise RuntimeError( + f"logits must be float32 (got {logits.dtype}); bf16/fp16 paths " + "are a follow-up — see the PR roadmap" + ) + for _nm, _t in (("row_starts", row_starts), ("row_ends", row_ends)): + if not (isinstance(_t, _TENSOR) and _t.is_cuda): + raise RuntimeError(f"{_nm} must be a CUDA tensor") + if _t.dtype is not _I32: + raise RuntimeError(f"{_nm} must be int32") + if _t.dim() != 1: + raise RuntimeError(f"{_nm} must be 1-D") + if not _t.is_contiguous(): + raise RuntimeError(f"{_nm} must be contiguous") + if len(logits.shape) != 2: + raise RuntimeError("logits must be 2-D") + num_rows = logits.shape[0] + if num_rows == 0: + return + if row_starts.shape[0] != num_rows or row_ends.shape[0] != num_rows: + raise RuntimeError( + f"row_starts/row_ends length must equal logits.shape[0]={num_rows}, " + f"got {row_starts.shape[0]}/{row_ends.shape[0]}" + ) + if not (logits.is_cuda and indices.is_cuda): + raise RuntimeError("all tensors must be CUDA") + if indices.dtype is not _I32: + raise RuntimeError("indices must be int32") + if len(indices.shape) != 2 or indices.shape[0] != num_rows: + raise RuntimeError(f"indices must be [num_rows={num_rows}, k], got {tuple(indices.shape)}") + if not indices.is_contiguous(): + raise RuntimeError("indices must be contiguous") + k = indices.shape[1] + if k < 4 or (k & 3): + raise RuntimeError(f"index_topk must be a multiple of 4 and >= 4, got {k}") + if indices.data_ptr() & 15: + raise RuntimeError("indices base must be 16-byte aligned") + if logits.stride(1) != 1: + raise RuntimeError("logits inner stride must be 1") + # DeepGEMM prefill rows are 1024B-aligned with >=256 float slack, so the + # row stride is valid for EVERY row count (the varlen 1-row shape[1] rule + # is a paged-MQA-arena quirk that would reject odd-width single-token + # prefill tiles — the common fully-cached follow-up turn). + npad = logits.stride(0) + if npad & 3: + raise RuntimeError(f"npad (logits row stride) must be a multiple of 4, got {npad}") + if logits.data_ptr() & 15: + raise RuntimeError("logits base must be 16-byte aligned") + d = logits.get_device() + if not 0 <= d < _GVR_MAX_DEV: + raise RuntimeError(f"device index out of range: {d}") + lg = logits + if logits.shape[1] != npad: + need = logits.storage_offset() + num_rows * npad + if logits.untyped_storage().size() // 4 < need: + raise RuntimeError("logits view storage too small to widen to its row stride") + lg = logits.as_strided((num_rows, npad), (npad, 1), logits.storage_offset()) + if workspace is not None: + validate_run_ws(workspace, logits) + ws = kernel_view(workspace) + else: + ws = _ws_hot.get(d) + if ws is None: + ws = default_workspace(logits) + n_env = _index(max_row_len) if max_row_len is not None else logits.shape[1] + n_env = min(max(n_env, 1), npad) + n_bucket = _prefill_bucket(n_env) + for r0 in range(0, num_rows, _PREFILL_ROW_SLAB): + r1 = min(r0 + _PREFILL_ROW_SLAB, num_rows) + tier = _prefill_tier(r1 - r0) + lc = _PREFILL_CACHE.get(_prefill_cache_key(tier, k, n_bucket)) + if lc is None: + if _is_capturing(): + raise RuntimeError( + "prefill launcher not compiled for this shape — warm up " + "before CUDA graph capture" + ) + lc = _prefill_launcher(tier, k, n_bucket) + _, fn, (scap, cmp_), tail = lc + # ABI parity with the varlen main call: pre_idx slot = row_ends, + # kv_lens slot = row_starts. The n / SMP / TGT / Q / SS2 / TGT2 launch + # scalars are dead (re-derived per row); only npad / k / SCAP_ / CMP_ + # matter, R=1. + pre = (0, npad, k, scap, cmp_, 1, 0, 0, 0, 0, 0) + fn(lg[r0:r1], row_ends[r0:r1], indices[r0:r1], ws, *pre, row_starts[r0:r1], *tail) + return + + __all__ = [ "route", "route_static", @@ -1684,7 +1868,9 @@ def run_varlen( "run", "run_ws", "run_varlen", + "run_prefill", "warmup_varlen", + "warmup_prefill", "workspace_bytes", "WS_BYTES", "default_workspace", @@ -1834,3 +2020,61 @@ def warmup_varlen( ) with _VARLEN_WARMUP_LOCK: _VARLEN_WARMUP_DONE.add(key) + + +_PREFILL_WARMUP_DONE: set = set() +_PREFILL_WARMUP_LOCK = threading.Lock() + + +def warmup_prefill( + top_k: int, + max_cols: int, + num_rows_list: Sequence[int] = (1, 149, 297), + row_stride: int | None = None, +) -> None: + """TESTING/INIT ONLY — compile the prefill engine set before serving. + + Six engines per k at most: the tier-0 (1024-thread) arm walks the pow2 + envelope buckets (U = 1/2/4/8), tiers 1/2 fix U so one launch each. One + tiny real launch per distinct ``(tier, k, bucket)`` cache key; ``ks=0``, + ``ke=n_env`` (all long rows). ``max_cols`` is the compressed max column + count (``get_indexer_max_seq_len``); the bucket caps at 32768 (U=8 above), + so envelopes past it share one key. The done-key gates only the GPU + launches — the ``_PREFILL_CACHE`` population is idempotent. + """ + dev = torch.cuda.current_device() + k = int(top_k) + max_cols = int(max_cols) + lo = _prefill_bucket(k + 1) + hi = _prefill_bucket(max_cols) + buckets = [] + b = lo + while b <= hi: + buckets.append(b) + b <<= 1 + if not buckets: + buckets = [hi] + keys = {} # cache_key -> (tier, bucket) representative for the launch + for rows in num_rows_list: + tier = _prefill_tier(int(rows)) + bset = buckets if tier == 0 else buckets[:1] + for bk in bset: + keys.setdefault(_prefill_cache_key(tier, k, bk), (tier, bk)) + done_key = (dev, k, max_cols, tuple(sorted(int(r) for r in num_rows_list)), row_stride) + with _PREFILL_WARMUP_LOCK: + if done_key in _PREFILL_WARMUP_DONE: + return + for tier, bk in keys.values(): + rows = _PREFILL_TIER_ROWS[tier] + stride = row_stride if row_stride is not None else ((bk + 256 + 255) // 256 * 256) + if stride < bk or stride % 4: + stride = (max(stride, bk) + 256 + 255) // 256 * 256 + logits = torch.zeros((rows, stride), dtype=torch.float32, device=dev) + ks = torch.zeros((rows,), dtype=torch.int32, device=dev) + ke = torch.full((rows,), bk, dtype=torch.int32, device=dev) + out = torch.empty((rows, k), dtype=torch.int32, device=dev) + run_prefill(logits[:, :bk], ks, ke, out, max_row_len=bk) + del logits, ks, ke, out + torch.cuda.synchronize() + with _PREFILL_WARMUP_LOCK: + _PREFILL_WARMUP_DONE.add(done_key) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 939756fbdbd1..def249712121 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -163,7 +163,44 @@ def _forward_prefill( row_ends, output_indices, ) - if self.prefill_implementation == TopKImplementation.CUTE_DSL_RADIX: + if self.prefill_implementation == TopKImplementation.CUTE_DSL_GVR: + # hint-free k derives from the output width; pin it to the module's k + assert output_indices.shape[1] == self.top_k + if not self.gvr_self_sampling: + # the temporal (hint) GVR engine has no prefill form + logger.warning_once( + "temporal GVR has no prefill engine; using the CUDA radix prefill Top-K.", + key="gvr_temporal_prefill_radix", + ) + elif scores.shape[1] <= self.top_k: + # every row is short (nv <= k): the exact radix path emits the + # identity/-1 answer without reading logits — cheaper than a + # zero-work self-sampling launch. Deliberate, no warning. + pass + elif self._selfsampling_prefill_ok(scores): + from ..cute_dsl_kernels.blackwell.top_k import selfsampling_topk_run_prefill + + logger.info_once( + "self-sampling GVR prefill top-K engaged " + f"(K={self.top_k}, cr={self.compress_ratio}, hint-free).", + key="selfsampling_topk_prefill_engaged", + ) + # ks/ke are already in compressed column units; run_prefill + # writes the local (column - ks) frame with -1 pad and no host + # reads (envelope from scores.shape[1]). + selfsampling_topk_run_prefill(scores, row_starts, row_ends, output_indices) + return output_indices + else: + # engine hardware-format gate missed (e.g. a non-fp4 layer with + # an odd DeepGEMM width, or a bf16 producer): exact radix. + logger.warning_once( + "self-sampling GVR prefill is selected but the scores do " + "not satisfy the engine's hardware-format gate " + f"(dtype={scores.dtype}, strides={tuple(scores.stride())}); " + "falling back to the CUDA radix prefill Top-K.", + key="selfsampling_topk_prefill_fallthrough", + ) + elif self.prefill_implementation == TopKImplementation.CUTE_DSL_RADIX: # Keep the op's reread policy default; only its copy width is tuned. torch.ops.trtllm.cute_dsl_indexer_topk_prefill_blackwell( scores, @@ -174,7 +211,7 @@ def _forward_prefill( _CUTE_DSL_PREFILL_COPY_BITS, ) return output_indices - if self.prefill_implementation != TopKImplementation.CUDA_RADIX: + elif self.prefill_implementation != TopKImplementation.CUDA_RADIX: raise NotImplementedError( f"{self.prefill_implementation.value} does not support prefill Top-K" ) @@ -187,6 +224,19 @@ def _forward_prefill( ) return output_indices + def _selfsampling_prefill_ok(self, scores: torch.Tensor) -> bool: + """Engine hardware-format gate for the self-sampling prefill Top-K. + + fp32 row-major scores with a float4-aligned row stride and a 16B base + (the DeepGEMM prefill logits arena, whose rows are 1024B-aligned). The + all-short tile case is handled by the caller before this check.""" + return ( + scores.dtype == torch.float32 + and scores.stride(1) == 1 + and scores.stride(0) % 4 == 0 + and scores.data_ptr() % 16 == 0 + ) + def _forward_decode( self, scores: torch.Tensor, 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 da1b9b5250ed..5b0b8bfb3ef7 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -310,6 +310,54 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( cute_dsl_radix.assert_not_called() +@pytest.mark.parametrize( + "use_self_sampling,sm_version,msl_c,should_warmup", + [ + (True, 100, 65536, True), + (True, 100, 30001, True), # odd msl_c must not skip the prefill leg + (False, 100, 65536, False), # temporal-hint layers: no prefill engine + (True, 90, 65536, False), # non-datacenter Blackwell + ], +) +def test_metadata_warmup_selfsampling_prefill_leg( + use_self_sampling, sm_version, msl_c, should_warmup +): + """The self-sampling warmup drives BOTH the decode (varlen) and the prefill + engines; the prefill leg sits before the DeepGEMM decode-stride guard so an + odd msl_c cannot skip it.""" + metadata = SimpleNamespace( + enable_gvr_topk=True, + use_self_sampling_topk=use_self_sampling, + sparse_mla_topk=512, + _indexer_compress_ratio=4, + kv_cache_manager=SimpleNamespace(), + get_indexer_max_seq_len=Mock(return_value=msl_c), + sparse_metadata_params=SimpleNamespace(use_cute_dsl_paged_mqa_logits=True), + num_sms=148, + ) + ss_host = ( + "tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_topk_decode_self_sampling_host" + ) + with ( + 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=sm_version, + ), + patch(f"{ss_host}.warmup_prefill") as warmup_prefill, + patch(f"{ss_host}.warmup_varlen"), + ): + DSAtrtllmAttentionMetadata.warmup_selfsampling_topk(metadata, next_n=1, batch_sizes=[8]) + + if should_warmup: + warmup_prefill.assert_called_once_with(512, max(msl_c, 32768)) + else: + warmup_prefill.assert_not_called() + + def test_kv_lens_row_reorder_threshold(): """Prepare row order only when CuTe DSL GVR has enough decode rows.""" num_sms = 16 @@ -632,6 +680,12 @@ def test_indexer_two_level_gvr_dispatch( assert indexer.top_k.gvr_self_sampling == use_self_sampling assert indexer.top_k.use_gvr_locality_domain assert indexer.top_k.needs_gvr_prior == (not use_self_sampling) + # Prefill uses the self-sampling engine on exactly the self-sampling + # layers; the temporal-hint layers keep the exact radix prefill. + expected_prefill = ( + TopKImplementation.CUTE_DSL_GVR if use_self_sampling else TopKImplementation.CUDA_RADIX + ) + assert indexer.top_k.prefill_implementation == expected_prefill @skip_pre_hopper diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index a9d319eac9f5..975df084ecf4 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -466,3 +466,121 @@ def test_gvr_emission_reset_parks_reused_slots(monkeypatch) -> None: assert torch.isfinite(lines[0]) and torch.isfinite(lines[3]), ( "untouched slots must keep their closed-loop state" ) + + +def _install_fake_prefill_runner(monkeypatch) -> Mock: + """Stub both self-sampling entries (a test may exercise decode and prefill + through the same lazily imported module); return the prefill Mock.""" + prefill = Mock() + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k", + SimpleNamespace( + selfsampling_topk_run_varlen=Mock(), + selfsampling_topk_run_prefill=prefill, + ), + ) + return prefill + + +def _prefill_call(top_k: TopK, scores: torch.Tensor, out_width: int = 2): + rows = scores.shape[0] + row_starts = torch.zeros(rows, dtype=torch.int32) + row_ends = torch.full((rows,), scores.shape[1], dtype=torch.int32) + output = torch.full((rows, out_width), -1, dtype=torch.int32) + top_k( + scores, + output, + is_prefill=True, + row_starts=row_starts, + row_ends=row_ends, + ) + return row_starts, row_ends, output + + +def test_gvr_v2_prefill_routes_to_selfsampling_runner(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + top_k = TopK( + 2, + prefill_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + ) + scores = torch.randn(3, 8) # fp32, stride(0)=8 %4==0, contiguous -> gate ok + row_starts, row_ends, output = _prefill_call(top_k, scores) + + runner.assert_called_once() + args, kwargs = runner.call_args + assert args[0] is scores and args[1] is row_starts and args[2] is row_ends + assert args[3] is output + assert kwargs == {} + + +def test_gvr_v2_prefill_format_gate_falls_back_to_radix(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + top_k = TopK( + 2, + prefill_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + ) + scores = torch.randn(3, 8, dtype=torch.bfloat16) # dtype gate miss + row_starts, row_ends, output = _prefill_call(top_k, scores) + + runner.assert_not_called() + radix.assert_called_once_with(scores, row_starts, row_ends, output, 2) + + +def test_gvr_v2_prefill_odd_stride_falls_back_to_radix(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + top_k = TopK(2, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + scores = torch.randn(3, 6) # stride(0)=6, 6 % 4 != 0 -> gate miss + + _prefill_call(top_k, scores) + + runner.assert_not_called() + radix.assert_called_once() + + +def test_gvr_v2_prefill_all_short_uses_radix(monkeypatch) -> None: + """scores.shape[1] <= top_k: every row is short, so the exact radix + identity/-1 path runs (no logits read) with no fallthrough warning.""" + runner = _install_fake_prefill_runner(monkeypatch) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + top_k = TopK(4, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + scores = torch.randn(3, 4) # shape[1] == top_k + + _prefill_call(top_k, scores, out_width=4) + + runner.assert_not_called() + radix.assert_called_once() + + +def test_gvr_v2_prefill_temporal_mode_uses_radix(monkeypatch) -> None: + """CUTE_DSL_GVR + gvr_self_sampling=False has no prefill engine -> radix.""" + runner = _install_fake_prefill_runner(monkeypatch) + radix = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_prefill", radix) + top_k = TopK( + 2, + prefill_implementation=TopKImplementation.CUTE_DSL_GVR, + gvr_self_sampling=False, + ) + scores = torch.randn(3, 8) + + _prefill_call(top_k, scores) + + runner.assert_not_called() + radix.assert_called_once() + + +def test_gvr_v2_prefill_rejects_output_width_mismatch(monkeypatch) -> None: + runner = _install_fake_prefill_runner(monkeypatch) + top_k = TopK(2, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) + scores = torch.randn(3, 8) + with pytest.raises(AssertionError): + _prefill_call(top_k, scores, out_width=3) + runner.assert_not_called() diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 3decf9b74303..cc0569e80a7b 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -968,3 +968,316 @@ def test_selfsampling_varlen_heterogeneous_lengths_main(): ref_v = torch.topk(row, k).values.sort().values got = row[out[r].long()].sort().values assert torch.equal(got, ref_v), f"row {r}: value multiset mismatch (n={n_r})" + + +# =========================================================================== +# ==== prefill: per-row [ks, ke) windows (run_prefill) ====================== +# =========================================================================== +# Contract: row r selects the Top-K of logits[r, ks:ke] (ks=row_starts[r], +# ke=row_ends[r], compressed column units), output in the LOCAL frame +# (column - ks) with a trailing -1 pad; nv=ke-ks <= k gives identity 0..nv-1. +# The base is rounded down to a 16B boundary and the <=3 lead lanes are masked, +# so ks % 4 in {1,2,3} (2nd+ request of a multi-request chunk) is exercised +# with poison (+inf/NaN/3e38/-inf) written at [ks-3, ks) and [ke, npad). + + +def _prefill_reference(logits, row_starts, row_ends, top_k): + rows, ncols = logits.shape + out = torch.full((rows, top_k), -1, dtype=torch.int32, device=logits.device) + ks, ke = row_starts.tolist(), row_ends.tolist() + for r in range(rows): + nv = max(min(ke[r], ncols) - ks[r], 0) + if nv == 0: + continue + if nv <= top_k: + out[r, :nv] = torch.arange(nv, dtype=torch.int32, device=logits.device) + else: + out[r] = torch.topk(logits[r, ks[r] : ks[r] + nv], top_k).indices.to(torch.int32) + return out + + +def _check_prefill_exact(logits, got, row_starts, row_ends, top_k): + """Tie-aware radix-parity check: trailing -1 pad from lengths; head unique + and in [0, nv); identity for nv <= k; exact index set when the k-th value + is unique, else strictly-above set + tie-class count (signed zeros and + genuine +/-inf compare like radix). NaN-in-window rows are structure-only.""" + assert got.shape == (logits.shape[0], top_k) and got.dtype == torch.int32 + ks, ke = row_starts.tolist(), row_ends.tolist() + got64 = got.to(torch.int64) + dev = logits.device + for r in range(logits.shape[0]): + nv = max(min(ke[r], logits.shape[1]) - ks[r], 0) + m = min(nv, top_k) + row = got64[r] + assert bool((row[m:] == -1).all()), f"row {r}: pad must be trailing -1 x{top_k - m}" + head = row[:m] + if m == 0: + continue + assert bool((head != -1).all()), f"row {r}: -1 inside the valid head" + assert int(head.min()) >= 0 and int(head.max()) < nv, f"row {r}: index outside [0,{nv})" + assert int(torch.unique(head).numel()) == m, f"row {r}: duplicate indices" + win = logits[r, ks[r] : ks[r] + nv] + if bool(torch.isnan(win).any()): + continue # NaN out of contract for both kernels; structure only + if nv <= top_k: + assert torch.equal(torch.sort(head).values, torch.arange(nv, device=dev)), ( + f"row {r}: short row must be identity" + ) + continue + vals = torch.sort(win, descending=True).values + v_k, v_next = vals[top_k - 1], vals[top_k] + got_vals = win[head] + if bool(v_k != v_next): + ref = (win >= v_k).nonzero(as_tuple=True)[0] + assert ref.numel() == top_k + assert torch.equal(torch.sort(head).values, ref), f"row {r}: index set mismatch" + else: + above = (win > v_k).nonzero(as_tuple=True)[0] + got_above = head[got_vals > v_k] + assert torch.equal(torch.sort(got_above).values, above), ( + f"row {r}: strictly-above set mismatch" + ) + assert int((got_vals == v_k).sum()) == top_k - above.numel(), ( + f"row {r}: wrong number of boundary-tied picks" + ) + assert bool((got_vals >= v_k).all()), f"row {r}: value below k-th selected" + + +def _make_prefill_case(rows, ncols, ks_list, ke_list, *, top_k, seed, dist="randn"): + """DeepGEMM-like storage: stride = align(ncols + 256, 256), column slice + [:, :ncols]; outside-window columns poisoned so an over-read/frame bug is + caught (+inf at [ks-3, ks), rotating NaN/inf/3e38/-inf elsewhere).""" + gen = torch.Generator(device=_DEV).manual_seed(seed) + stride = ((ncols + 256 + 255) // 256) * 256 + if dist == "randn": + full = torch.randn((rows, stride), generator=gen, dtype=torch.float32, device=_DEV) + elif dist == "equal": + full = torch.ones((rows, stride), dtype=torch.float32, device=_DEV) + elif dist == "twoval": + full = torch.randint(0, 2, (rows, stride), generator=gen, device=_DEV).float() + else: + raise ValueError(dist) + logits = full[:, :ncols] + row_starts = torch.tensor(ks_list, dtype=torch.int32, device=_DEV) + row_ends = torch.tensor(ke_list, dtype=torch.int32, device=_DEV) + cols = torch.arange(stride, device=_DEV).unsqueeze(0) + outside = (cols < row_starts.unsqueeze(1)) | (cols >= row_ends.unsqueeze(1)) + pat = torch.tensor([float("nan"), float("inf"), 3e38, float("-inf")], device=_DEV)[ + cols % 4 + ].expand(rows, -1) + full.masked_scatter_(outside, pat[outside]) + for r, ks in enumerate(ks_list): + full[r, max(ks - 3, 0) : ks] = float("inf") + return logits, row_starts, row_ends + + +@pytest.mark.parametrize("top_k", [512, 1024, 2048], ids=lambda k: f"k{k}") +def test_prefill_causal_ramp(top_k): + """Single-request causal ramp (ks=0): a run of short rows then long rows + in one launch straddles the k boundary. Covers nv < k, == k, > k.""" + rows = 148 + ks = [0] * rows + ke = list(range(1, rows + 1)) + lg, rs, re = _make_prefill_case(rows, rows, ks, ke, top_k=top_k, seed=top_k) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, top_k) + + +@pytest.mark.parametrize("lead", [1, 2, 3], ids=lambda x: f"lead{x}") +@pytest.mark.parametrize("top_k", [512, 2048], ids=lambda k: f"k{k}") +def test_prefill_packed_misaligned_ks(top_k, lead): + """Multi-request chunk: request 2 starts at ks % 4 == lead with +inf poison + at [ks-lead, ks). A leaked lead lane would become top-1 (wrong).""" + a = 300 + ks1 = ((a + 3) // 4) * 4 + lead + n1 = 4096 + 17 + rows = a + n1 + ncols = ks1 + n1 + ks = [0] * a + [ks1] * n1 + ke = list(range(1, a + 1)) + [ks1 + n1] * n1 + lg, rs, re = _make_prefill_case(rows, ncols, ks, ke, top_k=top_k, seed=top_k * 100 + lead) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + assert int(out.min()) >= -1, "negative index leaked (missed -lead correction / guard)" + _check_prefill_exact(lg, out, rs, re, top_k) + + +@pytest.mark.parametrize("top_k", [512, 1024], ids=lambda k: f"k{k}") +def test_prefill_short_rows(top_k): + """nv in {0, 1, k-1, k, k+1}: identity 0..nv-1 + trailing -1 (radix short + contract); nv==0 (ks==ke) -> all -1.""" + for nv in (0, 1, top_k - 1, top_k, top_k + 1): + rows = 4 + ncols = max(nv, 1) + 8 + ks = [0] * rows + ke = [nv] * rows + lg, rs, re = _make_prefill_case(rows, ncols, ks, ke, top_k=top_k, seed=nv + top_k) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, top_k) + + +@pytest.mark.parametrize("dist", ["equal", "twoval"], ids=lambda d: d) +@pytest.mark.parametrize("top_k", [512, 1024], ids=lambda k: f"k{k}") +def test_prefill_ties_degenerate(top_k, dist): + """All-equal (whole tie class) and two-valued (massive ties) rows drive the + degenerate A/B narrowing paths; tie-aware acceptance.""" + rows = 16 + n = 4096 + ks = [0] * rows + ke = [n] * rows + lg, rs, re = _make_prefill_case(rows, n, ks, ke, top_k=top_k, seed=top_k, dist=dist) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, top_k) + + +@pytest.mark.parametrize("lead", [1, 2, 3], ids=lambda x: f"lead{x}") +@pytest.mark.parametrize("top_k", [512, 1024], ids=lambda k: f"k{k}") +def test_prefill_neginf_tie_class(top_k, lead): + """nv > k with fewer than k finite values -> the k-th boundary is in the + -inf tie class (degen B), crossed with misaligned lead. A -inf-valued mask + would be emitted here as a negative index (== -lead); assert none leaks.""" + n_finite = top_k - 100 + nv = top_k + 400 + ks1 = ((37 + 3) // 4) * 4 + lead + ncols = ks1 + nv + rows = 5 + gen = torch.Generator(device=_DEV).manual_seed(top_k * 10 + lead) + stride = ((ncols + 256 + 255) // 256) * 256 + full = torch.full((rows, stride), float("-inf"), dtype=torch.float32, device=_DEV) + for r in range(rows): + full[r, ks1 : ks1 + n_finite] = torch.randn(n_finite, generator=gen, device=_DEV) + full[r, :ks1] = float("inf") + full[r, ks1 + nv :] = 3e38 + logits = full[:, :ncols] + rs = torch.tensor([ks1] * rows, dtype=torch.int32, device=_DEV) + re = torch.tensor([ks1 + nv] * rows, dtype=torch.int32, device=_DEV) + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(logits, rs, re, out) + torch.cuda.synchronize() + assert int(out.min()) >= -1, "negative index leaked in a -inf tie class" + _check_prefill_exact(logits, out, rs, re, top_k) + + +@pytest.mark.parametrize("top_k, n", [(512, 4099), (2048, 131075)], ids=lambda v: f"n{v}") +def test_prefill_deepgemm_single_row_odd_width(top_k, n): + """A 1-row tile with an odd num_k_tokens on a DeepGEMM-strided view must + use stride(0) (not shape[1]) and stay exact — the fully-cached follow-up + turn that the varlen 1-row rule would wrongly reject.""" + rows = 1 + ks = [0] + ke = [n] + lg, rs, re = _make_prefill_case(rows, n, ks, ke, top_k=top_k, seed=n) + assert lg.stride(0) % 256 == 0 and lg.shape[1] == n # DeepGEMM-like view + out = torch.full((rows, top_k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, top_k) + + +def test_prefill_slab_over_gridy_limit(): + """> 65535 rows in one call must be slabbed (gridDim.y <= 65535).""" + k = 512 + rows = 70000 + n = 2048 + stride = ((n + 256 + 255) // 256) * 256 + lg = torch.randn((rows, stride), dtype=torch.float32, device=_DEV)[:, :n] + rs = torch.zeros((rows,), dtype=torch.int32, device=_DEV) + re = torch.full((rows,), n, dtype=torch.int32, device=_DEV) + out = torch.full((rows, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out) + torch.cuda.synchronize() + idx = torch.tensor([0, 1, 32767, 32768, 65535, 65536, 69999], device=_DEV) + _check_prefill_exact(lg[idx], out[idx].contiguous(), rs[idx], re[idx], k) + + +def test_prefill_engine_key_distinct_from_decode(): + """The prefill compile shares the DSv3.2 decode varlen tuple (next_n=1, + cr_shift=0) but has a distinct prologue, so the compile keys must differ.""" + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_self_sampling as dev, + ) + + tpl = (256, 8, 4, 256, 2, False, False, 1, 0, 1) + a = dev.get_compiled(tpl, hint_free=True) + b = dev.get_compiled(tpl, hint_free=True, prefill=True) + assert a is not b + + +def test_prefill_guards(): + k = 512 + n = 4096 + stride = ((n + 256 + 255) // 256) * 256 + lg = torch.randn((3, stride), dtype=torch.float32, device=_DEV)[:, :n] + rs = torch.zeros((3,), dtype=torch.int32, device=_DEV) + re = torch.full((3,), n, dtype=torch.int32, device=_DEV) + out = torch.full((3, k), -7, dtype=torch.int32, device=_DEV) + with pytest.raises(RuntimeError, match="float32"): + ss_host.run_prefill(lg.to(torch.bfloat16), rs, re, out) + with pytest.raises(RuntimeError, match="row_starts"): + ss_host.run_prefill(lg, rs.to(torch.int64), re, out) + with pytest.raises(RuntimeError, match="row_starts/row_ends length"): + ss_host.run_prefill(lg, rs[:2], re, out) + with pytest.raises(RuntimeError, match="multiple of 4"): + ss_host.run_prefill(lg, rs, re, torch.full((3, k + 2), -7, dtype=torch.int32, device=_DEV)) + with pytest.raises(RuntimeError, match="16-byte aligned"): + ss_host.run_prefill(lg[:, 1:], rs, re, out) # base offset by 1 float + + +def test_prefill_warmup_idempotent_and_no_rejit(): + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_self_sampling as dev, + ) + + k = 512 + ss_host.warmup_prefill(k, 32768) + before = len(ss_host._PREFILL_WARMUP_DONE) + ss_host.warmup_prefill(k, 32768) + assert len(ss_host._PREFILL_WARMUP_DONE) == before, "warmup not idempotent" + orig = dev.get_compiled + calls = {"n": 0} + + def counting(*a, **kw): + calls["n"] += 1 + return orig(*a, **kw) + + dev.get_compiled = counting + try: + for rows in (1, 8, 37, 74, 100, 296, 297, 4096): + for nkv in (4096, 16384, 32768): + stride = ((nkv + 256 + 255) // 256) * 256 + lg = torch.zeros((rows, stride), dtype=torch.float32, device=_DEV)[:, :nkv] + rs = torch.zeros((rows,), dtype=torch.int32, device=_DEV) + re = torch.full((rows,), nkv, dtype=torch.int32, device=_DEV) + out = torch.empty((rows, k), dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out, max_row_len=nkv) + finally: + dev.get_compiled = orig + assert calls["n"] == 0, f"warmup missed keys: {calls['n']} live compiles" + + +def test_prefill_capture_no_host_sync(): + """After warmup, a run_prefill call captures under a CUDA graph (proves no + .item()/.max() host read).""" + k = 512 + n = 8192 + ss_host.warmup_prefill(k, max(n, 32768)) + stride = ((n + 256 + 255) // 256) * 256 + lg = torch.randn((64, stride), dtype=torch.float32, device=_DEV)[:, :n] + rs = torch.zeros((64,), dtype=torch.int32, device=_DEV) + re = torch.full((64,), n, dtype=torch.int32, device=_DEV) + out = torch.full((64, k), -7, dtype=torch.int32, device=_DEV) + ss_host.run_prefill(lg, rs, re, out, max_row_len=n) # compile outside capture + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + ss_host.run_prefill(lg, rs, re, out, max_row_len=n) + g.replay() + torch.cuda.synchronize() + _check_prefill_exact(lg, out, rs, re, k) From 50db9349a4b801bdb3884cdb1b4fd04c0885652d Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:15:09 +0000 Subject: [PATCH 09/10] [None][feat] add Rubin topology support to GVR V2 prefill Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 5 +- .../blackwell/top_k/__init__.py | 2 +- .../top_k/gvr_topk_decode_self_sampling.py | 2 +- .../gvr_topk_decode_self_sampling_host.py | 147 +++++++--- tensorrt_llm/_torch/modules/top_k.py | 5 +- .../attention/sparse/dsa/test_dsa_indexer.py | 96 ++++++- .../sparse/test_gvr_selfsampling_topology.py | 261 ++++++++++++++++++ tests/unittest/_torch/modules/test_top_k.py | 13 - 8 files changed, 455 insertions(+), 76 deletions(-) create mode 100644 tests/unittest/_torch/attention/sparse/test_gvr_selfsampling_topology.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 9281b5242e8e..06ccc0fad19e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -747,8 +747,9 @@ def __init__( 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}, " + "enable_heuristic_topk=True but the selected DSL GVR " + f"engine ({'self-sampling V2' if self._use_self_sampling_topk else 'temporal V1'}) " + f"is unavailable (cutlass_dsl={IS_CUTLASS_DSL_AVAILABLE}, " f"cutlass_dsl_rubin={IS_CUTLASS_DSL_RUBIN_AVAILABLE}, " f"sm={get_sm_version()}); using the exact radix decode " "top-K instead.", diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py index 25d06bff8bd0..3386089646fe 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CuTE DSL Top-K kernels for Blackwell architecture.""" +"""CuTE DSL Top-K kernels for Blackwell and compatible Rubin paths.""" from .filtered_top_k_decode_varlen import FilteredTopKKernelVarlenDecode from .filtered_top_k_varlen_util import FilteredTopKKernelVarlen diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index cee8b6bde066..2b3aec99a5c9 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Self-sampling GVR top-K decode kernels (CuTe DSL, SM100/103 and SM107). +"""Self-sampling GVR top-K kernels (CuTe DSL, SM100/103 and SM107). SM107/Rubin is enabled only when the installed CuTe DSL exposes its Rubin helpers. This is source support, not an R200 performance-tuning claim. diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py index e13070fda535..ed0a5be992c3 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py @@ -51,6 +51,7 @@ import operator import threading from collections.abc import Sequence +from dataclasses import dataclass import torch @@ -599,6 +600,8 @@ def route_streaming( k: int, force_main: bool = False, num_sms: int = DEFAULT_NUM_SMS, + *, + force_r_one: bool = False, ) -> dict[str, object]: """route() restricted to its STREAMING half (main / clus) — the varlen capture policy: per-row kernels must be picked from the families that are @@ -606,18 +609,22 @@ def route_streaming( skipped even when the envelope n would normally land on them. Where route() itself lands on main/clus this is IDENTICAL to route(). force_main additionally skips the clus rounding, so the raw - min(r1, r2) R matches the CUDA else-branch exactly.""" + min(r1, r2) R matches the CUDA else-branch exactly. ``num_sms`` defaults + to the B200 route constant; the prefill-only caller supplies its active + full-device or locality-domain SM count. ``force_r_one`` is reserved for + prefill, whose correctness contract is one CTA per row.""" if b < 1: raise RuntimeError(f"route_streaming requires b >= 1, got {b}") if num_sms < 1: raise RuntimeError(f"route_streaming requires num_sms >= 1, got {num_sms}") R = 1 - if b <= 32: - r1 = max(num_sms // b, 1) - r2 = max(((n >> 2) + 1023) // 1024, 1) - R = max(min(r1, r2), 1) - elif b <= min(num_sms // 2, MAX_SPLIT_ROWS) and (n >> 2) >= 16384 and k <= 1024: - R = 2 + if not force_r_one: + if b <= 32: + r1 = max(num_sms // b, 1) + r2 = max(((n >> 2) + 1023) // 1024, 1) + R = max(min(r1, r2), 1) + elif b <= min(num_sms // 2, MAX_SPLIT_ROWS) and (n >> 2) >= 16384 and k <= 1024: + R = 2 useclus = False if not force_main and 2 <= R <= 8 and k <= 1024: p2 = 1 @@ -802,21 +809,48 @@ def _available_num_sms(device: torch.device | int) -> int: """Return SMs available to launches in the current execution domain.""" return _execution_domain(device)[0] + # ---- prefill launcher cache ------------------------------------------------ -# Prefill routes always force R==1 (single CTA per row): route_streaming gives -# R>1 only for b<=74, so the representative row counts below (first row of each -# route band) pin R=1 and reduce the engine set to <=6 per k. The launcher -# compiled function depends only on the row TIER, k and the envelope bucket -# (which selects U on the tier-0 1024-thread arm; tiers 1/2 fix U), never on -# the exact row count (arbitrary q-tile / q-split remainders) or npad (a -# runtime scalar), so the cache stays bounded over a long-running server. +# Prefill explicitly forces R==1 (single CTA per row). The compiled function +# depends only on the active-compute topology, row tier, k and the envelope +# bucket (which selects U on the tier-0 1024-thread arm; tiers 1/2 fix U), +# never on the exact row count (arbitrary q-tile / q-split remainders) or npad +# (a runtime scalar), so the cache stays bounded over a long-running server. _PREFILL_CACHE = {} _PREFILL_ROW_SLAB = 32768 # gridDim.y <= 65535; slab so keys stay bounded -_PREFILL_TIER_ROWS = (75, 149, 297) # (rows<=148, 149..296, >296) band reps -def _prefill_tier(rows: int) -> int: - return 0 if rows <= 148 else 1 if rows <= 296 else 2 +@dataclass(frozen=True) +class _PrefillTopology: + """Active compute topology used by the prefill launch policy.""" + + locality_domain_id: int | None + active_num_sms: int + total_num_sms: int + + +def _prefill_topology(device: torch.device | int) -> _PrefillTopology: + """Read the caller's active full-device or locality-domain topology. + + A locality-domain context controls execution placement and has a smaller + SM partition. It does not imply that an arbitrary input allocation is + memory-local; allocation provenance remains the caller's responsibility. + """ + device_index = _device_ordinal(device) + active_num_sms, locality_domain_id = _execution_domain(device_index) + _, total_num_sms = _DEVICE_COMPUTE_INFO[device_index] + return _PrefillTopology(locality_domain_id, active_num_sms, total_num_sms) + + +def _prefill_tier_rows(topology: _PrefillTopology) -> tuple[int, int, int]: + """Return one representative row count for each active-SM wave tier.""" + num_sms = topology.active_num_sms + return (num_sms // 2 + 1, num_sms + 1, 2 * num_sms + 1) + + +def _prefill_tier(rows: int, topology: _PrefillTopology) -> int: + num_sms = topology.active_num_sms + return 0 if rows <= num_sms else 1 if rows <= 2 * num_sms else 2 def _prefill_bucket(n_env: int) -> int: @@ -825,25 +859,33 @@ def _prefill_bucket(n_env: int) -> int: return min(1 << max(int(n_env) - 1, 1).bit_length(), 32768) -def _prefill_cache_key(tier: int, k: int, n_bucket: int): +def _prefill_cache_key(topology: _PrefillTopology, tier: int, k: int, n_bucket: int) -> tuple: # tiers 1/2 fix U, so the bucket does not change their engine — collapse it # to one key so warmup covers them with a single launch. - return (tier, k, n_bucket if tier == 0 else 0) + return (topology, tier, k, n_bucket if tier == 0 else 0) -def _prefill_launcher(tier: int, k: int, n_bucket: int) -> tuple: +def _prefill_launcher(topology: _PrefillTopology, tier: int, k: int, n_bucket: int) -> tuple: """Capture-time prefill plan + compiled launcher (main family, R=1). Mirrors ``_varlen_launcher``'s main branch but with r_const=1, split=False (so tsh_en=0) and the prefill compile flag. SCAP_/CMP_/aim are envelope upper bounds; npad is filled per call in ``run_prefill``.""" - key = _prefill_cache_key(tier, k, n_bucket) + key = _prefill_cache_key(topology, tier, k, n_bucket) hit = _PREFILL_CACHE.get(key) if hit is not None: return hit - b_route = _PREFILL_TIER_ROWS[tier] + b_route = _prefill_tier_rows(topology)[tier] n_route = max(n_bucket, k + 1) - plan = route_streaming(b_route, n_route, n_route, k, force_main=True) + plan = route_streaming( + b_route, + n_route, + n_route, + k, + force_main=True, + num_sms=topology.active_num_sms, + force_r_one=True, + ) if plan["kernel"] != "main": raise RuntimeError(f"prefill route did not land on gvr_main: {plan['kernel']}") rt = plan["rt"] @@ -1765,6 +1807,13 @@ def run_prefill( under capture). Launches in ``<=65535``-row slabs so ``gridDim.y`` never overflows. + Execution locality: launches stay on the caller's current CUDA stream. + When called inside ``LocalityDomainRuntime.partition_context``, the row + tiers use that partition's actual SM count. This only localizes compute; + ``logits`` is memory-local only when its producer/allocation is local. + An explicit ``workspace`` follows the same zero-initialization and + per-in-flight ownership contract as ``run_varlen``. + KNOWN LIMITATION: rows containing NaN inside the window are out of contract (as for the radix reference — both order NaN implementation- specifically). DeepGEMM prefill logits are finite in-window. Trusted @@ -1828,27 +1877,29 @@ def run_prefill( if logits.untyped_storage().size() // 4 < need: raise RuntimeError("logits view storage too small to widen to its row stride") lg = logits.as_strided((num_rows, npad), (npad, 1), logits.storage_offset()) - if workspace is not None: - validate_run_ws(workspace, logits) - ws = kernel_view(workspace) - else: - ws = _ws_hot.get(d) - if ws is None: - ws = default_workspace(logits) n_env = _index(max_row_len) if max_row_len is not None else logits.shape[1] n_env = min(max(n_env, 1), npad) n_bucket = _prefill_bucket(n_env) + topology = _prefill_topology(d) + launchers = [] for r0 in range(0, num_rows, _PREFILL_ROW_SLAB): r1 = min(r0 + _PREFILL_ROW_SLAB, num_rows) - tier = _prefill_tier(r1 - r0) - lc = _PREFILL_CACHE.get(_prefill_cache_key(tier, k, n_bucket)) + tier = _prefill_tier(r1 - r0, topology) + lc = _PREFILL_CACHE.get(_prefill_cache_key(topology, tier, k, n_bucket)) if lc is None: if _is_capturing(): raise RuntimeError( "prefill launcher not compiled for this shape — warm up " "before CUDA graph capture" ) - lc = _prefill_launcher(tier, k, n_bucket) + lc = _prefill_launcher(topology, tier, k, n_bucket) + launchers.append((r0, r1, lc)) + + # Preflight every topology-specific launcher before consulting the + # default workspace. A graph-capture cache miss therefore fails before + # this entry can create a CUDA allocation. + ws = _workspace_for_varlen_launch(logits, workspace, topology.locality_domain_id) + for r0, r1, lc in launchers: _, fn, (scap, cmp_), tail = lc # ABI parity with the varlen main call: pre_idx slot = row_ends, # kv_lens slot = row_starts. The n / SMP / TGT / Q / SS2 / TGT2 launch @@ -2029,22 +2080,25 @@ def warmup_varlen( def warmup_prefill( top_k: int, max_cols: int, - num_rows_list: Sequence[int] = (1, 149, 297), + num_rows_list: Sequence[int] | None = None, row_stride: int | None = None, ) -> None: """TESTING/INIT ONLY — compile the prefill engine set before serving. Six engines per k at most: the tier-0 (1024-thread) arm walks the pow2 envelope buckets (U = 1/2/4/8), tiers 1/2 fix U so one launch each. One - tiny real launch per distinct ``(tier, k, bucket)`` cache key; ``ks=0``, - ``ke=n_env`` (all long rows). ``max_cols`` is the compressed max column - count (``get_indexer_max_seq_len``); the bucket caps at 32768 (U=8 above), - so envelopes past it share one key. The done-key gates only the GPU - launches — the ``_PREFILL_CACHE`` population is idempotent. + tiny real launch per distinct ``(topology, tier, k, bucket)`` cache key; + ``ks=0``, ``ke=n_env`` (all long rows). ``max_cols`` is the compressed max + column count (``get_indexer_max_seq_len``); the bucket caps at 32768 (U=8 + above), so envelopes past it share one key. The done-key gates only the + GPU launches — the ``_PREFILL_CACHE`` population is idempotent. """ dev = torch.cuda.current_device() + topology = _prefill_topology(dev) k = int(top_k) max_cols = int(max_cols) + if num_rows_list is None: + num_rows_list = (1, topology.active_num_sms + 1, 2 * topology.active_num_sms + 1) lo = _prefill_bucket(k + 1) hi = _prefill_bucket(max_cols) buckets = [] @@ -2056,16 +2110,23 @@ def warmup_prefill( buckets = [hi] keys = {} # cache_key -> (tier, bucket) representative for the launch for rows in num_rows_list: - tier = _prefill_tier(int(rows)) + tier = _prefill_tier(int(rows), topology) bset = buckets if tier == 0 else buckets[:1] for bk in bset: - keys.setdefault(_prefill_cache_key(tier, k, bk), (tier, bk)) - done_key = (dev, k, max_cols, tuple(sorted(int(r) for r in num_rows_list)), row_stride) + keys.setdefault(_prefill_cache_key(topology, tier, k, bk), (tier, bk)) + done_key = ( + dev, + topology, + k, + max_cols, + tuple(sorted(int(r) for r in num_rows_list)), + row_stride, + ) with _PREFILL_WARMUP_LOCK: if done_key in _PREFILL_WARMUP_DONE: return for tier, bk in keys.values(): - rows = _PREFILL_TIER_ROWS[tier] + rows = _prefill_tier_rows(topology)[tier] stride = row_stride if row_stride is not None else ((bk + 256 + 255) // 256 * 256) if stride < bk or stride % 4: stride = (max(stride, bk) + 256 + 255) // 256 * 256 diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index def249712121..75df8476be6d 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -187,7 +187,10 @@ def _forward_prefill( ) # ks/ke are already in compressed column units; run_prefill # writes the local (column - ks) frame with -1 pad and no host - # reads (envelope from scores.shape[1]). + # reads (envelope from scores.shape[1]). The runner preserves + # the current stream and, on Rubin, reads an enclosing + # locality-domain context for launch topology only. This call + # does not allocate or claim locality for ``scores``. selfsampling_topk_run_prefill(scores, row_starts, row_ends, output_indices) return output_indices else: 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 5b0b8bfb3ef7..7925ab3bf908 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -176,6 +176,49 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(use_self_samp assert metadata.needs_gvr_prior == (not use_self_sampling) +def test_metadata_sm107_temporal_gvr_does_not_allocate_prior() -> None: + 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=False, + ) + metadata = object.__new__(DSAtrtllmAttentionMetadata) + metadata.sparse_metadata_params = sparse_config.to_sparse_metadata_params() + metadata.kv_cache_manager = SimpleNamespace( + tokens_per_block=256, + compressed_block_sizes={}, + get_cache_indices=Mock(), + ) + metadata.is_cuda_graph = False + 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__" + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.IS_CUTLASS_DSL_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.IS_CUTLASS_DSL_RUBIN_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.get_sm_version", + return_value=107, + ), + ): + DSAtrtllmAttentionMetadata.__post_init__(metadata) + + assert not metadata.use_self_sampling_topk + assert not metadata.needs_gvr_prior + + @pytest.mark.parametrize( "kwargs,expected", [ @@ -186,6 +229,7 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(use_self_samp (dict(is_cute_dsl_available=False), False), (dict(sm_version=107), False), (dict(sm_version=107, is_cute_dsl_rubin_available=True), True), + (dict(sm_version=100, is_cute_dsl_rubin_available=False), True), (dict(sm_version=120), False), (dict(index_topk=256), False), (dict(compress_ratio=2), False), @@ -205,15 +249,30 @@ def test_use_self_sampling_gvr(kwargs, expected): assert use_self_sampling_gvr(**base) is expected -def test_sm107_gvr_admission_is_self_sampling_only(): - """Temporal GVR remains restricted to its validated SM100/103 path.""" - common = dict( - is_cute_dsl_available=True, - is_cute_dsl_rubin_available=True, - sm_version=107, +@pytest.mark.parametrize( + "sm_version,rubin_dsl,use_self_sampling,expected", + [ + (100, False, False, True), + (100, False, True, True), + (103, False, False, True), + (107, False, True, False), + (107, True, True, True), + (107, True, False, False), + (120, True, True, False), + ], +) +def test_gvr_hardware_gate_is_algorithm_specific( + sm_version, rubin_dsl, use_self_sampling, expected +): + assert ( + is_gvr_cute_dsl_supported( + is_cute_dsl_available=True, + is_cute_dsl_rubin_available=rubin_dsl, + sm_version=sm_version, + use_self_sampling_topk=use_self_sampling, + ) + is expected ) - assert is_gvr_cute_dsl_supported(**common, use_self_sampling_topk=True) - assert not is_gvr_cute_dsl_supported(**common, use_self_sampling_topk=False) @pytest.mark.parametrize("use_self_sampling_topk", [True, False]) @@ -311,16 +370,18 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( @pytest.mark.parametrize( - "use_self_sampling,sm_version,msl_c,should_warmup", + "use_self_sampling,sm_version,rubin_dsl,msl_c,should_warmup", [ - (True, 100, 65536, True), - (True, 100, 30001, True), # odd msl_c must not skip the prefill leg - (False, 100, 65536, False), # temporal-hint layers: no prefill engine - (True, 90, 65536, False), # non-datacenter Blackwell + (True, 100, False, 65536, True), + (True, 100, False, 30001, True), # odd msl_c must not skip the prefill leg + (False, 100, False, 65536, False), # temporal-hint layers: no prefill engine + (True, 107, False, 65536, False), # Rubin needs its CuTe DSL helpers + (True, 107, True, 65536, True), + (True, 90, False, 65536, False), ], ) def test_metadata_warmup_selfsampling_prefill_leg( - use_self_sampling, sm_version, msl_c, should_warmup + use_self_sampling, sm_version, rubin_dsl, msl_c, should_warmup ): """The self-sampling warmup drives BOTH the decode (varlen) and the prefill engines; the prefill leg sits before the DeepGEMM decode-stride guard so an @@ -343,6 +404,10 @@ def test_metadata_warmup_selfsampling_prefill_leg( "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.IS_CUTLASS_DSL_AVAILABLE", True, ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.IS_CUTLASS_DSL_RUBIN_AVAILABLE", + rubin_dsl, + ), patch( "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.get_sm_version", return_value=sm_version, @@ -689,7 +754,7 @@ def test_indexer_two_level_gvr_dispatch( @skip_pre_hopper -def test_indexer_sm107_temporal_gvr_falls_back_to_radix(): +def test_indexer_sm107_temporal_gvr_falls_back_to_radix() -> None: """SM107 support is V2-only; requesting temporal V1 must not select GVR.""" sparse_config = DeepSeekSparseAttentionConfig( index_head_dim=128, @@ -717,6 +782,7 @@ def test_indexer_sm107_temporal_gvr_falls_back_to_radix(): assert indexer.top_k.decode_implementation == TopKImplementation.CUDA_RADIX assert not indexer.top_k.gvr_self_sampling + assert indexer.top_k.prefill_implementation == TopKImplementation.CUDA_RADIX assert not indexer.top_k.needs_gvr_prior diff --git a/tests/unittest/_torch/attention/sparse/test_gvr_selfsampling_topology.py b/tests/unittest/_torch/attention/sparse/test_gvr_selfsampling_topology.py new file mode 100644 index 000000000000..876f517064b2 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_gvr_selfsampling_topology.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only topology policy tests for self-sampling GVR prefill.""" + +import importlib.util +import sys +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +_HOST_PATH = ( + Path(__file__).resolve().parents[5] + / "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling_host.py" +) +_HOST_SPEC = importlib.util.spec_from_file_location("_gvr_selfsampling_topology_host", _HOST_PATH) +assert _HOST_SPEC is not None and _HOST_SPEC.loader is not None +ss_host = importlib.util.module_from_spec(_HOST_SPEC) +sys.modules[_HOST_SPEC.name] = ss_host +_HOST_SPEC.loader.exec_module(ss_host) + + +def _topology( + active_num_sms: int, + *, + total_num_sms: int | None = None, + locality_domain_id: int | None = None, +) -> ss_host._PrefillTopology: + return ss_host._PrefillTopology( + locality_domain_id, + active_num_sms, + total_num_sms or active_num_sms, + ) + + +def test_prefill_b200_tier_parity() -> None: + """The parameterized policy reproduces the original 148-SM bands.""" + topology = _topology(148) + assert ss_host._prefill_tier_rows(topology) == (75, 149, 297) + assert [ss_host._prefill_tier(rows, topology) for rows in (1, 148, 149, 296, 297)] == [ + 0, + 0, + 1, + 1, + 2, + ] + + plans = [ + ss_host.route_streaming( + rows, + 32768, + 32768, + 512, + force_main=True, + num_sms=topology.active_num_sms, + force_r_one=True, + ) + for rows in ss_host._prefill_tier_rows(topology) + ] + assert [plan["rt"]["R"] for plan in plans] == [1, 1, 1] + assert [plan["block"] for plan in plans] == [1024, 512, 256] + + +@pytest.mark.parametrize("active_num_sms", [106, 212]) +def test_prefill_tiers_follow_active_sm_count(active_num_sms: int) -> None: + topology = _topology( + active_num_sms, + total_num_sms=212, + locality_domain_id=0 if active_num_sms < 212 else None, + ) + assert ss_host._prefill_tier(active_num_sms, topology) == 0 + assert ss_host._prefill_tier(active_num_sms + 1, topology) == 1 + assert ss_host._prefill_tier(2 * active_num_sms, topology) == 1 + assert ss_host._prefill_tier(2 * active_num_sms + 1, topology) == 2 + + plans = [ + ss_host.route_streaming( + rows, + 262144, + 262144, + 1024, + force_main=True, + num_sms=active_num_sms, + force_r_one=True, + ) + for rows in ss_host._prefill_tier_rows(topology) + ] + assert [plan["rt"]["R"] for plan in plans] == [1, 1, 1] + assert [plan["block"] for plan in plans] == [1024, 512, 256] + + +def test_prefill_topology_keeps_b200_on_full_device(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: SimpleNamespace(multi_processor_count=148, major=10, minor=0), + ) + ss_host._DEVICE_COMPUTE_INFO.clear() + + assert ss_host._prefill_topology(0) == _topology(148) + + +def test_prefill_topology_uses_full_device_outside_locality_domain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: SimpleNamespace(multi_processor_count=212, major=10, minor=7), + ) + monkeypatch.setattr(ss_host, "_current_locality_domain", lambda: None) + ss_host._DEVICE_COMPUTE_INFO.clear() + + assert ss_host._prefill_topology(0) == _topology(212) + + +def test_prefill_topology_uses_current_locality_partition( + monkeypatch: pytest.MonkeyPatch, +) -> None: + device_events: list[tuple[str, int]] = [] + + @contextmanager + def device_context(device: int) -> Iterator[None]: + device_events.append(("enter", device)) + try: + yield + finally: + device_events.append(("exit", device)) + + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: SimpleNamespace(multi_processor_count=212, major=10, minor=7), + ) + monkeypatch.setattr(torch.cuda, "device", device_context) + monkeypatch.setattr(ss_host, "_current_locality_domain", lambda: 1) + + def get_topology() -> tuple[tuple[int, int], ...]: + assert device_events == [("enter", 0)] + return ((106, 212), (106, 212)) + + monkeypatch.setattr(ss_host, "_locality_domain_topology", get_topology) + ss_host._DEVICE_COMPUTE_INFO.clear() + + assert ss_host._prefill_topology(0) == _topology( + 106, + total_num_sms=212, + locality_domain_id=1, + ) + assert device_events == [("enter", 0), ("exit", 0)] + + +def test_prefill_topology_rejects_a_different_device_total( + monkeypatch: pytest.MonkeyPatch, +) -> None: + @contextmanager + def device_context(_device: int) -> Iterator[None]: + yield + + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: SimpleNamespace(multi_processor_count=212, major=10, minor=7), + ) + monkeypatch.setattr(torch.cuda, "device", device_context) + monkeypatch.setattr(ss_host, "_current_locality_domain", lambda: 0) + monkeypatch.setattr(ss_host, "_locality_domain_topology", lambda: ((112, 224), (112, 224))) + ss_host._DEVICE_COMPUTE_INFO.clear() + + with pytest.raises(RuntimeError, match="topology.*target device"): + ss_host._prefill_topology(0) + + +def test_prefill_cache_key_includes_topology_identity() -> None: + full = _topology(212) + partition0 = _topology(106, total_num_sms=212, locality_domain_id=0) + partition1 = _topology(106, total_num_sms=212, locality_domain_id=1) + keys = { + ss_host._prefill_cache_key(topology, 0, 512, 32768) + for topology in (full, partition0, partition1) + } + assert len(keys) == 3 + + +def test_prefill_launcher_cache_is_topology_specific(monkeypatch: pytest.MonkeyPatch) -> None: + class _FakeDevice: + def get_compiled(self, tpl: tuple, **_kwargs: object) -> tuple: + return tpl + + monkeypatch.setattr(ss_host, "_PREFILL_CACHE", {}) + monkeypatch.setattr(ss_host, "_device", lambda: _FakeDevice()) + full = _topology(212) + partition = _topology(106, total_num_sms=212, locality_domain_id=0) + + full_launcher = ss_host._prefill_launcher(full, 1, 512, 32768) + partition_launcher = ss_host._prefill_launcher(partition, 1, 512, 32768) + + assert len(ss_host._PREFILL_CACHE) == 2 + assert full_launcher[1] == partition_launcher[1] + + +def test_prefill_capture_cache_miss_precedes_workspace_allocation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _FakeTensor: + def __init__( + self, + shape: tuple[int, ...], + dtype: object, + *, + strides: tuple[int, ...] | None = None, + ) -> None: + self.shape = shape + self.dtype = dtype + self.is_cuda = True + self._strides = strides or tuple(1 for _ in shape) + + def dim(self) -> int: + return len(self.shape) + + def is_contiguous(self) -> bool: + return True + + def data_ptr(self) -> int: + return 16 + + def stride(self, dim: int) -> int: + return self._strides[dim] + + def get_device(self) -> int: + return 0 + + float32 = object() + int32 = object() + logits = _FakeTensor((1, 1024), float32, strides=(1024, 1)) + row_starts = _FakeTensor((1,), int32) + row_ends = _FakeTensor((1,), int32) + indices = _FakeTensor((1, 4), int32) + allocation_attempted = False + + def allocate_workspace(_logits: _FakeTensor) -> object: + nonlocal allocation_attempted + allocation_attempted = True + return object() + + monkeypatch.setattr(ss_host, "_TENSOR", _FakeTensor) + monkeypatch.setattr(ss_host, "_F32", float32) + monkeypatch.setattr(ss_host, "_I32", int32) + monkeypatch.setattr(ss_host, "_PREFILL_CACHE", {}) + monkeypatch.setattr(ss_host, "_ws_hot", {}) + monkeypatch.setattr(ss_host, "_is_capturing", lambda: True) + monkeypatch.setattr(ss_host, "_prefill_topology", lambda _device: _topology(212)) + monkeypatch.setattr(ss_host, "default_workspace", allocate_workspace) + + with pytest.raises(RuntimeError, match="warm up before CUDA graph capture"): + ss_host.run_prefill(logits, row_starts, row_ends, indices, max_row_len=1024) + + assert not allocation_attempted diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 975df084ecf4..cf69bd18523e 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -425,19 +425,6 @@ def test_cute_dsl_prefill_dispatches_to_blackwell_kernel(monkeypatch) -> None: ) -def test_unsupported_prefill_implementation_raises() -> None: - top_k = TopK(1, prefill_implementation=TopKImplementation.CUTE_DSL_GVR) - - with pytest.raises(NotImplementedError, match="does not support prefill Top-K"): - top_k( - torch.ones(1, 1), - torch.empty(1, 1, dtype=torch.int32), - is_prefill=True, - row_starts=torch.zeros(1, dtype=torch.int32), - row_ends=torch.ones(1, dtype=torch.int32), - ) - - def test_gvr_emission_reset_parks_reused_slots(monkeypatch) -> None: """Cold-started rows must carry non-finite lines. From 406249909590215b210f5b50370a48b32b7ca617 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Mon, 7 Sep 2026 07:24:50 +0000 Subject: [PATCH 10/10] [None][infra] Restore blossom-ci allowlist to match main Stale-base artifact from the merge-from-main baseline, not a Rubin/GVR change; restore so the PR does not drop 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",