Warmup Guide — The Transformer From Scratch

Zero-to-senior primer for Phase 02. We start from "a token is an integer" and build the entire decoder-only transformer one mechanism at a time — embeddings, self-attention, the causal mask, multi-head, the KV-cache, RoPE, RMSNorm, the SwiGLU MLP, the residual stream, the full forward pass — and finish with FlashAttention as the IO-aware version of the attention you just wrote. Each chapter is the same shape: what it is → why it exists → how it works under the hood (mechanism, a diagram, a little math) → why it matters in production → the misconception that trips people up. By the end you can build a GPT forward pass from nothing and defend every line.

Table of Contents

Chapter 1: From Token IDs to Vectors — The Embedding Table

From zero. Phase 01 turned text into a list of integers — token IDs like [1, 5, 9, 2]. A neural network cannot multiply an integer "ID 5" by a weight matrix in any meaningful way; the ID is a name, not a quantity. The first thing every transformer does is replace each ID with a learned vector. That is the embedding table: a [vocab, d_model] matrix where row i is the vector for token i. "Embedding lookup" is literally X = [embed[tid] for tid in token_ids] — indexing rows.

Why it exists. We want tokens with related meaning to sit near each other in a continuous space so that the same downstream matmuls can act on "king" and "queen" coherently. The table is learned: gradient descent moves each row until the geometry is useful. d_model (the model width — 4096 for Llama-2-7B, 768 for GPT-2-small) is the dimensionality of that space and the single number that flows through the entire network unchanged.

Under the hood. The lab's gpt_forward starts exactly here:

token_ids = [1, 5, 9, 2]
X = [ embed[1],      # a d_model-length vector
      embed[5],
      embed[9],
      embed[2] ]     # X is now [T, d_model] = [4, d_model]

There is no matmul yet — just a gather. From here on, every position is a d_model vector and the sequence is a [T, d_model] matrix. (Many models tie the embedding table to the final LM head — the same [vocab, d_model] weights are reused, transposed, to turn the last hidden state back into logits, saving vocab·d_model parameters. The lab keeps them separate for clarity.)

Production significance. The embedding + LM-head rows are 2·vocab·d_model parameters — for a 128k-token vocab and d_model=4096 that is ~1B params before any transformer layer, which is why large vocabularies are expensive and why weight tying matters.

Common misconception. "The embedding encodes position." It does not — embed[5] is the same vector wherever token 5 appears. Position is injected separately (Chapter 7). The embedding encodes identity, not order.

Chapter 2: Self-Attention — Q, K, V and the Dot Product as Similarity

What it is. Self-attention lets every position look at every other position and pull in information from the relevant ones. It is the only operation in a transformer that mixes across positions — everything else (norm, MLP, residual) acts on each position independently. Without attention, the model would process each token in isolation.

Why Q, K, V. Borrow the analogy of a soft dictionary lookup. Each position emits three learned projections of its embedding:

  • a query q — "what am I looking for?"
  • a key k — "what do I offer to others searching?"
  • a value v — "what do I actually contribute if attended to?"

They are just three linear maps of the same input X:

$$ Q = X W_Q, \qquad K = X W_K, \qquad V = X W_V. $$

In the lab these are three linear(X, Wq/Wk/Wv) calls. Splitting the roles is what lets a position ask for something different from what it offers.

The dot product is similarity. To decide how much query i should attend to key j, take their dot product q_i · k_j. A dot product is large and positive when two vectors point the same way — so it is a learned, content-based similarity score. Computing all of them at once is a matmul:

$$ S = Q K^\top \quad (\text{shape } [T, T]),\qquad S_{ij} = q_i \cdot k_j. $$

Row i of S is "how much does query i match every key." We then turn each row into weights and take a weighted average of the value rows:

$$ \text{Attention}(Q,K,V) = \text{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}} + \text{mask}\right) V. $$

Under the hood — a worked picture.

                 keys (what each position offers)
                 k0    k1    k2
              ┌──────────────────┐
   query q1 ──┤ q1·k0 q1·k1  -∞  │  scores row for position 1 (k2 is the future → masked)
              └──────────────────┘
                     │ softmax over the row
                     ▼
              [ 0.6,  0.4,  0.0 ]   attention weights (sum to 1)
                     │
                     ▼  weighted sum of value rows
   output1 = 0.6·v0 + 0.4·v1 + 0.0·v2

That is exactly scaled_dot_product_attention in the lab: scores, scale, mask, softmax per row, then a weighted sum of V.

Production significance. This single op is where a model resolves coreference ("it" → which noun?), copies from context, does in-context learning, and attends to retrieved documents in RAG. When generation hallucinates or ignores the prompt, you are debugging these weights.

Common misconception. "Q, K, V are different inputs." In self-attention they are three projections of the same sequence X. (Cross-attention — encoder→decoder, or attending to image patches — is the variant where K/V come from a different source than Q. The lab is self-attention.)

Chapter 3: Why /√d_k and Why softmax

The softmax. We need to turn a row of raw scores into a set of non-negative weights that sum to 1 so the output is a convex combination of value rows (a weighted average, never an extrapolation). softmax does exactly that:

$$ \text{softmax}(s)_j = \frac{e^{s_j}}{\sum_k e^{s_k}}. $$

It is also differentiable and "soft" — small score changes nudge the weights smoothly, which is what gradient descent needs. The lab's softmax subtracts the max before exponentiating:

$$ \text{softmax}(s)_j = \frac{e^{s_j - \max(s)}}{\sum_k e^{s_k - \max(s)}}. $$

This changes nothing mathematically (softmax is shift-invariant) but everything numerically: e^{1002} overflows to inf and the result is NaN, while e^{1002-1002}=1 is fine. As a bonus, a masked -inf score maps to e^{-\infty}=0 — exactly zero weight, which is what the causal mask needs. That numerical-stability detail is part of the lesson, and test_softmax_numerically_stable_large_logits checks it.

Why divide by √d_k. The scores are dot products over d_k dimensions. If the query and key entries are roughly independent with unit variance, the dot product q·k = Σ_{i=1}^{d_k} q_i k_i has variance ≈ d_k and standard deviation ≈ √d_k. So as the head dimension grows, the raw scores grow too — and large scores push softmax toward a near one-hot distribution (it saturates). A saturated softmax has almost-zero gradient, so training stalls. Dividing by √d_k rescales the scores back to roughly unit variance regardless of d_k:

$$ \text{Var}!\left[\frac{q\cdot k}{\sqrt{d_k}}\right] \approx \frac{d_k}{(\sqrt{d_k})^2} = 1. $$

   raw scores (large d_k):  [ 8.0,  0.1, -7.9 ]  → softmax ≈ [1.00, 0.00, 0.00]  (saturated, dead gradient)
   /√d_k scaled:            [ 2.0,  0.0, -2.0 ]  → softmax ≈ [0.84, 0.11, 0.02]  (soft, trainable)

Production significance. This is the canonical "why √d_k" interview question. The honest answer is "to keep the softmax in its sensitive, high-gradient regime as the head dimension scales." test_sqrt_dk_scaling_is_present in the lab detects the factor by comparing to a hand-computed scaled softmax.

Common misconception. "√d_k is just an arbitrary normalization." It is a variance argument: without it, wider heads systematically saturate the softmax and the model trains worse. The constant comes from the standard deviation of a sum of d_k independent products.

Chapter 4: The Causal Mask & Autoregression

What it is. A decoder-only language model predicts the next token, so when computing the output at position i it must only see positions 0..i — never the future. The causal mask enforces this. It is a [T, T] matrix added to the scores before softmax: 0 on and below the diagonal, -∞ strictly above it.

$$ \text{mask}_{ij} = \begin{cases} 0 & j \le i \ -\infty & j > i \end{cases} $$

Why additive -∞. After adding the mask, every "future" score is -∞, and softmax(-\infty)=0. So position i puts exactly zero weight on every j > i — the future contributes nothing, and it contributes nothing to the gradient either. The lab builds this in causal_mask and softmax handles the -∞ cleanly because it subtracts the max.

Under the hood.

   causal_mask(4):                attention weights after softmax (row = a query):
   j:  0    1    2    3
   ┌─────────────────────┐        token 0 sees only {0}
 0 │  0   -∞   -∞   -∞   │        token 1 sees {0,1}
 1 │  0    0   -∞   -∞   │        token 2 sees {0,1,2}
 2 │  0    0    0   -∞   │        token 3 sees {0,1,2,3}
 3 │  0    0    0    0   │        the strict upper triangle is always 0 weight
   └─────────────────────┘

test_causal_attention_zero_in_upper_triangle proves the consequence: token 0 can only attend to V[0], so its output is exactly V[0].

Why it's called autoregression. "Autoregressive" means the model generates one token at a time, each conditioned only on what came before. The causal mask is what makes a parallel training pass (all positions at once) compute the same thing as the sequential generation it will later do — so you can train on a whole sequence in one shot and the model still learns next-token prediction honestly.

Production significance. Training on T positions in parallel — instead of T separate forward passes — is the entire reason transformers train efficiently. A bug here is catastrophic and silent: if the mask leaks even one future position, the model "cheats" during training (sees the answer), gets a deceptively low loss, and then generates garbage at inference because the future it relied on isn't there. This is a real, classic bug (Chapter: HITCHHIKERS war stories).

Common misconception. "The mask is set to a large negative number like -1e9." Many implementations do use a large finite negative (because -inf can produce NaN if a whole row is masked). The concept is -∞ → zero weight; the lab uses true -inf and guards the all-masked case by construction (the diagonal is always 0, so no row is ever fully masked).

Chapter 5: Multi-Head Attention — Many Subspaces at Once

What it is. Instead of one attention over the full d_model width, split it into h heads, each of width d_head = d_model / h, run attention independently in each, then concatenate the results and apply an output projection W_O. One attention can only express one weighted-average pattern per position; h heads express h of them at once.

Why multiple heads. Different relationships want different similarity functions. One head might track syntactic dependency (verb → its subject), another coreference (pronoun → antecedent), another positional/local patterns. With a single head, all of that has to be averaged into one set of weights and they interfere. Splitting into subspaces lets the model attend to several relationships in parallel and combine them.

Under the hood.

   X [T, d_model]
     │  project to Q, K, V  (each [T, d_model])
     ▼
   split each into h heads ──► Q_h, K_h, V_h   each [T, d_head]
     │
     ├─ head 0: SDPA(Q0, K0, V0, mask) ─► out0 [T, d_head]
     ├─ head 1: SDPA(Q1, K1, V1, mask) ─► out1 [T, d_head]
     └─ ...                                ...
     │  concat heads back to [T, d_model]
     ▼
   linear(concat, W_O)  ──►  [T, d_model]

That is multi_head_attention exactly: linear for the projections, _split_heads to slice each [T, d_model] into h matrices of [T, d_head], scaled_dot_product_attention per head, _concat_heads to glue them back, then the output projection W_O. Note d_model must be divisible by h — test_mha_rejects_indivisible_heads checks the guard.

Production significance. h is a real config knob (num_attention_heads). Mechanistic- interpretability work names specific heads ("induction heads" that do in-context copying). And the per-head split is precisely what GQA exploits next: share K/V across heads to shrink the cache.

Common misconception. "More heads = more parameters / more compute." Splitting d_model into heads is a reshape — the total Q/K/V projection size is unchanged. Heads change how the same budget is used (many narrow subspaces vs one wide one), not how big it is.

Chapter 6: MQA / GQA and the KV-Cache

The KV-cache (recap + mechanism). Phase 00 established the memory math; here is the mechanism. During generation ("decode"), the model emits one token at a time. To produce token t+1 it needs attention over all previous keys and values. Recomputing K/V for every past token at every step is O(T²) total work. Instead we cache the K and V vectors of every past token, so each new step computes K/V for just the one new token and reads the rest from the cache — O(T) total. The lab's attention_with_kv_cache is one such step: append the new K/V, attend the new query over the whole cache, return the context plus the grown cache.

The soul test. Why is the cache correct? Because under a causal mask, the last query attends to exactly tokens 0..t — which is exactly what the cache holds. So the incremental result is numerically identical (within float tolerance) to running full causal attention and taking the last row. test_kv_cache_equals_full_recompute proves this bit-for-(float)-bit, and it is the entire justification for the cache:

   full recompute:   SDPA(Q, K, V, causal_mask(T))  →  take row T-1
   kv-cache decode:  step t=0..T-1, attend new_q over cache  →  final context
   assertion:        context  ==  full[T-1]   (within approx)

MQA and GQA — shrinking the cache. The cache size is 2 · n_layers · T · d_model · bytes · batch (Phase 00). The lever is d_model in the cache — i.e. the K/V width. Two ideas reduce it:

  • MQA (Multi-Query Attention): all query heads share a single K and V head. The cache stores one head's worth of K/V instead of h — an h× reduction — at a small quality cost.
  • GQA (Grouped-Query Attention): the middle ground every modern model uses. Use g K/V heads (1 < g < h), each shared by a group of h/g query heads. With 8 K/V heads out of 32 query heads, the cache is 4× smaller with negligible quality loss.

$$ d_{kv} = d_{model}\cdot\frac{n_{kv}}{n_{head}} \quad\Rightarrow\quad \text{KV-cache} \propto n_{kv}. $$

   MHA   (h=4):  q0 q1 q2 q3        GQA (g=2):  q0 q1 | q2 q3      MQA (g=1):  q0 q1 q2 q3
                 k0 k1 k2 k3                     k0    |  k1                    k0
                 v0 v1 v2 v3                     v0    |  v1                    v0
                 cache = 4 K/V                   cache = 2 K/V                  cache = 1 K/V

Production significance. GQA is why models like Llama-2-70B and Mistral can serve long contexts at high concurrency without the KV-cache OOM from Phase 00. num_key_value_heads in a HF config is this exact knob. The lab implements plain multi-head; adding GQA is an extension (and ties straight back to Phase 00's cache math).

Common misconception. "GQA makes attention faster." Its main win is memory (a smaller KV-cache), which indirectly raises throughput because decode is memory-bandwidth-bound (Phase 00) — fewer bytes to read per token. The FLOPs barely change.

Chapter 7: Position — RoPE vs Learned Absolute

The problem. Self-attention is permutation-equivariant: shuffle the input positions and the outputs shuffle with them, unchanged in content. softmax(QKᵀ)V has no notion of "before" or "after." But language is ordered — "dog bites man" ≠ "man bites dog." So order must be injected.

Learned absolute positions (the old way). GPT-2 adds a learned vector pos_emb[t] to the embedding at position t. Simple, but: it can't represent positions past the trained maximum (no pos_emb[5000] if you trained to 1024), and it encodes absolute position when what attention really cares about is relative offset.

RoPE (Rotary Position Embedding, the modern default). Instead of adding a position vector, RoPE rotates the query and key vectors by an angle proportional to their position. Group the d dimensions into d/2 pairs (x_{2i}, x_{2i+1}); treat each pair as a 2-D vector and rotate it by angle pos · θ_i, where each pair has its own frequency:

$$ \theta_i = \text{base}^{-2i/d}, \qquad \begin{pmatrix} x'{2i} \ x'{2i+1} \end{pmatrix} = \begin{pmatrix} \cos(pos,\theta_i) & -\sin(pos,\theta_i) \ \sin(pos,\theta_i) & \cos(pos,\theta_i) \end{pmatrix} \begin{pmatrix} x_{2i} \ x_{2i+1} \end{pmatrix}. $$

That is _rope_vector in the lab. Low-i pairs rotate fast (high frequency, capture nearby offsets); high-i pairs rotate slowly (low frequency, capture long-range offsets) — a Fourier-feature ladder of position scales. base (default 10000; 1e6 for long-context models) sets how slowly the slowest pairs turn.

The key property — it encodes relative position. A rotation preserves length, so RoPE never changes a vector's norm (test_rope_preserves_norm), and position 0 is the identity (test_rope_position_dependent). The magic is in the dot product: rotating q by pos_q·θ and k by pos_k·θ, the attention score q'·k' depends only on the difference pos_q − pos_k, not on the absolute positions. Attention naturally cares about offset ("how far back is the thing I'm copying?"), and RoPE hands it relative offset for free.

   pair i of a vector, before and after rotation by angle pos·θ_i:

            ^                      the (even, odd) pair is a point in 2-D;
            │   • x'  (rotated)    rotation spins it by pos·θ_i around the origin.
            │  ╱                   |x'| = |x|  (norm preserved).
            │ ╱  ) angle = pos·θ_i  q rotated by pos_q·θ, k by pos_k·θ:
            │╱___________>          q'·k' depends only on (pos_q − pos_k).

Production significance. RoPE is in Llama, Mistral, Qwen, Gemma, and most current open models. Because it's relative and frequency-based, you can extend the context window after training by rescaling the frequencies — NTK-aware scaling and YaRN are exactly "tweak the RoPE base / frequencies so positions beyond the trained length still get sensible angles." That whole long-context literature is RoPE math.

Common misconception. "RoPE is added to the embedding like GPT-2's position embedding." It is a rotation applied to Q and K inside attention, not an additive vector on the residual stream. (The lab applies it once on the residual stream up front for simplicity — see the README's honest limits — but in real models the rotation lives inside each attention layer so it acts on the scores.)

Chapter 8: RMSNorm vs LayerNorm

Why normalize at all. Stacking dozens of layers, activations can drift to very large or very small magnitudes, which makes gradients explode or vanish and training unstable. A normalization layer rescales each position's vector to a controlled magnitude before each sublayer, keeping the numbers in a sane range so the network trains smoothly. It is a stability device, not a feature extractor.

LayerNorm (the original). For a vector x of width d, subtract the mean, divide by the standard deviation, then scale and shift with learned per-dimension weight and bias:

$$ y_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}},\gamma_i + \beta_i, \qquad \mu = \frac{1}{d}\sum_i x_i,\quad \sigma^2 = \frac{1}{d}\sum_i (x_i-\mu)^2. $$

It centers (zero mean) and scales (unit variance). The lab's layer_norm does exactly this; test_layer_norm_zero_mean_unit_var checks both.

RMSNorm (the modern default). Drop the mean-centering and the bias entirely; just divide by the root-mean-square and scale:

$$ y_i = \frac{x_i}{\sqrt{\frac{1}{d}\sum_j x_j^2 + \epsilon}},\gamma_i. $$

The lab's rms_norm does this; test_rms_norm_scales_to_unit_rms checks the output RMS is 1. It is cheaper (no mean, no variance subtraction, no bias term — fewer ops and fewer parameters) and, empirically, works just as well. That is why Llama, Mistral, and Gemma all use RMSNorm.

Pre-norm vs post-norm. The original transformer put the norm after the sublayer and residual add (post-norm); modern models put it before the sublayer (pre-norm), which is far more stable for deep stacks because the residual stream stays an unnormalized "clean" highway (Chapter 10). The lab is pre-norm.

   LayerNorm:  center (− mean) → scale (÷ std) → ×weight + bias      (4 reductions + bias)
   RMSNorm:                       scale (÷ rms) → ×weight            (1 reduction, no bias)
                ↑ skips mean-centering and the shift — cheaper, equally good

Production significance. It's a config choice with measurable wins in throughput and parameter count at scale, and "why did the field move LayerNorm → RMSNorm, pre-norm → post-norm" is a common architecture-depth interview question. The answer: cheaper + just as good (RMSNorm), and more stable for depth (pre-norm).

Common misconception. "RMSNorm is a lossy approximation of LayerNorm." It's a different normalization (no centering), not an approximation — and on transformer activations the centering turns out not to matter much, so you pay less for the same stability.

Chapter 9: The Feed-Forward / SwiGLU — Where the Parameters Live

What it is. After attention mixes information across positions, the feed-forward network (FFN / MLP) transforms each position independently — it's a per-token nonlinear function. The classic GPT FFN is two linear layers with a GELU in between: up to a wider hidden dimension d_ff (usually 4·d_model), a nonlinearity, then down back to d_model.

Why a nonlinearity (GELU). Without a nonlinearity between linear layers, the composition is just another linear layer — no extra expressive power. GELU is the classic smooth gate (the lab's gelu uses the tanh approximation): roughly the identity for large positive inputs, squashing negatives toward zero, test_gelu_shape_of_curve checks the shape.

SwiGLU (the modern FFN). Modern models replace the single up-projection + GELU with a gated unit. Two parallel projections of the input — a gate (passed through SiLU) and an up — are multiplied element-wise, then projected down:

$$ \text{SwiGLU}(x) = \big(\underbrace{\text{SiLU}(x W_{\text{gate}})}{\text{gate}} \odot \underbrace{(x W{\text{up}})}{\text{value}}\big), W{\text{down}}, \qquad \text{SiLU}(z) = z\cdot\sigma(z). $$

The gate lets the network modulate each hidden unit multiplicatively — a learned, data-dependent "how much of this feature passes" — which empirically beats a plain GELU MLP at equal parameters. That is swiglu_ffn (with _silu) in the lab. (Because SwiGLU has three matrices instead of two, real models shrink d_ff to ~8/3·d_model to keep the parameter count matched.)

Where the parameters live. Per layer, attention is four d_model × d_model matrices (4·d²), while the FFN is three matrices of roughly d_model × d_ff with d_ff ≈ 4·d_model (~8·d² for a plain MLP, ~3 · 8/3 · d² = 8d² for SwiGLU). So the FFN holds roughly 2/3 of every layer's parameters — and thus ~2/3 of the whole model's. This is why Phase 00's 2N/token rule treats the MLP as the dominant term, and why Mixture-of-Experts puts its experts in the FFN: that's where the parameters (and the capacity) are.

   per transformer layer (rough):
     attention:  Wq Wk Wv Wo            ≈ 4·d²          ~1/3 of the layer
     FFN:        W_gate W_up W_down      ≈ 8·d²          ~2/3 of the layer
                                        ───────
                              the MLP is where the model "stores what it knows"

Production significance. Quantization (P06), pruning, and MoE all target the FFN because that's the parameter mass. "Where are the parameters in a transformer?" with the 2/3 answer is a frequent senior interview question.

Common misconception. "Attention is the expensive, parameter-heavy part." Attention dominates the conversation and the long-context compute (O(T²)), but the FFN dominates the parameter count and most of the FLOPs at typical sequence lengths.

Chapter 10: The Residual Stream & the Pre-Norm Block

What it is. A transformer block does not replace its input; it adds a correction to it. The running sum that every block reads from and writes back to is the residual stream — a [T, d_model] highway running the full depth of the model.

The pre-norm block. Each block is two sublayers, each wrapped in norm → sublayer → add:

$$ X \leftarrow X + \text{MHA}(\text{RMSNorm}(X)), \qquad X \leftarrow X + \text{FFN}(\text{RMSNorm}(X)). $$

That is transformer_block in the lab: add_matrices(X, attn) and add_matrices(X, ffn) are the residual adds — the same add_matrices you wrote in Chapter 1's helpers.

Why the residual stream matters. Two reasons:

  1. Gradients flow. The +X means the gradient has a direct, unobstructed path back through every layer (the derivative of X + f(X) w.r.t. X includes the identity). This is what lets you train 80-layer models without vanishing gradients — the original ResNet insight, inherited by transformers.
  2. Layers compose by accumulation. Each sublayer reads the current stream, computes a small correction, and adds it back. The stream is a shared workspace; interpretability research views it as a "communication channel" where heads write features that later heads read.
   residual stream  ──┬─────────────────────────┬──────────────►  (continues to next block)
                      │                          │
                  RMSNorm                    RMSNorm
                      │                          │
                     MHA                        FFN
                      │                          │
                      └──►(+)                    └──►(+)
                      add back to the stream     add back to the stream

Production significance. Pre-norm + residual is the reason deep transformers are trainable; the choice of pre- vs post-norm visibly affects training stability and is a real architecture decision. Phase 03 (autograd) will differentiate exactly this structure, and the clean gradient path is why it works.

Common misconception. "Normalization rescales the residual stream." In pre-norm, the norm is applied to a copy fed into the sublayer; the residual stream itself stays un-normalized (that's the point — it's a clean highway). The lab does exactly this: n1 = rms_norm(X) is consumed by attention, but the add is X + attn, not n1 + attn.

Chapter 11: The Full Forward Pass — Embedding to Logits

Putting it together. Now every piece composes into gpt_forward:

$$ \text{ids} \xrightarrow{\text{embed}} X \xrightarrow{\text{+RoPE}} X \xrightarrow{\text{N blocks}} X \xrightarrow{\text{RMSNorm}} X \xrightarrow{\text{LM head}} \text{logits } [T, \text{vocab}]. $$

Step by step, exactly as the lab does it:

   1. embedding lookup:   X = [embed[tid] for tid in token_ids]      [T, d_model]
   2. position:           X = rope(X, [0,1,...,T-1])                 inject order
   3. mask:               mask = causal_mask(T)                      no peeking ahead
   4. blocks:             for block in blocks: X = transformer_block(X, block, mask)
   5. final norm:         X = [rms_norm(row, final_norm_w) for row in X]
   6. LM head:            logits = linear(X, lm_head)                [T, vocab]

What the output means. Row t of the [T, vocab] logits is the model's score for every possible next token, given tokens 0..t. To generate, you softmax-and-sample row t to pick token t+1 (sampling is the next phase's territory). During training, you compare all rows to the true next tokens at once — that's the parallelism the causal mask buys you (Chapter 4).

Determinism. build_params(seed=…) draws every weight from one seeded random.Random, so the same seed gives the same logits — test_gpt_forward_deterministic checks that same-seed runs match and different seeds differ. Determinism is a lab requirement (LAB-STANDARD) and a debugging superpower: a forward pass that's reproducible is one you can bisect.

Production significance. This is model(input_ids).logits in any framework. nanoGPT's GPT.forward is this exact sequence in ~30 lines; an HF LlamaForCausalLM is this plus GQA, real batching, and fused kernels. You now know what every one of those lines is doing.

Common misconception. "The model outputs a token." It outputs logits — a score per vocab entry, per position. Turning logits into a token (greedy / temperature / top-k / top-p sampling) is a separate, deliberate step, and a frequent source of "the model is repetitive / incoherent" bugs.

Chapter 12: FlashAttention — The IO-Aware Fused Version

The problem with naive attention. The attention you built materializes the full [T, T] score matrix in memory (scores = Q·Kᵀ), softmaxes it, then multiplies by V. For long sequences that T² matrix is huge — and worse, on a GPU it gets written to and read from slow HBM (high-bandwidth memory) multiple times. From Phase 00's roofline: attention is memory-bandwidth-bound, and the bottleneck is moving that T² matrix, not the FLOPs.

The FlashAttention idea. Never write the full score matrix to HBM at all. Tile Q, K, V into blocks that fit in fast on-chip SRAM, and compute attention block by block, maintaining a running ("online") softmax — keeping a running max and running normalizer so you can fold each new block into the output without ever having seen the whole row at once. The math is identical to the attention you wrote; the memory traffic is what changes.

   naive:    Q·Kᵀ → [T,T] in HBM → softmax in HBM → ·V        (reads/writes the T² matrix)
   flash:    for each block of K,V:                            (T² never touches HBM)
               load Q-block, K-block, V-block into SRAM
               update running max m, running sum l, running output o   (online softmax)
             → exact same result, a fraction of the HBM traffic

Why it's a big deal. Same answer, but it turns attention from HBM-bound to compute-bound by slashing memory traffic — multiple-× faster and O(T) memory instead of O(T²), which is what makes long context affordable. It is a pure systems optimization: no approximation, no change to the model, just IO-awareness. This is the bridge from "I understand the attention algorithm" (this lab) to "I understand why production attention is fast" (Phase 00's roofline made concrete).

Production significance. FlashAttention (and v2/v3) is the default attention kernel in PyTorch SDPA, vLLM, and every serious serving stack. When torch.nn.functional.scaled_dot_product_attention runs fast, this is usually why. The lab's "build a block/tiled softmax" extension is a from-scratch taste of the core trick.

Common misconception. "FlashAttention is a faster approximation of attention." It is exact — the online-softmax trick reorders the computation so the result is mathematically identical. It trades a bit of recomputation for far less memory traffic; it does not drop any terms.

Lab Walkthrough Guidance

The lab (lab-01-attention-forward-pass) turns these chapters into code. Suggested order (matches the file top-to-bottom; helpers precede their users):

  1. Linear algebra (matmul, transpose, add_vectors, add_matrices) — Chapter 1's helpers. add_matrices is the residual op. Get the validation (ragged/shape) right; everything builds on it.
  2. softmax — Chapter 3. The whole lesson is the max-subtraction; -inf must map to 0.
  3. Norms (layer_norm, rms_norm) — Chapter 8. RMSNorm has no mean and no bias.
  4. linear — the affine map used by every projection.
  5. causal_mask + scaled_dot_product_attention — Chapters 2–4. The 1/√d_k scale, the mask add, softmax per row, weighted sum of V. This is the core.
  6. Multi-head (_split_heads, _concat_heads, multi_head_attention) — Chapter 5. A reshape plus per-head SDPA plus W_O.
  7. RoPE (_rope_vector, rope) — Chapter 7. Rotate (even, odd) pairs; preserve the norm.
  8. FFN (gelu, _silu, swiglu_ffn) — Chapter 9. Gate ⊙ up, then down.
  9. Block + forward (transformer_block, build_params, gpt_forward) — Chapters 10–11.
  10. KV-cache (attention_with_kv_cache) — Chapter 6. The soul test — test_kv_cache_equals_full_recompute must pass to tolerance.

Run red → green: pytest test_lab.py -v, then LAB_MODULE=solution pytest test_lab.py -v, then python solution.py for the worked numbers.

Success Criteria

  • All tests pass against your lab.py and against LAB_MODULE=solution.
  • You can write softmax(Q·Kᵀ/√d_k + mask)·V from memory and explain every term — including the variance argument for √d_k and why softmax subtracts the max.
  • You can explain the causal mask as additive -∞ → zero softmax weight → autoregression, and why a mask leak is a silent, catastrophic training bug.
  • You can explain multi-head as parallel subspaces (a reshape, not extra params) and how GQA/MQA shrink the KV-cache (tying back to Phase 00's cache math).
  • You can explain RoPE's rotation, why it preserves the norm, and why the QK score is relative.
  • You can state where the parameters live (~2/3 FFN) and contrast RMSNorm with LayerNorm and pre- vs post-norm.
  • You can prove the KV-cache equals the full causal recompute's last row, and explain why (the last query under the mask sees exactly the cached tokens).
  • You can sketch why FlashAttention is faster (IO-aware, never materializes T²) and exact.

Interview Q&A

  • "Walk me through scaled dot-product attention." softmax(Q·Kᵀ/√d_k + mask)·V: Q/K/V are learned projections of the input; the dot product scores query-key similarity; √d_k keeps the softmax unsaturated; softmax makes a distribution; ·V is the weighted average of value rows.
  • "Why divide by √d_k?" The dot product over d_k dims has variance ≈ d_k; without the scale, wider heads produce large scores that saturate softmax to near one-hot and kill the gradient. /√d_k restores ~unit variance so training stays in the soft, high-gradient regime.
  • "What exactly does the causal mask do?" Adds -∞ to scores for future positions; softmax(-∞)=0, so position i puts zero weight on any j>i. It makes a parallel training pass compute honest next-token prediction; a leak lets the model cheat and then fail at inference.
  • "Why multiple heads?" To attend to several relationships in parallel subspaces (syntax, coreference, position) instead of averaging them into one. It's a reshape of the same parameter budget, not more parameters.
  • "RoPE vs learned absolute positions?" RoPE rotates Q/K by pos·θ_i so the attention score depends on relative offset, preserves the norm, and extrapolates to longer contexts (NTK/YaRN are RoPE-frequency tweaks). Learned absolute embeddings can't exceed their trained max and encode absolute, not relative, position.
  • "RMSNorm vs LayerNorm? Pre- vs post-norm?" RMSNorm drops mean-centering and bias — cheaper, equally good, the modern default. Pre-norm (normalize the sublayer input, keep a clean residual highway) is more stable for deep stacks than the original post-norm.
  • "Where do the parameters live?" ~2/3 in the feed-forward (~8d² vs attention's ~4d² per layer). That's why the 2N/token rule treats the MLP as dominant and why MoE puts experts there.
  • "Why is the KV-cache correct, and what does it save?" Under a causal mask the current query attends to exactly the past tokens, which is the cache — so incremental decode equals the full recompute's last row. It turns decode from O(T²) to O(T) by never recomputing past K/V.
  • "How does GQA help, given decode is memory-bound?" Fewer K/V heads → smaller KV-cache → fewer bytes read per token → higher throughput on a bandwidth-bound workload, at negligible quality cost.
  • "What is FlashAttention and is it an approximation?" An IO-aware, tiled, online-softmax attention that never writes the T² score matrix to HBM — same exact result, far less memory traffic, O(T) memory. Not an approximation.
  • "Why does the model output logits, not a token?" It scores every vocab entry per position; turning a logit row into a token (greedy/temperature/top-k/top-p) is a separate sampling step and a common source of generation-quality bugs.
  • "Why apply normalization at all?" Deep stacks let activation magnitudes drift, exploding or vanishing gradients; the norm rescales each position to a controlled magnitude so training stays stable.

References

  • Vaswani et al., Attention Is All You Need (2017) — the transformer, scaled dot-product and multi-head attention, the √d_k scale.
  • Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021) — RoPE.
  • Zhang & Sennrich, Root Mean Square Layer Normalization (2019) — RMSNorm.
  • Shazeer, GLU Variants Improve Transformer (2020) — SwiGLU; and Fast Transformer Decoding: One Write-Head is All You Need (2019) — MQA.
  • Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023).
  • Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022), and FlashAttention-2 (2023).
  • Karpathy, nanoGPT — the cleanest readable GPT forward pass; and Let's build GPT from scratch.
  • Alammar, The Illustrated Transformer — the canonical visual walkthrough of Q/K/V and multi-head.
  • PyTorch docs — torch.nn.functional.scaled_dot_product_attention, torch.nn.MultiheadAttention, torch.nn.RMSNorm.