FlashAttention: Exact Attention with Less GPU Memory

The bottleneck was never the arithmetic

The paper's starting observation is about hardware, not models. On an A100, high bandwidth memory offers 40 to 80GB at 1.5 to 2.0TB per second, while on-chip SRAM offers 192KB per streaming multiprocessor at around 19TB per second. On-chip memory is an order of magnitude faster and many orders of magnitude smaller, and compute speed has outpaced memory speed, so most operations in a Transformer are limited by memory traffic rather than by arithmetic. [1]

That diagnosis explains a failure the field had been living with. A generation of approximate attention methods reduced the arithmetic to linear or near-linear in sequence length and mostly failed to deliver wall-clock speedups, because they optimized a quantity that was not the constraint. The paper's term for what was missing is IO-awareness: accounting for reads and writes between levels of memory, a principle it borrows from database joins, image processing and numerical linear algebra. [1]

Tiling, and paying arithmetic to avoid memory

The goal is to never write the large attention matrix to high bandwidth memory at all. Two techniques get there. Tiling splits the inputs into blocks and performs the softmax reduction incrementally across several passes, so the reduction never needs the whole input at once. Recomputation stores only the softmax normalization factor from the forward pass and recomputes attention on-chip during the backward pass, which is faster than reading the stored intermediate matrix back. The whole thing is fused into a single CUDA kernel. [1]

Atlas interpretation: The trade is the interesting part: this does more arithmetic, not less. Recomputation is redundant work, added deliberately because the memory traffic it avoids costs more than the FLOPs it spends. The paper also proves a lower bound: no exact attention algorithm can beat its memory-access count across the whole range of SRAM sizes it considers. That is a claim about the problem rather than about one kernel, though the authors note it does not rule out an algorithm that wins at a particular SRAM size. [1]

What it measured

BERT-large reaching the MLPerf target accuracy took 17.4 minutes against the MLPerf 1.1 record's 20.0, a 15 percent end-to-end gain. GPT-2 small on OpenWebText took 2.7 days against HuggingFace's 9.5 and Megatron-LM's 4.7, at identical perplexity. The attention operation alone runs up to 7.6 times faster. On long-range arena the speedup reaches 2.4 times. [1]

Memory is where the asymptotic change happens: the footprint becomes linear in sequence length, up to 20 times more efficient than exact baselines. Runtime does not. The paper states plainly that runtime still grows quadratically with sequence length. This is a large constant-factor win on time and a complexity change only on memory, a distinction that is routinely lost in summaries. [1]

The same paper also ships a second, genuinely approximate algorithm, block-sparse FlashAttention. Several of the most-quoted results belong to it rather than to the exact version, including Path-256 at 63.1 percent accuracy, the 64K sequence length, and a further 2 to 4 times speedup. A claim that this work scaled Transformers to 64K tokens is, strictly, a claim about the approximate variant. [1]

Why the word "exact" is in the title

Atlas interpretation: The title advertises exactness because the obvious way to read this work is as one more approximate attention method, and it is not one. It computes the same function as standard attention. The paper's evidence is that GPT-2 reaches identical perplexity and identical training curves, since the model definition is unchanged. Everything it saves is in how the computation is scheduled against the memory hierarchy. [1]

The correction has needed repeating. The FlashAttention-2 paper restates that the speedups come with no approximation, and Tri Dao's own write-up of the third version says the same. Exact here means no algorithmic approximation, not bitwise equality: PyTorch's own documentation warns that fusing floating point operations means output can differ depending on which backend kernel runs. [2][4][5]

What it became

It is now a default rather than an option. PyTorch lists it as one of the backends behind scaled_dot_product_attention, with FlashAttention-2 integrated in PyTorch 2.2 for roughly a further doubling. The successor papers came from partly different teams: FlashAttention-2 in July 2023 is a single-author paper reporting around a 2 times gain and 50 to 73 percent of theoretical maximum FLOPs on A100, and FlashAttention-3 in July 2024 is led by Jay Shah, targeting Hopper asynchrony and low precision for a further 1.5 to 2 times. [5][2][3]

Atlas interpretation: The link most often drawn, that this work produced the long context windows that arrived in 2023, is real but softer than usually stated. FlashAttention-2's introduction makes the connection itself, listing 32K, 65K and 100K context models and attributing the enabling work to this line. It is the standard attention kernel underneath training and serving at those lengths. It is not the only contributor: MosaicML's MPT-7B, one of those long-context models, lists this kernel among the optimizations behind its training and inference speed while crediting ALiBi, a positional encoding method, for its ability to handle inputs longer than it was trained on. [2][6]