FlashAttention

FlashAttention is a way of computing attention that produces the exact same answer while never materialising the giant token-by-token score matrix in slow GPU memory.

The insight is about where work happens, and kitchens explain it. A cook with a tiny counter next to the stove and a big pantry across the room can work two ways. The slow way: carry every ingredient combination to the pantry and back as it is produced. The fast way: bring one batch of ingredients to the counter, do all the work on that batch there, and only walk finished dishes to the pantry. Same recipe, same dishes — far less walking.

On a GPU, the pantry is main memory (HBM, large but relatively slow) and the counter is on-chip SRAM (tiny but very fast). Standard attention writes the full n×n score matrix to the pantry — for 32,000 tokens, a billion-entry matrix — then reads it back to finish. FlashAttention tiles the computation: it processes blocks of the matrix entirely on the counter, using a running-softmax trick so blocks combine correctly, and the n×n matrix never exists in memory at all.

The payoff: memory grows linearly instead of quadratically with sequence length, and wall-clock speed improves several-fold, because attention was bottlenecked on memory traffic, not arithmetic. This is not an approximation — the output is exact. Published by Tri Dao et al. in 2022, it is now baked into PyTorch, vLLM and every serious serving stack, and it is a large part of why 100k-token context-windows became practical.

Where to go next#

• Full lesson: vLLM
• Related terms: attention, gpu, kv-cache, context-window