Module 9 — Quadratic cost and efficient attention

Every module so far paid attention only to correctness. The elephant in the room is cost: attention's memory and compute scale as $n^2$ in the sequence length, and that quadratic growth is what stops a plain Transformer from reading a book. This module makes the cost concrete, then surveys the four families of tricks the community uses to fight back.

Where the $n^2$ hides

Look at the attention formula from module 2:

$A = \text{softmax}(QK^\top / \sqrt{d_k}) \in \mathbb{R}^{n \times n}, \quad \text{output} = A V.$

$QK^\top$ is an $n \times n$ matrix. Even before we softmax it, we have to store it. Its size grows as the square of the sequence length. Two consequences:

• Memory to hold the attention scores: $\mathcal{O}(n^2)$ per head per layer. For $n = 4096$, a single head stores 16 million entries; twelve heads and twelve layers is more than two billion entries per batch element.
• FLOPs to compute the attention output: $\mathcal{O}(n^2 d)$ for the score matrix and its product with $V$.

The feed-forward blocks scale as $\mathcal{O}(n d^2)$: linear in $n$, quadratic in $d$. For short sequences they dominate. For long sequences, attention wins the race.

The break-even point in a base Transformer is around $n \approx d$. For $d = 512$, attention becomes the dominant cost around 512 tokens. For $d = 4096$, around 4000 tokens. The larger the model, the later attention takes over, but it always does eventually.

A worked memory estimate

Take a batch of 8, $d = 512$, $h = 8$ heads, $L = 12$ layers, $n = 8192$ tokens.

• Attention scores per layer: $8 \times 8 \times 8192 \times 8192 \times 4$ bytes (float32) $\approx 17$ GB.
• Multiply by twelve layers: more than 200 GB, just for the scores.

That is not a training bill you pay once — it is per forward pass. Even switching to bfloat16 halves it to about 100 GB, still well beyond a single accelerator. Yet vendors advertise 128k, 200k, 1M context windows. Something in the pipeline has to be different from the naive attention we built.

FlashAttention: don't store the scores

Tri Dao's FlashAttention (2022) is the most impactful implementation change of the last five years. It observes that the attention output is much smaller than the score matrix — $n \times d_k$ rather than $n \times n$ — and that we do not actually need to materialise the score matrix to compute the output.

The algorithm walks the sequence in blocks. For each block of queries, it iterates over blocks of keys and values, keeping running statistics for the softmax normalisation. Nothing but the current block's scores lives in memory at any time. Memory drops from $\mathcal{O}(n^2)$ to $\mathcal{O}(n)$; wall-clock speed doubles to triples because the tight inner loop stays in fast on-chip memory rather than paging to slower memory tiers.

FlashAttention is not a new operation. Its output is bit-for-bit compatible with standard attention (up to floating-point reduction order). It is a better implementation of the same math. Every serious training library (PyTorch since 2.0, JAX, TensorFlow's XLA) has it. In your own code:

import torch.nn.functional as F

# PyTorch 2.x picks FlashAttention automatically when the shapes fit.
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

Replacing our hand-rolled attention with scaled_dot_product_attention is the single biggest performance win in module 10's final Transformer.

Sparse and windowed attention

FlashAttention still computes $\mathcal{O}(n^2)$ operations, even if it does not store them. When $n$ is large enough that even the FLOPs bite, the answer is to compute fewer scores. Restrict the pattern of positions each token attends to.

The families are:

Family	Pattern	Cost	Trade-off
Local / sliding window	each token attends to a window of $w$ neighbours	$\mathcal{O}(n w)$	loses long-range dependencies
Strided / dilated	attend to positions at strides $s, 2s, 4s, \dots$	$\mathcal{O}(n \log n)$	hits long range with holes
Global tokens	a few tokens attend to everyone and vice versa	$\mathcal{O}(n g + n w)$	combines with local windows
Learned sparse	routing decides who attends to whom	$\mathcal{O}(n k)$	complex, hardware-dependent

Longformer and BigBird combine local windows with a small set of global tokens. Mistral's sliding-window attention takes the local route with a fixed window. All of them sacrifice something: with a strict window of 1024 tokens, a fact 5000 tokens away does not reach the current position in one hop — it needs several attention layers to propagate.

Linear-time approximations

A separate line of work approximates the softmax directly. Performer rewrites softmax attention as an inner product of feature maps, so the cost becomes linear in $n$. Linformer projects the keys and values to a fixed rank, so the score matrix has bounded size. RWKV and Mamba take the more radical route of dropping attention entirely and using a linear recurrent formulation.

These models are compelling on paper. In practice, they lose one to three points on downstream benchmarks compared to a plain Transformer at equal compute, which is why the industry defaults to FlashAttention + sliding window instead. The linear crowd is catching up fast — Mamba-2 and RWKV-6 are close to par — and may take over eventually.

What "long context" really costs

Vendors quote a maximum context length — 128k, 200k, 1M — and users interpret it as "the model reads all of that". Two effects moderate that reading.

• The KV cache from module 7 grows linearly with context length. For a 70B model at 128k tokens, the cache exceeds the model weights and needs paged attention, off-loading, or key-value compression. Each of those tricks costs quality.
• Effective attention decays with distance. Even with the right positional encoding, attention weights fall off long before the maximum. "Needle in a haystack" evaluations show that models routinely miss information placed in the middle of a long context, a phenomenon called lost in the middle.

Advertised context lengths are ceilings, not guarantees. If your application depends on retrieving information from position 60k in a 128k window, test it explicitly rather than trusting the number on the model card.

Quadratic attention is unforgiving

The most common surprise is going from a 1k context notebook demo to a 32k production request. Everything worked; suddenly the memory spikes to 32 times its previous level for the score matrix alone. Enabling FlashAttention on both training and inference paths is not optional at scale — it is the difference between fitting on one accelerator and spilling to four.

Measure the break-even for your model

Load your model, run forward passes at $n = 512, 1024, 2048, 4096, 8192$, and log peak memory and wall time. Plot both against $n$. You will see the linear part (feed-forward) hand off to the quadratic part (attention) at a specific length, and that length tells you where efficient-attention work will actually pay for itself.

In summary

• Attention costs $\mathcal{O}(n^2)$ in both memory and FLOPs, and dominates the total cost past a break-even length roughly equal to the model width.
• FlashAttention does not change the math; it computes the same result without ever storing the score matrix, cutting memory to $\mathcal{O}(n)$ and roughly doubling speed.
• Sparse, windowed and global-token patterns compute fewer scores at the cost of some long-range reach; linear-time approximations replace attention entirely but currently lose a small amount of quality.
• Advertised long context windows are ceilings: the KV cache and the "lost in the middle" effect mean that reaching the far end of the window in practice requires extra engineering and empirical validation.

Next module: we finally assemble every block written so far into a working Transformer and train it on the date-translation task announced in module 1.