Follow data through memory
GPU execution involves threads/warps, SM resources, registers, shared memory and HBM. Tensor Cores accelerate suitable matrix operations rather than replacing all general computation. TPUs combine matrix, vector/scalar resources and on-chip memory under compiler-driven execution. Generations differ; “GPUs are flexible, TPUs only multiply matrices” is not an adequate model.
Data movement is often expensive. Writing every intermediate to HBM and rereading it in another kernel can dominate modest FLOP counts. Tiling reuses loaded data in faster storage. Fusion reduces intermediate traffic but may increase register pressure or reduce parallelism.
Two Roofline bounds
Arithmetic intensity is . A simplified throughput bound is , with time lower bound . This is an idealized bound, not a latency predictor. Launches, synchronization, imbalance and communication add costs.
flops = 2 * 1024**3
bytes_moved = 3 * 1024**2 * 2 # Teaching assumption: each matrix transferred once.
intensity = flops / bytes_moved
print('FLOP/byte:', intensity)
Real GEMMs have cache reuse, layouts and dtype effects; this three-matrix estimate is not universally accurate traffic accounting.
Establish shapes before writing a kernel
Triton's tiled matmul tutorial introduces program instances and blocked, masked memory operations. JAX/Pallas offers lower-level kernel programming for supported TPU paths, with backend-specific constraints. Optimize in order: numerical correctness, safe tails, correct strides/layouts, then block sizes, warps and pipeline stages.
Compare random shapes and nondivisible dimensions against a reference before benchmarking. A kernel fast for 1024×1024 may lose on thin decode matrices. State whether a multiply-add counts as two FLOPs, whether transfers are timed and which precision is used.
Check your understanding
Question: A kernel reaches only 10% of peak compute. Is it necessarily poorly written?
Answer
No. It may be near the bandwidth bound at low arithmetic intensity. Compare the relevant Roofline first, then inspect memory efficiency, occupancy and synchronization. High occupancy is not the final objective either.Primary sources and further reading
Sources checked on 2026-09-10. Teaching examples are not production benchmarks; confirm APIs and model support against the linked version.
Inspect tensors, run experiments, and check the stated environment and execution status in the first cell.
Open in Colab ↗Download notebook