Peer reviewed
FlashAttention (1/2/3): IO-Aware Exact Attention
The kernel splits queries, keys, and values into blocks that fit in fast on-chip SRAM, computes softmax incrementally with running statistics, and recomputes intermediates during the backward pass instead of storing them, so memory scales linearly rather than quadratically in sequence length. This yields wall-clock speedups and lower memory use without approximation, giving identical results to standard attention. It enabled training and serving with longer context windows and became a default attention implementation in mainstream frameworks.