Lab 01 — Attention, RoPE & a Tiny GPT Forward Pass

Phase: 02 — The Transformer From Scratch Difficulty: ⭐⭐⭐☆☆ (the algebra is small; the mechanism is the whole point) Time: 4–6 hours

Every engineer says "I use transformers." Far fewer can build one with no framework — and that gap is exactly what a senior interview probes. This lab makes you write the mechanism by hand: scaled dot-product attention, the causal mask that makes it autoregressive, multi-head attention, RoPE rotary positions, RMSNorm/LayerNorm, a SwiGLU feed-forward, a pre-norm transformer block, the full GPT forward pass (embedding → blocks → norm → LM head → logits), and — the soul of the lab — the KV-cache, proven numerically identical to a full causal recompute for the last position. Pure stdlib lists and math; no numpy, no torch; deterministic from a seed. When generation is wrong, slow, or leaks the future, this is the muscle memory that lets you fix it.

What you build

  • matmul / transpose / add_vectors / add_matrices — the linear algebra a transformer is built from. add_matrices is the residual-stream op X + sublayer(X); we write the matmul triple loop so the m·n·k MAC cost is visible.
  • softmax — the numerically stable, max-subtracting softmax that turns scores into a distribution that sums to 1 and maps a masked -inf to exactly zero weight.
  • layer_norm / rms_norm — the two normalizations: LayerNorm (mean + variance + bias) and the cheaper, mean-free RMSNorm that Llama/Mistral/Gemma use.
  • linear — the affine projection y = x·W (+ b) used everywhere (Q/K/V/O, the MLP, the LM head).
  • causal_mask — the [T, T] additive mask (0 on/below the diagonal, -inf above) that is autoregression: token i may attend to j only if j <= i.
  • scaled_dot_product_attention — softmax(Q·Kᵀ / sqrt(d_k) + mask) · V: the core mechanism, including the 1/sqrt(d_k) scale that keeps the softmax from saturating.
  • multi_head_attention (with _split_heads / _concat_heads) — project, split into heads, attend per head in parallel subspaces, concat, output-project with Wo.
  • rope (with _rope_vector) — rotary position embedding: rotate each (even, odd) dim pair by pos·θ_i; preserves the norm and makes the attention dot product depend on relative offset.
  • gelu / swiglu_ffn (with _silu) — the activations and the gated down(silu(x·W_gate) ⊙ (x·W_up)) feed-forward where ~2/3 of the parameters live.
  • transformer_block — one pre-norm block: X += MHA(norm1(X)); X += FFN(norm2(X)).
  • build_params / gpt_forward — deterministically build a tiny GPT from a seed, then run the full forward pass: token ids → [T, vocab] logits.
  • attention_with_kv_cache — one incremental decode step against a key/value cache, proven equal to the full causal recompute's last row — the entire justification for the cache.

Key concepts

ConceptWhat to understand
softmax(Q·Kᵀ/√d_k)·Vdot product = similarity; /√d_k stops the softmax saturating; softmax → convex combination; ·V mixes value rows
The causal maskadditive -inf above the diagonal → exp(-inf)=0 weight → the future cannot leak into the present (autoregression)
Multi-headsplit d_model into h subspaces of width d_head = d_model/h; each head attends to a different relationship, then concat + Wo
RoPErotate (even, odd) pairs by pos·θ_i, θ_i = base^(-2i/d); rotation preserves norm, so the QK dot product depends only on relative position
RMSNorm vs LayerNormRMSNorm drops the mean subtraction and the bias — cheaper, equally good; the modern default
SwiGLU FFNgated MLP down(silu(gate) ⊙ up); the FFN holds ~2/3 of all parameters
Pre-norm residual streamnormalize the sublayer input, add the output back; the residual highway is why depth trains
KV-cachestore past K/V so decode is O(T) not O(T²); identical to full causal attention's last row
Numerical stabilitysoftmax subtracts the max so [1000, 1001, 1002] does not overflow — that detail is the lesson

Files

FilePurpose
lab.pyskeleton with # TODO markers — your implementation
solution.pycomplete reference; python solution.py runs a worked example
test_lab.pythe proof — run it red, make it green
requirements.txtpytest only (pure stdlib otherwise)

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

  • All tests pass against your implementation (and LAB_MODULE=solution).
  • You can explain why softmax subtracts the max and why test_softmax_numerically_stable_large_logits would NaN without it.
  • You can explain why a masked (-inf) score contributes exactly zero probability, and why that makes test_causal_attention_zero_in_upper_triangle give out[0] == V[0].
  • You can explain what the 1/sqrt(d_k) factor does and why test_sqrt_dk_scaling_is_present detects it (the unscaled distribution is sharper).
  • You can explain why RoPE leaves the vector norm unchanged (test_rope_preserves_norm) and why position 0 is the identity (test_rope_position_dependent).
  • You can explain why test_kv_cache_equals_full_recompute is the soul of the lab: the last query under a causal mask attends to exactly tokens 0..t, which is exactly the cache, so the incremental result must equal the full recompute's last row.
  • You can state where the parameters live (~2/3 in the FFN) and why GQA/MQA shrink the KV-cache.

How this maps to the real stack

The miniatureThe production mechanismWhere to verify it
scaled_dot_product_attentiontorch.nn.functional.scaled_dot_product_attention — same softmax(QKᵀ/√d_k + mask)·V, but fused and dispatched to a fast kernelPyTorch SDPA docs; is_causal=True for the mask
multi_head_attention (Wq/Wk/Wv/Wo, split/concat)torch.nn.MultiheadAttention / the attention module in nanoGPT — same project→split→attend→concat→WonanoGPT model.py CausalSelfAttention; HF LlamaAttention
causal_mask (-inf upper triangle)the attn_mask / is_causal flag every decoder passes to SDPA; torch.triu builds itPyTorch scaled_dot_product_attention(is_causal=True)
rope (_rope_vector)the RoPE in Llama/Mistral/Qwen/Gemma — apply_rotary_pos_emb, with a rope_theta baseHF LlamaRotaryEmbedding; the config's rope_theta (10000, or 1e6 for long context)
rms_normtorch.nn.RMSNorm / LlamaRMSNorm — the modern default over nn.LayerNormHF LlamaRMSNorm; the config's rms_norm_eps
swiglu_ffnthe gated MLP in Llama/Mistral (gate_proj, up_proj, down_proj with SiLU)HF LlamaMLP; ~2/3 of sum(p.numel() …)
transformer_blockone decoder layer (pre-norm + residual) stacked n_layers deepHF LlamaDecoderLayer; nanoGPT Block
gpt_forwardmodel(input_ids).logits — embedding → blocks → final norm → LM headnanoGPT GPT.forward; HF *ForCausalLM
attention_with_kv_cachethe past_key_values / DynamicCache every decode loop carries; vLLM's PagedAttention is the block-table version of thisHF use_cache=True, past_key_values; vLLM PagedAttention

Limits of the miniature (be honest in the interview): we apply RoPE once on the residual stream for simplicity, whereas real models rotate Q and K inside each attention layer (the rotation must touch the scores, not the residual); we use plain multi-head, not GQA/MQA (the lever that actually shrinks the KV-cache in production); attention here is the naive O(T²) materialize-the- scores form, not the IO-aware fused FlashAttention that never writes the full score matrix to HBM; there is no training, no dropout, no batching, and tiny dims; and our matmul is a Python triple loop, not a tensor-core GEMM. The algorithm is faithful; the systems engineering is the rest of the curriculum (P06 quantization, P09 serving/PagedAttention, P17 kernels).

Extensions (build these on real hardware)

  • Move RoPE inside the attention layer: rotate Q and K per head before the dot product, and confirm the QK score depends only on the relative offset (pos_q − pos_k).
  • Add GQA: give K/V fewer heads than Q (share K/V across query-head groups) and measure the KV-cache shrink (ties back to Phase 00's KV-cache math).
  • Implement a block (tiled) softmax that computes attention without ever materializing the full [T, T] score matrix — the core idea of FlashAttention's online softmax — and check it matches scaled_dot_product_attention to tolerance.
  • Add a greedy / temperature / top-k sampler on top of gpt_forward and generate a sequence token-by-token using attention_with_kv_cache; verify it equals the prefill-then-decode path.
  • Port one block to PyTorch and assert your gpt_forward logits match torch.nn.functional.scaled_dot_product_attention + nn.RMSNorm to a tight tolerance.

Interview / resume

  • Talking points: "Why divide attention scores by √d_k?" "How does the causal mask make a transformer autoregressive — what exactly is -inf doing?" "Why does RoPE encode relative position, and why is that better than learned absolute embeddings?" "Where do most of the parameters live, and why is the FFN ~2/3?" "Prove the KV-cache is correct — why is it identical to recomputing?"
  • Resume bullet: Implemented a transformer from scratch in pure Python — scaled dot-product attention with the causal mask, multi-head attention, RoPE, RMSNorm, a SwiGLU feed-forward, the full GPT forward pass, and an incremental KV-cache proven numerically identical to the full causal recompute — mapping each piece to its PyTorch / nanoGPT / FlashAttention equivalent.