II · THE IDEA · ARTIFICIAL INTELLIGENCE
Why Attention Gets Expensive
▶ Listen · narrated
A model reading one thousand tokens performs roughly a million comparisons. Read two thousand and that becomes four million. The arithmetic is simple; the consequences shape everything.
At a glance
- Complexity
- O(n²) in sequence length for both time and memory
- Dominant cost
- Moving data between memory levels, not the multiplications themselves
- FlashAttention gain
- 2–4× faster on long sequences by reducing memory round-trips
- Why it matters
- Attention cost determines context window limits and inference speed
Imagine a room where everyone must shake hands with everyone else. Ten people means 45 handshakes. Twenty people means 190. The number grows much faster than the headcount, because each new person must greet everyone already there. Attention works the same way: every token must compare itself to every other token to decide how much to care about it. Double the number of tokens and you quadruple the comparisons. This is why longer context windows are expensive — not just a bit more expensive, but much more expensive, in a way that grows rapidly. FlashAttention does not reduce the number of handshakes; it organises them so people do not have to walk across the room as often, which saves time even though the total work is the same.
Attention computes a score matrix S = QK^T of size (sequence_length × sequence_length), applies softmax row-wise, then multiplies by V to produce the output. For a sequence of length n, this is O(n²) in both time and space. The forward pass requires loading Q, K and V from HBM, computing S, writing S back to HBM, loading S and V again, then computing the output. Each HBM access is slow relative to the on-chip arithmetic.
FlashAttention tiles the computation: it divides Q, K, V into blocks that fit in SRAM, computes attention for each block entirely on-chip using an online softmax algorithm, and accumulates the output without materialising the full S matrix in HBM. This reduces HBM accesses from O(n²) to O(n² / M) where M is SRAM size, typically yielding a 2–4× speedup on sequences of several thousand tokens. FlashAttention-2 improves parallelism across the batch and head dimensions and reduces non-matmul operations.
The quadratic remains: longer sequences still cost n² in both FLOPs and memory, and very long contexts (100k+ tokens) remain expensive even with FlashAttention. Approximate methods — sparse attention, linear attention, state-space models — reduce the exponent but change the model's expressiveness. Exact attention is still standard for most applications because the tradeoffs of approximations are task-dependent and hard to predict.
Look closer
The quadratic is in comparisons, not just arithmetic
To compute attention for a sequence of n tokens, each token's query must be compared with every token's key — n² dot products in total. Then those scores are used to weight n² value contributions. The matrix holding those scores is n by n, so even storing the intermediate result requires memory that grows quadratically. This is not an implementation detail that can be optimised away; it is intrinsic to the mechanism of attending to every position.
Memory bandwidth is the real bottleneck
Modern accelerators can perform arithmetic far faster than they can move data between high-bandwidth memory (HBM) and on-chip cache. Standard attention implementations load the query, key and value matrices from HBM multiple times during a single forward pass, and each load is slow relative to the computation it enables. FlashAttention restructures the calculation so that blocks of the attention matrix are computed entirely in fast on-chip memory, written out once, then discarded. The arithmetic is identical; the memory access pattern is not.
The constant in front of n² is not small
Attention requires four matrices — queries, keys, values and output — each of size (sequence length × head dimension). For a model with 32 heads and head dimension 128, that is 4096 numbers per token, and every one must be read and written multiple times. At a sequence length of 8192 tokens, the intermediate attention scores alone occupy 256 megabytes per layer as 32-bit floats, before any activations or parameters are counted. The quadratic is real, but so is the size of the data being squared.
The story
Attention works by letting every token look at every other token and decide how much to care about each one. That decision is made by comparing the token's query vector to every other token's key vector, producing a score, then using those scores to mix together the value vectors. The result is a weighted sum: each token's output is built from contributions from every position, weighted by relevance.
The trouble is in the phrase "every other token". If you have ten tokens, that is ten queries each checking ten keys — one hundred comparisons. One hundred tokens means ten thousand comparisons. One thousand tokens means a million. The growth is quadratic, and it applies to both the number of operations and the amount of memory needed to hold the intermediate scores.
For years this was simply accepted as the cost of the mechanism. Attention was expensive, so context windows stayed short — often just 512 or 1024 tokens — and researchers worked on approximations that avoided computing the full n² matrix. Some methods sampled a subset of positions; others imposed structure, letting each token attend only to nearby tokens or to a fixed set of global ones. These reduced the exponent, but they also changed what the model could learn.
Then in 2022, Tri Dao and collaborators published FlashAttention, which made a different observation. The quadratic could not be eliminated, but the way the computation moved data could be changed. Standard implementations computed the full attention matrix, wrote it to memory, then read it back to compute the output. FlashAttention instead divided the work into blocks small enough to fit in the GPU's fast on-chip memory, computed each block completely, and never materialised the full matrix in slow memory. The number of arithmetic operations stayed the same, but the number of memory round-trips dropped sharply.
The speedup was substantial — often two to four times faster on sequences of several thousand tokens — and it came with a reduction in memory usage that made longer contexts feasible. FlashAttention-2, published in 2023, refined the parallelism and pushed the gains further. Both papers emphasised that the improvement came not from approximating attention but from respecting the hardware: computation is cheap, but moving data is not, and the cost structure of modern accelerators makes memory bandwidth the limiting factor.
This matters because attention cost is not an isolated problem. It determines how long a context window a model can handle, how quickly it can generate text, and how much it costs to run. Every token added to the context imposes a quadratic penalty, so a model that comfortably processes two thousand tokens may struggle at eight thousand, and sixteen thousand may be impractical without architectural changes or optimisations like FlashAttention. The exponent is the headline, but the constant in front of it — the actual size of the matrices, the memory bandwidth, the number of layers — is what turns the mathematics into an engineering constraint.
Why it mattered then
When the transformer architecture was introduced in 2017, attention's quadratic cost was noted but not yet a crisis. Early models worked with sequences of a few hundred tokens, and the benefits of attending to every position outweighed the expense. As models grew larger and tasks demanded longer contexts — full documents, conversations with history, code repositories — the quadratic became prohibitive. Approximate attention mechanisms proliferated, each trading off some expressiveness for better scaling, but none became standard because the tradeoffs were hard to predict and often task-dependent. FlashAttention mattered in 2022 because it was the first widely adopted method that kept exact attention while making it fast enough for sequences of thousands of tokens. It did not change the mathematics; it changed what was practical.
Why it matters now
Attention cost still determines the frontier of what models can do. Context windows have grown from a few thousand tokens to a hundred thousand or more, but every doubling remains expensive, and the quadratic has not gone away. Techniques like FlashAttention are now built into most training and inference frameworks, but they are optimisations, not solutions. The underlying scaling behaviour means that very long contexts — entire books, days of conversation history, large codebases — remain costly, and researchers continue to explore alternatives: sparse attention, linear attention, state-space models that avoid the quadratic entirely. Understanding why attention is expensive clarifies why those alternatives exist and what they give up. The quadratic is not a bug; it is the cost of letting every token see everything. Whether that cost is worth paying depends on the task, and on how much hardware you can afford.
The surprising detail
FlashAttention's core insight — that reordering memory access can double speed without changing the arithmetic — was not a deep learning innovation. It came from classical high-performance computing, where the gap between arithmetic throughput and memory bandwidth has been widening for decades. The attention mechanism is unusual in how clearly it exposes this gap, because the intermediate matrix is large and short-lived, written once and read once. Most neural network layers do not have this structure, so the same optimisation does not apply as cleanly elsewhere. The lesson generalises, though: on modern accelerators, the cost of an operation is often dominated by how many times you move the data, not by how many times you multiply it.
Remember this
Attention cost grows with the square of the sequence length, and memory bandwidth often matters more than arithmetic. This is why context windows have limits.
Test yourself
A model with 8192-token context uses exact attention. You want to double the context to 16384 tokens. Attention memory grows by what factor, and what is one architectural change that could reduce the exponent rather than optimising the constant?
Memory grows by a factor of four, because doubling sequence length squares the size of the attention matrix (from 8192² to 16384²). FlashAttention and similar methods reduce the constant by improving memory access, but they do not change the exponent. To reduce the exponent, you need sparse attention — letting each token attend to only a subset of positions, such as a local window or a fixed set of global tokens — or a different architecture entirely, such as state-space models or linear attention variants. These change what the model can learn, which is why they are not simply better: they trade expressiveness for efficiency, and the right choice depends on whether the task actually requires every token to see every other token.
Go deeper
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness · arXiv · Tri Dao et al. · 2022-05-27
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning · arXiv · Tri Dao et al. · 2023-07-17
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.