Sampling

The PyTorch backend supports a wide variety of features, listed below:

Forward Pass	Sampling Strategies	Sampling Features
No drafting	Greedy	Guided Decoding
Draft target model	TopP	Plugging Logits Post-Processor
Eagle 3	TopK	Temperature
Ngram	TopK + TopP	MinP
	Beam Search	Embedding / Logits Bias
	Best of / n (composable)	Stop criteria
	Rejection sampling (composable)	Return Logits
		Return LogProbs
		TopK LogProbs
		Penalties

General usage

There are two sampling backends available.

• Torch Sampler

• TRTLLM Sampler (deprecated)

Torch Sampler is used by default and supports a superset of features of TRTLLM Sampler. TRTLLM Sampler will be removed in release 1.4. One can specify which sampler to use explicitly with:

from tensorrt_llm import LLM

# Chooses TorchSampler explicitly
llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8',
          sampler_type="TorchSampler")

# Chooses TRTLLMSampler explicitly
llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8',
          sampler_type="TRTLLMSampler")

By default, the sampling backend is chosen to be auto. This will use Torch Sampler for all requests.

Here is an example to run a model with basic usage of sampling parameters. This example prepares two identical prompts which will give different results due to the sampling parameters chosen:

from tensorrt_llm import LLM, SamplingParams
llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8')
sampling_params = SamplingParams(
        temperature=1.0,
        top_k=8,
        top_p=0.5,
    )
llm.generate(["Hello, my name is",
            "Hello, my name is"], sampling_params)

It is also possible to specify different sampling parameters on a per-prompt basis:

from tensorrt_llm import LLM, SamplingParams
llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8')
sampling_params_0 = SamplingParams(
        temperature=1.0,
        top_k=8,
        top_p=0.5,
    )
sampling_params_1 = SamplingParams(
        top_k=4,
    )
llm.generate(["Hello, my name is",
            "Hello, my name is"],
            [sampling_params_0,
            sampling_params_1])

Model generation config defaults

The PyTorch backend can use compatible sampling defaults explicitly specified in a model’s generation_config.json. This behavior is opt-in:

from tensorrt_llm import LLM

llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8',
          generation_config='auto')

For trtllm-serve, enable it on the command line:

trtllm-serve nvidia/Llama-3.1-8B-Instruct-FP8 --generation-config auto

or in the server YAML configuration:

generation_config: auto

The generation_config option has two modes:

• trtllm (default) keeps the TRT-LLM sampling behavior and defaults.

• auto loads supported sampling values from the model’s generation_config.json.

In auto mode, values are resolved in this order:

1. A value explicitly specified by the request.

2. A value explicitly present in generation_config.json.

3. The existing default for the LLM API or serving protocol.

The supported fields are temperature, top_p, top_k, min_p, repetition_penalty, no_repeat_ngram_size, length_penalty, and early_stopping when its value is a boolean or integer. Defaults synthesized by Hugging Face Transformers for fields absent from the JSON file are not applied.

TRT-LLM’s existing model-specific handling of eos_token_id, BART forced_bos_token_id, and Whisper suppression tokens remains active in both modes.

LLM API sampling behavior when using Torch Sampler

• The sampling is controlled via SamplingParams.

• By default (temperature = top_p = top_k = None), greedy sampling is used (unless min-p or top-p decay is active, see below). With generation_config='auto', values explicitly specified in the model’s generation_config.json take the place of these defaults; see Model generation config defaults.

• If either temperature = 0, top_p = 0, top_k = 1, and/or min_p = 1, is specified, sampling is greedy, irrespective of the values of the remaining parameters.

• Otherwise, sampling proceeds according to the specified sampling parameter values and any unspecified parameters default to top_k = 0, top_p = 1, min_p = 0, temperature = 1.0:

  • The logits are scaled by 1/temperature before applying softmax to compute probabilities. Sampling is performed according to these probabilities.

  • If top_k = 0 (or top_k = vocab_size), top_p = 1 and min_p = 0, the output tokens are sampled from the entire vocabulary.

  • If 0 < min_p < 1 is specified, the sampling is restricted to the tokens whose probability is at least min_p times the probability of the most likely token (“min-p sampling”). When combined with top_k and/or top_p, min_p is applied first.

  • If 1 < top_k < vocab_size is specified, the sampling is restricted to the top_k highest-probability tokens.

  • If 0 < top_p < 1.0 is specified, the sampling is further restricted to a minimal subset of highest-probability tokens with total probability greater than top_p (“nucleus sampling”). In particular, the probability of the lowest-probability token in the selected subset is greater or equal than the probability of any not selected token. When combined with top_k, the probabilities of the tokens selected by top_k are rescaled such that they sum to one before top_p is applied.

  • The implementation does not guarantee any particular treatment of tied probabilities.

• Top-P decay is supported: if top_p_decay < 1 is specified, the effective top_p is multiplied by top_p_decay after every sampled token, bounded from below by top_p_min (default 1e-6), and reset to the initial top_p whenever the token top_p_reset_ids is sampled (default -1, which never matches a token). Out-of-range values (top_p_decay or top_p_min outside (0, 1], negative top_p_reset_ids) are rejected.

  • An active top-p decay implies top-p sampling even if top_p is unspecified or top_p = 1 (the initial top_p then defaults to 1). However, explicitly requested greedy sampling (temperature = 0, top_p = 0, and/or top_k = 1) takes precedence over top-p decay.

  • Top-P decay is not supported in combination with beam search or with speculative decoding modes that route draft tokens through the Torch Sampler; such requests are rejected.

• Positive Min-P is not supported in combination with one-model speculative decoding. Such requests are rejected at admission.

• Occurrence penalties are supported: repetition_penalty, presence_penalty and frequency_penalty discourage (or encourage) the model from reusing tokens it has already seen. All three rewrite the logits before temperature scaling, driven by the occurrence history of the prompt plus everything generated so far. Writing c for the number of times a token has occurred in that history:

  • repetition_penalty (default 1.0) rescales the logit of every token with c > 0: the logit is divided by the penalty when it is non-negative and multiplied by it when it is negative. The two branches move a positive and a negative logit the same way, so a value > 1 always pushes a seen token down, and a value < 1 always pulls it up. Must be > 0.

  • presence_penalty (default 0.0) subtracts the penalty itself from every token with c > 0. The amount does not depend on c, so it controls whether a token reappears, not how often.

  • frequency_penalty (default 0.0) subtracts the penalty multiplied by c, so the more often a token has already been produced, the harder it is pushed down.

  • prompt_ignore_length (default 0) excludes the first N prompt tokens from the presence and frequency counts. Those ignored tokens still count for repetition_penalty. Values <= 0 have no effect, and values larger than the prompt are clamped to the prompt length.

  • Occurrence penalties are not supported in combination with beam search; such requests are rejected.

• If no_repeat_ngram_size = n is specified, any token that would recreate an n-gram already present in the sequence (prompt included) is excluded from sampling. None or 0 disables the restriction.

Performance

The Torch Sampler leverages the optimized sampling kernels provided by FlashInfer, which is a required dependency for the Torch Sampler. The sampler also uses the sorting-free implementations whenever possible. This optimization does not compute the complete set of token sampling probabilities (after top-k / top-p masking etc.), which typically can be omitted unless requested by the user or required for speculative decoding (rejection sampling).

Moreover, Torch Sampler internally batches requests with compatible sampling parameters. This 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.

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.

To enable beam search, you must:

1. Enable the use_beam_search option in the SamplingParams object

2. Set the max_beam_width parameter in the LLM class to match the best_of parameter in SamplingParams

Parameter Configuration:

• best_of: Controls the number of beams processed during generation (beam width)

• n: Controls the number of output sequences returned (can be less than best_of)

• If best_of is omitted, the number of beams processed defaults to n

• max_beam_width in the LLM class must equal best_of in SamplingParams

The following example demonstrates beam search with a beam width of 4, returning the top 3 sequences:

from tensorrt_llm import LLM, SamplingParams
llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8',
          max_beam_width=4,   # must equal SamplingParams.best_of
    )
sampling_params = SamplingParams(
        best_of=4,   # must equal LLM.max_beam_width
        use_beam_search=True,
        n=3,         # return top 3 sequences
    )
llm.generate(["Hello, my name is",
            "Hello, my name is"], sampling_params)

Logits processor

Logits processors allow you to modify the logits produced by the network before sampling, enabling custom generation behavior and constraints.

To use a custom logits processor:

1. Create a custom class that inherits from LogitsProcessor and implements the __call__ method

2. Pass an instance of this class to the logits_processor parameter of SamplingParams

The following example demonstrates logits processing:

import torch
from typing import List, Optional

from tensorrt_llm import LLM, SamplingParams
from tensorrt_llm.sampling_params import LogitsProcessor

class MyCustomLogitsProcessor(LogitsProcessor):
    def __call__(self,
        req_id: int,
        logits: torch.Tensor,
        token_ids: List[List[int]],
        stream_ptr: Optional[int],
        client_id: Optional[int]
    ) -> None:
        # Implement your custom inplace logits processing logic
        logits *= logits

llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8')
sampling_params = SamplingParams(
        logits_processor=MyCustomLogitsProcessor()
    )
llm.generate(["Hello, my name is"], sampling_params)

You can find a more detailed example on logits processors here.