Flash Attention

Flash Attention is an IO-aware, exact attention algorithm that computes scaled dot-product attention using GPU on-chip SRAM tiling to dramatically reduce high-bandwidth memory access. It achieves 2–4x faster training than naive implementations while producing numerically identical results, without any approximation.

Flash Attention is an optimized algorithm for computing the attention mechanism at the core of transformer models. Introduced by Tri Dao et al. from Stanford in a 2022 paper, it produces results numerically identical to standard attention but rearranges computation to minimize data movement between GPU high-bandwidth memory (HBM) and fast on-chip SRAM. It is classified as an IO-aware algorithm—its design is explicitly motivated by the memory bandwidth bottleneck of modern GPU hardware rather than by reducing the arithmetic complexity of the attention operation itself.

Standard attention materializes the full N×N attention matrix in HBM, where N is the sequence length, resulting in O(N²) memory usage. Flash Attention uses a tiling strategy: it processes attention in blocks small enough to fit in SRAM, applying the numerically stable online softmax algorithm to accumulate partial results incrementally. The full attention matrix is never written to HBM. Flash Attention 2 (2023) improved parallelism across the sequence dimension and reduced non-matmul floating-point operations. Flash Attention 3 (2024) introduced asynchronous warp specialization and FP8 precision support targeting NVIDIA H100 and Hopper-class GPUs, pushing throughput closer to hardware theoretical peaks.

Attention is the dominant memory and compute bottleneck in transformer training and inference at long sequence lengths. Flash Attention reduces memory from O(N²) to O(N) and delivers 2–4x wall-clock speedups over unoptimized implementations, translating directly into lower training costs or more experimental iterations per compute budget. These gains compound at longer sequence lengths, making Flash Attention essential infrastructure for practical long-context models.

As of 2026, Flash Attention is integrated natively into PyTorch via torch.nn.functional.scaled_dot_product_attention, Hugging Face Transformers, vLLM, and virtually every major LLM training and inference framework. It is a training dependency for LLaMA 3, Mistral, and most other frontier models. Ring attention, which distributes Flash Attention tiling across multiple devices via all-to-all communication, extends its benefits to sequences that exceed single-GPU memory, enabling the million-token context windows now offered commercially.