Reading it like an operator
Six habits turn a wall of events into a diagnosis. Name the envelope first: the jit_<fn> event is wall time per call, and every other number is judged against it. Subtract the compute sum from the envelope: the difference is your unhidden transfer time, named exactly. League-table the ops: the top five carry the step, and the rest is noise until proven otherwise.
Then the subtler three. Check counts against iterations: an op at ×40 for 20 iterations is a pair, usually an async start and its span, and pairs are how you spot overlap. Treat custom calls as first-class suspects: they are opaque to every cost model, and the biggest finding this course made on hardware was a custom call nobody wrote. Close with the rooflineThe floor model: latency is at least the larger of FLOPs over peak compute and bytes over bandwidth. Predict first, measure second.taught in /l/tpu →: convert the big ops' times to bytes and flops through measured rates, place them against the ridgeThe FLOP-per-byte ratio where an op flips from memory-bound to compute-bound: about 240 on v5e, about 575 on v6e.taught in /l/tpu →, and say memory-bound or compute-bound with a number attached. At that point the profile is not a picture anymore; it is an argument.
Exercises
(4096, 512) bf16 row softmax on each rooflineThe floor model: latency is at least the larger of FLOPs over peak compute and bytes over bandwidth. Predict first, measure second.taught in /l/tpu →: memory-bound or compute-bound, chip by chip? (70, 300) array. The museum's lattice exhibit holds the answer's shape. jit envelope against the sum of its compute ops. State how many microseconds of transfer failed to hide, if any. copy-start.1 shows 40 events for 20 iterations, and where its 89.6 µs went. The envelope arithmetic is the proof. Check yourself
01 What is the envelope, and why does it come first?
The jit event for the function: wall time per call. Every other number is judged against it, so naming it first anchors the reading.
02 The compute events sum to far less than the envelope. What is the habit?
Find what fills the gap: the difference is unhidden transfer or overhead, and chasing it beats tuning a kernel that already looks fast.
Readings
- JAX profiling ↗ the tool the habits run on
- Scaling book · All about TPUs ↗ the constants a diagnosis quotes