What is a Tensor Processing Unit?

What is a Tensor Processing Unit?

A Tensor Processing Unit (TPU) is an application-specific integrated circuit (ASIC) designed by Google specifically to accelerate machine learning workloads. Unlike general-purpose processors, a TPU is hardwired to perform dense tensor arithmetic—primarily large-scale matrix multiplications and additions—at extreme speed and power efficiency. TPUs power Google services like Search, Translate, and Gemini, and are available to cloud developers for training and inference across massive compute clusters.

Think of a CPU as a skilled artisan carrying individual water buckets back and forth between a well and a fire, and a GPU as a fleet of hundreds of workers doing the same simultaneously. A TPU is like an automated pipe network: water flows directly from one valve to the next in continuous waves without ever having to return to the well until the operation is finished.

How it works

The efficiency of a TPU comes from eliminating the traditional von Neumann memory bottleneck during matrix math. It accomplishes this through a dedicated hardware pipeline:

1. The Systolic Matrix Multiply Unit (MXU): The core compute engine is a two-dimensional grid of multiply-accumulate arithmetic logic units (ALUs), typically arranged in dimensions of 128 × 128 or 256 × 256. Instead of writing intermediate results back to registers or memory caches after every multiplication, values stream directly across spatially adjacent ALUs in rhythmic cycles (a "systolic" flow). One matrix (such as model weights) is held stationary in the array while data vectors stream across it, accumulating products in place.
2. Vector and Scalar Units: Dedicated units handle operations that the MXU is not built for, such as activation functions (ReLU, GELU, Sigmoid), layer normalization, softmax calculations, and memory address computation.
3. High-Bandwidth Memory (HBM) and On-Chip Buffers: Model parameters and activations stream from high-bandwidth memory (e.g., 95 GiB HBM on TPU v5p with up to 2.765 TB/s bandwidth) directly into local SRAM buffers, keeping the MXUs saturated.
4. Inter-Chip Interconnect (ICI): Individual TPU chips connect directly to adjacent chips over custom high-speed optical and copper links without traversing standard data center Ethernet switches, forming multi-dimensional torus topologies (Pods) that scale to thousands of interconnected chips.
5. XLA (Accelerated Linear Algebra) Compiler: Code written in frameworks like JAX, PyTorch, or TensorFlow cannot execute raw x86 or CUDA machine instructions on a TPU. The OpenXLA/XLA compiler parses computation graphs, fuses operators, arranges tensor memory layouts, and generates optimized TPU machine code.

Tensor Processing Unit vs GPU

While both TPUs and Graphics Processing Units (GPUs) accelerate AI, they differ fundamentally in hardware architecture and software philosophy. A GPU is a massively parallel, instruction-driven processor composed of thousands of flexible cores (like NVIDIA Streaming Multiprocessors) that read and write intermediate results to shared register files, supporting broad general-purpose tasks from 3D graphics rendering to scientific simulations. A TPU is a domain-specific ASIC whose Matrix Multiply Units discard general instruction scheduling logic in favor of a fixed systolic array that passes values directly between arithmetic units, delivering higher compute density and performance-per-watt on rigid matrix math. Furthermore, while GPUs rely heavily on the CUDA runtime and manually tuned vendor kernels, TPUs rely on ahead-of-time compilation via XLA to optimize tensor graph execution.

A practical example

Developers interact with TPUs via Python frameworks like JAX or PyTorch (via torch_xla). The following example shows how to verify TPU hardware visibility and execute a distributed matrix multiplication across all available TPU cores using JAX on a Cloud TPU VM:

import jax
import jax.numpy as jnp

# Verify connected TPU devices
devices = jax.devices()
print(f"Active TPU Devices ({len(devices)} cores): {devices}")

# Define matrix multiplication mapped across TPU cores
@jax.jit
def parallel_matmul(a, b):
    return jnp.matmul(a, b)

# Create sample tensors in float32 / bfloat16
key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (8192, 8192), dtype=jnp.bfloat16)
y = jax.random.normal(key, (8192, 8192), dtype=jnp.bfloat16)

# Execute compiled systolic matrix multiplication
result = parallel_matmul(x, y)
result.block_until_ready()
print(f"Output shape: {result.shape}, dtype: {result.dtype}")

When executed on a v5litepod-8 instance, JAX outputs the active TPU core handles (e.g., [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), ...]), compiles the arithmetic graph via XLA into MXU systolic instructions, executes the 8,192 × 8,192 bfloat16 multiplication without register-spill stalls, and returns Output shape: (8192, 8192), dtype: bfloat16.

Tools and implementations

Tool	What it is for	Link
OpenXLA / XLA	Domain-specific compiler that converts ML graphs into TPU and GPU machine code	openxla.org
JAX	Autograd and XLA-driven array computing framework optimized for TPU hardware	github.com/jax-ml/jax
PyTorch / XLA (torch_xla)	PyTorch integration package enabling PyTorch models to compile and run on Cloud TPUs	github.com/pytorch/xla
MaxText	High-performance, highly scalable open-source LLM training and inference framework written in JAX for TPUs	github.com/google/maxtext
Google Cloud TPU API	Infrastructure service for provisioning TPU VMs, multi-host slices, and Pod topologies	cloud.google.com/tpu

Where it goes wrong

TPUs provide immense throughput, but their specialized design introduces strict architectural trade-offs:

• Compilation Overhead: Because TPUs rely on XLA compilation, dynamic tensor shapes (such as variable sequence lengths without padding) trigger frequent, slow recompilations at runtime that stall execution pipelines.
• Poor Irregular Compute: Workloads with branching control flow, sparse memory access, non-matrix calculations, or unvectorized scalar logic underutilize the MXU systolic arrays, performing significantly worse than on general-purpose CPUs or GPUs.
• Ecosystem Lock-in: Unlike the pervasive CUDA ecosystem with ubiquitous open-source C++ kernels (e.g., FlashAttention, Triton, custom CUDA extensions), writing raw custom assembly or low-level kernels for TPUs is constrained by Google's proprietary toolchain and XLA abstractions.
• Padding Waste: Systolic arrays are physically dimensioned (e.g., 128 × 128); if matrix dimensions are not multiples of the array tile size, the hardware pads dimensions with zeros, expending wasted FLOPs.

Why it matters now

As frontier AI models scale to hundreds of billions of parameters, training clusters require hundreds of megawatts of power and thousands of synchronized processors. Custom ASICs like the TPU offer superior compute density, energy efficiency, and low-latency interconnects compared to general-purpose hardware, making them central to the economic viability of AI infrastructure. Demonstrating this push into extreme environments, in October 2026, Google announced that its Project Suncatcher prototype satellite successfully executed the Gemma 4 model in orbit to generate its first response from space, with upcoming tests planned to evaluate TPU performance under orbital thermal stress and radiation exposure.

FAQs

Can I run custom CUDA kernels on a TPU? No, TPUs do not support NVIDIA's CUDA architecture or PTX assembly. Workloads must be written in frameworks that compile through XLA or StableHLO, such as JAX, PyTorch XLA, or TensorFlow.

What data types work best on a TPU? TPU Matrix Multiply Units are physically optimized for bfloat16 (Brain Floating Point) and int8/fp8 operations, which preserve the dynamic exponent range of 32-bit floats while halving memory footprint and multiplying ALU density.

How do TPUs connect across large clusters? TPUs use an Inter-Chip Interconnect (ICI), a custom optical circuit switching and copper network built directly into the silicon boards. This allows thousands of TPU chips to communicate in a 2D or 3D torus topology without using traditional data center switches.

What is the difference between a TPU VM and legacy TPU nodes? Legacy TPU architectures required a separate Compute Engine VM to drive the TPU over gRPC network calls, causing significant I/O latency. Modern TPU VMs run user code and the Python runtime directly on the host machine physically attached to the TPU board.