Top-p (Nucleus) Sampling, in depth

Why truncate the distribution at all

A language model's next-token distribution has a head of plausible tokens and a tail of tens of thousands of implausible ones. Each tail token is individually tiny, but together they can hold several percent of the mass. Sample a few hundred tokens from the full distribution and a tail token is all but guaranteed to appear somewhere, and once one nonsense token is in the context, the model conditions on it and the text derails.

Top-k truncation, keeping the k most probable tokens, fixes the tail but uses the same k everywhere. When the model is confident, k = 50 still admits 49 bad options; when many continuations are reasonable, k = 50 may cut good ones. Top-p makes the cut adaptive: it keeps however many tokens it takes to cover p of the probability mass. The choice of how much randomness to allow is covered in temperature; top-p decides what is eligible at all.

The nucleus, worked through

Take six candidate tokens with logits 4.0, 3.0, 2.2, 1.6, 1.0 and -2.0. Softmax gives the probabilities below. Sort descending, accumulate, and stop at the first prefix whose cumulative mass reaches p.

Two consequences are easy to miss. First, the kept set can overshoot p by almost a whole token's probability: at p = 0.5 the single token kept carries 0.5966. Second, adaptivity is dramatic. Over a 100-token vocabulary, a peaked distribution (logits 8, 3, 2, 1, then zeros) needs one token to cover 0.9, while a gently sloping one (logits evenly spaced from 1 to 0) needs 85. Top-k with any fixed k gets one of those two cases badly wrong.

Boundaries and ties: where libraries disagree

"Smallest prefix with mass at least p" sounds unambiguous, but implementations differ at the boundary and the differences are visible in tests.

• Inclusive or exclusive. Hugging Face's TopPLogitsWarper sorts ascending, takes the cumulative sum, and removes tokens whose cumulative mass is <= 1 - top_p. For probabilities 0.5, 0.3, 0.2 and p = 0.8, a descending-prefix implementation keeps two tokens (0.5 + 0.3 reaches 0.8 exactly), while the ascending formulation keeps all three, because 1 - 0.8 is 0.19999999999999996 in floating point and the 0.2 tail is not removed. Neither is wrong; they are different rules plus rounding.
• Ties. When equal probabilities straddle the cut, which one survives depends on the sort. Stable sorts keep the lower token ID, so the choice is deterministic but arbitrary. A rule based on "mass of strictly more probable tokens" keeps all tied tokens instead.
• Never empty. Rounding can make even the top token fail a strict comparison. Always force-keep at least one token; Hugging Face exposes this as min_tokens_to_keep, which protects the most probable tokens.
• p = 1.0 should be a no-op, and engines usually skip the filter entirely. Very small p makes that step greedy.

If you compare two engines token for token, these rules alone will produce differences on a small fraction of steps, long before any bug.

Where top-p sits in the sampler chain

Top-p is one stage in a sampler chain, and its meaning depends on what ran before it. A typical engine applies penalties (repetition, frequency, presence) to logits, applies grammar or JSON masks, divides by temperature, converts to probabilities, then applies top-k, top-p and min-p filters and samples. Temperature before top-p means a high temperature flattens the distribution and widens the nucleus; some libraries can be configured the other way round, and the two orders give different sets. Constraint masks must come before top-p, otherwise the nucleus can consist entirely of tokens the grammar then forbids. Temperature and sampling settings covers ordering from the API user's side and what each provider accepts; the sampling maths derives each filter.

Batched top-p with a different p per request

A serving engine samples for a whole batch at once, and every request can carry its own temperature and p. The vectorised form keeps a token when the mass of tokens strictly ahead of it in sorted order is below p, which is the smallest-prefix rule written so it needs no search. This is NumPy so you can run it; the PyTorch version uses the same sort, cumsum, gather and scatter operations.

import numpy as np


def top_p_mask(probs, top_p):
    """probs: [B, V] float32 rows summing to 1; top_p: [B]. Returns a bool keep-mask."""
    order = np.argsort(-probs, axis=-1, kind="stable")
    sorted_p = np.take_along_axis(probs, order, axis=-1)
    cum = np.cumsum(sorted_p, axis=-1, dtype=np.float32)
    before = cum - sorted_p                   # mass strictly ahead of each token
    keep_sorted = before < top_p[:, None]     # smallest prefix whose mass reaches p
    keep_sorted[:, 0] = True                  # never empty
    keep = np.zeros_like(keep_sorted)
    np.put_along_axis(keep, order, keep_sorted, axis=-1)
    return keep


def sample_top_p(logits, temperature, top_p, rng):
    z = logits / temperature[:, None]
    z = z - z.max(axis=-1, keepdims=True)
    probs = np.exp(z, dtype=np.float32)
    probs /= probs.sum(axis=-1, keepdims=True)
    probs = np.where(top_p_mask(probs, top_p), probs, 0.0)
    probs /= probs.sum(axis=-1, keepdims=True)
    u = rng.random((probs.shape[0], 1))
    return (probs.cumsum(axis=-1) < u).sum(axis=-1)

Against a float64 reference, this mask agreed on 3,000 random rows. Sampling the six-token example 20,000 times at p = 0.9 gave frequencies of 0.652, 0.236 and 0.112 for the three kept tokens and zero for the rest, matching the renormalised 0.6522, 0.2399 and 0.1078. A row with p = 0.5 in the same batch always returned "the".

The cost is the sort: O(V log V) per row, on vocabularies that now commonly exceed 100,000 tokens, every decoding step. At small batch sizes sampling is a small fraction of a step; at large batch sizes with small models it becomes noticeable, which is why engines either fuse it into custom kernels or avoid sorting altogether.

Precision: the float16 cumulative-sum collapse

Do the softmax and cumulative sum in float32 even when the model runs in bfloat16 or float16. The failure is not subtle rounding, it is collapse. With 128,000 tokens drawn from a realistic logit spread, the float64 and float32 cumulative sums agreed that p = 0.9 needs 13,483 tokens. The float16 cumulative sum stopped growing at exactly 0.5: once the running total reaches 0.5, the spacing between representable float16 values is about 0.0005, and every remaining tail probability is smaller than half of that, so adding it changes nothing. The sum never reaches 0.9, no prefix qualifies, and the filter keeps the whole vocabulary. Top-p silently turns into pure sampling, which is exactly the failure it exists to prevent, and no error is raised. bfloat16 has fewer mantissa bits, so it stalls sooner.

Sort-free top-p by rejection sampling

Sorting exists only to find the cut. FlashInfer's sampling kernels avoid the global sort with rejection sampling, documented in its API reference as sorting-free and run as multiple rounds inside one kernel. The underlying idea fits in a few lines; FlashInfer's own pivot strategy is more refined, so read this as the principle rather than its implementation:

def top_p_rejection(p, top_p, rng, max_rounds=64):
    pivot = -1.0
    for rounds in range(1, max_rounds + 1):
        cand = np.where(p > pivot, p, 0.0)
        t = rng.choice(p.size, p=cand / cand.sum())   # sample among tokens above the pivot
        if p[p > p[t]].sum() < top_p:                 # t is inside the nucleus: accept
            return t, rounds
        pivot = p[t]                                  # t and everything below it are out
    raise RuntimeError("no convergence")

A token belongs to the nucleus exactly when the mass of strictly more probable tokens is below p, so the acceptance test needs one reduction over the vocabulary, not a sort. Rejected draws raise the pivot and discard the rejected token with everything less probable. Because the nucleus always lies above the pivot, accepted tokens follow the renormalised nucleus distribution: 40,000 draws on the six-token example gave 0.650, 0.241 and 0.109, averaging 1.09 rounds. On the 128,000-token distribution, 300 draws averaged 1.08 rounds with a maximum of 2, because the first draw lands in a 0.9 nucleus nine times in ten. Note the tie rule: this test keeps every token tied at the boundary, unlike a stable sort.

Logprobs, speculative decoding and reproducibility

Logprobs. When an API returns token logprobs, they may describe the raw model distribution or the processed one after temperature and truncation, and engines differ; some let you choose. Under the processed distribution, every token outside the nucleus has probability zero. Evaluation code that scores sampled text with returned logprobs must know which one it got.

Speculative decoding. Speculative sampling accepts a draft token with probability min(1, p(x)/q(x)) and resamples from the normalised positive part of p - q on rejection, which reproduces the target distribution p exactly for any draft distribution q. With top-p, p must be the target's truncated and renormalised distribution at that position, and q must be the distribution the draft actually sampled from, including any truncation the draft applied. Use the raw target distribution and you sample from something other than what the user requested, with no visible error.

Reproducibility. A per-request seed makes the random draw repeatable, but not the boundary. Batch composition can change reduction order in floating point, which can move a near-boundary token in or out of the nucleus. Treat seeded sampling as approximately reproducible.

Failure modes

Failure modes, with the symptom you will see:

• Half-precision cumsum. Occasional nonsense tokens despite p = 0.9. Compute sampling in float32.
• Constraints applied after top-p. Structured output fails or falls back to a low-probability token. Mask first, then truncate.
• p too low with low temperature. Loops and repeated phrases, because each step is nearly greedy. Raise p or add a repetition penalty.
• High temperature with p near 1. The flattened distribution makes the nucleus enormous and quality collapses. Lower one of them.
• NaN logits. One NaN poisons the softmax and the cumulative sum, and every comparison is false. Check for NaN before sampling and fail the request.

Choosing p

Common defaults sit between 0.9 and 0.95 for open-ended text, and some provider APIs recommend changing either temperature or top-p but not both. Treat that as a starting point and tune on your own task: generate a few hundred outputs per setting, score them with your evaluator, and watch both quality and diversity, because lowering p always buys coherence with variety. For extraction, classification and code, a low temperature usually matters more than p. For creative or exploratory generation, top-p's adaptivity is the reason to prefer it over top-k, and min-p is a newer alternative that scales its cut to the top token's probability; sampling and decoding compares the family.

What to do next

1. Find where your stack applies top-p relative to penalties, grammar masks, temperature and top-k, and write the order down.
2. Confirm the sampler computes softmax and cumulative sums in float32; add a test that a 128,000-token distribution at p = 0.9 keeps far fewer than all tokens.
3. Write boundary tests: 0.5, 0.3, 0.2 at p = 0.8, ties at the cut, p = 1.0 and very small p, and decide which rule you want.
4. If you serve speculative decoding, verify that acceptance uses the truncated target distribution and the draft's actual sampling distribution.
5. Check whether your API returns raw or processed logprobs, and label evaluation data accordingly.
6. Tune p on your own task with an evaluator rather than copying defaults, changing one knob at a time.