Mathematically essential vs engineering optimization
The whole project splits cleanly into two layers, and keeping them separate is the difference between understanding TinyGPT and memorizing it.
- Mathematically essential — the code that defines what function the model computes. Change
it and the model computes something different. The oracle is
python_ref/model.py: the cleanest statement of the math, the thing every other backend is validated against. - Engineering optimization — the code that computes that same function faster, smaller, or on different silicon. Delete it and the model computes the identical function, just slower or bigger.
The litmus test: “If I removed this, would the model compute a different function, or the same function more slowly?” Different ⇒ essential. Same ⇒ optimization.
The two columns
| Mathematically essential (defines the function) | What it computes |
|---|---|
| Token + position embeddings | tokens → vectors; inject order |
Self-attention: softmax(QKᵀ/√d ⊙ causalmask) · V | mix information across positions |
| Multi-head | run attention in H subspaces, concat |
| LayerNorm + residuals | stabilize + let gradients skip |
| MLP + GELU | per-position non-linear transform |
| Output projection → logits | hidden → vocab scores |
| Cross-entropy loss | scores → a scalar to minimize |
| Backprop (chain rule through all the above) | loss → per-parameter gradients |
| Adam/AdamW update | gradients → weight step |
| Engineering optimization (preserves the function) | What it speeds/shrinks | Must preserve |
|---|---|---|
| Matmul tiling + 4×4 register blocking (study_guide §8) | arithmetic intensity → compute-bound | bit-parity matmul |
vec4<f32> packed loads (§9) | 128-bit memory transactions | same result (the bug that wasn’t) |
| FlashAttention-2 (online-softmax, §7) | O(T²)→O(T) memory; recompute on backward | identical attention math |
| KV cache | decode O(T²)→O(T) compute | same logits |
| Quantization (4-bit/8-bit) | weights smaller | approximately same weights (lossy!) |
| WASM / WebGPU backends | run on CPU-SIMD / GPU | parity with python_ref |
Memory64 (-sMEMORY64) | address heaps >4 GB | same kernels |
| Gradient checkpointing | recompute activations → less RAM | identical gradients |
Two subtleties worth internalizing:
- FA2 is optimization, not math. It computes the exact same attention output — it just never
materializes the
[B,H,T,T]matrix and recomputes scores on the backward pass. The online-softmax trick is algebraically identical to textbook softmax; the win is structural, not numerical. - Quantization is the one optimization that isn’t lossless. It approximates the weights, so it legitimately changes the output a little — which is exactly why it shows up as a real accuracy delta in evals (we measured bf16 vs 4-bit = +4 on tool-calling), not as a parity bug.
Loss drift — the number that polices the boundary
How do you prove an optimization preserved the math? Loss drift.
Loss drift = the percent difference between an optimized backend’s training loss and the Python reference’s loss at the same step N, same seed, same data, same hyperparameters.
drift = |loss_backend − loss_ref| / loss_ref.
It is the single number that says “this optimization is still computing the right function”:
- < ~2.5–5% drift at 50 steps = numerically equivalent. Floating-point addition isn’t associative, so a tiled GPU kernel sums partial products in a different order than the reference; that reorder produces small, bounded deviation. This is expected and benign.
- Larger drift = a real bug, not “just precision.” It means the optimization changed the
function — a wrong gradient, a transposed tile, a bad bind-group. Example (study_guide §9):
a
vec4integration passed its standalone kernel test but drove training loss to 88.67 while the reference sat at 2.94 — a ~30× drift. The cause was a bind-group access-mode mismatch returning wrong data; the one-line fix dropped drift below 1%.
This is also why parity testing beats standalone kernel benchmarks: a standalone benchmark checks the kernel in isolation (ideal shapes, fresh cache) and answers “is it fast + self-consistent?” Drift checks the kernel inside the real pipeline and answers “does it still compute the model’s function?” Only the second question is the one that ships. (See study_guide §9.)
Why this lens matters
- It tells you what mastery requires. You should be able to reimplement the essential column
from scratch, without AI — that’s the real exam (it’s
python_ref/model.py, ~300 lines). The optimization column you should understand in principle (why tiling, why FA2, why a KV cache), but any specific kernel is substitutable; the principle is the knowledge, the WGSL is an artifact. - It tells you where bugs live. A wrong output is usually an essential-layer bug (the math). A wrong speed/memory with correct output is an optimization-layer issue. A correct standalone kernel with wrong end-to-end loss is an optimization that broke the math — caught only by drift.
- It tells you what’s safe to change. Swapping a matmul kernel is safe if drift stays low. Changing the attention formula is a math change and needs re-deriving the backward pass.
References
- Oracle for the essential column:
python_ref/model.py. - Optimization-layer deep dives: study_guide.md §7 (FA2), §8 (register blocking), §9 (parity + the drift bug); online-softmax; WebGPU execution model.
- The math, taught from scratch: roadmap Phases 1–4; sessions 1–8 in this dir.