notebook № 036 plates · 22 min to read, longer to play
What a TPU is, and how it works
On a heart that beats in numbers
A TPU is a chip built around one equation. Nearly all the arithmetic in a neural network is matrix multiplication, so Google built a processor that does almost nothing else, and does it with a grid of multipliers that pass numbers along like a heartbeat. This notebook takes that grid apart, then follows it out to a whole pod.
plate I · the workload
A chip for one equation
A CPU is built to run anything. A GPU is built to run thousands of copies of anything. A TPU starts from a narrower and more ruthless question: what does a neural network actually spend its time on?
The answer is one line of algebra. A layer takes a batch of input vectors x, multiplies them by a weight matrix W, adds a bias, and squashes the result through an activation: y = f(Wx + b). The multiplication is a multiply-add for every input, every output and every example in the batch. The bias and the activation touch each output once.
Drag the sliders and watch the bill. Unless the layer is tiny, the matrix multiply is nearly all of the arithmetic: around 98% even with an elaborate activation like GELU, and 99.8% with a plain ReLU. That lopsided bar is the whole business case for a chip that is, mostly, a matrix multiplier with some plumbing attached.
Set inputs and outputs to 4,096 and the batch to 512, then work out how long you would sit there with a calculator.
plate II · the heartbeat
The systolic array
The obvious way to multiply matrices fast is to build many multipliers and feed each one from memory. The problem, as the CPU notebook showed, is that memory is slow and far away. Thousands of multipliers each fetching their own operands would choke on the wires long before they ran out of arithmetic.
A systolic array fixes this by letting the multipliers feed each other. Each cell holds one weight, preloaded and stationary. Inputs enter from the left edge and march one cell to the right every clock tick. Partial sums enter from the top and fall one cell down. At every tick each cell does exactly one thing: multiply the input passing through by its weight, add it to the sum arriving from above, and pass both along. The name comes from the heart: data pulses through in waves.
Below is the real thing, small enough to read. Each number you see is genuinely computed by the cell it sits in.
Two things to notice. First, the skew: each row of inputs starts one tick later than the row above, so that an input and the partial sum it belongs with arrive at the same cell at the same moment. The whole computation is a diagonal wavefront sweeping across the grid. Second, the fill and drain: at the start and the end, most cells sit idle while the wave enters and leaves. With a batch of one row, the array is mostly waiting. With a long stream of rows, the edges become a rounding error and every cell does useful work on every tick.
The payoff is in the wiring. Each operand is read from memory once and then reused by a whole row or column of cells over short, local wires. The first TPU packed 65,536 such cells into a 256×256 grid of 8-bit multipliers; later generations use 128×128 grids of bfloat16 multipliers, several per chip, and the newest have gone back to 256×256.
Set the array to 8×8 and the batch to 1, then to 8. Compare the utilisation. That ratio is why TPUs like big batches.
plate III · precision
Good-enough numbers
A multiplier’s size grows roughly with the square of the number of bits in its mantissa. Halve the bits and you can fit far more multipliers in the same silicon, move half as many bytes, and keep twice as many numbers on chip. So the question is: how few bits can a neural network get away with?
A floating-point number is three fields: a sign, an exponent that sets the scale, and a mantissa that sets the precision. The standard 16-bit float splits its bits 5 for the exponent and 10 for the mantissa: decent precision, but anything above 65,504 overflows to infinity, and gradients in training happily wander outside that range. bfloat16, the “brain” float, makes the opposite trade: it is simply float32 with the bottom 16 bits chopped off. Same 8-bit exponent, same range, only about two decimal digits of precision. Neural networks turn out to shrug at the lost precision and to care a great deal about the range.
The trick that makes it safe is that the matrix unit multiplies in bfloat16 but accumulates in float32, so long sums don’t drown in rounding error. The first TPU went further still and ran inference on 8-bit integers, which works once training is done and the range of each tensor is known in advance.
Store 257. Then 70,000. Then 1e-7. Each one breaks a different format in a different way.
plate IV · bandwidth
Feeding the beast
A matrix unit that can do hundreds of trillions of operations a second is only as fast as the memory that feeds it. TPUs sit next to stacks of HBM, high-bandwidth memory, which moves terabytes a second. That sounds like plenty until you divide the two numbers: the chip can do on the order of a hundred operations in the time it takes to fetch one byte.
So the real question for any matrix multiply is its arithmetic intensity: how many operations it does per byte it reads. Multiplying a batch of B vectors by a big weight matrix reads the weights once and reuses them B times. With a batch of one, every weight is fetched for a single multiply-add, and the matrix unit spends nearly all its time waiting. Grow the batch and the same weights pay for more and more work, until memory finally keeps up.
This is why serving a model to one user at a time is expensive, and why inference systems batch requests together so eagerly.
Find the smallest batch that keeps the MXU fully busy. Now switch to the small layer. Did the answer change much?
plate V · the compiler
Compiled, not scheduled
A CPU discovers what to do next at run time, one instruction at a time, with branch predictors and out-of-order engines guessing ahead. A TPU has almost none of that machinery. Instead, a compiler called XLA sees the whole computation graph before it runs, with every tensor shape known, and plans the chip’s work ahead of time: which values live in on-chip memory, when each transfer starts, which operations can be merged.
That last part, fusion, is where a lot of the speed hides. Run each operation as its own kernel and every intermediate result makes a round trip: written out to HBM, read straight back in by the next step. Fuse them and the intermediates never leave the chip.
Folding the elementwise ops into the matmul saves about forty per cent of the traffic; fusing the whole chain saves more than half, without changing a single result. This is the deeper bargain of the TPU: give up flexibility at run time, and in exchange a compiler that knows everything in advance can arrange the data so the hardware never waits. It is the same bet an ahead-of-time compiler makes for ordinary programs, which is the subject of its own notebook.
plate VI · scale
A pod is one computer
No single chip holds a large model or trains it in reasonable time, so TPUs are built to be used in hundreds and thousands. Each chip has dedicated inter-chip links, the ICI, wired straight to its neighbours in a grid whose edges wrap around: a torus. There is no network switch in the data path for most traffic; a message hops from chip to chip. The wrap-around links matter more than they look: they roughly halve the longest distance across the grid.
The traffic that dominates training is the all-reduce: after each step, every chip has computed gradients on its slice of the data, and every chip needs the sum. Sending everything to one chip would melt that chip’s links. The ring algorithm instead cuts each gradient into chunks and passes them around, adding as it goes (reduce-scatter), then passes the finished sums around again (all-gather). Every link is busy all the time, and each chip sends a little under twice its gradient, however many chips are in the ring.
Real pods do this in two or three dimensions at once. A TPU v4 pod joins 4,096 chips in a 3D torus, with optical switches that can rewire the topology. At that scale the pod stops looking like a cluster and starts behaving like one enormous, very particular computer. How it compares with a rack of GPUs is the subject of the comparison notebook.
On an 8×8 torus, find the two chips furthest apart. Then count the hops without the wrap-around links.
Questions people ask
- What is the difference between a TPU and a GPU?
A GPU is a general parallel processor: thousands of threads that can run any kernel, with tensor cores added for matrix maths. A TPU is built around large matrix units and leaves the planning to a compiler and the host CPU. It trades flexibility for efficiency on the matrix multiplications that dominate neural networks.
- What is a systolic array?
A grid of simple multiply-add cells, each wired only to its neighbours. Numbers enter at the edges and pulse through the grid one step per clock tick, like blood through a heart, which gives the design its name. Every cell does useful work on every tick without touching main memory.
- What is bfloat16?
A 16-bit floating-point format with the same 8-bit exponent as a 32-bit float but only 7 bits of mantissa. It keeps the wide range neural networks need while halving memory and bandwidth, and the lost precision rarely matters for training.
Related material
elsewhere in the notebooks
further reading
- In-Datacenter Performance Analysis of a Tensor Processing Unit Jouppi et al., 2017 The paper that introduced the first TPU: a 256×256 grid of 8-bit multipliers for inference, and the reasoning behind every design choice.
- How to Think About TPUs How to Scale Your Model The best modern explanation of the chip, its memory hierarchy and its interconnect, from people who scale models on it for a living.
- TPU system architecture Google Cloud Generations, chips, slices and pods, from the vendor’s own documentation.
- XLA OpenXLA The compiler behind plate V: it takes a whole computation graph and turns it into fused kernels ahead of time.
- Systolic array Wikipedia H. T. Kung’s idea from the late 1970s, long before anyone needed it for neural networks.