Skip to content

The Cost of Attention

The path from text to the next token in seven steps: tokens, embeddings, query/key/value, mask and heads, the block, the output and generation. ATTENTION TRANSFORMERS Tokenstext → ids Embeddingsids → X Q · K · Vone head Masktwo heads BlockLN · MLP Outputlogits → p Generationpick text → → next token

A side page. The running example has five tokens and costs nothing to run. This page is about what attention costs in a real model: the cache that generation keeps between one token and the next, and the \(O(n^2)\) of long contexts. ← Masks and heads · Overview

The KV cache is the thing nobody teaches and everybody pays for

A model generates text one token at a time, and every new token needs attention over all the previous ones. What does that require from each old token? Not its query: that only served to compute its own output, back then. But its key and value will be consulted by every token that comes after. Without keeping them, step \(t\) would rerun the whole network over the \(t\) tokens of the prefix just to rebuild them, and generating \(n\) tokens would cost \(1 + 2 + \dots + n = O(n^2)\) token passes instead of \(n\) — so nobody does it. Every K and V is computed once and kept.

That cache is now the dominant memory cost of serving a model, and it is the reason "128k context" is a pricing decision rather than an architecture decision:

\[ \text{cache bytes} \;=\; 2 \times L \times h_{kv} \times d_{\text{head}} \times n \times B \times \text{bytes} \]

The 2 counts K and V; \(L\) is the number of layers, \(n\) the context length and \(B\) the batch size. Plug in a real model, Llama 3 8B: \(L = 32\), \(h_{kv} = 8\), \(d_\text{head} = 128\), in 16 bits (2 bytes). Each token costs \(2 \times 32 \times 8 \times 128 \times 2 = 131{,}072\) bytes, or 128 KiB. A 128k-token context therefore takes 16 GiB — as much as the model's own weights — and that is for one user. The toy of these pages, with one layer and two heads of width 2, keeps \(2 \times 1 \times 2 \times 2 = 8\) numbers per token; GPT-2 small keeps \(2 \times 12 \times 12 \times 64 = 18{,}432\), 36 KiB per token in 16 bits.

Only one term in that product is a free design variable: \(h_{kv}\), the number of key/value heads. Queries can keep all \(h\) heads; K and V can be shared. Llama 3 8B has 32 query heads and only 8 key/value heads — with one per query the cache would be four times larger. In the shape table of masks and heads, K and V become \(B \times h_{kv} \times n \times d_\text{head}\), and each of them is repeated \(h / h_{kv}\) times to meet the queries; only the small version is cached. This produced the sequence every current model sits somewhere on:

Scheme KV heads Cache Quality Status in 2026
MHA — multi-head1 \(h\) 1× Baseline Historical. Nothing serves this at long context.
MQA — multi-query3 1 \(1/h\) Measurably worse Used where memory dominates everything
GQA — grouped-query4 \(g\) (typically 8) \(g/h\) Indistinguishable from MHA The default. Llama, Mistral, Qwen, Gemma
MLA — multi-head latent5 — ~\(1/57\) Reported at or above MHA DeepSeek-V2/V3; compresses KV to a low-rank latent and caches that

This is a memory-bandwidth problem, not a FLOPs problem

During generation you read the entire KV cache to produce one token. The arithmetic intensity is terrible, the accelerator is idle, and throughput is set by HBM bandwidth. Everything that speeds up inference — GQA, quantized KV caches, PagedAttention6, speculative decoding, larger batches — is an attack on bytes moved, not on operations performed. Recognizing which resource you are short of is most of practical LLM engineering.

Making \(O(n^2)\) affordable

Attention is \(O(n^2 d)\) in time and, written naively, \(O(n^2)\) in memory: the matrix \(S\) from the shape table has one entry per pair of tokens. At \(n = 128\text{k}\) it has \(1.7 \times 10^{10}\) entries — 34 GB in 16 bits, per head, per layer. Four families of answers exist, and they are not equally successful.

  • Same math, better execution

    FlashAttention2 never materializes the \(n \times n\) matrix. It tiles the computation so each block stays in SRAM, and uses the online-softmax trick to combine blocks. Memory drops to \(O(n)\), wall-clock drops several-fold, and the output is the naive one up to floating-point rounding.

    This is not an approximation and there is no quality trade-off. FlashAttention-2 reorganizes the work across the GPU and FlashAttention-3 tunes it for Hopper-class hardware. If you are writing attention by hand, you are almost certainly slower — use F.scaled_dot_product_attention.

  • Attend to less

    Sliding window: each token sees the last \(w\) only; stacking layers grows the effective reach exactly like a CNN's receptive field. Attention sinks7: keep the first few tokens always visible — models dump excess attention mass on them, and dropping them wrecks generation.

    Modern long-context models interleave: a few full-attention layers among many local ones, so global information still has a route while cost stays near-linear.

  • Change the math

    Linear attention replaces softmax with a kernel, making the operation associative and the cost \(O(n)\) — and giving up the sharp retrieval that softmax provides. State-space models (Mamba, Mamba-28) reach the same \(O(n)\) with a recurrent state and a constant-size cache.

    The honest 2026 verdict: pure linear models lose on exact recall from long context. Hybrids win — Jamba, Zamba, Nemotron-H and Falcon-H1 combine SSM layers with a small share of full attention (a minority of layers, or attention heads running beside the SSM) and report Transformer-level quality at a fraction of the cache.

  • Learn what to skip

    Trainable sparse attention — DeepSeek's NSA9 and Kimi's MoBA — pick which blocks of the past to attend to, with the selection trained end-to-end rather than fixed by a pattern.

    This is the most promising of the four right now: it keeps softmax's exact retrieval where it matters and pays near-linear cost everywhere else, and unlike earlier sparse schemes it is hardware-aligned enough to actually be faster.

What to actually use

Full attention with FlashAttention and GQA, RoPE for position, sliding-window or hybrid layers if your sequences are genuinely long. That combination is what essentially every model released in the last two years does, and departing from it needs a measured reason.

Check yourself — while the toy generates ␣roof after “the cat climbed the”, what is in the cache, and what is computed?

The keys and values of the five tokens of the prompt, for its one layer and two heads: \(5 \times 8 = 40\) numbers. For ␣roof the model computes only the new row — its embedding, its \(q\), \(k\) and \(v\), one row of scores against six keys, its MLP — and adds its \(k\) and \(v\) to the cache. Without the cache it would recompute all six rows, and generating \(n\) tokens would cost \(1 + 2 + \dots + n\) rows.

← Masks and heads · Overview · next chapter: Transformers

  1. Vaswani, A., et al. (2017). Attention Is All You Need — NeurIPS. Scaled dot-product attention, multi-head, sinusoidal positions. ↩

  2. Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — NeurIPS. See also FlashAttention-2 (2023) and FlashAttention-3 (2024). ↩

  3. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. Multi-query attention, and the first clear statement that decoding is memory-bound. ↩

  4. Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — EMNLP. The compromise that became the default. ↩

  5. DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. Multi-head latent attention: cache a low-rank latent instead of K and V. ↩

  6. Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention — SOSP. Virtual memory for the KV cache; the core of vLLM. ↩

  7. Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2024). Efficient Streaming Language Models with Attention Sinks — ICLR. Why the first token cannot be evicted. ↩

  8. Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality — ICML. Mamba-2, and the equivalence that makes hybrids designable. ↩

  9. Yuan, J., et al. (2025). Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention — ACL. Sparsity learned during pretraining rather than imposed after it. ↩