TL;DR: Divide logits by temperature, softmax in float64, sort descending, find the first index where the cumulative mass reaches p, and keep every token whose probability is at least the probability at that index, so tied tokens are kept or dropped together. Short-circuit p ≥ 1 and temperature 0, because a float32 cumulative sum over a 128,256-token vocabulary ended at 0.99974 in our runs and never reaches 1.
How to approach it
Write the definition first, in a form that does not mention sorting: keep token i if and only if the total probability of tokens strictly more likely than i is below p. That one sentence fixes the tie behavior (tied tokens have the same mass above them, so they share a fate), gives the test oracle, and shows that the kept set always contains the argmax. Then ask what the caller expects for temperature 0 (greedy is the usual contract), and whether p arrives per request, which it does in every serving engine, so the code is batched with a vector of p values.
A strong answer
import numpy as np
def top_p_filter(logits, temperature, top_p):
"""logits [B, V]; temperature, top_p [B]. Returns renormalized probs [B, V].
temperature 0 means greedy. Keeps every token tied with the cutoff token."""
logits = np.asarray(logits, dtype=np.float64)
B, V = logits.shape
out = np.zeros((B, V))
greedy = temperature <= 0
if greedy.any(): # argmax takes the lowest index on ties
out[greedy, logits[greedy].argmax(-1)] = 1.0
rows = np.flatnonzero(~greedy)
if rows.size == 0:
return out
z = logits[rows] / temperature[rows, None]
z = z - z.max(-1, keepdims=True)
p = np.exp(z)
p /= p.sum(-1, keepdims=True)
order = np.argsort(-p, axis=-1, kind="stable")
sp = np.take_along_axis(p, order, -1) # descending
cum = np.cumsum(sp, -1)
tp = top_p[rows, None]
k = np.minimum((cum < tp).sum(-1), V - 1) # first index with cum >= p
thresh = np.take_along_axis(sp, k[:, None], -1) # prob of the cutoff token
keep = (p >= thresh) | (tp >= 1.0) # ties kept; p=1 keeps all
p = np.where(keep, p, 0.0)
out[rows] = p / p.sum(-1, keepdims=True)
return out
def sample(probs, rng):
"""Inverse-CDF sampling, one token per row."""
u = rng.random((probs.shape[0], 1))
idx = (np.cumsum(probs, -1) < u).sum(-1)
last = probs.shape[1] - 1 - np.argmax(probs[:, ::-1] > 0, -1) # last nonzero
return np.minimum(idx, last)
Why each line is there:
- Temperature before softmax, in float64. Scaling probabilities instead of logits is a different and wrong operation. Subtracting the row max keeps
expfinite at temperature 0.5 on logits near 30. - The threshold, not the index. The usual implementation zeroes sorted positions after the cutoff index. With probabilities
[0.4, 0.2, 0.2, 0.2]and p = 0.5 that keeps the 0.4 token and whichever 0.2 token the sort placed first, so the output depends on token ids. Comparing every token to the cutoff probability keeps all three tied tokens, and the kept set matches the order-free definition. - p = 1 and the float sum. Summing sorted float32 probabilities over Llama 3's 128,256-token vocabulary, the final cumulative value came out between 0.99974 and 0.99975 for five random rows. A loop that looks for
cum >= pwith p = 1 never finds it and indexes past the end. Themin(..., V - 1)clamp and the explicittp >= 1keep-all handle it; float64 shrinks the drift but does not make it zero. - Greedy is a separate branch. Dividing by a temperature of 0 produces inf and NaN. Greedy also has its own tie rule (lowest index), which matters when two requests must agree.
- Masked tokens. A logit of
-inf(a banned token or a grammar mask) becomes probability 0 and can never cross the threshold, since the threshold is positive.
Tested, not asserted: 3,000 random batches of 4 rows, vocabularies of 1 to 39 tokens, temperatures 0, 0.5, 1 and 2 and p values from 0.05 to 1.0, with a third of the batches drawn from a five-value integer logit set so ties are common. 8,970 sampled rows were compared token by token against the order-free definition above, with renormalized probabilities matching to 1e-12. 613 of 178,800 tokens sat exactly on the boundary (mass above equal to p within 1e-9) and were excluded, because the definition itself is ambiguous there under floating point. Separate checks: the tied example keeps all three tied tokens whichever position the larger one occupies, p = 1 keeps all 128,256 tokens, -inf tokens get zero probability, an empty batch returns shape (0, V), and 400,000 inverse-CDF samples from a filtered 5-token distribution matched it with total variation distance 0.00009 and never produced a filtered-out token.
example: p = [0.30, 0.25, 0.25, 0.12, 0.08], top_p = 0.6
sorted cumsum = 0.30, 0.55, 0.80 first index with cum >= 0.6 is 2
cutoff prob = 0.25 -> keep {0.30, 0.25, 0.25}, mass 0.80
renormalized = 0.375, 0.3125, 0.3125
Note that the kept mass, 0.80, exceeds p. That is correct: top-p keeps the smallest set reaching p, and a tie can push it past.
The reversal condition: if the contract you must reproduce defines top-p by sort position (keep the first k sorted tokens, whatever their ties), match it exactly, because a reproducibility test against it will fail otherwise. The tie-inclusive definition is the one that does not depend on token ids, which is why it is the better default for a fresh implementation.
What interviewers probe next
"Sorting 128K entries per row per step is slow on a GPU. Avoid it." Find the threshold by binary search on the probability value: count the mass above a candidate threshold with one reduction, bisect about 30 times, then mask. That is O(V) per iteration with no sort, and it gives the tie-inclusive set by construction.
"Add top-k and min-p in the same function." Apply them as further masks on the same probabilities (top-k by the k-th largest value, again tie-inclusive; min-p by p >= min_p * p.max()), then renormalize once at the end. The order of top-k and top-p matters only when both bind, so document it.
"Make two runs with the same seed produce the same tokens." Use a per-request generator seeded from the request, not a global one, so batch composition does not change which uniform each row draws.
How it's scored
Weak: applies temperature to probabilities, zeroes positions after the sorted cutoff so ties depend on token id, and crashes or drops tokens at p = 1.
Hire: temperature on logits, stable softmax, a vectorized cutoff per row, the greedy branch, and renormalization.
Strong Hire: states the order-free definition, uses it as the test oracle, explains the float cumulative-sum drift with a number, handles masked tokens and empty batches, and proposes the sort-free threshold search for GPUs.
Common mistakes
- Dropping the token that crosses p instead of keeping it, which can leave an empty set when the top token alone exceeds p.
- Comparing
cum > pon the unshifted cumulative sum, which drops the crossing token and is off by one from the definition. - Sharing one random generator across the batch, so adding an unrelated request changes another request's sample.
- Testing only on distinct random logits, where ties never occur and the bug is invisible.