[None][feat] MTP one-model advanced_sampling_mode: skip redundant top-k / top-p filter kernels with additional config enum

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.

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,46 +261,60 @@ 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 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]:

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,

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)
 

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,