From 4fc76dc94698684f4497fe7a11703c54b31b849d Mon Sep 17 00:00:00 2001 From: Jhao-Ting Chen Date: Fri, 24 Jul 2026 10:53:19 -0700 Subject: [PATCH] [feat] advanced_sampling_mode: skip redundant top-k/top-p filter kernels in the one-model speculative sampler Add a deploy-time `advanced_sampling_mode` enum on `DecodingBaseConfig` (default FULL) that lets the one-model speculative sampler skip flashinfer's top-k mask and top-p renorm kernels when the deployment disables those filters. When a skipped filter is already disabled the output matches FULL bit-for-bit, so this is a lossless throughput optimization for fixed-config / temperature-only deployments (measured ~8-28% OTPS/gpu gain on Qwen3.6-35B-A3B-NVFP4, B200, MTP draft_len=3). - Enum end-to-end: `AdvancedSamplingMode.skips_top_k` / `.skips_top_p` are the single source of truth; pydantic parses the config string into the enum, `SpecMetadata` carries it, and `resolve_advanced_sampling_filters` resolves it once on the speculative side so the ops stay mode-agnostic. - Unified sampler: `sample_from_logits_op` (top_k mask -> softmax -> top_p sample / plain sample) replaces the `sampling_batch_spec_dec_one_model` wrapper; a `None` top_k / top_p skips that filter's kernel. - Greedy rows are handled natively via the `DISABLE_TEMP_VAL` sentinel temperature (softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches with no special-case guard. - All modes work with `use_rejection_sampling` on or off; the two are independent. - Regenerated `tensorrt_llm/usage/llm_args_golden_manifest.json` (adds `speculative_config.advanced_sampling_mode`). - Unit tests in `tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` and docs in `docs/source/features/sampling.md`. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jhao-Ting Chen --- docs/source/features/sampling.md | 43 +++++ .../pyexecutor/sampler/ops/flashinfer.py | 61 ++++--- .../_torch/speculative/eagle3_dynamic_tree.py | 6 +- tensorrt_llm/_torch/speculative/interface.py | 70 ++++---- tensorrt_llm/_torch/speculative/utils.py | 1 + tensorrt_llm/llmapi/llm_args.py | 33 ++++ .../usage/llm_args_golden_manifest.json | 12 ++ .../test_advanced_sampling_mode.py | 158 ++++++++++++++++++ 8 files changed, 328 insertions(+), 56 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 29f75abde6ba..1e526f16c004 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -133,6 +133,49 @@ Moreover, Torch Sampler internally batches requests with compatible sampling par can greatly reduce the overall latency of the sampling step when request batches are comprised of requests with very heterogeneous sampling strategies (e.g. a mix of requests using greedy and top-p-after-top-k sampling). +## Advanced sampling mode (speculative decoding) + +For one-model speculative decoding (e.g. MTP-Eagle one-model), the per-request +advanced sampler applies a `top_k` mask, a temperature softmax, and a `top_p` +filter before sampling each draft/target token. When a deployment fixes its +sampling configuration such that a filter is always disabled (`top_k = 0` / +`top_k = vocab_size`, or `top_p = 1`), that filter's kernel is pure overhead. + +`advanced_sampling_mode` (on `DecodingBaseConfig`, so it is available to any +speculative config) lets you skip those redundant kernels for a fixed deploy +config. The output is identical to `FULL` whenever the skipped filter is already +disabled, so this is a lossless throughput optimization for advanced use cases: + +| Mode | `top_k` kernel | `top_p` kernel | +|---|---|---| +| `full` (default) | applied | applied | +| `no_topk` | **skipped** | applied | +| `no_topp` | applied | **skipped** | +| `no_topk_no_topp` | **skipped** | **skipped** | + +Notes: + +* `full` is the default and always safe; the specialization is opt-in. +* `advanced_sampling_mode` and `use_rejection_sampling` are independent: every mode + works with rejection sampling on or off; the flag no longer gates the mode choice. +* `no_topp` and `no_topk_no_topp` disable `top_p`, switching the sampler from the + fused `top_p_sampling_from_probs` to the cheaper `sampling_from_probs`; `no_topk` + keeps `top_p`. +* Greedy requests are handled natively (via a sentinel temperature that makes the + softmax collapse to a one-hot argmax), so any mode supports mixed greedy + + sampling batches without a special case. +* `advanced_sampling_mode` is a deploy-time choice; it is *not* part of the CUDA + graph key, so it adds no extra warmup graphs. + +```python +from tensorrt_llm.llmapi import MTPDecodingConfig + +spec_config = MTPDecodingConfig( + max_draft_len=3, + advanced_sampling_mode="no_topk_no_topp", # temperature-only deploy config +) +``` + ## Beam search Beam search is a decoding strategy that maintains multiple candidate sequences (beams) during text generation, exploring different possible continuations to find higher quality outputs. Unlike greedy decoding or sampling, beam search considers multiple hypotheses simultaneously. diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index 572b6166034b..89c013cb568e 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py @@ -26,13 +26,14 @@ instead of a per-call break from tracing flashinfer's lazy JIT bootstrap. """ -from typing import Any, Callable, Optional, TypeVar, Union, cast +from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union, cast import torch from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl -from .vanilla import GREEDY_TEMPERATURE_THRESHOLD +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode _OpT = TypeVar("_OpT", bound=Callable[..., Any]) @@ -260,7 +261,8 @@ def compute_probs_from_logits( """Compute filtered+normalized probs via flashinfer (hard dependency). ``temperatures``, ``top_k``, ``top_p`` are per-request tensors matching the - spec-decoding call site in interface.py. + spec-decoding call site in interface.py. A ``None`` top_k / top_p skips that + filter's kernel. """ if top_k is not None: top_k = sanitize_top_k(top_k, logits.shape[-1]) @@ -268,38 +270,51 @@ def compute_probs_from_logits( return compute_probs_from_logits_op(logits, temperatures, top_k, top_p) +def resolve_advanced_sampling_filters( + advanced_sampling_mode: "AdvancedSamplingMode", + top_k: Optional[torch.Tensor], + top_p: Optional[torch.Tensor], +) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Resolve advanced_sampling_mode to effective (top_k, top_p) tensors. + + A filter the mode disables (via ``skips_top_k`` / ``skips_top_p``), or one already + None, becomes None so the downstream op skips that kernel; kept filters pass through + unchanged (the op sanitizes top_k internally). + """ + eff_top_k = None if advanced_sampling_mode.skips_top_k or top_k is None else top_k + eff_top_p = None if advanced_sampling_mode.skips_top_p or top_p is None else top_p + return eff_top_k, eff_top_p + + @torch.compile(options={"max-autotune": True}) -def sampling_batch_spec_dec_one_model( +def sample_from_logits_op( logits: torch.Tensor, temperatures: torch.Tensor, - top_k: torch.Tensor, - top_p: torch.Tensor, + top_k: Optional[torch.Tensor] = None, + top_p: Optional[torch.Tensor] = None, seed: Optional[torch.Tensor] = None, offset: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """CUDA-graph compatible sampling; supports mixed sampling params. Returns sampled tokens.""" - top_k = sanitize_top_k(top_k, logits.shape[-1]) - # Greedy rows (temperature <= threshold) reduce to top_k=1 sampling: their - # divisor is clamped to 1.0 below (order-preserving for those rows), so - # flashinfer deterministically returns the max-probability token, i.e. the - # argmax of the original logits. All ops remain branch-free (no - # data-dependent control flow), so this stays CUDA-graph safe. - is_greedy = temperatures <= GREEDY_TEMPERATURE_THRESHOLD - top_k = torch.where(is_greedy, torch.ones_like(top_k), top_k) - top_p = torch.where(is_greedy, torch.ones_like(top_p), top_p) - # Dividing by a greedy row's temperature (0 / <= threshold) would blow the - # logits up to inf/nan; clamp those rows to 1.0 to keep the division safe. - safe_temp = torch.where(is_greedy, torch.ones_like(temperatures), temperatures) - logits = logits.div_(safe_temp.unsqueeze(dim=1)) - return top_k_top_p_sampling_from_logits_op(logits, top_k, top_p, seed=seed, offset=offset) + """CUDA-graph compatible one-model sampler; returns sampled tokens. + + ``top_k`` / ``top_p`` are None when the caller's advanced_sampling_mode disables that + filter. + """ + if top_k is not None: + top_k = sanitize_top_k(top_k, logits.shape[-1]) + logits = top_k_mask_logits_op(logits, top_k) + probs = softmax_op(logits, temperatures) + if top_p is not None: + return top_p_sampling_from_probs_op(probs, top_p, seed=seed, offset=offset) + return sampling_from_probs_op(probs, seed=seed, offset=offset) @torch.compile(options={"max-autotune": True}) def sampling_batch_spec_dec_one_model_for_rejection( logits: torch.Tensor, temperatures: torch.Tensor, - top_k: torch.Tensor, - top_p: torch.Tensor, + top_k: Optional[torch.Tensor], + top_p: Optional[torch.Tensor], seed: Optional[torch.Tensor] = None, offset: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index f7db871c856e..f8b68fc67a3c 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -24,7 +24,7 @@ from tensorrt_llm._utils import get_sm_version, nvtx_range from ..attention_backend import AttentionMetadata -from ..pyexecutor.sampler.ops.flashinfer import sampling_batch_spec_dec_one_model +from ..pyexecutor.sampler.ops.flashinfer import sample_from_logits_op from .eagle3 import Eagle3OneModelWorker if TYPE_CHECKING: @@ -772,7 +772,7 @@ def _sample_and_accept_dynamic_tree( self.offset = torch.tensor([0], dtype=torch.int64, device=logits.device) self.seed.add_(1).remainder_(2**31) top_ks = spec_metadata.top_ks[:num_flat_tokens] - sampled = sampling_batch_spec_dec_one_model( + sampled = sample_from_logits_op( logits, spec_metadata.temperatures[:num_flat_tokens], top_ks, @@ -880,7 +880,7 @@ def _sample_and_accept_dynamic_tree_rejection( # Context tokens bypass the rejection kernel — sample them directly. if num_contexts > 0: top_ks_ctx = spec_metadata.top_ks[:num_contexts] - sampled_ctx = sampling_batch_spec_dec_one_model( + sampled_ctx = sample_from_logits_op( logits[:num_contexts], spec_metadata.temperatures[:num_contexts], top_ks_ctx, diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 5b7129b0395b..e680036e3c54 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -42,9 +42,11 @@ if IS_FLASHINFER_AVAILABLE: import flashinfer +from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode + from ..pyexecutor.sampler.ops.flashinfer import ( - compute_probs_from_logits, sampling_batch_spec_dec_one_model, - sampling_batch_spec_dec_one_model_for_rejection) + compute_probs_from_logits, resolve_advanced_sampling_filters, + sample_from_logits_op, sampling_batch_spec_dec_one_model_for_rejection) from ..pyexecutor.sampler.ops.vanilla import greedy_search_sampling_batch @@ -538,6 +540,8 @@ class SpecMetadata: group_all_greedy_sample: Optional[bool] = None # Whether to use rejection sampling for one-model speculative decoding. use_rejection_sampling: bool = False + # Advanced-sampling specialization (deploy-time; from DecodingBaseConfig.advanced_sampling_mode). + advanced_sampling_mode: AdvancedSamplingMode = AdvancedSamplingMode.FULL # Sampling parameters for non-greedy sampling (per-request) temperatures: Optional[torch.Tensor] = None top_ks: Optional[torch.Tensor] = None @@ -847,7 +851,6 @@ def _normalize_request_sampling_params( # resurrect the local value. if self.group_all_greedy_sample is not None: self.is_all_greedy_sample = self.group_all_greedy_sample - return per_request_normalized, per_request_slot_ids @property @@ -1009,7 +1012,7 @@ def __init__(self, use_separate_draft_kv_cache: bool = False): self.force_num_accepted_tokens: float = get_force_num_accepted_tokens_float( ) # One-model speculative sampling goes through flashinfer unconditionally - # (sampling_batch_spec_dec_one_model), so flashinfer>=0.6.4 is a hard + # (sample_from_logits_op), so flashinfer>=0.6.4 is a hard # dependency here. Fail at construction with a clear error instead of # crashing mid-inference on the first non-greedy sampling step. if not IS_FLASHINFER_AVAILABLE or Version( @@ -1503,7 +1506,7 @@ def advanced_sample_draft(self, With rejection enabled and a ``draft_step``, samples via ``sampling_batch_spec_dec_one_model_for_rejection`` and scatters this step's proposal distribution into the slot-indexed ``draft_probs`` - buffer; otherwise uses ``sampling_batch_spec_dec_one_model`` (tokens + buffer; otherwise uses ``sample_from_logits_op`` (tokens only). Returns tokens in draft-vocab space (the caller applies d2t). Expects 2D ``[batch_size, vocab]`` logits (one row per request). """ @@ -1512,13 +1515,15 @@ def advanced_sample_draft(self, top_ps = spec_metadata.request_top_ps[:batch_size] self._update_advance_draft_sampling_seed(logits.device) + eff_top_ks, eff_top_ps = resolve_advanced_sampling_filters( + spec_metadata.advanced_sampling_mode, top_ks, top_ps) if spec_metadata.use_rejection_sampling and draft_step is not None: draft_tokens, probs = ( sampling_batch_spec_dec_one_model_for_rejection( logits, temperatures, - top_ks, - top_ps, + eff_top_ks, + eff_top_ps, seed=self.seed, offset=self.offset)) # Scatter probs into the slot-indexed buffer so each request's data @@ -1532,12 +1537,12 @@ def advanced_sample_draft(self, spec_metadata.draft_probs[batch_slots, draft_step, :vocab] = probs spec_metadata.draft_probs_last_dim = vocab else: - draft_tokens = sampling_batch_spec_dec_one_model(logits, - temperatures, - top_ks, - top_ps, - seed=self.seed, - offset=self.offset) + draft_tokens = sample_from_logits_op(logits, + temperatures, + eff_top_ks, + eff_top_ps, + seed=self.seed, + offset=self.offset) return draft_tokens.type(torch.int32) @@ -1687,6 +1692,8 @@ def _sample_and_accept_draft_tokens_rejection( spec_metadata.top_ks[gen_start:gen_end]) top_ps = (None if spec_metadata.skip_top_p else spec_metadata.top_ps[gen_start:gen_end]) + top_ks, top_ps = resolve_advanced_sampling_filters( + spec_metadata.advanced_sampling_mode, top_ks, top_ps) target_probs_flat = compute_probs_from_logits( gen_logits, temperatures, top_ks, top_ps) @@ -1894,7 +1901,7 @@ def advanced_sample_draft_block(self, gen_logits: torch.Tensor, With rejection enabled, samples via ``sampling_batch_spec_dec_one_model_for_rejection`` and scatters the K proposal rows into ``draft_probs[gen_slot_ids, 0:K, :]``; otherwise uses - ``sampling_batch_spec_dec_one_model`` (tokens only). Only called for a + ``sample_from_logits_op`` (tokens only). Only called for a non-greedy batch (the all-greedy path is handled by the caller). Returns ``[num_gens, K]`` int32 tokens in draft-vocab space (the caller applies d2t); stored probs likewise stay in draft-vocab space. @@ -1916,14 +1923,16 @@ def advanced_sample_draft_block(self, gen_logits: torch.Tensor, self._update_advance_draft_sampling_seed(gen_logits.device) flat_logits = gen_logits.reshape(num_gens * K, vocab) + eff_top_ks, eff_top_ps = resolve_advanced_sampling_filters( + spec_metadata.advanced_sampling_mode, top_ks, top_ps) if getattr(spec_metadata, "use_rejection_sampling", False): flat_tokens, flat_probs = ( sampling_batch_spec_dec_one_model_for_rejection( flat_logits, temps, - top_ks, - top_ps, + eff_top_ks, + eff_top_ps, seed=self.seed, offset=self.offset)) # Scatter the K prob rows per gen request into its stable slot row. @@ -1937,12 +1946,12 @@ def advanced_sample_draft_block(self, gen_logits: torch.Tensor, spec_metadata.draft_probs[gen_slot_ids, :K, :vocab] = probs spec_metadata.draft_probs_last_dim = vocab else: - flat_tokens = sampling_batch_spec_dec_one_model(flat_logits, - temps, - top_ks, - top_ps, - seed=self.seed, - offset=self.offset) + flat_tokens = sample_from_logits_op(flat_logits, + temps, + eff_top_ks, + eff_top_ps, + seed=self.seed, + offset=self.offset) return flat_tokens.reshape(num_gens, K).type(torch.int32) @@ -2189,7 +2198,7 @@ def _sample_tokens_for_batch( # Use logits.shape[0] directly: for PARD under CUDA graph capture # runtime_draft_len may reflect the PARD-max while the captured # graph was built for a shorter draft_len, causing a shape mismatch - # in sampling_batch_spec_dec_one_model (which is torch.compiled). + # in sample_from_logits_op (which is torch.compiled). num_tokens = logits.shape[0] temperatures = spec_metadata.temperatures[:num_tokens] @@ -2207,13 +2216,14 @@ def _sample_tokens_for_batch( self.seed += 1 self.seed %= (2**31) - sampled_tokens = sampling_batch_spec_dec_one_model( - logits, - temperatures, - top_ks, - top_ps, - seed=self.seed, - offset=self.offset) + eff_top_ks, eff_top_ps = resolve_advanced_sampling_filters( + spec_metadata.advanced_sampling_mode, top_ks, top_ps) + sampled_tokens = sample_from_logits_op(logits, + temperatures, + eff_top_ks, + eff_top_ps, + seed=self.seed, + offset=self.offset) else: sampled_tokens = torch.argmax(logits, dim=-1) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 62189035cafb..bffa8833058c 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -111,6 +111,7 @@ def get_spec_metadata(spec_config, hidden_size=model_config.hidden_size, max_num_tokens=max_num_tokens, use_rejection_sampling=use_rejection_sampling, + advanced_sampling_mode=spec_config.advanced_sampling_mode, vocab_size=vocab_size, num_seq_slots=num_seq_slots, draft_vocab_size=draft_vocab_size, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 68d9fdbeb207..06edf8a616df 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1687,6 +1687,32 @@ class CalibConfig(StrictBaseModel): "The maximum sequence length to initialize tokenizer for calibration.") +class AdvancedSamplingMode(StrEnum): + """Deploy-time specialization of the one-model advanced sampler. + + FULL - per-row tensor top_k/top_p (default; mixed per-request sampling). + NO_TOPK - top_k disabled, top_p honored. Skips the top_k mask kernel. + NO_TOPP - top_p disabled, top_k honored. Skips the top_p renorm kernel. + NO_TOPK_NO_TOPP - both disabled (pure temperature sampling). Skips both kernels. + """ + FULL = "full" + NO_TOPK = "no_topk" + NO_TOPP = "no_topp" + NO_TOPK_NO_TOPP = "no_topk_no_topp" + + @property + def skips_top_k(self) -> bool: + """Single source of truth: does this mode disable the top_k filter?""" + return self in (AdvancedSamplingMode.NO_TOPK, + AdvancedSamplingMode.NO_TOPK_NO_TOPP) + + @property + def skips_top_p(self) -> bool: + """Single source of truth: does this mode disable the top_p filter?""" + return self in (AdvancedSamplingMode.NO_TOPP, + AdvancedSamplingMode.NO_TOPK_NO_TOPP) + + class DecodingBaseConfig(StrictBaseModel): max_draft_len: Optional[NonNegativeInt] = Field( default=None, description="The maximum number of draft tokens.") @@ -1768,6 +1794,13 @@ class DecodingBaseConfig(StrictBaseModel): "in a future release. Non-greedy sampling is now auto-detected per " "request; this flag no longer has any effect.") + advanced_sampling_mode: AdvancedSamplingMode = Field( + default=AdvancedSamplingMode.FULL, + description= + "Deploy-time specialization of the one-model advanced sampler that skips disabled " + "filter kernels. FULL (default): per-row top_k/top_p. NO_TOPK: skip top_k. " + "NO_TOPP: skip top_p. NO_TOPK_NO_TOPP: skip both.") + # If set, drafting is allowed to use chain drafter. _allow_chain_drafter: bool = PrivateAttr(True) # If set, drafting uses greedy sampling, irrespective of sampling parameters. diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f20e02169d62..42a8dd01cb3c 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1636,6 +1636,18 @@ "kind": "value", "path": "speculative_config.acceptance_rate_window_size" }, + { + "allowed_values": [ + "full", + "no_topk", + "no_topp", + "no_topk_no_topp" + ], + "annotation": "", + "converter": "", + "kind": "categorical", + "path": "speculative_config.advanced_sampling_mode" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py b/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py new file mode 100644 index 000000000000..9c760526ff79 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py @@ -0,0 +1,158 @@ +# Copyright (c) 2025-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. +"""Unit tests for one-model ``advanced_sampling_mode``. + +Covers the config contract (enum + skip properties + the use_rejection_sampling +requirement for the top-p-disabling modes), ``resolve_advanced_sampling_filters`` +mode resolution, a CUDA check that NO_TOPK yields the same distribution as FULL +when top_k is disabled, and native greedy handling (greedy rows return argmax). +""" + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.sampler.ops import flashinfer as su +from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import GREEDY_TEMPERATURE_THRESHOLD +from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode, DecodingBaseConfig, MTPDecodingConfig + + +def test_enum_skip_properties(): + """Enum members + the skip properties (single source of truth for filter skipping).""" + M = AdvancedSamplingMode + assert [m.value for m in M] == ["full", "no_topk", "no_topp", "no_topk_no_topp"] + assert (M.FULL.skips_top_k, M.FULL.skips_top_p) == (False, False) + assert (M.NO_TOPK.skips_top_k, M.NO_TOPK.skips_top_p) == (True, False) + assert (M.NO_TOPP.skips_top_k, M.NO_TOPP.skips_top_p) == (False, True) + assert (M.NO_TOPK_NO_TOPP.skips_top_k, M.NO_TOPK_NO_TOPP.skips_top_p) == (True, True) + + +def test_advanced_sampling_mode_on_base_config(): + """The field lives on DecodingBaseConfig (not MTP-specific) and defaults to FULL.""" + assert "advanced_sampling_mode" in DecodingBaseConfig.model_fields + assert MTPDecodingConfig(max_draft_len=1).advanced_sampling_mode == AdvancedSamplingMode.FULL + + +def test_all_modes_construct_regardless_of_rejection(): + """Every mode constructs with or without rejection sampling (no config gating).""" + for mode in ("full", "no_topk", "no_topp", "no_topk_no_topp"): + for rej in (False, True): + cfg = MTPDecodingConfig( + max_draft_len=1, advanced_sampling_mode=mode, use_rejection_sampling=rej + ) + assert cfg.advanced_sampling_mode.value == mode + + +@pytest.mark.parametrize( + "mode,expect_top_k_none,expect_top_p_none", + [ + ("full", False, False), + ("no_topk", True, False), + ("no_topp", False, True), + ("no_topk_no_topp", True, True), + ], +) +def test_resolve_advanced_sampling_filters(mode, expect_top_k_none, expect_top_p_none): + """Mode resolution None-ifies disabled filters (so the op skips that kernel) + and passes kept filters through unchanged.""" + top_k = torch.zeros(2, dtype=torch.int32) + top_p = torch.ones(2) + eff_top_k, eff_top_p = su.resolve_advanced_sampling_filters( + AdvancedSamplingMode(mode), top_k, top_p + ) + assert (eff_top_k is None) is expect_top_k_none + assert (eff_top_p is None) is expect_top_p_none + if not expect_top_k_none: + assert eff_top_k is top_k + if not expect_top_p_none: + assert eff_top_p is top_p + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA + flashinfer sampling kernels" +) +@pytest.mark.parametrize("top_p_val", [1.0, 0.9]) +def test_no_topk_matches_full(top_p_val): + """With top_k disabled, NO_TOPK skips the top_k mask kernel (a no-op at k=vocab) + and yields the same sampling distribution as FULL. We compare the resulting + probability distributions rather than the sampled tokens: the flashinfer top_k mask + at k=vocab injects ~1e-8 fp noise that leaves the distribution unchanged but can flip + an individual sampled token across GPU archs, so exact-token equality is not portable. + A real (non-no-op) filter would move mass by orders of magnitude, far above atol.""" + dev = "cuda" + torch.manual_seed(0) + batch, vocab = 64, 32000 + logits = torch.randn(batch, vocab, device=dev, dtype=torch.float32) * 2.0 + temperatures = torch.full((batch,), 0.7, device=dev, dtype=torch.float32) + top_k = torch.zeros(batch, device=dev, dtype=torch.int32) # disabled + top_p = torch.full((batch,), top_p_val, device=dev, dtype=torch.float32) + + ek_full, ep_full = su.resolve_advanced_sampling_filters( + AdvancedSamplingMode.FULL, top_k.clone(), top_p + ) + ek_nt, ep_nt = su.resolve_advanced_sampling_filters( + AdvancedSamplingMode.NO_TOPK, top_k.clone(), top_p + ) + probs_full = su.compute_probs_from_logits(logits.clone(), temperatures, ek_full, ep_full) + probs_no_topk = su.compute_probs_from_logits(logits.clone(), temperatures, ek_nt, ep_nt) + assert torch.allclose(probs_full, probs_no_topk, atol=1e-5, rtol=0) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA + flashinfer sampling kernels" +) +@pytest.mark.parametrize("mode", ["no_topk", "no_topk_no_topp"]) +def test_greedy_row_returns_argmax_natively(mode): + """Greedy rows carry the sentinel temperature, so the sampler returns their + argmax token even in a mixed batch -- this is why no mixed-batch guard is needed.""" + dev = "cuda" + torch.manual_seed(0) + batch, vocab = 8, 4096 + logits = torch.randn(batch, vocab, device=dev, dtype=torch.float32) * 3.0 + disable = GREEDY_TEMPERATURE_THRESHOLD / 10 # sentinel for greedy rows + temperatures = torch.full((batch,), 0.7, device=dev, dtype=torch.float32) + temperatures[0] = disable # greedy rows mixed with sampled rows + temperatures[1] = disable + top_k = torch.zeros(batch, device=dev, dtype=torch.int32) + top_p = torch.ones(batch, device=dev, dtype=torch.float32) + seed = torch.tensor([7], dtype=torch.int64, device=dev) + offset = torch.tensor([0], dtype=torch.int64, device=dev) + + eff_top_k, eff_top_p = su.resolve_advanced_sampling_filters( + AdvancedSamplingMode(mode), top_k, top_p + ) + tokens = su.sample_from_logits_op( + logits, temperatures, eff_top_k, eff_top_p, seed=seed, offset=offset + ) + argmax = logits.argmax(dim=-1) + assert tokens[0].item() == argmax[0].item() + assert tokens[1].item() == argmax[1].item() + + +def test_advanced_mode_accepted_on_all_spec_paths(): + """The MTP-one-model-only gate was removed (the field is on the base config), + so non-FULL modes construct on any spec path instead of raising at config time.""" + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=MTPDecodingConfig( + max_draft_len=1, use_mtp_vanilla=True, advanced_sampling_mode="no_topk" + ), + ) + assert args.speculative_config.advanced_sampling_mode == AdvancedSamplingMode.NO_TOPK + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])