arXiv. arXiv:2205.14135 (Submitted 27 May 2022, last revised 23 Jun 2022, v2)
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× speedup on GPT-2 (seq. length 1K), and 2.4× 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).
Summary
The paper keeps attention mathematically exact and changes where intermediate tensors live. Standard attention writes an N × N score matrix to HBM. FlashAttention tiles Q, K, and V into SRAM, maintains an online softmax, and never materializes that matrix. Extra arithmetic in the backward pass is accepted because HBM traffic is more expensive. The 2022 measurements report faster BERT-large, GPT-2, and Long Range Arena training, plus the first Transformer-style results above chance on Path-X (16K) and Path-256 (64K). Arithmetic remains O(N²). Later fused kernels will not automatically reproduce the original wall-clock multiples.
Detailed reading
IO-awareness versus FLOP reduction
The abstract locates the difficulty in movement between GPU memory levels rather than in attention arithmetic alone. Standard attention materializes S = QK^T, applies a row-wise softmax, then computes O = PV. For length N the intermediates contain N² entries. Quadratic arithmetic is real. On an A100, however, HBM bandwidth is about 1.5–2.0 TB/s, while on-chip SRAM is an order of magnitude faster and only about 192 KB per streaming multiprocessor. Attention can therefore be IO-bound even when its FLOP count appears acceptable.
Although these methods reduce the compute requirements to linear or near-linear in sequence length, many of them do not display wall-clock speedup against standard attention and have not gained wide adoption.
Approximate methods that lower theoretical multiplies may still lose on wall-clock time if they introduce extra kernels, irregular reads, or materialized intermediates. The paper treats wall-clock time, not FLOP reduction alone, as the relevant criterion.
Tiling and exact softmax
(i) We restructure the attention computation to split the input into blocks and make several passes over input blocks, thus incrementally performing the softmax reduction (also known as tiling).
Q, K, and V are partitioned. A query block and successive key/value blocks are loaded from HBM into SRAM. The kernel updates running output statistics and discards the tile. The full N × N attention matrix is never written to HBM.
A row-wise softmax depends on every key, so independent per-tile normalization would be incorrect. FlashAttention maintains an online softmax: a running maximum, a running sum of exponentials, and a partially weighted value accumulator. When a later tile raises the maximum, earlier accumulators are rescaled. After all tiles are visited, the result matches standard softmax apart from ordinary floating-point reordering. The algorithm is exact, not sparse and not low-rank.
Recomputation and IO complexity
In contrast, even with more FLOPs, our recomputation speeds up the backward pass due to reduced HBM accesses (Fig. 2).
A conventional backward pass stores the attention probability matrix. FlashAttention stores compact row statistics (m, ℓ) and rebuilds score tiles on chip. Additional arithmetic is accepted because HBM traffic is more expensive. Figure 2 of the paper shows higher FLOP count, fewer HBM bytes, and shorter runtime.
Standard attention (Algorithm 0) requires Θ(Nd+N²) HBM accesses, while FlashAttention (Algorithm 1) requires Θ(N²d²M^{-1}) HBM accesses.
With SRAM size M and typical head dimensions, d² is much smaller than M, so traffic falls substantially. The paper also argues that no exact algorithm can asymptotically improve on this bound for every SRAM size. Arithmetic remains O(N²); locality changes. Complexity tables that count only FLOPs omit this distinction.
Reported measurements
We train BERT-large (seq. length 512) 15% faster than the training speed record in MLPerf 1.1, GPT2 (seq. length 1K) 3× faster than baseline implementations from HuggingFace and Megatron-LM, and long-range arena (seq. length 1K-4K) 2.4× faster than baselines.
On eight A100 GPUs, BERT-large training time is reported as 17.4 ± 1.4 minutes versus 20.0 ± 1.5 for the cited MLPerf 1.1 record. GPT-2 medium is about 3× relative to HuggingFace and about 1.7–1.8× relative to Megatron-LM, with unchanged perplexity because the model definition is unchanged. GPT-2 small at context 4K is still faster than Megatron at 1K, with 0.7 better perplexity. Path-X (16K) and Path-256 (64K) are reported as the first Transformer-style results above chance on those Long Range Arena tasks.
Evidence boundary
These multiples belong to the 2022 hardware and software stack. Later fused kernels in modern frameworks will not automatically reproduce them. FlashAttention still evaluates dense pairs, so arithmetic remains quadratic. Sparse designs such as NSA change which pairs exist; the two approaches may be composed, but they answer different questions. The lasting design principles are: fuse operations so intermediates do not round-trip to HBM; tile to on-chip capacity; recompute when movement costs more than arithmetic; and report wall-clock time together with the memory hierarchy.