FlashAttention Explained: Faster, Memory-Aware Attention

FlashAttention is an algorithm for computing the attention step in transformer models that produces exactly the same result as standard attention but runs faster and uses far less memory. It gets there not by doing less math but by moving less data. Standard attention writes a huge intermediate matrix out to the GPU’s main memory and reads it back; FlashAttention breaks the computation into tiles small enough to stay in fast on-chip memory and never materializes the full matrix at all. It is now the default attention implementation in most training and inference stacks.

A quick refresher on attention

In self-attention, each token produces three vectors: a query, a key, and a value. For a sequence of length N, the model computes:

1. Scores — multiply the query matrix Q by the transpose of the key matrix K, giving an N × N matrix S of how much each token should attend to every other token.
2. Softmax — normalize each row of S into a probability distribution P.
3. Output — multiply P by the value matrix V to get each token’s new representation.

The N × N matrix is the problem. Double the sequence length and it grows fourfold. For long contexts, that matrix is enormous, and a naive implementation stores all of it.

The real bottleneck is memory traffic

A GPU has two very different kinds of memory. High-bandwidth memory (HBM) is large, holding tens of gigabytes, but comparatively slow to access. On-chip SRAM, the shared memory and registers inside each streaming multiprocessor, is tiny but much faster. The gap between them mirrors the general trade-off covered in SRAM vs DRAM: fast memory is small and expensive, big memory is slow.

A standard attention implementation runs as separate steps. It computes S and writes it to HBM, reads S back to compute softmax and writes P to HBM, then reads P back to multiply by V. The arithmetic itself is cheap relative to the cost of shuttling an N × N matrix in and out of HBM multiple times. In other words, standard attention is memory-bound: the compute units spend much of their time waiting for data.

FlashAttention’s authors framed the problem in exactly these terms. The goal is to be IO-aware, minimizing reads and writes to HBM even if that means doing some extra computation.

How FlashAttention works

The algorithm rests on two techniques.

Tiling

Q, K, and V are split into blocks small enough to fit in SRAM. FlashAttention loads a block of queries, then loops over blocks of keys and values. For each pair, it computes the partial scores, applies softmax, and accumulates into the output, all inside SRAM. Only the final output for that block of queries is written back to HBM.

The full N × N matrix never exists in HBM. At any moment, only one tile of it lives in fast memory.

Online softmax

Tiling has a catch. Softmax normalizes over an entire row, which seems to require seeing every score in that row before producing any output. FlashAttention sidesteps this with a technique often called online softmax.

As each new block of scores arrives, the algorithm tracks two running statistics per row: the maximum score seen so far and the running sum of exponentials. When a new block raises the maximum, previously accumulated output is rescaled by a correction factor. After the last block, one final normalization yields exactly the softmax result a full-row computation would have produced. Nothing is approximated.

Recomputation in the backward pass

Training needs the attention probabilities to compute gradients. Rather than storing the N × N matrix for the backward pass, FlashAttention stores only the small per-row statistics and recomputes the attention tiles during backpropagation. Recomputing is cheaper than reading a giant matrix back from HBM, so this extra arithmetic still comes out ahead.

Standard attention vs FlashAttention

	Standard attention	FlashAttention
Result	Exact	Exact, same output
Stores N × N matrix in HBM	Yes	No
Extra memory for attention	Grows with N²	Grows with N
Main bottleneck	HBM reads and writes	Much closer to compute-bound
Kernel structure	Several separate operations	One fused kernel
Long-context feasibility	Limited by memory	Far longer sequences on the same hardware

The compute cost is still quadratic in sequence length. FlashAttention does not change the asymptotic number of operations. It changes how efficiently the hardware executes them and removes the quadratic memory requirement.

Later versions

Successive versions of FlashAttention kept the same core idea while tuning how work is split across the GPU. Improvements included parallelizing over the sequence dimension, not just across batches and heads, so long sequences with small batches still keep the GPU busy; reducing non-matrix-multiply operations, which run much slower than matrix math on tensor cores; and taking advantage of asynchronous data movement and lower-precision formats on newer accelerators.

The details are tied to specific GPU generations, but the principle is stable: fuse the operations, keep tiles in SRAM, and minimize trips to HBM.

How it fits with other attention optimizations

FlashAttention is complementary to most other techniques for making transformers cheaper:

• Grouped-query attention reduces the number of key and value heads, shrinking memory and bandwidth for the K and V tensors. FlashAttention kernels support it directly.
• The KV cache avoids recomputing keys and values for earlier tokens during generation. Inference-focused variants of FlashAttention are designed to read efficiently from a KV cache, including caches split into pages.
• Quantization shrinks the numbers themselves, which pairs naturally with a kernel already optimized around memory traffic.

It also illustrates why memory bandwidth is such a recurring theme in AI hardware. The reason accelerators invest so heavily in HBM versus GDDR is the same reason FlashAttention matters: for many workloads, moving data is the expensive part, not the arithmetic.

Do you need to do anything to use it?

Usually not. Major deep learning frameworks ship fused attention operations that dispatch to a FlashAttention-style kernel when the hardware and input shapes support it. If you write a model using the framework’s built-in scaled dot-product attention rather than hand-coding the matrix multiplications, you are likely already benefiting.

Where it does pay to pay attention: custom attention masks, unusual head dimensions, or older hardware may force a fallback to a slower path. If long-context training runs out of memory unexpectedly, checking which attention backend is active is a good first step.

The takeaway

FlashAttention computes exact attention while never writing the full N × N score matrix to GPU main memory. Tiling keeps small blocks of queries, keys, and values in fast on-chip SRAM, online softmax lets it normalize rows incrementally, and recomputation avoids storing the matrix for training. The math is unchanged; the memory traffic is not. That is what turned long-context transformers from a memory problem into a practical default.