GPUs, TPUs

The systems unit begins: making GPUs less magic

Tatsu opens the systems portion of the course with a confession — “I am not a systems person. Surprise.” — and a promise. The promise is a plot: throughput achieved on square matrix multiplies as a function of matrix dimension. Bigger matrices should mean more throughput, and broadly they do, but the plot is shot through with strange patterns — lines that split apart, periodic dips, sudden cliffs. By the end of the lecture you will be able to name every feature on it: compute intensity, tiling, wave quantization.

Why GPUs at all? Because the old road ran out. Through the 1990s you made computers faster by making serial instructions execute faster — Dennard scaling — and that tapped out in the 2000s. Transistors kept multiplying, but smaller transistors stopped meaning faster clocks. The alternative is horizontal: keep the clock, multiply the workers. GPU FLOP/s have grown more than 1000x in ten years — “there is no LLM scaling without GPU scaling” — driven at the V100 era by tensor cores, then by structured sparsity and ever-lower-precision number formats. He recommends Horace He’s blog, the GPU Mode group, and the JAX scaling book as the sources this lecture leans on.

Inside the chip: SMs, the memory hierarchy, and warps

A CPU is built for latency: big control units, a few ALUs, complex branching, each thread finishing fast. A GPU is built for throughput: hundreds of lightweight compute units, any individual task possibly slow, aggregate work enormous. The basic unit is the streaming multiprocessor (SM) — a GA100 has 128 of them — each independently programmable, each containing streaming processors that run threads in parallel, plus tensor cores, the specialized matmul circuits.

The compute story is half the picture; the memory story dominates this lecture. Memory comes in a hierarchy, and the closer it sits to the SM, the faster it is:

Memory	Where it lives	Latency (cycles)
Shared memory	inside the SM	~23/19 (ld/st)
L1 cache	inside the SM	~33
L2 cache	on die	~200
Global memory (HBM/DRAM)	chips next to the GPU	~290

Global memory is what the spec sheet means by “80 GB” — and it is roughly 10x slower than L1. Why not build the whole chip from fast SRAM? It’s about 100x more expensive and power-hungry; Grok-style all-SRAM accelerators exist for inference, but practical chips respect the hierarchy. (Shared memory and L1 are both SRAM — the difference is that the cache manages itself, while shared memory is programmable: you decide what goes in it.)

The model’s strengths: scaling is easy (want more throughput? add SMs), SIMT programming is “deceptively easy” (one instruction stream, many inputs — functional programming, essentially), and threads are lightweight enough that the scheduler can swap warps the instant one stalls. People exploited this before NVIDIA blessed it — an early paper hacked programmable shaders into fast matrix multiplies — but since the V100’s tensor cores, matmuls are the one privileged operation: more than 10x faster than any other floating-point op you can do on a GPU.

TPUs: convergent evolution

Two slides on TPUs, because they are the alternative evolution of the same idea. Build an energy-efficient ML accelerator and you converge: lightweight control, a big fast matmul unit, and a fast-memory/slow-memory hierarchy (HBM plus a local VMEM). The JAX book’s mapping is nearly exact — SM to Tensor Core, warp scheduler to VPU, SMEM to VMEM, tensor core to MXU — and the matmul units are even the same circuit, a systolic array. One naming trap: TPUs call their processors tensor cores; GPUs call their matmul units tensor cores. Disambiguate by context.

The real difference is count and size: an H100 has 132 SMs and 528 small matmul units; a TPU v5p has 2 tensor cores and 8 big MXUs. Many small units buy flexibility; few big ones lock you into large matmuls — Tatsu’s own batch-size sweep once stopped at 64 because the TPU refused anything smaller. That, and networking (next unit), are the differences; everything else in this lecture transfers.

Making it fast, tricks 1–4: divergence, precision, fusion, recomputation

One more hardware fact frames everything: compute FLOP/s have grown about 60,000x over twenty years, DRAM bandwidth about 100x, interconnect about 30x. The gap widens every generation, which is why the roofline model matters — below a critical operational intensity you are memory-bound and extra compute buys nothing; the goal of every trick below is to get onto (or make the most of) the flat, compute-bound part. Six tricks; the first is the odd one out, the rest are memory.

Control divergence. SIMT means every thread in a warp executes the same instruction — so an if statement forces both branches to execute, with threads on the wrong side masked out and idle. This is why GPU code multiplies by masks instead of branching.

Low precision. Fewer bits means fewer bytes to move. A nontrivial slice of the 1000x plot is just number representation: fp32 to bf16 to int8.

At fp8 there is no one canonical format — E4M3 (more mantissa) and E5M2 (more range) are used in different places — and the frontier formats get exotic. MXFP8 keeps elements in E4M3 but replaces the single fp32 scaling factor with one E8M0 scale per 32 elements, so different parts of a matrix can live at different magnitudes.

Operator fusion. Think of the GPU as a factory (compute) fed by a warehouse (memory) over a conveyor belt. Computing $\sin^2 x + \cos^2 x$ naively launches five CUDA kernels, each reading from and writing to global memory. A fused kernel reads once, does everything inside the SM, writes once. Easy fusions like this, torch.compile does automatically.

Recomputation. Stack three sigmoids and backprop: storing activations costs 1 read

• 3 writes forward, 3 reads + 1 write backward — 8 memory accesses for almost no math. Throw the activations away and recompute them on the fly during backward: 5 accesses total, 5/8 the memory traffic for a little extra compute. When compute is abundant and memory is precious, deliberately redoing work is optimal.

Coalescing, tiling, and the matrix mystery

Coalescing is DRAM trivia that stops being trivia the moment you multiply matrices. DRAM is read in bursts: accessing one location delivers its whole burst section (128 bytes or more) essentially for free, because the slow step is copying a row to the sense amplifier.

Tiling is the big one: cut the matrices into submatrices, load a tile of each input into shared memory once, and compute all the partial sums that tile supports before loading the next — repeated reads hit shared memory, not global.

But tile sizes interact with everything. If the matrix dimension isn’t divisible by the tile, you launch skinny tiles that waste their blocks; if rows don’t align with burst boundaries, reading one tile triggers two tiles’ worth of transactions (the fix is padding). PyTorch’s max-autotune flag spends fifteen minutes benchmarking exactly these tile choices. This is why Karpathy got a ~25% nanoGPT speedup by increasing the vocab size from 50,257 to 50,304 — the nearest multiple of 64.

Now the mystery plot reads itself. The rising diagonal is compute intensity (roofline). The lines fan out by divisibility: matrices divisible by 16 or 32 sit on top, odd sizes at the bottom — burst alignment, via tiling. And the periodic cliffs are the last effect:

FlashAttention: the victory lap

“At this point, we understand everything we need in order to reinvent our own FlashAttention.” The paper says it applies two established techniques — tiling and recomputation — to compute exact attention in subquadratic HBM accesses. Attention is three matmuls (the $QK^\top$ scores, times $V$) with a softmax in the middle. The matmuls we know how to do: Figure 1 of the paper is literally a tiled matrix multiply, copying blocks of $Q$, $K$, $V$ into SRAM and computing there.

The obstacle is the softmax — a global operation that couples every tile. The key trick is the online softmax (Milakov and Gimelshein, 2018): keep a running maximum and a running normalizer, and correct retroactively whenever a bigger value arrives,

$m_j = \max(m_{j-1}, x_j), \qquad d_j = d_{j-1}\, e^{m_{j-1} - m_j} + e^{x_j - m_j},$

then divide by $d_V$ at the end. Because this telescoping sum needs only the current tile plus two running scalars, the softmax computes tile by tile, the exponential fuses into the same kernel, and the final multiply by $V$ happens in tile form too. The backward pass adds recomputation: rather than store the $N^2$ attention matrix, throw it away and rebuild it tile by tile.

Check yourself

Napkin exercises in the lecture’s spirit — each one is a mechanism from the lecture with fresh numbers.

Problem 5.1 wave quantization

An A100 has 108 SMs and your matmul kernel uses 256 x 128 tiles, one tile per SM at a time. How many waves does a 3584-square matmul take? What happens at 3585, and roughly what does it cost?

Show solution

At 3584: $\frac{3584}{256} \times \frac{3584}{128} = 14 \times 28 = 392$ tiles. Waves: $\lceil 392/108 \rceil = 4$ (the last wave runs 68 of 108 SMs). At 3585 both tile counts round up: $15 \times 29 = 435$ tiles, and $\lceil 435/108 \rceil = 5$ waves — the fifth wave runs just 3 tiles while 105 SMs idle. One extra row and column adds about 0.06% more work but ~25% more wall-clock time (4 waves to 5). Same mechanism as the lecture’s 1792-to-1793 cliff.

Problem 5.2 online softmax update

Your running online-softmax state after some tiles is $m = 2$, $d = 3$. The next tile contributes a single logit $x = 4$. Compute the new state $(m', d')$, and explain why this trick is what makes attention tileable.

Show solution

New max: $m' = \max(2, 4) = 4$. The old normalizer was accumulated relative to $m = 2$, so rescale it and add the new term: $d' = d \cdot e^{m - m'} + e^{x - m'} = 3e^{-2} + e^{0} \approx 0.406 + 1 = 1.406$. The state is two scalars per row regardless of how many tiles have been processed — so each tile of $QK^\top$ scores can be exponentiated and folded in without ever materializing the full row, which is exactly what lets FlashAttention keep everything in SRAM and write only the final normalized output to HBM.

Problem 5.3 low-precision accounting

How many bytes does a 4096 x 4096 weight matrix occupy in bf16, and in MXFP8 (E4M3 elements, one E8M0 scale factor per 32 elements)? Now recall what MXFP8 training actually stores for each quantized matrix — how does the comparison change?

Show solution

Elements: $4096^2 \approx 16.78 \times 10^6$. bf16: 2 bytes each $\approx 33.6$ MB. MXFP8: 1 byte per element plus 1 byte per 32 elements = $16.78\times10^6 \cdot (1 + 1/32) \approx 17.3$ MB — roughly half. But because the per-32-element scaling pattern doesn’t survive transposition, the framework keeps a second quantized copy for the transpose: $2 \times 17.3 \approx 34.6$ MB, slightly more storage than bf16. The win is in bandwidth and tensor-core throughput during the matmuls, not in footprint — and quantize/dequantize overhead is why the end-to-end gain is 20–30%, not 2x.

Problem 5.4 coalescing

A row-major fp32 matrix of size 4096 x 4096 sits in global memory with 128-byte burst sections. A warp of 32 threads makes one load per step. Layout A: at each step, thread $t$ reads element $(i, t)$ — consecutive threads read consecutive elements of row $i$. Layout B: thread $t$ reads element $(t, i)$ — consecutive threads read the same column position of consecutive rows. How many DRAM transactions per step does each layout cost?

Show solution

Layout A: the 32 threads touch 32 consecutive fp32 values = 128 contiguous bytes — exactly one burst section. One transaction, fully coalesced. Layout B: consecutive threads are one full row apart, $4096 \times 4 = 16{,}384$ bytes — every thread lands in a different burst section, so 32 transactions per step, a 32x traffic blowup for the same useful data. This is the lecture’s mnemonic in action: for row-major storage, warps should span within a row; per-thread traversal should run along the row so that the warp’s simultaneous accesses stay contiguous.