Google TPU — Tensor Processing Units

Two answers to the same question

Both machines exist to multiply large matrices quickly. They differ in what sits at the center. A GPU’s center is the SIMT execution model: thousands of threads grouped into warps, resident on streaming multiprocessors, with a hardware scheduler swapping between them to hide memory latency. Matrix hardware — tensor cores — is an instruction that warps issue, a powerful feature bolted onto a general substrate.

A TPU’s center is the matrix unit itself. The vector unit, the scalar control core, and the memory-movement machinery exist to keep that array fed. There is no warp scheduler arbitrating thousands of independent instruction streams, no occupancy target to hit, no register-pressure cliff to fall off. There is a very wide arithmetic pipe and a compiler whose job is to keep data marching into it. Nearly every practical difference below is a consequence of that inversion rather than an independent design decision.

Inside a systolic array — data flows, instructions do not

A systolic array is a fixed grid of small multiply-accumulate cells wired directly to their neighbors. Operands are pushed in at the edges and march through the grid; each cell multiplies, accumulates, and passes its result to the next cell. Weights typically stay resident in the cells while activations stream past, so a single value loaded once is reused across an entire row or column of the computation.

What has been deleted is the per-operation overhead. A GPU issuing matrix instructions still fetches and decodes them, reads and writes a register file, and moves operands through shared memory. In the array, the wiring is the dataflow: no instruction fetch per multiply, no register file traffic between adjacent cells, no address computation. The energy and area saved go straight back into more multipliers.

What the array is exceptionally good at

The array’s ideal workload is a large, dense matrix multiply with high arithmetic reuse and shapes that are known in advance. Big transformer projections, feed-forward layers, and large-batch training steps are exactly this. Every operand entering the grid is consumed by many cells, so the ratio of arithmetic to memory traffic is high and the machine sits in the compute-bound regime a roofline analysis wants it in.

Because the dataflow is fixed, utilization is also unusually predictable. There is no tail effect from a badly chosen block size, no occupancy regression when a kernel spills registers, no sensitivity to launch overhead. A shape that maps cleanly onto the array runs at close to the machine’s ceiling, consistently. For steady training work that predictability is worth real money.

What the array is bad at — small, ragged, and irregular work

The same rigidity is the failure mode. The array has a fixed square dimension, and operands that do not fill it are padded. A matrix multiply whose inner dimension is a fraction of the array width still occupies the whole grid; the unused cells burn cycles on zeros. Small matmuls, thin batches, and heads or experts that get sliced too finely all pay this tax, and it is not a few percent — it scales with how badly the shape misses the tile.

Anything that is not a dense matmul is worse. Gather and scatter, sorting, top-k, sparse or ragged structures, custom pointwise kernels with data-dependent addressing — these fall to the vector unit, where the TPU no longer has a structural advantage. On a GPU the same work is just another kernel. The question to ask of a workload is not “is it big?” but “what fraction of it is dense, well-shaped matmul?”

Compiler-centric by design — XLA ahead of time

The programming models diverge just as sharply. The TPU path runs through XLA, which compiles a whole traced program ahead of time against a fixed shape signature. It sees the entire graph, so it can fuse aggressively, choose layouts globally, schedule on-chip buffers, and overlap communication with compute as a planning problem rather than a runtime one. What you get back is essentially one large, optimized program.

CUDA works the other way. Kernels are launched at runtime; shapes can change between launches; a Python if can pick a different kernel on each iteration; you can drop into a hand-written kernel for one operator and leave the rest alone. Optimization is incremental and local, and you can always escape to a lower level.

Note that XLA itself is not TPU-exclusive — it compiles for GPUs too. The real contrast is whole-program compilation against static shapes versus flexible dispatch at runtime.

Variable sequence lengths and the recompilation tax

This is where the abstract difference becomes a concrete engineering problem. XLA specializes compiled programs by shape signature, so a genuinely new shape means a new compilation. Serving a language model with arbitrary prompt lengths would, taken naively, trigger a compile on nearly every request — and compilation is far more expensive than the step it is compiling.

The standard fix is bucketing: quantize sequence lengths to a small set of buckets, pad each request up to its bucket, and warm the cache by compiling all buckets ahead of time. It works, but note what you have bought it with — padding waste on every request that lands mid-bucket, plus a warm-up phase to manage.

Dynamic control flow is not the problem; XLA has loop and conditional constructs. Dynamic shapes are. Workloads with wide, unpredictable shape distributions fit the flexible-runtime model more comfortably.

Memory: compiler-scheduled buffers vs caches and occupancy

Both machines hang HBM off the package and both have fast on-chip memory. The difference is who decides what lives there. A GPU gives you a hardware cache hierarchy plus a programmer-managed scratchpad, and it hides the latency it cannot avoid by keeping many warps resident and switching between them — latency hiding through concurrency.

A TPU leans on the compiler instead. XLA plans on-chip buffer residency and the prefetches that fill them as part of compiling the program, so latency is hidden through scheduling. When the schedule is good this is excellent: no cache thrash, no occupancy trade-off, no surprises. When the access pattern is data-dependent and cannot be planned statically, the compiler has less to work with than a hardware cache would have had at runtime — the same trade in a different costume.

Pod scale — a torus changes what collectives cost

At cluster scale, GPU deployments are usually tiered: a fast scale-up domain of a handful of GPUs joined by NVLink, then a slower scale-out network beyond it. Collectives are shaped around that cliff, and the topology of the outer network is a datacenter design choice that varies per site.

TPU pods take a different approach: chips are joined by dedicated inter-chip interconnect (ICI) in a regular mesh with wraparound links — a torus — that is part of the machine rather than assembled per deployment. The wraparound matters. Every chip has the same number of neighbors, there is no edge, and rings can be embedded in the topology in multiple dimensions, which is precisely what bandwidth-optimal all-reduce and all-gather want. Collective cost becomes a predictable function of slice geometry, and the compiler can plan around it because it knows the topology at compile time.

The cost is rigidity: your parallelism layout must suit the mesh you were allocated.

The portability cost of committing

Choosing a TPU is choosing a software stack, and that bill arrives later than the hardware decision. JAX is the most natural fit and PyTorch runs through an XLA backend, so mainstream model code is usually portable. The friction is at the edges: any custom CUDA kernel, any dependency on a CUDA-only library, any profiling workflow built on NVIDIA tooling, and any operator your framework only implements as a hand-written GPU kernel must be replaced rather than recompiled.

There is also an ecosystem asymmetry. New techniques tend to ship as CUDA kernels first, so a TPU shop is sometimes waiting for a port — and TPUs are available from one cloud, which narrows your options for capacity and pricing leverage. None of this is disqualifying, but it is a recurring engineering cost that belongs in the comparison.

The honest verdict — match the machine to the shape

Neither machine is generally better; they are tuned for different distributions of work. The TPU is strongest when the workload is large, dense, homogeneous, and shape-stable: pretraining a big model, long steady training runs, or serving at high volume where you can bucket shapes and amortize compilation. Those workloads reward static planning and a tightly coupled fabric, and they are exactly what the architecture was built for.

The GPU is strongest when the workload is varied, irregular, or moving: research where models change weekly, inference with unpredictable shapes, pipelines full of custom kernels and non-matmul operators, or anything that needs to run in more than one place. Flexibility costs peak efficiency, and you pay it in exchange for not having to know the future. Assess the workload honestly, and let its shape pick the machine.