Inference & Sampling

Your program will generate text the way every LLM does: ask the model for scores, choose one token, append it, ask again. You will write the loop and the two simplest ways to choose: always the best (greedy), or a random draw from the model's distribution, sharpened or flattened by a temperature.

The generation loop

tokens = the prompt
repeat up to MAX times:
    logits = model(tokens)[last position]       one score per vocabulary entry
    next   = choose(logits)
    append next to tokens
    stop if next is the end-of-sequence token

The model is called once per new token, and each call sees everything generated so far. In this exercise the "model" is a stand-in: a table with one row of logits per token id, and the row used is the one for the last token. A real transformer reads the whole context, but the loop around it is exactly this one. In lesson 14 you replace the table with the block stack you built.

Two small rules matter. The stop token is checked on what was generated: an end-of-sequence token that is already in the prompt does not end generation. And MAX bounds the number of new tokens, so MAX 0 returns the prompt unchanged.

Greedy: pick the largest

The simplest choice is the highest score. It is deterministic and cheap, and it tends to get stuck: the likeliest continuation of a repeated phrase is more of the phrase. When two scores are equal, take the lower id so the result is reproducible.

Temperature: reshape before sampling

To sample, turn the logits into probabilities with softmax, but divide the logits by a temperature T first:

p = softmax( logits / T )

For logits [2, 1, 0]:

T = 1    ->  [0.6652, 0.2447, 0.0900]     the model's own distribution
T = 0.5  ->  [0.8668, 0.1173, 0.0159]     sharper: the favourite takes more
T = 2    ->  [0.5065, 0.3072, 0.1863]     flatter: closer to a coin toss

As T goes to 0 the distribution collapses onto the maximum, which is greedy; T = 0 itself would divide by zero, so treat it as the greedy special case. As T grows, every token approaches probability 1/V.

Drawing from a distribution

A random number generator gives a number u in [0, 1). To turn it into a token, lay the probabilities end to end on [0, 1) and see where u lands: the answer is the first index whose running total exceeds u.

p   = [0.6652, 0.2447, 0.0900]
cum = [0.6652, 0.9099, 1.0000]
u = 0.30 -> token 0      u = 0.70 -> token 1      u = 0.95 -> token 2

Rounding can leave the final running total a hair below 1, so fall back to the last index if no total exceeds u. To keep this exercise reproducible, the draws arrive as input: one u per sampled step, taken in order. A greedy run needs none.

Why generation is slow, and what real systems add

Every step recomputes attention over the whole context, so a naive loop costs quadratic work in the length. The KV cache keeps the keys and values of past tokens (they never change) and computes only the new token's, at the price of memory: layers x heads x length x head width x 2 numbers. Speculative decoding lets a small draft model propose several tokens that the big model verifies in one pass. The chat context window is the limit of the position table from lesson 2.

Common mistakes

• Dividing the probabilities by T instead of the logits, or applying T after the softmax.
• Using <= against the running total, or comparing u with the probability of each token instead of the running total.
• Reusing one draw for every step, or reading the draws in the wrong order.
• Stopping when the last prompt token is the stop token, or not stopping when a generated one is.

Your exercise

Write pick (greedy or temperature draw for one step) and generate (the loop). The starter parses the table, the prompt and the draws.