Generation

A trained decoder-only model answers one question: given these tokens, how likely is each next token? Generation turns that into text, one token at a time. This module covers the loop, the cache that stops it redoing work, and the rules for choosing each token. Module 23 returns to the loop when serving many users at once.

12.1 The Autoregressive LoopRun this section

Generation is a loop around the model:

1. Tokenize the prompt.
2. Run the model on the tokens so far and take the logits at the last position: the scores for the next token.
3. Choose one token from them, with one of the decoding rules of section 12.3.
4. Append it to the sequence.
5. Stop if the token is the special end-of-sequence token, or a length limit is reached, or a stop string has appeared. Otherwise go back to step 2.

tokens = tokenizer.encode(prompt)
for _ in range(max_new_tokens):
    logits = model(tokens)[-1]          # scores for the next token only
    next_token = choose(logits)         # greedy, top-k, top-p ... (section 12.3)
    tokens.append(next_token)
    if next_token == tokenizer.eos_id:
        break
text = tokenizer.decode(tokens)

Two features drive the rest of the module. Each step depends on the previous token, so steps run strictly in order. And as written, step $t$ reruns the model over all $t$ tokens though only the last is new; section 12.2 removes that waste.

The first pass is different: the prompt's tokens are all known, so they are processed in one parallel pass, as in training. This is prefill; the one-token steps after it are decode. Section 23.2 shows why their costs differ so much.

12.2 The KV CacheRun this section

Rerunning the whole sequence at step $t$ costs on the order of $t^2$ operations, since attention compares every token with every other. Over $T$ steps:

$\mathcal{O}(\cdot)$ ("big O") describes how cost grows, ignoring constant factors. Doubling the output length makes generation about eight times slower.

Most of that work is repeated. Under the causal mask, earlier tokens never see later ones, so the keys and values of tokens $1 \dots t-1$ are the same at every step. Computed once and stored, they form the KV cache.

In words. At step t, only the newest token becomes a query, key and value. The key and value are appended to the cache, and the one new query attends over all of it.

Mathematical KV Cache Derivation

At step $t$, compute a query for the newest token only:

Append its key and value as new rows of the caches (the semicolon stacks the new row under the old):

The superscript $(t)$ labels the step, not a power. The output is ordinary attention, one query row against the whole cache:

This is exactly the $\mathbf{y}_t$ of full recomputation, with no approximation. Each step projects one token instead of $t$ and computes one row of scores instead of $t \times t$, so $T$ tokens cost on the order of $T^2$ instead of $T^3$.

In words. Twelve generation steps, shading the scores each computes. Without the cache, every step rebuilds the whole triangle and re-projects every key and value; with it, each step adds one row and one key and value. That is 78 scores instead of 364.

The new token must also get its real position (index $t-1$ counting from 0), not position 0. The reference positional_embedding.py numbers its input from 0 on every call, so with a cache the caller must pass the offset; RoPE (section 15.1) needs the same. A cache that forgets the offset gives every new token the first token's position, and output quietly degrades.

Measured (lab 3a). A cache built on this repository's own MultiHeadAttention, reusing its weights and changing no projection, reproduces the uncached logits position by position. The largest absolute difference is $7.15 \times 10^{-7}$, and the top token agrees at every position ($24/24$). Over a 24-token decode, K/V projections fall from $468$ to $31$, a $15.1\times$ reduction. Wall-clock time improves by only $1.56\times$ at 112k parameters, because in a model this small the fixed cost of each call outweighs the arithmetic. The projection count is the measure that carries over to larger models.

Caveat. Do not test a cache by comparing generated text. The two paths can differ in the last bit of a float, and when two tokens nearly tie, that picks a different token and the outputs diverge: a tie-break, not a bug. Instead feed both paths the same tokens and compare logits at each position (teacher forcing). A difference there is real, usually because the cached path does not mirror the block exactly, such as normalizing in a different place (section 8.5).

The cache trades arithmetic for memory: one key and value per token, layer and head. For long contexts and many users, that memory is the main cost of serving (section 23.1).

12.3 Decoding: From Logits to TextRun this section

A decoding strategy chooses the next token from the logits $\mathbf{z}$. Most also apply a temperature (section 4.3) first: low sharpens toward the top token, high flattens.

1. Greedy decoding. Always pick $w_{\text{next}} = \arg\max_i z_i$. Deterministic, so often used for code and maths, where there is usually one right answer, but it tends to loop and its prose is flat.

2. Top-$k$ sampling. Keep the $k$ highest-scoring tokens (say $k=50$), set the other logits to $-\infty$, softmax and sample. This cuts the long tail of unlikely tokens, where most nonsense comes from.

3. Top-$p$ (nucleus) sampling (Holtzman et al., 2020). Keep the smallest set $V^{(p)}$ of top tokens whose probabilities reach $p$, renormalize and sample:

   $$\sum_{i \in V^{(p)}} P(w_i) \ge p \quad (\text{typically } p = 0.90)$$

   What does this say?

   The set adapts: it grows when the model is unsure and shrinks to one or two tokens when the next word is obvious, which a fixed $k$ cannot do.

4. Repetition penalty (Keskar et al., 2019). Every token already in the output gets a less favorable logit, with a strength $\theta$ slightly above 1 (a chosen number, not the model parameters):

   $$z_i \longleftarrow \begin{cases} z_i / \theta, & \text{if } z_i > 0 \\ z_i \cdot \theta, & \text{if } z_i < 0 \end{cases} \quad (\theta > 1.0, \text{e.g. } \theta = 1.15)$$

   What does this say?introduces 2 notations

   Positive logits shrink toward zero and negative ones grow more negative: either way the token becomes less likely, discouraging loops.

5. Beam search. Keep the several most probable partial outputs (a beam of 4 or 5), extend each by every token, and keep the best by total log-probability. It suits tasks with one right answer, such as translation, but open-ended text comes out bland, because the most probable text is rarely the most natural, so chat models sample instead. Summed log-probabilities favor short outputs, so scores are divided by a power of the length.

In words. The ten most likely next tokens for two prompts. With one right answer, top-p keeps just that token; for an open-ended prompt the same p keeps many. Top-k keeps the same number in both, greedy always one, and a higher temperature makes every rule keep more.

Measured (lab 2b and tests/test_generation.py). The tests check three properties of the nucleus sampler in src/infra/generation.py. The kept set is the smallest whose total probability reaches $p$. The kept probabilities sum to 1 within $10^{-6}$. The top token is always kept, even as $p \to 0$, so torch.multinomial never receives a row of all zeros. The tests also confirm that sampling becomes greedy decoding as the temperature goes to $0^{+}$, and that raising the temperature increases the number of distinct 8-token completions produced from 40 seeds.

Exercises

Core Academic References

• Holtzman et al. (2020): The Curious Case of Neural Text Degeneration (nucleus sampling). ICLR 2020. arXiv:1904.09751
• Keskar et al. (2019): CTRL: A Conditional Transformer Language Model for Controllable Generation (repetition penalty). arXiv:1909.05858
• Pope et al. (2023): Efficiently Scaling Transformer Inference on TPU v4 (KV caching). MLSys 2023. arXiv:2211.05102
• Freitag & Al-Onaizan (2017): Beam Search Strategies for Neural Machine Translation. arXiv:1702.01806