Vol. IV · Attention Plate 7 of 9attention cost, quadratic, FLOPs, memory, sequence length, long contextC

What Attention Costs.

Double the length of the text and the projections and MLPs do twice the work. Attention does four times.

For one layer and one sequence of T tokens of width d, the matrix work splits into two kinds. The projections — queries, keys, values and output — and the MLP touch each token separately: about 24 T d² floating-point operations, linear in T. Attention itself compares every query with every earlier key and averages the values: about 4 T² d operations (half that with the causal mask, but counted in full by most kernels), quadratic in T. Per token, the first is 24 d², fixed, and the second is 4 T d, growing with the length.

The two are equal when T = 6d. For a model of width 4 096, that is a context of about 24 000 tokens: below it, attention is a minor cost; above it, it dominates. Memory is worse. A naive implementation materialises the full T × T score matrix for every head, which at long lengths is far larger than the model — the reason for FlashAttention (plate XVI.4), which never stores it.

Below: the floating-point work per token of one layer, split between attention and everything else, as the context grows from 512 tokens to a million, for the width chosen. The second panel shows the memory a naive implementation would need for one layer’s attention scores.

What you will learn

Count the FLOPs of attention and of the rest of a layer.
Find the length at which attention dominates.
Compute the memory of a materialised score matrix.
Explain why long contexts need special kernels and architectures.

Workshop

Left: FLOPs per token for one layer against context length (log scales): attention (red), projections and MLP (blue). Right: memory to hold one layer’s attention scores for one sequence in 16-bit numbers, against length.

Linear, then quadratic

At width 4 096, a layer spends 4.0 × 10⁸ FLOPs per token on its projections and MLP whatever the context, and attention adds 4 T d. At 8 192 tokens attention is 25% of the layer’s work; the two are equal at 24,576 tokens; at 128 thousand tokens attention is 84%. A materialised score matrix for one layer with 32 heads at 128 thousand tokens would need 1.1 TB in 16-bit numbers, far beyond any single accelerator’s memory — for one layer of one sequence.

These numbers explain much of the engineering in later volumes. FlashAttention (plate XVI.4) computes attention in tiles without ever storing the full score matrix, removing the memory problem but not the quadratic work. Sliding-window attention (plate VI.7), latent and grouped attention (plates VII.2–VII.3) and linear-time alternatives (plate VII.8) attack the work itself. And at inference time the cost that grows with length is not the scores but the cache of keys and values, which plate VI.9 and Vol. XXIII measure.

The 6d crossover is a rule of thumb for dense models with a 4× MLP; models with grouped-query or latent attention shift it, and causal kernels that skip the masked half halve attention’s work.

Terms

Quadratic cost: Work that grows with the square of the sequence length.
Projection FLOPs: The work of computing queries, keys, values and outputs: linear in length.
Score matrix: The T × T matrix of query–key scores for one head.
Materialise: Store an intermediate result in memory in full, rather than computing it piece by piece.

Attention costs 4 T d per token per layer against 24 d² for the rest: negligible for short texts, dominant beyond about T = 6d, and its score matrix is too large to store at long lengths.