Making attention cheaper

Attention’s two costs are the KV cache, which grows with every token, and the score matrix, which grows with the square of the length. Modern models share keys and values between heads (GQA, MLA), let most layers look only at recent tokens, compute attention in tiles that never store the full matrix (FlashAttention), and pick a few keys to read per query (sparse attention).

Key idea: Fewer K/V heads (MQA, GQA) or one compressed latent (MLA) shrink the cache; local windows bound the work; FlashAttention tiles the softmax so the T × T matrix is never stored; sparse attention reads only the top-k keys
Real world: A study group. Instead of every student photocopying the whole library (one copy each), they share a few copies (GQA), keep a summary sheet they can expand when needed (MLA), mostly look at this week’s notes (sliding window), and only fetch the handful of books the index says matter (sparse).
In an LLM: Llama 3 8B would need 69 GB of KV cache for one 128K-token conversation with full multi-head attention. With GQA it needs 17 GB. DeepSeek-V3, a far bigger model, needs about 9 GB thanks to MLA.

Two different costs

The KV cache lesson showed that a model generating text keeps every earlier token’s keys and values, in every layer. The cost lesson showed that attention compares every query with every key. These are two separate problems:

Memory. The KV cache grows by a fixed amount per token. At long contexts and many users it, not the weights, fills the GPU.
Work and memory traffic. The score matrix has $T \times T$ entries. Double the length and it has four times as many.

Each technique on this page attacks one or both.

Multi-head attention (MHA) gives every head its own query, key and value projections. During generation, the queries are used once and thrown away, but the keys and values are cached. So the number of key/value heads decides the size of the cache.

Multi-query attention (MQA, Shazeer, 2019): all query heads share one key head and one value head. The cache shrinks by the number of heads, but quality drops a little.
Grouped-query attention, (Ainslie et al., 2023): the query heads are split into groups, and each group shares a key/value head. Llama 3 has 32 query heads and 8 key/value heads in its 8B model, and 64 query heads with the same 8 key/value heads in its 70B model. This is now the default.
Multi-head latent attention, (DeepSeek-V2, 2024): each token stores one small compressed vector, the latent, and every head rebuilds its own key and value from it with learned matrices. DeepSeek-V3 caches 512 latent numbers plus a 64-number key that carries the RoPE rotation: 576 numbers per token per layer, where its 128 full heads would need 40,960.

context128K tokensusers at once1

✓ = what Llama 3 8B actually uses. Its weights alone are about 16 GB in 16-bit.

numbers per token per layer2,048

per token, all layers128 KiB

vs MHA4.0× smaller

grouped-query attention: the 32 query heads are split into 8 groups and each group shares one key and one value head. The cache is 4 times smaller than MHA and quality barely changes. Llama, Qwen, Mistral and Gemma all do this.

Top: which query heads share which keys and values (8 query heads drawn). Bottom: the KV cache each design needs for a real model, in 16-bit numbers. Pick a design, a model and a context length.

Notice that MLA is not always the smallest. With few heads, GQA’s cache is already tiny. MLA pays off in models with many large heads, which is why DeepSeek and Kimi K2 use it, and why GLM-5 adopted it too.

example · sharing textbooks

A class of 32 students could each buy every textbook (MHA). Or groups of four could share one copy (GQA): an eighth of the books, and everyone can still read what they need. Or the class could share a single copy (MQA): cheapest, but now everyone queues for the same pages. MLA is a different idea: keep one compact study guide per chapter, from which any student can rebuild the part of the book they need.

example · serving many users

A chatbot server running Llama 3 70B holds the 141 GB of weights once, but needs a KV cache per conversation. With 8 key/value heads, a 32K-token conversation needs about 11 GB. With 64 (full MHA) it would need 86 GB, and the server could hold eight times fewer conversations at once. That is why GQA spread so fast: it is a direct multiplier on how many users one GPU can serve.

Looking only nearby

Most of what a token needs is close by: the previous few words, the current sentence. Sliding-window attention, , lets each token see only the last $w$ tokens (Longformer popularised it for LLMs; Mistral 7B used a 4,096-token window). Work per token and cache per layer stop growing.

Information can still travel further than $w$ , one window per layer, like a rumour passed along a line. But some tasks need a direct look far back, so most models interleave: several local layers, then one global layer with full attention.

Gemma 2: alternates local (4,096-token window) and global layers one to one.
Gemma 3: five local layers (1,024-token window) for every global one, which cut its long-context KV cache to a fraction.
gpt-oss: alternates a 128-token window with full attention.

window6 tokens

scores computed129 of 300 (43%)

token 24 looks at6 keys

keys this layer must keep6

Each token only looks at the last 6 tokens, itself included. Work and memory per token stop growing: a band of fixed width instead of a triangle. Information can still travel further, one window per layer: after 10 layers, up to 50 tokens back. Mistral 7B used a 4,096-token window in every layer.

24 tokens. Row = the token doing the looking (query), column = the token looked at (key). Filled squares are scores the layer computes. Click a row to follow one token.

FlashAttention: the same answer, less traffic

A GPU has a large, slow main memory (HBM, tens of gigabytes) and a tiny, fast memory on each processor (SRAM, a few hundred kilobytes). Moving numbers between them often costs more time than the arithmetic. Written naively, attention writes the full score matrix to main memory, reads it back to take the softmax, writes the weights, and reads them again to multiply by the values. At 8,192 tokens, one head’s score matrix is 67 million numbers.

(Dao et al., 2022) computes exactly the same result without ever storing that matrix. It loads a block of queries into fast memory, streams blocks of keys and values past it, and keeps three running numbers per row:

$m$ : the largest score seen so far,
$\ell$ : the sum of $e^{\text{score} - m}$ so far,
$o$ : the sum of $e^{\text{score} - m} \times$ value so far.

When a new block contains a bigger score, $m$ rises, and the old $\ell$ and $o$ are multiplied by $e^{m_\text{old} - m_\text{new}}$ to put them on the new scale. At the end, $o / \ell$ is the softmax-weighted average, identical to the ordinary computation. This trick is called the online softmax.

start

follow row:

scores (only this tile is in memory)4 × 8

k1k2k3k4k5k6k7k8q1q2q3q4

values, one column per key2 × 8

v1v2v3v4v5v6v7v8v·11.00.00.5−1.00.80.2−0.41.0v·20.01.00.50.5−0.61.0−0.41.0

→

O (written once)4 × 2

q1q2q3q4

row	max m	sum l	output o
q1	−∞	0.000	[0.000, 0.000]
q2	−∞	0.000	[0.000, 0.000]

row q1, this step

Every row starts with m = −∞, l = 0, o = [0, 0]. Nothing has been read yet.

memory traffic so far (numbers)0

standard attention needs176

Standard attention computes the whole 4 × 8 score matrix, writes it to the GPU’s main memory, reads it back for the softmax, writes the weights, and reads them again to mix the values. FlashAttention never stores it. Press play.

FlashAttention on 4 queries and 8 keys, in tiles of 2 × 2. Only the current tile of scores ever exists; each row keeps a running max m, a running sum l and a running output o. Play it, or step with the arrows.

FlashAttention is not an approximation: the answer matches to rounding error. It is faster because the expensive part, memory traffic, drops sharply because the $T \times T$ matrix is never written to or read from main memory. It also makes long training sequences possible, because the $T \times T$ matrix would not fit at all. Newer versions tune the same idea for each GPU generation: FlashAttention-2 (2023), FlashAttention-3 for Hopper GPUs (2024) and FlashAttention-4 for Blackwell GPUs (2026).

example · adding up a long receipt

You can total a receipt by first copying every line onto a separate sheet, then adding up the sheet. Or you can keep a running total in your head as you read. FlashAttention keeps a running total. The only wrinkle is the softmax’s “minus the max”: if a much bigger item turns up halfway, you rescale the total so far before carrying on, like converting a running total from pounds to pence when you meet a price in pence.

example · a 128K-token prompt

Reading a 128K-token document, standard attention would need a score matrix of 128K × 128K, about 17 billion numbers per head per layer: 34 GB in 16-bit, for one head of one layer. With FlashAttention the GPU never holds more than one tile of it at a time. Every major serving and training library (PyTorch’s scaled_dot_product_attention, vLLM, SGLang) uses kernels of this kind.

Sparse attention: read only what matters

FlashAttention makes full attention efficient, but it still compares every query with every key. Sparse attention skips most comparisons on purpose. The 2025–26 designs learn which keys to read:

Native Sparse Attention (NSA, Yuan et al., 2025, DeepSeek and Peking University) runs three branches per query: attention over compressed blocks of keys (a coarse summary), attention over the few blocks that summary rates highest, and a sliding window for recent tokens.
DeepSeek Sparse Attention (DSA, in DeepSeek-V3.2, 2025) adds a small, fast “lightning indexer” that scores every earlier token, then runs MLA on only the top 2,048. GLM-5 (2026) uses the same design.
DeepSeek-V4 (April 2026) goes further with a hybrid of compressed sparse attention and heavily compressed attention layers. At a 1M-token context, DeepSeek reports that V4-Pro needs 27% of the per-token compute and 10% of the KV cache of V3.2.

The indexer still has to look at every key, but it is so much cheaper than full attention that the total cost falls sharply on long inputs.

Which model uses what

Model	Attention	Key/value heads	Local layers
GPT-2 (2019)	MHA	12 of 12 (small)	none
Mistral 7B (2023)	GQA	8 of 32	all layers, 4,096 window
Llama 3 8B / 70B (2024)	GQA	8 of 32 / 8 of 64	none
Gemma 3 (2025)	GQA	varies by size	5 of every 6, 1,024 window
DeepSeek-V3 (2024)	MLA	one 512-number latent	none
gpt-oss-120b (2025)	GQA	8 of 64	every other layer, 128 window
DeepSeek-V3.2 (2025), GLM-5 (2026)	MLA + DSA	latent, top 2,048 keys	none

going deeper

KV cache formulas

Per token, per layer, in numbers (multiply by 2 bytes for 16-bit and by the layer count):

\text{MHA: } 2 n_h d_h \qquad \text{GQA: } 2 n_{kv} d_h \qquad \text{MQA: } 2 d_h \qquad \text{MLA: } d_c + d_h^R

with $n_h$ heads of size $d_h$ , $n_{kv}$ key/value heads, latent size $d_c$ and rotated-key size $d_h^R$ . For DeepSeek-V3, $d_c = 512$ and $d_h^R = 64$ . Llama 3 8B: $2 \times 8 \times 128 \times 32$ layers $\times$ 2 bytes $= 131{,}072$ bytes, 128 KiB per token.

MLA can skip rebuilding the keys at inference by folding the up-projection into the query: $q^\top (W^{UK} c) = (W^{UK\top} q)^\top c$ . That is why the rotation is kept in a separate 64-number part: a position-dependent rotation would get in the way of this folding (DeepSeek-V2, section 2.1.3).

The online softmax update

For one query row and a new block of scores $s_j$ with values $v_j$ :

m' = \max(m, \max_j s_j), \qquad \ell' = e^{m - m'}\,\ell + \sum_j e^{s_j - m'}, \qquad o' = e^{m - m'}\, o + \sum_j e^{s_j - m'} v_j

and finally output $o / \ell$ . FlashAttention-2 delays the division to the end exactly as shown; the original paper divided at every step.

In code

import torch.nn.functional as F
# Picks a FlashAttention kernel when it can. GQA: pass fewer K/V heads with enable_gqa=True.
out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

In papers

MQA: Shazeer, 2019. GQA: Ainslie et al., 2023.
MLA: DeepSeek-V2, 2024; sizes in DeepSeek-V3, 2024.
Sliding windows: Longformer, 2020; Mistral 7B, 2023; Gemma 2, 2024; Gemma 3, 2025.
FlashAttention: 1, 2, 3, 4.
Sparse: NSA, 2025; DeepSeek-V3.2 (DSA), 2025; DeepSeek-V4 model card, 2026.