Skip to main content

TPU overview

Exam guide§2.1

While GPUs are general-purpose parallel processors, TPUs take a distinct approach to AI acceleration: they are purpose-built for the matrix math at the heart of machine learning.

What are Cloud TPUs?

TPU stands for Tensor Processing Unit. TPUs are custom-developed, application-specific integrated circuits (ASICs) built by Google specifically to accelerate the intensive computations in machine learning - purpose-built engines for AI.

Code that runs on a TPU is compiled by the Accelerated Linear Algebra (XLA) compiler. XLA is a just-in-time (JIT) compiler: it takes the computational graph produced by an ML framework and translates the linear algebra operations, loss functions, and gradient calculations into machine code for the TPU. The rest of your program runs on the TPU host machine, and the XLA compiler itself ships inside the TPU VM image on that host.

ML frameworkPyTorch / JAXComputationalgraphXLA JIT compilerlinear algebra ·loss · gradientsTPU machine coderuns on the TPUTPU host machineRuns the rest of your program · the XLA compiler ships inside the TPU VM image
XLA compilation flow - the framework's computational graph is JIT-compiled by XLA into TPU machine code; everything else runs on the TPU host.

Google runs this infrastructure at planet scale - the chips are wired together with a custom Inter-Core Interconnect and, across a data center, reconfigured on the fly by optical circuit switches into the machine shape a workload needs. TPUs power Google's own services, including Photos and Search.

Inside Google's Cloud TPU data centers

When to use TPUs

Cloud TPUs shine on specific, demanding AI tasks where their specialized architecture pays off:

  • Training massive deep learning models - large, complex models, especially large language models (LLMs), whose training is dominated by matrix calculations.
  • Models relying on embeddings - TPUs feature SparseCores, dataflow processors built to accelerate embedding-heavy models like recommendation systems.
  • Scientific and healthcare AI - computationally intensive workloads such as protein folding and drug discovery, where complex simulations and deep learning matter.
FactsSparseCores
  • SparseCores are dedicated dataflow processors on the TPU for accelerating embedding operations - the lookups and updates behind recommendation systems and other sparse models.

In short: large-scale neural networks, complex embedding operations, or advanced scientific modeling are the sweet spot for Cloud TPUs.

Why choose Cloud TPU?

Cloud TPUs cover the full arc from initial training to fine-tuning to high-volume inference. Five reasons they stand out:

ReasonWhat it means
Cost-efficient scalingEngineered for high performance while staying cost-effective on demanding AI workloads - scale up without prohibitive expense.
Versatile framework supportNot locked to one ecosystem: the XLA compiler accelerates PyTorch and JAX alike, so teams choose what fits.
Integrated orchestrationIntegrate with Google Kubernetes Engine (GKE) for large-scale AI; the Dynamic Workload Scheduler provisions all needed accelerators at once, improving scalability and avoiding wasteful spend.
Simplified development with Agent PlatformDeeply integrated into Agent Platform, Google Cloud's fully-managed AI platform - it handles the underlying infrastructure so you focus on the workload.
Powered by reconfigurable OCSTPU v4 was the first supercomputer to deploy Optical Circuit Switches (OCS), which dynamically reconfigure the interconnect topology between chips.
GotchaOCS vs. InfiniBand

Optical Circuit Switches are significantly cheaper, use less power, and provide more bandwidth than traditional InfiniBand. The efficiency comes from avoiding power-hungry optical-to-electrical-to-optical conversions and complex network packet switches. Reconfigurable OCS improves scale, availability, utilization, modularity, deployment, security, power efficiency, and performance.