Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jhaotingc wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions docs/source/features/sampling.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
61 changes: 39 additions & 22 deletions tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)

@zhaoyangwang-nvidia zhaoyangwang-nvidia Jul 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please following compute_probs_from_logits_op, add a sample_from_logits_op; resolve advanced_sampling_mode once on the speculative side and call it directly, removing the sampling_batch_spec_dec_one_model wrapper layer.


@torch.compile(options={"max-autotune": True})
def sample_from_logits_op(
    logits: torch.Tensor,
    temperatures: 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:
    if top_k is not None:
        logits = flashinfer.sampling.top_k_mask_logits(logits, top_k)
    probs = flashinfer.sampling.softmax(
        logits, temperatures, enable_pdl=get_env_enable_pdl())
    if top_p is not None:
        return flashinfer.sampling.top_p_sampling_from_probs(
            probs, top_p, seed=seed, offset=offset)
    return flashinfer.sampling.sampling_from_probs(probs, seed=seed, offset=offset)

@jhaotingc jhaotingc Jul 21, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@zhaoyangwang-nvidia I was having concern the minute diff when user turn on no_topp for non-rejection path.

In your suggestion, if user decide to turn on no_topp and pass in top_p 1.0, the kernel will be SamplingFromProbKernel instead of TopPSamplingFromProbKernel. The two kernels handle sorting differently.

Before #15775 where it removes drafter sampling for default case, the diff is more significant. So I disabled no_topp use case when rejection sampling False to make it consistent. The no_rejection sampling path is already having TopP kernel, switching to sampling from prob does not help much because 1 kernel still needs to run. The rejection path already has TopK/TopP added by you so I reused the same logic for no_topk_no_topp.

If you're ok with the behavior changes for turning on no_topp in no_rejection path, now that draft step sampling is removed, I'm ok with it as well.

@jhaotingc jhaotingc Jul 21, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, another reason of not able to discard this function is this line. This line makes ALL full path having one GREEDY result calculated and one sampled result calculated, just to make sure we take care of when temperatures <= vanilla.GREEDY_TEMPERATURE_THRESHOLD

Using this function, you suggested changes the behavior completely, and the same issue isn't considered in #15775

@jhaotingc jhaotingc Jul 21, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#15542 seems to add this greedy & advanced together, is this intended? should this be kept?

@jhaotingc jhaotingc Jul 21, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm ok didn't see the comments here. I'll change the behavior, thanks.

https://github.com/NVIDIA/TensorRT-LLM/pull/16561/changes/BASE..c8eb118b7afd05bf6055408baf4bc47e21bb5ae0#r3619819010

@zhaoyangwang-nvidia zhaoyangwang-nvidia Jul 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for the late reply — timezone difference on my side.

Both behavior differences you raised are valid observations:

SamplingFromProbKernel vs TopPSamplingFromProbKernel producing different tokens for the same distribution when no_topp is enabled;
the greedy rows no longer going through an explicit argmax + torch.where selection, but instead being handled natively via the DISABLE_TEMP_VAL sentinel temperature.
That said, I'd argue both are correct behaviors that differ only by floating-point accumulation order / kernel implementation details, not accuracy bugs. In case 1, top_p=1.0 means no filtering, so any properly implemented sampling kernel is a valid realization of the same distribution. In case 2, softmax(logits / DISABLE_TEMP_VAL) collapses to a one-hot at the argmax, so sampling from it returns the argmax token — the greedy semantics are preserved.

Do you have any real cases where these differences caused an actual accuracy problem (e.g., a measurable eval-score regression, or a rejection-rate change)? If so, let's discuss further and I'm happy to reconsider. If not, I'd prefer to keep the code simple and unified — one mode-agnostic sampling op with the filters resolved once on the speculative side — rather than maintaining separate paths just for bit-for-bit token equivalence, which we don't guarantee across versions anyway.

@jhaotingc jhaotingc Jul 22, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you! Changed to what's suggested on comment, will click resolve before merge.

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]:
Expand Down
6 changes: 3 additions & 3 deletions tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
68 changes: 39 additions & 29 deletions tensorrt_llm/_torch/speculative/interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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).
"""
Expand All @@ -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
Expand All @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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)

Expand Down Expand Up @@ -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]
Expand All @@ -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)

Expand Down
1 change: 1 addition & 0 deletions tensorrt_llm/_torch/speculative/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading
Loading