Google TPU

The idea in one minute#

A GPU is thousands of small programmable cores that happen to have matrix units. A TPU is the opposite: a few very large matrix units (MXUs) with just enough programmable logic around them to keep them fed. Each MXU is a systolic array, a grid of multiply-accumulate cells through which data pulses in lock-step, so each value fetched from memory is reused hundreds of times without touching a register file or cache.

TPUs are also networked differently. Instead of switches, chips connect directly to their neighbours in a 3D torus, 64 chips to a rack-sized cube, and optical circuit switches join 144 cubes into a 9,216-chip pod that a compiler (XLA) programs as one machine. Google designs them for its own models and rents them on Google Cloud. Facts checked on 16 October 2026.

A picture#

How it really works#

The systolic array#

An MXU on recent TPUs is a 256 × 256 grid of cells. For a matrix multiply Y = X · W:

1. A 256 × 256 tile of weights is loaded into the cells, one weight per cell.
2. Activations enter from the left edge, skewed by one clock per row, and move one cell right per clock.
3. Each cell multiplies the activation passing through by its weight, adds the partial sum arriving from the cell above, and passes the new sum down.
4. Finished dot products leave the bottom edge, one column per clock.

Every weight is fetched once and used for every row of X; every activation is fetched once and used by 256 cells. No cell reads a register file or a cache. That is where TPUs get their energy efficiency, and it has a consequence the code below measures: the array takes 2 × 256 clocks to fill and drain, so it is efficient only when X has many rows, a large batch.

Around the MXUs#

• Vector unit: elementwise work (activations, normalisation, softmax), on 2D vector registers.
• Scalar unit: control flow and address generation; the TPU issues very long instruction words that drive all units at once.
• VMEM: software-managed on-chip memory; the compiler decides what lives there.
• SparseCores (since v4): small dataflow cores for embedding lookups, the scattered, memory-bound part of recommendation and ranking models that ruins a systolic array’s efficiency.
• HBM: 192 GB of HBM3e at 7.37 TB/s on Ironwood.
• ICI (inter-chip interconnect): direct links to neighbours, about 1.2 TB/s per chip.

Ironwood is Google’s first TPU with two compute dies in one package, each with its own TensorCore, SparseCores and HBM, joined by a die-to-die link.

Torus, cubes and optics#

A torus wires each chip to its neighbours in each dimension and wraps the edges around, so a 4 × 4 × 4 cube has no edge chips: each chip has exactly six links. There are no switch chips at all, which saves cost and power. The trade is hops: two far-apart chips are several hops away, whereas NVL72’s switches make every pair one hop.

Between cubes, Google uses optical circuit switches (OCS): MEMS mirrors that steer a beam of light from any input fibre to any output fibre. They do not read packets, so they add almost no latency or power, and re-pointing them changes the pod’s shape in seconds: a training job can get a torus of exactly the size it asked for, and a failed rack can be routed around. A pod is 144 cubes: 9,216 chips with 1.77 PB of HBM.

The generations#

The split is the news: Google concluded that training and inference want different chips, the same conclusion that led NVIDIA to add LPX racks for decode.

Software#

TPUs are programmed through XLA, a compiler that takes a whole computation graph and fuses, tiles and schedules it for the MXUs and the torus. JAX is the native front end; PyTorch runs through PyTorch/XLA; Pallas is the kernel language for hand-written kernels; vLLM serves on TPUs through its TPU backend. XLA’s sharding annotations let one program span a whole pod: the collectives over ICI are generated for you. Anthropic, Apple and others train on TPUs; Google trains Gemini on them.

Code#

A cycle-accurate model of a weight-stationary systolic array. It checks the answer against an ordinary matrix multiply, then measures how full a 256 × 256 array is at different batch sizes.

// systolic.go — a weight-stationary systolic array, clock by clock.
package main

import "fmt"

// run multiplies X (B×n) by W (n×n) on an n×n array; returns Y and the clocks taken.
func run(X, W [][]float64) ([][]float64, int) {
	B, n := len(X), len(W)
	xreg := make([][]float64, n) // activation held in each cell, moving right
	preg := make([][]float64, n) // partial sum held in each cell, moving down
	for i := range xreg {
		xreg[i], preg[i] = make([]float64, n), make([]float64, n)
	}
	Y := make([][]float64, B)
	for i := range Y {
		Y[i] = make([]float64, n)
	}
	clocks := B + 2*n - 1
	for t := 0; t < clocks; t++ {
		nx, np := make([][]float64, n), make([][]float64, n)
		for k := 0; k < n; k++ {
			nx[k], np[k] = make([]float64, n), make([]float64, n)
			for j := 0; j < n; j++ {
				var xin, pin float64
				if j == 0 { // left edge: row k gets X[b][k], skewed by k clocks
					if b := t - k; b >= 0 && b < B {
						xin = X[b][k]
					}
				} else {
					xin = xreg[k][j-1]
				}
				if k > 0 {
					pin = preg[k-1][j]
				}
				nx[k][j] = xin
				np[k][j] = pin + xin*W[k][j]
			}
		}
		xreg, preg = nx, np
		// the bottom row holds Y[b][j] at clock b + j + n - 1
		for j := 0; j < n; j++ {
			if b := t - j - (n - 1); b >= 0 && b < B {
				Y[b][j] = preg[n-1][j]
			}
		}
	}
	return Y, clocks
}

func main() {
	n, B := 4, 3
	X, W := make([][]float64, B), make([][]float64, n)
	for b := range X {
		X[b] = make([]float64, n)
		for k := range X[b] {
			X[b][k] = float64(b + k + 1)
		}
	}
	for k := range W {
		W[k] = make([]float64, n)
		for j := range W[k] {
			W[k][j] = float64((k*n+j)%5) - 2
		}
	}
	Y, clocks := run(X, W)
	ok := true
	for b := 0; b < B; b++ {
		for j := 0; j < n; j++ {
			want := 0.0
			for k := 0; k < n; k++ {
				want += X[b][k] * W[k][j]
			}
			ok = ok && want == Y[b][j]
		}
	}
	fmt.Printf("4x4 array, batch 3: correct=%v in %d clocks; Y[0] = %v\n\n", ok, clocks, Y[0])

	fmt.Println("256x256 MXU   batch   clocks   cells busy")
	for _, batch := range []int{1, 8, 64, 256, 1024, 8192} {
		c := batch + 2*256 - 1
		fmt.Printf("%14d %8d %10.1f%%\n", batch, c, 100*float64(batch)/float64(c))
	}
}

The small run proves the dataflow computes a real matrix product. The table is the important part: at batch 1 a 256 × 256 array is busy 0.2% of the time; it needs batches in the thousands of rows to approach full use. Real MXUs overlap one tile’s drain with the next tile’s fill, so steady-state use is better than this one-shot table, but the core fact holds: each weight loaded into the array is used once per row of the batch. Compilers therefore fold sequence length and batch into the rows (prefill and training have plenty), and decode at small batch is where TPUs, like GPUs, are limited by memory instead.

Remember this#

• A TPU is a few huge systolic matrix units (256 × 256 cells) plus vector, scalar and sparse helpers.
• Systolic arrays reuse every fetched value across a whole row or column of cells; they need large batches to fill.
• Chips connect directly in a 3D torus (64 per cube); optical circuit switches join 144 cubes into a 9,216-chip Ironwood pod.
• XLA compiles one program for the whole pod; JAX is the native language, Pallas the kernel language.
• The eighth generation splits into TPU 8t (training) and 8i (inference, 384 MB SRAM).

Try it#

1. Run systolic.go. Change the array to 8 × 8 and the batch to 1. How many clocks does one matrix-vector product take, and what fraction of cells did useful work?
2. In a 4 × 4 × 4 torus, what is the largest number of hops between two chips? Without the wrap-around links?
3. Why does a torus suit the all-reduce of data-parallel training better than the all-to-all of expert parallelism?

Check yourself#

1. In a weight-stationary array, what moves and what stays still?
2. Why does a systolic array waste most of its cells at batch 1?
3. What does an optical circuit switch do, and what does it not do?
4. Name one advantage and one cost of a torus compared with a switched NVLink domain.

Sources#

• Google Cloud documentation: TPU7x (Ironwood)
• Google: three things to know about Ironwood
• Google Ironwood at Hot Chips 2025 (ServeTheHome)
• Google unveils TPU 8t and 8i (DCD)
• TPU v4: an optically reconfigurable supercomputer (Jouppi et al., ISCA 2023)
• In-datacenter performance analysis of a TPU (Jouppi et al., 2017)