Write the sampling step of a generation loop
RolesLLM Product EngineeringInference EngineeringDeveloper Experience and Education
The question
Given the logits of the last position, a temperature, a value of k and a seeded random generator, write the function that draws the next token. State the order of the operations, how ties are broken, and how the draw uses exactly one random number.
Answer it first, aloud or on paper. Then open the answer.
The answer
Show short answer
Divide the logits by the temperature, keep the k largest and set the rest to negative infinity, take the softmax, and draw one index from the result. Ties keep the lower index. The draw calls the generator once, takes the cumulative sum of the probabilities in vocabulary order, and returns the first index whose cumulative value exceeds the random number.
Show full answer
The code
def top_k_filter(logits, k):
if k >= logits.shape[0]:
return logits.copy()
order = np.argsort(-logits, kind="stable")
keep = order[:k]
out = np.full_like(logits, -np.inf)
out[keep] = logits[keep]
return out
def sample(logits, rng, temperature=1, top_k=None):
z = logits / temperature
if top_k is not None:
z = top_k_filter(z, top_k)
probs = softmax(z)
u = rng.random()
c = np.cumsum(probs)
hits = np.nonzero(u < c)[0]
if hits.size:
return int(hits[0])
return int(np.nonzero(probs > 0)[0][-1])The order
Temperature first, then the filter, then the softmax. The filter has to see the scaled logits, because the scale changes which tokens hold the mass when a top-p cutoff is added later. The softmax runs last so that a removed token has probability exactly zero and the kept ones renormalize.
Ties
Sorting the negated logits with a stable sort keeps the lower index first among equal values. Sorting ascending and reversing the result flips that, so the higher index wins and the output differs from the reference on every tie.
One random number
A seeded generator gives reproducible samples only if every call consumes the same number of draws. The function reads one uniform number, compares it with the cumulative sum, and returns the first crossing. Rounding can leave the last cumulative value a hair below one, so the fallback returns the last token with positive probability instead of failing.
Do not change the input
The filter writes into a new array. Writing negative infinity into the caller's logits destroys the values that the next step, or a log, still needs.
Follow-ups
L2WorkingWhat does top-k one produce?
Show expected answer
greedy decoding, the argmax every time.
L2WorkingWhy must the filter run before the softmax and not after?
Show expected answer
zeroing probabilities after the softmax leaves a distribution that does not sum to one unless it is renormalized; masking the logits does that for free.
L3SeniorThe assembly lesson keeps the context inside a fixed length during generation. How?
Show expected answer
when the prompt plus generated tokens exceed the context length, the oldest tokens are dropped so the active window fits.
L3SeniorHow do you make two runs with one seed produce the same text?
Show expected answer
the same generator, one draw per step, the same tie rule and the same temperature; any extra draw shifts every later token.
Wrong answers
Show the wrong answers that interviewers hear
“Call the generator's choice function with the probabilities.”
It works, but it consumes draws differently across versions and breaks the one-draw contract that the tests pin.
“Apply the temperature after the filter.”
The filter then selects tokens at temperature one and the setting changes less than intended.
“Sort descending by reversing an ascending sort.”
Equal logits then keep the higher index, which contradicts the tie rule.
Taught in
The GPT assembly lesson describes the generation step as temperature, top-k mask, softmax and multinomial draw, and the sampling practice pins the order, the tie rule and the single draw.