FlashAttention: Why Tiling Attention Makes Transformers Fast Without Changing the Result

Transformers did not get slow because attention needs too many multiplications. They got slow because attention moves too much data. FlashAttention is a lesson in that distinction. It does not change what attention computes. It changes where the intermediate results live, how often they are read and written, and in what order the work happens. The result is the same attention output, computed with far less traffic to slow memory, and therefore much faster in practice, especially for long sequences.

This article explains the memory wall behind standard attention, the online softmax trick that makes tiling possible, the forward and backward passes of FlashAttention, and the limits that remain. If you have ever watched a model run out of memory at a modest sequence length, or wondered why a theoretically quadratic operation can still be made practical, this is the mechanism to understand.

The problem is memory, not arithmetic

Standard scaled dot-product attention starts with three matrices for one head: queries Q , keys K , and values V . For a sequence of length N and head dimension d , each is N×d . The output is:

The score matrix S=QKT/d​ is N×N . That is the first problem. For N=8192 , one head already produces about 67 million scores. Store that in 16-bit precision and you need roughly 134 megabytes for a single head, for a single layer, for a single sequence. Multiply by heads, layers, and batch size, and the score matrix dominates memory long before the weights do.

The second problem is worse than the storage itself: the traffic. A standard implementation typically materialises S in high-bandwidth memory (HBM), reads it back to apply the softmax, writes the probability matrix P back to HBM, and then reads P again to multiply by V . Every one of those passes crosses the slowest link in the system.

Modern GPUs have a memory hierarchy, and the gap between its levels is the whole story. HBM is large but relatively slow. On-chip SRAM is tiny but very fast, often an order of magnitude higher in bandwidth. A standard attention implementation is therefore memory-bound: the arithmetic units spend much of their time waiting for data to arrive from HBM. Counting floating point operations misses this completely. Two algorithms with the same operation count can run at very different speeds if one of them reads and writes N2 values to HBM and the other does not.

FlashAttention, introduced by Dao and co-authors, attacks exactly this gap. Its goal is to never materialise the full N×N matrix in HBM at all.

Why softmax seems to block tiling

Tiling is the classic fix for a memory-bound operation: break the matrices into blocks that fit in fast SRAM, compute on one block at a time, and keep the partial results on chip. Matrix multiplication tiles beautifully. Softmax does not, at first glance, because it is a row-wise operation that needs the whole row before it can produce any output.

For one row of scores x=(x1​,…,xN​) , the softmax is:

The subtraction of the maximum m is there for numerical stability: it keeps the exponentials at or below 1 and prevents overflow. But it creates a dependency. You cannot compute the denominator, the sum of exponentials, until you have seen every element of the row and found its maximum. If you process the row in blocks, the first block cannot be finished until the last block has been seen. That seems to force either storing the whole row or making two passes over it.

Online softmax removes that obstacle. The trick is to keep running statistics and to correct earlier partial results when new information arrives, rather than waiting for the complete row.

Online softmax: correct as you go

Suppose you process a row in blocks. After the first block, you have a running maximum m1​ and a running sum of exponentials l1​ :

Now a second block arrives with a larger maximum, so the new running maximum becomes m2​>m1​ . Every exponential computed against m1​ is now off by a common factor, because:

So the old running sum can be rescaled, not recomputed:

The same rescaling applies to the partial output. Attention output for a row is a weighted sum of value vectors, where the weights are the softmax probabilities. If you have accumulated a partial output O1​ from the first block, normalised by the old sum l1​ , you can rescale it by the same factor and add the contribution of the new block. In general, after processing blocks up to block t , with running maximum mt​ and running sum lt​ , the partial output is:

In practice the implementation keeps the unnormalised accumulator and divides once at the end, which avoids a division per block. The key point is conceptual: softmax is no longer a barrier to tiling. Each block can be processed, folded into running statistics, and discarded. No block needs to be stored after its contribution has been absorbed.

This is the same idea as computing a running mean or a running maximum in a single streaming pass. The only extra ingredient is the rescaling factor, which is a single scalar per row, so correcting all earlier work costs almost nothing.

The tiled forward pass

With online softmax in hand, the forward pass of FlashAttention looks like this for one query block:

1. Load a block of queries Qi​ into SRAM, together with running statistics mi​ and li​ and the partial output Oi​ .
2. Loop over key and value blocks (Kj​,Vj​) , one pair at a time:
   • Load Kj​ and Vj​ into SRAM.
   • Compute the block scores Sij​=Qi​KjT​/d​ on chip.
   • Find the block maximum, update the running maximum, and rescale li​ and Oi​ .
   • Compute the block softmax weights and accumulate Pij​Vj​ into Oi​ .
   • Discard Sij​ and Pij​ . They never reach HBM.
3. After all key blocks have been processed, divide Oi​ by li​ and write only the final output block back to HBM.

The block sizes are chosen so that a query block, a key block, a value block, and their score block all fit in SRAM at once. That constraint is what sets the tile sizes in a real implementation, and it is why FlashAttention is written as a fused GPU kernel rather than as a sequence of ordinary library calls. A library that materialises each intermediate between calls cannot achieve the same effect, because the intermediate would already have been written to HBM.

The memory complexity of the forward pass drops from quadratic to linear in sequence length. Only Q , K , V , and the output need to live in HBM, plus the small running statistics. The N×N matrix exists only transiently, one tile at a time, on chip.

Note what has not changed. The number of multiplications is still quadratic in N . FlashAttention is not a sparse or low-rank approximation. It computes the exact same attention output, up to floating point rounding, as the standard implementation. The speedup comes from doing the same arithmetic with far fewer slow memory accesses, and from keeping the arithmetic units fed.

The backward pass: recompute instead of storing

Training needs gradients, and the backward pass raises an obvious objection. Standard backpropagation stores the attention probabilities P during the forward pass so the backward pass can reuse them. If FlashAttention never stores P , how does it compute gradients?

The answer is recomputation, and it is a good trade. FlashAttention stores only the output O and the per-row statistics m and l from the forward pass. In the backward pass, it recomputes each score block Sij​ from Qi​ and Kj​ on chip, reconstructs the block probabilities from the stored statistics, and accumulates the gradients for Q , K , and V block by block.

Recomputation costs extra arithmetic, but arithmetic is cheap relative to HBM traffic. Storing P would cost N2 memory and the reads and writes to move it. Recomputing it costs a second pass over the same tiles, all on chip. For a memory-bound operation, that trade is a clear win, and it is the same logic behind gradient checkpointing, applied at the granularity of a single kernel instead of a whole layer.

Concretely, the backward pass needs the gradient of the loss with respect to the output, dO . From that, the gradient with respect to the value block is dVj​=PijT​dOi​ , accumulated over query blocks. The gradient with respect to the score block follows from the softmax Jacobian applied to dPij​=dOi​VjT​ , and then dQi​ and dKj​ follow from dSij​ by the usual matrix product rules. Every quantity in that chain is recomputed or accumulated tile by tile, so the backward pass also avoids materialising an N×N matrix.

A small numerical example

The rescaling step is easy to get wrong in intuition, so consider one row with scores split into two blocks: block 1 is (1,3) and block 2 is (2,5) . The full softmax, computed directly, has maximum 5 and denominator:

Online, after block 1 the running maximum is 3 and the running sum is e1−3+e3−3≈0.135+1.000=1.135 . Block 2 raises the maximum to 5, so the old sum is rescaled by e3−5≈0.135 , giving 0.135×1.135≈0.153 . Adding block 2 gives 0.153+e2−5+e5−5≈0.153+0.050+1.000=1.203 , exactly the direct denominator. The partial output is rescaled by the same factor, so the final weighted sum is identical.

That is the whole trick, and it generalises to any number of blocks. Nothing is approximated. The running maximum only ever increases, the rescaling factor only ever shrinks earlier contributions, and the exponentials stay numerically safe throughout.

Using it in practice

In most codebases you do not implement FlashAttention yourself. You enable it through a library that dispatches to a fused kernel when the shapes, dtype, and mask allow it. In PyTorch, the scaled dot-product attention operator will typically select a flash-style kernel automatically for suitable inputs:

import torch
import torch.nn.functional as F

# q, k, v: (batch, heads, seq_len, head_dim), float16 or bfloat16
q = torch.randn(2, 8, 4096, 64, device="cuda", dtype=torch.float16)
k = torch.randn(2, 8, 4096, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 8, 4096, 64, device="cuda", dtype=torch.float16)

out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(out.shape)  # torch.Size([2, 8, 4096, 64])

A few practical points matter more than the call itself:

• Dtype matters. Flash kernels are written for 16-bit types. In float32, the dispatch usually falls back to a slower path, and the memory benefit shrinks.
• Causal masking is cheap, arbitrary masks are not. A causal mask can be applied tile by tile with a simple rule about block positions. An arbitrary attention mask may force materialisation or a fallback kernel, so check what your implementation actually dispatches to.
• Head dimension has limits. Fused kernels support a restricted set of head dimensions, commonly up to 256. Unusual head sizes can silently fall back.
• Verify the dispatch. Do not assume the fast path is active. Profiling tools and backend introspection will tell you which kernel ran. A silent fallback gives correct results at standard speed, which is the most confusing kind of performance bug.

FlashAttention 2 refined the original design with better work partitioning across thread blocks and a reduced number of non-matrix-multiplication operations, and later variants added support for grouped-query attention and other patterns. The core idea did not change: tile, keep intermediates on chip, rescale running statistics, recompute in the backward pass.

What FlashAttention does not fix

It is important to be clear about the limits, because FlashAttention is sometimes described as if it made attention linear. It did not.

First, the arithmetic is still quadratic in sequence length. For very long sequences, the operation count itself eventually dominates, and no amount of tiling removes that. FlashAttention pushes the practical limit much further out, from a few thousand tokens to tens or hundreds of thousands depending on the model and hardware, but it does not change the asymptotic scaling.

Second, the key-value cache in inference still grows linearly with sequence length. FlashAttention speeds up computing attention over a long context, but it does not compress the cache. Long-context inference still needs cache management, eviction, or compression strategies on top.

Third, tiling helps most when the operation is memory-bound. For short sequences, the score matrix is small, the overhead of a fused kernel can dominate, and the speedup is modest. The benefit grows with sequence length, which is exactly where standard attention hurts most.

Finally, exactness cuts both ways. Because FlashAttention computes the same result as standard attention, it inherits the same modelling properties, including the same failure modes. If attention itself is the wrong inductive bias for a task, a faster implementation will not fix it. Sparse attention, linear attention, and state space models address that different question, and they trade exactness for a different scaling law.

The broader lesson

FlashAttention is a hardware-aware algorithm, and that phrase deserves unpacking. Most algorithm analysis counts operations and ignores where data lives. On modern accelerators, where data lives often matters more. An algorithm that is optimal in operation count can be slow in wall-clock time if it streams large intermediates through slow memory, while an algorithm with extra arithmetic can be fast if it keeps its working set on chip.

That lesson generalises well beyond attention. Whenever you profile a model and find the arithmetic units idle, ask where the data is going, not how many operations you are doing. Fusing operations, tiling to fit fast memory, and recomputing cheap values instead of storing expensive ones are the recurring moves. FlashAttention just applies them to the operation that needed them most.

Takeaway

Standard attention is slow for long sequences mainly because it materialises an N×N matrix in slow memory and reads it back several times. FlashAttention tiles the computation into SRAM-sized blocks, uses online softmax with running maxima and rescaled partial outputs to make tiling possible, and recomputes score blocks in the backward pass instead of storing them. The output is exact, the memory is linear in sequence length, and the speedup comes from cutting HBM traffic, not from approximating attention. When you next reach for a longer context window, that is the mechanism making it affordable.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial