The GPU, for contrast
Set the TPU aside for a moment and look at the machine most kernel authors learn on first: an H100-class GPU, by the numbers on its public spec sheet. These are approximate, since binning and clocks vary card to card: about 9.9e14 dense bf16 FLOP/s, about 3.35e12 bytes per second of HBM3 bandwidth. Divide one by the other and you get a 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 → point near 295 flops per byte, the same 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 → arithmetic that gave the TPU chapters their four numbers in the first place.
Line 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 → points up and the contrast turns concrete instead of vague. v5e sits near 240, v6e near 575, an H100-class part near 295. A kernel that is compute bound on a v6e at a given arithmetic intensity can be bandwidth bound on an H100-class GPU, because the GPU's ridge sits lower: its FLOP/s grew less relative to its bandwidth than v6e's did. The point is not which chip wins. It is that the same compute-bound-or-bandwidth-bound question gets a different answer depending which chip you asked it on.
The rest of the spec sheet describes a differently shaped machine, not just a faster one. An H100-class part has 132 streaming multiprocessorsOne of a GPU’s many cores; each keeps many warps resident and switches among them to hide memory latency.taught in /l/tpu →, each running threads in groups of 32 called warps. Shared memory, the GPU's equivalent of a software-managed scratchpad, is configurable up to about 228 KB per SM: small next to VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu →'s roughly 128 MiB, but replicated 132 times instead of held once. Above that sits an L2 cache around 50 MB, a level of the memory hierarchy the TPU side of this course never had to name. Same job, keep compute fed from something close, solved with many small pools instead of one large one.
Two ways to hide latency
A TPU chip runs one big core, not thousands of small ones, and its grid executes sequentially: one step's compute has to be scheduled before its results are needed, while the next step's data is already moving in behind it. That is the software pipeline earlier chapters described. Latency gets hidden by overlapping DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas → transfers with compute inside a scratchpad you can see and size, VMEM. Because the whole schedule reduces to a handful of constants, FLOP/s, bandwidth, 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 → point, tile shape, you can reason about whether a kernel is compute bound or bandwidth bound before you ever run it.
A GPU hides latency a completely different way. Instead of one sequential pipeline, it keeps thousands of warps32 GPU threads scheduled as one unit; the GPU hides latency by switching among resident warps rather than by pipelining a scratchpad.taught in /l/tpu → resident on its streaming multiprocessorsOne of a GPU’s many cores; each keeps many warps resident and switches among them to hide memory latency.taught in /l/tpu → at once. When one warp stalls waiting on memory, the hardware scheduler switches to another warp that is ready, with no software pipeline required to make that happen. This is why occupancyHow many warps stay resident on a GPU core; the currency of GPU latency hiding, the way pipeline depth is on TPU.taught in /l/tpu →, how many warps stay resident, and register pressureHow many registers each GPU thread claims; it caps how many warps fit on a core, and with them the latency-hiding budget.taught in /l/tpu →, how many registers each thread claims and therefore how many warps fit, are the currencies that matter on a GPU the way tile shape and pipeline depth matter on a TPU.
Both machines still leave the same gap above them, worth naming plainly: neither scheduler rewrites your algorithm. A TPU pipelining DMAsAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas → and a GPU switching warps32 GPU threads scheduled as one unit; the GPU hides latency by switching among resident warps rather than by pipelining a scratchpad.taught in /l/tpu → will both faithfully execute a naive attention kernel that materializes the full sequence-length by sequence-length score matrix, and both will pay for it in bandwidth. That is why flash attention, the algorithm that never materializes that matrix, exists as hand-written source on both machines rather than as something either compiler produces on its own. Hiding latency well and choosing what to compute are not the same job.
Check yourself
01 Both machines wait on HBM. What hides the wait on each?
The GPU switches among resident warps at runtime, so another thread computes while one waits. The TPU runs one sequential core, so the compiler and the pipeline overlap transfers against compute ahead of time.
02 Why does the TPU approach demand more of the compiler?
There is no runtime scheduler to cover a miss: any wait not planned away at compile time is a stall the machine simply takes.
Readings
- Scaling book · How to think about GPUs ↗ the thread-switching half of the contrast
- Scaling book · All about TPUs ↗ the pipelined half