#!/usr/bin/env python3 """GPU memory, roofline, and cost arithmetic for the 'design ChatGPT' round. The output is a SCRIPT for what you say out loud. Derive it, do not recite it — in the round you will not have this, so run it until the arithmetic is reflex. python3 gpu_math.py --list python3 gpu_math.py --model llama-70b --gpu h100 python3 gpu_math.py --model llama-70b --gpu h200 --seq-len 8192 --batch 64 python3 gpu_math.py --model llama-8b --gpu h100 --dtype fp8 All hardware figures are vendor datasheet values; prices are 2026-reported order-of-magnitude anchors and should be re-checked before you quote them. """ from __future__ import annotations import argparse GIB = 2**30 # name -> (params_B, layers, hidden, attn_heads, kv_heads, head_dim) MODELS = { "llama-8b": (8.0, 32, 4096, 32, 8, 128), "llama-70b": (70.6, 80, 8192, 64, 8, 128), "llama-405b": (405.0, 126, 16384, 128, 8, 128), "mistral-7b": (7.2, 32, 4096, 32, 8, 128), "qwen-32b": (32.5, 64, 5120, 40, 8, 128), "gpt-oss-20b": (20.0, 48, 6144, 48, 8, 128), } # name -> (memory_GB, bandwidth_TBs, bf16_TFLOPs, fp8_TFLOPs, usd_per_hour) # # DENSE figures throughout. Vendor datasheets headline the *with-sparsity* number, # which is 2x the dense one and assumes 2:4 structured sparsity. LLM inference # weights are dense, so the sparsity path never applies and quoting it inflates # every machine-balance and MFU number by 2x. A100 has no FP8; the column holds # its dense INT8 rate instead. GPUS = { "a100-80": (80, 2.039, 312, 624, 1.80), "h100": (80, 3.35, 989.5, 1979, 2.50), "h200": (141, 4.8, 989.5, 1979, 3.80), "b200": (192, 8.0, 2250, 4500, 6.50), } DTYPE_BYTES = {"fp32": 4, "fp16": 2, "bf16": 2, "fp8": 1, "int8": 1, "int4": 0.5} def rule(title: str) -> None: print(f"\n{title}\n{'-' * max(len(title), 52)}") def main() -> None: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--model", default="llama-70b") parser.add_argument("--gpu", default="h100") parser.add_argument("--dtype", default="fp16", choices=sorted(DTYPE_BYTES)) parser.add_argument("--kv-dtype", default=None, choices=sorted(DTYPE_BYTES)) parser.add_argument("--seq-len", type=int, default=4096) parser.add_argument("--batch", type=int, default=None, help="override the fitted batch") parser.add_argument("--gpus", type=int, default=1, help="tensor-parallel degree") parser.add_argument("--activation-gib", type=float, default=4.0, help="activations + framework overhead reserve") parser.add_argument("--list", action="store_true") args = parser.parse_args() if args.list: print("\nModels:") for name, (p, layers, hidden, heads, kv, hd) in MODELS.items(): print(f" {name:<14} {p:>6.1f}B params {layers:>3} layers " f"{heads:>3} heads / {kv} kv head_dim {hd}") print("\nGPUs:") for name, (mem, bw, bf16, fp8, cost) in GPUS.items(): print(f" {name:<10} {mem:>4} GB {bw:>5.2f} TB/s " f"{bf16:>5.0f} TFLOP/s bf16 ${cost}/hr") return if args.model not in MODELS: raise SystemExit(f"unknown model {args.model}; try --list") if args.gpu not in GPUS: raise SystemExit(f"unknown gpu {args.gpu}; try --list") params_b, layers, hidden, heads, kv_heads, head_dim = MODELS[args.model] mem_gb, bw_tbs, bf16_tflops, fp8_tflops, usd_hr = GPUS[args.gpu] w_bytes = DTYPE_BYTES[args.dtype] kv_bytes = DTYPE_BYTES[args.kv_dtype or args.dtype] params = params_b * 1e9 print(f"\n{'=' * 60}") print(f"{args.model} @ {args.dtype} on {args.gpus} x {args.gpu}") print(f"{'=' * 60}") # ---- weights ---------------------------------------------------------- rule("1. Weights") weight_bytes = params * w_bytes weight_gib = weight_bytes / GIB per_gpu_weight_gib = weight_gib / args.gpus print(f" {params_b:.1f}B params x {w_bytes} B/param = {weight_gib:,.1f} GiB") if args.gpus > 1: print(f" tensor-parallel over {args.gpus} -> {per_gpu_weight_gib:,.1f} GiB per GPU") total_mem_gib = mem_gb * GIB / GIB print(f" GPU has {total_mem_gib:.0f} GiB -> " f"{'FITS' if per_gpu_weight_gib < total_mem_gib else 'DOES NOT FIT'}") if per_gpu_weight_gib >= total_mem_gib: need = -(-weight_gib // (total_mem_gib - args.activation_gib)) print(f" need at least {int(need)} GPUs just to hold the weights") # ---- kv cache --------------------------------------------------------- rule("2. KV cache — the thing that actually limits your batch") per_token = 2 * layers * kv_heads * head_dim * kv_bytes print(f" per token = 2 (K and V) x {layers} layers x {kv_heads} kv-heads") print(f" x {head_dim} head_dim x {kv_bytes} B = " f"{per_token:,.0f} B/token") print(f" per sequence at {args.seq_len:,} tokens = " f"{per_token * args.seq_len / GIB:.3f} GiB") mha_per_token = 2 * layers * heads * head_dim * kv_bytes print(f"\n Without GQA ({heads} kv-heads instead of {kv_heads}) it would be " f"{mha_per_token:,.0f} B/token") print(f" -> GQA is a {mha_per_token / per_token:.1f}x reduction in KV cache. " f"That is the") print(f" architectural decision that makes long context affordable at all.") # ---- batch ------------------------------------------------------------ rule("3. What is left for batch") available = (total_mem_gib - per_gpu_weight_gib - args.activation_gib) * args.gpus print(f" {total_mem_gib:.0f} GiB x {args.gpus} - {weight_gib:.1f} weights " f"- {args.activation_gib * args.gpus:.1f} activations/overhead") print(f" = {available:,.1f} GiB for KV cache") if available <= 0: print(" NEGATIVE — this model does not fit on this configuration.") return max_batch = int(available * GIB / (per_token * args.seq_len)) print(f" / {per_token * args.seq_len / GIB:.3f} GiB per sequence " f"= {max_batch:,} concurrent sequences at {args.seq_len:,} tokens") batch = args.batch or max(max_batch, 1) if args.batch: print(f" (using --batch {batch})") print("\n Say this in the round: the KV cache, not the weights, is the") print(" capacity constraint. Doubling context halves your batch.") # ---- roofline --------------------------------------------------------- rule("4. Roofline — why decode is bandwidth-bound") # Under tensor parallelism each rank holds 1/N of the weights and 1/N of the # KV heads and reads its own shard CONCURRENTLY, so the resources that matter # are the AGGREGATE bandwidth and FLOP/s of the group. Dividing aggregate # bytes by one GPU's bandwidth would overstate every latency by N -- and # therefore overstate $/token by N, which is the number people actually quote. bw_bytes = bw_tbs * 1e12 * args.gpus tflops = fp8_tflops if args.dtype in ("fp8", "int8") else bf16_tflops flops = tflops * 1e12 * args.gpus step_bytes = weight_bytes + per_token * args.seq_len * batch mem_time = step_bytes / bw_bytes decode_flops = 2 * params * batch compute_time = decode_flops / flops print(f" ONE DECODE STEP at batch {batch}:") print(f" must read {step_bytes / GIB:,.1f} GiB " f"(weights + KV) -> {mem_time * 1000:,.2f} ms at " f"{bw_tbs * args.gpus:,.2f} TB/s aggregate") print(f" must compute {decode_flops / 1e12:,.1f} TFLOP " f"-> {compute_time * 1000:,.2f} ms at " f"{tflops * args.gpus:,.0f} TFLOP/s aggregate (dense)") ratio = mem_time / compute_time if compute_time else float("inf") verdict = "MEMORY-BOUND" if ratio > 1 else "compute-bound" print(f" ratio {ratio:,.1f}x -> {verdict}") print(f" theoretical floor: {1 / mem_time:,.0f} steps/s = " f"{batch / mem_time:,.0f} tokens/s aggregate") if args.gpus > 1: print(f" NB: assumes perfect TP scaling. Real TP{args.gpus} pays an " f"all-reduce per layer;\n" f" budget 10-20% off this on NVLink, much worse across nodes.") prefill_flops = 2 * params * args.seq_len prefill_time = prefill_flops / flops prefill_mem = weight_bytes / bw_bytes print(f"\n ONE PREFILL of {args.seq_len:,} tokens (single sequence):") print(f" compute {prefill_flops / 1e12:,.1f} TFLOP -> " f"{prefill_time * 1000:,.1f} ms") print(f" memory {weight_bytes / GIB:,.1f} GiB -> " f"{prefill_mem * 1000:,.1f} ms") pverdict = "compute-bound" if prefill_time > prefill_mem else "memory-bound" print(f" -> prefill is {pverdict.upper()}") print(""" THIS IS THE WHOLE ROUND IN TWO LINES: prefill is compute-bound, decode is memory-bound. They are different workloads sharing one accelerator, which is why a naive scheduler lets one long prefill block everyone's decode — and why chunked prefill exists. Say that and you have earned the follow-up.""") # ---- cost ------------------------------------------------------------- rule("5. Cost per million tokens") tokens_per_s = batch / mem_time realized = tokens_per_s * 0.4 # a deliberately conservative efficiency factor cost_per_s = usd_hr * args.gpus / 3600 print(f" theoretical {tokens_per_s:,.0f} tok/s -> at 40% realized " f"efficiency {realized:,.0f} tok/s") print(f" ${usd_hr}/hr x {args.gpus} GPU = ${cost_per_s * 1e6:,.2f} per 1M GPU-seconds") print(f" cost per 1M output tokens = " f"${cost_per_s / realized * 1e6:,.2f}") print(""" Caveat you must state when quoting this: the 40% factor is a placeholder for scheduler overhead, ragged batches, and attention cost that this napkin ignores. Real numbers come from a benchmark, not from this script. Quoting a measured number beats quoting a modelled one every time — which is exactly why the Track D drill list includes 'benchmark a claim'.""") if __name__ == "__main__": main()