Flash Attention Explained: Same Maths, Far Less Memory

Flash attention is one of the rare optimisations that produces exactly the same output as the thing it replaces. It is not an approximation, it does not trade quality for speed, and there is no setting where you would prefer the naive version. Understanding why it works explains something more general: on modern accelerators, the bottleneck is usually moving data, not computing on it.

The cost that actually bites

A GPU has two very different kinds of memory. There is a large pool of high-bandwidth memory measured in tens of gigabytes, and a much smaller pool of on-chip memory measured in tens of kilobytes per compute unit. The on-chip pool is roughly an order of magnitude faster to reach.

Arithmetic throughput on these chips has grown far faster than memory bandwidth for years. The consequence is that a great many kernels finish their arithmetic and then sit idle waiting for the next block of data to arrive.

Attention is a textbook case. The comparisons it performs are simple, and the volume of intermediate data it generates is enormous. It spends most of its time on memory traffic.

What naive attention does to memory

The standard implementation proceeds in stages. It multiplies queries against keys to produce a full score matrix, writes that matrix out, reads it back to apply softmax, writes the normalised weights out, reads them back to multiply against values, and writes the result.

That score matrix has one entry for every pair of tokens. At 8,000 tokens that is 64 million entries per attention head, per layer. At 100,000 tokens it is 10 billion. Each of those entries is written to slow memory and read back at least once.

The arithmetic itself is trivial by comparison. Nearly all the wall-clock time goes into shuttling a matrix that exists only to be consumed immediately. It is also why attention historically dominated the memory footprint of long-context training, before this was fixed.

The idea: never materialise the matrix

Flash attention refuses to build the full score matrix at all. It splits the queries, keys and values into blocks small enough to fit in on-chip memory, and processes one block pair at a time.

For each pair it computes the scores, applies the weighting and accumulates into a running output — all without the intermediate ever leaving fast memory. When the loop finishes, the output is complete and the score matrix was never written anywhere.

The obstacle is that softmax normalises across an entire row, and a block only sees part of that row. The solution is to carry a running maximum and a running sum alongside the partial output, and rescale the accumulated result whenever a new block reveals a larger value. The final numbers come out identical to computing the row in one go.

Why it is not a trade-off

This is worth stating plainly because so many inference optimisations do involve a trade. Quantisation trades numerical precision for memory. Speculative decoding trades extra compute for latency. Aggressive sampling settings trade coherence for variety.

Flash attention trades nothing. It reorganises the order of operations to respect the memory hierarchy, and the mathematics is unchanged. Two models with identical weights, one served with flash attention and one without, produce the same output.

What differs is that the memory required for attention grows with sequence length rather than with its square. That single change is a large part of why context windows moved from a few thousand tokens to the 1M windows now shipping on Kimi K3, GLM-5.2 and both DeepSeek V4 variants.

Where you actually feel it

The gain concentrates in prefill — the phase where the model reads your prompt. Prefill runs attention across every input token at once, so it is exactly the workload that generates the huge score matrix.

During decode, the model attends over one new query against a cached set of keys and values, so the score matrix is a single row and there was never much to save. Generation speed is limited by other things. The KV cache is what dominates there, and inference latency explained covers how the two phases combine into the numbers you measure.

So the practical signature is a large improvement in time to first token on long prompts, and little change in tokens per second once generation is underway.

What it does not solve

Flash attention reduces memory traffic. It does not reduce the number of comparisons, which is still quadratic in sequence length. A 200K-token prompt still performs vastly more attention work than a 20K one, and still costs more and takes longer.

It also does nothing for the KV cache, which is a separate store holding one key and one value per token per layer for the whole session. That is the constraint that limits how many concurrent long-context requests a machine can serve, and it is attacked by different techniques — grouped-query attention and the latent-attention schemes used in the Kimi line.

Nor does it help with the softer failure of long context, where the model accepts everything you send but uses the middle of it less reliably than the ends. That is a property of attention itself, not of how it is implemented.

The takeaway

If you serve models yourself, confirm your runtime has a flash-attention kernel enabled for your hardware. Almost every serious inference stack ships one, but it can be silently disabled by an unsupported GPU, an odd head dimension or a fallback path, and the symptom is simply that long prompts are slower and use more memory than they should.

If you use a hosted API, this is already done for you and there is no knob. The useful part is the mental model: when a long prompt is slow, the cause is usually data movement, and the fix is usually sending less rather than tuning harder. Choosing a model for self-hosting covers the rest of that sizing exercise.