Cloud TPU

Overview

Cloud TPUs are Google's custom chips for machine learning, offered as a service on Google Cloud. TPU stands for tensor processing unit. Unlike general-purpose GPUs, TPUs were designed from the start for one job: the large matrix multiplications at the heart of neural networks. They are available as single machines with a few chips, as slices of many chips, and as pods with thousands of chips joined by a dedicated high-speed network. Google trains and serves its own Gemini models on TPUs, and other companies use them through Google Cloud.

An everyday picture: a bakery that only makes bread. A general kitchen can cook anything, but switching between dishes takes time and equipment. A specialised bread bakery has ovens, mixers and conveyor belts laid out for bread alone, so it produces loaves faster and more cheaply. TPUs are that specialised bakery for neural network maths, while GPUs are versatile kitchens that also happen to be very good at bread.

TPUs are usually programmed through compilers rather than hand-written kernels. Frameworks such as JAX, PyTorch and TensorFlow describe the computation, and the XLA compiler turns it into instructions for the TPU. That makes the programming experience different from CUDA on NVIDIA GPUs.

What it is

Chips in a slice talk to each other over ICI without going through the data-centre network, which makes collective operations such as all-reduce very fast. Several slices can be combined over the data-centre network for even larger jobs.

Where it came from

Google began designing TPUs around 2013, when it calculated that the growth of neural network services, such as speech recognition, would otherwise need far more data centres. The first TPU, focused on inference, went into Google's data centres in 2015 and was announced publicly in 2016.

The second generation, introduced in 2017, added training support and became the first Cloud TPU available to customers, in 2018. Later generations followed: v3 with liquid cooling, v4 with optical circuit switches to reconfigure pod connections, v5e for cost-efficient training and serving, v5p for large-scale training, and Trillium, the sixth generation, in 2024.

Ironwood, the seventh generation, was announced in April 2025 with a focus on inference, and became generally available in late 2025. Each Ironwood chip has 192 GB of HBM, and pods scale to 9,216 chips. In April 2026 Google previewed an eighth generation split into separate chips for training and for inference.

How it works

The systolic array

The core of a TPU is the matrix multiply unit, a grid of multiply-accumulate cells. Data flows through the grid in waves, like blood pumped through a body, which is where the word "systolic" comes from. Each cell multiplies, adds and passes values to its neighbour, so data loaded once is reused many times without going back to memory. This makes TPUs very efficient at dense matrix maths in formats such as BF16, a number format Google introduced, and in lower precisions such as INT8 and FP8 on newer chips.

Compilation with XLA

A framework builds a graph of operations. XLA fuses operations, chooses memory layouts, schedules work and generates TPU code. Because compilation happens ahead of execution, TPUs favour static shapes: a new input shape can trigger recompilation. Programs are therefore often written with fixed batch sizes and padded sequences.

Scaling across chips

Large models are split across chips by sharding. In JAX, a user describes a mesh of devices and how each array is partitioned, and the compiler inserts the communication. ICI carries that traffic directly between neighbouring chips. Libraries and frameworks such as MaxText, Pathways and PyTorch/XLA provide ready-made recipes for large training and serving jobs.

Using Cloud TPUs

Users create TPU VMs or request slices of a chosen size, either directly, through Google Kubernetes Engine, or through managed services such as Vertex AI. Capacity can be reserved, used on demand where available, or used as cheaper preemptible capacity.

import jax
print(jax.devices())  # lists the TPU cores visible to this host

Key concepts and terms

A worked example

This scenario is illustrative. A research team wants to pre-train a 2-billion-parameter language model on 200 billion tokens.

1. They write the model in JAX, using an open training framework that supports TPUs, and test it on a single small TPU VM with a tiny dataset.
2. They fix the sequence length at 4,096 tokens and use a constant global batch size, so XLA compiles the training step once.
3. For the full run they request a slice of 64 chips. They describe a device mesh with data parallelism across most chips and tensor sharding within small groups; the compiler inserts the collectives.
4. The first step takes several minutes to compile; later steps run in a steady time, and profiling shows the MXUs busy most of the time.
5. A data pipeline bug causes variable-length batches for an hour, triggering repeated recompilation and slowing training. Padding all batches to the fixed length restores speed.
6. Checkpoints go to Cloud Storage every 1,000 steps, so the job can resume if preemptible capacity is reclaimed.

Strengths and limitations

Strengths:

• Very efficient for dense matrix maths in neural networks.
• Fast, direct chip-to-chip interconnect for large-scale jobs.
• Large pods suit training and serving of very large models.
• Strong integration with JAX and XLA, and growing PyTorch support.
• Competitive price-performance for suitable workloads.

Limitations:

• Available only on Google Cloud.
• Compiler-driven model favours static shapes; dynamic workloads need care.
• Smaller ecosystem of custom kernels and tools than CUDA.
• Some PyTorch code paths are less mature than on GPUs.
• Capacity for the newest generations may need reservations.

How it compares with alternatives

Practical guidance and best practices

1. Start with JAX or a framework recipe that already supports TPUs.
2. Keep input shapes static; pad sequences and fix batch sizes.
3. Test on a small TPU VM before requesting a large slice.
4. Use BF16 for training, which TPUs handle natively.
5. Feed data fast enough; profile the input pipeline as well as compute.
6. Checkpoint regularly to Cloud Storage, especially on preemptible capacity.
7. Use the TPU profiler to check MXU utilisation and communication time.
8. Plan capacity early with reservations for long or large runs.

Common mistakes

• Sending varied input shapes and paying for constant recompilation.
• Porting GPU code that relies on custom CUDA kernels without alternatives.
• Ignoring host-side data loading, which leaves TPUs idle.
• Requesting a huge slice before the code works on a small one.
• Forgetting to shard large arrays, causing out-of-memory errors on one chip.
• Treating TPU and GPU performance numbers as directly comparable without testing.

Frequently asked questions

Can I buy a TPU?

Data-centre TPUs are not sold as hardware; they are rented through Google Cloud. Google has sold small Edge TPU devices for on-device inference, which are a different product.

Do TPUs work with PyTorch?

Yes, through PyTorch/XLA and related projects. JAX is the most common choice for TPUs, and support for PyTorch has been improving, but some features and libraries still work best on GPUs.

Are TPUs faster than GPUs?

It depends on the workload, model size, software and generation. TPUs excel at large, dense, regular workloads at scale. GPUs are more flexible and have a larger software ecosystem. Benchmark your own job.

What is a TPU pod?

A pod is the largest group of TPU chips connected directly by the inter-chip interconnect. Users usually rent a slice of a pod sized to their job.

Why does my first TPU step take so long?

The first step includes XLA compilation of the whole program. Later steps reuse the compiled code, as long as input shapes and the program stay the same.

Glossary

• Accelerator: a chip specialised for certain computations, such as ML.
• Multiply-accumulate: multiplying two numbers and adding the result to a total.
• Recompilation: compiling again because shapes or code changed.
• Preemptible capacity: cheaper compute that can be reclaimed at short notice.
• Device mesh: a logical arrangement of chips used for sharding.
• PyTorch/XLA: the project that runs PyTorch on XLA devices such as TPUs.
• Edge TPU: a small Google chip for on-device inference.

Summary

Cloud TPUs are Google's custom accelerators for neural networks, built around systolic matrix units, high-bandwidth memory and a direct inter-chip interconnect that links thousands of chips into pods. First deployed internally in 2015 and offered to customers since 2018, they have reached their seventh generation, Ironwood, with an eighth generation previewed in 2026. They reward workloads that are large, regular and written for compilation, typically in JAX, and they are available only through Google Cloud.

Related

It belongs to Hardware in my stack, under MLOps and AI infrastructure. Others in the same group: