A transformer decoder from scratch: attention, RoPE, KV cache and decoding

Before this lesson: lesson 07 (the from-scratch routine in PyTorch). The theory of each piece: attention, positional encodings (RoPE), normalization and residual connections and the transformer block in the Deep Learning track; Linear Algebra 15 (attention in matrix form); the categorical distribution from the Probability and Statistics track. Time: about 105 minutes, plus the notebook.

---

Why this matters #

"Implement multi-head attention" is one of the most frequent prompts in ML coding rounds for LLM roles. The usual follow-ups are: add a causal mask, add a KV cache and generate, add RoPE, replace LayerNorm by RMSNorm, then sample with top-p or run beam search. Stanford CS336, assignment 1, asks for RMSNorm, RoPE and causal multi-head attention. The round tests whether you can write it fast, with the right shapes, and show that it works. The bugs are silent: a reshape that mixes tokens, a mask aligned to the wrong corner, positions that restart at 0, an unbiased variance, a top-p rule that is off by one. Each one gives a tensor of the right shape with wrong numbers. The same block runs in vision transformers and in scientific foundation models over patch tokens, and the KV cache is the starting point of every LLM serving question (lesson 13).

1. Predict first #

2. The picture #

2.1 One head on one sequence #

There are $T$ tokens with vectors $x_1, \dots, x_T \in \mathbb{R}^D$. Each token makes a query $q_i$, a key $k_i$ and a value $v_i \in \mathbb{R}^{d_h}$, by three linear maps. Token $i$ scores every key, $s_{ij} = q_i^\top k_j / \sqrt{d_h}$. A softmax over $j$ (over the keys) turns the scores into weights $a_{ij} \ge 0$ that sum to 1, and the output of token $i$ is the weighted average $o_i = \sum_j a_{ij} v_j$. In a decoder, token $i$ may only look at keys $j \le i$: this is the causal mask.

Interactive figure: loading…

2.2 A batch and several heads change only the shapes #

Notation: $B$ is the batch size, $T$ the sequence length, $D$ the model width, $H$ the number of heads, and $d_h = D/H$ the head size. The heads are one more batch axis, and @ batches over all leading axes, so there is no Python loop over heads:

Step	Code	Shape
input	x	$(B, T, D)$
fused projection	qkv = x @ W_qkv.T	$(B, T, 3D)$
split	q, k, v = qkv.split(D, dim=-1)	$(B, T, D)$ each
head split	q.view(B, T, H, d_h).transpose(1, 2)	$(B, H, T, d_h)$
scores	q @ k.transpose(-2, -1) / sqrt(d_h)	$(B, H, T, T)$
mask, then softmax over keys	.masked_fill(~mask, -inf).softmax(-1)	$(B, H, T, T)$
weighted values	probs @ v	$(B, H, T, d_h)$
merge heads	.transpose(1, 2).reshape(B, T, D)	$(B, T, D)$
output projection	@ W_o.T	$(B, T, D)$

The head split is the step that goes wrong most often. The $D$ features of one token are next to each other in memory. So the split must first cut the last axis into $(H, d_h)$, and only then move the head axis forward.

Interactive figure: loading…

2.3 The other pieces each change one row of the table #

• LayerNorm or RMSNorm rescales each token's $D$ features before attention and before the feed-forward network (the pre-norm block).
• RoPE rotates $q$ and $k$ after the head split, before the scores.
• The KV cache keeps $k$ and $v$ of past tokens (after RoPE). A new token computes only its own $q, k, v$, appends its $k, v$ to the cache, and attends over all cached keys.

RoPE, as a picture. Cut the query and the key into pairs of numbers. Turn each pair by an angle that grows with the token's position. A query at position $m$ is turned by $m\theta$, a key at position $n$ by $n\theta$. The dot product depends only on the angle between the two vectors, and that angle changes by $(n - m)\theta$. So the score sees only the offset $n - m$.

Interactive figure: loading…

The KV cache, as a picture. Generation runs the prompt once (the prefill), then adds one token per step (decode). In a decode step only the new rows of the attention matrix are computed; the keys and values of the old tokens come from the cache. The mask for the new rows must be aligned to the bottom-right corner of the score matrix, because the new queries are the last positions.

Interactive figure: loading…

3. Definitions and main results #

3.1 Attention with a boolean mask #

For $Q \in \mathbb{R}^{T_q \times d_h}$, $K \in \mathbb{R}^{T_k \times d_h}$, $V \in \mathbb{R}^{T_k \times d_h}$ and a boolean mask $M \in \{\text{True}, \text{False}\}^{T_q \times T_k}$:

• Convention. In F.scaled_dot_product_attention (SDPA) and in CS336, True means "may attend". In nn.MultiheadAttention, True in attn_mask or key_padding_mask means "blocked". Write your convention in the docstring.
• Scale by $\sqrt{d_h}$, the head size, not by $\sqrt{D}$. SDPA's default scale is $1/\sqrt{E}$ with $E$ = q.size(-1), which is $d_h$ after the head split.
• Mask before the softmax. Then the masked weights are exactly 0, and the visible weights of each row still sum to 1.
• A row with every key masked gives $0/0$ = NaN in your own softmax (some PyTorch versions return zeros instead; do not rely on either). Make sure every query sees at least one key; this matters for left padding.
• Padding mask: shape $(B, 1, 1, T_k)$, broadcast over heads and queries; combine it with the causal mask by a logical AND.

Why $1/\sqrt{d_h}$. If the entries of $q$ and $k$ are independent with mean 0 and variance 1, then each product $q_i k_i$ has mean 0 and variance $\mathbb{E}[q_i^2]\,\mathbb{E}[k_i^2] = 1$, and the $d_h$ products are uncorrelated. So $q^\top k = \sum_{i=1}^{d_h} q_i k_i$ has variance $d_h$: for $d_h = 128$ the scores have a standard deviation of about 11. A softmax over such large numbers is almost one-hot, and its gradient is almost 0. Dividing by $\sqrt{d_h}$ brings the variance back to 1.

3.2 Multi-head and grouped-query attention #

• Parameters: $W_Q, W_K, W_V, W_O \in \mathbb{R}^{D \times D}$, so $4D^2$ (plus $4D$ with biases) for every $H$. Heads split the width; they add no parameters.
• FLOPs per layer for $T$ tokens (2 FLOPs per multiply–add): projections $2T \cdot 4D^2 = 8TD^2$; $QK^\top$ and $AV$ each $H \cdot 2T^2 d_h = 2T^2 D$.
• The head split must be view(B, T, H, d_h) and then transpose(1, 2). view(B, H, T, d_h) gives the same shape but cuts the wrong pieces (figure above); outputs can even depend on future tokens.
• The merge: after transpose(1, 2) the tensor is not contiguous in memory, so use reshape (or .contiguous().view).
• Grouped-query attention (GQA, Ainslie et al. 2023). $n_h$ query heads share $n_{kv} < n_h$ key/value heads, in groups of $g = n_h / n_{kv}$. Query head $h$ uses KV head $\lfloor h/g \rfloor$: in code, k.repeat_interleave(g, dim=1) (same for $v$), or SDPA with enable_gqa=True. Multi-query attention (MQA) is $n_{kv} = 1$. The $K$ and $V$ projections shrink to $2 D\, n_{kv} d_h$ parameters, and the KV cache shrinks by the factor $g$. Llama 2 70B and Llama 3 8B use 8 KV heads.

3.3 LayerNorm, RMSNorm and the pre-norm block #

For one token $x \in \mathbb{R}^D$, with gain $\gamma \in \mathbb{R}^D$, shift $\beta \in \mathbb{R}^D$ and a small $\varepsilon > 0$:

Here $\mathbf 1$ is the all-ones vector and $\odot$ the elementwise product.

• Normalize over the last axis only. LayerNorm has $2D$ parameters, RMSNorm $D$.
• The variance is the biased one (divide by $D$). torch.var divides by $D - 1$ by default; use correction=0.
• $\varepsilon$ goes inside the root: $10^{-5}$ in nn.LayerNorm and CS336, $10^{-5}$ or $10^{-6}$ in LLaMA models. Pass $\varepsilon$ explicitly in a test; library defaults differ.
• Compute in fp32 and cast back. In fp16, $x_i^2$ overflows (above 65,504) once $|x_i|$ is above about 256; in bf16 the mean of squares loses precision.

Fact. With $\gamma = \mathbf 1$, $\beta = 0$ and the same $\varepsilon$: $\operatorname{LN}(x) = \operatorname{RMSNorm}(x - \mu\mathbf{1})$.

Proof.

1. The mean square of $x - \mu\mathbf 1$ is $\frac1D\sum_i (x_i - \mu)^2 = \sigma^2$.
2. So $\operatorname{RMSNorm}(x - \mu\mathbf 1) = (x - \mu\mathbf 1)/\sqrt{\sigma^2 + \varepsilon} = \operatorname{LN}(x)$. $\blacksquare$

So RMSNorm is LayerNorm without the centering and without $\beta$: one reduction instead of two, with similar quality. LLaMA, T5 and CS336 use it.

The pre-norm block. With $h$ the hidden states $(B, T, D)$:

and one final norm before the output projection to the vocabulary. The FFN (feed-forward network) acts on each token separately (in LLaMA a SwiGLU MLP). The residual path $h \to h$ is never normalized, which makes deep stacks train stably.

3.4 RoPE #

The head size $d_h$ is even, and $\Theta$ is the base (10,000 in LLaMA 1 and 2, 500,000 in LLaMA 3). The frequencies are $\omega_i = \Theta^{-2i/d_h}$ for $i = 0, \dots, d_h/2 - 1$. A query or key at position $m$ is cut into $d_h/2$ pairs $(a_i, b_i)$, and pair $i$ is turned by the angle $\phi = m\omega_i$:

As a matrix, $x \mapsto R_m x$, where $R_m$ is block diagonal with the $2 \times 2$ rotations $R(m\omega_i)$.

Theorem (relative positions). For all $q, k \in \mathbb{R}^{d_h}$ and positions $m, n$: $\langle R_m q, R_n k \rangle = q^\top R_{n-m}\, k$. The score depends on $q$, $k$ and $n - m$ only.

Proof.

1. A $2 \times 2$ rotation satisfies $R(\alpha)^\top = R(-\alpha)$ and $R(\alpha)R(\beta) = R(\alpha + \beta)$.
2. Block-diagonal matrices multiply block by block, so $R_m^\top R_n$ has the blocks $R(-m\omega_i)R(n\omega_i) = R((n - m)\omega_i)$; that is, $R_m^\top R_n = R_{n-m}$.
3. $\langle R_m q, R_n k\rangle = (R_m q)^\top R_n k = q^\top R_m^\top R_n k = q^\top R_{n-m} k$. $\blacksquare$

Rules for the code that follow from the theorem:

• RoPE goes on $q$ and $k$, never on $v$. Rotated values would carry absolute positions into the output, and these rotations do not cancel. Every head uses the same rotation.
• Positions are absolute token positions. A common shift of $m$ and $n$ cancels, but $q$ and $k$ must use the same origin. $R_0 = I$, and rotations keep norms.
• Build the cos and sin tables once, in fp32, with shape $(\text{max\_len}, d_h/2)$, and index their rows with a positions tensor. Near position 4000 the bf16 spacing is 16, so position 4001 would be stored as 4000.
• Two layouts. Interleaved pairs $(2i, 2i+1)$: RoFormer, Meta's LLaMA code, CS336. Half-split pairs $(i, i + d_h/2)$: GPT-NeoX and the Hugging Face LLaMA code (rotate_half). They differ by one fixed permutation of the coordinates of $q$ and $k$. Scores do not change if you permute the rows of $W_Q$ and $W_K$ in the same way, which is what the Hugging Face conversion script does. Load weights with the wrong layout and the model runs, but its outputs are garbage.

3.5 The KV cache #

Prefill: run the prompt through the model in one pass and keep $k$ and $v$ of every layer. Decode: one new token per step; compute $q, k, v$ for it only, append $k, v$ to the cache, and attend over all cached keys.

Theorem (a causal decoder can be cached). In a decoder-only transformer with causal masks, let $h_j^{(\ell)}$ be the hidden state of token $j$ after layer $\ell$. Then $h_j^{(\ell)}$ depends only on tokens $1, \dots, j$. So the keys and values of a token never change once they are computed, and decoding with a cache gives the same outputs as a full forward pass over the whole sequence (up to float rounding).

Proof (induction on $\ell$).

1. $\ell = 0$: $h_j^{(0)}$ is the embedding of token $j$.
2. Assume the claim for layer $\ell - 1$. In layer $\ell$, $q_j$, $k_j$, $v_j$ are per-token maps of $h_j^{(\ell-1)}$ (a norm, a linear map, the rotation $R_j$), so they depend only on tokens $1, \dots, j$.
3. The causal mask lets query $j$ combine only $k_i, v_i$ with $i \le j$, and these depend only on tokens $1, \dots, i$, with $i \le j$.
4. The residual add, the norm and the FFN act on each token separately. So $h_j^{(\ell)}$ depends only on tokens $1, \dots, j$. $\blacksquare$

Each step of the proof is a rule for the code. Let $P$ be the number of cached tokens and $T$ the number of new tokens (0-based positions from here on).

• Step 2, positions. Rotate each key by its absolute position before you store it. The new tokens have positions $P, P + 1, \dots, P + T - 1$, not $0, \dots, T - 1$.
• Step 3, the mask. There are $T$ queries and $P + T$ keys. Query $i$ has position $P + i$ and may see key $j$ exactly when $j \le P + i$. In PyTorch: torch.ones(T, P + T, dtype=torch.bool).tril(diagonal=P) (True = attend). For $T = 1$ nothing is masked. SDPA's is_causal=True aligns the triangle to the upper-left corner, which is wrong whenever $P > 0$.
• Step 4, no other mixing. No operation outside attention may mix tokens. BatchNorm over the time axis, for example, would break the theorem.

Cost per layer and decode step, with $t$ tokens in the context. With the cache: projections and FFN for one token, $O(D^2)$, plus attention over $t$ keys, $O(tD)$. Without the cache all $t$ tokens are recomputed: $O(tD^2 + t^2 D)$. The work per step drops by a factor of about $t$. Decoding then becomes memory-bound: every step reads all the weights and the whole cache.

Memory per token, with $L$ layers and $b$ bytes per number:

where the 2 counts $K$ and $V$.

Two implementations. torch.cat on the time axis is simple, but each step copies the whole cache: $O(T^2)$ copying over $T$ steps. A preallocated buffer $(B, n_{kv}, \text{max\_len}, d_h)$ with a fill pointer has static shapes, as production code needs. In an interview, write cat first and mention the buffer.

3.6 Decoding: greedy, temperature, top-k and top-p #

At each step the model outputs a logit vector $z \in \mathbb{R}^{|\mathcal V|}$, one number per token of the vocabulary $\mathcal V$ (about $3 \cdot 10^4$ to $2.5 \cdot 10^5$ tokens). The softmax turns logits into probabilities $p_i = e^{z_i}/\sum_j e^{z_j}$. A decoding rule picks the next token, appends it, and repeats until an end-of-sequence token or a maximum length.

• Greedy decoding takes $\arg\max_i p_i$ at each step. It is deterministic.
• Temperature. Sample from $p_i(\tau) = e^{z_i/\tau}/\sum_j e^{z_j/\tau}$, with temperature $\tau > 0$. $\tau < 1$ makes the distribution sharper, $\tau > 1$ flatter; the order of the tokens never changes. $\tau \to 0$ gives greedy decoding; $\tau \to \infty$ gives the uniform distribution. Implement $\tau = 0$ as argmax, not as a division by zero.
• Top-k. Keep the $k$ largest logits, set the others to $-\infty$, take the softmax, sample. Weakness: $k$ is fixed. When the distribution is flat, a small $k$ cuts good tokens; when it is peaked, a large $k$ keeps junk.
• Top-p (nucleus) sampling (Holtzman et al. 2020). The nucleus $\mathcal V^{(p)}$ is the smallest set of tokens with total probability at least $p$. Sample from $p$ restricted to $\mathcal V^{(p)}$ and divided by the kept mass. Its size adapts to the shape of the distribution: one token for a peaked distribution, many for a flat one.

The top-p algorithm for one row of logits:

1. Compute $p = \operatorname{softmax}(z/\tau)$.
2. Sort in decreasing order, $p_{(1)} \ge p_{(2)} \ge \dots$, and remember the original indices.
3. Compute the cumulative sums $c_j = p_{(1)} + \dots + p_{(j)}$, with $c_0 = 0$.
4. Keep sorted token $j$ when the mass before it is below $p$: $c_{j-1} < p$. This keeps the token that crosses $p$, and it always keeps the top token.
5. Set the logits of the removed tokens to $-\infty$, scatter back to the original order, take the softmax, sample.

For a batch, do this on a $(B, |\mathcal V|)$ tensor with sort, cumsum and scatter along the last axis; never loop over the vocabulary in Python. Order: temperature first, then top-k, then top-p (the order of the Hugging Face logits processors). Temperature changes the nucleus, so the order matters.

Interactive figure: loading…

Variants. Min-p (Nguyen et al. 2024) keeps token $i$ when $p_i \ge p_{\min} \max_j p_j$. The Gumbel-max trick: with $g_i = -\log(-\log u_i)$ and independent $u_i \sim \text{Uniform}(0, 1)$, $\arg\max_i (z_i + g_i)$ is an exact sample from $\operatorname{softmax}(z)$.

3.7 Beam search #

Goal. Find a sequence $y = (y_1, \dots, y_n)$ with a high log-probability $\log P(y) = \sum_{t=1}^{n} \log p(y_t \mid y_{<t})$. Exact search over all sequences is impossible. Beam search keeps the $n_b$ best prefixes (the beam; $n_b$ is the beam width) at every step. Work with log-probabilities: a product of many numbers below 1 underflows, a sum of logs does not.

Algorithm (the interview version).

1. Start with one hypothesis: the prompt, with score 0.
2. For every live hypothesis $b$ with score $s_b$, compute the log-probabilities $\ell_b$ of the next token. Every pair $(b, v)$ is a candidate with score $s_b + \ell_b[v]$.
3. Flatten the $n_b \times |\mathcal V|$ scores and take the top $n_b$. For a flat index $i$, the parent beam is $i \,/\!/\, |\mathcal V|$ and the token is $i \bmod |\mathcal V|$.
4. A hypothesis that ends with the end-of-sequence token is finished: keep it with its score and stop expanding it.
5. Reorder the KV cache: row $j$ of the new cache is the row of the parent beam of hypothesis $j$ (a gather along the batch axis).
6. Stop when all hypotheses are finished or at the maximum length. Return the finished hypothesis with the best length-normalized score.

Two details that interviewers check. At step 1 all beams would be identical, so start with a single beam (or give beams $2, \dots, n_b$ the score $-\infty$). And production code (Hugging Face, for example) takes the top $2n_b$ candidates, moves finished ones to a separate list and refills the beam with live ones, so the beam never shrinks.

Length normalization. Every extra token adds a term $\log p < 0$, so raw scores favour short outputs. Normalize the final score:

where $|y|$ is the number of generated tokens and $\alpha \in [0, 1]$. With $\alpha = 1$ the score is the mean log-probability per token (the default length_penalty = 1.0 of Hugging Face). GNMT divides by $((5 + |y|)/6)^{\alpha}$ instead.

Interactive figure: loading…

Cost. Each step runs the model on $n_b$ rows (one batch), scores $n_b |\mathcal V|$ candidates and keeps $n_b$ KV caches. $n_b = 1$ is greedy decoding.

When it helps and when it fails. Use beam search when one correct output exists (translation, speech recognition, structured outputs) or when you want a ranked top-$n_b$ list. For open-ended text, beam search finds text that is more probable than human text, but repetitive and generic; once a phrase repeats, the model gives the next repetition an even higher probability (Holtzman et al. 2020). Very large beams can even make translation worse (Koehn and Knowles 2017).

3.8 Testing each piece fast #

• Reference: the same inputs through SDPA, nn.MultiheadAttention (copy the weights), F.layer_norm or F.rms_norm, in float64, with a tight tolerance such as atol=1e-10.
• No look-ahead: change the tokens after position $t$; the outputs up to $t$ must not move.
• RoPE: a common shift of all positions keeps every score; position 0 is the identity.
• Cache: prefill $P$ tokens, then decode one at a time; every output equals the full pass.
• Norms: use a small $D$ (the unbiased-variance bug changes the output by $\sqrt{(D-1)/D}$, which is invisible for $D = 4096$).
• Sampling: $\tau \to 0$ equals argmax; $k = |\mathcal V|$ and $p = 1$ change nothing; frequencies over $10^5$ draws match the renormalized distribution. Beam: width 1 equals greedy; with a large width the result equals brute force on a toy model.

Answers to section 1. (1) With is_causal=True, SDPA builds the mask ones(1, 5).tril() = [True, False, False, False, False]: the triangle starts at the upper-left corner, so the new token attends only to key 0, the first prompt token. It should see all 5 keys ($j \le 4 + 0$); a single decode query needs no mask at all. (2) No. In an encoder, the layer-1 output of token 3 attends to token 10, so the keys and values of token 3 in layers $2, \dots, L$ change when token 10 arrives; cached values would be stale. Only the first layer's keys and values (projections of the embeddings) stay valid. In a causal decoder they never change: this is the theorem of section 3.5.

4. Worked example #

(a) Attention by hand. One head, $d_h = 2$, $T = 3$, causal, and $Q = K = V = X$ with rows $x_1 = (1, 0)$, $x_2 = (0, 1)$, $x_3 = (1, 1)$. The scores are $S = XX^\top/\sqrt 2$.

1. Row 1 sees only key 1. Output $(1, 0)$.
2. Row 2: scores $(x_2 \cdot x_1, x_2 \cdot x_2)/\sqrt2 = (0, 0.707)$. With $e^{0.707} = 2.028$ the weights are $1/3.028 = 0.330$ and $2.028/3.028 = 0.670$. Output $0.330\,(1, 0) + 0.670\,(0, 1) = (0.330, 0.670)$.
3. Row 3: scores $(1, 1, 2)/\sqrt2 = (0.707, 0.707, 1.414)$. With $e^{1.414} = 4.113$ the sum is $2.028 + 2.028 + 4.113 = 8.169$, and the weights are $(0.248, 0.248, 0.503)$. Output $0.248\,(1, 0) + 0.248\,(0, 1) + 0.503\,(1, 1) = (0.752, 0.752)$.

Row 3 is the same with and without the mask, because the last query sees every key anyway. That is the decode step of a cache. Changing $x_3$ changes only row 3: no look-ahead.

(b) The offset mask. The cache holds $P = 3$ tokens, and $T = 2$ new tokens arrive (a chunked prefill, or two draft tokens to verify). The keys are $0, \dots, 4$, and the queries have positions 3 and 4, which are also their RoPE positions. ✓ = may attend:

Mask	query	key 0	key 1	key 2	key 3	key 4
correct, tril(diagonal=3)	position 3	✓	✓	✓	✓	·
correct, tril(diagonal=3)	position 4	✓	✓	✓	✓	✓
wrong, is_causal=True	position 3	✓	·	·	·	·
wrong, is_causal=True	position 4	✓	✓	·	·	·

(c) RoPE by hand. $d_h = 4$ (two pairs) and $\Theta = 100$, so $\omega_0 = 1$ and $\omega_1 = 100^{-1/2} = 0.1$. Take $q = k = (1, 0, 1, 0)$ in the interleaved layout. A pair $(1, 0)$ turned by $\phi$ is $(\cos\phi, \sin\phi)$.

1. Query at $m = 3$: $(\cos 3, \sin 3, \cos 0.3, \sin 0.3) = (-0.990, 0.141, 0.955, 0.296)$.
2. Key at $n = 1$: $(\cos 1, \sin 1, \cos 0.1, \sin 0.1) = (0.540, 0.841, 0.995, 0.100)$.
3. Score: for two unit pairs, the dot product is the cosine of the angle between them, so the score is $\cos(3 - 1) + \cos(0.3 - 0.1) = \cos 2 + \cos 0.2 = -0.416 + 0.980 = 0.564$.
4. Shift both positions by 10 ($m = 13$, $n = 11$): again $0.564$. Same position ($m = n$): $2.0 = q^\top k$, as without RoPE.

The fast pair ($\omega_0 = 1$) changes sign within a few positions and separates neighbours. The slow pair ($\omega_1 = 0.1$) changes slowly and encodes coarse, long-range position.

(d) LayerNorm and RMSNorm. $x = (1, 3, 5, 7)$, $\varepsilon = 0$, $\gamma = \mathbf 1$, $\beta = 0$.

• LayerNorm: $\mu = 4$, $\sigma^2 = (9 + 1 + 1 + 9)/4 = 5$, output $(x - 4)/\sqrt 5 = (-1.342, -0.447, 0.447, 1.342)$.
• RMSNorm: $\sqrt{(1 + 9 + 25 + 49)/4} = \sqrt{21} = 4.583$, output $(0.218, 0.655, 1.091, 1.528)$.
• With the unbiased variance $20/3 = 6.667$, the LayerNorm output would be multiplied by $\sqrt{5/6.667} = \sqrt{0.75} = 0.866$: a 13% error. With $D = 4096$ the same bug gives the factor $\sqrt{4095/4096} = 0.99988$, which a loose test misses.

(e) KV-cache memory. Llama-2-7B has $L = 32$ layers, $n_{kv} = n_h = 32$ heads and $d_h = 128$. In fp16: $2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 524{,}288$ bytes = 512 KiB per token. One 4096-token sequence needs 2 GiB, and eight of them need 16 GiB: more than the 13.5 GB of fp16 weights (6.74 billion parameters × 2 bytes). With GQA and 8 KV heads (Llama 3 8B, same $L$ and $d_h$) it is 128 KiB per token, 4 times less.

(f) Sampling. Five tokens $a, b, c, d, e$ with logits $z = (2, 1, 0.5, 0, -1)$.

	$a$	$b$	$c$	$d$	$e$
$\tau = 1$	0.563	0.207	0.126	0.076	0.028
$\tau = 0.5$	0.829	0.112	0.041	0.015	0.002
$\tau = 2$	0.375	0.227	0.177	0.138	0.084

• Top-k with $k = 2$, $\tau = 1$: keep $a, b$ (mass 0.770). Renormalized: $(0.731, 0.269)$.
• Top-p with $p = 0.8$, $\tau = 1$: the cumulative sums are $0.563, 0.770, 0.896, \dots$ The mass before $c$ is $0.770 < 0.8$, so $c$ stays. The mass before $d$ is $0.896 \ge 0.8$, so $d$ goes. Nucleus $\{a, b, c\}$, renormalized $(0.629, 0.231, 0.140)$.
• Temperature, then top-p with $p = 0.8$: at $\tau = 0.5$ the nucleus is $\{a\}$ alone ($0.829 \ge 0.8$), so sampling becomes greedy. At $\tau = 2$ it is $\{a, b, c, d\}$. The same $p$ keeps 1, 3 or 4 tokens.

(g) Beam search by hand. Vocabulary $\{A, B, E\}$, where $E$ is the end token. The toy model (every two-token prefix other than $BA$ has the same row as $BA$; at most 3 tokens):

prefix	$p(A)$	$p(B)$	$p(E)$
(start)	0.5	0.4	0.1
$A$	0.25	0.15	0.6
$B$	0.9	0.05	0.05
$BA$	0.1	0.1	0.8

1. Greedy: $A$ (0.5), then $E$ (0.6). Output "$A\,E$" with probability $0.30$.
2. Beam with $n_b = 2$, step 1: keep $A$ ($\log 0.5 = -0.693$) and $B$ ($\log 0.4 = -0.916$).
3. Step 2: six candidates: $AA$ 0.125, $AB$ 0.075, $AE$ 0.30, $BA$ 0.36, $BB$ 0.02, $BE$ 0.02. Keep $BA$ (0.36, live) and $AE$ (0.30, finished). Greedy dropped $B$ at step 1, but $BA$ is the best prefix of length 2.
4. Step 3: expand $BA$: $BAE = 0.36 \cdot 0.8 = 0.288$; $BAA$ and $BAB$ have 0.036 each. Keep $AE$ (0.30) and $BAE$ (0.288). Both are finished, so stop.
5. Pick the winner. Raw log-probabilities: $\log 0.30 = -1.204$ beats $\log 0.288 = -1.245$, so the answer is "$A\,E$", the same as greedy. With $\alpha = 1$: $-1.204/2 = -0.602$ against $-1.245/3 = -0.415$, so the answer is "$B\,A\,E$". Length normalization changed the result.

5. In code #

PyTorch's own functions check the claims above. The notebook rebuilds every piece in NumPy.

import numpy as np
import torch
import torch.nn.functional as F

f64 = torch.float64
sdpa = F.scaled_dot_product_attention

# (a) one head, T = 3, d_h = 2, Q = K = V = X, causal
X = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype=f64)[None, None]  # (1, 1, 3, 2)
out = sdpa(X, X, X, is_causal=True)
print(out[0, 0].numpy().round(4))           # rows (1, 0), (0.3302, 0.6698), (0.7517, 0.7517)

# no look-ahead: change the last token; rows 0 and 1 do not move
X2 = X.clone()
X2[0, 0, 2] = torch.tensor([5.0, -3.0], dtype=f64)
print(torch.equal(sdpa(X2, X2, X2, is_causal=True)[0, 0, :2], out[0, 0, :2]))   # True

# (b) the cache theorem for one layer: cached computations = rows of the full causal pass
torch.manual_seed(0)
B, H, T, Dh = 2, 4, 6, 8
q, k, v = (torch.randn(B, H, T, Dh, dtype=f64) for _ in range(3))
full = sdpa(q, k, v, is_causal=True)                                    # (B, H, T, Dh)
offset = torch.ones(2, 6, dtype=torch.bool).tril(diagonal=4)            # P = 4, T = 2
checks = {
    "decode 1 query, 5 keys": (sdpa(q[:, :, 4:5], k[:, :, :5], v[:, :, :5]), full[:, :, 4:5]),
    "chunk 2, offset mask":   (sdpa(q[:, :, 4:], k, v, attn_mask=offset), full[:, :, 4:]),
    "chunk 2, is_causal":     (sdpa(q[:, :, 4:], k, v, is_causal=True), full[:, :, 4:]),  # wrong
}
for name, (got, ref) in checks.items():
    print(f"{name:24s} {(got - ref).abs().max().item():.1e}")
# decode 1 query, 5 keys   6.7e-16
# chunk 2, offset mask     0.0e+00
# chunk 2, is_causal       2.7e+00

# GQA: 8 query heads share 2 KV heads (groups of 4)
q8 = torch.randn(B, 8, T, Dh, dtype=f64)
k2, v2 = torch.randn(B, 2, T, Dh, dtype=f64), torch.randn(B, 2, T, Dh, dtype=f64)
a = sdpa(q8, k2, v2, is_causal=True, enable_gqa=True)
b = sdpa(q8, k2.repeat_interleave(4, dim=1), v2.repeat_interleave(4, dim=1), is_causal=True)
print(f"{(a - b).abs().max().item():.1e}")    # 0.0e+00

# MHA parameters do not depend on the number of heads: 4 D^2 + 4 D
print([sum(p.numel() for p in torch.nn.MultiheadAttention(512, h).parameters()) for h in (1, 8, 16)])
# [1050624, 1050624, 1050624]

# (d) LayerNorm vs RMSNorm, and the variance trap
x = torch.tensor([1.0, 3.0, 5.0, 7.0], dtype=f64)
ln, rms = F.layer_norm(x, (4,), eps=0.0), F.rms_norm(x, (4,), eps=0.0)
print(ln.numpy().round(4), rms.numpy().round(4))
# [-1.3416 -0.4472  0.4472  1.3416] [0.2182 0.6547 1.0911 1.5275]
print(torch.allclose(ln, F.rms_norm(x - x.mean(), (4,), eps=0.0)))   # True: LN = RMSNorm of x - mean
print(round(x.var().item(), 4), x.var(correction=0).item())          # 6.6667 5.0

# (c) RoPE scores, d_h = 4, base 100, q = k = (1, 0, 1, 0)
omega = 100.0 ** (-np.arange(0, 4, 2) / 4)                      # [1.0, 0.1]
def rotated_unit(m):                                            # each pair (1, 0) turned by m * omega_i
    return np.stack([np.cos(m * omega), np.sin(m * omega)], axis=1).ravel()
for m, n in [(3, 1), (13, 11), (5, 5)]:
    print(m, n, round(float(rotated_unit(m) @ rotated_unit(n)), 5))   # 0.56392, 0.56392, 2.0

# (f) top-p with p = 0.8 on z = (2, 1, 0.5, 0, -1), and the Gumbel-max trick
z = np.array([2.0, 1.0, 0.5, 0.0, -1.0])
p = np.exp(z - z.max()); p /= p.sum()
keep = (np.cumsum(p) - p) < 0.8                 # mass before each (already sorted) token
print(keep, (p * keep / (p * keep).sum()).round(3))
# [ True  True  True False False] [0.629 0.231 0.14  0.    0.   ]
rng = np.random.default_rng(0)
g = -np.log(-np.log(rng.uniform(size=(200_000, 5))))
freq = np.bincount(np.argmax(z + g, axis=1), minlength=5) / 200_000
print(np.abs(freq - p).max() < 5e-3)           # True

The three cache checks agree to float64 rounding, except is_causal=True on the chunk, which is off by about 2.7 because it uses the wrong mask. The notebook checks the beam search of 4(g) against a brute force over all sequences.

6. Where it shows up #

• LLM architecture. The LLaMA-style block is pre-norm RMSNorm, RoPE, GQA and a SwiGLU FFN. CS336 assignment 1 builds it piece by piece; nanoGPT is the same block with LayerNorm and learned absolute positions.
• LLM serving. The KV cache is the reason decoding is memory-bound, and the reason servers page it (vLLM's PagedAttention), quantize it, and share it between requests with a common prefix (lesson 13, ML Systems: inference serving). Sliding-window attention bounds the cache (Mistral 7B uses a window of 4096 tokens). FlashAttention computes the same attention without storing the $T \times T$ matrix (ML Systems: FlashAttention).
• Chat models and RL. Deployments usually combine a temperature below 1 with top-p near 0.9. Regression tests use greedy decoding, which is deterministic. RL fine-tuning samples rollouts, and the temperature sets how much the policy explores.
• Speech and translation. Whisper decodes with beam search (5 beams) and falls back to sampling at a higher temperature when the output looks repetitive. Machine translation has used beam search with length normalization since GNMT.
• ML for science. Weather, climate and other scientific transformers use patches of a grid as tokens. 2-D (axial) RoPE turns half of the pairs with the row index and half with the column index. Missing observations or land–sea masks become key-padding masks. Sequence models for molecules and proteins use temperature and top-p to set the diversity of generated candidates, and beam search for ranked lists (reaction prediction, retrosynthesis).

7. Common mistakes #

• Softmax over the queries (dim=-2) instead of the keys (dim=-1).
• Scaling by $\sqrt{D}$ instead of $\sqrt{d_h}$.
• view(B, H, T, d_h) instead of view(B, T, H, d_h).transpose(1, 2): right shape, mixed tokens.
• view after transpose: an error, because the tensor is not contiguous. Use reshape.
• Mask polarity: masked_fill(mask, -inf) with a mask where True means "attend", or the SDPA convention mixed with the nn.MultiheadAttention one.
• is_causal=True with a KV cache, or a square-matrix rule such as triu(diagonal=1) applied to a $T \times (P + T)$ score matrix.
• RoPE positions that restart at 0 in each decode step; keys cached before RoPE and then rotated again; RoPE on $v$; angles computed in bf16; mixing the interleaved and half-split layouts.
• LayerNorm with x.var() (unbiased), $\varepsilon$ outside the root, or the wrong axis; RMSNorm with a mean subtraction (that is LayerNorm without $\beta$).
• A query row with every key masked (left padding): NaN in your own softmax.
• Top-p off by one: keeping only the tokens with $c_j \le p$ drops the token that crosses $p$, and removes every token when $p_{(1)} > p$. The rule is $c_{j-1} < p$.
• Masking without renormalizing, temperature applied after top-p, and $\tau = 0$ as a division.
• Beam bookkeeping: identical beams at step 1; i // V and i % V swapped; the KV cache not reordered by parent beam; products of probabilities instead of sums of logs; no length normalization.
• A test that checks only shapes, or that compares with your own earlier code instead of a library reference.

8. Check yourself #

1. A model has $L = 24$ layers, $D = 2048$, $n_h = 16$ query heads, $n_{kv} = 4$ KV heads, and a bf16 cache. How many bytes of KV cache per token? For one sequence of 32,768 tokens?
   Answer

   $d_h = 2048/16 = 128$. Per token: $2 \cdot 24 \cdot 4 \cdot 128 \cdot 2 = 49{,}152$ bytes = 48 KiB. For 32,768 tokens: $49{,}152 \cdot 32{,}768$ bytes = 1.5 GiB.
2. Your LayerNorm uses x.var(-1, keepdim=True). With $D = 4$, how wrong is the output? Why can a test with $D = 4096$ miss the bug?
   Answer

   The unbiased variance is $\frac{D}{D-1}\sigma^2$, so the output is multiplied by $\sqrt{(D-1)/D}$: $\sqrt{0.75} = 0.866$ for $D = 4$. For $D = 4096$ the factor is $0.99988$, inside a loose tolerance.
3. Why does RoPE rotate $q$ and $k$ but not $v$, and why does every head use the same rotation?
   Answer

   The theorem needs both $q$ and $k$ rotated; then the score depends only on $n - m$. On $v$, the output $\sum_j a_{ij} R_j v_j$ would carry the absolute positions of the keys, and these rotations do not cancel. Position belongs to the token, not to the head, so all heads use the same tables.
4. Your test "cached decoding = full forward pass" gives a maximum error of $10^{-15}$ in float64 but $2 \cdot 10^{-2}$ in bf16. Is there a bug?
   Answer

   Probably not. bf16 has 8 significant bits (machine epsilon $2^{-7} \approx 0.008$), and the two paths add numbers in different orders, so differences of order $10^{-2}$ are expected. Check the logic in float64 with a tight tolerance; run bf16 only as a smoke test with a loose tolerance.
5. Four tokens have equal logits. With top-p and $p = 0.6$, how many tokens are kept?
   Answer

   Three. Each has probability 0.25. The mass before tokens 1–4 is 0, 0.25, 0.5, 0.75, and only the first three values are below 0.6. Which three are kept depends on how the sort breaks ties.
6. Does temperature change greedy decoding? Does it change top-p sampling?
   Answer

   Greedy: no, because $z/\tau$ has the same argmax for every $\tau > 0$. Top-p: yes, a lower $\tau$ sharpens $p$, so fewer tokens are needed to reach the mass $p$ (worked example (f): 1, 3 or 4 tokens).

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write scaled dot-product attention with a mask.

   Show answer

2. Why divide the scores by $\sqrt{d_h}$?

   Show answer

3. Write RMSNorm for one token $x \in \mathbb{R}^D$.

   Show answer

4. State the relative-position property of RoPE.

   Show answer

5. How many bytes of KV cache does one token need?

   Show answer

6. Why is a KV cache exact in a causal decoder but not in a bidirectional encoder?

   Show answer

7. With $P$ cached tokens and $T$ new ones, which mask do the new queries need?

   Show answer

8. State the top-p (nucleus) rule.

   Show answer

9. What does temperature do as $\tau \to 0$ and as $\tau \to \infty$?

   Show answer

10. Why does beam search need length normalization?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

Recognize the task from the wording.

The prompt says	What to write first
"implement multi-head attention"	the shape legend, fused QKV, head split by reshape, a test against SDPA in float64
"make it causal", "GPT-style"	a boolean mask (True = attend), applied before the softmax; a no-look-ahead test
"generate text fast", "add a cache"	forward(x, cache) returning (out, cache); offset positions and the tril(diagonal=P) mask; a cached = full test
"LLaMA-style block"	RMSNorm (pre-norm), RoPE on $q$ and $k$, GQA with repeat_interleave, SwiGLU
"sample from the model"	temperature, then top-k, then top-p on a $(B, \lvert\mathcal V\rvert)$ tensor; $\tau = 0$ as argmax
"best sequence", "translation"	beam search with log-probabilities, finished list, length normalization

What to say aloud before coding attention (about 30 seconds): "I use batch-first tensors: $x$ is $(B, T, D)$, $H$ heads, $d_h = D/H$. One fused projection to $3D$, split into $q, k, v$, then view(B, T, H, d_h).transpose(1, 2) so each head stays inside one token. Scores are $(B, H, T, T)$, scaled by $\sqrt{d_h}$, masked with $-\infty$ where the boolean mask is False, softmax over the last axis. Then merge with transpose and reshape, and the output projection. I will test it against SDPA in float64 and with a no-look-ahead test."

A plan for 45 minutes.

1. 0–3 min: clarify (batch-first? biases? mask convention? may you use F.softmax and nn.Linear? is a cache needed?); write the signature and a legend of the shape letters.
2. 3–8 min: the skeleton below, with a shape comment per step.
3. 8–20 min: attention forward; the test against SDPA and the no-look-ahead test.
4. 20–27 min: RMSNorm; test against F.rms_norm with the same $\varepsilon$.
5. 27–37 min: RoPE and the cache; test that cached decoding equals the full pass at every position.
6. 37–45 min: complexity, memory, follow-ups.

import torch
import torch.nn as nn


class CausalSelfAttention(nn.Module):
    """Causal self-attention with RoPE, GQA and a KV cache.
    Legend: B batch, T new tokens, P cached tokens, D = d_model, H query heads,
    H_kv KV heads (H % H_kv == 0), Dh = D // H. Masks: True = may attend."""

    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int | None = None,
                 rope_base: float = 10_000.0, max_len: int = 4096):
        super().__init__()
        ...  # W_qkv: D -> (H + 2 H_kv) Dh, no bias; W_o: D -> D; cos/sin buffers (max_len, Dh/2), fp32

    def forward(self, x: torch.Tensor, cache: tuple | None = None) -> tuple:
        # x: (B, T, D); cache: (k, v), each (B, H_kv, P, Dh), keys already rotated
        # 1. qkv = W_qkv(x)                               (B, T, (H + 2 H_kv) Dh)
        # 2. split, view(B, T, heads, Dh).transpose(1, 2)  q (B, H, T, Dh); k, v (B, H_kv, T, Dh)
        # 3. RoPE on q and k at positions P + arange(T)
        # 4. k, v = cat(cache, new) on dim 2              (B, H_kv, P + T, Dh)
        # 5. repeat k, v to H heads; s = q k^T / sqrt(Dh)  (B, H, T, P + T)
        # 6. key j visible to query i iff j <= P + i; probs = softmax of s over the last axis
        # 7. (probs @ v).transpose(1, 2).reshape(B, T, D), then W_o; return out, (k, v)
        ...

Typical questions and short model answers.

1. "Why divide by $\sqrt{d_h}$?" — "If the entries of $q$ and $k$ have mean 0 and variance 1, the dot product has variance $d_h$. For $d_h = 128$ that is a standard deviation of about 11, so the softmax is nearly one-hot and its gradient nearly zero. Dividing by $\sqrt{d_h}$ keeps the scores at unit scale at initialization."
2. "What does attention cost?" — "For $T$ tokens and width $D$: the projections are $O(TD^2)$, the scores and the weighted sum $O(T^2 D)$ over all heads, and the score matrix needs $O(T^2)$ memory per head. FlashAttention computes the same result in tiles without storing that matrix, so the memory becomes $O(T)$, but the FLOPs stay $O(T^2 D)$."
3. "Why a KV cache, and why is it exact?" — "In a causal decoder, token $j$'s hidden states depend only on tokens up to $j$, so its keys and values never change. Caching them turns each decode step from $O(tD^2 + t^2 D)$ into $O(D^2 + tD)$, for $2 L n_{kv} d_h$ stored numbers per token. Two details: new tokens get positions $P, P+1, \dots$ for RoPE, and the mask is aligned to the bottom-right corner."
4. "MHA, GQA or MQA?" — "Fewer KV heads mean a smaller cache and less memory traffic per decode step, so faster decoding and larger batches, for a small loss in quality. GQA with 8 KV heads is the common middle ground; MQA with one KV head is the extreme."
5. "Why RoPE?" — "It puts relative position into the score with no parameters, because rotations satisfy $R_m^\top R_n = R_{n-m}$. Each key is rotated once by its absolute position, so it fits the cache. For longer contexts you can interpolate positions or raise the base."
6. "Top-k or top-p?" — "Top-p adapts to the shape of the distribution: one token when the model is sure, many when it is not. Top-k keeps a fixed number, which is too many in the first case and too few in the second. I apply temperature first, then the filters, and renormalize."

Typical follow-ups. "Prompts of different lengths in one batch?" (pad on the left, take positions from the attention mask with cumsum - 1, mask the pad keys, and make sure no query row is fully masked). "Why pre-norm?" (the residual path is never normalized, so deep stacks train stably without a long warmup). "Beam search with a KV cache?" (gather the cache rows by parent beam after each step). "How does speculative decoding stay exact?" (a draft token $x$ from the draft distribution $q$ is accepted with probability $\min(1, p(x)/q(x))$; after a rejection, sample from $\max(0, p - q)$ renormalized; the output distribution is exactly the target $p$).

10. Practice #

The notebook for this lesson runs in Colab, in NumPy. Each exercise is tested against a reference:

1. mha_forward: causal multi-head attention, checked against a loop over heads in float64, with a no-look-ahead test and a test that catches the $\sqrt{D}$ scale.
2. rms_norm: checked against the formula and against LayerNorm of the centered input.
3. apply_rope: interleaved RoPE with a positions argument; the relative-position property is checked numerically.
4. decode_step: one decode step of a tiny 2-layer decoder with a KV cache; every step must equal the full forward pass.
5. top_p_filter and beam_search: the nucleus on a batch of rows, and beam search checked against the hand-solved example and a brute force.

Timed practice (25–30 minutes each, talking aloud): "implement causal multi-head attention" (also on Deep-ML: Implement Multi-Head Attention, Implement Masked Self-Attention); "implement LayerNorm for sequence data"; "add RoPE with a positions argument"; "add a KV cache with GQA and generate"; "implement top-p sampling for a batch"; "implement beam search given a step(tokens, cache) -> logits function". Then run the CS336 assignment 1 tests for RMSNorm, RoPE, scaled dot-product attention and multi-head self-attention with RoPE on your code.

Say it aloud drill. Without notes, in 2 minutes: the shape of every tensor in the attention forward pass. Then in 1 minute: why cached decoding is exact, and the two rules (positions and mask) that follow from the proof.

11. Watch and read #

Let's build GPT: from scratch, in code, spelled out.Andrej Karpathy · 1 h 56 min

Self-attention, the causal mask, multi-head attention and the decoder block, written from an empty file.

Attention in transformers, step-by-step | Deep Learning Chapter 63Blue1Brown · 26 min

Queries, keys and values as pictures; the masking and the multi-head idea of section 2.

LLaMA explained: KV-Cache, Rotary Positional Embedding, RMS Norm, Grouped Query Attention, SwiGLUUmar Jamil · 1 h 11 min

The LLaMA block piece by piece: RMSNorm, RoPE, the KV cache and grouped-query attention.

• Stanford CS336 (Spring 2025), Lecture 3: Architectures, hyperparameters (video on YouTube; embedding is disabled): pre-norm, RMSNorm, RoPE, GQA and the choices of current LLMs.
• Stanford CS336, assignment 1 handout and tests: RMSNorm, RoPE, softmax and scaled dot-product attention (the boolean mask convention), causal multi-head self-attention.
• PyTorch, scaled_dot_product_attention: the reference implementation, the mask convention, is_causal, enable_gqa.
• Sebastian Raschka, Understanding and Coding the KV Cache in LLMs from Scratch.
• Vaswani et al. (2017), Attention Is All You Need, §3.2 (scaled dot-product and multi-head attention).
• Su et al. (2021), RoFormer: Enhanced Transformer with Rotary Position Embedding, §3.4.
• Zhang and Sennrich (2019), Root Mean Square Layer Normalization.
• Ainslie et al. (2023), GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.
• Holtzman et al. (2020), The Curious Case of Neural Text Degeneration, §3 (nucleus sampling, top-k, temperature).