FlashAttention removed the $T\times T$ matrices. It left alone the other memory cost of attention, the one that grows with every token a model generates. To produce token $T+1$, every layer needs the keys and values of all $T$ earlier tokens, so a serving system keeps them in the KV cache (§6.5). For Llama 2 7B in bf16 that is half a megabyte per token. Run it on 32K tokens and one sequence would hold 16 GiB of cache, more than the model's own weights, all of which every decoding step must read from HBM.

And generation keeps getting longer. Reasoning models write thousands of tokens of chain of thought before they answer (§18.5), and reinforcement learning from verifiable rewards spends much of its time generating such rollouts (§18.4), so the cost of decoding, and with it the cache, sits on the critical path of training as well as serving.

So the question for this section is how to make the cache smaller without making the model worse. The answers form a family. Share one key and value head among all query heads (multi-query attention), share them among groups of query heads (grouped-query attention), or compress keys and values together into a small latent vector that every head decompresses for itself (multi-head latent attention). We derive each, implement each, and put them side by side on real models, with the quality each one costs.

The cache, recalled, and why it sets the speed of decoding

§6.5 counted the cache. With $L$ layers, $n_{kv}$ key-value heads of width $d_k$ (take $d_v = d_k$), a context of $T$ tokens, a batch of $B$ sequences, and $b$ bytes per number (§6.5 writes $s$), it holds

one key and one value vector per key-value head per layer for every token, and §6.5 showed that at long context and large batch it outgrows the weights. For ordinary multi-head attention $n_{kv} = h$, the number of query heads, and $h\,d_k = d$. Everything in this section changes one part of that product, the $2n_{kv}d_k$ numbers that each token leaves behind in each layer.

The cache is not only a capacity problem. It sets the speed of decoding. At each step a query head does $2d_k$ FLOPs per cached token for its score and $2d_k$ for its weighted value, $4d_k$ in all, and the cache supplies $2d_k$ numbers per token per key-value head. With $h$ query heads sharing $n_{kv}$ key-value heads, the arithmetic intensity of decode attention is

which in bf16 is just $h/n_{kv}$, the number of query heads that share each key-value head. For multi-head attention that is 1 FLOP per byte, more than a hundred times below an A100's ridge point of about 153 (§12.1). Decode attention is about as memory-bound as a computation can be, and its time per token is the size of the cache divided by the HBM bandwidth. Shrinking the cache speeds it up in direct proportion. Shazeer (2019) made this argument when he introduced the first fix, with the ratio of memory access to arithmetic for incremental multi-head attention growing like $\Theta(T/d + 1/B)$ in our notation, the $T/d$ term coming from the cache.

Notice what is missing from that formula, the batch size. For the weight matrices, batching is the cure for low intensity, because one weight read from HBM serves every sequence in the batch, so the intensity of a linear layer is about $B$ FLOPs per byte (§11.2). The cache gets no such reuse. Each sequence attends to its own keys and values, so a batch of $B$ sequences reads $B$ times the bytes for $B$ times the FLOPs, and the intensity of decode attention stays at $h/n_{kv}$ however large the batch. Larger batches make the weights cheaper per token and leave attention exactly as memory-bound as before, which is why at long context attention, not the weights, ends up dominating decode time. The only levers on that intensity are architectural, how many query heads share each cached vector, which is what the rest of this section is about.

The spread is a factor of about 67 between GPT-3 and DeepSeek-V3, even though DeepSeek-V3 has almost four times as many parameters. GPT-3's 4.5 MiB per token means a single 128K-token sequence would need more than half a terabyte of cache. Nearly every large model released since 2023 uses one of the methods below, and the rest of the section explains them.

def kv_cache_bytes(cfg, T, batch=1, bytes_per=2):
    return 2 * cfg["layers"] * cfg["heads"] * cfg["head_dim"] * T * batch * bytes_per   # MHA for everything

Multi-query attention, one key head for everyone

Shazeer's idea was the most drastic one available. Keep $h$ query heads, each with its own projection, but give all of them a single key head and a single value head,

with $\mathbf{W}_K, \mathbf{W}_V \in \R^{d\times d_k}$ and the mask $\mathbf{M}$ of §4.4. In plain words, every head still asks its own question and gets its own pattern of weights, but all heads look up the same keys and read the same values. The cache shrinks by a factor of $h$, and the decode intensity rises to $h$ FLOPs per byte, because one cached key now serves $h$ dot products.

What does it cost? Shazeer trained Transformer translation models on WMT14 English to German, widening the feed-forward layer of the multi-query model (from 4,096 to 5,440) so both had the same number of parameters. The table collects his numbers (Tables 1 and 2 of the paper, times in TPUv2 microseconds per output token, encoder plus decoder).

Decoding became about 12 times faster for the decoder, at a cost of 0.2 dev BLEU and no loss in test BLEU. On a billion-word language-modeling benchmark, perplexity went from 29.9 to 30.2. PaLM (Chowdhery et al. 2022) adopted multi-query attention for all its sizes, reporting a neutral effect on quality and training speed and a large saving at inference, and so did Falcon 7B.

The catch showed up later and at scale. Ainslie et al. (2023) found multi-query models less stable. T5-Large models trained with it from scratch had frequent loss spikes and diverged immediately when fine-tuned on long-input tasks. Its quality also trails multi-head attention, by a little on their tasks and by more on the hard benchmarks of DeepSeek's later ablation (below). A single key head is a real bottleneck. All $h$ heads must find what they need in the same $d_k$-dimensional keys.

Grouped-query attention, the middle ground

The obvious compromise is to have some key-value heads, but fewer than query heads. Split the $h$ query heads into $n_{kv}$ groups of $g = h/n_{kv}$ consecutive heads, and give each group one key-value head. Query head $i$ (counting from 0) uses key-value head

This is grouped-query attention (GQA, Ainslie et al. 2023). With $n_{kv} = h$ it is multi-head attention, and with $n_{kv} = 1$ it is multi-query attention, so it interpolates between the two. The cache shrinks by $g$, and the decode intensity is $g$ FLOPs per byte in bf16. The mapping is contiguous (heads $0, \dots, g-1$ share key-value head 0), the convention of Llama's repeat_kv and of PyTorch's enable_gqa, which both repeat each key-value head $g$ times in a row. A checkpoint trained with one mapping gives nonsense under another.

Ainslie et al. also showed how to get a GQA model without pretraining from scratch. Take a multi-head checkpoint, build each group's key and value projections by mean-pooling the projections of the heads in the group, and continue pretraining (uptraining) for a small fraction of the original steps, $\alpha = 5\%$. Mean-pooling worked better than keeping the first head of each group, which worked better than random initialization, in the order of how much information each preserves. On T5-XXL, with GQA applied to the decoder's self- and cross-attention, they reported these results (Table 1, average over summarization, translation, and question-answering tasks).

GQA with 8 groups was almost as good as multi-head attention and almost as fast as multi-query attention. The authors chose 8 groups as "a favorable middle ground", and the number stuck. Llama 2 70B, every Llama 3 model, Mistral 7B, and many others use 8 key-value heads, whatever their number of query heads.

Implementing it without copying

The simplest implementation repeats each key-value head $g$ times and runs ordinary multi-head attention. That works, and wastes the memory and bandwidth the method was meant to save. The better way is a reshape. View the queries of shape (B, h, T, d_k) as (B, n_kv, g, T, d_k) and the keys as (B, n_kv, 1, T, d_k), and let broadcasting pair each group of queries with its own key-value head.

Notice what does and does not change. The cache and the key and value projections shrink by $g$. The attention core does not, since every query head still computes $T$ scores against $T$ keys. GQA saves memory and bandwidth, not FLOPs, which is exactly what decoding needs. Decode kernels use the same reshape on the GPU. At one new token per sequence, the $g$ query heads of a group become $g$ rows of a single query tile against one key tile, so each key loaded from HBM serves $g$ dot products, and a small matrix product on tensor cores replaces $g$ matrix-vector products. That is where the intensity $h/n_{kv}$ of the previous section actually comes from in a kernel.

Sharing meets tensor parallelism

Tensor parallelism splits a layer's heads across GPUs (§10.4), and the key-value heads have to split with them. With $t$ GPUs, each holds $h/t$ query heads and, when $t$ divides $n_{kv}$, $n_{kv}/t$ key-value heads, so its share of the cache shrinks by $t$ as well. When $t$ exceeds $n_{kv}$ there are not enough key-value heads to go around. Each must be stored on $t/n_{kv}$ GPUs, and the cache per GPU stops shrinking at one key-value head. Llama 2 70B's 8 key-value heads split evenly over the 8 GPUs of a typical node, while a 16-way split would store every one of them twice. Multi-query attention is the extreme case, its single head copied to every GPU. MLA's latent, below, is shared by all heads in the same way, which is why serving systems such as SGLang offer data-parallel attention for MLA models, giving each GPU different sequences rather than different heads.

import numpy as np

def gqa(Q, K, V, causal=False):
    B, h, T_q, d_k = Q.shape
    n_kv = K.shape[1]
    idx = np.arange(h) % n_kv                     # interleaved mapping: wrong convention
    S = Q @ K[:, idx].swapaxes(-1, -2) / np.sqrt(d_k)
    P = np.exp(S - S.max(-1, keepdims=True)); P /= P.sum(-1, keepdims=True)

import numpy as np

def mha_to_gqa(W_K, W_V, h, n_kv):
    d, hd = W_K.shape
    d_k = hd // h
    return W_K[:, : n_kv * d_k], W_V[:, : n_kv * d_k]       # keeps the first heads, not the groups

Multi-head latent attention, compress instead of share

Sharing throws information away. Every head in a group must use the same keys. DeepSeek-V2 (DeepSeek-AI 2024) asked a different question. Is there a small vector per token from which every head's key and value can be reconstructed? If the keys and values of all heads together have low rank, the answer is yes, and that vector is all the cache needs to hold.

Low-rank joint compression

For token $t$ with hidden state $\mathbf{h}_t\in\R^d$, multi-head latent attention (MLA) computes a latent vector and decompresses it into per-head keys and values,

with the down-projection $\mathbf{W}^{DKV}\in\R^{d_c\times d}$, up-projections $\mathbf{W}^{UK}, \mathbf{W}^{UV}\in\R^{hd_k\times d_c}$, and $d_c \ll hd_k$. (DeepSeek writes vectors as columns, so their matrices act from the left. We follow them in this subsection.) Only $\mathbf{c}_t$ is cached. Queries are compressed the same way, $\mathbf{c}^Q_t = \mathbf{W}^{DQ}\mathbf{h}_t$ and $\mathbf{q}^C_t = \mathbf{W}^{UQ}\mathbf{c}^Q_t$, which saves activation memory in training but does nothing for the cache. DeepSeek-V2 uses $h = 128$ heads of width $d_k = 128$, so $hd_k = 16{,}384$, and compresses keys and values into $d_c = 512$ numbers, with $d_c' = 1{,}536$ for queries.

Decompressing $\mathbf{k}^C_s$ for every cached token at every step would cost $2hd_kd_c$ FLOPs per cached token per step, far more than the $4hd_k$ of the attention itself. The trick that makes MLA practical is that you never have to. Look at head $i$'s score between the query at $t$ and the key at $s$. Write $\mathbf{W}^{UK}_i \in \R^{d_k\times d_c}$ for the rows of $\mathbf{W}^{UK}$ that produce head $i$'s key, so $\mathbf{k}^C_{s,i} = \mathbf{W}^{UK}_i\mathbf{c}_s$. Then

In plain words, instead of decompressing every cached latent into a key, move the decompression matrix onto the query. The query is computed once per step, the keys are many, so this is far cheaper, and the score is a dot product between a $d_c$-dimensional "absorbed" query and the cached latent itself. DeepSeek-V2 says $\mathbf{W}^{UK}$ "can be absorbed into" the query projection.

The values work the same way from the other side. Head $i$'s output is $\sum_s a_{ts,i}\mathbf{W}^{UV}_i\mathbf{c}_s = \mathbf{W}^{UV}_i\big(\sum_s a_{ts,i}\mathbf{c}_s\big)$, a weighted sum of latents, decompressed once, and since the output projection multiplies it next, $\mathbf{W}^{UV}_i$ can be folded into $\mathbf{W}_O$. So at inference MLA is multi-query attention, with one shared "key" and one shared "value" per token, both equal to the latent $\mathbf{c}_s$ of width 512 (the rotary part of the next subsection adds 64 more numbers to the key), and with per-head queries that live in the latent space. That is why its decode intensity is high. Every one of the 128 heads reads the same cached vector.

Why RoPE gets in the way, and the decoupled fix

Rotary embeddings (§5.5) rotate queries and keys by position-dependent matrices $\mathbf{R}_t$ and $\mathbf{R}_s$. Apply them to the decompressed key and the score becomes

A matrix that depends on the pair of positions now sits between the query and $\mathbf{W}^{UK}_i$. The absorbed query $\mathbf{W}^{UK\top}_i\mathbf{q}$ would have to be recomputed for every cached position, because matrix products do not commute and $\mathbf{R}_{s-t}$ cannot be moved out of the way. The paper's verdict is that keys for all prefix tokens would have to be recomputed during inference, which defeats the purpose.

DeepSeek's fix is to decouple position from content. Each query gets a few extra dimensions $\mathbf{q}^R_{t,i} = \operatorname{RoPE}(\mathbf{W}^{QR}_i\mathbf{c}^Q_t)\in\R^{d_r}$ that carry RoPE, and each token gets one extra rotary key $\mathbf{k}^R_t = \operatorname{RoPE}(\mathbf{W}^{KR}\mathbf{h}_t)\in\R^{d_r}$ shared by all heads, MQA-style. The full query and key of head $i$ are concatenations, and the score is a sum of a content part and a position part,

The content part has no rotation, so it can be absorbed. The rotary part is a small ordinary dot product with a cached, already rotated key. The cache per token per layer is $\mathbf{c}_t$ and $\mathbf{k}^R_t$, $d_c + d_r$ numbers. DeepSeek-V2 uses $d_r = 64 = d_k/2$ and $d_c = 512 = 4d_k$, so the cache is $\tfrac{9}{2}d_k = 576$ numbers per token per layer, which the paper notes is the same as GQA with only 2.25 groups.

The cell checks the whole construction. It builds a small MLA layer, computes one decoding step the explicit way (decompress every cached key and value) and the absorbed way (attend in the latent space), and compares.

The two agree to machine precision. In a real model the absorbed matrices are precomputed once, and the absorbed path never forms a per-head key at all.

import numpy as np

def mla_step(x, pos, cache_c, cache_kr):
    return mla_reference(x, pos, cache_c, cache_kr)       # decompresses the whole cache

How good is it?

DeepSeek-V2 reports two ablations that, together, are the best public evidence on this trade-off (Appendix D). With dense 7B models trained on 1.33T tokens, and parameter counts matched by adjusting depth, multi-head attention beat GQA with 8 groups, which beat MQA, on every benchmark they list.

Then, comparing MLA with MHA in mixture-of-experts models (§13.5), MLA used 14% of MHA's cache in a 16B-parameter model (15.6K against 110.6K elements per token) and 4% in a 250B model, and scored higher on most benchmarks, for example MMLU 50.0 against 48.7 in the small model and 59.0 against 57.5 in the large one, though not all (C-Eval in the small model was 50.9 against 51.6). These are results from one lab's training setup, not a law. They do suggest that compressing keys and values loses less than sharing them.

The headline number needs care. DeepSeek-V2 reports a 93.3% smaller KV cache than DeepSeek 67B. That figure includes quantizing the cache to about 6 bits per element on average. By element count, MLA alone shrinks it by about 82%, as the cell computes from the two configurations.

The last lines show the other half of the story. Decode attention with GQA-8 runs at 4 FLOPs per byte, and with MLA at over 200, close to an H100's ridge point. MLA turns decoding from a pure bandwidth problem into something close to a balanced one, which is why it pairs so well with large batches. Its costs are complexity (absorbed weights, the decoupled rotary part, a kernel specialized for a 576-wide key and 512-wide value) and more FLOPs per cached token, which matters for prefill. For that reason serving implementations typically use the absorbed form only for decoding. In prefill, which is compute-bound anyway, they decompress the keys and values and run ordinary multi-head attention, with query-key width $d_k + d_r = 192$ and value width 128 per head. DeepSeek-V3 and Kimi K2 kept MLA with the same $d_c = 512$ and $d_r = 64$.

def decode_attention_time(cfg, T, batch, peak, bw, bytes_per=2):
    L, h, d_k = cfg["layers"], cfg["heads"], cfg["head_dim"]
    flops = L * h * 4 * d_k * T * batch
    nbytes = L * h * d_k * T * bytes_per                 # forgets K and V, the kv heads, and the batch
    return flops / peak + nbytes / bw, "memory"

The trade-offs, side by side

Two other levers multiply with these, and are covered elsewhere. Fewer tokens in the cache, by attending only to a window of recent tokens plus a few sinks, is the subject of §12.6. Fewer bits per number, storing the cache in 8 or 4 bits, and paging it into fixed-size blocks so that sequences of different lengths share memory without fragmentation (the PagedAttention of §12.4), are what serving systems built around paged attention and quantized KV caches do. Sharing keys and values across adjacent layers (cross-layer attention, Brandon et al. 2024) is a further option some models use.

In PyTorch

A grouped-query attention layer is a few lines with enable_gqa=True, which passes the smaller key and value tensors straight to the kernel. The module below caches keys and values for decoding and reports how many bytes its cache holds.

Read more

• Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need. Multi-query attention, with the memory-bandwidth analysis of incremental decoding that motivates everything here.
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Grouped-query attention, uptraining, and the comparison of conversion methods.
• DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. Section 2.1 derives MLA and decoupled RoPE, and Appendix D has the MHA, GQA, MQA, and MLA ablations.
• Chowdhery et al. (2022). PaLM: Scaling Language Modeling with Pathways. The largest model to use multi-query attention, with notes on why.
• Brandon et al. (2024). Reducing Transformer Key-Value Cache Size with Cross-Layer Attention. Sharing keys and values between adjacent layers, a complement to sharing them between heads.
• Pope et al. (2022). Efficiently Scaling Transformer Inference. How the KV cache, multi-query attention, and partitioning interact when serving very large models.

Watch next

▶

▶

▶