FlashAttention

Standard attention implementations are memory-bound, spending most of their time reading and writing large intermediate matrices to GPU High Bandwidth Memory (HBM). FlashAttention is IO-aware, meaning it explicitly accounts for the speed hierarchy between slow HBM and fast on-chip SRAM. By minimizing HBM reads and writes, it shifts the bottleneck from memory bandwidth to compute utilization, achieving wall-clock speedups of 2-4x over optimized baselines like PyTorch's scaled_dot_product_attention.

To avoid materializing the full N x N attention matrix in HBM, FlashAttention uses tiling. The algorithm splits the input sequences (Q, K, V) into blocks that fit entirely within SRAM. Attention is computed incrementally block-by-block, with intermediate softmax statistics carefully tracked and rescaled. This allows exact attention computation without ever storing the full quadratic intermediate matrix, reducing peak memory from O(N²) to O(N).

In the backward pass, standard attention saves the large attention matrix from the forward pass for gradient calculation. FlashAttention discards this matrix and instead recomputes it on-the-fly during the backward pass using the stored softmax normalization statistics. This trades a modest increase in FLOPs for a dramatic reduction in memory footprint, enabling training with 4-8x longer sequences on the same hardware without running out of memory.

Unlike sparse or low-rank attention approximations (e.g., Longformer, Linformer), FlashAttention computes mathematically exact self-attention. The output is numerically identical to standard attention up to floating-point rounding errors. This means there is no accuracy trade-off—models trained with FlashAttention achieve identical perplexity and downstream task performance while benefiting from the speed and memory advantages.

The second iteration, FlashAttention-2, significantly improves GPU utilization by redesigning the work partitioning strategy. It reduces non-matmul FLOPs, parallelizes over the sequence length dimension (previously parallelized only over batch and heads), and distributes work more evenly across thread blocks. This achieves 2x speedup over the original FlashAttention and pushes GPU utilization to near theoretical limits, reaching up to 73% of peak FLOPs on A100 GPUs.

FlashAttention-3 targets NVIDIA's Hopper GPU architecture (H100), exploiting new hardware features like WGMMA instructions for asynchronous Tensor Core operations and TMA (Tensor Memory Accelerator) for hardware-accelerated data movement. It introduces interleaved block-wise matmul and softmax scheduling to hide memory latency. On H100 GPUs, FlashAttention-3 achieves 1.5-2.0x speedup over FlashAttention-2 and reaches up to 740 TFLOPS, representing 75% of theoretical peak performance.