Tensor processing units: how a matrix-focused chip accelerates AI
Tensor processing units (TPUs) are Google's custom application-specific integrated circuits designed to accelerate machine learning workloads. They focus on matrix operations, using a systolic array of multiply-accumulators to perform computations without frequent memory access. This design makes them faster and more energy-efficient than general-purpose CPUs and GPUs for neural network tasks.
What is a tensor processing unit?
A tensor processing unit (TPU) is a custom application-specific integrated circuit (ASIC) developed by Google to accelerate machine learning workloads, as described in the Google Cloud TPU introduction. Unlike general-purpose CPUs, which handle a wide range of tasks, TPUs are specialized for the large matrix operations common in neural networks.
The IBM overview of AI accelerators explains that AI accelerators, including TPUs, are built to speed up AI neural networks, deep learning, and machine learning. They use parallel processing to perform billions of calculations at once, which is critical for handling the large amounts of data in AI applications.
TPUs are part of a broader category of AI hardware that includes GPUs, FPGAs, and other ASICs. They are designed to complement CPUs and GPUs, offering a specialized option for workloads dominated by matrix computations.
How TPUs work: systolic arrays and matrix multiplication
The core of a TPU is a matrix processor that excels at multiply-accumulate operations, which are fundamental to neural network computations. According to the TPU system architecture documentation, TPUs contain thousands of multiply-accumulators connected directly to each other in a systolic array, forming a large physical matrix.
In a systolic array, data flows through the array of processing elements, and each multiplication result is passed to the next accumulator. This eliminates the need for memory access during matrix multiplication, avoiding the von Neumann bottleneck that limits CPUs and GPUs. The TPU loads parameters and data from high-bandwidth memory (HBM) into the matrix multiply unit (MXU), performs the computation, and streams results out.
The MXU is composed of a systolic array of multiply-accumulators, with sizes varying by TPU version. For example, Cloud TPU v3 contains two 128x128 systolic arrays, while newer versions like v6e use 256x256 arrays. Each MXU can perform a large number of multiply-accumulate operations per cycle, using bfloat16 inputs and accumulating in FP32 format.
This architecture allows TPUs to achieve high computational throughput on neural network calculations, as the hardware is optimized for the specific pattern of matrix operations rather than general-purpose processing.
TPU performance and efficiency
The arXiv paper on TPU performance provides an early analysis of a TPU deployed in datacenters. It reports that the TPU's matrix multiply unit, with 65,536 8-bit MACs, offers a peak throughput of 92 TeraOps per second, and that the TPU was on average 15 to 30 times faster than contemporary GPUs or CPUs for neural network inference, with 30 to 80 times higher TOPS per watt.
This performance advantage comes from the TPU's deterministic execution model, which is a better match for the 99th-percentile response-time requirements of production neural network applications than the time-varying optimizations of CPUs and GPUs, such as caches and out-of-order execution. The lack of these features also makes the TPU relatively small and low power.
The Google Cloud TPU product page highlights that TPUs are engineered for the entire AI lifecycle, from training to inference, and are powered by Google Cloud's AI Hypercomputer architecture. They are used to power Google's flagship products, including Search, Photos, and Maps, and support frameworks like PyTorch and JAX.
When to use TPUs versus CPUs or GPUs
The Cloud TPU introduction provides guidelines for choosing hardware. TPUs are best suited for models dominated by matrix computations, those with large effective batch sizes, and workloads that train for weeks or months. They are also ideal for ultra-large embeddings common in ranking and recommendation systems.
CPUs are recommended for quick prototyping, simple models, or models with many custom operations. GPUs are a good choice for models with significant custom PyTorch or JAX operations, or when TensorFlow ops are not available on TPUs.
TPUs are not suited for workloads that require frequent branching, high-precision arithmetic, or custom operations in the main training loop. Additionally, models with dynamic shapes are problematic because the XLA compiler compiles the graph just in time for the first batch, and recompiling for shape changes is too slow.
The XLA compiler is essential for running code on TPUs. It takes the graph from an ML framework and compiles the linear algebra, loss, and gradient components into TPU machine code, while the rest of the program runs on the TPU host machine.
TPU system architecture and scaling
A TPU chip contains one or more TensorCores, each consisting of matrix multiply units, a vector unit, and a scalar unit. The MXUs provide the bulk of compute power, while the vector unit handles activations and softmax, and the scalar unit manages control flow and memory addresses, as detailed in the TPU architecture page.
TPUs can be connected in groups called slices, which are collections of chips within a TPU Pod connected by high-speed inter-chip interconnects. This allows scaling workloads with little to no code changes, as noted in the Cloud TPU introduction.
The Google Cloud TPU product page mentions that TPU clusters can scale up to 1 million chips, maximizing goodput for training frontier models. This scalability is supported by the AI Hypercomputer architecture, which combines purpose-built hardware with open software.
For developers, the TPU software stack includes JAX, an open-source library for numerical computation, and TorchTPU, a PyTorch backend that allows running PyTorch workloads on TPUs with standard syntax, as described on the Google Cloud TPU page.
| Hardware | Best for | Not suited for |
|---|---|---|
| CPU | Quick prototyping, simple models, small batch sizes | Large matrix-heavy models |
| GPU | Custom PyTorch/JAX ops, medium-to-large models | Models with ops not available on TPU |
| TPU | Matrix-dominated models, large batch sizes, long training | Frequent branching, high-precision arithmetic, dynamic shapes |
Frequently asked questions
What is a tensor processing unit?
A tensor processing unit (TPU) is a custom application-specific integrated circuit (ASIC) developed by Google to accelerate machine learning workloads, particularly those involving large matrix operations.
How does a TPU differ from a GPU?
While GPUs have thousands of ALUs and are general-purpose processors, TPUs are specialized matrix processors with a systolic array of multiply-accumulators, which reduces memory access and increases throughput for neural network computations.
What is a systolic array?
A systolic array is a network of processing elements where data flows through in a rhythmic fashion. In TPUs, it consists of multiply-accumulators directly connected to each other, allowing matrix multiplication without frequent memory access.
When should TPUs be used instead of CPUs or GPUs?
TPUs are best for models dominated by matrix computations, with large batch sizes, and that train for weeks or months. They are not suited for workloads with frequent branching, high-precision arithmetic, or dynamic shapes.
What software is needed to run code on a TPU?
Code must be compiled by the XLA compiler, which converts the ML framework graph into TPU machine code. TPUs support frameworks like PyTorch and JAX, with tools like TorchTPU for PyTorch workloads.