Modern AI Engineering

Lesson 17.3 · 26 min · video lesson

Google TPUs: Purpose-Built Hardware for Neural Networks

What if, instead of building a chip that can do anything, we built one that does almost nothing except multiply matrices — but does that thousands of times more efficiently?

In short: A TPU (Tensor Processing Unit) is Google's custom chip for neural networks. Its heart is a systolic array: a grid of simple multiply-add cells through which data pulses in lockstep, so each number fetched from memory is reused many times without being re-read. That makes matrix multiplication very fast and power-efficient, at the cost of flexibility: TPUs shine on large, regular, compiled workloads and struggle with irregular or dynamic ones.

What you will learn

  • What is a TPU
  • Why Google built the TPU
  • A quick refresher: CPU and GPU
  • The one operation that matters most
  • The big idea: Systolic Array
  • How data flows through a TPU
  • The full journey of a TPU computation
  • Why a TPU is so fast and power efficient
  • Where TPUs are used
  • Limitations of a TPU

In this lesson

  1. What is a TPU?
  2. Why Google built the TPU
  3. A quick refresher: CPU and GPU
  4. The one operation that matters most
  5. The big idea: the systolic array
  6. How data flows through a TPU
  7. The full journey of a TPU computation
  8. Why a TPU is so fast and power-efficient
  9. Where TPUs are used
  10. Limitations of a TPU
  11. Worked example, step by step
  12. Practice: try it yourself

Key terms in this lesson

TPU
Tensor Processing Unit: Google's custom chip for neural-network computation.
ASIC
Application-specific integrated circuit: a chip designed for one kind of task.
Systolic array
A grid of simple processing elements through which data flows rhythmically between neighbours.
Multiply-accumulate (MAC)
The operation acc = acc + a × b, the building block of matrix multiplication.
Weight-stationary
A systolic dataflow where each cell holds a fixed weight while activations and partial sums move.
XLA
Accelerated Linear Algebra: the compiler that turns ML programs into optimised TPU (and GPU) code.
bfloat16
A 16-bit floating-point format with FP32's exponent range and fewer mantissa bits.

This lesson is part of the Generative AI Engineering and Complete AI Engineer tracks. See plans and pricing to open the full lesson, its video, labs and quiz, or try the free lesson first (free account needed).

← 17.2 CUDA Kernels: Writing Parallel Code for NVIDIA GPUs · 17.4 Language Processing Units: A New Approach to LLM Inference →