7.9 Sampling Strategies & Decoding Algorithms¶

Every token an LLM produces is chosen by a decoding algorithm — a procedure that converts the raw logit vector from the final transformer layer into a discrete token. The choice of algorithm is not a minor implementation detail. It governs the fundamental tradeoff between diversity and fidelity: too deterministic, and the model parrots its training distribution; too stochastic, and it drifts into incoherence. This chapter covers that tradeoff with precision — from the mathematics of each algorithm to production-grade implementation tricks.

We assume you have read about the autoregressive decoding loop in The Anatomy of LLM Inference: Prefill, Decode & The KV Cache and that you understand how Speculative Decoding: Draft Models, Medusa, EAGLE & Lookahead interacts with the sampling step. We also touch on connections to RL fine-tuning discussed in Policy Gradients & PPO for Language Models.

The Probability Pipeline: Logits to Token¶

The model outputs a real-valued vector \(\mathbf{z} \in \mathbb{R}^{|V|}\) called logits, where \(|V|\) is the vocabulary size (32,000 for LLaMA-2, 128,256 for LLaMA-3). The pipeline from logits to a sampled token is:

Logit processors transform the raw logits before the softmax — applying temperature scaling, vocabulary filtering, or penalty terms. The sampler draws from the resulting distribution. This separation is important: processors compose and stack, while the sampler is usually just a multinomial draw.

import torch
import torch.nn.functional as F
from typing import List, Optional

def sample_token(
    logits: torch.Tensor,          # shape: (vocab_size,), raw model outputs
    processors: List["LogitProcessor"],
    do_sample: bool = True,
) -> int:
    """
    Apply a list of logit processors in sequence, then sample one token.
    Returns the integer token id.
    """
    logits = logits.clone().float()  # always upcast to fp32 for numerical safety

    for proc in processors:
        logits = proc(logits)        # each processor modifies logits in-place or returns a new tensor

    if not do_sample:
        # greedy: just the argmax
        return int(logits.argmax())

    # convert to probabilities and sample
    probs = F.softmax(logits, dim=-1)
    return int(torch.multinomial(probs, num_samples=1))

Everything in this chapter is a specialisation of this loop. Let’s build each piece.

The Gumbel-Max Trick: Sampling as a Perturbed Argmax¶

torch.multinomial is convenient but it is not what high-throughput engines run. Internally it normalises, builds a cumulative distribution over the whole vocabulary, and does a search — awkward on GPU for a 128k-entry vector, and it allocates in a way that is unfriendly to CUDA graphs. There is an exact, sort-free alternative.

Let \(g_1,\ldots,g_{|V|}\) be i.i.d. samples from the standard Gumbel distribution, obtained from uniforms as \(g = -\log(-\log u)\) with \(u \sim \mathrm{Uniform}(0,1)\). Then

exactly. Sketch: \(z_t + g_t\) is Gumbel with location \(z_t\), and \(P(\max_t \{z_t+g_t\} \le y) = \prod_t \exp(-e^{z_t - y}) = \exp(-e^{\log \sum_t e^{z_t} - y})\), so the max is itself Gumbel with location \(\log\sum_t e^{z_t}\); carrying the same algebra through the argmax gives \(P(\arg\max = k) = e^{z_k}/\sum_t e^{z_t}\). Note this works on the unnormalised logits — no softmax, no partition function, no sort. Three consequences matter in practice:

1. It is one fused elementwise kernel plus an argmax reduction, with static shapes — exactly what a CUDA-graph-captured decode step wants.
2. Truncation composes for free. Tokens masked to \(-\infty\) stay at \(-\infty\) after adding finite noise, so top-k/top-p/min-p masks work unchanged.
3. Temperature is just a pre-scale: \(\arg\max_t(z_t/T + g_t)\) samples from \(P_T\). And \(T \to 0\) recovers greedy, which is why greedy is the zero-noise limit rather than a separate code path.

An algebraically equivalent form — the exponential race — divides probabilities by i.i.d. \(\mathrm{Exponential}(1)\) noise instead, which is how engines such as vLLM implement random sampling:

import torch

def gumbel_sample(logits: torch.Tensor, generator=None) -> torch.Tensor:
    """
    Exact sample from softmax(logits) with no sort and no cumsum.
    logits: (..., vocab). Returns integer token ids of shape (...).
    """
    u = torch.rand(logits.shape, device=logits.device, generator=generator)
    u = u.clamp_min(1e-20)                  # torch.rand is [0,1); guard log(0)
    g = -torch.log(-torch.log(u))           # inverse-CDF sample from Gumbel(0,1)
    return (logits + g).argmax(dim=-1)


def exponential_race_sample(probs: torch.Tensor, generator=None) -> torch.Tensor:
    """Equivalent form used in production samplers: argmax_t p_t / E_t, E_t ~ Exp(1)."""
    q = torch.empty_like(probs).exponential_(1.0, generator=generator)
    return probs.div(q).argmax(dim=-1)


if __name__ == "__main__":
    torch.manual_seed(0)
    logits = torch.tensor([2.0, 1.0, 0.0, -1.0])
    target = torch.softmax(logits, dim=-1)
    draws = gumbel_sample(logits.expand(200_000, -1))
    empirical = torch.bincount(draws, minlength=4).float() / draws.numel()
    print("target   :", target.tolist())
    print("empirical:", [round(x, 3) for x in empirical.tolist()])  # matches to ~1e-3

The trick also buys per-request reproducibility under continuous batching: give each sequence its own torch.Generator seed and its token stream no longer depends on which other requests happen to share the batch. (Bitwise determinism additionally requires batch-invariant matmul/attention kernels, since a different batch composition changes reduction order in the model, not just the sampler — see Continuous Batching & Request Scheduling.) Speculative decoding is the one place you still need explicit normalised probabilities, because its accept/reject rule compares \(p(t)\) and \(q(t)\) numerically — see Speculative Decoding: Draft Models, Medusa, EAGLE & Lookahead.

Greedy Decoding and Its Failure Modes¶

Greedy decoding picks the most probable token at every step:

It is deterministic, fast, and not equivalent to finding the most probable sequence. The globally most probable sequence — the maximum a posteriori (MAP) sequence — requires an exponential search, because token choices interact: the best next token often forecloses the best continuation.

In practice, greedy decoding produces repetitive, low-entropy text. Once the model writes “The cat sat on the”, it assigns high probability to “mat”, which then makes “mat mat mat” the greedy continuation. This is not a defect in the model; it is a consequence of choosing the locally optimal token at each step without lookahead.

def greedy_decode(
    model,
    input_ids: torch.Tensor,   # (1, seq_len)
    max_new_tokens: int = 100,
    eos_token_id: int = 2,
) -> torch.Tensor:
    """Minimal greedy decode loop."""
    generated = input_ids
    for _ in range(max_new_tokens):
        with torch.no_grad():
            logits = model(generated).logits[:, -1, :]   # (1, vocab)
        next_token = logits.argmax(dim=-1, keepdim=True)  # (1, 1)
        generated = torch.cat([generated, next_token], dim=1)
        if next_token.item() == eos_token_id:
            break
    return generated

Use greedy when: (a) you need exact reproducibility, (b) the task has a single correct answer and the model is well-calibrated (e.g., code completion with high-temperature training), or © you are the draft model in speculative decoding and the verifier handles stochasticity.

Temperature Scaling¶

Temperature \(T > 0\) is the single most important hyperparameter in decoding. It divides every logit by \(T\) before the softmax:

• \(T \to 0\): distribution collapses to a point mass at the argmax (greedy).
• \(T = 1\): the model’s trained distribution, unchanged.
• \(T \to \infty\): distribution becomes uniform over the vocabulary.

Dividing by \(T < 1\) sharpens the distribution (amplifies differences between logits); dividing by \(T > 1\) flattens it (suppresses differences).

class TemperatureProcessor:
    def __init__(self, temperature: float):
        assert temperature > 0, "Temperature must be positive"
        self.temperature = temperature

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        # Dividing logits by T is equivalent to multiplying log-probs by 1/T,
        # which is exactly the "softmax temperature" formulation.
        return logits / self.temperature

Temperature and the Training Distribution¶

A crucial subtlety: the model was trained with teacher-forcing at \(T = 1\). When you use \(T < 1\) at inference, you are sampling from a distribution that is sharper than what the model saw during training. This is usually fine — the probabilities just become more concentrated. When you use \(T > 1\), you are sampling tokens the model considers unlikely, which can produce creative but also hallucinated or incoherent text.

Temperature interacts deeply with RL fine-tuning (see Policy Gradients & PPO for Language Models). During PPO rollouts, the policy temperature determines exploration; a too-cold policy underfits the reward landscape, a too-hot policy produces garbage completions that confuse the reward model. Similarly, distillation objectives (see Distillation, Model Compression & Knowledge Transfer) often match soft targets — temperature-scaled probability distributions — from the teacher, where a higher temperature exposes the “dark knowledge” in the teacher’s off-peak probabilities.

Top-k Sampling¶

Top-k sampling restricts the vocabulary to the \(k\) tokens with the highest logits before sampling:

Setting logits to \(-\infty\) ensures those tokens get zero probability after softmax. This prevents the long tail of improbable tokens from polluting the distribution.

class TopKProcessor:
    def __init__(self, top_k: int):
        assert top_k >= 1
        self.top_k = top_k

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        # kth_val is the minimum value among the top-k
        kth_val = torch.topk(logits, self.top_k).values[-1]
        # Mask everything below the threshold
        return logits.masked_fill(logits < kth_val, float('-inf'))

Limitation: \(k\) is absolute. With \(k = 50\), you always consider 50 tokens whether the distribution is tight (one token dominates) or flat (many tokens are reasonable). Top-p addresses this.

Top-p / Nucleus Sampling¶

Holtzman et al. (“The Curious Case of Neural Text Degeneration”, 2020) introduced nucleus sampling: sample from the smallest set of tokens whose cumulative probability exceeds \(p\).

Tokens outside \(V_p\) are suppressed. The nucleus adapts: when the model is confident (the top token alone has probability 0.9), the nucleus is tiny; when the model is uncertain (probability spread over many tokens), the nucleus expands.

class TopPProcessor:
    def __init__(self, top_p: float):
        assert 0 < top_p <= 1.0
        self.top_p = top_p

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        # Sort logits descending to build cumulative sum
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        probs = F.softmax(sorted_logits, dim=-1)
        cumulative_probs = torch.cumsum(probs, dim=-1)

        # Remove tokens once cumulative probability exceeds top_p.
        # We shift by one so the token that crosses the threshold is kept.
        remove_mask = cumulative_probs - probs > self.top_p
        sorted_logits[remove_mask] = float('-inf')

        # Scatter back to original token ordering
        logits_filtered = torch.full_like(logits, float('-inf'))
        logits_filtered.scatter_(0, sorted_indices, sorted_logits)
        return logits_filtered

A typical production default is top_p=0.9 with temperature=0.8. Note that top-k and top-p compose: apply top-k first to cut the long tail cheaply, then apply top-p for adaptive nucleus selection.

Min-p Sampling¶

Min-p (Nguyen et al., 2024) takes a different angle: instead of keeping a top fraction by mass, it keeps all tokens whose probability is at least \(p_{\min}\) times the maximum token probability:

This scales the threshold relative to the peak, so it automatically tightens when the model is confident and loosens when it is uncertain — similar to top-p, but parameterised differently and often producing smoother behaviour at extreme temperatures. By 2026 the temperature-plus-min-p pair (e.g. temperature=0.7–1.0, min_p=0.05–0.1) has become a common default among open-model serving stacks, while commercial APIs still expose top-p as their primary truncation knob. A newer refinement, top-nσ (Tang et al., 2024), thresholds on the pre-softmax logits — keeping tokens within \(n\) standard deviations of the maximum logit — and is provably temperature-invariant: its candidate set does not change as you scale temperature, avoiding the noise-token blow-up that top-p and min-p can suffer at high \(T\).

class MinPProcessor:
    def __init__(self, min_p: float):
        assert 0 < min_p < 1.0
        self.min_p = min_p

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        probs = F.softmax(logits, dim=-1)
        # Absolute threshold: min_p * p_max
        threshold = self.min_p * probs.max()
        logits = logits.masked_fill(probs < threshold, float('-inf'))
        return logits

Typical Sampling¶

Meister et al. (“Typical Decoding for Natural Language Generation”, 2023) approach diversity from an information-theoretic angle. They observe that human text is typically drawn from the centre of the entropy distribution: tokens whose surprisal \(-\log P(t)\) is close to the conditional entropy \(H = -\sum_t P(t) \log P(t)\).

Define the typicality of token \(t\):

Typical sampling keeps only tokens satisfying this condition, discarding both the most probable (low surprisal, repetitive) and the least probable (high surprisal, incoherent).

class TypicalProcessor:
    def __init__(self, mass: float = 0.9):
        """Keep the smallest set of 'typical' tokens covering `mass` probability."""
        self.mass = mass

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        probs = F.softmax(logits, dim=-1)
        neg_log_probs = -torch.log(probs + 1e-10)

        # Conditional entropy (expected surprisal)
        H = (probs * neg_log_probs).sum()

        # Typicality: distance of each token's surprisal from entropy
        typicality = (neg_log_probs - H).abs()

        # Sort by typicality (most typical first), then take the nucleus covering `mass`
        sorted_typicality, sorted_indices = torch.sort(typicality)
        sorted_probs = probs[sorted_indices]
        cumsum = torch.cumsum(sorted_probs, dim=-1)
        remove_mask = cumsum - sorted_probs > self.mass

        sorted_logits = logits[sorted_indices]
        sorted_logits[remove_mask] = float('-inf')

        logits_filtered = torch.full_like(logits, float('-inf'))
        logits_filtered.scatter_(0, sorted_indices, sorted_logits)
        return logits_filtered

Repetition and Frequency Penalties¶

Even with good sampling, models can fall into repetitive loops. Two complementary penalties address this.

Repetition penalty (Keskar et al., “CTRL”, 2019) discounts any token that has already appeared in the context:

for penalty \(\theta > 1\). This reduces the probability of repeated tokens without suppressing them entirely.

Frequency penalty (used in OpenAI’s API) subtracts a penalty proportional to how many times the token has appeared:

Presence penalty is a simpler binary version — subtract a fixed \(\beta\) if the token appears at all:

class RepetitionPenaltyProcessor:
    """
    Multiplicative repetition penalty (CTRL-style).
    theta > 1 reduces repetition; theta = 1 is no-op.
    """
    def __init__(self, penalty: float, input_ids: torch.Tensor):
        assert penalty >= 1.0
        self.penalty = penalty
        # Track unique tokens seen in the context
        self.seen = set(input_ids.flatten().tolist())

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        logits = logits.clone()
        for token_id in self.seen:
            if logits[token_id] > 0:
                logits[token_id] /= self.penalty
            else:
                logits[token_id] *= self.penalty
        return logits


class FrequencyPenaltyProcessor:
    """
    Additive frequency penalty: subtract alpha * count(t).
    Also supports presence penalty (binary).
    """
    def __init__(
        self,
        frequency_penalty: float,
        presence_penalty: float,
        input_ids: torch.Tensor,
    ):
        self.frequency_penalty = frequency_penalty
        self.presence_penalty = presence_penalty
        # Count occurrences of each token in context
        token_list = input_ids.flatten().tolist()
        from collections import Counter
        self.counts = Counter(token_list)

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        logits = logits.clone()
        for token_id, count in self.counts.items():
            logits[token_id] -= (
                self.frequency_penalty * count
                + self.presence_penalty
            )
        return logits

Beam Search¶

Beam search maintains a beam of \(B\) partial hypotheses, expanding each at every step and keeping only the top-\(B\) by cumulative log-probability:

The length normalization exponent \(\alpha\) (typically 0.6–0.8) prevents the beam from preferring shorter sequences.

import heapq
from dataclasses import dataclass, field
from typing import Tuple

@dataclass(order=True)
class BeamHypothesis:
    score: float          # negative log-prob (min-heap)
    tokens: list = field(compare=False)

def beam_search(
    model,
    input_ids: torch.Tensor,     # (1, prefix_len)
    beam_size: int = 4,
    max_new_tokens: int = 100,
    eos_id: int = 2,
    length_alpha: float = 0.6,
) -> List[int]:
    """
    Minimal beam search. Returns the best sequence as a list of token ids.
    NB: production implementations use KV-cache for each beam; this toy
    version re-encodes each step for clarity.
    """
    prefix = input_ids[0].tolist()
    # heap items: (neg_score, token_list)
    active_beams: List[Tuple[float, List[int]]] = [(0.0, prefix[:])]
    finished: List[Tuple[float, List[int]]] = []

    for _ in range(max_new_tokens):
        if not active_beams:
            break
        candidates: List[Tuple[float, List[int]]] = []

        for neg_score, tokens in active_beams:
            ids = torch.tensor([tokens], dtype=torch.long)
            with torch.no_grad():
                logits = model(ids).logits[0, -1, :]   # (vocab,)
            log_probs = F.log_softmax(logits, dim=-1)

            # Expand: consider all tokens in the vocabulary
            topk_logprob, topk_ids = torch.topk(log_probs, beam_size)
            for lp, tid in zip(topk_logprob.tolist(), topk_ids.tolist()):
                new_tokens = tokens + [tid]
                new_neg_score = neg_score - lp   # minimise negative log-prob
                if tid == eos_id:
                    # Normalise by length
                    n = len(new_tokens) - len(prefix)
                    normalised = new_neg_score / (n ** length_alpha)
                    finished.append((normalised, new_tokens))
                else:
                    candidates.append((new_neg_score, new_tokens))

        # Keep the best beam_size active hypotheses
        candidates.sort(key=lambda x: x[0])
        active_beams = candidates[:beam_size]

    # Fall back to active beams if none finished
    if not finished:
        n_prefix = len(prefix)
        for neg_score, tokens in active_beams:
            n = len(tokens) - n_prefix
            normalised = neg_score / max(n, 1) ** length_alpha
            finished.append((normalised, tokens))

    finished.sort(key=lambda x: x[0])
    return finished[0][1]   # best hypothesis

When to use beam search: summarisation, machine translation, and tasks where the output has a clear quality metric. For open-ended generation, beam search amplifies repetition and produces bland text — sampling is better. For structured generation (see Structured & Constrained Generation), beam search is often combined with constraint masks.

Diverse beam search (Vijayakumar et al., 2018) adds a dissimilarity penalty between beams, encouraging them to explore different branches — useful when you want \(N\)-best diverse outputs.

Contrastive Decoding and DoLa¶

Two recent methods depart from purely probability-based selection.

Contrastive Decoding¶

Li et al. (“Contrastive Decoding”, 2023) observe that an amateur model (smaller, weaker) and the expert model (the LLM) assign high probability to the same fluent but factually incorrect completions. The contrastive score subtracts the amateur’s log-probability from the expert’s:

Tokens that the expert prefers over the amateur are amplified; tokens both models agree on (fluency) but the expert is unsure about are suppressed. In practice, the amateur is often the same model at a smaller context window or with its early layers.

DoLa (Decoding by Contrasting Layers)¶

Chuang et al. (“DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models”, 2023) push this idea inside a single model. The observation: factual knowledge is represented in the later transformer layers, while surface-level fluency is settled earlier. DoLa computes:

to select the premature layer that diverges most from the final layer, then uses:

as the decoding score. This amplifies tokens the final (knowledge-rich) layers prefer over intermediate layers, reducing hallucination without an external model.

def dola_logits(
    model,
    input_ids: torch.Tensor,
    premature_layer: int,         # e.g. layer 16 in a 32-layer model
    alpha: float = 0.1,           # mixing coefficient
) -> torch.Tensor:
    """
    Compute DoLa-adjusted logits by contrasting the final layer
    against a premature intermediate layer.

    The model must expose intermediate hidden states, e.g. via
    output_hidden_states=True in HuggingFace.
    """
    with torch.no_grad():
        outputs = model(
            input_ids,
            output_hidden_states=True,
        )
    hidden_states = outputs.hidden_states   # tuple of (batch, seq, d_model)

    # Intermediate hidden states are NOT normalised, so we must apply the
    # model's final norm before the LM head — the head was trained on
    # normalised activations, and skipping this makes the premature-layer
    # distribution garbage (the classic logit-lens / DoLa re-implementation bug).
    # In HuggingFace, hidden_states[-1] is already post-norm, so the final-layer
    # logits are just outputs.logits.
    norm = model.model.norm            # LLaMA-style; adjust for other architectures
    lm_head = model.lm_head
    premature_logits = lm_head(norm(hidden_states[premature_layer][:, -1, :]))
    final_logits     = outputs.logits[:, -1, :]

    # Contrastive combination
    dola_log_probs = (
        F.log_softmax(final_logits, dim=-1)
        - alpha * F.log_softmax(premature_logits, dim=-1)
    )
    return dola_log_probs.squeeze(0)

The Bias-Diversity Tradeoff¶

Every decoding choice sits on a single underlying tradeoff surface. Let us make it concrete.

Diversity refers to the entropy or variety of the generated text — how many different continuations the sampler would produce for the same prompt. Bias refers to the systematic difference between the decoded distribution and the model’s true distribution.

Greedy decoding has zero variance (perfectly reproducible) but extreme bias — it always picks the mode, ignoring the full shape of the distribution. Uniform sampling has maximum diversity but ignores everything the model learned. Temperature is the most direct knob on this axis.

The ideal point depends on the task:

One important empirical finding: for most instruction-tuned models, the training process implicitly calibrates the output distribution for temperature around 0.7–1.0. Going significantly below 0.3 causes the model to confidently hallucinate because the distribution was not trained to be sharp; going above 1.5 on instruction models often causes grammatical collapse.

Logit Processor Pipelines in Practice¶

Production systems (HuggingFace Transformers, vLLM, SGLang) implement logit processing as a composable pipeline. The full processing order typically is:

Below is a composable, HuggingFace-compatible LogitsProcessor implementation:

from transformers import LogitsProcessor, LogitsProcessorList
import torch

class CompositeLogitsProcessor(LogitsProcessor):
    """
    A single LogitsProcessor that applies a pipeline of sub-processors,
    each expecting (input_ids, scores) -> scores.
    This follows the HuggingFace LogitsProcessor protocol.
    """

    def __init__(self, processors: list):
        self.processors = processors

    def __call__(
        self,
        input_ids: torch.LongTensor,     # (batch, seq_len)
        scores: torch.FloatTensor,       # (batch, vocab_size)
    ) -> torch.FloatTensor:
        for proc in self.processors:
            scores = proc(input_ids, scores)
        return scores


# Example: build a typical production pipeline
def build_logits_processor_list(
    temperature: float = 0.8,
    repetition_penalty: float = 1.1,
    top_k: int = 50,
    top_p: float = 0.9,
) -> LogitsProcessorList:
    from transformers import (
        TemperatureLogitsWarper,
        RepetitionPenaltyLogitsProcessor,
        TopKLogitsWarper,
        TopPLogitsWarper,
    )

    procs = LogitsProcessorList()
    if repetition_penalty != 1.0:
        procs.append(RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty))
    if temperature != 1.0:
        procs.append(TemperatureLogitsWarper(temperature=temperature))
    if top_k > 0:
        procs.append(TopKLogitsWarper(top_k=top_k, min_tokens_to_keep=1))
    if top_p < 1.0:
        procs.append(TopPLogitsWarper(top_p=top_p, min_tokens_to_keep=1))
    return procs

Recent transformers also ships MinPLogitsWarper (and EpsilonLogitsWarper / TypicalLogitsWarper), so min-p and typical sampling are one-liners rather than custom code; the LogitsWarper base class was folded into LogitsProcessor, but the concrete classes above kept their names. In normal use you never build this list by hand — you pass the knobs to model.generate(..., do_sample=True, temperature=0.8, top_p=0.9) and transformers assembles the same list for you; you construct it explicitly only when you need a custom processor in the chain (generate(..., logits_processor=procs)).

The min_tokens_to_keep=1 argument is critical: it prevents cases where the filter masks all tokens (which would cause a nan after softmax), by always keeping at least the top token.

The Same Knobs in Real Serving Engines¶

Every production engine exposes the same conceptual pipeline under slightly different names. Knowing the mapping is most of the job.

# vLLM: sampling parameters travel with the request, not with the engine,
# so different requests in one continuous batch can use different settings.
from vllm import LLM, SamplingParams

llm = LLM(model="Qwen/Qwen3-1.7B", max_model_len=8192)

chat = SamplingParams(
    temperature=0.7,        # exactly 0.0 is special-cased to greedy
    top_p=0.95,             # set to 1.0 to disable
    top_k=-1,               # -1 disables here; HuggingFace uses 0
    min_p=0.0,              # peak-relative floor; 0.0 disables
    repetition_penalty=1.0, # CTRL-style multiplicative
    frequency_penalty=0.0,  # OpenAI-style additive, per occurrence
    presence_penalty=0.0,   # OpenAI-style additive, binary
    max_tokens=512,
    stop=["</answer>"],     # string stop sequences (detokenised, not token ids)
    seed=1234,              # per-request RNG -> reproducible regardless of batchmates
    logprobs=5,             # return top-5 logprobs per step (for eval / debugging)
)

# RL rollouts and best-of-n: n>1 shares the prompt's KV cache across samples.
rollout = SamplingParams(n=8, temperature=1.0, top_p=1.0, max_tokens=1024, seed=None)

for out in llm.generate(["Explain nucleus sampling in one sentence."], chat):
    print(out.outputs[0].text)

Three engine-specific details worth memorising:

• vLLM applies penalties before temperature, matching the HuggingFace order, and treats temperature=0.0 as a request for greedy rather than as a division by zero. Its sampler is fully vectorised over the batch: penalties are applied with a scatter_add_-built count matrix rather than the Python for loops in this chapter’s teaching code, and truncation is a batched sort. Never ship a per-token Python loop over the vocabulary in a serving path — at 128k vocab and 50 requests it dominates the decode step.
• SGLang takes the same fields on its sampling_params dict (temperature, top_p, top_k, min_p, frequency_penalty, presence_penalty, repetition_penalty) and additionally lets a program change them mid-generation — low temperature inside a JSON block, higher outside. See SGLang: RadixAttention & Structured Programs.
• llama.cpp (the GGUF/laptop path) models decoding literally as a sampler chain built from llama_sampler_init_* pieces — top_k, top_p, min_p, typical, temp / temp_ext, penalties, dist — pushed onto a llama_sampler_chain in the order you want them applied, with dist (the multinomial draw) or greedy last. It is the cleanest existing implementation of exactly the abstraction this chapter builds — and the one you will use to run a quantized model on a laptop (see Quantization II: INT4/INT8/FP8, GGUF, bitsandbytes & QAT).

Both OpenAI-compatible servers (vLLM, SGLang, TGI) accept temperature, top_p, frequency_penalty, presence_penalty, seed and logprobs on /v1/chat/completions; top_k, min_p and repetition_penalty are non-standard extensions passed via an extra_body field, which is the usual reason a min-p setting silently does nothing when you switch clients.

Custom Logit Processors¶

Custom processors enable powerful behaviours:

• Grammar enforcement: mask all tokens incompatible with a context-free grammar at the current parser state (see Structured & Constrained Generation).
• Watermarking: John Kirchenbauer et al. (“A Watermark for Large Language Models”, 2023) add a small positive bias to a randomly chosen half of the vocabulary, seeded by the previous token. This creates a statistically detectable fingerprint in the output.
• Token healing: vLLM’s token healing re-processes the last token of the prompt to avoid boundary artefacts when the prompt ends mid-token.
• Vocabulary bias: force specific tokens (e.g., “yes”/”no” for classification prompts) by setting all other logits to \(-\infty\).

Temperature and RL/Distillation Interactions¶

This section connects decoding to training, two topics that interact subtly.

Policy Temperature in RL Fine-Tuning¶

During PPO (see Policy Gradients & PPO for Language Models), the policy generates rollouts. The sampling temperature during rollout generation serves as an exploration parameter:

• Low \(T\): the policy stays close to the current mode, producing similar completions. Good for exploiting known-good strategies.
• High \(T\): the policy explores more, which may discover high-reward completions but also floods the reward model with low-quality samples.

The KL penalty term in PPO (\(\beta \cdot D_{\text{KL}}(\pi \| \pi_{\text{ref}})\)) is computed between the policy logits and reference model logits at \(T = 1\). If you sample rollouts at \(T \ne 1\) but compute the KL at \(T = 1\), you are optimising a different objective than the one you are generating samples from. Most correct implementations compute both log-probabilities at \(T = 1\) (or consistently at the same \(T\)). See Advantage Estimation, KL Control & Stability Tricks for details.

Knowledge Distillation and Soft Targets¶

Hinton et al. (“Distilling the Knowledge in a Neural Network”, 2015) showed that training a student on soft targets from a teacher at temperature \(T > 1\) transfers more information than one-hot labels. The teacher’s distribution at high temperature exposes similarities between classes (or tokens) — its “dark knowledge”.

For LLM distillation, we match:

The \(T^2\) factor re-scales the gradient to have the same magnitude regardless of temperature (since softmax gradients scale as \(1/T^2\) when divided by \(T\)). In practice, temperatures of 2–5 are common for token-level distillation.

Putting It All Together: A Production Sampler¶

Here is a complete, self-contained sampler that combines all the techniques above:

import torch
import torch.nn.functional as F
from collections import Counter
from typing import List, Optional


class ProductionSampler:
    """
    A complete, composable sampler implementing:
      - Temperature scaling
      - Repetition penalty (CTRL-style)
      - Frequency + presence penalties (OpenAI-style)
      - Top-k filtering
      - Top-p (nucleus) filtering
      - Min-p filtering
      - Greedy fallback

    Usage:
        sampler = ProductionSampler(temperature=0.8, top_p=0.9, rep_penalty=1.1)
        next_token = sampler(logits, context_ids)
    """

    def __init__(
        self,
        temperature: float = 1.0,
        top_k: int = 0,                  # 0 = disabled
        top_p: float = 1.0,              # 1.0 = disabled
        min_p: float = 0.0,              # 0.0 = disabled
        repetition_penalty: float = 1.0, # 1.0 = disabled
        frequency_penalty: float = 0.0,
        presence_penalty: float = 0.0,
        do_sample: bool = True,
    ):
        self.temperature = temperature
        self.top_k = top_k
        self.top_p = top_p
        self.min_p = min_p
        self.repetition_penalty = repetition_penalty
        self.frequency_penalty = frequency_penalty
        self.presence_penalty = presence_penalty
        self.do_sample = do_sample

    def __call__(
        self,
        logits: torch.Tensor,            # (vocab_size,) raw logits
        context_ids: Optional[List[int]] = None,
    ) -> int:
        logits = logits.clone().float()

        # 1. Repetition penalty (multiplicative)
        if self.repetition_penalty != 1.0 and context_ids:
            for tid in set(context_ids):
                if logits[tid] > 0:
                    logits[tid] /= self.repetition_penalty
                else:
                    logits[tid] *= self.repetition_penalty

        # 2. Frequency + presence penalties (additive)
        if (self.frequency_penalty != 0.0 or self.presence_penalty != 0.0) and context_ids:
            counts = Counter(context_ids)
            for tid, cnt in counts.items():
                logits[tid] -= (
                    self.frequency_penalty * cnt + self.presence_penalty
                )

        # 3. Temperature scaling
        if self.temperature != 1.0:
            logits = logits / self.temperature

        # Remember the pre-truncation argmax: it is both the greedy answer and
        # the only safe fallback if the filters below mask every token.
        best_token = int(logits.argmax())

        if not self.do_sample:
            return best_token

        # 4. Top-k filter
        if self.top_k > 0:
            kth = torch.topk(logits, min(self.top_k, logits.size(-1))).values[-1]
            logits = logits.masked_fill(logits < kth, float('-inf'))

        # 5. Top-p (nucleus) filter
        if self.top_p < 1.0:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            probs = F.softmax(sorted_logits, dim=-1)
            cumprobs = torch.cumsum(probs, dim=-1)
            remove = cumprobs - probs > self.top_p
            sorted_logits[remove] = float('-inf')
            logits = torch.full_like(logits, float('-inf'))
            logits.scatter_(0, sorted_idx, sorted_logits)

        # 6. Min-p filter
        if self.min_p > 0.0:
            probs = F.softmax(logits, dim=-1)
            threshold = self.min_p * probs.max()
            logits = logits.masked_fill(probs < threshold, float('-inf'))

        # 7. Safety: ensure at least one valid token. An all -inf vector would
        # softmax to nan and crash torch.multinomial; this is the equivalent of
        # HuggingFace's min_tokens_to_keep=1 guard.
        if bool(torch.isinf(logits).all()):
            return best_token

        # 8. Sample
        probs = F.softmax(logits, dim=-1)
        return int(torch.multinomial(probs, num_samples=1))


# ── Quick test ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    torch.manual_seed(42)
    vocab_size = 32_000

    # Simulate a logit vector with a clear peak at token 7
    logits = torch.randn(vocab_size) * 2.0
    logits[7] = 10.0

    sampler = ProductionSampler(temperature=0.8, top_p=0.9, repetition_penalty=1.1)
    context = [1, 7, 42, 7, 7]        # token 7 has appeared 3 times

    counts = Counter()
    for _ in range(1000):
        t = sampler(logits.clone(), context)
        counts[t] += 1

    # With repetition penalty, token 7 should appear less than greedy would predict
    print(f"Token 7 selected {counts[7]} / 1000 times with rep_penalty=1.1")
    greedy_sampler = ProductionSampler(temperature=0.8, top_p=0.9, repetition_penalty=1.0)
    counts_nopenalty = Counter()
    for _ in range(1000):
        t = greedy_sampler(logits.clone(), context)
        counts_nopenalty[t] += 1
    print(f"Token 7 selected {counts_nopenalty[7]} / 1000 times without penalty")

Further Reading¶

• Holtzman et al., “The Curious Case of Neural Text Degeneration” (ICLR 2020) — introduced nucleus (top-p) sampling and the analysis of degenerate modes.
• Meister et al., “Typical Decoding for Natural Language Generation” (ACL 2023) — information-theoretic motivation for typical sampling.
• Keskar et al., “CTRL: A Conditional Transformer Language Model for Controllable Generation” (2019) — introduced the repetition penalty.
• Li et al., “Contrastive Decoding: Open-ended Text Generation as Optimization” (ACL 2023) — expert vs. amateur contrastive scores.
• Chuang et al., “DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models” (ICLR 2024) — layer-contrastive decoding for hallucination reduction.
• Vijayakumar et al., “Diverse Beam Search: Decoding Diverse Solutions from Neural Sequence Models” (AAAI 2018) — diversity-encouraging beam search.
• Kirchenbauer et al., “A Watermark for Large Language Models” (ICML 2023) — logit-based watermarking via vocabulary partitioning.
• Hinton et al., “Distilling the Knowledge in a Neural Network” (NeurIPS 2015 Workshop) — foundational work on temperature-scaled soft targets.
• HuggingFace transformers — src/transformers/generation/logits_process.py is the canonical reference implementation of every processor discussed here.

Exercises¶

1. Greedy decoding picks \(\arg\max_t P(t \mid t_{<i})\) at every step. Explain, with a small two-step example, why this does not in general return the globally most probable sequence (the MAP sequence). Construct a toy case with a vocabulary of two tokens \(\{A, B\}\) and two decoding steps where greedy chooses a first token that leads to a lower-probability full sequence than a different first token would.

2. A model emits three top logits \(z = [4.0,\, 2.0,\, 1.0]\) for tokens \(A, B, C\) (all other tokens are negligible). Using \(e^4 \approx 54.60\), \(e^2 \approx 7.389\), \(e^1 \approx 2.718\), \(e^8 \approx 2980.96\):

(a) Compute \(P(A), P(B), P(C)\) at temperature \(T = 1\). (b) Compute them at \(T = 0.5\). © Give the ratio \(P(A)/P(B)\) at both temperatures using the closed form, and state in one sentence what this shows about temperature.

3. After softmax, a step produces the sorted probability distribution

$\(\mathbf{p} = [0.50,\, 0.20,\, 0.15,\, 0.10,\, 0.05].\)$

(a) Using the chapter’s top-p rule (keep the token that crosses the threshold; formally, remove token \(i\) iff \(\big(\sum_{j\le i} p_j\big) - p_i > p\)), which tokens survive top-p with \(p = 0.9\)? (b) Which tokens survive min-p with \(p_{\min} = 0.30\) (keep \(t\) iff \(P(t) \geq p_{\min}\cdot \max_{t'}P(t')\))? © The two filters keep different sets here. Explain what property of min-p causes the difference.

4. The chapter’s RepetitionPenaltyProcessor penalises every token that has ever appeared in the context. In a long chat, this can wrongly suppress common function words that legitimately recur (e.g. “the”, “of”). Implement a WindowedRepetitionPenaltyProcessor that applies the CTRL-style multiplicative penalty only to tokens appearing in the last window tokens of the context, keeping the chapter’s sign convention (\(z > 0 \Rightarrow z/\theta\), \(z < 0 \Rightarrow z\cdot\theta\)).

import torch

class WindowedRepetitionPenaltyProcessor:
    """
    CTRL-style multiplicative repetition penalty restricted to a
    sliding window of the most recent `window` context tokens.
    theta > 1 discourages recent repeats; theta = 1 is a no-op.
    """
    def __init__(self, penalty: float, input_ids: torch.Tensor, window: int = 64):
        assert penalty >= 1.0
        self.penalty = penalty
        self.window = window
        # Only the last `window` tokens are eligible for penalisation.
        recent = input_ids.flatten().tolist()[-window:]
        self.seen = set(recent)

    def __call__(self, logits: torch.Tensor) -> torch.Tensor:
        logits = logits.clone()
        for token_id in self.seen:
            if logits[token_id] > 0:
                logits[token_id] /= self.penalty
            else:
                logits[token_id] *= self.penalty
        return logits

5. Implement contrastive decoding (Li et al., 2023) as a logit processor. Given raw expert logits (the LLM) and raw amateur logits (a weaker model) at the same step, return a score vector suitable for sampling that implements

$\(\text{CD}(t) = \log P_{\text{expert}}(t) - \log P_{\text{amateur}}(t),\)$

restricted to plausible tokens only — those whose expert probability is at least \(\alpha\) times the expert’s max probability (reuse the min-p style threshold from the chapter). Implausible tokens must be masked to \(-\infty\) so they cannot be sampled.

import torch
import torch.nn.functional as F

class ContrastiveDecodingProcessor:
    """
    Contrastive decoding: score = log P_expert - log P_amateur,
    restricted to the expert's plausible set (min-p style threshold).

    Construct with the amateur's raw logits for THIS step, then call
    with the expert's raw logits. Returns a score vector to sample from.
    """
    def __init__(self, amateur_logits: torch.Tensor, alpha: float = 0.1):
        assert 0.0 < alpha < 1.0
        self.amateur_log_probs = F.log_softmax(amateur_logits.float(), dim=-1)
        self.alpha = alpha

    def __call__(self, expert_logits: torch.Tensor) -> torch.Tensor:
        expert_log_probs = F.log_softmax(expert_logits.float(), dim=-1)

        # Plausibility constraint (adaptive, peak-relative like min-p):
        # keep only tokens with P_expert(t) >= alpha * max_t P_expert(t).
        expert_probs = expert_log_probs.exp()
        threshold = self.alpha * expert_probs.max()
        implausible = expert_probs < threshold

        # Contrastive score, then mask the implausible tail to -inf.
        cd = expert_log_probs - self.amateur_log_probs
        cd = cd.masked_fill(implausible, float('-inf'))
        return cd