FlashAttention

Definition

FlashAttention is a way of computing attention that gives exactly the same answer as the textbook version but runs far faster — because it stops treating attention as a maths problem and starts treating it as a memory problem. The insight behind it, from Tri Dao and colleagues in 2022, is that the slow part of attention on a GPU is not the multiplications. It is the traffic: standard attention builds a giant n-by-n table of scores in the GPU's main memory and reads it back several times, and the chip spends most of its time waiting for those bytes to move rather than doing arithmetic.

The word to hold onto is exact. FlashAttention is not an approximation like sparse or linear attention, which drop or compress interactions to save work and change the output in the process. It reorganises the identical computation so the huge intermediate table never has to be written down, and it returns the same numbers the slow method would — down to floating-point rounding. You get the speed and the memory saving for free, with no cost in model quality. That combination is why it went from a research paper to the default attention kernel inside PyTorch, Hugging Face and most training stacks within about a year.

How It Works

Attention is memory-bound, not compute-bound

Self-attention works by comparing every token to every other token. For a sequence of length n, that comparison produces an n-by-n grid of scores — token 5's relevance to token 900, and so on for every pair. The standard implementation computes this grid, applies a softmax across each row, and uses the result to take a weighted average of the value vectors.

The problem is where that grid lives. A GPU has two very different kinds of memory, and the gap between them is the whole story. As the FlashAttention paper describes the A100 it was written on: it has "40-80GB of high bandwidth memory (HBM) with bandwidth 1.5-2.0TB/s and 192KB of on-chip SRAM per each of 108 streaming multiprocessors with bandwidth estimated around 19TB/s." Read those two numbers together. HBM is the big pool the model lives in, but at roughly 2 TB/s it is ten times slower than the tiny on-chip SRAM. And the SRAM is genuinely tiny — 192 KB per core, a little over 20 MB in total across the whole chip.

Standard attention writes the entire n-by-n score grid out to HBM, reads it back to run the softmax, writes the softmax result back, and reads it again for the final weighted average. The arithmetic in between is cheap. The chip is bandwidth-starved, doing round trips to slow memory for a table it will throw away moments later. This is the same memory wall that governs inference generally: moving the bytes costs more than computing on them.

The n-by-n matrix is enormous — and it is pure waste

Here is why that traffic hurts so much. Take a single attention head with a head dimension of 128, in 16-bit precision (2 bytes per number), and watch the score matrix grow as the sequence gets longer:

The score matrix is n-squared × 2 bytes: at 64K tokens that is 65,536 × 65,536 × 2 = about 8 GiB for one head — and a model has dozens of heads across dozens of layers. Meanwhile the actual inputs and output for that head — the query, key, value and output vectors — are each just n × 128 × 2 bytes, about 16 MiB apiece, 64 MiB together. The intermediate the GPU is laboriously shuttling to and from HBM is 128 times bigger than the data that matters, and it exists only to be summed away by the softmax. Worse, it grows with the square of the sequence length while the real data grows only linearly: every doubling of context quadruples the matrix but merely doubles the inputs. That divergence is what makes long context expensive, and it is exactly the cost FlashAttention removes.

Tiling plus online softmax: computing the grid without storing it

FlashAttention's trick is to never hold a whole row of the score matrix at once. It cuts the queries, keys and values into blocks small enough to fit in the fast SRAM, and streams through them. For each block of keys and values, it computes just that block's slice of the scores inside SRAM, folds it into a running result, and discards it before moving on. The full n-by-n grid is computed piece by piece but never assembled in HBM.

The obstacle is the softmax. A softmax needs to divide by the sum of the exponentials across an entire row — which naively means you have to see the whole row before you can normalise any of it. FlashAttention gets around this with online softmax: it keeps a running maximum and a running denominator as it goes, and when a later block turns up a larger score than anything seen so far, it rescales the partial result already accumulated so the final answer comes out identical to a one-shot softmax. This is the piece that makes the whole thing exact rather than approximate.

Because the tiles are small — a 128-by-128 block in 16-bit is just 32 KiB, comfortably inside the 192 KB of SRAM per core — the expensive HBM round trips for the score matrix disappear entirely. The paper reports "up to 9× fewer" memory accesses than standard attention, and proves a matching lower bound: no exact attention algorithm can do asymptotically better across all SRAM sizes. The memory footprint drops from O(n-squared) to O(n) — the paper measures FlashAttention as "up to 20× more memory efficient than exact attention baselines" — because the thing that scaled with n-squared is simply never written down.

What the reader should take away

The speed does not come from cutting corners. It comes from noticing that the corners were being cut in the wrong place: the GPU was optimised to do arithmetic fast and then made to spend its time moving a throwaway table across a slow bus. FlashAttention is IO-aware — it is written around the memory hierarchy instead of ignoring it — and that single change of perspective is why it is both faster and lighter without touching the model.

Real-World Applications

It is the default attention kernel almost everywhere. Since PyTorch 2.0 (2023), calling torch.nn.functional.scaled_dot_product_attention dispatches to a FlashAttention kernel when the inputs allow it — so a large share of all transformer training and inference on NVIDIA hardware runs FlashAttention whether or not the author ever typed the name. Hugging Face Transformers exposes it via an attn_implementation="flash_attention_2" flag, NVIDIA's Megatron-LM uses it for large-scale pretraining, and inference servers such as vLLM build on it. This is unusually broad adoption for a technique this recent, and it happened because there was no trade-off to weigh: the output is identical, so switching it on is close to free.

It is a large part of why long context became practical. When a lab ships a model with a 100K- or million-token context window, the O(n-squared) memory blowup of standard attention would have made training and serving those lengths prohibitive — an 8 GiB score matrix per head at 64K is a wall you hit fast. FlashAttention removed that particular wall by making attention's memory linear in the sequence length. It is not the only ingredient — the KV cache and the arithmetic of serving still dominate inference economics, and techniques such as grouped-query attention and disaggregated serving address those — but the ability to compute attention over a very long sequence without materialising its matrix is a precondition for the rest.

It set a training-speed record on release. In the 2022 paper the authors used FlashAttention to train BERT-large (sequence length 512) about 15% faster than the then-current MLPerf 1.1 speed record set by NVIDIA (17.4 minutes against 20.0, averaged over ten runs on eight A100s), and reported a 3x end-to-end speedup training GPT-2 (sequence length 1K) and a 2.4x speedup on the Long-Range Arena benchmark. Those particular percentages are perishable — they are tied to that hardware, those baselines and that year, and later versions moved them — but the reason behind them is the durable part: fewer trips to slow memory.

Challenges

The speedup is hardware-specific and version-specific. FlashAttention is hand-tuned around a particular GPU's memory hierarchy and instruction set, so a kernel written for one generation does not automatically deliver the same gain on the next. This is why the technique keeps being rewritten rather than finished — see Future Trends — and why a quoted speedup number without a chip and a date attached to it means little.

It is a fused kernel, not a Python function you can casually modify. The whole point is to keep the intermediate results in registers and SRAM and never spill them, which means the tiling, the online-softmax bookkeeping and the backward pass are welded together in low-level CUDA. Adding a new twist to attention — a custom mask, a new positional scheme, a different scoring function — often means writing or waiting for a matching fused kernel rather than just editing a few lines of model code. The convenience of the standard, un-fused implementation was that you could change it; the speed of FlashAttention is bought partly by giving that up.

It solves the score matrix, not the KV cache. FlashAttention removes the O(n-squared) intermediate of attention, and it is easy to conclude from that that long-context memory is a solved problem. It is not. During generation the model still has to store a key and value vector for every past token — the KV cache — and that cost is linear in context, per user, and often the real ceiling on how many sessions a GPU can serve. FlashAttention makes computing attention cheap in memory; it does nothing about remembering the past, which is a separate bill.

Future Trends

The clearest trend is that FlashAttention is a moving target that gets re-implemented for each new GPU generation rather than a fixed artifact. FlashAttention-2 (Dao, 2023) kept the algorithm but reworked how the computation is partitioned across a GPU's cores and warps, cutting the non-matmul overhead and raising utilisation. FlashAttention-3 (2024) is written specifically for NVIDIA's Hopper architecture, exploiting its asynchronous memory and tensor-core instructions and its FP8 support to overlap data movement with computation more aggressively. The through-line is that each version chases the same goal — do exact attention with as few memory stalls as the hardware allows — against a memory hierarchy that keeps changing shape, which is why the numbers keep moving while the principle does not.

The second is that the IO-aware idea has spread beyond attention. Treating a kernel as a memory problem first and an arithmetic problem second — fuse the operations, keep the working set in SRAM, count the HBM accesses — is now a standard lens for writing fast deep-learning kernels generally, and compilers and libraries increasingly try to produce fused kernels of this shape automatically rather than leaving each one to be hand-written.

Code Example

FlashAttention's exactness rests on online softmax: the claim that you can compute a softmax block by block, without ever holding the whole row, and get the same answer. That claim is easy to check. The snippet below computes attention for one query the standard way — materialising the full length-n score row — and then again the FlashAttention way, streaming over blocks and keeping only a running max, denominator and output. The two results should be identical up to floating-point noise.

import numpy as np

np.random.seed(0)
n, d = 2048, 64
Q = np.random.randn(d)          # one query row
K = np.random.randn(n, d)
V = np.random.randn(n, d)

# Standard attention: build the full length-n score row (a row of the n x n
# matrix), softmax it, then weight V. The whole row lives in memory at once.
scores = K @ Q
p = np.exp(scores - scores.max())
out_standard = (p / p.sum()) @ V

# FlashAttention style: stream over blocks, never hold the whole row.
# Keep a running max (m), running denominator (l) and running output (acc).
block = 256
m, l, acc = -np.inf, 0.0, np.zeros(d)
for i in range(0, n, block):
    s = K[i:i+block] @ Q               # this block's scores only
    m_new = max(m, s.max())
    corr = np.exp(m - m_new)           # rescale everything seen so far
    p = np.exp(s - m_new)
    l = l * corr + p.sum()
    acc = acc * corr + p @ V[i:i+block]
    m = m_new
out_flash = acc / l

print("max abs difference between the two outputs:", np.abs(out_standard - out_flash).max())

Running it prints:

max abs difference between the two outputs: 1.1657341758564144e-15

That difference — about one part in a quadrillion — is floating-point rounding, not a real disagreement. The block-wise version never held more than 256 scores at a time, yet it reproduced the full-row softmax exactly. That is the entire promise of FlashAttention in miniature: the real kernel does this on 2D tiles inside GPU SRAM instead of a 1D row in NumPy, and adds a matching backward pass, but the reason it is exact rather than approximate is precisely the rescaling step (corr) that fixes up the running total whenever a later block reveals a bigger score.