Flash Attention
Flash Attention is the IO-aware, exact attention algorithm that computes self-attention in O(N) memory (instead of O(N²)) by tiling the computation to exploit GPU SRAM (shared memory) locality, avoiding materialization of the full N×N attention matrix in HBM (high-bandwidth memory) — delivering 2-4× wall-clock speedup and enabling much longer sequence lengths.
Standard attention computes: Attention(Q,K,V) = softmax(QK^T / √d) · V. The naive implementation materializes the N×N attention score matrix S = QK^T in GPU HBM, reads it back for softmax, then multiplies by V. For sequence length N=8192 and batch×heads=32, this intermediate matrix is 32 × 8192 × 8192 × 2 bytes ≈ 4 GB — far exceeding GPU SRAM capacity and requiring expensive HBM round-trips.
The IO Bottleneck: Modern GPUs have ~20 MB of SRAM (shared memory per SM) with ~19 TB/s bandwidth, versus ~80 GB of HBM with ~3 TB/s bandwidth. Standard attention is memory-bandwidth-bound because it reads/writes the N×N matrix from/to slow HBM multiple times. Flash Attention restructures the computation to keep working data in fast SRAM.
Tiling Algorithm:
1. Divide Q into blocks of size Br × d, and K,V into blocks of size Bc × d (where Br, Bc fit in SRAM)
2. For each Q block, iterate over all K,V blocks
3. Compute block attention scores S_ij = Q_i · K_j^T in SRAM
4. Track running softmax statistics (row-max m and row-sum l) using the online softmax trick
5. Update the output block O_i incrementally: O_i += exp(S_ij - m_new) · V_j, adjusting for the running normalization
6. After all K,V blocks are processed, normalize O_i by the final softmax denominator
Online Softmax Trick: The key mathematical insight. Standard softmax requires two passes (find max, then compute exp/sum). Flash Attention maintains running max and sum across blocks, rescaling previous partial results when a new block produces a larger maximum. This enables single-pass softmax computation across tiled blocks.
Flash Attention 2 Improvements: Reduced non-matmul FLOPs by restructuring the algorithm to minimize rescaling operations; improved parallelism by distributing work across both the sequence and head dimensions; and optimized warp-level scheduling to reduce shared memory bank conflicts. Result: ~2× faster than Flash Attention 1, approaching theoretical peak FLOPS.
Flash Attention 3 (Hopper): Exploits NVIDIA H100 features: asynchronous GEMM via TMA (Tensor Memory Accelerator) for overlapping data loading with computation; FP8 support for quantized attention with 2× throughput; and warp specialization (producer-consumer warp groups) for better instruction-level parallelism.
Impact on Sequence Length: By reducing memory from O(N²) to O(N), Flash Attention makes training with 64K-1M+ token sequences practical. Context windows expanded from 2K (GPT-3) to 128K+ (Claude, GPT-4 Turbo) largely because Flash Attention removed the memory wall.
Flash Attention is perhaps the single most impactful systems optimization in modern deep learning — by recognizing that attention's bottleneck is memory bandwidth rather than computation, it unlocked longer contexts, faster training, and lower inference costs without any approximation or quality loss.