Why attention is HBM-bound

Textbook attention is easy to write and expensive to run. For each head you build an n × n score matrix, softmax it, and multiply by V. The math is correct. For long sequences the memory traffic is the problem: those scores live in GPU HBM, and moving them costs more than the FLOPs.

The big idea

Imagine printing a seating chart for every pair of people at a stadium, carrying it to a warehouse, then bringing it back to seat them. The warehouse is HBM. The seats are SRAM. Classic scientific computing called this "blocking" or "tiling." FlashAttention is that idea applied to attention.

The analogy breaks on FLOPs. Full attention still does O(n² d) math. IO-aware kernels do not make attention linear. They stop warehousing the chart.

Walking through it

For one head, sequence length n, head dimension d:

S = Q K^T / sqrt(d)     # n × n scores
P = softmax(S)          # n × n probabilities
O = P V                 # n × d output

Bytes scale as O(n²) for S and P, plus O(nd) for Q, K, V and O. FLOPs are O(n² d), so arithmetic intensity for the score stage is roughly O(d). That is often not enough to saturate tensor cores once S round-trips through HBM.

Sequence length n = 4096, head dim d = 128, FP16:

Full score matrix per head: 4096² × 2 bytes ≈ 32 MB

At n = 8192 the same matrix is 8192² × 2 ≈ 128 MB per head. A block with 32 heads that wrote full maps would move multiple GB of attention matrices through HBM on one layer, before MLP traffic.

Training also stores activations for the backward pass. Materialized n × n maps blow up memory at long context and force smaller batches.

Common mistakes

• Blaming FLOPs when the kernel is waiting on HBM. Intensity near O(d) is the giveaway.
• Thinking a faster GPU fixes n × n writes. Bandwidth helps, but the quadratic still wins as n grows.
• Confusing this with approximate attention. The problem here is layout and traffic, not the formula.

Check yourself

1. Why is the score stage often memory-bound even though it does O(n² d) FLOPs? Answer: It also moves O(n²) bytes, so intensity is only about O(d), which is below the ridge on modern GPUs.
2. How big is one FP16 score matrix at n = 4096? Answer: 4096² × 2 ≈ 32 MB per head.

Key takeaways

• Standard attention materializes n × n scores in HBM.
• Bytes grow with n². FLOPs grow with n² d. Intensity stays about O(d).
• SRAM is the fast path. HBM is the warehouse you should not use for the score matrix.
• This is why FlashAttention exists. It does not change the FLOP count of full attention.

Further learning

Next, online softmax is the identity that lets you stream the row in tiles.