# Module 12: Generation

> The autoregressive decoding loop, how the KV cache avoids recomputation and what it costs, and greedy, beam, temperature, top-k and nucleus sampling.

- Part III: Building an LLM · core track · Intermediate · 3 hours
- By the end you can implement a decoding loop with a KV cache and choose a decoding strategy for a task.
- Builds on: [Module 8: Three Families: Encoder, Encoder-Decoder, Decoder](https://transformers-and-math.drmhse.com/md/08-three-families.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/12-generation
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

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 Loop

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.

```python
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 Cache

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:

$$
\operatorname{Time}_{\text{naive}}(T) = \sum_{t=1}^T \mathcal{O}(t^2) = \mathcal{O}(T^3) \text{ operations!}
$$

$\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**.

```mermaid
graph LR
    subgraph StepT["Generation step t"]
        TokenT["Newest token x_t"] --> Q["query q_t = x_t W^Q"]
        TokenT --> K["key k_t = x_t W^K"]
        TokenT --> V["value v_t = x_t W^V"]
        K --> AppendK["Append k_t to K_cache"]
        V --> AppendV["Append v_t to V_cache"]
        Q --> Attn["Attend: q_t against all of K_cache, weights applied to V_cache"]
        AppendK --> Attn
        AppendV --> Attn
        Attn --> Out["Output y_t, used to choose token x_{t+1}"]
    end
```

> **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:

$$
\mathbf{q}_t = \mathbf{x}_t W^Q \in \mathbb{R}^{1 \times d_k}
$$

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

$$
K_{\text{cache}}^{(t)} = \left[ K_{\text{cache}}^{(t-1)}; \, \mathbf{x}_t W^K \right] \in \mathbb{R}^{t \times d_k}
$$

$$
V_{\text{cache}}^{(t)} = \left[ V_{\text{cache}}^{(t-1)}; \, \mathbf{x}_t W^V \right] \in \mathbb{R}^{t \times d_v}
$$

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

$$
\mathbf{y}_t = \operatorname{Softmax}\left(\frac{\mathbf{q}_t (K_{\text{cache}}^{(t)})^\top}{\sqrt{d_k}}\right) V_{\text{cache}}^{(t)} \in \mathbb{R}^{1 \times d_v}
$$

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$.

> **Interactive figure** `kv-cache-steps`, on the web page: https://transformers-and-math.drmhse.com/modules/12-generation#figure-kv-cache-steps
>
> **Try:** Without a cache, how many attention scores are computed in total to generate 12 tokens, if step $t$ computes all $t \times t$ scores of the causal triangle (counting only the allowed ones)? With a cache, each step computes just one new row.
>
> **What happens:** Without a cache, step $t$ computes $t(t+1)/2$ scores, and summing over 12 steps gives 364. With a cache, step $t$ computes only its $t$ new scores, $1 + 2 + \dots + 12 = 78$ in total. The gap grows like $T^3$ against $T^2$.

> **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 Text

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)
   $$
   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)
   $$
   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.

> **Interactive figure** `decoding`, on the web page: https://transformers-and-math.drmhse.com/modules/12-generation#figure-decoding
>
> **Try:** Set top-p to 0.9 and compare the two prompts. Before you look, predict which prompt keeps more candidate tokens, and what that says about why top-p was invented.
>
> **What happens:** The open-ended prompt keeps many tokens; the prompt with one obvious answer keeps one or two. Top-p adapts the candidate set to how unsure the model is, which a fixed top-k cannot do.

> **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

**Exercise (check).** Why can the prompt be processed in one parallel pass, but the generated tokens cannot?

*Solution:* All the prompt's tokens are known in advance, so the causal mask lets every position be computed at once, as in training. Each generated token depends on the one chosen before it, which does not exist until the previous step has finished, so generation steps must run in sequence.

**Exercise (apply).** Without a cache, how many key/value projections does it take to generate 24 tokens after a prompt of 1 token, if step $t$ projects all $t$ tokens? (Count one projection per token per step.)

*Answer:* 300 (within 0)

*Solution:* The steps process $1, 2, \dots, 24$ tokens, so $1 + 2 + \dots + 24 = 24 \cdot 25 / 2 = 300$ projections. With a cache, each token is projected once: 24.

**Exercise (apply).** A next-token distribution is $[0.5, 0.3, 0.1, 0.05, 0.05]$ (already sorted). Which tokens does top-p keep at $p = 0.9$, and what is the renormalized probability of the first one?

*Answer:* 0.5556 (within 0.001)

*Solution:* The running totals are $0.5, 0.8, 0.9$, so the smallest set reaching 0.9 is the first three tokens. Their total is 0.9, so the first one's renormalized probability is $0.5 / 0.9 \approx 0.556$.

**Exercise (apply).** With a repetition penalty of $\theta = 1.25$, what does a logit of $2.0$ become, and what does a logit of $-2.0$ become, for a token that has already appeared?

*Solution:* Positive logits are divided: $2.0 / 1.25 = 1.6$. Negative logits are multiplied: $-2.0 \times 1.25 = -2.5$. Both moves make the token less likely.

**Exercise (extend).** Explain why comparing the generated text of a cached and an uncached model is a poor test of a KV cache, and describe a better one.

*Solution:* The two paths can differ in the last bits of a floating-point number. When two candidate tokens are nearly tied, that tiny difference can flip the choice, and from then on the two generations follow different paths, so the text differs even though the cache is correct. The better test feeds both paths the same tokens (teacher forcing) and compares their logits at every position; a correct cache agrees to within floating-point rounding.

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