How do you implement LLM sampling from logits?

In this article

1. Before you start
2. Why interviewers ask this
3. Sampling from logits from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page belongs to the machine learning (ML) coding round guide. You only need three ideas, all explained in a clause below.

Why interviewers ask this

Writing a sampler from scratch is a short coding task that exposes whether a candidate understands what a language model actually emits. It has been reported by candidates in coding rounds at AI labs and large tech companies, sometimes phrased as "implement top-k or nucleus sampling" and sometimes as "sample from model logits".

A weak answer calls torch.multinomial on the raw logits, or applies top-p without saying what the cutoff keeps. A strong answer orders the steps (scale, softmax, truncate, renormalize, draw), knows why each exists and handles the edge cases: temperature zero, ties, and the token that crosses the top-p threshold. A Staff-level answer also says when not to sample at all: for tasks with one right answer, greedy or beam search is the better tool.

Sampling from logits from first principles

Top-k sampling

Top-k sampling is a decoding method that keeps only the k highest-probability tokens, renormalizes their probabilities so they sum to one and draws one token at random from them. The count k is fixed whatever the context.

Top-p sampling

Top-p sampling, also called nucleus sampling, is a decoding method that keeps the smallest set of highest-probability tokens whose probabilities add up to at least p, renormalizes and draws from that set. The set size changes with the model's confidence.

In plain words

Picture a raffle. Every possible next token owns tickets, and the model decides how many. Greedy decoding always gives the prize to the token with the most tickets. Sampling draws one ticket at random. A token with a quarter of the tickets wins a quarter of the time. Temperature changes how unequal the ticket counts are. Top-k and top-p throw away the tokens with very few tickets before the draw. The analogy breaks in one place. The tickets are not whole numbers. They are probabilities that must sum to one, so after discarding tokens you must rescale the rest.

In one precise sentence: sampling draws the next token from a probability distribution built from the logits, after optionally rescaling it with a temperature and truncating its tail.

Step 1: The model outputs logits, one raw score per token

Logits are the raw, unnormalized scores a language model outputs, one real number per vocabulary token, and they can be negative or larger than one. For one generation step you take the logits of the last position: a vector whose length is the vocabulary size.

Input: that vector. Output: nothing yet; the numbers are not probabilities. The distinction matters in code, because undefined treats its input as probabilities or weights, so the rows do not need to sum to one but must be non-negative, finite and have a non-zero sum. Raw logits can be negative, so they must go through softmax first. Greedy decoding is the simplest use of logits: take the index of the largest one and stop. Everything below is what you do when you want randomness instead. For how the logits are produced, see building generative pre-trained transformer (GPT) from scratch.

Step 2: Temperature rescales the logits, then softmax makes probabilities

Temperature $T$ divides every logit before the softmax, so token $i$ gets probability $p_i=\frac{e^{z_i/T}}{{\sum }_j{e}^{z_j/T}}$, where $z_i$ is its logit. Hugging Face describes its temperature warper as exponential scaling of the output probability distribution that controls randomness; the default value is 1.0 in the v4.26.1 documentation snapshot, which changes nothing.

Input: the logits and $T$. Output: a probability distribution. Subtract the largest scaled logit before exponentiating so that large values do not overflow; that is the whole point of a numerically stable softmax. Dividing by a number below one stretches the gaps between logits, so the top token gains probability; dividing by a number above one shrinks the gaps. Division by zero is not defined, so implementations special-case a temperature of zero as greedy argmax; treat that as a convention, not a mathematical fact. You now have a full distribution to cut.

Step 3: Top-k keeps a fixed number of tokens

Top-k sampling keeps the k highest-probability tokens and discards the rest. Fan et al. (2018) call it top-k random sampling and sample from the 10 most likely candidates; the Hugging Face top_k parameter defaults to 50 in the v4.26.1 snapshot (check the current documentation for today's default).

Input: probabilities and k. Output: the same vector with every token outside the top k set to zero. Then divide by the sum of what is left, because the kept probabilities no longer add up to one. The drawback is that k is blind to context. When the model is nearly certain, a large k lets a few junk tokens in; when many continuations are reasonable, a small k cuts good ones. That is the gap top-p fills. Ties at the k-th place need a rule: a sort-based implementation keeps exactly k, while a threshold-based one keeps every tie.

Step 4: Top-p keeps the smallest set that reaches p

Top-p sampling keeps the smallest set of most probable tokens whose cumulative probability reaches p; the Hugging Face documentation words it as the smallest set whose probabilities add up to top_p or higher. Holtzman et al. (2020) call the kept set the nucleus and describe it as dynamic, so it grows and shrinks with the context. A Hugging Face code comment says 0.9 to 0.95 is the practical range for p.

Input: probabilities and p. Sort descending, take the running total, and keep every token up to and including the first one that pushes the total to p or beyond. The off-by-one is dropping that crossing token: with a mask of "running total greater than p" you remove the token that reaches p and keep a set whose mass is below p. The safe mask removes a token only when the mass before it already reached p, which is cumsum - probs >= p. The worked example shows a case where this decides the set.

Step 5: Renormalize and draw one token

After truncation, divide the surviving probabilities by their sum and draw one index with a random number generator. A multinomial draw with one sample is a weighted random choice: pick a uniform random number in zero to one and return the token whose slice of the cumulative sum contains it.

Input: the truncated, renormalized distribution. Output: one token index, which you append to the sequence before the next step; the logits for the next position come from running the model again on the longer sequence (a key-value (KV) cache avoids recomputing the earlier positions). Seed the generator when you need a reproducible run. Hugging Face wires these stages as separate logits processors, so the same chain exists in generate(); its documentation also warns that undefined must be set or the temperature value has no effect.

1 · Logits2 · Temperature, softmax3 · Top-k4 · Top-p5 · DrawReplay

Step-by-step description

1. The model outputs one raw score per token.
2. Scale the logits, then make probabilities.
3. Keep a fixed count of tokens.
4. Keep the smallest set that reaches p.
5. Renormalize and draw one token.

• Logits go through temperature, softmax, top-k, top-p and a random draw in that order.
• Top-p keeps the smallest set of tokens that reaches probability p.

The whole chain fits in one short function. It uses NumPy, applies temperature before softmax, builds both truncation masks from the sorted probabilities and renormalizes once before the draw.

import numpy as np

def softmax(x):
    x = x - x.max()                      # stability: subtract the max
    e = np.exp(x)
    return e / e.sum()

def sample_next_token(logits, temperature=1.0, top_k=None, top_p=None, rng=None):
    rng = rng or np.random.default_rng()
    if temperature < 0:
        raise ValueError("temperature must be >= 0")
    if top_k is not None and top_k < 1:
        raise ValueError("top_k must be >= 1")
    if top_p is not None and not 0 < top_p <= 1:
        raise ValueError("top_p must be in (0, 1]")
    if temperature == 0:                 # convention: temperature zero means greedy
        return int(np.argmax(logits))
    probs = softmax(np.asarray(logits, dtype=float) / temperature)
    order = np.argsort(-probs)           # token indices, most probable first
    keep = np.ones(len(probs), dtype=bool)
    if top_k is not None:
        keep[order[min(top_k, len(probs)):]] = False      # clamp k to the vocabulary
    if top_p is not None:
        before = np.cumsum(probs[order]) - probs[order]   # mass before each token
        keep[order[before >= top_p]] = False              # keeps the crossing token
    keep[order[0]] = True                # always keep the most probable token
    probs = np.where(keep, probs, 0.0)
    probs /= probs.sum()                 # renormalize once
    return int(rng.choice(len(probs), p=probs))

Worked example

Take six tokens A to F with illustrative logits 2.0, 1.0, 0.5, 0.0, -1.0 and -2.0. At temperature 1 the softmax gives A 0.5573, B 0.205, C 0.1243, D 0.0754, E 0.0277 and F 0.0102.¹ At temperature 0.5 token A rises to 0.829 and F falls to 0.0003; at temperature 2.0 token A falls to 0.3565 and F rises to 0.0482.¹

Top-k with k = 3 keeps A, B and C, whose combined mass is 0.8866, and renormalizing divides each by that sum.¹ For top-p with p = 0.9 the running totals in sorted order are 0.5573, 0.7623, 0.8866 and 0.962.¹ The third total is still below 0.9, so token D must be kept: the nucleus has four tokens. A mask that removes every token whose running total exceeds 0.9 would keep three tokens and a mass of only 0.8866, which is the off-by-one from step 4.¹

The nucleus also adapts. At temperature 0.5 the same p = 0.9 keeps 2 tokens, and the renormalized probabilities are about 0.88 and 0.12; at temperature 2.0 it keeps 5.¹ Top-k with k = 3 would keep three tokens in every case.

For the draw, a seeded random number of 0.3238 falls in the first slice of the renormalized top-p distribution, which spans 0 to 0.5793, so the sampled token is A.¹

Text description of the diagram

• Token A holds 0.829 of the probability at temperature 0.5 and 0.3565 at temperature 2.0.
• The order of the six tokens never changes.

Token A holds 0.829 of the probability at temperature 0.5 and 0.3565 at temperature 2.0, while the order of the six tokens never changes.¹

Trade-offs and edge cases

Truncated sampling sits between two failure modes. Holtzman et al. (2020) argue that maximization, such as greedy or beam search, is an inappropriate objective for open-ended text and can give bland or repetitive output, while sampling from the entire tail lets unreliable tokens in. Truncation removes that tail. For tasks with one right answer, such as translation or code with a test, beam search or greedy decoding often fits better; see beam search explained.

Text description of the diagram

• Greedy decoding keeps one token.
• Top-k keeps a fixed count of tokens.
• Top-p keeps the smallest set of tokens that reaches probability p.

Greedy decoding keeps one token and is deterministic, top-k keeps a fixed count, and top-p keeps the smallest set that reaches p, so only top-p adapts to the model's confidence.

Edge cases a reviewer probes:

• Temperature has no effect without sampling. In Hugging Face, do_sample defaults to False, which means greedy decoding, so changing temperature alone does nothing.
• Temperature zero. Divide by zero is undefined; return the argmax.
• Ties. Decide whether equal logits at the k-th place are all kept or cut, and say so.
• Order of operations. This page applies temperature first, then the cuts, then one renormalization. Libraries differ in the order and in whether they renormalize between top-k and top-p, so read the order in the library you use.
• Top-p equal to one keeps everything, and top-k larger than the vocabulary must be clamped.
• Never leave nothing to sample. Always keep at least the top token, so a tiny p or k = 0 cannot leave nothing to sample; the function above also rejects invalid values.
• Related controls. A repetition penalty (which lowers the logits of tokens already generated) and min-p sampling (a cutoff relative to the top probability) fit into the same chain; mention them if the interviewer asks, but they are not covered here.

How to answer it in an interview

The spoken structure: name the order of steps, say what each truncation keeps, flag the crossing-token bug and the zero-temperature convention, then say when you would not sample.

Common mistakes

• "I pass the logits to torch.multinomial." (wrong) It expects non-negative probabilities or weights, and logits can be negative. (right) "I apply softmax first, then sample."
• "Top-p keeps tokens whose probability is at most p." (wrong) That describes single-token probabilities. (right) "It keeps the smallest set whose cumulative probability reaches p."
• "I drop every token whose running total passes p." (wrong) That removes the crossing token and keeps too little mass. (right) "I keep a token unless the mass before it already reached p."
• "I set temperature=0.7 and the output stays the same." (wrong) Without do_sample=True the library decodes greedily. (right) "Sampling must be switched on, then temperature applies."
• "Top-k and top-p are the same." (wrong) One keeps a fixed count and the other a variable one. (right) "Top-p adapts to confidence; top-k does not."

Follow-up questions

What decoding and text-generation methods exist for a decoder?

Greedy decoding takes the highest-probability token, beam search keeps several candidate sequences, and sampling draws at random, with temperature, top-k and top-p as the usual controls. Hugging Face's guide describes multinomial sampling as randomly selecting a token from the distribution over the whole vocabulary.

How does nucleus sampling compare to top-k?

Both truncate the tail before the draw. Top-k uses a fixed count, which can be too many when the model is sure and too few when it is not; top-p's count changes with the context, which is the argument Holtzman et al. (2020) make for it.

Why can't you pass logits straight to torch.multinomial?

The function treats its input as probabilities or weights, which must be non-negative and finite with a non-zero sum, as the PyTorch documentation states. Logits can be negative, so apply softmax first.

What happens at temperature zero?

The division is undefined, so implementations treat it as greedy decoding. The limit makes sense: as the temperature shrinks, the top token's probability approaches one.

Why does temperature do nothing in my Hugging Face generate call?

Because do_sample is False by default, the call decodes greedily and ignores temperature. Set do_sample=True, as the generation utilities documentation says.

Where do the logits come from, and how does the mask relate?

The logits are the output of the decoder's last layer for the final position; the causal mask keeps each position from seeing the future during training, covered in causal masks and cross-attention and multi-head attention in PyTorch.

Sources

1. The Curious Case of Neural Text Degeneration (arXiv 1904.09751) — arXiv (Holtzman et al.; ICLR 2020), 22 Apr 2019
2. Hierarchical Neural Story Generation — ACL Anthology (Fan, Lewis, Dauphin), 15 Jul 2018
3. Text generation strategies — Hugging Face Transformers documentation, 26 Oct 2026
4. Utilities for generation (logits processors) — Hugging Face Transformers documentation, 26 Oct 2026
5. GenerationConfig parameters (v4.26.1 snapshot) — Hugging Face Transformers documentation, 1 Feb 2023
6. torch.multinomial — PyTorch documentation, 26 Oct 2026