FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

The bottleneck was memory traffic, not math

Standard attention computes $S = QK^\top$, applies a softmax to get $P$, and multiplies $PV$. The naive implementation writes the full $N \times N$ score matrix $S$ out to GPU high-bandwidth memory (HBM), reads it back to apply the softmax, writes $P$, and reads it again for the value multiply. For a sequence of length $N$, that is several round trips of an $N \times N$ object through HBM. The crucial insight FlashAttention starts from is that on modern GPUs attention is bottlenecked by this memory traffic, not by the arithmetic. The matmuls are cheap relative to the cost of shuttling the giant intermediate matrix between the fast on-chip SRAM and the slow, capacious HBM. The GPU spends its time waiting on memory, and the $O(N^2)$ HBM reads and writes are where the wait comes from.

This reframes the goal. The win is not a cleverer approximation of attention — it is computing the exact same attention while moving far fewer bytes between HBM and the compute units. FlashAttention is IO-aware: it counts memory accesses, not just FLOPs, and minimizes them.

Tiling: compute attention without ever storing the scores

The method tiles the computation. $Q$, $K$, and $V$ are split into blocks small enough that a block's worth of the computation fits in SRAM. The algorithm loops over blocks of keys and values, and for each it loads the corresponding $Q$ block, computes that tile of the score matrix inside SRAM, and immediately folds it into a running output — all without ever writing the score tile back to HBM. The full $N \times N$ matrix $S$ is never materialized anywhere in slow memory; only the small tiles transit SRAM, and only the final output and a little bookkeeping land in HBM. Memory consumption drops from $O(N^2)$ to $O(N)$, which on its own makes long sequences feasible that would otherwise exhaust GPU memory just holding the scores.

The trick that makes tiling work: online softmax

Tiling attention has a catch. The softmax normalizes across an entire row of the score matrix — it needs the sum of exponentials over all keys — but tiling processes the keys a block at a time, so you never have a whole row in hand at once. FlashAttention resolves this with online softmax: it maintains, for each query, a running maximum and a running sum of exponentials, and rescales the partial output as each new key block arrives. When a block reveals a larger score than any seen so far, the accumulated result is corrected by the appropriate exponential factor so that the final answer is exactly what a single full-row softmax would have produced. This running-renormalization is the mathematical heart of the method: it is what lets a block-by-block computation reproduce the global softmax with no approximation. The recomputation it costs is repaid many times over by the HBM traffic it avoids — and in the backward pass, FlashAttention recomputes the score tiles on the fly rather than storing them, trading cheap arithmetic for expensive memory just as it does in the forward direction.

What FlashAttention-1 delivered

The payoff is exact attention that runs 2–4× faster than the standard PyTorch implementation and uses memory that grows linearly rather than quadratically in sequence length. Because the speedup comes from reducing memory traffic, it grows with sequence length — the longer the context, the more HBM round trips the naive kernel pays and the more FlashAttention saves. That linear-memory property is what made training and serving at long context lengths practical; models could attend over far longer sequences without the score matrix blowing the memory budget. FlashAttention was adopted almost immediately across the training and inference stacks precisely because it is a drop-in replacement: same outputs, fewer bytes moved.

FlashAttention-2: closing the gap to peak

FlashAttention-1 was fast but still left the GPU's matmul units underutilized relative to a pure GEMM. FlashAttention-2 (a follow-up by Tri Dao) reworks the algorithm to get closer to peak hardware throughput, with three main changes. First, it reduces non-matmul FLOPs — operations like the rescaling are far slower per-FLOP on a GPU than the tensor-core matmuls, so minimizing them matters more than the raw FLOP count suggests. Second, it improves parallelism, parallelizing the attention computation along the sequence-length dimension in addition to batch and heads, which keeps the GPU's many streaming multiprocessors busy even when the batch is small. Third, it refines the work partitioning between warps within each thread block to cut the shared-memory communication they need to exchange. Together these roughly double the throughput of the original, pushing exact attention to a large fraction of the GPU's theoretical matmul peak.

Why it matters

FlashAttention is the canonical example of an algorithm whose win comes from respecting the memory hierarchy rather than from changing the math. A long line of prior work attacked attention's $O(N^2)$ cost with approximations — sparse patterns, low-rank factorizations, kernel tricks — accepting some quality loss to dodge the quadratic. FlashAttention showed that you often do not need to approximate at all: by being careful about which bytes go where, you can compute the exact operation faster and with linear memory. It quietly became the default attention kernel for both training and inference, and it is one of the reasons the long-context era was possible at all — it is what keeps the score matrix out of HBM in prefill-heavy and long-sequence workloads.

Limitations and continuation

The gains are hardware-specific: FlashAttention is tuned to the GPU memory hierarchy and tensor-core layout, so each generation of accelerator needs the kernel re-tuned to its SRAM sizes and instruction set — a line continued by FlashAttention-3, which targets the newer Hopper architecture and its asynchronous, low-precision features. The method addresses the memory-traffic cost of attention but not the fundamental $O(N^2)$ compute, so at extreme sequence lengths the arithmetic itself eventually dominates and genuinely sub-quadratic or sparse approaches re-enter the picture. And being a hand-written, hardware-aware kernel, it is more involved to implement and maintain than a few lines of framework code — the price of running at the metal. None of that has dislodged it; for the sequence lengths that matter in practice, exact attention with minimal IO is simply the right default.