Two machines, one job

Term for term, the two machines line up better than the marketing suggests. An SM plays the part of a TensorCore, a warp32 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 → scheduler the part of the VPUThe vector unit for elementwise work, organized as (8, 128) lanes; the origin of the tiling lattice every layer above obeys.taught in /l/tpu →, a CUDA core the part of a VPU ALU, SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → the part of 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 →, and a GPU Tensor Core the part of the MXU. HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu → is HBM on both. Once you can translate in both directions, most GPU documentation stops being foreign.

The counts are where the family resemblance ends. A TPU v5p has 2 TensorCores with 8 MXUs between them; an H100 has 132 SMs carrying 528 Tensor Cores. Each TPU TensorCore has one big VPUThe vector unit for elementwise work, organized as (8, 128) lanes; the origin of the tiling lattice every layer above obeys.taught in /l/tpu → built from 4 independently programmable 8x128 units, 4096 ALUs in total, while the H100 has 528 independent 32-wide SIMD units, about 16k ALUs. Counting individual lanes, an H100 has 132 * 4 * 32 = 16,896 CUDA cores against a v5p's 2 * 4 * 8 * 128 = 8192 ALUs, running at roughly the same frequency.

The memory asymmetry runs the other way, and it is bigger than the table makes it look. An H100's on-chip fast memory is 32 MB of SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → against a TPU's 128 MB of VMEM. Bandwidth widens the gap again: TPU 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 → runs at around 40 TB/s, while the closest GPU equivalent by size, the L2, has been measured near 5.5 TB/s and is not under your control. That is the single fact behind most claims that TPUs are better at inference, since weights that live in VMEM load fast enough to change the roofline.

Both chips spend a transistor budget; they spend it on different problems. The GPU buys many small independent units plus the hardware to decide, at runtime, which warp32 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 → runs next on which subpartition. That machine is hardware-scheduled. The TPU buys a few very large units with a single thread of control and only VPUThe vector unit for elementwise work, organized as (8, 128) lanes; the origin of the tiling lattice every layer above obeys.taught in /l/tpu →-wide vector instructions, which leaves the compiler to place every load and every matmul in a schedule ahead of time. That machine is compiler-scheduled.

The GPU schedules at runtime. The TPU schedules at compile time.

The TPU side is cheaper to build and simpler to reason about, and it moves the whole burden into the compiler, which must pipeline every memory load against MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → and VPUThe vector unit for elementwise work, organized as (8, 128) lanes; the origin of the tiling lattice every layer above obeys.taught in /l/tpu → work or the machine stalls with nothing to hide the wait. The GPU side asks much less of its compiler. You can launch dozens of unrelated kernels and each one lands on an independent SM, and they will run.

They may also run badly. Kernels that thrash the shared L2 or fail to coalesce their loads are slow for reasons the source does not show, and because the hardware owns so much of the runtime it is hard to see which reason applies. TPUs more often reach close to 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 → with less work, precisely because less was left to runtime in the first place.

On a GPU, the tuning surface is 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 → and locality. You watch registers per thread against the 64 resident 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 → an SM can hold, knowing that 256 registers per thread leaves room for only 8. You size shared-memory tiles against 256 kB of SMEM. You order your accesses so that threads in a warp touch adjacent addresses, and you think about whether SMs are cooperating or fighting over the shared L2.

On a TPU the surface is the schedule you write down. The BlockSpec index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → decides which block each grid step sees, the grid order decides what streams in behind the current step, the (8, 128) lattice decides whether your shapes compile well at all, and the 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 → budget decides how large a block you may ask for. Stage 0 and stage 1 of the path are almost entirely this vocabulary.

Underneath, it is the same question on both machines. An H100's bf16 matmul 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 → is 9.9e14 / 3.4e12, about 290 FLOPs per byte, the same shape of number as the v5e's 240 from stage 0. What differs is who is responsible for staying above it.

Attention is the clearest case of the same idea landing twice. The intermediate score matrix is the thing you cannot afford to write out to HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu → and read back, on either machine, which is why the algorithm gets restructured to keep it on chip. The original paper frames the whole method as IO-awareness: tiling to reduce reads and writes between HBM and the on-chip SRAM, rather than reducing FLOPs.

On a GPU that on-chip memory is one SM's SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu →, so the kernel is written as tiles staged into shared memory with threads inside a warp32 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 → cooperating on each tile, and correctness of the online softmax expressed at thread granularity. On a TPU the same algebra becomes a grid over KV blocks: a BlockSpec stages each block into 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 →, the running maximum, running sum, and accumulator are carried across grid steps, and the pipeline hides the next block's transfer under the current block's math.

The algebra is identical in both cases. What changes is who does the staging, which is the same split you have been reading about since the top of this lesson.

You can now read a GPU spec sheet without translating twice: SM against TensorCore, Tensor Core against MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu →, SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → against 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 →, NVLink domain against torus axis, and a measured 370 GB/s against a claimed 450. That is enough hardware to hold both machines in your head while reading anyone's kernel.

It is also where the two chips stop being the whole story. One question is left at this scale, which is what happens when the thing you shard across is a rack rather than a chip, and the last lesson of the stage answers it before the path turns to authorship.