GPU vs TPU for AI

A graphics processing unit (GPU) is a chip with many small arithmetic cores that was first built to draw pixels in parallel and was later opened up for general computing. A tensor processing unit (TPU) is an application-specific integrated circuit (ASIC), meaning a chip designed for one class of work, that Google built to accelerate neural networks. Both speed up the same central operation of deep learning, multiplying large grids of numbers called matrices, but they make different trade-offs. GPUs keep broad programmability and are sold to anyone; TPUs trade some flexibility for dense matrix throughput and energy efficiency, and outside Google they are rented through Google Cloud rather than bought.

A modern NVIDIA data-center GPU is a collection of streaming multiprocessors (SMs), each of which runs many threads at once. NVIDIA's Hopper architecture write-up says the H100 SXM5 has 132 SMs with four Tensor Cores each, 528 in total. A Tensor Core is a block that performs a small matrix multiply-accumulate in one step at reduced precision. Hopper adds a Transformer Engine that chooses, layer by layer, between 8-bit floating point (FP8) and 16-bit math. The chip has 80 billion transistors on TSMC's 4N process. Programmers reach all of this through NVIDIA's CUDA programming model, and frameworks such as PyTorch call CUDA libraries for them, which is why most published research code runs on GPUs with no changes.

A TPU puts most of its silicon into matrix units (MXUs) arranged as a systolic array, a grid of multiply-accumulate cells that pass numbers to their neighbors on every clock tick so data flows through the grid instead of going back to memory after each step. Google Cloud's architecture page says MXUs are 128 by 128 cells on TPU v5p and earlier, about 16,000 multiply-accumulates per cycle, and 256 by 256 on TPU v6e and the seventh generation. Each TensorCore on the chip also has a vector unit for operations such as activations and softmax and a scalar unit for control flow. Programs do not target the chip directly: the XLA compiler turns the linear-algebra parts of a JAX, TensorFlow or PyTorch graph into TPU machine code, and the rest runs on the host computer. Chips are wired into slices and pods by a dedicated inter-chip interconnect (ICI).

The history explains the split. The first TPU paper, presented at ISCA 2017, describes an inference-only chip in service since 2015 with a 65,536-cell 8-bit matrix unit, 92 trillion operations per second at peak and 28 MiB of on-chip memory. Its authors report that it ran about 15 to 30 times faster than a contemporary NVIDIA K80 GPU or Intel Haswell CPU on Google's production workloads, with 30 to 80 times better performance per watt. Later generations added training. Google's bfloat16 article explains that TPU v2 and v3 multiply in bfloat16, a 16-bit format with the same 8-bit exponent as 32-bit floats, and accumulate in 32-bit precision. The 2023 TPU v4 paper describes a 4,096-chip supercomputer linked by optical circuit switches and reports it 1.2 to 1.7 times faster than an NVIDIA A100 system of similar size while using 1.3 to 1.9 times less power. In April 2025 Google announced Ironwood, quoting 4,614 peak FP8 teraflops and 7.37 TB/s of memory bandwidth per chip and 42.5 exaflops for a 9,216-chip pod.

Comparing headline numbers needs care. NVIDIA lists the H100 SXM at 3,958 FP8 teraflops and 1,979 BF16 teraflops, but those figures assume structured sparsity, a pattern in which half the weights are zero; the Hopper write-up confirms sparsity doubles the dense rate, so dense FP8 is about half the listed figure. Google's pod figures are theoretical peaks, and vendor benchmarks are chosen by the vendor. The TPU v4 versus A100 comparison, for example, comes from Google's own paper. Real speed depends on how well a model keeps the matrix units busy, how much data moves between chips and how mature the compiler is for that model. Independent tests on the actual workload are the only reliable guide.

Google's own guidance is a useful decision rule. Its introduction to Cloud TPU recommends TPUs for models dominated by matrix math, models without custom operations inside the training loop, models that train for weeks or months, large models with large batch sizes, and models with very large embeddings such as ranking and recommendation. It recommends GPUs for models with many custom PyTorch or JAX operations, models that use operations TPUs do not support, and medium to large models. It recommends CPUs for quick prototypes, small models and jobs limited by input or network speed. It also lists what TPUs handle poorly: code with frequent branching, many element-wise operations, high-precision arithmetic, and custom operations in the main loop.

Each option also has failure modes. On TPUs, shapes that do not tile neatly onto the 128 or 256 grid waste capacity, and the same Google page notes memory performs best with dimensions that are multiples of 8. Every program must pass through the XLA compiler, so debugging relies on compiler tools rather than the CUDA profilers most practitioners already know. On GPUs, peak throughput also depends on using Tensor Core-friendly precisions and shapes, and multi-GPU training relies on fast links such as NVLink, which tops out at 900 GB/s per H100. Lock-in cuts both ways: CUDA code does not run on TPUs, and TPU access is limited to Google's cloud and Google's own products.

In practice the requirements differ more than the arithmetic. A GPU can be bought, rented from many clouds or installed in a workstation, and an H100 SXM draws up to 700 W. A TPU is reached through Google Compute Engine, Google Kubernetes Engine or Google's managed AI platform, and large jobs are sized in slices of chips; Google says a full Ironwood pod draws nearly 10 megawatts. For a small team, the deciding questions are usually which framework the code already uses, whether it relies on custom kernels, and which provider offers the needed capacity at an acceptable price on the day.