"""Roofline, MFU Budget & the Latency Napkin — YOUR implementation. Fill in every `# TODO`. Signatures, docstrings and validation contracts are given. pytest test_lab.py -v # red until you implement LAB_MODULE=solution pytest test_lab.py -v # the reference (must be green) python solution.py # the worked example ORDER OF WORK 1. arithmetic_intensity, ridge_point, roofline_throughput, roofline_report 2. mfu, hfu, mfu_budget 3. prefill_seconds, decode_seconds, interactive_latency, chips_for_latency_budget, weights_fit_chips 4. tile_efficiency, kv_cache_bytes, gqa_saving, decode_batch_intensity, depth_vs_width THE MONEY TESTS test_prefill_reproduces_the_talks_number ~5.8 s on one v5e chip test_a_4x4_station_brings_prefill_under_the_half_second_limit test_a_smaller_model_needs_fewer_chips <- the whole economic argument test_decode_stays_memory_bound_even_at_huge_batch UNITS: seconds are seconds, bytes are bytes, FLOP/s is a rate. Say the unit out loud. """ import math # -------------------------------------------------------------------------------------- # Hardware — given to you. peak_flops is DENSE bf16 (no structured sparsity). # -------------------------------------------------------------------------------------- HARDWARE = { # name peak bf16 FLOP/s HBM bytes HBM bandwidth B/s board watts "H100": (990e12, 80e9, 3.35e12, 700), "A100-80": (312e12, 80e9, 2.03e12, 400), "TPU v5e": (197e12, 16e9, 0.819e12, 170), "TPU v5p": (459e12, 95e9, 2.77e12, 600), } def _chip(chip): if chip not in HARDWARE: raise ValueError(f"unknown chip {chip!r}; known: {sorted(HARDWARE)}") return HARDWARE[chip] def arithmetic_intensity(flops, bytes_moved): """FLOPs performed per byte moved. The x-axis of the roofline. This single number decides which resource you are fighting. It is a property of the ALGORITHM (and its blocking), not of the chip. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def ridge_point(chip): """peak FLOP/s / HBM bandwidth — the arithmetic intensity at which a kernel stops being memory-bound and becomes compute-bound. H100 ~ 296 FLOP/byte. Anything below that is waiting on memory, and adding FLOPs is free while adding bytes is not. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def roofline_throughput(chip, intensity): """Achievable FLOP/s at a given arithmetic intensity. below the ridge: bandwidth * intensity (the sloped roof) above the ridge: peak (the flat roof) """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def roofline_report(chip, flops, bytes_moved): """Full roofline analysis of one operation.""" # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def mfu(model_flops, seconds, chip, n_chips): """Model FLOPs Utilization: useful model FLOPs achieved / peak available. 'Useful' means the 6ND arithmetic the MODEL requires — NOT including recomputation from activation checkpointing. That distinction is MFU vs HFU (below). """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def hfu(model_flops, seconds, chip, n_chips, recompute_factor=8.0 / 6.0): """Hardware FLOPs Utilization: counts recomputation as useful work. Full activation checkpointing adds roughly one extra forward pass, taking hardware FLOPs per token from ~6N to ~8N. HFU is therefore ALWAYS >= MFU. When someone quotes a utilization number, ask which one they mean — it is a ~33% difference for free. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def mfu_budget(matmul_s, vector_s, memory_s, comms_s, optimizer_s): """Decompose wall-clock into where the matmul unit was NOT busy. THE POINT OF THIS FUNCTION: an MFU of 35% is not a failure grade, it is an accounting identity. A chip is several units with wildly different throughput, and a transformer must use all of them: matmul unit ~1000 TFLOP/s <- 'peak' is THIS number alone vector unit ~50-100x slower <- gelu, softmax, norms HBM ~3 TB/s <- activations in and out interconnect ~0.05-0.9 TB/s <- collectives The breakdown is your optimization agenda: comms high -> overlap them; memory high -> fuse kernels; vector high -> fuse the norm into the matmul epilogue. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def prefill_seconds(n_tokens, n_params, chip, n_chips, mfu_frac=1.0): """Prefill: process the whole prompt at once. 2N FLOPs per token, forward only. COMPUTE bound — you amortize each weight read across thousands of tokens, so arithmetic intensity is high. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def decode_seconds(n_tokens, n_params, chip, n_chips, bytes_per_param=2, bandwidth_efficiency=1.0, kv_bytes_per_token=0.0): """Decode: generate one token at a time. MEMORY-BANDWIDTH bound — every generated token re-reads EVERY weight to perform a trivial amount of arithmetic. With the model sharded over n_chips, each chip reads its 1/n_chips slice. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def interactive_latency(prefill_tokens, decode_tokens, n_params, chip, n_chips, scaffolding_s=0.25, prefill_mfu=1.0): """End-to-end latency for one interactive turn, itemized. Reproduces the structure of Feinberg's napkin math: an agent turn is a prefill of the incremental context plus a short decode, inside a hard latency budget, with a fixed overhead for scaffolding, load balancing, request validation and KV retrieval. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def chips_for_latency_budget(budget_s, prefill_tokens, decode_tokens, n_params, chip, scaffolding_s=0.25, max_chips=4096, prefill_mfu=1.0): """Smallest power-of-two chip count that meets a latency budget. Returns None if even `max_chips` cannot do it — which is the signal that the answer is not more hardware, it is a smaller model. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def weights_fit_chips(n_params, chip, bytes_per_param=2): """Minimum chips just to HOLD the weights — before any latency consideration.""" # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def tile_efficiency(dim, tile=128): """Fraction of a padded matmul dimension that is real work. Systolic arrays and tensor cores operate on fixed tiles. A dimension that is not a multiple of the tile gets padded, and you pay for the padding. This is why production models use dimensions like 4096, 8192, 11008. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def kv_cache_bytes(n_layers, n_kv_heads, d_head, seq_len, batch, bytes_per_elem=2): """2 (K and V) * L * n_kv * d_head * T * B * bytes.""" # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def gqa_saving(n_layers, n_query_heads, n_kv_heads, d_head, seq_len, batch): """What sharing KV heads buys, as a ratio and in bytes. THE cleanest example of inference co-design: an architecture decision made BEFORE training that changes serving throughput by the group factor, at negligible quality cost. Once trained, it is frozen forever. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def decode_batch_intensity(n_params, n_layers, n_kv_heads, d_head, seq_len, batch, bytes_per_param=2): """Arithmetic intensity of a decode step at a given batch size. Weights are read ONCE and shared across the batch; the KV cache is per request. This is exactly why batching is the primary decode optimization — it is the only lever that moves you rightward on the roofline. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def depth_vs_width(n_params_target, d_model, n_layers, params_per_layer_fn=None): """Serial-depth cost of a shape choice. Depth is SERIAL: layer k+1 cannot start until layer k finishes, so depth costs decode latency directly and adds pipeline stages to synchronize. Width is parallel and matmul-friendly. Co-design usually means: as wide as quality allows, as shallow as quality tolerates. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError if __name__ == "__main__": print("Fill in the TODOs, then compare against: python solution.py")