Why is standard attention memory-bound, and how does FlashAttention fix it?

In this article

1. Before you start
2. Why interviewers ask this
3. Memory-bound attention from first principles
4. The Staff-level answer
5. The rubric
6. Three variants of this question
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page sits in the transformer architecture path and builds on the neural network basics.

Why interviewers ask this

The question separates candidates who know attention as a formula from candidates who know how it runs on a GPU. Anyone can say attention is quadratic in sequence length. The interview signal is whether you can say that the quadratic matrix is the problem because of where it is stored, not because of how many multiplications it needs.

A weak answer says "FlashAttention is a faster approximation of attention." A strong answer says it is exact, names the slow and the fast memory, and explains that it removes memory traffic rather than arithmetic. A Staff-level answer adds the mechanism that makes this possible, an online softmax, and the cost: a more complex kernel and a backward pass that recomputes values instead of storing them.

Memory-bound attention from first principles

IO-aware

IO-aware is a design property of an algorithm that accounts for the cost of moving data between slow and fast memory, so it counts memory reads and writes, not only arithmetic, and reorganises the work to keep data in fast memory.

In plain words

Picture a chef who cooks very fast but keeps all ingredients in a warehouse across town. A small counter sits next to the stove. If every dish needs a trip to the warehouse, the chef spends the day driving. Cooking speed stops mattering.

A GPU is the chef. The warehouse is large, slow memory. The counter is small, fast memory. Standard attention makes many trips for a few chops each. The analogy breaks in two places. Ingredients are used up, but data in memory is not, so the same bytes can be fetched again and the cost is the trips, not the stock. And a GPU has thousands of chefs sharing the same road, so the road itself is the limit.

In one precise sentence, standard attention is memory-bound because its large intermediate matrices are written to and read from slow memory, and each byte moved supports only a few operations.

Step 1: Standard attention writes and rereads the N×N matrices

Scaled dot-product attention, defined by Vaswani et al. (2017), is the softmax of the scaled product of the query and key matrices, multiplied by the value matrix, where Q, K and V are the query, key and value matrices and N is the number of tokens.

A GPU has two memory levels that matter here. HBM (high bandwidth memory) is the large, slower memory that holds tensors. SRAM (static random-access memory) is the small, fast on-chip memory (shared memory) that tiles are staged in.

Standard attention runs three separate kernels. The first multiplies Q by K and writes the N×N score matrix S to HBM. The second reads S, applies the softmax and writes the N×N weight matrix P. The third reads P, multiplies by V and writes the output. For N = 4,096 that matrix has 16,777,216 entries, which is 32 mebibytes per head in fp16 (illustrative setup, one head).¹

1 · Matrices in HBM2 · Waiting3 · Tiles4 · OutputReplay

Step-by-step description

1. Standard attention writes S and P to HBM.
2. Few operations per byte, so the chip waits.
3. Tiles stay in SRAM with a running max and sum.
4. The output is written once.

• Standard attention writes the score matrix and the weight matrix to HBM.
• The chip then waits on memory.
• FlashAttention keeps a tile in SRAM with a running max and sum.

Step 2: Few operations per byte moved make it memory-bound

The chip's machine balance is its peak floating-point operations (FLOPs) per second divided by its memory bytes per second. Take an illustrative accelerator with 300 trillion FLOPs per second and 2 terabytes per second: its balance is 150 FLOPs per byte.¹ A kernel needs an arithmetic intensity above that to be compute-bound.

The softmax step is far below it. If each element costs about 5 FLOPs (illustrative: subtract the max, exponentiate, add to the sum, divide) and moves 4 bytes (read 2, write 2 in fp16), its intensity is 1.25 FLOPs per byte.¹

These figures come from a simplified model that ignores masks, dropout and overlap of compute with transfers.¹ The sizes use two units: four passes over a 32-mebibyte matrix come to 134.2 megabytes, and Q, K, V and the output add about 2 megabytes, giving 136.3.¹

The payoff of the argument is in the last two rows. At both lengths the tiled path's memory time (13.1 µs at N = 4,096) is below the arithmetic time (14.3 µs), so in this model tiling turns attention from memory-bound to compute-bound.¹

Step 3: Tiling with an online softmax keeps the matrix on chip

Input/output (IO) cost is the reads and writes between levels of GPU memory. IO-aware is a design approach that treats that cost as the quantity to minimise, and FlashAttention is an exact attention algorithm built that way. Dao et al. (2022) use tiling, which is splitting Q, K and V into blocks small enough to sit in SRAM, to cut the number of HBM reads and writes.

The obstacle is the softmax, which needs a whole row of scores before it can normalise. The fix is the online softmax, a streaming method from Milakov and Gimelshein (2018) that keeps a running maximum and a running sum and rescales them as each new block arrives. When a new block raises the maximum from m_old to m_new, FlashAttention multiplies both the running sum and the partial output (the running weighted sum of the V rows) by exp(m_old - m_new), then adds the new block's contribution. Dividing the final output by the final sum gives exactly softmax(S)V.

A toy run, and the traffic it saves

A toy run shows the online softmax matching a single pass. Scores 2, 1, 0 and 3 (illustrative) arrive in two blocks: the first block sets the maximum at 2, and the second raises it to 3, so the earlier sum and partial output are rescaled. The final running maximum is 3 and the running sum is 1.553, the same as one pass over all four, and the weights are 0.2369, 0.0871, 0.0321 and 0.6439.¹

Each tile of scores is used and discarded inside SRAM, so the N×N matrix never exists in HBM. With blocks of 256 columns (illustrative, from an SRAM of 65,536 numbers), N = 4,096 needs 16 passes and about 26.2 megabytes of traffic, roughly 5.2 times less than the standard path.¹

Text description of the diagram

• Standard attention moves 136.3 megabytes at 4,096 tokens and 541.1 megabytes at 8,192 tokens.
• Tiled attention moves 26.2 megabytes at 4,096 tokens and 102.8 megabytes at 8,192 tokens.

Step 4: The result is exact, and the backward pass recomputes

The output is the same as standard attention up to ordinary floating-point rounding, because the online softmax is algebraically identical to the full one. The paper's abstract calls the algorithm exact and says it needs fewer HBM accesses than standard attention. The arithmetic is not reduced: the work is still quadratic in N, which the page on why attention is quadratic covers, along with the methods that do trade quality for lower compute.

For training, the backward pass needs the scores. Rather than store the N×N matrices, FlashAttention keeps small per-row statistics and recomputes the scores tile by tile in SRAM. That adds arithmetic but removes the large HBM transfers, a good trade when memory is the limit. The memory held for attention then grows linearly with N instead of quadratically, and the output is written to HBM once, at the end.

The Staff-level answer

What makes it Staff-level: it names the bottleneck as data movement, states what does not change (exactness, quadratic compute) and gives the cost of the fix.

The rubric

Three variants of this question

Is FlashAttention an approximation?

No. FlashAttention returns the same attention output as the standard formula up to floating-point rounding. Dao et al. (2022) describe it as an exact attention algorithm. Only the block-sparse extension in the same paper is approximate.

Does FlashAttention make attention linear in sequence length?

Memory use becomes linear, but compute does not. Every query still meets every key, so the arithmetic is quadratic. Longer prompts still cost more; see the context window and its cost.

What if the sequence length doubles?

The score matrix quadruples. The standard path's traffic does too: in the model above it rises from 136.3 to 541.1 megabytes.¹ The tiled path's traffic also rises about fourfold, from 26.2 to 102.8 megabytes, but it stays about five times lower and no N×N matrix is stored.¹

Common mistakes

• "FlashAttention approximates attention to save time." (wrong) It would then change model outputs. (right) "It is exact; it reorders the computation to cut memory traffic."
• "It reduces the FLOPs, so attention is no longer quadratic." (wrong) The arithmetic is unchanged. (right) "Compute is still quadratic; memory traffic and footprint drop."
• "Attention is slow because matrix multiplication is slow." (wrong) In the model above the matmul FLOPs alone would take 14.3 µs at N = 4,096, while moving the standard traffic takes 68.2 µs.¹ (right) "The softmax and the round trips through slow memory dominate."
• "Recomputing scores in the backward pass makes it slower." (wrong) It adds arithmetic but removes large transfers. (right) "Recomputation trades cheap FLOPs for expensive memory traffic."

Follow-up questions

How does FlashAttention speed up attention without changing the math?

It reorders the same computation into on-chip tiles processed with an online softmax, so the result is identical up to floating-point rounding. The saving is fewer HBM reads and writes, not fewer operations.

What is online softmax?

Online softmax is a method that computes a softmax in one streaming pass by keeping a running maximum and running sum, and, when a larger maximum appears, rescaling the sum and the partial output by exp(m_old - m_new). It lets each tile contribute to the normaliser without seeing the whole row.

Why does the backward pass recompute the scores?

Storing the N×N scores would put the quadratic matrix back in HBM. Recomputing them from Q, K and small saved statistics costs extra arithmetic but saves more in memory traffic.

What does FlashAttention-2 change?

FlashAttention-2 improves parallelism and work partitioning on the GPU. The idea of tiling to avoid the N×N matrix in HBM stays the same. FlashAttention-3, published in July 2024, adapts the approach to Hopper-generation GPUs using asynchrony and 8-bit floating point (FP8) low precision.

Does FlashAttention help when generating one token at a time?

It helps little with the matrix problem. When one token is generated the query length is 1, so the score matrix is a single row and there is no N×N matrix to avoid. That step is memory-bound on reading the cached keys and values instead, and the kernel work becomes splitting the key-value sequence across thread blocks to keep the GPU busy.

Does the traffic grow with the number of heads?

Yes. Every head builds its own score matrix, and so does every sequence in a batch, so the traffic multiplies by the number of heads and the batch size; see multi-head attention and the reason for dividing scores by the square root of d_k. More questions of this kind are in the large language model (LLM) interview questions.

Sources

1. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv 2205.14135) — arXiv (Dao, Fu, Ermon, Rudra, Ré), 27 May 2022
2. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (arXiv 2307.08691) — arXiv (Dao), 17 Jul 2023
3. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arXiv 2407.08608) — arXiv (Shah et al.), 11 Jul 2024
4. Online normalizer calculation for softmax (arXiv 1805.02867) — arXiv (Milakov and Gimelshein), 8 May 2018
5. Attention Is All You Need (arXiv 1706.03762) — arXiv (Vaswani et al.), 12 Jun 2017