Attention Is Memory-Bound, and Nobody Noticed for Five Years

Between roughly 2018 and 2022, a large body of work attacked attention's quadratic cost by reducing its arithmetic. Sparse attention patterns computing only some entries. Low-rank approximations. Kernel methods. Locality-sensitive hashing to attend only to likely-relevant positions.

Many of these reduced the operation count substantially, from quadratic to linear or near-linear, and were correct about the mathematics.

Most of them were not faster in wall-clock time. Some were slower than the dense implementation they replaced, at the sequence lengths people actually used.

That outcome was treated for a while as an engineering problem: the approximations needed better kernels, better libraries, more tuning. It was not. It was a diagnosis error.

The operations were never the bottleneck. A GPU running standard attention spends most of its time waiting on memory, and an algorithm that halves the arithmetic while leaving the memory traffic alone does not halve the runtime. It changes almost nothing.

FlashAttention's contribution was to measure the right quantity, and then to reduce it while computing attention exactly.

The fact that makes this work is that a GPU has two very different places to keep data, and their properties are close to opposite.

High-bandwidth memory is the large pool quoted on a spec sheet, tens of gigabytes, with bandwidth on the order of one and a half to three terabytes per second. Large and, by GPU standards, slow.

On-chip SRAM is the memory inside each streaming multiprocessor. On an A100 there are 108 of these, each holding around 192 kilobytes, for roughly twenty megabytes in total. Bandwidth is around nineteen terabytes per second.

So SRAM is about an order of magnitude faster and about three orders of magnitude smaller.

Every byte a kernel touches comes from one of these. Data in SRAM is nearly free to access; data in HBM costs roughly ten times more per byte, and there is no way to make the trip cheaper.

The consequence is that a kernel's speed is often decided not by what it computes but by how many times it moves data across that boundary. Optimising a GPU kernel is frequently a data movement problem wearing an arithmetic disguise.

Now trace what a textbook attention implementation does, in terms of that boundary.

Read Q and K from HBM. Compute their product, producing an n-by-n score matrix. That matrix does not fit in SRAM, so write it to HBM.

Read the score matrix back from HBM. Apply the softmax. Write the result to HBM.

Read the softmax output from HBM again, read V, multiply, write the final output to HBM.

The n-by-n matrix has crossed the boundary four times, twice written and twice read, and it is the largest object in the computation by a wide margin. Q, K and V are each only n by d, where d is 64 or 128, so the score matrix is larger by a factor of n over d, typically a factor of a hundred or more.

Meanwhile the arithmetic is two matrix multiplications and an elementwise operation, which tensor cores dispatch quickly.

So the profile is a kernel doing modest arithmetic separated by enormous memory transfers. The arithmetic units idle while the transfers complete.

The abstraction hides the scale, so put numbers on the score matrix in 16-bit, for a single attention head.

At 2,048 tokens it is 8.4 MB. At 4,096 it is 33.6 MB. At 8,192 it is 134.2 MB. At 16,384 it is 536.9 MB.

Now multiply by the heads and the batch, because they all exist at once. A batch of 8 with 32 heads at 8,192 tokens needs 34.4 gigabytes for score matrices alone, before weights, activations or optimizer state.

That number explains something that puzzled people at the time. Long-context training was failing with out-of-memory errors that seemed disproportionate to model size, and the reason was an intermediate value that appears nowhere in the parameter count.

It also explains why the traffic dominates. Moving 34 gigabytes across the memory boundary four times is 138 gigabytes of transfer, and at roughly 2 terabytes per second that is about 70 milliseconds of pure waiting, for one attention layer, in one forward pass.

The arithmetic in the same layer takes a small fraction of that.

There is a standard way to make this precise, and it generalises far beyond attention.

Arithmetic intensity is the ratio of operations performed to bytes moved. Divide a kernel's FLOPs by its memory traffic and compare against the hardware's own ratio of peak FLOPs to peak bandwidth.

If your kernel's intensity is above the hardware's, you are compute-bound: the arithmetic units are the constraint and reducing operations helps. If it is below, you are memory-bound: the units are idle waiting on data, and reducing operations helps not at all.

Modern accelerators have very high ratios, hundreds of FLOPs per byte, because arithmetic has grown far faster than memory bandwidth for decades. So the bar for being compute-bound is high, and most kernels fall below it.

Standard attention falls well below. It is a chain of operations each of which reads a large matrix, does a modest amount of work per element, and writes a large matrix.

This is the roofline model, and reading a kernel against it is the first diagnostic step. The catalogue's lesson on why most code is memory-bound develops the general case; this path is what it looks like when applied to one operation seriously.

The diagram contrasts the two shapes, and the difference is entirely about how often data crosses the boundary.

The standard path is four separate kernel launches, each reading its input from HBM and writing its output back. The score matrix and the softmax output both make the round trip, and both are the large n-by-n object.

The fused path loads blocks of Q, K and V into SRAM, performs the entire computation there, and writes only the final output back. The score matrix is never written to HBM at any point, because it never exists in full.

That is the whole idea, stated once. The n-by-n matrix is an artefact of how the computation is written down, not something the answer requires. Attention's output is n by d, which is small. Everything larger is scaffolding, and scaffolding does not need to be materialised.

What makes it non-trivial is the softmax in the middle, which appears to require the whole row before it can produce anything, and therefore appears to require the whole matrix to exist. Removing that obstacle is the next lesson.

The property that separated this from everything before it is worth stating plainly, because it is easy to assume otherwise.

FlashAttention computes exactly the same function as standard attention. Not an approximation whose error is bounded, not a sparse pattern that skips entries, not a low-rank surrogate. The same numbers, to within floating point reassociation.

That matters for three reasons.

Adoption becomes trivial. A drop-in replacement that changes no results needs no retraining, no revalidation, and no discussion about acceptable degradation. Approximate methods each required all three.

It composes with everything. Because the interface and semantics are unchanged, it works with any model, any position encoding, any attention variant built on top.

And it removed the excuse. Once an exact method was faster than every approximation, approximations had to justify themselves on quality-per-unit-time rather than on operation counts, and most could not.

The broader lesson generalises past attention: before approximating a computation to make it faster, check whether the exact version is being executed badly.

The technique underneath is not specific to attention, and recognising it elsewhere is most of the transferable value.

A sequence of operations each reading from and writing to main memory pays the round trip once per operation. Fusing them into a single kernel means intermediate values stay in registers or on-chip memory and never make the trip.

The canonical small example is an elementwise chain. Multiply a tensor by a scalar, add a bias, apply an activation. As three kernels that is three reads and three writes of the whole tensor. Fused, it is one read and one write, and the intermediate values live in registers. The arithmetic is identical and the kernel runs roughly three times faster, because the arithmetic was never the cost.

This is why compilers for deep learning spend most of their effort on fusion decisions rather than on the operations themselves.

Attention is the hard case of the same idea. The chain is matrix multiply, softmax, matrix multiply, and the obstacle is that the middle step is a reduction over a full row rather than an elementwise operation. Fusing across a reduction is what required a new algorithm rather than a compiler pass.

The results are worth having in mind, because they set expectations for what this class of optimisation is worth.

The original FlashAttention reported speedups up to 7.6 times on GPT-2 training, and reduced attention's memory footprint from quadratic to linear in sequence length. The second number is the more consequential one: it is what made long-context training feasible rather than merely faster.

FlashAttention-2 restructured the work partitioning and reduced non-matrix-multiply operations, reaching 1.7 to 3.0 times faster than the original and 3 to 10 times faster than a standard implementation, at up to 230 TFLOPs per second, which is 73 percent of theoretical peak on an A100.

On newer hardware the same code was leaving performance behind: FlashAttention-2 reached only about 35 percent utilisation on an H100, because it did not exploit that generation's asynchrony.

FlashAttention-3 addressed that with warp specialisation, overlap of matrix multiplication with softmax, and low-precision support, reaching 1.5 to 2.0 times faster than its predecessor, up to 840 TFLOPs per second in BF16, about 85 percent utilisation, and 1.3 PFLOPs per second in FP8.

Note the pattern in that sequence: each version is a response to a specific hardware generation.

It is fair to ask why any of this matters when the implementation is a library you import.

The direct answer is that reasoning about the memory hierarchy transfers to everything else you will profile. The pattern that made attention slow, an intermediate larger than its inputs and outputs making unnecessary round trips, recurs constantly. Once you look for it you find it in loss functions, normalisation layers, custom heads, data preprocessing, and anywhere someone wrote three tensor operations in a row.

The practical answer is that the library only covers the standard case. The moment your model does something slightly unusual, a custom attention bias, an unusual masking pattern, a novel position encoding, a new sparsity structure, the fast path stops applying and you are back to the four-kernel version without necessarily noticing.

And the diagnostic answer is that knowing what a well-optimised attention kernel achieves tells you whether yours is one. Roughly 70 to 85 percent of theoretical peak is what good looks like. A profile showing 15 percent is not a hardware limit; it is a fusion opportunity.

The next lesson shows how the softmax obstacle was removed, and the one after that is about writing kernels of your own.