ML Performance mental-math drills

Block 1 practice — TPU / GPU Performance interviews. Say the reasoning aloud; target under 2:00 per problem.

Memorize these specs (dense BF16, no sparsity)

Chip	TFLOP/s (BF16)	HBM BW (TB/s)	Ridge (FLOP/B)	HBM (GB)	Shared/vector mem
NVIDIA H100 SXM	989	3.35	~295	80	228 KB smem/SM · 132 SMs
NVIDIA B200	2,250	7.7	~292	180	228 KB smem/SM
AMD MI300X	1,307	5.3	~247	192	64 KB LDS/CU · 304 CUs
TPU v5p	459	2.76	~166	95	~128 MB CMEM/VMEM-class on-chip
TPU v6e (Trillium)	918	1.64	~560	32	large on-chip VMEM
NVIDIA A100 (anchor)	312	2.0	~156	80	164 KB smem/SM · 108 SMs

Ridge point = peak FLOP/s ÷ bytes/s. Arithmetic intensity above the ridge → compute-bound; below → memory-bound. FP8 doubles the FLOP number, so the FP8 ridge is ~2×. Verify exact numbers against current datasheets — interviewers accept round figures but not wrong ratios.

Formulas — say these without thinking

GEMM (M,N,K), bf16: FLOPs = 2·M·N·K bytes ≈ 2·(MK + KN + MN) ⇒ AI ≈ MNK / (MK+KN+MN) (square N×N: AI ≈ N/3)

Decode step, batch 1: bytes/token ≈ params·dtype_bytes + KV read, KV = 2·layers·kv_heads·head_dim·seqlen·kv_bytes ⇒ tok/s ceiling = BW / bytes_per_token

Prefill attention: AI grows with seqlen → compute-bound. Decode attention/GEMV: AI ≈ 1–2 FLOP/byte → hopelessly memory-bound; batching raises AI ≈ batch size.

Tile fit: smem_bytes = stages·(Bm·Bk + Bk·Bn)·dtype_bytes · accum regs/thread = Bm·Bn·4 / threads (limit 255 regs, 64K/SM)

Ring all-reduce: time = 2·(n−1)/n · bytes / link_BW ≈ 2·bytes/BW for large n (bandwidth-optimal, latency ignored)

Training: FLOPs = 6·P·tokens (fwd 2 + bwd 4) · MFU = achieved / peak · step time = 6·P·batch_tokens / (MFU·peak·chips)

Memory: inference ≈ P·dtype_bytes + KV · Adam mixed-precision training ≈ 16 B/param (2 wt + 2 grad + 12 optim/master) + activations

Drill

—

2:00

Press New problem to start. Work it on paper, out loud, then reveal.