Attention and Its Computational Cost

For a 16,384-token sequence with GPT-2’s twelve heads in half precision, one attention score tensor is 6.0 GiB. A straightforward implementation of attention needs two of them at once. On this machine it asked for 12.3 GiB, on a card with 12, and did not fail. I will explain below why it did not, because the explanation is itself instructive. PyTorch’s own attention function, given the same inputs, used 24 MiB, which is the size of its output. The two agreed on the result to within the rounding error of 16-bit floating point, and the second was eleven times faster at a length of 8,192.

Same mathematics, same answer, five hundred times less memory. The difference is not a cleverer algorithm in the sense of fewer operations. It is that one implementation writes an enormous intermediate result to memory and reads it back several times, and the other never forms it.

This chapter measures what attention costs in arithmetic, in memory and in time, and it separates two very different shapes of the same computation: the one where many queries attend to many keys, and the one where a single new query attends to everything that came before. The second is what a generation loop does for almost its whole life.

Where this sits

Chapter 4 identified attention as the only part of the model that moves information between positions, and the only tensor whose size grows with the square of the sequence length. Chapter 5 showed that scoring a known sequence is parallel while choosing tokens is serial. Chapter 7 gave the roofline and the idea that an operation is either limited by arithmetic or by bytes. This chapter combines the three.

The Inference Lab piece added here is an attention cost model, attention_flops, score_bytes and kv_bytes, tested without a GPU, and a comparison harness that runs implementations on the same inputs against a float64 reference. It adds one rule that matters from now on: a benchmark of attention must record which kernel actually ran. PyTorch’s attention function is a semantic entry point and not an implementation, and on this GPU two of the four backends it knows about are not available.

For the Pi series, the cost of context is where the retrieval and memory volumes meet inference. Choosing how much retrieved text to put in a prompt is a choice about \(T\), and the numbers here say what \(T\) costs in memory and time. The memory book’s persistent knowledge is a way of keeping \(T\) small.

What is being computed

For one head, with queries \(Q \in \mathbb{R}^{T_q \times d}\), keys and values \(K, V \in \mathbb{R}^{T_{kv} \times d}\):

Two matrix products and a softmax. The scores \(S\) and the probabilities \(P\) are both \(T_q \times T_{kv}\). The products cost \(2T_qT_{kv}d\) each, so across \(B\) sequences and \(h\) heads: For causal self-attention over \(T\) tokens, position \(i\) sees \(i\) keys, so only \(T(T+1)/2\) of the \(T^2\) pairs are needed, and the count is almost halved. Softmax, masking and rescaling add elementwise work that is small in FLOPs and, as Chapter 4 showed, not small in time.

The two shapes are:

Shape	\(T_q\)	\(T_{kv}\)	FLOPs grow as	Scores per head
Prefill (a whole prompt)	\(T\)	\(T\)	\(T^2\)	\(T \times T\)
Decode (one new token, cached keys)	1	\(T\)	\(T\)	\(1 \times T\)

Doubling the prompt quadruples the first and doubles the second, and the unit tests assert both. For GPT-2 at a context of 1,024 the attention products are 13 per cent of the matmul arithmetic of a forward pass (Chapter 4). That is small. What the arithmetic count hides is what the next sections measure.

The cost of writing down the scores

A direct implementation follows the formulas: compute \(S\), mask it, softmax it to get \(P\), multiply by \(V\). Every step reads the previous step’s whole result from memory and writes its own. With \(h = 12\) heads in fp16, one \(T \times T\) tensor is:

\(T\)	One score tensor (fp16)
512	6 MiB
1,024	24 MiB
2,048	96 MiB
4,096	384 MiB
8,192	1.5 GiB
16,384	6.0 GiB

Four times as large per doubling, per sequence, and the batch multiplies it again. The direct implementation holds two of these live at the same time, the masked scores and then the probabilities, so its peak should be about twice the table. Measured peak extra memory over the inputs, for the three implementations that exist on this machine:

\(T\)	Naive	Math backend	Memory-efficient backend
512	21 MiB (3.5×)	34 MiB (5.7×)	1 MiB
2,048	200 MiB (2.1×)	472 MiB (4.9×)	3 MiB
8,192	3,148 MiB (2.05×)	7,264 MiB (4.7×)	12 MiB
16,384	12,568 MiB (2.05×)	28,864 MiB (4.7×)	24 MiB

The multiples in parentheses are the peak over one score tensor. The naive peak settles at 2.05, matching the prediction. PyTorch’s reference “math” backend, which is a general fallback, needs 4.7 tensors. The memory-efficient backend’s peak grows only with the output: at 16,384 it is 24 MiB, which is exactly \(12 \times 16384 \times 64 \times 2\) bytes. It never forms the score matrix. It processes the keys in blocks, keeping a running maximum and a running sum so that the softmax can be completed without having seen all the scores at once. That idea is the subject of Chapter 13, and the reason it is exact and not an approximation is worth saying now: the output is the same function.

The 28 GiB row did not fail, and neither did the 12.3 GiB one. The card has 12 GiB. I took this to be the Windows display driver model, which lets GPU memory be oversubscribed into system memory by paging. The allocation succeeded and the peaks above are real, but any timing taken at that point would measure paging. I did not time those cases. On a machine where oversubscription is not available, these rows are out-of-memory errors, and that is how the failure usually shows up.

Which kernel ran

PyTorch’s scaled_dot_product_attention chooses an implementation at run time from a set of backends. I asked each in turn:

Backend	fp16	fp32
Math	runs	runs
Memory-efficient	runs	runs
FlashAttention	unavailable: no kernel	unavailable
cuDNN attention	unavailable	unavailable

PyTorch reports the reasons when asked to explain a refusal. For FlashAttention: “Torch was not compiled with flash attention”, so it is this Windows build of PyTorch that lacks the kernel, whatever the GPU could run. For cuDNN attention: it supports only compute capabilities 8.0 to 9.0, and this GPU is 7.5. Nothing in this chapter is a FlashAttention measurement, and the FlashAttention chapter will have to say so unless it uses different software. The memory-efficient backend is a close relative: it is also an exact, block-wise, never-materialise-the-scores kernel, and what it shows is representative of the approach. It is not a statement about FlashAttention’s specific numbers.

The general point is more important than the specific one. The same line of Python can run any of these kernels depending on hardware, data type, mask and head dimension. A report that says “I used SDPA” has not said what it measured.

Same answer?

Each implementation was compared to the naive algorithm run in float64 on the same inputs, for 12 heads of dimension 64:

Precision	Length	Naive	Math	Memory-efficient
fp32	64	\(8.3 \times 10^{-7}\)	\(1.0 \times 10^{-6}\)	\(8.9 \times 10^{-7}\)
fp32	2,048	\(7.1 \times 10^{-7}\)	\(7.1 \times 10^{-7}\)	\(9.4 \times 10^{-7}\)
fp16	64	\(1.6 \times 10^{-3}\)	\(9.7 \times 10^{-4}\)	\(9.8 \times 10^{-4}\)
fp16	2,048	\(1.8 \times 10^{-3}\)	\(9.7 \times 10^{-4}\)	\(9.7 \times 10^{-4}\)

In fp32 all three are at rounding level. In fp16 the library kernels reach about \(10^{-3}\), which is the resolution of 16-bit numbers on outputs of order one. The naive implementation is about twice as far off, because it rounds its intermediate score and probability tensors to 16 bits and the library kernels keep them in higher precision. Here the fast implementation is also the more accurate one, and the reason is the same: it never stores the intermediate in low precision.

Time

The same three-times-repeated measurement protocol as Chapter 7, with the clock ramp and CUDA events. Repeats agreed to within 1 per cent. Medians:

\(T\)	fp16 naive	fp16 efficient	fp32 naive	fp32 efficient
512	0.29 ms	0.04 ms	0.55 ms	0.16 ms
1,024	1.07 ms	0.14 ms	2.18 ms	0.52 ms
2,048	5.71 ms	0.43 ms	8.83 ms	1.79 ms
4,096	18.7 ms	1.52 ms	33.3 ms	6.74 ms
8,192	67.0 ms	5.97 ms	155.8 ms	26.0 ms

Both are quadratic. The fitted exponent for \(T \ge 1{,}024\) was 1.96 (fp16) and 2.04 (fp32) for naive, and 1.82 and 1.88 for the efficient kernel. The efficient kernel falls short of 2 because it has not saturated at the lower lengths: its throughput was still rising from 11.7 TFLOP/s at 1,024 to 17.3 at 8,192. That reflects underuse of the hardware at the small end, not a shortage of arithmetic at the large end. Attention was not made subquadratic. It was made faster per unit of quadratic work, by about eleven times in fp16 at 8,192.

Counting only the causal half of the arithmetic, the efficient fp16 kernel reaches 17.3 TFLOP/s, about 55 per cent of the best matrix multiplication rate measured in Chapter 7, and the naive one 1.5 TFLOP/s, about 5 per cent. So the naive implementation is not compute-bound. A rough count shows why: it writes \(S\), reads and rewrites it for masking, reads and rewrites it for softmax, and reads \(P\) for the second product. That is three writes and three reads of a \(T^2h\)-element tensor, 12 bytes per element in fp16. At \(T = 8{,}192\), 805 million elements times 12 bytes is 9.7 GB, which would take about 35 ms at the 278 GB/s copy bandwidth. The measured 67 ms is roughly twice that, so traffic over the score tensor accounts for about half of the naive time, with its full (not causal) matrix products and the mask construction most plausibly making up the rest. This is an estimate, built from a hand count of passes and a bandwidth measured on a different operation, and I would not defend it beyond the claim that the naive implementation is dominated by moving the score tensor.

What the causal mask saves

The causal mask removes half of the pairs, so one might expect causal attention to cost half as much as full attention. I timed both:

Implementation	Causal time over full time (T = 1,024 / 2,048 / 4,096)
Work ratio (arithmetic)	0.50
Memory-efficient kernel	0.71 / 0.59 / 0.58
Math backend	1.27 / 1.32 / 1.30
Naive implementation	1.69 / 1.68 / 1.74

The efficient kernel gets close to the arithmetic saving, because it skips whole blocks that lie entirely above the diagonal. The ratio improves with length as the skipped blocks outnumber the partial ones along the diagonal. The math backend and the naive implementation are slower with a causal mask than without, by 30 per cent and 70 per cent. They compute the full \(T \times T\) product anyway and then pay an extra pass over it to apply the mask. For them, the mask saves no arithmetic and adds memory traffic.

A mask is a specification and not an optimisation. Whether it saves work depends entirely on whether the kernel is written to exploit it.

In the whole model

Does any of this change a real forward pass? Hugging Face GPT-2 can be loaded with either its eager attention or PyTorch’s SDPA. In fp32, batch 1:

Tokens	Eager	SDPA	Faster by	Peak extra memory	Max logit difference
128	7.6 ms	7.1 ms	7%	34 MiB, both	\(9 \times 10^{-5}\)
512	31.1 ms	26.8 ms	14%	148 MiB, both	\(1.1 \times 10^{-4}\)
1,024	71.4 ms	54.7 ms	23%	272 MiB, both	\(1.4 \times 10^{-4}\)

The saving grows with length, as the attention share of the model does, but it is a fraction of a forward pass and not a factor. And the peak memory is the same. This is the number that most surprises. Attention is not the largest tensor in a forward pass at 1,024 tokens. The logits are: \(1024 \times 50257 \times 4\) bytes is 206 MB, against a 272 MiB peak, because the naive loop computes the output projection at every position and keeps all of it: almost three quarters of the peak. That is the waste noted in Chapter 4: only the last row is used. SDPA’s saving shows up at the transient score tensor (50 MB per layer in fp32 at 1,024 tokens), which is freed before the logits are formed.

This should temper any simple story about attention being the memory problem. At the lengths where GPT-2 can run, it is not. At 16,384 tokens, for models whose context allows it, it is. Which tensor is largest depends on \(T\), the vocabulary and the head count, and the answer is found by measuring.

The other shape: one query, many keys

During generation, each step brings one new query. Everything before it is already in the KV cache of Chapter 10: the keys and values of all earlier positions. The attention for that step is a \(1 \times T\) problem. Its arithmetic is \(4hTd\) FLOPs, about three million at \(T = 1{,}024\). Its memory traffic is not small: it has to read every cached key and value once.

The bytes in a cache, from kv_bytes, for GPT-2 small (12 layers, 12 heads of 64, fp16):

which is 36 MiB per sequence at 1,024 tokens and 144 MiB at 4,096. The weights of the model in fp16 are 249 MB. The two traffic sources per decode step therefore compare as follows: weights are read once per step whatever the batch size, and the cache is read once *per sequence*.

Batch	Context	Cache read per step	Compared with weights (249 MB)
1	1,024	36 MiB	15%
7	1,024	252 MiB	106%
16	1,024	576 MiB	243%
16	4,096	2.3 GiB	970%

At a context of 1,024 the cache traffic equals the weight traffic at a batch of about seven. Beyond that, attention over the cache is the larger part of every decode step. Batching, which Chapter 7 showed to be nearly free for the weights, is not free for the cache. This is the reason a serving system’s capacity is set by cache memory and cache bandwidth well before it is set by weights, and it is why the later chapters on paging, quantised caches and grouped-query attention exist. Grouped-query attention, for example, reduces the number of key and value heads and therefore the factor of 12 in the formula above.

I measured one query against a cache of \(T\) keys (fp16, 12 heads) at three batch sizes:

Batch	Context \(T\)	Time	Cache read (one layer)	Achieved bandwidth
1	128	25.6 µs	0.4 MiB	15 GB/s
1	1,024	64.7 µs	3 MiB	49 GB/s
1	16,384	1,007 µs	48 MiB	50 GB/s
16	1,024	205 µs	48 MiB	246 GB/s
16	16,384	3,112 µs	768 MiB	259 GB/s
64	1,024	779 µs	192 MiB	258 GB/s
64	16,384	12,226 µs	3,072 MiB	263 GB/s

Time is linear in \(T\), as the cost model says. And the bandwidth tells two stories. At batch 16 and 64 the kernel reads the cache at 250 to 263 GB/s, close to the ceiling from Chapter 7: decode attention at useful batch sizes is a pure bandwidth problem. At batch 1 it reaches only about 50 GB/s, a fifth of the ceiling, whatever the length. My hypothesis is that one sequence offers only 12 independent heads to a 34-multiprocessor GPU, so most of the chip is idle, and that more parallelism across the cache length would help. I did not test it. If it is right, it describes a kernel limitation and not a fundamental one, and Chapter 13 returns to techniques that split the key dimension across the chip.

The shape of the result matches what Chapter 7 found for weight matrices: a lone decode step wastes the hardware, and the remedy is to serve more than one sequence at a time. Chapter 9 puts the two phases together.

What the numbers do not show

• No FlashAttention. It was unavailable on this GPU in this build, so every result about “optimised attention” is for the memory-efficient kernel.
• One consumer GPU under a display driver. Memory oversubscription, clock behaviour and launch costs are specific to it. The shapes of the curves are general; the positions of the knees are not.
• GPT-2’s shape only. 12 heads of dimension 64, multi-head (not grouped) attention. Models with 64 heads, grouped key/value heads or head dimension 128 shift every number here.
• No masks beyond causal. Padding masks, sliding windows (Chapter 22) and arbitrary masks can disable the fast kernels.
• Random inputs. Attention over trained models’ real activations, with their large outliers (Chapter 4), may stress low-precision accumulation differently.
• The decode kernel was measured in isolation. Inside a real step it competes with the weight reads, and the numbers there will differ.
• The batch-1 bandwidth explanation is a hypothesis.
• Hand-counted traffic estimate for the naive kernel is an estimate.

Inference Lab increment

The Lab now has:

• attention_flops(batch, heads, q_len, kv_len, head_dim, causal), score_bytes(...), kv_bytes(...) and fit_exponent(...): the cost model, tested for the doubling properties above.
• naive_attention and sdpa(backend=...): implementations that can be named and compared with a float64 reference.
• backends, memory, scaling, causal, decode, cpu and model commands that produce the tables in this chapter.

The rule it adds to the Chapter 6 rule: a change to attention is accepted when it matches the float64 reference within a stated tolerance, records the backend that actually ran, and reports memory as well as time. A time-only comparison would have missed the 500-fold difference that opened this chapter.

What you should now be able to do

• State the FLOPs of attention for any \(T_q\), \(T_{kv}\), batch and head count, for dense and causal cases.
• Compute the size of the score tensor and of a KV cache from model dimensions.
• Explain what is quadratic in attention (arithmetic) and what a block-wise exact kernel removes (memory traffic and storage), and what it does not.
• Say why a causal mask may not save any time, and how to tell.
• Choose between prefill-shaped and decode-shaped analysis, and say which is compute-bound and which is bandwidth-bound.
• Find the batch size at which KV-cache traffic overtakes weight traffic for a given model and context.
• Report which attention kernel was used in a benchmark.

Key terms

• Score matrix: the \(T_q \times T_{kv}\) tensor of scaled query-key products.
• Materialise: write an intermediate tensor to memory.
• Prefill-shaped / decode-shaped attention: many queries against many keys, or one query against a cache.
• KV cache: the stored keys and values of earlier positions.
• Exact attention: attention computed to the same function as the direct formula, differing only by rounding.
• SDPA backend: one of the kernels PyTorch’s attention function can dispatch to.
• Online softmax: computing a softmax over blocks with running statistics.
• Achieved bandwidth: bytes actually read per second by a kernel.

Reader exercises

1. Compute, for a model with 32 layers, 32 heads of 128 and a context of 8,192, the KV cache for one sequence in fp16. At what batch size does it exceed the weight memory of a 7-billion-parameter fp16 model?
2. Repeat memory for \(h = 32\) and \(d = 128\). Which sizes now fail on your GPU, and what does the failure look like?
3. Write the three-line change that makes the naive implementation skip masking for the first row of each block. Does it change the time?
4. Run decode with the sequence length split across several thread blocks if you can write the kernel, or with torch.compile. Does batch-1 bandwidth improve?
5. Test whether SDPA still dispatches to the efficient kernel when you pass an explicit boolean mask instead of is_causal.
6. Extend the whole-model comparison to fp16 and report where peak memory finally differs.
7. Replace the naive implementation’s fp16 intermediates with fp32 ones and measure how its error and time change.
8. Estimate, with the cost model, the context at which attention’s share of GPT-2’s forward pass matmul FLOPs reaches 50 per cent.

References

• Attention Is All You Need — defines scaled dot-product attention.
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — the block-wise exact algorithm discussed here and in Chapter 13. It was not run in this chapter.
• PyTorch: scaled_dot_product_attention — official documentation, including the backend selection rules.

Evidence in this chapter

Claim	Category
Attention costs \(4BhT_qT_{kv}d\) FLOPs; causal counts \(T(T+1)/2\) pairs	Established background
Block-wise exact attention avoids materialising the scores	Established background, from the FlashAttention paper
Naive peak memory is 2.05 score tensors; efficient peak is the output size	Book result
FlashAttention and cuDNN attention unavailable on this GPU in this build	Book result
Efficient kernel 11× faster than naive in fp16 at \(T=8192\); both near-quadratic	Book result, three runs within 1%
Naive implementation is dominated by score-tensor traffic	Book result for the share; the hand count is an estimate
Causal mask costs 30% to 70% extra in the math and naive paths	Book result
SDPA is 7% to 23% faster for whole GPT-2 and saves no peak memory	Book result, fp32, batch 1
KV traffic equals weight traffic at batch ≈ 7 for GPT-2 at 1,024	Arithmetic from the formulas, not measured end to end
Decode attention at batch 16 and 64 reaches 250–263 GB/s	Book result
Batch-1 decode attention is limited by parallelism	Book hypothesis, untested
FlashAttention numbers, GQA, paged caches, other masks, other GPUs	Not measured

Implementation status

• Manuscript: content/books/inference/08-chapter.md
• Runnable examples: examples/ch08/attention.py (original), examples/ch08/attention_cost.py (cost model and measurements)
• Tests: tests/test_ch08_attention.py, tests/test_ch08_attention_cost.py (cost formulas, causality and equivalence properties; one GPU test skips without CUDA; model commands are not exercised by tests)
• Evidence: research/ch08/sources.md, research/ch08/results.json, research/evidence-register.md
• Verification command: python -m unittest discover -s tests -v

A correction to the earlier version of this chapter: the original example reported naive attention taking 1.0 ms against 0.14 ms for SDPA, a sevenfold gap, from one un-warmed call of each. Repeated with a warm-up on the same shape the ratio is 2.0; the first-call ratio was 6.1. The sevenfold figure was mostly start-up cost.

The chapter is a draft with draft = false. It makes no claim about FlashAttention, grouped-query attention, other hardware or production serving. Chapter 9 joins the two shapes of this chapter into the two phases of a request: prefill and decode.