Flash Attention
Flash attention computes exact attention faster by tiling the calculation to keep intermediate results in fast on-chip memory, never materializing the full score matrix.
The memory bottleneck
Standard attention computes a score matrix of size sequence-length squared, applies a softmax, and multiplies by the values. On a modern accelerator the arithmetic is cheap relative to moving that large matrix between high-bandwidth memory and the fast on-chip memory near the compute units. Attention is therefore memory-bound: the bottleneck is reading and writing the score matrix, not the multiplications themselves. Flash attention removes this bottleneck without changing the result.
Tiling and the online softmax
Flash attention splits the queries, keys, and values into blocks and processes them in tiles that fit in on-chip memory. For each query block it walks over key blocks, accumulating the output incrementally. The challenge is the softmax, which normally needs the whole row to compute its denominator. Flash attention uses an online softmax that maintains a running maximum and running normalization as it sees each block, rescaling the partial output so the final result is exactly the standard softmax attention.
- Exact attention, not an approximation, with identical numerical outputs up to floating-point order
- Memory use grows linearly, not quadratically, in sequence length
- The large score matrix is never written to slow memory
- Fewer memory reads and writes yield large wall-clock speedups
# conceptual online softmax accumulation per key block
m_new = max(m, block_max)
l = exp(m - m_new)*l + exp(block_max - m_new)*block_sum
acc = exp(m - m_new)*acc + exp(block_scores - m_new) @ V_block
m = m_newImpact
Flash attention made longer context lengths practical by cutting both the memory footprint and the time of the attention layer, and it also reduces memory in the backward pass by recomputing the scores from stored statistics instead of storing the whole matrix. It is a systems-level optimization: the mathematics of attention is unchanged, but the schedule of computation respects the memory hierarchy. Where an approximation is acceptable, sparse and linear attention reduce cost further by changing what is computed.