Chapter 3.1 β€” The Cost of Attention

Contents

  1. Picking up from Part II: architecture solved, cost remains
  2. Why self-attention is quadratic, precisely
  3. Sparse attention: Longformer and BigBird
  4. Linear attention: reformulating the softmax away
  5. FlashAttention: the same answer, computed honestly
  6. FlashAttention-2 and the trend it represents
  7. Interview angle
  8. Self-check questions
  9. Sources

1. Picking up from Part II: architecture solved, cost remains

Part II left us with a specific, settled answer to "what is a Transformer": a stack of blocks, each combining multi-head self-attention with a position-wise feedforward network, wrapped in some normalization scheme, with positional information injected either additively into the input or directly into the attention computation via something like RoPE or ALiBi. That's a complete architectural recipe, and it's the recipe underneath essentially every large language model in production today. But a recipe that works on paper and a recipe that works at the scale of a hundred-thousand-token context window are two different things, and the gap between them is the subject of this chapter.

The gap comes from a single computational fact that Part II mentioned only in passing: self-attention, as originally formulated, requires every token in a sequence to compute a compatibility score against every other token. That's fine when your sequences are a few hundred tokens long. It becomes the dominant cost in the entire model β€” dwarfing the feedforward layers, dwarfing everything else β€” once sequences run into the tens or hundreds of thousands of tokens, which is exactly the direction the field has been pushing as it chases longer documents, longer conversations, and longer chains of reasoning. This chapter is about why that happens and what has been done about it; the next chapter picks up a related but distinct problem, which is that even if you could compute attention over an arbitrarily long sequence for free, the positional encodings you trained with may simply not know what to do with positions they never saw.

2. Why self-attention is quadratic, precisely

Recall the mechanics of scaled dot-product attention: given queries $Q$, keys $K$, and values $V$, each of shape $(n, d)$ for a sequence of length $n$ and head dimension $d$, attention computes $\text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V$. The term $QK^T$ is where the trouble lives. $Q$ is $(n, d)$ and $K^T$ is $(d, n)$, so their product is an $(n, n)$ matrix β€” one entry for every pair of tokens in the sequence. Computing this matrix takes $O(n^2 d)$ time, and simply storing it takes $O(n^2)$ memory. Every other operation in a Transformer block β€” the linear projections, the feedforward layers, the normalization β€” scales linearly in $n$. Attention alone scales quadratically, and once $n$ is large enough, quadratic beats linear no matter how small the constant factor is.

It's worth being precise about where this quadratic cost actually comes from, because it's not a quirk of one particular implementation β€” it's structural. Every token needs to know, for every other token, how relevant that other token is before it can form a weighted combination of values. That's an all-pairs comparison by definition, and all-pairs comparisons over $n$ items cost $O(n^2)$ no matter how cleverly you implement the arithmetic. Doubling the context length doesn't double the attention cost, it quadruples it β€” go from 2K to 128K tokens, a 64x increase in length, and you've bought yourself a roughly 4,096x increase in attention compute and memory, holding everything else fixed β€” the cost scales as $(n'/n)^2$, and $64^2 = 4{,}096$. This is why "just make the context window longer" was never going to be simply a matter of changing a config value; it forces you to either accept a punishing cost, approximate the computation, or find a fundamentally cheaper way to compute the same thing.

The two families of fix in this chapter attack this problem from different philosophical angles. Sparse attention patterns (Longformer, BigBird) accept an approximation: they compute attention over a carefully chosen subset of token pairs rather than all of them, betting that most of the relevant signal is captured by that subset. Linear attention reformulates the entire computation via kernel feature maps so that it never needs to materialize the $(n,n)$ matrix at all, at the cost of giving up the exact softmax nonlinearity. FlashAttention takes a third position entirely: it changes nothing mathematically, computes the exact same attention output, but restructures how the computation touches GPU memory so that it runs dramatically faster in practice. Understanding why these are three genuinely different strategies β€” approximate the pattern, reformulate the algebra, or optimize the hardware mapping β€” is the core of this chapter.

3. Sparse attention: Longformer and BigBird

The most direct way to avoid computing an $(n,n)$ attention matrix is to simply decide, in advance, that most of those $n^2$ pairs don't need to be computed at all. This is the idea behind sparse attention patterns, and Longformer is a clean illustration of it. Longformer restricts most tokens to a local sliding window: each token attends only to a fixed number of neighbors on either side, which brings the cost of that portion of attention down to $O(n \cdot w)$ for window size $w$ β€” linear in sequence length rather than quadratic. A sliding window alone would be too restrictive, though, since it means information can only propagate a few tokens per layer, and some tasks genuinely need a token to attend to something far away (a classification token needs to see the whole document, a question token needs to see the answer span wherever it sits). Longformer's fix is to designate a small number of "global" tokens β€” task-specific tokens like a [CLS] token, or in some setups specific important tokens β€” that attend to, and are attended to by, every other token in the sequence, unrestricted by the window. This hybrid of cheap local attention plus a sprinkling of expensive-but-rare global attention gives most of the modeling benefit of full attention at a small fraction of the cost.

BigBird generalizes this idea by combining three attention patterns rather than two: a local window (like Longformer), a set of global tokens (like Longformer), and additionally a random attention pattern where each token also attends to a small number of randomly chosen other tokens. The random component matters for a reason that's easy to underappreciate: it improves the graph-theoretic connectivity of the attention pattern. If you think of the attention pattern as a graph where an edge means "this token can pass information to that token," local-plus-global attention alone can still leave certain pairs of tokens many hops apart, meaning information has to route through several layers to connect them. Adding sparse random edges shortens those paths, and the BigBird authors show this connectivity property is what lets a sparse pattern retain the theoretical expressiveness of full attention β€” specifically, they show their sparse attention mechanism is a universal approximator of sequence-to-sequence functions and Turing complete, properties that a naive local-window-only pattern cannot claim.

Both approaches accept an approximation in exchange for going from $O(n^2)$ to $O(n)$ or $O(n \log n)$ in the sequence length, and both were validated on genuinely long-document tasks where the full $O(n^2)$ cost was previously prohibitive. The tradeoff is real, though: choosing which pattern to use (window size, how many global tokens, how much randomness) is itself a design decision that has to fit the task, and getting it wrong means silently discarding attention connections that mattered.

4. Linear attention: reformulating the softmax away

Sparse attention keeps the softmax exactly as it was and simply computes it over fewer pairs. Linear attention takes a different route: it changes the algebra of attention itself so that the $(n,n)$ matrix is never formed, exact or approximate, over any subset. The key move, due to Katharopoulos et al., is to notice that the softmax's role in attention is to turn a similarity score $Q_iK_j^T$ into something like a nonnegative weight, and that this can be replaced with a more general similarity function $\text{sim}(q, k)$ expressed as a kernel: $\text{sim}(q,k) = \phi(q)^T\phi(k)$ for some feature map $\phi$. Once similarity is expressed this way, the attention output for a given query can be rewritten using the associativity of matrix multiplication: instead of computing $(\phi(Q)\phi(K)^T)V$ β€” which still forms an $(n,n)$ intermediate β€” you compute $\phi(Q)(\phi(K)^T V)$, where $\phi(K)^T V$ is a $(d, d)$ matrix that doesn't grow with sequence length at all. That reordering is the entire trick: the same three matrices, multiplied in a different order, never require an $(n,n)$ object to exist in memory, and the total cost drops to $O(n d^2)$ β€” linear in $n$.

This reformulation has a second, arguably more important consequence: it reveals that autoregressive linear attention can be computed exactly like a recurrent neural network. Because the $(d,d)$ matrix $\phi(K)^T V$ can be accumulated incrementally as new tokens arrive β€” add the new token's outer product $\phi(k_t)^T v_t$ to a running sum β€” linear attention transformers can be run at inference time with constant per-token cost and constant memory, independent of how many tokens have been generated so far, exactly like an RNN's hidden state. Katharopoulos et al. frame this explicitly as "transformers are RNNs," and it's a genuinely useful mental model: the quadratic-cost, fully parallel training-time view of attention and the constant-cost, sequential recurrent view of attention are two faces of the same computation, connected by whether or not you exploit the associativity of matrix multiplication.

The price of this efficiency is that you no longer have the exact softmax. Softmax attention has a particular, useful property: it's a normalized, sharply peaked distribution that can, with enough capacity, effectively behave like a hard lookup, attending almost entirely to one or two highly relevant tokens. Kernel feature maps that approximate softmax similarity generally produce smoother, less peaked weightings, and the choice of feature map $\phi$ becomes a real design decision with real quality consequences β€” a poor feature map can measurably hurt a model's ability to do the kind of sharp, selective retrieval that attention is good at. Linear attention variants have seen adoption in specific efficiency-critical settings, but full softmax attention has remained the default in most frontier models, precisely because that peakedness tends to matter for quality in ways that are hard to fully recover with a kernel approximation.

That verdict was accurate as of 2020 and is worth revisiting, because the recurrent view opened up in this section turned out to be the more consequential half of the contribution. If autoregressive linear attention is an RNN with a fixed-size matrix state, then the interesting question isn't which kernel feature map to pick β€” it's what rule you use to update that state, and Katharopoulos et al.'s "add the new outer product to a running sum" is only the simplest of many options. That question is what Mamba, RWKV, and the delta-rule models are all answers to, and Chapter 3.6 takes it up in full.

5. FlashAttention: the same answer, computed honestly

Sparse and linear attention both change what gets computed. FlashAttention changes nothing about what gets computed β€” it computes the exact same softmax attention output as the original formulation, bit for bit up to floating-point tolerance β€” and instead changes how the computation is mapped onto GPU hardware. To see why that alone can yield a large speedup, you need to separate two things that are easy to conflate: how many floating-point operations an algorithm performs, and how long it actually takes to run. Modern GPUs have a small amount of extremely fast on-chip memory (SRAM) and a much larger amount of slower off-chip memory (high-bandwidth memory, HBM). A naive implementation of attention computes the full $(n,n)$ score matrix in HBM, applies softmax to it in HBM, and then multiplies by $V$, reading and writing that entire $(n,n)$ matrix to slow memory multiple times over. For long sequences, this movement of data between HBM and the GPU's compute units β€” not the arithmetic itself β€” is the actual bottleneck. Attention on modern hardware is memory-bandwidth-bound, not compute-bound, which means an algorithm that does the exact same number of floating-point operations but moves less data can be substantially faster in wall-clock time.

FlashAttention's fix is to never materialize the full $(n,n)$ matrix in HBM at all. It works by tiling: it splits the queries, keys, and values into blocks small enough to fit in fast SRAM, and processes attention block by block, computing partial softmax statistics for each block and combining them incrementally as it moves through the sequence. This requires a bit of numerical care β€” you can't compute softmax's normalization until you've seen every score, but you also don't want to hold every score in memory β€” and FlashAttention solves this with an online softmax technique that maintains a running maximum and a running sum, rescaling previously accumulated partial results as new blocks arrive, so that the final result is mathematically identical to a standard softmax computed all at once. The net effect is that FlashAttention performs the same total FLOPs as standard attention (in fact, it does somewhat more, because of recomputation it performs during the backward pass to avoid storing the whole score matrix) but drastically fewer HBM reads and writes, and since HBM bandwidth is the true bottleneck, this translates into large wall-clock speedups and a memory footprint that scales linearly in sequence length rather than quadratically.

This is a genuinely different category of optimization from sparse or linear attention, and it's worth being explicit about the distinction in your own mental model, because interviewers will probe exactly this point: FlashAttention is not an approximation and does not change the model's outputs at all; it is a systems-level optimization of an unchanged mathematical object. That means it composes with everything else in this chapter β€” you can, and people do, apply IO-aware tiling techniques to sparse attention patterns too β€” and it means adopting FlashAttention costs you nothing in model quality, which is a large part of why it was adopted essentially universally, immediately, across the field.

6. FlashAttention-2 and the trend it represents

FlashAttention-2 refines the same core idea rather than replacing it, and the refinements are worth knowing at a high level because they illustrate how much performance was still being left on the table by the original's engineering, separate from any change to the underlying algorithm. FlashAttention-2 improves work partitioning across the GPU's parallel thread blocks and warps, reduces the number of non-matrix-multiply operations (which are comparatively expensive on GPU tensor cores optimized for matrix multiplication), and restructures the loop order so that more of the GPU's streaming multiprocessors stay busy throughout the computation, particularly on the workloads β€” like long-sequence, modest-batch-size inference β€” where the original FlashAttention left GPU utilization noticeably below peak. None of this changes what's being computed; it's a second pass of squeezing more of the theoretical hardware throughput out of the same tiled, IO-aware algorithm.

FlashAttention-3 continues the same lineage onto NVIDIA's Hopper architecture (H100-class GPUs) specifically, exploiting hardware asynchrony that FlashAttention-2's design couldn't target: overlapping the tensor cores doing matrix multiplication with the separate units moving data through memory, instead of strictly alternating between compute and data movement, plus native support for low-precision FP8 execution in the parts of the computation where the accuracy cost is negligible. The result gets close to the peak throughput Hopper's tensor cores can actually deliver on attention, and it's the same pattern FlashAttention-2 already established: each new hardware generation exposes new asynchrony and precision options that the same underlying tiled, IO-aware algorithm can be re-engineered to exploit, without changing what's actually being computed.

The broader lesson to take from FlashAttention and FlashAttention-2 together is that "the cost of attention" is not a single fixed number determined purely by the $O(n^2)$ complexity β€” it's a moving target shaped by how well an algorithm is matched to the memory hierarchy and parallelism of the hardware it runs on. That's a theme that will recur through the rest of this book: architectural decisions and systems decisions are not separable in practice, and a chapter on "the cost of attention" that only discussed asymptotic complexity would miss half of what actually determines whether a given context length is practical to serve. With both the sparse/linear approximation route and the exact IO-aware route on the table, the next chapter turns to a problem these techniques don't solve on their own: even with cheap-enough attention, a model's positional encodings need to behave sensibly at lengths it never trained on, which is a question about generalization, not about compute cost.

7. Interview angle

"Why is self-attention quadratic, and what exactly gets bigger as the quadratic term grows?" A strong answer identifies the $QK^T$ product specifically: computing pairwise compatibility scores between every pair of the $n$ tokens necessarily produces an $(n,n)$ matrix, and both computing and storing it costs $O(n^2)$. It should also note that this is structural β€” it follows from the definition of "every token compares itself against every other token" β€” and is not an artifact of a particular implementation, which is precisely why FlashAttention can fix the constant factor and the memory access pattern but doesn't change the underlying $O(n^2)$ asymptotic FLOP count.

"FlashAttention is often described as making attention 'free.' Is that accurate?" No, and a strong candidate pushes back on the framing: FlashAttention reduces wall-clock time and memory footprint by being IO-aware, not by reducing the number of floating point operations (it actually does somewhat more FLOPs due to recomputation in the backward pass). The saving comes from recognizing that attention on modern GPUs is memory-bandwidth-bound, so minimizing HBM reads/writes via tiling and fused kernels yields large real-world speedups even though the asymptotic compute complexity is unchanged.

"When would you reach for a sparse attention pattern like Longformer's versus FlashAttention?" These solve different problems and a strong answer says so explicitly: FlashAttention is a drop-in exact optimization with zero quality cost, so there's essentially never a reason not to use it. Sparse attention patterns are a genuine approximation that trades some modeling capacity for a change in asymptotic complexity, and are worth it specifically when sequence lengths are long enough, and full-attention memory/compute costs high enough, that even an IO-aware exact implementation is still too expensive β€” and when the task's structure (e.g., long documents where locality plus a few global anchor tokens plausibly captures most of the relevant dependencies) makes the specific sparsity pattern a reasonable inductive bias.

"What do you give up by switching to linear attention, and why hasn't it replaced softmax attention in frontier models?" A strong answer names the mechanism: linear attention replaces the softmax similarity with a kernel feature map so that $\phi(Q)(\phi(K)^TV)$ can be computed without ever forming the $(n,n)$ matrix, giving $O(n)$ cost and an equivalent recurrent formulation for inference. The cost is that most practical kernel feature maps produce smoother, less peaked attention distributions than an exact softmax, which tends to hurt the model's ability to do sharp, selective retrieval over long contexts β€” a capability that in practice has mattered enough that most frontier models have stuck with exact (FlashAttention-accelerated) softmax attention rather than linear attention, reserving linear/recurrent-style attention for settings where extreme sequence lengths make the tradeoff worthwhile.

"BigBird claims Turing completeness for its sparse attention pattern. Why does that matter, and why doesn't a local sliding window alone have this property?" A strong answer connects this to graph connectivity: treating the attention pattern as a graph over tokens, a pure local window keeps information paths short only between nearby tokens, so two far-apart tokens may need many layers to exchange information, effectively limiting the model's expressiveness for tasks requiring long-range dependencies. BigBird's combination of local, global, and random edges keeps the graph's diameter small (any two tokens can reach each other in few hops), which the authors use to argue their sparse pattern retains the same theoretical expressive power (universal approximation, Turing completeness) as full attention, unlike a naive windowed pattern.

8. Self-check questions

  1. Starting from the definition of scaled dot-product attention, explain precisely which operation produces the $O(n^2)$ cost and why it cannot be avoided by a smarter implementation of the same computation.
  2. If you double the context length of a model using standard full attention, by roughly what factor does the attention compute and memory grow, and why does this make "just increase the context window" a nontrivial engineering decision?
  3. Explain how Longformer's local-window-plus-global-token design and BigBird's local-plus-global-plus-random design differ, and why BigBird's random component specifically improves information flow across the sequence.
  4. Walk through the algebraic reordering that makes linear attention linear in sequence length, and explain why this reordering also yields an equivalent recurrent (RNN-like) formulation useful at inference time.
  5. What does it mean for attention to be "memory-bandwidth-bound" rather than "compute-bound" on a GPU, and how does this fact motivate FlashAttention's entire design?
  6. Why is FlashAttention described as computing "the same result faster" rather than as an approximation, and what does that distinction imply about whether it changes model outputs or training dynamics?
  7. Name one respect in which FlashAttention-2 improves on FlashAttention, and explain why that improvement is about hardware utilization rather than about the underlying tiled/IO-aware algorithm.

9. Sources

  • Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The Long-Document Transformer. arXiv:2004.05150. https://arxiv.org/abs/2004.05150
  • Zaheer, M., et al. (2020). Big Bird: Transformers for Longer Sequences. NeurIPS 2020. arXiv:2007.14062. https://arxiv.org/abs/2007.14062
  • Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML 2020. arXiv:2006.16236. https://arxiv.org/abs/2006.16236
  • Dao, T., Fu, D., Ermon, S., Rudra, A., & RΓ©, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135. https://arxiv.org/abs/2205.14135
  • Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024. arXiv:2307.08691. https://arxiv.org/abs/2307.08691
  • Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608. https://arxiv.org/abs/2407.08608

Self-check quiz

Test your understanding of this chapter with a short quiz β€” a mix of concept questions and quick calculations.

Precisely, which operation in scaled dot-product attention produces the $O(n^2)$ cost?

Explanation: Q is (n,d) and Kα΅€ is (d,n), so QKα΅€ is an (n,n) matrix β€” one entry per token pair. That all-pairs comparison is structural, not an implementation quirk.

What is the key difference between how sparse attention (Longformer/BigBird) and linear attention reduce attention's cost?

Explanation: Sparse attention picks a subset of pairs to compute exactly. Linear attention reorders Ο†(Q)(Ο†(K)α΅€V) so no (n,n) object is ever materialized, over any subset, dropping cost to O(ndΒ²).

Why does FlashAttention speed up attention without changing its mathematical output at all?

Explanation: FlashAttention actually performs somewhat more FLOPs (due to backward-pass recomputation) but drastically fewer HBM reads/writes β€” since bandwidth, not arithmetic, is the real bottleneck, this yields large real speedups.

Why does BigBird's random attention component matter, beyond what a local-window-plus-global-token pattern alone provides?

Explanation: Random edges shorten the graph's diameter, which BigBird's authors use to argue their sparse pattern retains full attention's theoretical expressiveness (universal approximation, Turing completeness).

Self-attention computes an $(n,n)$ matrix of pairwise scores. For a sequence of $n = 1000$ tokens, how many entries does this matrix contain?

Explanation: 1000 Γ— 1000 = 1,000,000 pairwise entries.

If you double a model's context length, by what factor does attention compute (and memory) grow, since cost scales as $O(n^2)$?

Explanation: Doubling n means (2n)Β² = 4nΒ² β€” a 4x increase, not 2x.

Going from a 2K-token context to a 128K-token context is a 64x increase in sequence length. By roughly what factor does attention compute increase, since cost scales as $O(n^2)$?

Explanation: 64Β² = 4096 β€” a 64x longer sequence costs roughly 4,096x more in attention compute and memory.

Longformer uses a local window of size $w = 512$ for a sequence of length $n = 100{,}000$. Approximately how many times more expensive is full $O(n^2)$ attention compared to Longformer's $O(n \cdot w)$ local attention β€” i.e., what is $n / w$? (Round to the nearest whole number.)

Explanation: 100,000 / 512 β‰ˆ 195.3, so full attention costs roughly 195x more than the local-window portion of Longformer's attention.