An exact implementation of attention that is faster because it never writes the full n × n attention matrix to GPU memory. It computes attention in tiles that fit in fast on-chip memory, using an online softmax, so it moves far less data and its memory use grows linearly with sequence length.
Card 1 of 292. Answer: An exact implementation of attention that is faster because it never writes the full n × n attention matrix to GPU memory. It computes attention in tiles that fit in fast on-chip memory, using an online softmax, so it moves far less data and its memory use grows linearly with sequence length.
About · Quizzes · History of AI · To suggest new questions or report an error send me a dm.
Standard attention writes the score matrix QKᵀ and its softmax to the GPU's main memory and reads them back, and those transfers, not the arithmetic, dominate the time. FlashAttention, by Tri Dao and colleagues (2022), fuses the steps into one kernel: it loads blocks of queries, keys and values into on-chip SRAM, keeps a running maximum and sum so the softmax can be rescaled block by block, and in the backward pass recomputes attention instead of storing it.
Example: at 8,192 tokens, one head's score matrix has 8,192² ≈ 67 million entries, 134 MB at 16-bit precision, and with 32 heads that is 4.3 GB for a single layer. FlashAttention never stores it. The paper reported GPT-2 training about 3 times faster than the standard implementation.
Interview points: the result is exact, not an approximation like sparse or linear attention; it does slightly more arithmetic, because of recomputation, yet runs faster, which is the lesson of IO-aware algorithms; and compute is still quadratic in sequence length, only memory becomes linear. FlashAttention-2 (2023) improved parallelism, and FlashAttention-3 (2024) targets H100 GPUs, including 8-bit precision.
