ML//transformer//attention//Flash Attention

Flash Attention is an implementation of attention that computes exactly the same result as the standard one without ever writing the full \(N \times N\) matrix to GPU memory, which solves the memory problem of attention over long sequences. The mathematics is unchanged; only the way the data moves is.

Flash Attention is an implementation of attention that computes exactly the same result as the standard one without ever writing the full N×NN \times NN×N matrix to GPU memory, which solves the memory problem of attention over long sequences. The mathematics is unchanged; only the way the data moves is.

Standard attention computes QKTQK^{\mathsf T}QKT, writes the full N×NN \times NN×N matrix to HBM, reads it back and multiplies it by VVV. For N=4096N = 4096N=4096 that is 16M values, and for N=32768N = 32768N=32768 it is 1B values, at which point the GPU runs out of memory. Flash Attention divides QQQ, KKK and VVV into small blocks that fit in SRAM (fast on-chip memory), computes attention block by block and accumulates the result.

Standard1Load QQQ, KKK, VVV complete2Write the N×NN \times NN×N matrix to RAM3Read it back4Write the output

Flash1Load a small block2Compute the partial result in SRAM3Move to the next blocknever writes the full matrix

Because the full matrix is never materialized in HBM, memory goes from O(N2)O(N^2)O(N2) to O(N)O(N)O(N) while the result stays exact. It enables longer context windows at the same hardware cost, and it is a key enabler of models with contexts of 100K tokens and more.