FlashAttention: The IO-Aware Algorithm That Rewrote GPU Training Speed Records

The attention mechanism at the core of every modern Transformer is quadratic in sequence length. Standard self-attention for N tokens takes O(N^2) time and allocates O(N^2) memory, which becomes a practical wall around 2,048 tokens for most GPU configurations. Before FlashAttention, the standard answer to that wall was approximate attention: sparse patterns, low-rank decompositions, linear approximations. Those methods reduce FLOP counts substantially, yet many failed to deliver measurable wall-clock improvements. The reason, laid out in 'FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness' by Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Re (arXiv:2205.14135, submitted May 2022), is that FLOPs are not the bottleneck. Memory bandwidth is.

Why attention is memory-bound, not compute-bound

Attention decomposes into a sequence of operations: query-key matrix multiply, scaling, masking, softmax, dropout, then value multiply. The matrix multiply steps are compute-bound, meaning the GPU's arithmetic units are the limiting factor. Softmax, masking, and dropout are elementwise or reduction operations with low arithmetic intensity: the GPU runs out of data to process long before it runs out of compute capacity. Because those elementwise ops dominate attention's runtime despite accounting for only a small fraction of its total FLOPs, cutting FLOPs in the matrix multiply steps does little for wall-clock time. Approximate attention methods that sparsify or decompose the score matrix still read and write HBM for every softmax and dropout element. The memory bottleneck remains even when FLOPs drop significantly.

HBM and SRAM: the two-tier memory problem

Modern data-center GPUs have two relevant memory tiers. High Bandwidth Memory (HBM) is the large off-chip pool: an NVIDIA A100 80GB carries 80 GB of HBM at roughly 2 TB/s bandwidth. On-chip SRAM is the small fast pool inside each streaming multiprocessor (SM): an A100 SM has 192 KB of shared memory, and across its 108 SMs the total on-chip capacity is around 20 MB at bandwidth approaching 19 TB/s per chip. Standard attention writes the full N-by-N attention matrix to HBM, reads it back for softmax, writes it again after dropout, and reads it once more to compute the output. Each of those round trips crosses the slow HBM bus. For a sequence of 2,048 tokens in a 12-head model, the attention score matrix alone is roughly 48 MB in float32, written and read multiple times per forward pass. This is where the time goes.

Tiling: keeping computation inside SRAM

FlashAttention avoids materializing the full NxN score matrix in HBM by tiling the computation. The key insight is that softmax can be computed incrementally using the log-sum-exp trick, which allows the algorithm to process the attention matrix in blocks that fit entirely within SRAM. Each tile loads a block of queries, keys, and values from HBM into SRAM, computes a partial attention output alongside a running softmax normalizer, then updates the output accumulator before loading the next block. The final output never requires storing the full score matrix anywhere. The algorithm is mathematically identical to standard attention: it produces the exact same outputs, not an approximation. The paper's IO complexity analysis shows FlashAttention requires O(N^2 / M) HBM accesses, where M is SRAM size, versus O(N^2) for standard attention. For typical SRAM values (tens of kilobytes per SM), that is roughly an order-of-magnitude reduction in HBM traffic.

Benchmark results: speed records without any approximation

The benchmark numbers from the paper are striking precisely because no quality tradeoff is involved. Every result is exact attention, not an approximation.

• BERT-large training (sequence length 512): 15% faster end-to-end than the MLPerf 1.1 training speed record at the time of publication.
• GPT-2 (sequence length 1,024): 3x faster than baseline implementations from HuggingFace and Megatron-LM.
• Long-range arena (sequence length 1,024 to 4,096): 2.4x faster than baselines.
• GPT-2 perplexity: 0.7 improvement from training on longer sequences within the same compute budget.
• Long-document classification: 6.4 percentage-point improvement.
• Path-X (sequence length 16,384): the first Transformer to achieve better-than-chance performance, reaching 61.4% accuracy.
• Path-256 (sequence length 65,536): 63.1% accuracy, a task no prior exact attention implementation could even run at scale.

The block-sparse extension of FlashAttention, which applies the same tiling technique to sparse attention patterns, also outperformed every existing approximate attention method on speed benchmarks.

FlashAttention-2: closing the gap to peak hardware throughput

Dao followed up with FlashAttention-2 (arXiv:2307.08691, July 2023), motivated by a concrete measurement: FlashAttention-1 reached only 25 to 40% of the A100's theoretical maximum FLOPs/s. Profiling revealed the cause: suboptimal work partitioning between CUDA thread blocks and warps, leading to either low GPU occupancy or unnecessary shared-memory reads and writes. FlashAttention-2 addressed this with three targeted changes. First, it restructures the online softmax rescaling step to cut non-matmul FLOPs. Second, it parallelizes attention computation across thread blocks even within a single attention head, keeping occupancy high even for short sequences. Third, it redistributes work between warps inside each thread block to minimize shared-memory communication overhead. Together, these push hardware utilization to 50 to 73% of theoretical peak FLOPs/s on A100, roughly a 2x improvement over FlashAttention-1. In end-to-end GPT-style model training, FlashAttention-2 reaches up to 225 TFLOPs/s per A100 GPU, corresponding to 72% model FLOPs utilization.

Adoption and the broader lesson

FlashAttention became a de facto infrastructure standard with unusual speed. Llama 2 and most subsequent Meta open-weight models train with it. Mistral AI built its sliding-window attention on top of FlashAttention kernels. The HuggingFace Transformers library, Megatron-LM, and virtually every serious long-context training stack adopted it within months of the original release. A third version targeting H100 GPUs (with FP8 support and asynchronous warp specialization) followed as H100 deployments scaled up. The GitHub repository at github.com/Dao-AILab/flash-attention has become one of the most-cited infrastructure components in LLM development. The broader lesson extends well beyond this one algorithm: the bottleneck in deep learning is increasingly memory bandwidth, not arithmetic throughput, and kernel-level IO optimization can unlock capabilities that no amount of approximate algorithmic cleverness achieves at the same quality level. Scaling to 16K or 64K context windows became practically feasible not because of a new model architecture, but because of careful accounting of reads and writes between two levels of GPU memory. PromptingIndex covers FlashAttention alongside PagedAttention, speculative decoding, and grouped-query attention as part of its series on LLM inference and training efficiency.