Understanding Transformers · Part 16 of 16

Hardware Accelerators: GPUs, TPUs, and JAX

Every lesson so far has been matrix multiplication in disguise. This lesson is about the chips built to do that one operation as fast as possible.

The series was matmul all along

Look back at what actually ran in this series' code: embedding lookups, Q KT, weights @ V, the feedforward's W₁x and W₂h, the final projection to logits. Every one of them is a matrix multiplication, and a transformer of any size is, computationally, an enormous stack of them. That single fact is why the hardware underneath matters as much as the architecture on top of it — a chip that does matrix multiplication fast makes every part of this series faster, uniformly, without changing a line of the model.

Why a naive loop is slow

A matrix multiply of two n×n matrices is n3 individual multiply-and-add operations. Written as nested loops, a CPU does them essentially one at a time. Written as a single vectorized call, the same CPU can do many of those multiply-adds together — its instruction set has operations built to apply one arithmetic step across several numbers in one cycle. Same chip, same math, very different speed:

import numpy as np, time

n = 200
A, B = np.random.randn(n, n), np.random.randn(n, n)

# Naive: one multiply-add at a time
start = time.time()
C = np.zeros((n, n))
for i in range(n):
    for j in range(n):
        s = 0.0
        for k in range(n):
            s += A[i, k] * B[k, j]
        C[i, j] = s
loop_time = time.time() - start

# Vectorized: the same n^3 multiply-adds, done in bulk
start = time.time()
C_fast = A @ B
vector_time = time.time() - start
# loop_time is typically 100-1000x vector_time -- same CPU, same result

Every NumPy snippet in this series has quietly relied on that gap — A @ B instead of a triple loop is already an appeal to hardware built for bulk arithmetic. GPUs and TPUs are the next step of the same idea, taken much further.

GPUs: thousands of small cores, one instruction

A CPU has a handful of powerful cores optimized for varied, branching, sequential work — good at running a web server or a Python interpreter. A GPU was originally built to compute color values for millions of independent pixels at once, which turned out to be the same shape of problem as multiplying large matrices: enormous numbers of independent, identical arithmetic operations with little to no branching. A modern GPU has thousands of small arithmetic cores that all execute the same instruction on different data simultaneously (an approach called SIMT — single instruction, multiple threads). Handing a matmul to a GPU means splitting its n3 multiply-adds across thousands of cores instead of a handful, which is why training and running large transformers moved to GPUs as soon as they got big enough to need it.

TPUs: a chip that does only this

A TPU (Tensor Processing Unit, Google's accelerator) goes a step further: instead of a general-purpose parallel chip that happens to be good at matmul, it's an ASIC — a chip designed for essentially one job. Its core is a systolic array: a grid of multiply-accumulate units wired directly to their neighbors, so partial results flow from one unit to the next without being shuttled back out to memory between steps. Memory movement, not arithmetic, is usually what limits speed at scale, and the systolic array is built specifically to avoid it for exactly the matmul-heavy pattern a transformer produces. The trade-off is generality: a TPU is far less flexible than a GPU, but for large-scale transformer training and inference, it can be substantially more efficient per watt.

JAX: compiling the whole function, not just running it

Every PyTorch line in the companion notebook runs eagerly — each operation executes the instant Python reaches it, exactly as written, one kernel launch at a time. JAX, also from Google, is built around a different idea: jax.jit traces a whole function once, hands it to the XLA compiler — the same compiler backend TPUs use internally — and gets back one compiled, fused program. Fusing a dozen small operations into one larger kernel cuts both the number of kernel launches and the amount of intermediate memory that has to be written out and read back: the same memory-movement bottleneck the TPU section above described, now attacked in software instead of silicon. That's where JAX's speed comes from — not a faster chip, a smarter compiled program.

jax.jit isn't the only function transformation JAX offers. jax.grad takes an ordinary function and returns a new function computing its gradient — the same calculus from Part 2, expressed as a transformation instead of PyTorch's .backward() walking a graph built while the forward pass ran. jax.vmap takes a function written for a single example and turns it into one that runs across a batch automatically, no explicit loop required — the companion notebook's JAX training step uses exactly this instead of looping over the batch by hand. jit, grad, and vmap compose: that training step wraps a vmap'd, gradient-taking function in a single jax.jit, and the whole thing compiles into one program.

Compiling through XLA has a side effect worth knowing about: the same compiled program can target CPU, GPU, or TPU, since XLA is the compiler backend all three understand. That's not, on its own, a reason to reach for JAX — PyTorch already lets you pick a device with one line (device = torch.device("cuda" if torch.cuda.is_available() else "cpu"), right there in the companion notebook). The real difference is that JAX gets you a faster, fused, differentiable, auto-batched program as the ordinary way of writing the code, and portability comes along for free with it — rather than something you get by moving tensors between devices by hand.

Open the notebook in Colab →