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× 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).

DarcStar commentary

Our perspective on this work and how it informs our research — not part of the original paper.

If we had to point at the single idea our long-context work is built on, it's this one. FlashAttention's move looks almost too simple in hindsight: attention wasn't slow because of arithmetic, it was slow because of memory traffic — shuttling a giant N×N score matrix to and from the GPU's big, far memory. Stop materializing that matrix, tile the computation so it stays in fast on-chip memory, and combine the pieces with an online softmax that never needs a whole row at once. Same answer, exactly — no approximation — for a fraction of the memory movement.

That last clause is the one people undersell. For years the field chased approximate attention to cut FLOPs, and mostly traded quality for a speedup that didn't show up on the clock. FlashAttention changed the question from “how do we approximate this cheaply” to “how do we compute it exactly without moving so much memory” — and won. That reframing is the ground we stand on: the load-bearing primitive under every attention variant we care about isn't “attention,” it's this exact, associative online softmax over pieces of the key space. Get it right once and sparse, dilated, and ring attention all become configurations of it.

So we don't file this as a speedup — we file it as the enabler. Without an IO-aware exact kernel, long context is a memory-bandwidth non-starter before any of the clever scaling tricks get a turn. This is the paper that made the rest of the shelf possible.