From 64cee60c90dab71eb2740d6f8606ad3adcadd288 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/sampling_utils.py | 61 ++++--- .../_torch/speculative/eagle3_dynamic_tree.py | 6 +- tensorrt_llm/_torch/speculative/interface.py | 68 ++++---- 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 | 157 ++++++++++++++++++ 8 files changed, 327 insertions(+), 54 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 72ad4a62ad41..1107cc3e8717 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -129,6 +129,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/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index 20bc59316844..dbb304f0f1eb 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -22,7 +22,7 @@ import sys from collections.abc import Hashable from dataclasses import dataclass -from typing import Any, Literal, Optional, Type, TypeAlias, TypeVar, cast +from typing import TYPE_CHECKING, Any, Literal, Optional, Type, TypeAlias, TypeVar, cast import torch @@ -53,6 +53,9 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.sampling_params import SamplingParams +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import AdvancedSamplingMode + # Ops imported above are re-exported for dependent modules (sampler, drafting # loops, tests). mypy runs in strict mode (no implicit re-export), so they must # be listed here. @@ -931,45 +934,59 @@ 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]) - return flashinfer.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: with the - # divisor clamped to 1.0 by safely_apply_temperature_inplace (order-preserving - # for those rows), 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 <= vanilla.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) - logits = vanilla.safely_apply_temperature_inplace(logits, temperatures) - return flashinfer.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 b988e6366c4a..bf51c59d5d63 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.sampling_utils import sampling_batch_spec_dec_one_model +from ..pyexecutor.sampler.sampling_utils 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 50e23d087ca7..00184655a5bc 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.sampling_utils import ( compute_probs_from_logits, greedy_search_sampling_batch, - sampling_batch_spec_dec_one_model, + resolve_advanced_sampling_filters, sample_from_logits_op, sampling_batch_spec_dec_one_model_for_rejection) @@ -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 d09e0b576b5b..34b03608678b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1688,6 +1688,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.") @@ -1769,6 +1795,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 fb87dbc2b432..b2ef3dad851f 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1635,6 +1635,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..de64564b8ab3 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py @@ -0,0 +1,157 @@ +# 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 import sampling_utils as su +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 = su.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"])