The film
About 5:07, hosted by Pylon. It asks you one question and waits for your answer before it shows you. The guide's voice is not recorded yet, so for now the film plays with captions; the full words are also below.
Nothing loads until you press play.
The lab: the cannery
Pick a kernel and sizes; the emulator counts every byte to and from memory and checks the answer against an exact one; then, if your browser has WebGPU, time the real kernels on your own chip. In the film, at 4,096 words and 16 numbers per head, the plain way moved 285,736,960 bytes and fused 17,301,504; at 64 per head, 538,968,064 against 270,532,608.
Counting runs in this page. Timing runs the kernels on your own graphics chip, only when you press Time. Nothing is sent anywhere.
Everything the film says
Fused attention
Fused attention.
Attention, the step where each word looks back at the words before it, is where the lighthouse meets the cannery. Tonight Pylon shows how to do it in one trip, without ever writing the whole table of scores to the warehouse.
The table of scores
For every word, attention scores that word and every earlier one, turns the scores into shares, and mixes the words by those shares.
Done the plain way on a graphics chip, that takes three passes, and the table crosses the warehouse door four times. Write the scores. Read them back and write the shares. Read the shares back to mix. For 4,096 words, one head's scores alone fill 64 mebibytes, in 32-bit numbers.
Shares without the whole row
The trouble is the shares. To turn a score into a share you divide by a total over the whole row, and you seem to need every score before you can start.
Maxim Milakov and Natalia Gimelshein showed in 2018 that the largest score and the total can be built in one pass, as the scores go by. When a bigger score arrives, shrink the total so far to match, and carry on.
On our model's own sentence, the tile-by-tile answer and the plain one both land within float32 rounding of the exact answer, on all four heads.
One trip
Tri Dao and colleagues built FlashAttention on this in 2022. Our kernel follows the same idea. A crew loads a tile of questions and walks tiles of keys and values past it on its bench. It keeps the running largest score and total, and writes only the finished answers. Dao's version loops the other way round, and also saves the running largest and total for training.
Count one head's warehouse traffic for 4,096 words, with 64 numbers per head: 514 mebibytes the plain way, 258 fused. The bench is small, so the keys and values are fetched again for every tile of questions. With 16 numbers per head, more fits on the bench: 272.5 against 16.5.
Four thousand words, sixteen numbers per head. How many times less traffic, fused? (the film pauses here)
Forget one step
The one step that is easy to forget: shrinking the running total when a bigger score arrives.
We left it out on purpose. On our model's sentence, some answers came out wrong on every one of the four heads, and the checker caught it. It is right only when the biggest score happens to arrive in the first tile.
Why keep the largest
The running largest score does a second job. Raising e to a score makes big numbers fast, and fp16, the common 16-bit format, overflows above 65,504.
The largest score in our sentence's last row is 10.32, and e to that is 30,320. That fits, but a score of 11.09 would not. Across the held-out rows, the largest score reaches 23.41, and e to some of those scores would be infinity in fp16. Subtracting the running largest first keeps every term at 1 or below.
What the count leaves out
These are counts of the memory requests our own kernels make, not times, and a chip's caches may serve some of them. Fused attention also does a little more arithmetic, the rescaling, to save the trips.
The lab runs both versions on your own chip, if your browser allows it, and times them beside the count.
Goodnight
Do the whole job on the bench, keep a running largest and a running total, and never build the wall of scores.
Next, with Tallow: what the 16-bit numbers in those trays can and cannot hold. Goodnight.
Sources
Every claim in the film, with where it comes from and what that source does not prove. The numbers about the model are computed from its weights, which you can load and check: facts.json.
Night School's attention emulator and kernels, built for this film. media/school/fused/facts.json, written by tools/school/fill_fused.mjs from system/school/fused*.mjs: online softmax and tiled attention checked against a float64 reference on random inputs and the attention model's own Q, K and V; WGSL kernels for the lab; an exact emulator of global memory traffic for unfused and fused attention at N = 64 to 4,096, d = 16 and 64; a broken-rescale control; the overflow arithmetic. https://github.com/HarperZ9/HarperZ9.github.io/tree/main/media/school/fused
Used for: N 4,096: d 64 unfused 514 MiB against fused 258 MiB; d 16, 272.5 against 16.5. The broken rescale fails on all four probe heads. Probe row largest score 10.32, exp 30,320 in fp16; on held-out rows the largest score is 23.41.
Does not prove: Any timing; caches and arithmetic cost are not modelled.
Dao and others 2022, FlashAttention. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., Re, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135. https://arxiv.org/abs/2205.14135
Used for: Tiling with the online softmax avoids writing the N x N attention matrix to high-bandwidth memory; Theorem 2 gives O(N^2 d^2 / M) memory accesses.
Does not prove: Anything about this emulator.
Milakov and Gimelshein 2018, Online normalizer calculation for softmax. Milakov, M., Gimelshein, N. (2018). Online normalizer calculation for softmax. arXiv:1805.02867. https://arxiv.org/abs/1805.02867
Used for: The softmax normaliser can be computed in one pass with a running maximum and a rescaled running sum.
Does not prove: Anything about this model.
What this lesson leaves out
Counts of bytes asked for, not times; caches can serve repeated reads, and the film's kernels were not timed on a graphics chip.
The check tolerance is a first-order float32 error bound derived for this lesson, not a published one.
The model computes in float32 and subtracts the largest score, so it never overflows; 870 of 307,200 held-out scores would overflow fp16 without that subtraction.
Before this: Matmul tiling (K22). Next, in the lamp room: Mixed precision (K24).