![\"A](\"/b0f33048b0fa307d4b162edbb8b81774/memory-hierarchy.svg\")
When I wrote about why GPU memory runs out during long LLM runs, the culprit was the KV cache: memory that grows with every token you keep around. This post covers a different memory problem in the same layer. It shows up during the attention computation itself, before anything is cached.
\nHere is the puzzle FlashAttention set out to solve. The attention math in a transformer is not that many floating-point operations (FLOPs), and a modern GPU does those operations blindingly fast. Yet for a long time the attention layer was one of the slowest parts of training and inference, and it got quadratically slower as the sequence grew. If the GPU has compute to spare, where does the time go?
\nThe answer is that attention is memory-bound, not compute-bound: the GPU spends most of its time moving data, not doing arithmetic. This post is for engineers who serve or train transformer models and want to understand why. I walk through:
\nThe reference throughout is the original paper, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022).
\n“Memory-bound” only means something once you know that a GPU has more than one kind of memory, and that they differ in speed by more than an order of magnitude.
\nHBM (high-bandwidth memory) is the large pool you think of as “the GPU’s memory.” On an A100 that is 40 to 80 GB, running at roughly 1.5 to 2.0 TB/s. When you read that a model “doesn’t fit in 40 GB,” this is the 40 GB in question.
\nSRAM is the on-chip memory right next to the compute units. It is tiny: about 192 KB per streaming multiprocessor (one of the GPU’s many independent compute blocks) and around 20 MB across the whole A100. But it runs at roughly 19 TB/s, an order of magnitude faster than HBM. (These figures come from the FlashAttention paper, Section 2.) SRAM is where computation actually happens fast; HBM is the warehouse you fetch from.
\nThe gap between those two bandwidths is the whole story. If an algorithm keeps shuttling data to and from HBM, the compute units sit idle waiting for it, no matter how many teraFLOPs the chip is rated for. This is a memory hierarchy in exactly the same sense as the CPU cache hierarchy: small and fast at the top, large and slow below. The same performance lesson applies. You win by keeping working data in the fast tier and touching the slow tier as little as possible.
\n
Attention takes three matrices, the queries Q, keys K, and values V, and computes softmax(Q Kᵀ) V. Written out the obvious way, following the original transformer formulation, it takes three steps:
S = Q Kᵀ. For a sequence of length N, this is an N × N matrix.P = softmax(S), still N × N.O = P V to get the output.The problem is steps 1 and 2 together. The N × N matrix is written out to HBM, read back to apply the softmax, written out again, and read back once more for the final multiply. At a sequence length of 4,000, that matrix has 16 million entries per attention head, and it sits in HBM for the whole round trip.
Here is the part that surprises people: the matrix multiplies are not the bottleneck; the softmax is. Softmax, dropout, and masking are elementwise operations. Each does cheap arithmetic per element, so its speed is limited almost entirely by how fast it can read and write its operands. The paper measures this directly: the elementwise steps dominate wall-clock time even though they do a tiny fraction of the FLOPs. The GPU is starved for data, not for compute.
\nThe memory cost is also quadratic. Materializing (fully storing) the N × N matrix means attention uses O(N²) memory, which is why long-context models were historically so expensive. Double the sequence and you quadruple the matrix.
Once you see it, the insight is almost annoyingly simple: if the N × N matrix is what’s killing you, don’t create it. FlashAttention computes exactly the same output without ever materializing the full attention matrix in HBM.
It does this with tiling, which means working on the matrices in small blocks:
\nQ, K, and V into blocks small enough to fit in SRAM.The scores never leave the chip.
\n![\"The](\"/6f0c14b2658e46554a164ba44f21e0e0/tiling-loop.svg\")
The catch, and the clever part, is the softmax. Softmax normalizes each row by a sum over the entire row, and it first subtracts the row maximum for numerical stability. If SRAM only ever holds one block of the row at a time, you don’t know the true maximum or the true sum until you’ve seen every block. So how can a single pass produce an exact result?
\nThe answer is the online softmax, a running-statistics trick that predates FlashAttention. It comes from Milakov and Gimelshein’s Online normalizer calculation for softmax (2018). Instead of waiting for the whole row, you keep three running values for each query row:
\nm, the maximum score seen so far,ℓ, the running sum of exp(score − m), which is the softmax denominator,O, the running weighted sum of the value vectors.When a new block contains a score larger than the current m, you rescale the values you already have by the correct factor and carry on. Because every partial result is rescaled to a consistent reference whenever the max moves, the value you end up with after the last block is identical to a full-row softmax.
That property is why you can drop FlashAttention into an existing model with no retraining. It is exact attention, not an approximation: same numbers, different memory schedule.
\nTraining needs one more piece. The backward pass normally uses the N × N attention matrix that the forward pass produced, and FlashAttention never stored it. So it recomputes the relevant blocks on the fly during the backward pass, from Q, K, and V, which are cheap to reload.
Spending a bit of extra compute to save a lot of memory traffic is a good deal precisely because the layer is memory-bound. The recomputed matmuls are close to free compared with the HBM reads and writes they save.
\nThe gains come in two forms, and they compound.
\nMemory drops from quadratic to linear in sequence length. Because the big matrix is never stored, FlashAttention’s memory use is O(N) instead of O(N²). The paper reports roughly 10× less memory at a sequence length of 2K and 20× at 4K (FlashAttention paper, Section 4). That is the difference between fitting a long context and hitting a CUDA out-of-memory error. It’s the same failure I traced to the KV cache in the earlier post, showing up here for a different reason.
\nWall-clock time drops too, even though recomputation adds FLOPs, because FLOPs were never the bottleneck. The paper reports:
\nIt also made sequence lengths practical that weren’t before, reaching better-than-chance results on the 16K-token Path-X task.
\nThe follow-ups keep the core trick and get more out of the GPU.
\nFlashAttention-2 (2023) reworks how the computation is split across the GPU’s parallel units:
\nThe result is up to about 230 TFLOPs/s on an A100 forward pass, roughly 73% of the theoretical peak, up from the 25% to 40% the first version reached.
\nFlashAttention-3 (2024) targets the Hopper architecture (H100). It overlaps the matmuls with the softmax using the chip’s asynchronous instructions and warp specialization, and it adds FP8 support. Together these give roughly 1.5 to 2 times the throughput of FlashAttention-2 in FP16, and push FP8 attention toward 1.2 PFLOPs/s.
\nEach version justified a rewrite because “memory-bound” is a moving target. As you claw back HBM traffic, the next bottleneck appears somewhere else: first warp scheduling, then instruction-level overlap. Each release chases the newly exposed one.
\nFlashAttention is not a free win everywhere. Three things are worth knowing before you assume it is.
\nIt’s a kernel, and kernels are hardware-specific. The speedups above depend on particular GPUs and particular data types. The optimized FP8 path lives on Hopper; older cards fall back to slower paths. If your model runs on mixed hardware, which is normal in a shared HPC cluster, you get a range of speedups, not one number. Check what your actual GPUs support instead of quoting the headline figure.
\nIt speeds up attention, not the KV cache. These are two separate memory stories in the same layer. FlashAttention avoids materializing the attention matrix during the computation. It does nothing about the KV cache that grows as you generate tokens. On long autoregressive runs, the KV cache is still what eats your HBM, and the fixes for that (shorter context, cache eviction, a quantized cache) are a different toolbox.
\nYou usually consume it rather than write it. Unless you do kernel work, FlashAttention reaches you through a framework: PyTorch’s scaled_dot_product_attention, an inference server like vLLM, or the flash-attention library directly. The useful skill is knowing whether it’s actually being used. A silent fallback to the naive path (wrong dtype, an unsupported mask, a head dimension the kernel doesn’t handle) quietly costs you the speedup, and profiling is the only way to find out.
Most of my work sits a layer or two above the attention kernel. Archi, the RAG copilot I built for CMS computing operations at CERN, and tools like LLM DevMate live well above it. But the abstraction leaks downward the moment latency or GPU memory becomes the constraint, and on shared HPC hardware it always eventually does. Knowing that attention is memory-bound lets you read a profiler trace and tell “the model is too big” apart from “the kernel isn’t the one you think it is.”
\nThe lesson generalizes past attention. When something is slower than the arithmetic says it should be, the answer is usually in the memory hierarchy, not the FLOP count. That reframing, from counting operations to counting data movement, is the most useful thing FlashAttention taught me, and it applies far outside the transformer.
\nDiagrams by M. Hassan Ahmed, released under CC0. No external image was used for this post; the figures are original work by the author.
","frontmatter":{"title":"How FlashAttention Speeds Up the Attention Layer","date":"2026-08-11T00:00:00.000Z","description":"Attention is memory-bound, not compute-bound. Here's how FlashAttention uses tiling and online softmax to skip the N×N matrix and run exact attention faster.","thumbnail":{"childImageSharp":{"fluid":{"base64":"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABQAAAALCAIAAADwazoUAAAACXBIWXMAAAsSAAALEgHS3X78AAAB0klEQVQoz0XQa6+aQBAGYNK0gsIuiwJ7X64iyE2PtrX1cjym3/p/+vs7FpMmTybDsu/sgrU1HeiS4bq7fRwfj6+P++EO9f3tDs79GRbPw+W2v13G623/PplS1szXtq+hsnzc7G7VeCm6n3n3oxqv0Gfb02Z3zdpTvX8v+jNsWPcXsznOsIKU5QVmgsOMxIUf5UFcLmkJjy7RbmBsJGwkHSznvpo4WCx8uSDaeu4gGi0TU+51Psi0U9lgijHk1QzxMCoS06fJIGUbszrmNWU1FQ2XLaWVtXgOk3OiHF/ZWH5B4pPLfGLCVQoy3dWb07g95/kx5i0THRW91KPJDkr1EJYeUctVItla8srw6lf9XYktWmUk0IY3KqqyeMNWZUI3QETrkjcla+AK1hwLQpSKiyEb+rT/Vuz+nH9rPbhhYSPKWZWbXZ7stGhTPQDFt6kZ82QPB0CYO4g7UDH8FWFj8RnxBVFwI3iFiYLzg2UCojgXqqG8AqtV6gfamiMGFli4RE284Fkh7HgxDnORv6Fo7YUlYTXhjU+flYgWRQWE6T9s4YsnLF6NLxw3CmhV9neWHUCc7CEW6gFAH8A3w/gXRCfTOMej07rtxnPMXhD9v9+L/wIza0l/29XgfQAAAABJRU5ErkJggg==","aspectRatio":1.899441340782123,"src":"/static/2d6b9a4a8b053eec65ba5b226b112ebe/40a76/hero.png","srcSet":"/static/2d6b9a4a8b053eec65ba5b226b112ebe/c972b/hero.png 340w,\n/static/2d6b9a4a8b053eec65ba5b226b112ebe/27625/hero.png 680w,\n/static/2d6b9a4a8b053eec65ba5b226b112ebe/40a76/hero.png 1360w,\n/static/2d6b9a4a8b053eec65ba5b226b112ebe/ed396/hero.png 2000w","sizes":"(max-width: 1360px) 100vw, 1360px"}}}}}},"pageContext":{"slug":"/2026-08-11-how-flashattention-speeds-up-attention/","previous":"blog/2026-08-12-matryoshka-embeddings-rag/","next":"blog/2026-08-14-fastapi-token-bucket-rate-limiting/"}},"staticQueryHashes":["32046230"]}