What a GPU is fast at, why moving bytes is the real cost, and why FlashAttention exists.
Chapter 8 priced Beacon's run in operations: 3.8 × 10²⁵. Chapter 9 built the corpus those operations consume. This chapter is about the machines that execute them, and it has one idea at its centre: arithmetic is cheap and moving data is expensive. Every engineering decision in the rest of Part 2 and in Chapter 19 descends from that. By the end you can look at any operation in the model, from Chapter 2's matmul to Chapter 3's T² scores, and say whether a GPU will spend its time computing or waiting.
The question this chapter answers: what does a GPU actually do well, what limits it, and how does the model's arithmetic have to be arranged to keep it busy?
A kitchen with a thousand cooks and one small pantry door. Each cook can chop a vegetable in a second. But every vegetable comes through the door, and the door passes one crate a minute. If each crate holds a thousand vegetables, the cooks are busy: a crate arrives, a thousand cooks chop for a second, done, next crate. If each crate holds ten vegetables, ten cooks chop for a second and nine hundred and ninety stand idle for fifty-nine seconds. The kitchen is not limited by cooks. It is limited by the door.
The obvious fix is to bring in fewer, fuller crates. The clever fix is to notice that some recipes need the same vegetable chopped, mixed, chopped again, and to keep it on the counter between steps instead of sending it back through the door each time. A cook's counter is tiny, but it is right there. A tray in the pantry is bigger but behind the door. The warehouse across town holds everything and takes an hour.
The cooks are the GPU's cores. The door is memory bandwidth. Crates per minute is bytes per second; vegetables per crate is how much arithmetic you get to do per byte fetched. The counter is on-chip SRAM, the pantry is HBM, the warehouse is the other GPUs and the network. FlashAttention is the recipe rewritten so the vegetable never goes back through the door.
| In the kitchen | In the machine | The word we will use |
|---|---|---|
| A thousand cooks chopping in lockstep | Thousands of simple cores executing the same instruction on different data | SIMT, streaming multiprocessor (SM) |
| A cook who chops a whole tray in one motion | A unit that multiplies a small matrix block in one instruction | tensor core |
| The pantry door | Bytes per second between HBM and the cores | memory bandwidth |
| Vegetables per crate | FLOPs performed per byte moved | arithmetic intensity |
| Counter · pantry · warehouse | Registers and SRAM · HBM · other GPUs and network | memory hierarchy |
| "Are we limited by cooks or by the door?" | A plot of achievable FLOP/s against intensity | roofline model |
| The recipe rewritten to keep the vegetable on the counter | Attention computed in tiles that never leave SRAM | FlashAttention |
| Fraction of the hour the cooks were chopping | Achieved FLOP/s over peak, for the whole run | model FLOPs utilisation (MFU) |
A CPU has a handful of cores, each built to run one thread as fast as possible: deep pipelines, branch prediction, large caches. A GPU has thousands of cores, each simple, arranged in groups (streaming multiprocessors, SMs) that execute one instruction across many data elements at once. An H100 has 132 SMs, and each SM runs thousands of threads in lockstep groups public. Ask it to do a thousand different things and it crawls. Ask it to do the same thing to a thousand numbers and it flies.
That is a description of a matmul. Chapter 2 showed that every entry of XW is an independent dot product; a [16384, 8192] × [8192, 28672] product is 470 million independent dot products of length 8192, which is exactly the shape of work a GPU was built for. Since about 2017 GPUs have gone further and added tensor cores: hardware that multiplies a small block, such as 16 × 16 by 16 × 16, in a single instruction public. Nearly all of a GPU's headline "FLOP/s" number is tensor-core throughput on matmul-shaped work at reduced precision. Feed it anything else and you get a fraction.
| GPU | Dense BF16 matmul | HBM size | HBM bandwidth | On-chip SRAM per SM | Evidence |
|---|---|---|---|---|---|
| A100 (2020) | 312 TFLOP/s | 80 GB | 2.0 TB/s | 192 KB | public |
| H100 SXM (2022) | 989 TFLOP/s | 80 GB | 3.35 TB/s | 228 KB | public |
| B200 (2024) | ≈ 2.25 PFLOP/s | ≈ 180–192 GB | ≈ 8 TB/s | ≈ 228 KB | public, spec sheets vary by SKU |
Read across a row. From A100 to H100, matmul throughput tripled; bandwidth went up 1.7×. The gap between "how fast can it compute" and "how fast can it be fed" widened, and it keeps widening with every generation. That trend is why the rest of this chapter is about bytes.
Data lives at several distances from the cores, and each step outward is roughly ten times larger and several times slower.
Two levels matter most for the model. HBM is where the weights and activations live between operations; every matmul reads its inputs from there and writes its outputs back. Shared memory (SRAM on each SM) is where a kernel keeps the block it is working on; it is tiny, about 228 KB, but roughly six times the bandwidth of HBM and right next to the cores. The art of a fast kernel is to load a block from HBM into SRAM once, do as much arithmetic on it as possible, and write back once.
Beyond the GPU, NVLink connects the eight GPUs of a node at about 900 GB/s per GPU on H100 public, and InfiniBand or Ethernet connects nodes at 400 Gb/s per port, 50 GB/s, about 70 times slower than HBM public. Chapter 11 is about arranging a training run so that the slow links carry as little as possible.
Here is the tool that turns the kitchen picture into a number. For any operation, count the FLOPs it performs and the bytes it must move between HBM and the cores. Their ratio is the operation's arithmetic intensity, in FLOPs per byte. Compare it to the GPU's ratio of peak FLOP/s to bandwidth, the ridge point:
H100: ridge = 989 × 10¹² FLOP/s ÷ 3.35 × 10¹² B/s ≈ 295 FLOP per byte intensity < 295 → memory-bound: the cores finish before the next bytes arrive; time = bytes / bandwidth intensity > 295 → compute-bound: bytes arrive before the cores are done; time = FLOPs / peak
An operation needs to do about 300 floating-point operations on every byte it fetches, or the H100's cores wait. That is a demanding number, and most of the operations in a transformer do not meet it.
Three operations from a Llama 3 70B block, in bf16 (2 bytes per number). Shapes: C = 8192, MLP hidden H = 28672, sequence T = 8192, two sequences in the batch.
1 · the MLP expand matmul during prefill: [16384, 8192] × [8192, 28672]
FLOPs = 2 × 16384 × 8192 × 28672 ≈ 7.7 × 10¹²
bytes = 2 × (16384×8192 + 8192×28672 + 16384×28672) ≈ 1.7 × 10⁹ (read X, read W, write Y)
intensity ≈ 4,600 FLOP/B ≫ 295 compute-bound. Time ≈ 7.7×10¹² / 989×10¹² ≈ 7.8 ms
2 · the same matmul during decode: one token per sequence: [2, 8192] × [8192, 28672]
FLOPs = 2 × 2 × 8192 × 28672 ≈ 9.4 × 10⁸
bytes = 2 × (2×8192 + 8192×28672 + 2×28672) ≈ 4.7 × 10⁸ (the weight matrix dominates)
intensity ≈ 2 FLOP/B ≪ 295 memory-bound. Time ≈ 4.7×10⁸ / 3.35×10¹² ≈ 140 µs, cores idle 99%
3 · attention scores for one head, one sequence, the matrix written out: [8192, 128] × [128, 8192]
FLOPs = 2 × 8192 × 8192 × 128 ≈ 1.7 × 10¹⁰
bytes = 2 × (2×8192×128 + 8192×8192) ≈ 1.4 × 10⁸ (writing the 8k × 8k result is most of it)
intensity ≈ 124 FLOP/B < 295 memory-bound, and the softmax that follows reads it all back again
$ python code/ch10/roofline.pyridge point = 295 FLOP/byte matmul [16k, 8k]×[8k, 28k] (prefill) intensity 4587.5 FLOP/B compute 7782.2 µs memory 500.8 µs → compute-bound matmul [2, 8k]×[8k, 28k] (decode) intensity 2.0 FLOP/B compute 0.9 µs memory 140.3 µs → memory-bound Q·Kᵀ one head, T=8k, materialised intensity 124.1 FLOP/B compute 17.4 µs memory 41.3 µs → memory-bound RMSNorm over [16k, 8k] intensity 1.0 FLOP/B compute 0.5 µs memory 160.3 µs → memory-bound softmax over [8k, 8k] intensity 1.2 FLOP/B compute 0.3 µs memory 80.1 µs → memory-bound
The same matrix multiply is compute-bound at prefill and memory-bound at decode, by a factor of a thousand in intensity, because the weight matrix is read once either way but multiplied against 16,384 rows in one case and 2 in the other. This is Chapter 7's "decode is memory-bound", now with the numbers. And every elementwise operation (norms, activations, softmax) sits at intensity around 1: they are pure memory traffic, and a naive implementation that runs them as separate passes over HBM costs more time than the matmuls they surround.
Three consequences that the rest of the book keeps using:
Chapter 3's attention, executed literally: compute S = QKᵀ, a [T, T] matrix, and write it to HBM. Read it back, compute the row softmax, write P. Read P back, multiply by V. Three passes over a matrix with T² entries. At T = 128k that matrix is 33 GB per head per layer in bf16, larger than most GPUs' memory, and each of the three passes moves it at 3.35 TB/s. The arithmetic, by contrast, is a modest 4·T²·d: high intensity when done in one go, but the literal implementation never does it in one go.
The obstacle to doing it in one go is the softmax. Each row's probabilities need the row's maximum (for numerical safety) and the row's sum, and both depend on all T scores in the row. You cannot blend the values with the first hundred keys' probabilities until you know the denominator, which needs the remaining 127,900. So it seems the whole row must exist before any of it can be used.
FlashAttention's insight, from a 2022 paper by Dao and colleagues public, is that the softmax can be computed incrementally. Keep three running quantities per query row: the maximum seen so far m, the sum of exponentials so far ℓ, and the unnormalised output so far o. Load a tile of keys and values into SRAM, compute the tile's scores, and fold them in. When a new tile raises the maximum, multiply the old ℓ and o by e^(m_old − m_new) to rescale them to the new reference point, then add the tile's contributions. At the end, divide o by ℓ. The result is bit-for-bit the same softmax-weighted blend; the [T, T] matrix was never stored anywhere.
For one query row, after processing keys 1 … j:
m_j = max(m_{j−1}, s_j) running maximum of scores
ℓ_j = ℓ_{j−1} · e^(m_{j−1} − m_j) + e^(s_j − m_j) running sum, rescaled to the new max
o_j = o_{j−1} · e^(m_{j−1} − m_j) + e^(s_j − m_j) · v_j running output, rescaled the same way
final output = o_T / ℓ_T = Σ_j e^(s_j − m_T) v_j / Σ_j e^(s_j − m_T) = softmax(s) · V
The rescaling factor e^(m_old − m_new) is what makes the trick exact: every earlier term was computed relative to an old maximum, and multiplying by the factor converts it to the new one. In practice tiles of 64 to 128 keys are folded in at once rather than one key at a time, but the three lines are the whole idea. The backward pass (Chapter 5) recomputes the tile's scores from Q and K instead of storing P; that costs extra FLOPs and saves the T² memory again, which at these intensities is the right trade.
What changes and what does not:
| Attention as written (Chapter 3) | FlashAttention | |
|---|---|---|
| FLOPs, forward | 4·T²·d per head | 4·T²·d per head (same) |
| Memory for scores | T² per head, in HBM | one tile, in SRAM |
| HBM traffic | ≈ 3·T² reads/writes | ≈ T·d, read Q, K, V once |
| Bound by | memory | compute (on H100, about 75% of peak in FA3 public) |
| Max context on one GPU | a few thousand tokens | limited by the KV cache, not by the scores |
The arithmetic did not shrink. Attention is still quadratic in T in work, and Chapter 12's sparse and linear variants exist because of that. But the memory went from quadratic to linear, and the traffic through the pantry door fell by orders of magnitude, so the cores are finally the limit. Every frontier model since 2023 is trained and served with a FlashAttention-style kernel public, and the 128k and million-token contexts advertised today would be impossible without it.
Beacon at 128k tokens, one layer, forward.
full score matrix, one head T² × 2 bytes = 128,000² × 2 ≈ 32.8 GB × 128 heads ≈ 4.2 TB per layer (the H100 has 80 GB) FlashAttention working set, one tile (Q tile 128×128 + K tile + V tile + score tile) × 2 bytes ≈ 128 KB (SRAM has 228 KB) reduction in resident memory ≈ 250,000× arithmetic, unchanged 4 × T² × d = 4 × 128,000² × 128 ≈ 8.4 TFLOP per head per layer
$ python code/ch10/flash_memory.pyfull score matrix, one head, one layer : 32.8 GB × 128 heads : 4.2 TB per layer FlashAttention working set, one tile : 128.0 KB (fits in 228 KB of SRAM per SM) reduction factor : 250000× arithmetic per head, unchanged : 8.39 TFLOP (the T² work stays; the T² memory goes)
Four terabytes per layer is the size of the thing that must never exist. And 8.4 TFLOP per head per layer, times 128 heads, times 126 layers, is 135 PFLOP of attention arithmetic for one forward pass over one 128k-token sequence: about two minutes of an H100 at peak. That is why long-context training is expensive even with the memory problem solved, and why Chapter 12 looks at attention variants that cut the work itself.
Put it together over a whole run. Chapter 5 gave the arithmetic a training step must do: about 6·N FLOPs per token. Multiply by tokens per second achieved and divide by the cluster's peak, and you have model FLOPs utilisation: the fraction of the time the tensor cores were doing the model's matmuls rather than waiting for bytes, for other GPUs, or for the CPU.
MFU = (6 × N × tokens per second) / (number of GPUs × peak FLOP/s per GPU)
Llama 3 405B, 16,000 H100s: reported 38–43% MFU public
⇒ 16,000 × 989 TFLOP/s × 0.40 ≈ 6.3 × 10¹⁸ useful FLOP/s
⇒ 3.8 × 10²⁵ FLOPs ÷ 6.3 × 10¹⁸ ≈ 70 days of pure training time
Forty percent is a good number. The other sixty percent goes to attention's non-matmul work, to elementwise operations that were not fully fused, to communication between GPUs (Chapter 11), to pipeline gaps, and to the occasional restart (Chapter 13). Small models on one GPU reach 50–60%; frontier runs across thousands of GPUs rarely exceed 45% public across the reports that state it. When someone quotes a cluster's peak in exaFLOP/s, divide by two and a half before planning anything.
| Quantity | Value | Evidence |
|---|---|---|
| Training hardware, Llama 3 405B | 16,000 H100 GPUs, 8 per node, NVLink within, RoCE Ethernet between | public |
| Reported MFU | 38–43% | public |
| Attention kernel | FlashAttention-style, fused | public for open models; universal inferred |
| Weights resident in HBM at serving, bf16 | 810 GB across ≥ 5 GPUs, in practice a full 8-GPU node | computed, Chapter 4 |
| Closed frontier labs | Cluster sizes reported in the tens of thousands of GPUs; MFU, interconnect, and kernels not disclosed | unknown |
Run the model on a CPU. A server CPU manages a few TFLOP/s on matmul and a few hundred GB/s of bandwidth: both a hundred times worse than an H100, and the ratio between them is similar, so the same operations are bound the same way, only slower. Beacon's decode would produce a token every few seconds per conversation. Training would take centuries. GPUs are not a convenience; the arithmetic budget of Chapter 8 is unreachable without them.
Keep the GPUs but halve HBM bandwidth. Peak FLOP/s unchanged, ridge point doubles to 590 FLOP/B. Prefill matmuls at 4,600 FLOP/B barely notice. Decode, norms, softmax, and unfused attention all take twice as long. Serving throughput halves; training MFU falls by whatever fraction of step time was memory-bound. Bandwidth, not FLOP/s, is the spec that decides most of what a GPU can do for a language model, which is why HBM capacity and speed are what each generation advertises.
Materialise the attention matrix. Context length is capped by HBM: on an 80 GB GPU with 32 heads, a single layer's scores in bf16 exceed memory at about 11k tokens. Every context longer than that is impossible, and shorter ones spend most of their attention time moving the matrix through the door three times. This is the world before 2022.
Run every elementwise operation as its own kernel. Each RMSNorm, activation, residual add, and rotary rotation becomes a separate read-and-write of the full activation tensor. For a block that has roughly ten such operations around two matmuls, the elementwise traffic exceeds the matmul traffic and the block runs at a fraction of peak. Fusing them into the matmul epilogues is the difference between 20% and 40% MFU.
Put the batch's sequences on different nodes and attend across them. Attention needs every key and value for a sequence on the GPU computing that sequence's queries. Split a sequence across the 50 GB/s link and each layer's attention waits for keys to arrive at one-seventieth of HBM speed. Chapter 11's parallelism schemes are designed precisely so that the heavy traffic stays inside the node.
Say it back. A GPU is thousands of simple cores running the same instruction on different data, with tensor cores that multiply small matrix blocks in one step; it is fast at matmul and at nothing else. Its data sits in a hierarchy: registers and SRAM on the chip, tiny and fast; HBM beside it, large and slower; other GPUs over NVLink; other nodes over a network seventy times slower than HBM. Arithmetic intensity, FLOPs per byte moved from HBM, decides whether an operation is limited by the cores or by the door, and the H100's ridge is about 300 FLOP per byte. Big-batch matmuls clear it; decode matmuls, norms, softmax, and attention as literally written do not. Attention's problem is that its T² score matrix must be written and read three times; FlashAttention removes the matrix by computing the softmax incrementally with a running max, sum, and output, folding in one tile of keys and values at a time inside SRAM, leaving the arithmetic quadratic but the memory linear. Over a run, the fraction of peak the model's matmuls actually achieve is the MFU, about 40% for the best public frontier runs. Beacon trained on 16,000 H100s at that utilisation, and its 128k-token attention would need four terabytes per layer if anyone tried to store it.
[1.0, 3.0, 2.0, 3.5] with values [1, 2, 3, 4] (scalars), processing one score at a time. Track m, ℓ, o after each step, including the rescale factor. Confirm o/ℓ at the end equals the ordinary softmax-weighted average computed directly.[M, 8192] × [8192, 28672], what batch size M makes it compute-bound on the B200, and on the H100? What does that imply about the minimum batch a serving system must gather to use the newer GPU fully?roofline.py with a function that takes a model's (C, H, L, V) and a batch of M tokens and reports total bytes and FLOPs for one forward pass through all matmuls (attention projections, MLP, unembedding), then the overall intensity. Run it for Llama 3 8B at M = 1, 64, 4096. At which M does the whole forward pass cross the H100's ridge?