GPU vs TPU: flexible, or purpose-built?

A GPU (Nvidia, dominant via the CUDA ecosystem) is a general-purpose massively-parallel processor — originally built for graphics, now the default for nearly all AI workloads. A TPU (Google's Tensor Processing Unit) is an ASIC — application-specific silicon — purpose-built around a systolic array optimized specifically for matrix multiplication, available only on Google Cloud.

A single training run happens on one or the other, not both. But many teams DO use both across their overall stack — train on TPU pods for cost/scale where the workload fits, while keeping GPU infrastructure for portability, edge deployment, or the parts of the pipeline (research iteration, custom kernels) that don't fit TPU's constraints well. Frameworks like JAX increasingly aim for code portability across both, turning this into more of an infra/cost decision made per workload than a full rewrite each time.

The architectural reason a TPU is efficient for large matrix multiplies is the systolic array: a grid of multiply-accumulate units where data flows through the grid and each loaded weight gets reused across many computations before being evicted, rather than being fetched from memory again for each use. Since matrix multiplication IS mostly repeated multiply-accumulate operations over the same operands, this design minimizes the memory-bandwidth pressure for exactly the operation that dominates transformer training and inference — and memory bandwidth, not raw compute, is very often the actual bottleneck in large-scale training.

A GPU's tensor cores also accelerate matmul well, but the GPU as a whole is a more general SIMT (single-instruction-multiple-thread) processor that wasn't built around ONLY that operation — which is exactly why it remains better for the long tail of non-matmul, custom, research-oriented workloads: novel attention variants, custom loss functions, sparse operations, anything that needs a hand-written CUDA kernel rather than compiling cleanly through XLA. The TPU's specialization is a real efficiency win precisely where the workload matches it, and a real limitation the moment it doesn't.

A lab designing a new transformer architecture with experimental custom attention kernels wants GPUs — during this phase, the ability to write and iterate on custom CUDA kernels while the architecture itself is still changing is worth more than raw matmul efficiency, because the workload isn't yet a fixed, well-understood shape a systolic array is optimized for.

Once that architecture is settled and the lab is running repeated, large, well-understood pretraining passes of the now-finalized architecture — a workload that IS mostly large matrix multiplies, at Google-infrastructure scale — TPU pods' cost/performance for that specific, now-stable, matmul-heavy workload becomes genuinely attractive, if the lab is already on GCP. The right accelerator tracks how settled and matmul-shaped the workload has become, not a fixed rule about which is "better."