Attention, and How to Shrink It

The last post ended on a number: one Llama 2 7B sequence needs about 2 GB of KV cache, and at any real concurrency the cache, not the weights, fills the GPU first. We left the cache written as a product:

$$\text{cache} = \text{num\_layers} \times \text{num\_heads} \times 2 \times \text{seq\_len} \times d_\text{head} \times \text{bytes}$$

Most of those factors are nailed down. $\text{num\_layers}$ and $d_\text{head}$ are the model’s shape; $\text{seq\_len}$ is the request; $\text{bytes}$ is precision, and quantizing it is a real lever but a separate post. The term the first half of this post attacks is $\text{num\_heads}$ — or more precisely, the number of distinct keys and values you keep per token. Hence the first three techniques: MQA, GQA, and MLA — the first two shrink the cache by sharing keys across heads; the third shrinks it by compressing them into a latent and decoding per head. The fourth, FlashAttention, ignores this formula entirely and goes after a different cost. Two independent levers — shrink what you cache, then change how you compute it — and I’ll take them in that order.

The keys don’t need their own head

Every successor to multi-head attention is built on one observation: a token needs many query heads to ask different questions, but the keys and values those queries read don’t have to be unique to each head.

Recall the shape from last post. Vanilla multi-head attention gives each of the $h$ heads its own $Q$, $K$, and $V$ — $h$ queries, $h$ keys, $h$ values, every one of them a separate projection. The queries are where the expressiveness lives: each head aims its query at a different kind of relationship, one chasing the subject of a verb, another a coreference like it → cat. But each head also computes and then caches its own $K$ and $V$, and that — the $\text{num\_heads}$ factor — is what makes the cache wide. The whole family of optimizations keeps the query heads diverse and shrinks what they read. There are two ways to do that: share — keep fewer distinct sets of $K$/$V$ and make the heads split them — or compress — keep one small code per token and rebuild every head’s keys and values from it. Sharing is cheap to build and costs diversity; compressing keeps the diversity and costs machinery.

Sharing: MQA and GQA

MQA and GQA are the two production forms of sharing: one $K$/$V$ for all heads, or one per group.

Multi-Query Attention (Shazeer, 2019, Fast Transformer Decoding: One Write-Head Is All You Need) takes the idea to its limit: keep all $h$ query heads, but give them a single shared $K$ and $V$. The cache shrinks by $\text{num\_heads}$× — 32 query heads, $1/32$ the keys and values. The win lands exactly where decode hurts: decode is memory-bandwidth bound, so reloading a thirty-second of the bytes each step is most of the speedup. The cost is a small but real quality hit, plus the training instability the later GQA work documented.

Grouped-Query Attention (Ainslie et al., 2023) is the middle ground, and the one almost everyone landed on. Split the $h$ query heads into $G$ groups; each group shares one $K$/$V$. The cache is $G/\text{num\_heads}$ of full MHA — $G=1$ is MQA, $G=h$ is MHA. Llama 2 70B runs 64 query heads over 8 KV groups: $1/8$ the cache, and on the GQA paper’s own benchmark it nearly matched full MHA quality (47.1 vs 47.2) while beating MQA. You can even convert an existing MHA checkpoint by mean-pooling each group’s heads and uptraining on ~5% of the original compute. Mistral 7B, Llama 3, and most of the open-weight field ship GQA.

That’s the whole sharing story: a single parameter $G$ running from MHA ($G=h$) down to MQA ($G=1$), with GQA the default stop in between.

Compressing: MLA

MLA stops sharing and starts compressing: instead of keeping fewer distinct sets of keys and values, it keeps none — a single compressed latent vector per token goes in the cache, and every head’s keys and values are rebuilt from it on the fly. DeepSeek built it for DeepSeek-V2 and kept it for V3: a 60-layer, 128-head model whose per-token KV cache is more than an order of magnitude smaller than an equivalent multi-head model’s, at quality DeepSeek reports as better than full MHA.

The mechanism is a low-rank bottleneck. Each token’s hidden state $h_t$ is down-projected to a small latent $c^{KV}_t$ of dimension $d_c$ — in DeepSeek-V2, $d_c = 512$ against a full $\text{num\_heads} \times d_\text{head} = 128 \times 128 = 16{,}384$. At attention time, two learned up-projection matrices expand $c^{KV}_t$ back into per-head keys and values. The trick that makes this nearly free is absorption: because those up-projections are fixed weights, they fold into the query and output matrices ($W_q$, $W_o$) algebraically, so the full $K$ and $V$ are never materialized to be stored. Only $c^{KV}_t$ goes in the cache — one short vector per token instead of $\text{num\_heads}$ keys and $\text{num\_heads}$ values.

That makes the per-token cache $(d_c + d_h^R) = (512 + 64) = 576$ numbers per layer, versus $2 \times \text{num\_heads} \times d_\text{head} = 32{,}768$ for full MHA — about a 57× smaller per-layer footprint, or roughly what GQA would cost at 2.25 groups. DeepSeek reports the V2 design cuts the KV cache 93.3% against their dense 67B baseline while raising max generation throughput 5.76× — and, unlike any amount of sharing, does it while improving benchmark quality rather than conceding some. The catch is engineering complexity: a handful of extra projections, the decoupled-RoPE path, and a kernel that knows how to do the absorption. It’s not a config flag you flip on an existing model — which is exactly the problem the next wave of work set out to solve.

Converting: MLA for models that weren’t born with it

MLA’s complexity used to mean you only got it if you trained for it — DeepSeek and almost no one else. The interesting 2025 development is that MLA has become a conversion target: take a model trained with GQA, and turn it into an MLA model after the fact.

The clean theoretical result comes from TransMLA (a NeurIPS 2025 spotlight): for the same KV-cache budget, MLA strictly contains GQA. GQA forces every query head in a group to read the identical replicated key and value; MLA spends the same bytes on a shared latent but lets each head learn its own up-projection from it — same cache, strictly larger function class. Sharing is the degenerate case of compressing. So any GQA checkpoint can be re-parameterized as an MLA model and then compressed, recovering the original quality with a short fine-tune (the paper reports ~6B tokens to claw back a 93%-compressed Llama-2-7B).

It’s no longer one paper. MHA2MLA reaches MLA from plain MHA checkpoints on a fraction of a percent of the original pretraining data; X-EcoMLA (from AMD) initializes the MLA weights by SVD-decomposing the existing attention and distills them back to parity — a 6.4× cache cut on Llama-3.2-1B for about 70 GPU-hours, against the hundreds of thousands it took to pretrain the model. The recipe has even jumped modalities: MHA2MLA-VLM (AAAI 2026) converts a vision-language model, cutting ~95% of the KV cache on Qwen2.5-VL-7B at no average-score loss. Every one of these is a single-lab result with author-reported numbers, so I’d hold the specifics loosely — but the direction is unmistakable: MLA is becoming the format existing models migrate to, across vendors and even modalities, not just a DeepSeek house style.

The other lever: FlashAttention

Everything so far shrank what you cache. FlashAttention (Dao et al., 2022) is the other lever from the intro: it changes how you compute attention and touches the cache not at all. It produces the exact same output — not an approximation — while cutting the attention op’s memory from $O(\text{seq}^2)$ to $O(\text{seq})$.

The trick is to respect the memory hierarchy. A GPU has a large, slow pool (HBM, ~2 TB/s on an A100) and a tiny, fast one (on-chip SRAM, ~19 TB/s). Standard attention writes the full $\text{seq} \times \text{seq}$ score grid out to HBM and reads it back — and at long sequence length that round-trip, not the matmul, is the bottleneck. FlashAttention tiles $Q$, $K$, and $V$ into blocks that fit in SRAM and uses an online softmax — a running max and sum — to accumulate each tile’s contribution without ever materializing the whole grid. The big intermediate never leaves the chip. The reported kernel speedups run ~2–4×; the lineage continued with FlashAttention-2 (better work partitioning, ~73% of A100 peak) and FlashAttention-3 (Hopper-specific async and FP8, ~75% of H100 peak).

Two scope notes worth stating plainly. The win lands where the grid exists — prefill and training; at decode each new token scores a single row against the cache, and the bottleneck is streaming the cache itself, which is the first half of this post. And the memory FlashAttention saves is activation memory — a different pool from the KV cache. Run it on top of any model and the cache formula from the top of this post is unchanged. It’s the perfect complement to the cache tricks, not a substitute.

The frontier: new axes

Sharing and compressing moved one factor — $\text{num\_heads}$. The work since MLA is mostly about finding other axes to push, and it’s where most of the 2024–2026 research energy has gone. Three directions are worth knowing, in rough order of how far they depart from vanilla attention.

Share the cache down the stack. The head tricks leave $\text{num\_layers}$ untouched, and that’s a whole axis. Cross-Layer Attention (CLA, MIT, 2024) computes $K$/$V$ in only some layers and lets neighbors reuse them, buying ~2× on top of MQA at near-equal quality. YOCO (Microsoft, 2024) goes further with a decoder-decoder split that stores a single global cache, making cache memory roughly independent of depth — author-reported ~9.4× total-memory savings and dramatic prefill speedups at million-token context. Both stack cleanly with everything in the first half.

Attend to less, and learn what to skip. Sparse attention is old; the 2025 shift is making it trainable from scratch instead of bolted onto a frozen model. DeepSeek’s Native Sparse Attention (NSA, ACL 2025 best paper) runs three gated branches per query — coarse compressed blocks, a few full-resolution selected blocks, and a local sliding window — designed so gradients flow through the selection and the kernels stay GPU-friendly. It reports ~9× faster forward and ~11.6× faster decode at 64K context while matching full attention on quality. Moonshot’s MoBA reaches a similar place from the MoE direction, routing each query to a top-$k$ of key blocks, and is reportedly already serving Kimi’s long-context traffic.

Drop the quadratic entirely. The most radical axis replaces softmax attention with a linear-time recurrence that keeps a fixed-size state instead of a cache that grows with every token. That trades away exact recall — a fixed state can’t losslessly hold unbounded history — which is why production systems hybridize. MiniMax-01 interleaves mostly linear (Lightning Attention) layers with a periodic full-softmax layer across a 456B model to reach multi-million-token context; Jamba mixes Mamba (state-space) layers with a few attention layers at a 7:1 ratio for an 8× smaller cache at 256K. The recurring shape is the same: a cheap linear majority, with just enough full attention sprinkled back in to recover the recall the linear layers lose.

Axis	What it cuts	Representative work	Honest status
Share heads	num_heads	MQA, GQA	shipped everywhere
Compress to a latent	num_heads → latent	MLA; TransMLA / MHA2MLA to convert	shipped (DeepSeek); conversion is new
Share across layers	num_layers	CLA, YOCO	research, author-reported
Trainable sparsity	tokens attended	NSA, MoBA	NSA peer-reviewed; both single-lab
Linear / hybrid	the cache itself	MiniMax-01, Jamba	shipped at scale; recall tax
IO-aware kernel	memory traffic, not cache	FlashAttention 1/2/3	universal, exact

What comes next

Two levers on one formula: sharing and compressing shrink what you cache — MQA, GQA, MLA — and FlashAttention changes how you compute it, exactly. The frontier attacks the terms this post left alone: layers, sparsity, the quadratic itself.

Where it’s heading is already visible: DeepSeek-V3.2 ships MLA’s latent cache paired with sparse attention (DSA), in production. The next post takes the term this one skipped — bytes — and quantizes it. More notes as I go.