The Ultimate Tug-of-War Between Cache and Performance: From MHA, MQA, GQA to MLA
A few days ago, DeepSeek-V2, released by High-Flyer, sparked heated discussion. The first thing that shocked everyone was the price: 1 yuan per million tokens, generally two orders of magnitude cheaper than existing competing APIs — to the point that someone joked "at this price, even if it outputs garbage, I'd consider that garbage a form of art." Second, judging from the technical report, one of the key technologies behind this remarkably low price is the newly proposed MLA (Multi-head Latent Attention), which is an improvement over GQA. It's said to be more efficient and better-performing than GQA, and it has drawn a lot of attention from readers.
In this post, we'll walk through the evolution from MHA, MQA, GQA to MLA together, with a particular focus on explaining the design ideas behind MLA.
MHA
MHA (Multi-Head Attention) is the form of attention proposed in the foundational work Attention is All You Need, and it can be said to be the cornerstone of today's mainstream LLMs. Mathematically, multi-head attention is equivalent to the concatenation of several independent single-head attentions. Suppose the input sequence of (row) vectors is $\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_l$, where $\boldsymbol{x}_i\in\mathbb{R}^d$; then MHA can be formally written as
\begin{equation} \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}_{\leq t}^{(s)} ,\boldsymbol{v}_{\leq t}^{(s)}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}_i^{(s)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\\ \boldsymbol{k}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d\times d_k} \\ \boldsymbol{v}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d\times d_v} \end{gathered} \end{equation}
For simplicity, we omit the scaling factor of the attention matrix here. In practice, a common setting is $d_k = d_v = d / h$; for LLaMA2-7B we have $d=4096, h=32, d_k = d_v = 128$, and for LLaMA2-70B, $d=8192,h=64, d_k = d_v = 128$.
Since here we only consider the Causal Attention used by mainstream autoregressive LLMs, when generating tokens one by one, the newly predicted $t+1$-th token does not affect the already-computed $\boldsymbol{k}_{\leq t}^{(s)} ,\boldsymbol{v}_{\leq t}^{(s)}$. So we can cache these results for reuse in subsequent generation steps, avoiding unnecessary repeated computation — this is the so-called KV Cache.
The later developments of MQA, GQA, and MLA are all products of the same theme: "how to reduce the KV Cache while preserving performance as much as possible."
The Bottleneck
A natural question is: why is reducing the size of the KV Cache so important?
As we all know, LLM inference generally runs on GPUs, and a single GPU's memory is limited. Part of it is used to store the model parameters and forward-pass activations, which depends on the size of the model — once the model is fixed, this is essentially a constant. Another part is used to store the KV Cache, which depends not only on the model size but also on the input length, meaning it grows dynamically during inference. When the context length is long enough, its size can come to dominate, potentially exceeding the total memory of a single card or even a whole machine (8 cards).
The general principle for deploying models on GPUs is: if it can be deployed on a single card, don't spread it across multiple cards; if it can be deployed on a single machine, don't spread it across multiple machines. This is because "intra-card communication bandwidth > inter-card communication bandwidth > inter-machine communication bandwidth." Due to the "bucket effect," the more devices a model deployment spans, the more it is dragged down by inter-device communication bandwidth. In fact, even within a single H100 card, where the bandwidth between SRAM and HBM already reaches 3TB/s, this is still the bottleneck for inference with short contexts — let alone the slower inter-card and inter-machine communication.
So the purpose of reducing the KV Cache is to enable inference over longer contexts on fewer devices, or to allow a larger batch size for inference at the same context length, thereby achieving faster inference speed or greater total throughput. Ultimately, of course, the goal is to reduce inference cost.
For a more detailed understanding of this issue, readers may further consult FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, A Guide to LLM Inference and Performance, LLM Inference Speed of Light, and other articles — I won't elaborate further here (mainly because my own understanding is limited, and I'd rather not risk saying too much and getting it wrong).
MQA
MQA, or "Multi-Query Attention," is a very simple and straightforward attempt to reduce the KV Cache, first proposed in Fast Transformer Decoding: One Write-Head is All You Need, a paper from 2019. This shows that, well before the current LLM boom, reducing the KV Cache was already a topic of great interest to researchers.
The idea behind MQA is simple: let all attention heads directly share the same K and V. In formula terms, this means dropping the head-index superscript ${}^{(s)}$ from all the $\boldsymbol{k},\boldsymbol{v}$ in MHA:
\begin{equation}\require{cancel} \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}_{\leq t}^{\color{#ccc}{\smash{\bcancel{(s)}}}} ,\boldsymbol{v}_{\leq t}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)\boldsymbol{v}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\\ \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \boldsymbol{x}_i\boldsymbol{W}_k^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_k} \\ \boldsymbol{v}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \boldsymbol{x}_i\boldsymbol{W}_v^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_v} \end{gathered} \end{equation}
Models that use MQA include PaLM, StarCoder, and Gemini. Clearly, MQA directly reduces the KV Cache to $1/h$ of its original size, which is a very substantial saving — from a pure memory-saving perspective, this is essentially the ceiling.
In terms of performance, so far it seems that the loss is fairly limited for most tasks, and advocates of MQA believe this loss can be made up through further training. Additionally, because MQA shares K and V, the number of attention parameters is reduced by nearly half; to keep the total parameter count unchanged, the size of the FFN/GLU is usually correspondingly increased, which can also help compensate for some of the performance loss.
GQA
However, some worried that MQA compresses the KV Cache too aggressively, to the point that it might hurt the model's learning efficiency and final performance. To address this, an intermediate version between MHA and MQA — GQA (Grouped-Query Attention) — was proposed, in the paper GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, from last year.
Looking back, the idea behind GQA is also quite simple: it divides all the heads into $g$ groups (where $g$ divides $h$ evenly), with each group sharing the same pair of K and V. Mathematically, this is expressed as
\begin{equation} \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}_{\leq t}^{\color{red}{(\lceil sg/h\rceil)}} ,\boldsymbol{v}_{\leq t}^{\color{red}{(\lceil sg/h\rceil)}}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{red}{(\lceil sg/h\rceil)}}{}^{\top}\right)\boldsymbol{v}_i^{\color{red}{(\lceil sg/h\rceil)}}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{red}{(\lceil sg/h\rceil)}}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\\ \boldsymbol{k}_i^{\color{red}{(\lceil sg/h\rceil)}} = \boldsymbol{x}_i\boldsymbol{W}_k^{\color{red}{(\lceil sg/h\rceil)}}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{\color{red}{(\lceil sg/h\rceil)}}\in\mathbb{R}^{d\times d_k} \\ \boldsymbol{v}_i^{\color{red}{(\lceil sg/h\rceil)}} = \boldsymbol{x}_i\boldsymbol{W}_v^{\color{red}{(\lceil sg/h\rceil)}}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{\color{red}{(\lceil sg/h\rceil)}}\in\mathbb{R}^{d\times d_v} \end{gathered} \end{equation}
Here $\lceil\cdot\rceil$ denotes the ceiling function. GQA provides a natural transition from MHA to MQA: when $g=h$ it reduces to MHA, and when $g=1$ it reduces to MQA. When $1 < g < h$, it only compresses the KV Cache to $g/h$ — a lower compression ratio than MQA, but it also offers greater flexibility and stronger guarantees on performance. The best-known user of GQA is probably Meta's open-sourced LLaMA2-70B, as well as the entire LLaMA3 series. Other models using GQA include TigerBot, DeepSeek-V1, StarCoder2, Yi, ChatGLM2, and ChatGLM3 — considerably more models than use MQA (although ChatGLM describes itself as using MQA, it's actually GQA with $g=2$).
In LLaMA2/3-70B, GQA uses $g=8$, and other similarly-sized models using GQA have basically kept this same setting. This is not a coincidence, but rather is again motivated by inference efficiency considerations. As we know, a model of this size (70B) cannot be deployed on a single card (A100/H100 80G) without extreme quantization. If a single card won't do, then it has to be a single machine — and generally a machine can hold 8 cards. As mentioned earlier, each attention head is actually computed independently and then concatenated; when $g=8$, each card can be made responsible for computing exactly the attention heads corresponding to one K,V group, which maximizes the diversity of K, V while minimizing inter-card communication as much as possible.
MLA
With the groundwork of MHA, MQA, and GQA laid, understanding MLA (Multi-head Latent Attention) becomes relatively easier. DeepSeek-V2's technical report introduces MLA from the perspective of low-rank projection, which has led some readers to ask questions like "LoRA has been around for so long — why did it take until MLA for anyone to apply low-rank decomposition to the KV Cache?"
However, I think that "low-rank projection" is not really the essential angle here, because if we're talking about low-rank projection, in fact simply stacking together all the K, V of GQA already reveals that GQA is itself effectively doing low-rank projection:
\begin{equation}\underbrace{\left[\boldsymbol{k}_i^{(1)},\cdots,\boldsymbol{k}_i^{(g)},\boldsymbol{v}_i^{(1)},\cdots,\boldsymbol{v}_i^{(g)}\right]}_{\boldsymbol{c}_i\in\mathbb{R}^{g(d_k+d_v)}} = \boldsymbol{x}_i \underbrace{\left[\boldsymbol{W}_k^{(1)},\cdots,\boldsymbol{W}_k^{(g)},\boldsymbol{W}_v^{(1)},\cdots,\boldsymbol{W}_v^{(g)}\right]}_{\boldsymbol{W}_c\in\mathbb{R}^{d\times g(d_k+d_v)}}\end{equation}
Here we concatenate all the $\boldsymbol{k}_i^{(s)},\boldsymbol{v}_i^{(s)}$ together and denote it as $\boldsymbol{c}_i$, and likewise concatenate the corresponding projection matrices as $\boldsymbol{W}_c$. Note that generally $d_c = g(d_k+d_v) < d$ holds, so the transformation from $\boldsymbol{x}_i$ to $\boldsymbol{c}_i$ is itself a low-rank projection. So the essential improvement of MLA is not low-rank projection per se, but rather what happens after the low-rank projection.
Part 1
What does GQA do after the projection? First, it splits the vector in half to serve as K and V respectively; then each half is further divided into $g$ parts, and each part is duplicated $h/g$ times, in order to "make up" the K, V needed for $h$ attention heads. We know that splitting and duplication are just simple linear transformations, so MLA's first idea is to replace these simple linear transformations with general linear transformations, to enhance the model's capacity:
\begin{equation} \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}_{\leq t}^{(s)} ,\boldsymbol{v}_{\leq t}^{(s)}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}_i^{(s)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\\ \boldsymbol{k}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d_c\times d_k} \\ \boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_c\times d_v} \\[10pt] \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \end{gathered} \end{equation}
However, although in theory this would increase model capacity, we shouldn't forget that GQA's main purpose is to reduce the KV Cache. For the sake of saving computation and communication costs, what we generally cache is the projected $\boldsymbol{k}_i, \boldsymbol{v}_i$, not the pre-projection $\boldsymbol{c}_i$ or $\boldsymbol{x}_i$. But with this approach, MLA reintroduces distinct projection matrices that once again make all the K, V heads different from one another — so the size of the KV Cache goes right back to being as large as in MHA, defeating the whole purpose of GQA.
In response, MLA discovered that we can exploit the specific form of dot-product attention and sidestep this problem via a simple yet clever identity transformation. First, during training, things proceed as usual, and there isn't much room for optimization here. Then, during inference, we make use of
\begin{equation}\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top} = \left(\boldsymbol{x}_t\boldsymbol{W}_q^{(s)}\right) \left(\boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\right){}^{\top} = \boldsymbol{x}_t\left(\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}\right)\boldsymbol{c}_i^{\top} \end{equation}
This means that, at inference time, we can merge $\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}$ together to serve as the projection matrix for Q, so that $\boldsymbol{c}_i$ takes the place of the original $\boldsymbol{k}_i$. Likewise, following $\boldsymbol{o}_t$ there is another projection matrix, so the $\boldsymbol{W}_v^{(s)}$ of $\boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}$ can also be absorbed into the projection matrix that follows it; equivalently, $\boldsymbol{v}_i$ can also be replaced by $\boldsymbol{c}_i$. In other words, the KV Cache now only needs to store all the $\boldsymbol{c}_i$, rather than storing all of $\boldsymbol{k}_i^{(s)}$ and $\boldsymbol{v}_i^{(s)}$ separately. Notice that $\boldsymbol{c}_i$ has nothing to do with ${}^{(s)}$, meaning it is shared across all heads — that is, at inference time, MLA can be identity-transformed into an MQA.
Let me emphasize again: the theme of this post is, throughout, reducing the KV Cache. So what has MLA achieved up to this point? The answer is: it enhances the capacity of GQA by using distinct projection matrices, while still being able to keep the same KV Cache size at inference time. Conversely, if we only need capacity comparable to GQA, could we shrink the KV Cache even further? In other words, there is no need for $d_c$ to equal $g(d_k+d_v)$ — it could instead take a smaller value (DeepSeek-V2 sets it to 512), further compressing the KV Cache. This is the core idea of MLA.
Additional notes:
1. The identity transformation that merges $\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}$ into a single matrix, in theory, only holds exactly under infinite numerical precision. In practice, if we use single precision — and especially BF16 — the precision loss introduced by this transformation can be quite noticeable, and it can be amplified to a considerable degree after accumulating over many layers.
2. In practice, we generally don't compute Q according to $\boldsymbol{x}_t\left(\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}\right)$, but rather according to $\left(\boldsymbol{x}_t\boldsymbol{W}_q^{(s)}\right)\boldsymbol{W}_k^{(s)}{}^{\top}$. Although this is sequential (rather than fused), under the low-rank assumption it requires less computation and also incurs less theoretical precision loss. That said, in this article we'll continue to describe things in terms of merging $\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}$ into a single matrix, for simplicity of presentation.
Part 2
Everything seems perfect so far — it looks like we're about to arrive at an ideal design that's both effective and cheap. But not so fast — if we think a bit more deeply, we'll find that the version of MLA described so far has an unavoidable flaw: it is incompatible with RoPE (Rotary Position Embedding).
As mentioned earlier, the key step that allows MLA to keep a KV Cache the same size as GQA's is "merging $\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}$ into a single (position-independent) matrix to serve as the projection matrix for Q." But if RoPE is added, this step becomes impossible. This is because RoPE is a position-dependent, $d_k\times d_k$ block-diagonal matrix $\boldsymbol{\mathcal{R}}_m$ satisfying $\boldsymbol{\mathcal{R}}_m\boldsymbol{\mathcal{R}}_n^{\top}=\boldsymbol{\mathcal{R}}_{m-n}$. Once RoPE is added, MLA ends up with an extra term $\boldsymbol{\mathcal{R}}_{t-i}$ inserted between $\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}$:
\begin{equation} \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\quad,\quad\boldsymbol{k}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i} \\ \boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top} = \left(\boldsymbol{x}_t\boldsymbol{W}_q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_t}\right) \left(\boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right){}^{\top} = \boldsymbol{x}_t\left(\boldsymbol{W}_q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_{t-i}}\boldsymbol{W}_k^{(s)}{}^{\top}\right)\boldsymbol{c}_i^{\top} \end{equation}
Here, $\boldsymbol{W}_q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_{t-i}}\boldsymbol{W}_k^{(s)}{}^{\top}$ can no longer be merged into a single fixed projection matrix (since it depends on the position difference $t-i$), which means MLA's approach cannot be combined with RoPE.
Some time ago, I had the honor of discussing this issue with the DeepSeek team, but it turns out to be a genuinely deep problem, and at the time I wasn't able to offer any particularly useful suggestions. The simplest approach would be to abandon RoPE altogether and use some other attention-bias-based positional encoding, such as ALiBi, but DeepSeek's experiments showed it performs noticeably worse than RoPE (note: it's not that MLA cannot add RoPE at all — it's that once RoPE is added, the identity-transformation trick for reducing the KV Cache no longer works). I also suggested trying Sandwich, which, unlike ALiBi, doesn't monotonically decay to negative infinity, and might perform somewhat better — but it feels like treating the symptom rather than the underlying cause. Another compromise would be to also change the input of $\boldsymbol{q}_i$ to $\boldsymbol{c}_i$, and then apply RoPE after $\boldsymbol{c}_i$, i.e.
\begin{equation}\boldsymbol{q}_i^{(s)} = \boldsymbol{c}_i\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\boldsymbol{W}_q^{(s)},\quad\boldsymbol{k}_i^{(s)} = \boldsymbol{c}_i\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\boldsymbol{W}_k^{(s)}\end{equation}
In this way, $\boldsymbol{\mathcal{R}}_i$ could be absorbed into $\boldsymbol{c}_i$, but then there would no longer be any $\boldsymbol{\mathcal{R}}_m\boldsymbol{\mathcal{R}}_n^{\top}=\boldsymbol{\mathcal{R}}_{m-n}$ term. In this case RoPE would no longer be realizing relative position via absolute position; it would simply be adding absolute position information onto Q and K, letting the model figure out relative positional information on its own.
The version of MLA that was ultimately released adopted a hybrid approach: each attention head's Q and K get $d_r$ additional dimensions for adding RoPE, and among these, the extra dimensions added to K are shared across heads:
\begin{equation} \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}_{\leq t}^{(s)} ,\boldsymbol{v}_{\leq t}^{(s)}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}_i^{(s)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \left[\boldsymbol{x}_i\boldsymbol{W}_{qc}^{(s)}, \boldsymbol{x}_i\boldsymbol{W}_{qr}^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_k + d_r},\quad \boldsymbol{W}_{qc}^{(s)}\in\mathbb{R}^{d\times d_k},\boldsymbol{W}_{qr}^{(s)}\in\mathbb{R}^{d\times d_r}\\ \boldsymbol{k}_i^{(s)} = \left[\boldsymbol{c}_i\boldsymbol{W}_{kc}^{(s)}, \boldsymbol{x}_i\boldsymbol{W}_{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_k+d_r},\quad \boldsymbol{W}_{kc}^{(s)}\in\mathbb{R}^{d_c\times d_k}, \boldsymbol{W}_{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_r} \\ \boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_c\times d_v} \\[10pt] \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \end{gathered} \end{equation}
This way, the dimensions without RoPE can repeat the operation from "Part 1," so at inference time the KV Cache only needs to store $\boldsymbol{c}_i$, while the newly added RoPE-carrying dimensions supply the positional information. And since these are shared across all heads, only $d_r$ extra dimensions need to be added to the K Cache. The original paper sets $d_r = d_k / 2 = 64$, which is a modest increase relative to the original $d_c=512$.
Part 3
One last detail: in the final version of MLA, the input for Q was also changed to a low-rank projection form. This is unrelated to reducing the KV Cache — it's mainly meant to reduce the memory footprint of the parameters and the corresponding gradients during training (the original paper says "activations," though personally I find that a bit unclear):
\begin{equation} \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}_{\leq t}^{(s)} ,\boldsymbol{v}_{\leq t}^{(s)}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}_i^{(s)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \left[\boldsymbol{c}_i'\boldsymbol{W}_{qc}^{(s)}, \boldsymbol{c}_i'\boldsymbol{W}_{qr}^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_k + d_r},\quad \boldsymbol{W}_{qc}^{(s)}\in\mathbb{R}^{d_c'\times d_k},\boldsymbol{W}_{qr}^{(s)}\in\mathbb{R}^{d_c'\times d_r}\\ \boldsymbol{k}_i^{(s)} = \left[\boldsymbol{c}_i\boldsymbol{W}_{kc}^{(s)}, \boldsymbol{x}_i\boldsymbol{W}_{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_k+d_r},\quad \boldsymbol{W}_{kc}^{(s)}\in\mathbb{R}^{d_c\times d_k}, \boldsymbol{W}_{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_r} \\ \boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_c\times d_v} \\[10pt] \boldsymbol{c}_i' = \boldsymbol{x}_i \boldsymbol{W}_c'\in\mathbb{R}^{d_c'},\quad \boldsymbol{W}_c'\in\mathbb{R}^{d\times d_c'} \\ \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \\ \end{gathered} \label{eq:mla-mha}\end{equation}
Note that in the second term of $\boldsymbol{k}_i^{(s)}$, the RoPE part, the input is still $\boldsymbol{x}_i$ rather than $\boldsymbol{c}_i$ — this follows the setting in the original paper and is not a typo. The original paper's value for $d_c'$ is 1536, different from $d_c=512$. For comparison, here is the RoPE-equipped MHA formulation:
\begin{equation} \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}_{\leq t}^{(s)} ,\boldsymbol{v}_{\leq t}^{(s)}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}_i^{(s)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\\ \boldsymbol{k}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_k^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d\times d_k} \\ \boldsymbol{v}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d\times d_v} \end{gathered} \end{equation}
We can see that, during training, aside from the extra low-rank projection step and applying RoPE to only part of the dimensions, MLA is essentially no different from MHA with the Q, K head size changed from $d_k$ to $d_k + d_r$.
During the decoding stage, MLA switches to an MQA form:
\begin{equation} \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}\boldsymbol{W}_v^{(1)}, \boldsymbol{o}_t^{(2)}\boldsymbol{W}_v^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\boldsymbol{W}_v^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}_{\leq t}^{\color{#ccc}{\smash{\bcancel{(s)}}}} ,\boldsymbol{c}_{\leq t}\right)\triangleq\frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)\boldsymbol{c}_i}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \left[\boldsymbol{c}_i'\boldsymbol{W}_{qc}^{(s)}\boldsymbol{W}_{kc}^{(s)}{}^{\top}, \boldsymbol{c}_i'\boldsymbol{W}_{qr}^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_c + d_r}\\ \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \left[\boldsymbol{c}_i, \boldsymbol{x}_i\boldsymbol{W}_{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_c+d_r}\\ \boldsymbol{W}_{qc}^{(s)}\in\mathbb{R}^{d_c'\times d_k},\boldsymbol{W}_{kc}^{(s)}\in\mathbb{R}^{d_c\times d_k},\boldsymbol{W}_{qr}^{(s)}\in\mathbb{R}^{d_c'\times d_r},\boldsymbol{W}_{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_r} \\[10pt] \boldsymbol{c}_i' = \boldsymbol{x}_i \boldsymbol{W}_c'\in\mathbb{R}^{d_c'},\quad \boldsymbol{W}_c'\in\mathbb{R}^{d\times d_c'} \\ \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \\ \end{gathered} \label{eq:mla-mqa}\end{equation}
At this point the head size of Q, K becomes $d_c + d_r$, and the head size of V becomes $d_c$ — under the original paper's settings, these are 4 times $d_k$ and $d_v$ respectively. So actually, this conversion that MLA performs during decoding, although it effectively reduces the KV Cache, actually increases the amount of computation involved in decoding.
So why does it still improve inference efficiency? This brings us back to the issue discussed in the "Bottleneck" section. We can divide LLM inference into two parts: generating the first token (Prefill) and generating each subsequent token (Generation). The Prefill stage involves parallel computation over all input tokens, after which the corresponding KV Cache is stored; this part is bottlenecked by compute, bandwidth, and memory all at once, so we can use the MHA form of MLA $\eqref{eq:mla-mha}$ for this computation. But the Generation stage, since it computes only one token per step, is actually mostly bottlenecked by bandwidth and memory; here we can use the MQA form of MLA $\eqref{eq:mla-mqa}$, which noticeably speeds up Generation.
There's another detail that fully reflects this design characteristic. Typically, LLM architecture parameters satisfy $h \times d_k = d$, i.e., num_heads * head_size = hidden_size, but DeepSeek-V2 is different: it has $d_k=128,d=5120$, yet $h=128$ — three times the usual setting! This is because MLA's KV Cache size has nothing to do with $h$, so increasing $h$ only increases computation and boosts model capacity, without increasing the KV Cache — and therefore without creating a speed bottleneck.
Summary
This post has given a brief overview of the evolution of multi-head attention, in particular the shift in design philosophy from MHA to MQA, GQA, and finally to MLA, with a detailed exposition of MLA at the end. In this post, MLA is presented as a generalization of GQA: it replaces GQA's splitting-and-duplication with projection matrices, introduces an identity-transformation trick to further compress the KV Cache, and adopts a hybrid approach to remain compatible with RoPE. All in all, MLA can rightly be considered a highly practical variant of attention.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.