FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Abstract
Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3$\times$ speedup on GPT-2 (seq. length 1K), and 2.4$\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
FlashAttention introduces an I/O-aware tiling algorithm that computes exact self-attention without materializing the full attention matrix in GPU memory. By rigorously optimizing data movement between HBM and SRAM, the method achieves substantial training speedups, reduces memory footprint, and unlocks previously infeasible long-context capabilities, fundamentally reshaping how transformer models are implemented and scaled across the machine learning ecosystem.
Methodology
The paper introduces a paradigm shift in attention implementation by prioritizing I/O complexity over raw FLOP counts. Rather than materializing the full $N \times N$ attention matrix in High Bandwidth Memory (HBM), the authors employ a tiling strategy that partitions queries, keys, and values into blocks that fit within on-chip SRAM. By leveraging a numerically stable online softmax formulation, the algorithm incrementally computes attention scores and accumulates outputs without ever storing intermediate attention weights. The theoretical analysis rigorously bounds HBM accesses at $O(N^2d/M)$, demonstrating asymptotic optimality for practical SRAM sizes. The methodology is deeply hardware-aware, carefully orchestrating memory loads, compute kernels, and register usage to maximize GPU occupancy while minimizing memory bandwidth bottlenecks. The extension to block-sparse attention further demonstrates the flexibility of the tiling framework, enabling approximate attention with significantly lower overhead than prior sparse methods.
Results & evaluation
The empirical evaluation is comprehensive and strategically chosen to highlight both speed and capability gains. Benchmarks span standard NLP workloads (BERT-large, GPT-2), long-range sequence modeling (Long Range Arena), and extreme-context challenges (Path-X, Path-256). The paper reports meaningful wall-clock speedups (15% on BERT, 3× on GPT-2, 2.4× on LRA) alongside substantial memory reductions that enable longer sequence training without gradient checkpointing. Crucially, the experiments demonstrate that exact attention, when made computationally feasible, outperforms approximate alternatives in both perplexity and downstream accuracy. The inclusion of challenging pathfinding tasks provides compelling evidence that the algorithm unlocks previously intractable sequence lengths, validating the core hypothesis that memory efficiency directly translates to model capability.
Reproducibility
High. The algorithm is mathematically exact and deterministic, eliminating the stochastic variance often associated with approximate attention methods. The authors provide detailed CUDA implementation strategies, including kernel fusion, register allocation, and tiling parameter selection. The open-source release (initially built on NVIDIA Apex, later standalone) has been extensively adopted and integrated into major frameworks (PyTorch, HuggingFace Transformers, DeepSpeed), confirming its robustness and ease of reproduction across diverse hardware configurations and training pipelines.
Limitations
The initial implementation is heavily optimized for NVIDIA GPUs via CUDA, limiting immediate portability to AMD or TPU ecosystems. Tiling introduces minor overhead for very short sequences where standard attention is already memory-bound. The block-sparse variant requires heuristic mask selection, which may not generalize optimally to all data distributions. Additionally, the original formulation focuses primarily on training throughput rather than autoregressive inference latency, a gap later addressed in subsequent iterations (FlashAttention-2/3). Support for certain attention variants (e.g., cross-attention with highly mismatched sequence lengths) required additional engineering post-publication.
Broader impact
This work fundamentally democratizes long-context transformer training, drastically reducing compute costs, energy consumption, and hardware barriers for academic and industrial researchers. By shifting the community's optimization focus from FLOPs to memory bandwidth, it has catalyzed a new wave of I/O-aware ML systems research. The algorithm enables novel applications in genomics, long-form video understanding, and document-level reasoning, while its widespread adoption has become a de facto standard in modern LLM training stacks, directly accelerating the pace of foundation model development. FlashAttention introduces an I/O-aware tiling algorithm that computes exact self-attention without materializing the full attention matrix in GPU memory. By rigorously optimizing data movement between HBM and SRAM, the method achieves substantial training speedups, reduces memory footprint, and unlocks previously infeasible long-context capabilities, fundamentally reshaping how transformer models are implemented and scaled across the machine learning ecosystem.
Curator note
Dao et al.; 2-4x speedup; enabled longer contexts; universally adopted
Discussion
Talks & videos
Related later work
Author of this paper? Send a correction.