Lab 01 — Roofline, MFU Budget & the Latency Napkin

Build the arithmetic that turns "we want a real-time agent" into "the model must be smaller" — and reproduce the napkin math from Feinberg's Princeton talk.

The problem

A product manager says "sub-second responses." A serving engineer says "we'll add GPUs." Both are guessing. The roofline and the latency napkin turn that conversation into arithmetic, and the arithmetic usually says something neither of them expected: no amount of hardware fixes this; the model has to be smaller.

That conclusion is why a pre-training lead owns inference co-design. The levers that matter — n_kv_heads, depth vs width, tile alignment, parameter count — are all frozen the moment training starts.

What you build

Key concepts

Files

Run

pytest test_lab.py -v                       # your lab.py — red until you implement
LAB_MODULE=solution pytest test_lab.py -v   # the reference — must be green (50 passed)
python solution.py                          # the worked example

Where to start

1. arithmetic_intensity → ridge_point → roofline_throughput. The roofline is one min(). Get test_roofline_is_continuous_at_the_ridge green — the two roof segments must meet exactly.
2. mfu_budget, then test_perfect_mfu_when_only_matmul_runs. That boundary case is Feinberg's point: 100% requires a pure matmul loop, which is not a neural network.
3. prefill_seconds / decode_seconds. Then test_prefill_reproduces_the_talks_number — if you get ~5.8 s you have it right.
4. chips_for_latency_budget. Note it must return None when the budget is unreachable.
5. The co-design levers.

The traps:

• Prefill is 2N per token, forward only. 6N is training — a 3× error.
• decode_seconds must not contain peak. Decode is driven by bytes and bandwidth. If peak FLOP/s appears in it, the model is wrong.
• chips_for_latency_budget must terminate and return None rather than loop. "More hardware does not fix this" is a real answer and the function has to be able to give it.
• A budget below the scaffolding overhead is impossible — raise, don't return a number.
• gqa_saving must reject n_kv_heads > n_query_heads and non-divisible configurations.

Success criteria

• LAB_MODULE=solution pytest test_lab.py -v → 50 passed.
• Your lab.py reaches 50 passed.
• python solution.py runs and you can explain all nine sections.
• You reproduce ~5.8 s single-chip prefill and the 4×4 station conclusion.
• You can explain why decode beats prefill at batch 1, and the batching caveat.
• You can state why H100's ridge point is higher than A100's, and why that matters.

How this maps to the real stack

What is not a lie: the ridge-point definition, 2N per token for prefill, the KV-cache formula, the GQA ratio, and the fact that decode is bandwidth-bound. Those are exact, and those are what get asked about.

Extensions

1. Add the attention T² term. At 128k context the sequence-dependent attention matmuls dominate prefill (Phase 00, Break 1). Extend prefill_seconds and watch the napkin's conclusions change completely for long-context products.
2. Model continuous batching properly. Add a queue with arrival rates and compute p50/p99 latency versus throughput. That curve — the latency/throughput frontier — is what serving teams actually optimize, and this lab only shows its two endpoints.
3. Add speculative decoding. Model an acceptance rate α and a draft model of size N_d; compute the effective intensity gain and find where it stops paying.
4. Validate against real hardware. Run a 7B model in vLLM, measure tokens/sec at several batch sizes, and compare to decode_batch_intensity. The gap between your model and reality is the lesson.
5. Build the co-design search. Given a latency budget and a compute budget, search over (N, n_layers, d_model, n_kv_heads) for the configuration that maximizes predicted quality (using Phase 01's scaling law) subject to meeting the napkin. That is the actual job, and it is a genuinely good portfolio piece.

Interview / resume bullets

• "Built a roofline and MFU-accounting toolkit for LLM training and serving — arithmetic intensity versus ridge point with an explicit optimize-bytes-or-FLOPs verdict, MFU/HFU disambiguation, and wall-clock decomposition across matmul, vector, memory, collective and optimizer time to rank optimization work."
• "Reproduced Google DeepMind's published inference-scaling analysis from first principles: ~5.8 s single-chip prefill for a 70B model on TPU v5e, a 4×4 prefill station to meet a 0.5 s API limit, and the batch-1 result that decode costs ~3.8× prefill — concluding quantitatively that meeting an interactive latency budget requires a smaller model rather than more accelerators."
• "Quantified the irreversible inference co-design levers set at pre-training time: GQA group size (8× concurrent-request throughput at 8 KV heads versus 64), tile-aligned matrix dimensions, and depth-versus-width serial latency cost."
• Interview-ready: "Decode never becomes compute-bound at any realistic batch size — even 1024 reaches 17% of the ridge — because the KV cache grows with the batch while the weight read does not. That one fact explains batching, GQA, quantization and speculative decoding."