Lab 01 — Transformer FLOPs / Memory / Cost & Tradeoff Modeler

Phase: 00 — Foundations: Scaling Laws, FLOPs/Memory Math & Tradeoffs Difficulty: ⭐⭐☆☆☆ (the arithmetic is easy; the judgment is ⭐⭐⭐⭐⭐) Time: 3–4 hours

Round-1 of an AI engineering interview — and the first slide of every design review — is back-of-envelope arithmetic worn as architecture. This lab builds the calculator: training and inference FLOPs (6ND, 2N), weight memory across precisions, the KV-cache (with GQA), Chinchilla-optimal data, the roofline (compute- vs memory-bound), the memory-bandwidth ceiling on decode throughput, the serving cost per million tokens, and a weighted tradeoff resolver whose entire point is that the same two designs flip winners when the workload's priorities change — which is why "best model" is a meaningless phrase.

What you build

  • training_flops / forward_flops_per_token / inference_flops — the compute budget: 6·N·D to train, 2·N per token to run. The numbers that size a cluster and a bill.
  • weight_memory_bytes — N × bytes_per_param across fp32/fp16/int8/int4 — the lever that decides how many GPUs you need.
  • kv_cache_bytes — 2 · L · T · d_model · bytes · batch, with the GQA branch (d_kv = d_model · n_kv/n_head) — the memory wall everyone forgets and the thing that actually OOMs a server.
  • chinchilla_optimal_tokens / chinchilla_optimal_params — the compute-optimal D≈20N allocation (and its inverse).
  • arithmetic_intensity / roofline_ridge_intensity / roofline_attainable_flops / is_memory_bound — the roofline: min(peak, bw·I), ridge = peak/bw, and the check that proves decode is memory-bandwidth-bound.
  • decode_tokens_per_sec / cost_per_million_tokens / mfu — throughput ceiling (and how batching amortizes the weight read), $/1M tokens, and the honest efficiency metric.
  • Candidate (five dials) + resolve_tradeoff — score two designs under the workload's weights and return the winner and the deciding dial (the one sentence for your decision record).

Key concepts

ConceptWhat to understand
6ND train, 2N runforward = 2N/token (one MAC/param); backward ≈ 2× forward ⇒ 6N/token ⇒ 6ND total
Weights ≠ the memory boundweights are fixed; the KV-cache grows with batch × context and usually dominates
GQA shrinks the cachefewer K/V heads ⇒ d_kv scaled by n_kv/n_head ⇒ proportionally smaller KV memory
Decode is memory-boundI≈1 FLOP/byte vs an A100 ridge of ~156 ⇒ batch & quantize, don't buy FLOPs
Batching amortizes the readone weight read serves B sequences ⇒ throughput ~∝ batch
Chinchilla D≈20Ntraining-optimal; deployed models over-train past it because inference cost is forever
The deciding dial"best model" is a category error; best for these weights is engineering

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 training is 6ND and inference is 2N/token (MACs × passes).
  • You can compute the KV-cache for (L=32, T=4096, d=4096, fp16) ≈ 2 GB in your head and why GQA 8/32 makes it 0.54 GB.
  • You can explain why test_decode_tokens_per_sec_and_batching shows ~linear scaling with batch, and why that is the entire economic argument for batched serving.
  • You can explain why is_memory_bound(1.0, …) is True for decode and why that one fact governs inference optimization.
  • You can explain why test_weights_flip_the_winner is the whole point of the lab.
  • Given any "deploy model X" prompt you can, in two minutes, write FLOPs + weight/KV memory + a $/1M estimate + the deciding dial on paper.

How this maps to the real stack

The miniatureThe production mechanismWhere to verify it
training_flops / inference_flopsthe 6ND / 2N rules used to size every training run and serving fleet; DeepSpeed's FLOPs profiler and the Chinchilla paper use exactly thisdeepspeed.profiling.flops_profiler; the model card's "training compute"
weight_memory_bytesthe GPU-count decision; fp16→int8→int4 is the difference between 2 GPUs and 1model.get_memory_footprint() (HF); bitsandbytes 8/4-bit
kv_cache_bytes (+ GQA)the actual serving memory wall; vLLM reserves blocks for it and num_key_value_heads (GQA) is the real configvLLM gpu_memory_utilization, max_model_len; the config's num_key_value_heads
arithmetic_intensity / rooflinethe model behind every "are we compute- or memory-bound" decision; NVIDIA Nsight and the roofline paperNVIDIA Nsight Compute roofline; torch.profiler
decode_tokens_per_secthe bandwidth ceiling vLLM/TensorRT-LLM throughput is measured against; batching is the levervLLM throughput logs; the HBM bandwidth in the GPU datasheet
cost_per_million_tokensthe build-vs-buy number you compare against API pricingcloud GPU hourly price ÷ measured tok/s
resolve_tradeoffa scored model/serving decision with the deciding dial written into an ADRyour own design-review template

Limits of the miniature (be honest in the interview): the FLOP rules ignore attention's O(T²) term (matters at long context), embeddings, and the LM head; real KV-cache adds block fragmentation and over-allocation (which is why PagedAttention exists); decode throughput is a ceiling — real systems lose to kernel launch overhead, sampling, and Python; and a tradeoff score is a structured conversation, not a number you can sum mechanically across every workload.

Extensions (build these on real hardware)

  • Add the attention O(T²) FLOPs term and find the context length where it overtakes the MLP.
  • Add training_memory_bytes(N, optimizer="adamw", dtype="fp16", activation_checkpointing=bool) using the ~16 bytes/param rule, and show why a 7B model needs ZeRO to train on one GPU.
  • Add paged_kv_waste(seq_lens, block_size) modeling fragmentation, and show how PagedAttention's block allocator reclaims it (preview of P09).
  • Plot a real roofline for your GPU (peak_flops, bandwidth from the datasheet) and mark where prefill, decode, and a big GEMM land on it.
  • Pull params with sum(p.numel() for p in model.parameters()) on a real HF model and check your weight_memory_bytes against model.get_memory_footprint().

Interview / resume

  • Talking points: "How many GPUs to serve a 70B at 1k QPS, and what limits it?" "Is decode compute- or memory-bound — prove it." "Bigger model or distill a smaller one — defend with numbers." "Compose the $/1M for self-hosting vs an API."
  • Resume bullet: Built a transformer cost-and-tradeoff model that computes training/inference FLOPs, weight and KV-cache memory (with GQA), the compute-vs-bandwidth roofline, decode throughput and $/1M-token serving cost, and produces explainable, priority-weighted model-selection decisions for design reviews.