#!/usr/bin/env python3 """ roofline.py — arithmetic intensity, ridge points, and "is this kernel memory-bound or compute-bound?" for Projects 1, 13 and 14. The model, from first principles -------------------------------- A processor has two hard ceilings: * peak compute P [FLOP/s] -- how fast it can multiply-add * peak bandwidth B [byte/s] -- how fast it can pull operands from DRAM A kernel does W FLOPs and moves Q bytes between DRAM and the chip. Its *arithmetic intensity* is I = W / Q [FLOP/byte] Time is bounded below by BOTH ceilings, so achievable performance is perf = min(P, I * B) [FLOP/s] This is the roofline (Williams, Waterman & Patterson, CACM 2009). The two regimes meet at the *ridge point* I_ridge = P / B [FLOP/byte] Below I_ridge you are memory-bound: the multipliers idle waiting for operands, and buying more FLOP/s buys you nothing. Above it you are compute-bound. Q is *DRAM* traffic, not total loads. A value read from L2 costs no DRAM traffic. That is why tiling works: it does not change W, it shrinks Q by making each loaded byte serve more FLOPs, which slides the kernel right along the roofline until it hits the compute ceiling. Hardware numbers and their asterisks ------------------------------------ Vendor headline FLOP/s are frequently quoted *with 2:4 structured sparsity*, which doubles the number and does not apply to a dense GEMM. Every entry below records the dense figure and names the trap. Always cross-check a datasheet footnote before you use its number as a denominator. Usage ----- python3 roofline.py table python3 roofline.py gemm --m 4096 --n 4096 --k 4096 --dtype bf16 --hw h100 python3 roofline.py decode --params 7e9 --hw h100 """ from __future__ import annotations import argparse # name -> (dense peak FLOP/s at the stated dtype, DRAM bandwidth byte/s, note) HW = { # NVIDIA H100 SXM5: 989.4 TFLOP/s dense BF16/FP16 tensor core. # The widely-quoted 1979 TFLOP/s is the *sparse* (2:4) figure -- do not use # it for a dense GEMM. HBM3 at 3.35 TB/s. "h100": (989.4e12, 3.35e12, "dense BF16 tensor core; 1979 TF/s figure is 2:4 sparse"), # NVIDIA A100 80GB SXM: 312 TFLOP/s dense BF16, HBM2e 2.039 TB/s. "a100": (312e12, 2.039e12, "dense BF16 tensor core; 624 TF/s figure is 2:4 sparse"), # A modern server CPU socket, AVX-512, ~32 cores at ~2.5 GHz sustained: # 32 cores * 2 FMA units * 16 fp32 lanes * 2 flop * 2.5e9 = 5.1 TFLOP/s. # 8 channels of DDR5-4800 = 8 * 4.8e9 * 8 byte = 307 GB/s. "cpu-server": (5.1e12, 307e9, "AVX-512 fp32, 32c @2.5GHz sustained; 8ch DDR5-4800"), # An Apple M-series performance cluster, for the machine this journey is # most likely to be run on. Figures are order-of-magnitude, measure yours. "laptop": (1.0e12, 200e9, "order-of-magnitude only -- MEASURE your own machine"), } DTYPE_BYTES = {"fp32": 4, "tf32": 4, "bf16": 2, "fp16": 2, "fp8": 1, "int8": 1} def ridge(hw: str) -> float: p, b, _ = HW[hw] return p / b def classify(intensity: float, hw: str) -> str: r = ridge(hw) if intensity < r * 0.9: return f"MEMORY-bound (I={intensity:.1f} < ridge {r:.0f})" if intensity > r * 1.1: return f"COMPUTE-bound (I={intensity:.1f} > ridge {r:.0f})" return f"at the ridge (I={intensity:.1f} ~ {r:.0f})" def gemm(m: int, n: int, k: int, dtype: str, hw: str) -> dict: """C[m,n] = A[m,k] @ B[k,n]. W = 2*m*n*k FLOPs (one multiply + one add per inner-product term). Q depends entirely on whether the operands fit in cache: * best case -- each matrix streamed exactly once: Q_min = (m*k + k*n + m*n) * bytes * worst case -- no reuse at all, B re-read for every row of A: Q_max = (m*k + m*k*n + m*n) * bytes The gap between these two is the entire subject matter of tiling, and the reason a naive triple loop can be 100x off peak while doing exactly the same arithmetic. """ b = DTYPE_BYTES[dtype] w = 2 * m * n * k q_min = (m * k + k * n + m * n) * b q_max = (m * k + m * k * n + m * n) * b p, bw, _ = HW[hw] i_min, i_max = w / q_max, w / q_min return { "flops": w, "q_best_bytes": q_min, "q_worst_bytes": q_max, "intensity_best": i_max, "intensity_worst": i_min, "verdict_best": classify(i_max, hw), "verdict_worst": classify(i_min, hw), "t_compute_bound_s": w / p, "t_memory_bound_best_s": q_min / bw, "t_memory_bound_worst_s": q_max / bw, } def decode(params: float, batch: int, dtype: str, hw: str) -> dict: """One autoregressive decode step of a dense transformer. Per generated token, every weight is used exactly once, so W = 2 * N * batch (2 FLOPs per multiply-accumulate) The weights are read from HBM once per step and shared across the whole batch, so (ignoring KV cache, which is second-order at small context) Q = N * bytes_per_param Therefore I = 2 * N * batch / (N * bytes) = 2 * batch / bytes Read that again: **arithmetic intensity in decode depends on batch size and nothing else.** Not on model size. Not on how clever your kernel is. This single identity is why continuous batching exists, why a batch-1 chatbot wastes ~99% of an H100's multipliers, and why the prefill and decode phases of the same model want completely different hardware treatment. """ b = DTYPE_BYTES[dtype] w = 2 * params * batch q = params * b p, bw, _ = HW[hw] intensity = w / q t_mem = q / bw t_cmp = w / p return { "intensity": intensity, "verdict": classify(intensity, hw), "t_memory_s": t_mem, "t_compute_s": t_cmp, "t_bound_s": max(t_mem, t_cmp), "tokens_per_s": batch / max(t_mem, t_cmp), "mfu": (w / max(t_mem, t_cmp)) / p, } def batch_for_compute_bound(dtype: str, hw: str) -> float: """Solve I = 2*batch/bytes = ridge => batch = ridge * bytes / 2.""" return ridge(hw) * DTYPE_BYTES[dtype] / 2.0 def _table() -> None: print(f"{'hardware':<12} {'peak dense':>14} {'bandwidth':>12} {'ridge':>10} note") print("-" * 100) for name, (p, b, note) in HW.items(): print(f"{name:<12} {p/1e12:>11.1f} TF/s {b/1e9:>9.0f} GB/s " f"{p/b:>7.0f} F/B {note}") print() print("Batch size required to leave the memory-bound regime in decode:") print(f"{'hardware':<12} {'bf16':>8} {'fp16':>8} {'fp8':>8} {'fp32':>8}") print("-" * 48) for name in HW: row = " ".join(f"{batch_for_compute_bound(d, name):>6.0f}" for d in ("bf16", "fp16", "fp8", "fp32")) print(f"{name:<12} {row}") print() print("Quantizing weights to fp8 HALVES the batch you need to saturate the") print("multipliers, because it halves Q while leaving W unchanged. That is") print("the real reason low precision wins at inference -- not the faster") print("arithmetic, the smaller operands.") def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="cmd", required=True) sub.add_parser("table") g = sub.add_parser("gemm") g.add_argument("--m", type=int, default=4096) g.add_argument("--n", type=int, default=4096) g.add_argument("--k", type=int, default=4096) g.add_argument("--dtype", default="bf16", choices=DTYPE_BYTES) g.add_argument("--hw", default="h100", choices=list(HW)) d = sub.add_parser("decode") d.add_argument("--params", type=float, default=7e9) d.add_argument("--batch", type=int, default=1) d.add_argument("--dtype", default="bf16", choices=DTYPE_BYTES) d.add_argument("--hw", default="h100", choices=list(HW)) a = ap.parse_args() if a.cmd == "table": _table() elif a.cmd == "gemm": r = gemm(a.m, a.n, a.k, a.dtype, a.hw) print(f"GEMM {a.m}x{a.n}x{a.k} {a.dtype} on {a.hw}") print(f" work {r['flops']/1e9:12.2f} GFLOP") print(f" DRAM traffic best {r['q_best_bytes']/1e6:12.2f} MB " f"I={r['intensity_best']:8.1f} {r['verdict_best']}") print(f" DRAM traffic worst {r['q_worst_bytes']/1e6:12.2f} MB " f"I={r['intensity_worst']:8.1f} {r['verdict_worst']}") print(f" t if compute-bound {r['t_compute_bound_s']*1e3:12.3f} ms") print(f" t if memory-bound {r['t_memory_bound_best_s']*1e3:12.3f} ms (best)" f" / {r['t_memory_bound_worst_s']*1e3:.3f} ms (no reuse)") elif a.cmd == "decode": print(f"Dense decode, {a.params/1e9:.0f}B params, {a.dtype}, on {a.hw}") print(f"{'batch':>7} {'I (F/B)':>9} {'step ms':>9} {'tok/s':>10} " f"{'MFU':>7} regime") for bs in (1, 2, 4, 8, 16, 32, 64, 128, 256, 512): r = decode(a.params, bs, a.dtype, a.hw) print(f"{bs:>7} {r['intensity']:>9.1f} {r['t_bound_s']*1e3:>9.3f} " f"{r['tokens_per_s']:>10.0f} {r['mfu']*100:>6.1f}% " f"{r['verdict'].split(' ')[0]}") return 0 if __name__ == "__main__": raise SystemExit(main())