FlashAttention

FlashAttention computes the same numbers as the attention you already know. Not approximately the same — the same function, argued for in a theorem. It is faster because it moves less memory, and it moves less memory by never writing down an intermediate result that the textbook version writes down, reads back, writes again and reads again.

Two pieces of vocabulary, because the whole argument is about them. HBM is High Bandwidth Memory — the stack of DRAM next to the GPU die, the 80 GB an A100 advertises, delivering about 2 TB/s. SRAM is the small pool of memory on the chip itself: 192 KB per streaming multiprocessor on an A100, roughly 19 TB/s, about ten times faster and four hundred thousand times smaller. Those two figures are from the FlashAttention paper, and the gap between them is the entire idea.

A textbook attention kernel for a sequence of N tokens runs three separate passes over HBM. It computes S = QKᵀ and writes an N × N matrix of scores out to HBM. It reads that matrix back, applies the softmax, writes the N × N result back. It reads that back one more time and multiplies by V. For an 8k-token sequence that matrix is 67 million entries, 128 MiB per head, and it gets crossed four times.

The panel below is an accounting of every byte the kernel moves. Pick a sequence length, then switch the algorithm and watch which line item disappears.

Leave everything at its default and switch between Standard and FlashAttention-2. HBM traffic falls from 16.3 GiB to 4.19 GiB, and the reason is visible in the tape: of those 16.3 GiB, 16.0 GiB is the score matrix crossing the bus four times — written, read, written, read — and 256 MiB is Q, K, V and the output, the tensors anyone actually asked for. Peak intermediate memory falls from 4.00 GiB to 65.0 MiB. Now drag sequence length to 32,768: the standard kernel wants to allocate 64 GiB for its score matrices, which is where the familiar torch.OutOfMemoryError comes from, while FlashAttention needs 260 MiB and keeps running.

Why it is allowed to be exact

The obstacle to tiling attention is the softmax. Every output element depends on the sum of exponentials of an entire row of scores, and you cannot normalise by a sum you have not finished computing. That is the reason the naive kernel materialises the row: it needs all of it before it can divide.

The way out is to carry two extra numbers per row. Keep the largest score seen so far, call it m, and the running sum of exponentials, call it ℓ. When a new tile arrives with a larger maximum, every partial result computed before it was scaled against the old maximum — so multiply the running sum and the running output by exp(m_old − m_new) to put them on the new scale, then add the new tile. At the end, divide the accumulated output by ℓ. The answer is the same as if you had seen the whole row at once, because multiplying numerator and denominator of a fraction by the same constant does not change the fraction.

That is the entire trick, and the next panel runs it in front of you. The scores are on the left, the tiles are as wide as you make them, and the two checkboxes let you break the algorithm in each of the two ways it can be broken.

Move tile width through all five settings. The table changes completely — different maxima, different corrections, different partial sums — and the result does not move in any digit that float32 can represent. That is what "exact" means here. It does not mean bit-identical to a specific reference implementation; adding the same numbers in a different order can change the last bit, and the readouts show the one-shot version missing the double-precision truth by about as much as the tiled version does. It means the algorithm computes the function, with no term dropped, no sampling, and no low-rank stand-in. The paper states it as a theorem, and the difference between that and Linformer or Performer — which really do compute something else — is the difference between a compiler optimisation and a modelling decision.

Now break it. Uncheck rescale the accumulator: the running sum keeps contributions that were scaled against a maximum that no longer applies, and the answer is simply wrong. This is the bug everyone writes the first time they implement it. Then re-check it and uncheck subtract the running maximum instead. At magnitude ×6 nothing happens at all — the subtraction is algebraically a no-op, since exp(s−m)/Σexp(s−m) equals exp(s)/Σexp(s). Drag magnitude up to ×30 and the result becomes NaN, because the largest score is then 90 and exp(90) is 1.2 × 10³⁹, while float32 stops at 3.4 × 10³⁸. The maximum subtraction is not part of the mathematics; it is the thing that keeps the mathematics representable, and it is why the running maximum has to be carried alongside the running sum.

What it does not do

It does not reduce arithmetic. It adds some. The rescaling multiplications are extra work, and in training the backward pass throws away the score matrix and recomputes it from Q, K and V — the paper keeps only the output and the two statistics per row, then rebuilds S and P on demand. The authors are direct about the trade: "even with the increased FLOPs due to recomputation, our algorithm both runs faster … thanks to the massively reduced amount of HBM access." A kernel that does more arithmetic and finishes sooner is the clearest possible statement that arithmetic was never the constraint. If that idea is new, it is worth reading arithmetic intensity next, because it is the general form of the argument.

It does not shrink the KV cache. This is the most expensive misunderstanding in the list, because it leads people to expect a concurrency win that never arrives. The KV cache holds keys and values for every past token so decoding does not have to recompute them; FlashAttention removes an activation, the N × N scores, which lives for the duration of one kernel. The cache lives for the duration of the request. Set the simulation to FlashAttention-2 and read the peak intermediate memory: that number, and only that number, is what the kernel saves. The lever on cache bytes is grouped-query attention, and the other one is quantization.

It does almost nothing for a decode step. When you are generating one token, the query is a single row: there is no N × N matrix, there is an N-element vector of scores, and the kernel is reading the whole cache to do one dot product per head. Nothing is being materialised that could be avoided. The FlashAttention machinery is still used at decode time — the library ships a dedicated entry point that reads the cache in place — but what it buys there is parallelism across the sequence dimension when the batch is too small to fill the GPU, which is a different problem with a different name (flash-decoding). Set sequence length to 512 in the panel above to see the win shrink: with the N² term small, the traffic reduction drops from about 4× to under 3×, and it keeps shrinking as N does.

It is not always available. Real kernels support a specific set of head dimensions and dtypes. FlashAttention rejects float32, which is why some vision backbones in a multimodal model must be left on another implementation, and support for a given head_dim depends on the version you installed. When PyTorch cannot dispatch to it, you get the warning that 43,000 people have looked up — Torch was not compiled with flash attention — and a silent fall back to a slower path, at which point your traffic is back on the first row of the panel.

Checking it on a real system

In Transformers, the implementation is selected by name, and you can switch it at runtime:

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="flash_attention_2"
)
model.set_attn_implementation("sdpa")

"eager" is the three-kernel version this lesson starts with. "sdpa" is PyTorch's dispatcher, which picks a backend for you and usually picks a flash kernel when the shapes and dtype allow. "flash_attention_2" and "flash_attention_3" call the library directly. The dispatcher's choice is the part that silently changes under you, so pin it while measuring:

from torch.nn.attention import SDPBackend, sdpa_kernel

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    outputs = model.generate(**inputs)

If the backend is genuinely unavailable for your shapes, that block raises rather than quietly falling back, which is exactly what you want in a benchmark.

Three measurements tell you whether it is working, in increasing order of effort. First, peak memory: torch.cuda.max_memory_allocated() around a forward pass at two sequence lengths. If it grows with the square of the sequence length you are on the eager path, whatever the config says; if it grows linearly, you are not. This is the cheapest and most conclusive test, because the quadratic allocation is the one thing a flash kernel cannot be doing.

Second, the profiler. torch.profiler with CUDA activities will name the kernel: something containing flash_fwd or fmha for the flash path, versus a sequence of gemm, softmax and gemm for the eager one. Three separate kernels where you expected one is the whole story.

Third, achieved throughput. Take the useful attention FLOPs — 4 × N² × head_dim × heads × batch, halved for a causal model — divide by the measured time, and compare against your GPU's dense tensor-core rate. Below 20% of peak means memory-bound, which for a long prompt means the score matrix is going to HBM. Above 50% means you are getting what the paper got. On an A100 the FlashAttention-2 paper reports up to 230 TFLOP/s against a dense peak of 312, and it is worth knowing what good looks like before you spend a day tuning.

One caveat on all three: none of this touches decode. If your workload is long generations from short prompts, attention is a small share of your time and the honest answer is that this kernel is not your problem — chunked prefill and the cache are. Profile before installing anything.

Next: the roofline the whole argument rests on, arithmetic intensity; the step being tiled, scaled dot-product attention; and the methods that really are approximations, linear attention.