The TPU chip

A CPU spends most of its area deciding what to do next: branch predictors, reorder buffers, cache hierarchies that guess your access pattern. A TPU spends almost none. The designers looked at what training and serving actually run and found matrix multiplication, over and over, at every scale. So the chip commits. One huge matrix unit does that work, a vector unit handles the arithmetic around it, and everything else on the die exists to keep those two fed.

The commitment shows up in what is missing. There is no branch predictor to speak of, no out-of-order engine, no hardware cache deciding for you what stays close to the compute. When this course later has you choreograph data by hand with a BlockSpec, this lesson is the reason: the hardware chose not to guess, which means someone has to say. On a TPU that someone is the compiler, and in Pallas it is you.

Picture the multiply as a piece of machinery rather than a loop. Weights get parked in a square grid of cells, one value per cell, and stay there. Activations enter from one edge, one diagonal per cycle, and every cell they pass does a single multiply-accumulate: take the incoming value, multiply by the parked weight, add to the partial sum flowing through, pass both along. Results drain out the far edge. Nothing fetches an instruction per operation and nothing asks a cache for its operand; the schedule is the geometry. An array built to pulse data through itself this way is called a systolic array, and the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → is one: 128x128 cells on most generations, 256x256 on v6e.

The geometry explains the costs you will keep meeting. The array has a fill phase while the first diagonals march in, a steady state where all cells work every cycle, and a drain phase at the end. Small matmuls live mostly in fill and drain, which is why they waste the unit; the arithmetic below a (8, 128)-shaped tile can't even occupy one edge. Big matmuls amortize the ramp and run the array flat out. When stage 0 had you predict that a skinny matmul lands memory-bound and a square one compute-bound, this grid is the machinery behind the prediction.

The schedule is the geometry: data moves, and the movement is the computation.

Softmax needs an exponential, layernorm needs a square root, and neither is a matrix product. That work goes to 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 vector unit that applies the same operation across wide registers of data. It is the second citizen of the chip by area and the first by variety: adds, multiplies, exponentials, comparisons, casts, everything elementwise a kernel does between matmuls. On v5p each core carries 64 32-bit vector registers to stage that work.

The split matters for how you read a kernel's cost. A fused attention kernel is not one workload; it is matmuls on the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → with exponentials and rescaling on 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 → stitched between them. If the elementwise work is thin, it hides behind the matmuls. If it is thick, the VPU becomes the bottleneck while the MXU idles, and no 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 → over FLOPs alone will tell you. The habit to build now: when you meet a kernel, sort its lines into MXU lines and VPU lines before you estimate anything.

Between the compute units and 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 → sits 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 →, roughly 128 MiB of on-chip memory on v5e. Calling it a cache would miss the design. A cache decides for you what stays close; VMEM holds exactly what software staged into it, nothing more. Every block a kernel touches was placed there by a DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas → that something explicitly issued, and the compute units read only from there. Kernel engineering on TPU is mostly the choreography of that staging: which block arrives when, and whether the next transfer overlaps the current compute.

Two smaller memories complete the picture. SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → holds scalars: loop bounds, block indices, flags, the values a kernel branches on. And a scalar core runs alongside the vector units executing the control flow your kernel compiles to, issuing the DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas → descriptors that move blocks from 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 → into VMEM. The guide sections below walk these two in detail, with the diagram most block diagrams omit; here it is enough to hold the division: 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 → compute, the scalar core decides and fetches, 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 → is where the two worlds meet.

One more unit sits on the die and no kernel in this course touches it. SparseCore exists for embedding lookups, the wide scattered reads that recommendation models hammer and that a dense 128x128 grid is exactly wrong for. It earns its area on those workloads and stays dark on ours. It belongs in your picture of the machine anyway, as the exception that proves the bet: when a workload class mattered enough and fit the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → badly enough, it got its own silicon rather than bending the array.