Why attention is limited by memory rather than arithmetic, how FlashAttention computes it exactly in tiles, and how modern models shrink the KV cache, sparsify attention, and stretch to long contexts.

Here is a result that should bother you. Dao et al. (2022) timed one attention layer of GPT-2 medium (sequence length 1,024, 16 heads of width 64, batch 64, forward and backward) on an A100 GPU. The standard PyTorch implementation did 66.6 GFLOPs of arithmetic and took 41.7 ms. Their new kernel, FlashAttention, did more arithmetic, 75.2 GFLOPs, because it recomputes part of the forward pass during the backward pass. It took 7.3 ms. Thirteen percent more work, almost six times faster.

If the GPU's time were set by its arithmetic, that could not happen. The same figure gives the answer in its middle row. The standard kernel moved 40.3 GB between the GPU's main memory and its compute units, and FlashAttention moved 4.4 GB. At the A100's memory bandwidth of 1.5 to 2 TB/s, 40.3 GB takes 20 to 27 ms on its own, most of the 41.7 ms, while 4.4 GB takes about 2 to 3 ms. Attention, as it is usually written, does not wait for arithmetic. It waits for memory.

This section makes that statement exact. We count every number standard attention reads and writes, put each of its kernels on the roofline, define the IO complexity that the rest of the chapter is about, and find the one obstacle, the softmax's normalizer, that stops the obvious fix. The next two sections remove the obstacle.

Two memories, one fast and one big

A GPU has two kinds of memory that matter here. They differ by a factor of thousands in size and about ten in speed. The numbers below are for the A100, from Section 2.1 of the FlashAttention paper, and for the H100, from Table 1 of the FlashAttention-3 paper.

High-bandwidth memory (HBM) holds the model, the activations, and every tensor you create. It is big and, by the standards of the chip, slow. SRAM is the scratchpad inside each streaming multiprocessor, together with the registers, and it is where arithmetic actually happens. Across the whole A100 it adds up to about 20 MB, thousands of times smaller than HBM and about ten times faster. (How SMs, warps, and the memory hierarchy fit together is the subject of §11.1.)

Every GPU kernel follows the same pattern. It loads its inputs from HBM into SRAM and registers, computes, and writes its outputs back to HBM. The next kernel starts by reading them again. So a chain of operations written as separate kernels pays for a round trip to HBM at every link, whether or not the intermediate result was ever needed outside the chain. Keep that picture in mind, because it is the whole story of this section.

The roofline, in one paragraph

A kernel that does $F$ floating-point operations and moves $N_{\text{IO}}$ bytes to and from HBM needs at least $F/\pi$ seconds of arithmetic and $N_{\text{IO}}/\beta$ seconds of memory traffic, where $\pi$ is the peak arithmetic rate in FLOP/s and $\beta$ the memory bandwidth in bytes per second. The two overlap, so a good estimate is

In plain words, what matters is the arithmetic intensity $I$, FLOPs per byte moved. If it is below the machine's ridge point $I^\ast$, the kernel is memory-bound and runs at the speed of HBM, no matter how fast the arithmetic units are. Above it, the kernel is compute-bound.

The ridge point of an A100 is $312\times10^{12} / 2.039\times10^{12} \approx 153$ FLOPs per byte (taking the 80 GB card's 2.039 TB/s), and of an H100 $989/3.35 \approx 295$. A kernel must do 150 to 300 FLOPs for every byte it touches to keep the tensor cores busy. Most kernels do not. The derivation of the roofline and the many ways it is used for tuning are in §11.2. We only need the formula above.

Notice that the ridge point grew from one generation to the next. Arithmetic has been getting faster more quickly than memory, so a kernel that was balanced on one GPU is memory-bound on the next. This is why IO-aware algorithms matter more every year, not less.

Standard attention, kernel by kernel

Fix one head, with queries $\mathbf{Q}\in\R^{T\times d_k}$, keys $\mathbf{K}\in\R^{T\times d_k}$, and values $\mathbf{V}\in\R^{T\times d_k}$ (take $d_v = d_k$ and self-attention, $T_q = T_k = T$, to keep the counts short). The FlashAttention papers write $N$ for the sequence length and $d$ for the head width, so their $N$ is our $T$ and their $d$ is our $d_k$. Attention computes, as in §4.3,

(We follow the FlashAttention papers in writing $\mathbf{P}$ for the attention weights, which §4.3 called $\mathbf{A}$. They are the same matrix.)

The standard implementation, Algorithm 0 in Dao et al. (2022), runs these as three kernels.

1. Read $\mathbf{Q}$ and $\mathbf{K}$ from HBM, compute $\mathbf{S}$, write $\mathbf{S}$ to HBM.
2. Read $\mathbf{S}$, compute $\mathbf{P}$, write $\mathbf{P}$ to HBM.
3. Read $\mathbf{P}$ and $\mathbf{V}$, compute $\mathbf{O}$, write $\mathbf{O}$ to HBM.

Now count, in numbers moved (multiply by 2 for bytes in 16-bit precision). Kernel 1 reads $2Td_k$ and writes $T^2$. Kernel 2 reads $T^2$ and writes $T^2$. Kernel 3 reads $T^2 + Td_k$ and writes $Td_k$. Adding up,

In plain words, the inputs and output are only $4Td_k$ numbers, but the two $T\times T$ matrices are each written once and read once, and for any $T$ much larger than $d_k$ that is almost all the traffic.

We can check the count by running the algorithm with a counter that records every array that crosses between HBM and SRAM. The HBM class below is the instrument we will use for the whole chapter.

At $T = 1{,}024$ and $d_k = 64$, the two square matrices are 94% of all the traffic. At $T = 8{,}192$ they are over 99%.

Real code moves even more

Three kernels is the best case for a non-fused implementation. Code written the obvious way in PyTorch launches a separate kernel for each elementwise step on the $T\times T$ matrix, and each one is another full read and write. A typical training forward pass does the matrix product, a scale by $1/\sqrt{d_k}$, a masked fill for the causal mask, the softmax, and dropout before the final product. With $k$ elementwise kernels on the score matrix (the softmax counts as one), the square-matrix traffic is one write of $\mathbf{S}$, $2k$ for the elementwise kernels, and one read before $\mathbf{P}\mathbf{V}$,

With scale, mask, softmax, and dropout ($k = 4$) that is $10T^2$, two and a half times the best case. The softmax kernel itself is fused (it computes the row maximum, the sum, and the division while the row sits in SRAM), but only because one row of a few thousand numbers fits on chip. That detail will matter in a moment.

And the backward pass is worse. Running the backward pass of §4.3 as one kernel per matrix product means reading $\mathbf{P}$ to form $\mathrm{d}\mathbf{V} = \mathbf{P}^\top\mathrm{d}\mathbf{O}$, writing $\mathrm{d}\mathbf{P} = \mathrm{d}\mathbf{O}\,\mathbf{V}^\top$, reading $\mathbf{P}$ and $\mathrm{d}\mathbf{P}$ to write $\mathrm{d}\mathbf{S}$, and reading $\mathrm{d}\mathbf{S}$ twice more for $\mathrm{d}\mathbf{Q}$ and $\mathrm{d}\mathbf{K}$. The next cell counts it.

Seven passes over a $T\times T$ matrix in the backward pass, on top of the forward pass's four (or ten). Add the memory to store $\mathbf{P}$ between the passes, $hT^2$ numbers per layer, which for Llama 3 8B at 8,192 tokens is 4 GiB per layer (§4.7), and you have the two costs that FlashAttention removes.

Every kernel is below the ridge

§11.2 computed the intensity of naive attention as a whole, about 65 FLOPs per byte for $d_k = 128$ in bf16, well below the ridge. Here we go one level deeper and ask the question kernel by kernel, because the answer shows where the time goes. How much arithmetic does each kernel do per byte? Use $b$ bytes per number ($b = 2$ in bf16) and assume $T \gg d_k$.

Scores, $\mathbf{Q}\mathbf{K}^\top$. This is $2T^2d_k$ FLOPs (one multiply and one add per term) for $b(2Td_k + T^2)$ bytes, so

In plain words, each score costs $2d_k$ FLOPs to compute and $b$ bytes to write, so the intensity can never exceed $2d_k/b$, which is $d_k$ in 16-bit. With $d_k = 64$ or 128 that is 64 or 128 FLOPs per byte, below the A100's ridge of 153 and far below the H100's 295. A matrix multiplication is supposed to be the compute-bound operation par excellence. This one is not, because its output is huge compared with its inner dimension. It is a thin product, like an outer product, and writing the output dominates.

Weighted values, $\mathbf{P}\mathbf{V}$. The same counts in reverse, reading $T^2$ and writing $Td_k$, so the same intensity, about $d_k$ FLOPs per byte.

Softmax. Per element, a comparison for the maximum, a subtraction, an exponential, an addition for the sum, and a division, about 5 operations, for $2b$ bytes (one read and one write). That is about 1.25 FLOPs per byte, more than a hundred times below the ridge. The masking, scaling, and dropout kernels are lower still.

Now put real numbers in. Take one attention layer of Llama 3 8B, with 32 heads of width 128, one sequence of 8,192 tokens, bf16, on an A100, and the minimal three-kernel version.

The softmax, with 1% of the FLOPs, takes about half the time, 4.2 of 8.6 ms. Both matrix products are memory-bound too, though only just, because $d_k = 128$ puts them at 124 FLOPs per byte, close to the ridge. And the last line is the punch line. A single fused kernel that read $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$ once and wrote $\mathbf{O}$ once would have an intensity of $4T^2d_k / (4bTd_k) = T/b$, about 4,000 FLOPs per byte, far above the ridge. It would be compute-bound and take 3.5 ms, the time of the arithmetic alone. Everything above that, more than half of the standard kernel's time even in this minimal form, is waiting for HBM.

Decoding is the extreme case. When a model generates one token at a time, each attention head has a single query, so the "score matrix" is one row of $T$ numbers, and the kernel does about $4d_k$ FLOPs per cached token while reading that token's key and value (with ordinary multi-head attention, one key-value head per query head), $2d_k$ numbers or $4d_k$ bytes in bf16. That is about one FLOP per byte, far below any ridge point, no matter how cleverly the kernel is fused, and the time per generated token is set by how fast the cache streams out of HBM. Fusion cannot fix that. Only a smaller cache can, which is the subject of §12.5.

This is an empirical pattern across the whole Transformer, not only attention. Ivanov et al. (2021) profiled a BERT encoder layer in PyTorch and found that tensor contractions (matrix products) were 99.8% of the FLOPs but only 61% of the runtime. Normalization operators such as softmax and LayerNorm were 0.17% of the FLOPs and 25.5% of the runtime, and elementwise operators 0.03% of the FLOPs and 13.5% of the runtime. FLOPs are a poor predictor of time for anything that is not a large, square matrix product. Fusing chains of memory-bound kernels into one is the general cure, and §11.7 shows how compilers such as torch.compile do it automatically.

def standard_attention_traffic(T, d_k, k=1):
    raise NotImplementedError

def min_head_width(T, ridge, b=2):
    return 128   # placeholder

IO complexity, the quantity that matters

The examples suggest a better yardstick than FLOPs. In the external memory model (Aggarwal and Vitter 1988, building on the red-blue pebble game of Hong and Kung 1981), a computation runs on a fast memory that holds $M$ numbers, backed by a slow memory of unlimited size, and its cost is the number of values moved between the two. On a GPU, the fast memory is an SM's SRAM and registers and the slow memory is HBM. The count is the IO complexity of the algorithm, or its number of HBM accesses.

Theorem 2 of Dao et al. (2022) states the IO complexity of standard attention as

which is what we counted, $4T^2 + 4Td_k$, up to the constant. The $Td_k$ term is unavoidable, since any algorithm must read its inputs and write its output. The $T^2$ term is not reading or writing anything the caller asked for. It is the price of materializing $\mathbf{S}$ and $\mathbf{P}$, intermediate results that the algorithm writes to slow memory only because they do not fit in fast memory.

The same model explains tiled matrix multiplication. Multiplying two $n\times n$ matrices naively, one dot product at a time, moves $\Theta(n^3)$ values, and tiling with blocks that fit in fast memory reduces it to $\Theta(n^3/\sqrt{M})$, which Hong and Kung proved is optimal (§11.5). FlashAttention is the same move applied to attention. Choose tiles of $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ that fit in SRAM, compute each tile's contribution to $\mathbf{O}$ on chip, and never send $\mathbf{S}$ or $\mathbf{P}$ to HBM at all. We will prove in §12.3 that this costs $\Theta(T^2d_k^2/M)$ accesses, and that no exact algorithm can do better for every $M$.

def attention_layer_times(T, d_k, heads, peak, bw, b=2, k=1, softmax_flops=5):
    return {"standard": 0.0, "fused": 0.0}

The obstacle, a normalizer that needs the whole row

So why not fuse the three kernels? Tiling the first product is easy. A tile of queries and a tile of keys produce a tile of $\mathbf{S}$, exactly as in a tiled matrix multiply. The trouble is the next step. Entry $(i, j)$ of $\mathbf{P}$ is

where $m_i$ is the row maximum that keeps the exponentials from overflowing (§4.3) and $\ell_i$ the normalizer. Both depend on every score in row $i$. So before you can turn even one tile of scores into weights, you need statistics over the whole row, and the whole row is $T$ numbers per query, which at $T = 128{,}000$ is far more than an SM's SRAM can hold for a useful number of queries.

There seem to be only two ways out, and both are bad. Make two passes, computing $m_i$ and $\ell_i$ in the first and the output in the second, which means computing every score twice. Or write the scores to HBM and read them back, which is standard attention. The way out of the dilemma is an observation about the softmax itself. Its statistics can be computed one tile at a time, as long as you are willing to correct what you have already accumulated whenever a new tile changes the running maximum. That is the online softmax, and it is the subject of §12.2.

The real way

In PyTorch the fused kernel is one call, torch.nn.functional.scaled_dot_product_attention. On a GPU it dispatches to a FlashAttention kernel when the inputs allow it, and on a CPU to a fused, tiled kernel of the same design. The block below times the three-step version against it on your CPU. The absolute numbers depend on the machine, but the CPU has its own fast and slow memories (caches and DRAM), and the gap you see has the same cause.

The outputs agree to about $10^{-6}$ (float32 rounding), so this is the same function. On a typical laptop the fused version is several times faster, and the gap widens with $T$.

Read more

• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Section 2 describes the GPU memory hierarchy and standard attention's IO cost, and Figure 2 is the measurement this section opens with.
• Williams, Waterman & Patterson (2009). Roofline: An Insightful Visual Performance Model for Multicore Architectures. The original roofline paper, short and still the clearest statement of the model.
• Ivanov et al. (2021). Data Movement Is All You Need: A Case Study on Optimizing Transformers. Profiles BERT operator by operator and shows how memory-bound operators dominate runtime despite their tiny FLOP counts.
• Rabe & Staats (2021). Self-attention Does Not Need O(n^2) Memory. Chunked attention with constant memory per query, the memory half of the FlashAttention story.
• Hong & Kung (1981). I/O Complexity: The Red-Blue Pebble Game. The model of computation behind IO complexity and the lower bound for matrix multiplication.
• NVIDIA (2020). NVIDIA A100 Tensor Core GPU Architecture. The source for SM counts, shared memory sizes, and peak throughputs.

Watch next

▶

▶