Why can FlashAttention be faster with the same formula?

The bottleneck

Naive attention may materialize and repeatedly read/write T×T scores and probabilities in HBM. FlashAttention tiles the operation and performs local work in on-chip SRAM, avoiding full intermediate matrices in device memory.

The mechanism

Online softmax maintains each row’s maximum m, normalization sum l and weighted output. Each new tile rescales previous accumulators to the updated maximum and merges its contribution. Backward can recompute intermediates, trading compute for IO. The mathematical attention is unchanged, not necessarily bitwise floating-point results.

m′ = max(m, max(S_tile)); l′ = e^(m−m′)l + Σe^(S_tile−m′)

Examples and evidence

The sketch moves Q/K/V tiles from HBM into SRAM, writing required outputs and statistics. Avoiding full T² intermediate storage does not eliminate T² dense interactions. Report sequence length, head dimension, precision and hardware when measuring gains.

New constraints

On-chip capacity, thread layout and kernel support define applicability. Small sequences or special masks may gain less. FlashAttention optimizes execution; it does not replace long-context evaluation or KV allocation management.