Phase 02 — The Transformer From Scratch

The phase where "I use transformers" becomes "I can build one." Phase 00 taught you to treat a model as a physical system with a compute, memory, and bandwidth budget; Phase 01 turned text and images into token IDs. Now we open the box those tokens flow through and build the actual mechanism — attention, the causal mask, multi-head, RoPE, RMSNorm, the SwiGLU MLP, the residual stream, the full forward pass, and the KV-cache — by hand, in pure Python, with no framework hiding the arithmetic. This is the architecture every later phase trains, quantizes, serves, and attacks.

Why this phase exists

You cannot debug, optimize, or reason about a thing you can only call. The senior who can answer "why √d_k," "what exactly is the mask doing," "where do the parameters live," and "why is the KV-cache correct" is the one trusted to own the model stack. Every one of those answers is a few lines of code you will write here.

The transformer is also the load-bearing abstraction of the entire curriculum. Autograd (P03) differentiates this forward pass. Fine-tuning (P05) injects adapters into these linear layers. Quantization (P06) compresses these weights. Serving (P09) caches these keys and values with PagedAttention. Long-context, FlashAttention, GQA, MoE — all are modifications to the block you build in this lab. Get the forward pass in your hands and the rest of the track is changes to a thing you already understand, not new magic.

Three ideas do the heavy lifting, and you should leave the phase able to defend each from first principles:

1. Attention is content-based, weighted retrieval. softmax(Q·Kᵀ/√d_k)·V scores every query against every key (similarity), normalizes to a distribution, and mixes the value rows by it. The causal mask makes it autoregressive; multiple heads make it multi-relational.
2. Position is injected, not inherent. Attention is permutation-equivariant on its own, so order has to be added — by RoPE (rotate Q/K so the score depends on the relative offset), the modern default over learned absolute embeddings.
3. The residual stream is the highway, and most of the model is the MLP. Pre-norm blocks read from the residual stream, compute a correction, and write it back; ~2/3 of the parameters live in the feed-forward, not in attention.

Concept map — a transformer block

   token ids ──► embedding lookup ──► + position (RoPE) ──► residual stream X  ┐
                                                                               │
   ┌───────────────────── one pre-norm block (× n_layers) ────────────────────┤
   │                                                                          │
   │      X ──► RMSNorm ──► Multi-Head Attention ──► (+) ──► X'               │
   │                          │  split d_model into h heads    ▲              │
   │                          │  softmax(Q·Kᵀ/√d_k + mask)·V   │ residual     │
   │                          └───────────────────────────────┘              │
   │                                                                          │
   │      X' ──► RMSNorm ──► SwiGLU FFN (~2/3 of params) ──► (+) ──► X''      │
   │                          down(silu(x·W_gate) ⊙ (x·W_up))    ▲           │
   │                          └──────────────────────────────────┘ residual  │
   └──────────────────────────────────────────────────────────────────────────┘
                                       │
        final RMSNorm ──► LM head (linear to vocab) ──► logits [T, vocab]
                                       │
                       (the causal mask: token i attends only to j ≤ i)

The lab

The lab is a runnable, test-verified miniature — see the lab standard. Run it red (pytest test_lab.py), make it green, then read solution.py. Pure stdlib, deterministic from a seed, offline.

Integrated scenario ideas

• Prove the KV-cache to a skeptic. Run full causal attention, then run the same input through attention_with_kv_cache step by step, and show the last-position output is identical — then explain why (the last query under a causal mask sees exactly tokens 0..t, which is the cache).
• Find where the future leaks. Remove the causal mask and show a token's prediction changes when you edit a later token — the canonical training-data-leakage bug — then put the mask back.
• Locate the parameters. Count params in attention vs the FFN for a real config (e.g. Llama-style d_model, d_ff ≈ 8/3·d_model) and show the FFN is ~2/3 — connect to Phase 00's 2N/token.
• RoPE is relative. Take two tokens, RoPE them at positions (p, q) and at (p+k, q+k), and show the QK dot product is unchanged — the property that lets RoPE extrapolate context length.
• Norm ablation. Swap RMSNorm for LayerNorm in the block and show the forward pass still runs; discuss why the field moved to RMSNorm (cheaper, no bias, equally good).

Deliverables checklist

• lab.py passes pytest test_lab.py -v (and LAB_MODULE=solution does too).
• You can derive softmax(Q·Kᵀ/√d_k)·V and explain every term, including why √d_k is there.
• You can explain the causal mask as additive -inf → zero softmax weight → autoregression.
• You can explain why multiple heads exist and what _split_heads/_concat_heads do.
• You can explain RoPE's rotation, why it preserves the norm, and why the score is relative.
• You can state where the parameters live (~2/3 FFN) and contrast RMSNorm with LayerNorm.
• You can prove the KV-cache equals the full causal recompute's last row and say why.

Key takeaways

• Attention is weighted retrieval, the mask makes it causal. softmax(Q·Kᵀ/√d_k + mask)·V is the whole engine; the -inf mask is the one line that makes a transformer a language model instead of an encoder.
• Position is added on purpose. Self-attention is order-blind; RoPE injects relative position by rotation, which is why modern long-context tricks (NTK/YaRN scaling) are RoPE-base tweaks.
• Most of the model is the MLP. ~2/3 of the parameters are in the feed-forward — which is why the 2N/token rule from Phase 00 holds and why MoE puts its experts there.
• The KV-cache is correct because of the causal mask. Decode never recomputes past K/V, and the result is bit-for-(float)-bit the full recompute's last row — the identity that makes serving (P09) possible and PagedAttention worth building.

Next: Phase 03 — Autograd & Training Dynamics.