"""Roofline, MFU Budget & the Latency Napkin — reference solution. The arithmetic that turns a product requirement ("respond in under a second") into an architecture constraint ("the model must be smaller"): * the roofline model: arithmetic intensity vs the ridge point * MFU decomposed into where the matmul unit was NOT busy — an agenda, not a grade * MFU vs HFU (activation checkpointing counted as useful work) * prefill vs decode as two different machines * shape co-design: tile quantization, GQA, depth-vs-width * Feinberg's Llama3-70B-on-v5e napkin math, reproduced Pure stdlib, deterministic. Run `python solution.py` for a worked example. """ import math # -------------------------------------------------------------------------------------- # Hardware. Single source of truth. peak_flops is DENSE bf16 (no structured sparsity). # Sources: NVIDIA H100/A100 whitepapers; Google Cloud TPU v5e/v5p documentation. # -------------------------------------------------------------------------------------- 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] # ====================================================================================== # 1. The roofline # ====================================================================================== 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. """ if bytes_moved <= 0: raise ValueError("bytes_moved must be positive") if flops < 0: raise ValueError("flops cannot be negative") return flops / bytes_moved 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. """ peak, _hbm, bw, _w = _chip(chip) return peak / bw 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) """ if intensity < 0: raise ValueError("intensity cannot be negative") peak, _hbm, bw, _w = _chip(chip) return min(peak, bw * intensity) def roofline_report(chip, flops, bytes_moved): """Full roofline analysis of one operation.""" ai = arithmetic_intensity(flops, bytes_moved) ridge = ridge_point(chip) achievable = roofline_throughput(chip, ai) peak, _hbm, bw, _w = _chip(chip) return { "arithmetic_intensity": ai, "ridge_point": ridge, "bound_by": "compute" if ai >= ridge else "memory", "achievable_flops": achievable, "fraction_of_peak": achievable / peak, "seconds": flops / achievable if achievable > 0 else float("inf"), # What to fix: below the ridge, cutting bytes helps; above it, cutting FLOPs does. "lever": "reduce bytes moved" if ai < ridge else "reduce FLOPs", } # ====================================================================================== # 2. MFU — an accounting identity, not a grade # ====================================================================================== 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). """ if seconds <= 0 or n_chips <= 0: raise ValueError("seconds and n_chips must be positive") peak, _hbm, _bw, _w = _chip(chip) return model_flops / (seconds * n_chips * peak) 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. """ if recompute_factor < 1.0: raise ValueError("recompute_factor must be >= 1 (you cannot compute less than once)") return mfu(model_flops * recompute_factor, seconds, chip, n_chips) 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. """ parts = {"matmul": matmul_s, "vector": vector_s, "memory": memory_s, "comms": comms_s, "optimizer": optimizer_s} for k, v in parts.items(): if v < 0: raise ValueError(f"{k} time cannot be negative") total = sum(parts.values()) if total <= 0: raise ValueError("total time must be positive") out = {f"{k}_fraction": v / total for k, v in parts.items()} out["mfu"] = matmul_s / total out["total_seconds"] = total # The single biggest non-matmul consumer — where to spend your next week. losses = {k: v for k, v in parts.items() if k != "matmul"} out["biggest_lever"] = max(losses, key=losses.get) return out # ====================================================================================== # 3. Prefill and decode are two different machines # ====================================================================================== 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. """ if not 0 < mfu_frac <= 1: raise ValueError("mfu_frac must be in (0, 1]") if n_tokens <= 0 or n_chips <= 0: raise ValueError("n_tokens and n_chips must be positive") peak, _hbm, _bw, _w = _chip(chip) return (2.0 * n_params * n_tokens) / (n_chips * peak * mfu_frac) 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. """ if not 0 < bandwidth_efficiency <= 1: raise ValueError("bandwidth_efficiency must be in (0, 1]") if n_tokens <= 0 or n_chips <= 0: raise ValueError("n_tokens and n_chips must be positive") _peak, _hbm, bw, _w = _chip(chip) bytes_moved = n_tokens * (n_params * bytes_per_param + kv_bytes_per_token) return bytes_moved / (n_chips * bw * bandwidth_efficiency) 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. """ if scaffolding_s < 0: raise ValueError("scaffolding_s cannot be negative") p = prefill_seconds(prefill_tokens, n_params, chip, n_chips, prefill_mfu) d = decode_seconds(decode_tokens, n_params, chip, n_chips) return {"prefill_s": p, "decode_s": d, "scaffolding_s": scaffolding_s, "total_s": p + d + scaffolding_s, "prefill_fraction": p / (p + d) if (p + d) > 0 else 0.0} 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. """ if budget_s <= scaffolding_s: raise ValueError("budget must exceed the fixed scaffolding overhead") n = 1 while n <= max_chips: r = interactive_latency(prefill_tokens, decode_tokens, n_params, chip, n, scaffolding_s, prefill_mfu) if r["total_s"] <= budget_s: return n n *= 2 return None def weights_fit_chips(n_params, chip, bytes_per_param=2): """Minimum chips just to HOLD the weights — before any latency consideration.""" _peak, hbm, _bw, _w = _chip(chip) return math.ceil(n_params * bytes_per_param / hbm) # ====================================================================================== # 4. Shape co-design # ====================================================================================== 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. """ if dim <= 0 or tile <= 0: raise ValueError("dim and tile must be positive") return dim / (math.ceil(dim / tile) * tile) 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.""" for name, v in (("n_layers", n_layers), ("n_kv_heads", n_kv_heads), ("d_head", d_head), ("seq_len", seq_len), ("batch", batch)): if v <= 0: raise ValueError(f"{name} must be positive") return 2 * n_layers * n_kv_heads * d_head * seq_len * batch * bytes_per_elem 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. """ if n_kv_heads > n_query_heads: raise ValueError("n_kv_heads cannot exceed n_query_heads") if n_query_heads % n_kv_heads != 0: raise ValueError("n_query_heads must be divisible by n_kv_heads") mha = kv_cache_bytes(n_layers, n_query_heads, d_head, seq_len, batch) gqa = kv_cache_bytes(n_layers, n_kv_heads, d_head, seq_len, batch) return {"mha_bytes": mha, "gqa_bytes": gqa, "ratio": mha / gqa, "group_size": n_query_heads // n_kv_heads, "bytes_saved": mha - gqa} 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. """ flops = 2.0 * n_params * batch kv = kv_cache_bytes(n_layers, n_kv_heads, d_head, seq_len, batch) return arithmetic_intensity(flops, n_params * bytes_per_param + kv) 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. """ if n_layers <= 0 or d_model <= 0: raise ValueError("n_layers and d_model must be positive") return {"n_layers": n_layers, "d_model": d_model, "serial_steps_per_token": n_layers, "params_per_layer": n_params_target / n_layers, "aspect_ratio": d_model / n_layers} # ====================================================================================== # Worked example # ====================================================================================== def main(): print("=" * 78) print("ROOFLINE, MFU BUDGET & THE LATENCY NAPKIN") print("=" * 78) print("\n[1] The ridge point: where each chip stops being memory-bound") for chip in HARDWARE: peak, hbm, bw, _w = HARDWARE[chip] print(f" {chip:10s} peak {peak/1e12:6.0f} TFLOP/s bw {bw/1e12:5.2f} TB/s" f" -> ridge {ridge_point(chip):6.0f} FLOP/byte") print("\n[2] Roofline of three operations on an H100") ops = [ ("big matmul (8192^3)", 2 * 8192 ** 3, 3 * 8192 * 8192 * 2), ("decode, batch=1", 2 * 70e9, 70e9 * 2), ("layernorm (8k x 8k)", 5 * 8192 * 8192, 2 * 8192 * 8192 * 2), ] for name, f, b in ops: r = roofline_report("H100", f, b) print(f" {name:22s} AI={r['arithmetic_intensity']:8.1f} " f"{r['bound_by']:7s} {r['fraction_of_peak']:6.1%} of peak " f"-> {r['lever']}") print("\n[3] MFU is an accounting identity, not a grade") b = mfu_budget(matmul_s=100, vector_s=45, memory_s=60, comms_s=50, optimizer_s=25) print(f" MFU (matmul busy) {b['mfu']:6.1%}") for k in ("vector", "memory", "comms", "optimizer"): print(f" lost to {k:12s} {b[k + '_fraction']:6.1%}") print(f" -> biggest lever: {b['biggest_lever']}. That is your next week's work.") print("\n[4] MFU vs HFU — always ask which you are being shown") model_flops, secs, chips = 6 * 70e9 * 1e9, 3600.0, 1000 print(f" MFU (honest) {mfu(model_flops, secs, 'H100', chips):6.1%}") print(f" HFU (counts recompute as useful) {hfu(model_flops, secs, 'H100', chips):6.1%}") print(" Same run. Same hardware. ~33% apart, for free.") print("\n[5] THE NAPKIN: Llama3-70B on TPU v5e, one agent turn") print(" 128k context, 8k incremental prefill, 128 decode tokens,") print(" 1.0 s budget, 250 ms of it scaffolding. Compute-bound prefill assumed.") N = 70e9 print(f" chips just to HOLD the weights: {weights_fit_chips(N, 'TPU v5e')}") for chips in (1, 4, 16, 64, 128): r = interactive_latency(8192, 128, N, "TPU v5e", chips) ok = "OK " if r["total_s"] <= 1.0 else "MISS" print(f" {chips:4d} chips: prefill {r['prefill_s']:7.3f}s " f"decode {r['decode_s']:7.3f}s total {r['total_s']:7.3f}s {ok}") print(" -> 1 chip gives ~5.8s of prefill; a 4x4 = 16-chip station brings prefill") print(" under 0.5s, matching the talk. But batch-1 DECODE is ~3.8x prefill,") print(" so the full budget needs ~64 chips. Two phases, two bottlenecks.") print("\n[6] The conclusion that pays for this whole phase: shrink the model") for n in (70e9, 35e9, 17e9, 8e9): need = chips_for_latency_budget(1.0, 8192, 128, n, "TPU v5e") print(f" N={n/1e9:5.1f}B -> {str(need):>5s} chips for a 1.0s turn" f" ({weights_fit_chips(n, 'TPU v5e')} just to hold weights)") print(" Halving N halves BOTH columns. This is the economic case for Flash.") print("\n[7] Shape co-design I: tile quantization") for d in (11000, 11008, 8192, 4097, 4096): print(f" d_ff={d:6d} -> {tile_efficiency(d):6.1%} of the padded matmul is real work") print("\n[8] Shape co-design II: GQA, decided before training and frozen forever") for kv in (64, 8, 1): g = gqa_saving(80, 64, kv, 128, 8192, 32) print(f" {kv:2d} kv heads (group {g['group_size']:2d}): " f"{g['gqa_bytes']/1e9:7.1f} GB of KV cache " f"{g['ratio']:5.1f}x smaller than MHA") print("\n[9] Shape co-design III: batching is the only decode lever") ridge = ridge_point("H100") print(f" H100 ridge point: {ridge:.0f} FLOP/byte") for batch in (1, 8, 64, 256, 1024): ai = decode_batch_intensity(70e9, 80, 8, 128, 8192, batch) print(f" batch={batch:5d}: AI={ai:7.1f} FLOP/byte " f"{'MEMORY' if ai < ridge else 'COMPUTE'} bound " f"({ai/ridge:5.1%} of the way to the ridge)") print(" -> even batch 1024 is memory-bound. Decode never becomes compute-bound") print(" at realistic sizes, which is why GQA and quantization matter so much.") print("=" * 78) if __name__ == "__main__": main()