1. Opening

This note uses one lens throughout: write attention as tensor contractions, while keeping every dimension explicit. From this view, several seemingly separate MLA phenomena can be understood more deeply and more uniformly. Specifically, we get a few interesting conclusions:

• In prefill, MLA reduces the FLOPs of the linear projections, but it cannot absorb too early: the quadratic core-attention term should still be computed in $d_k,d_v$, not $d_c$.
• In decode, the best order changes because the cache already exists. After absorbing key/value projections, the length-$t$ cache computation can go through the latent state instead of reconstructed per-head $K,V$.
• RoPE gets in the way exactly because it makes the latent-to-key map position-dependent, which breaks full key absorption.
• Q compression mainly reduces training activation memory and prefill/training projection FLOPs, while contributing little to decode.
• MTP-style speculative decoding behaves differently on MLA because MLA decode no longer reads the full per-head KV cache; it reads the compressed latent and RoPE caches and performs additional contractions instead.

2. Tensor View

Takeaway. Unlike a flattened matrix view, the tensor view keeps every meaningful axis explicit. FLOPs can be read directly from the output and contracted axes, while different computation orders can be compared in the same notation.

The tensor view is just this: every tensor contraction is determined by three kinds of axes — output axes, contracted axes, and broadcast axes. Once these axes are explicit, FLOPs become shape accounting.

• $\underset{(h,T,d_k)}{Q}=\underset{(T,d)}{X}\underset{(h,d,d_k)}{W_Q}\;\Longleftrightarrow\;\mathrm{einsum}(\texttt{"td,sdk->stk"},X,W_Q).$

• $\underset{(h,T,T)}{S}=\underset{(h,T,d_k)}{Q}\underset{(h,d_k,T)}{K^\top}\;\Longleftrightarrow\;\mathrm{einsum}(\texttt{"stk,suk->stu"},Q,K).$

General Form. For a batched contraction whose output axes are $B_1,\ldots,B_r,m,n$ and whose contracted axis is $k$, the FLOPs count is:

$$\mathrm{FLOPs} \;=\; 2 \cdot B_1 \cdots B_r \cdot m \cdot n \cdot k .$$

In words: multiply every output dimension by the contracted dimension, then multiply by $2$ for multiply-adds (the standard GEMM convention counts one fused multiply-add as one multiplication plus one addition). Thus $Q=XW_Q$ costs $2hTdd_k$, while $S=QK^\top$ costs $2hT^2d_k$.

3. Notation and Preliminaries

Takeaway. This section fixes notation by annotating shapes directly on the equations, then gives a basic introduction to MHA, MLA, and the RoPE rotation used later.

For broader background from MHA/MQA/GQA to MLA, see Jianlin Su's note. Here we do not try to survey those variants. We use MHA as the clean algebraic reference, MLA as the object of the derivation, and introduce only the RoPE identities needed for the later absorption discussion.

To keep the derivations simple, we use the following conventions:

• Elementwise attention scores are always written as $S_{s,t,u}$: $h$ is the number of heads, $s=1,\ldots,h$ selects one head, $t$ is the query position, and $u$ is the key position. Thus rows of the score matrix correspond to $t$, and columns correspond to $u$.
• RoPE is omitted from the MLA equations until Section 5.
• We ignore the causal mask; readers can replace the full attention matrix with its causal form directly.
• Softmax is omitted from the FLOPs comparison. MHA and MLA apply the same softmax to the same-shaped attention scores, while model FLOPs use the dominant GEMMs and tensor contractions as the comparable measure. Optimized implementations such as PyTorch SDPA also fuse softmax into scaled-dot-product attention. Section 5 still counts the explicit addition of the content and RoPE scores, because it is an extra operation introduced by the decoupled branch and scales with the full attention matrix.
• $A^\top$ means transposing only the last two axes, leaving batch axes untouched.

---

MHA. The reference form is:

$$\underset{(h,T,d_k)}{Q} = \underset{(T,d)}{X} \underset{(h,d,d_k)}{W_Q}, \qquad \underset{(h,T,d_k)}{K} = \underset{(T,d)}{X} \underset{(h,d,d_k)}{W_K}, \qquad \underset{(h,T,d_v)}{V} = \underset{(T,d)}{X} \underset{(h,d,d_v)}{W_V}.$$

Then attention is:

$$\underset{(h,T,T)}{S} = \underset{(h,T,d_k)}{Q} \underset{(h,d_k,T)}{K^\top}, \qquad \underset{(h,T,T)}{P} = \operatorname{softmax}\!\left(\frac{S}{\sqrt{d_k}}\right),$$ $$\underset{(h,T,d_v)}{O} = \underset{(h,T,T)}{P} \underset{(h,T,d_v)}{V}, \qquad \underset{(T,d)}{Y} = \underset{(h,T,d_v)}{O} \underset{(h,d_v,d)}{W_o}.$$

The last line contracts both the head axis and the $d_v$ axis. This is equivalent to the usual implementation that concatenates all heads and then applies one output projection: flattening $(h,d_v)$ into a single $hd_v$ axis does not change the sum,

$$Y_{t,d}=\sum_{s=1}^{h}\sum_{v=1}^{d_v}O_{s,t,v}(W_o)_{s,v,d} =\sum_{j=1}^{hd_v}(O_{\mathrm{concat}})_{t,j}(W_{o,\mathrm{concat}})_{j,d}.$$

The FLOPs are also identical. In the tensor-contraction form, the output axes are $T$ and $d$, and the contracted axes are $h$ and $d_v$, giving $2Tdh d_v$. After concatenation, the multiplication has shapes $(T,hd_v)(hd_v,d)$, giving $2T(hd_v)d=2Thd_vd$.

---

MLA. The KV path replaces per-head $K,V$ projections with a shared latent. To keep the derivation simple, we temporarily omit RoPE from the MLA form until Section 5:

$$\underset{(T,d_c)}{C^{KV}} = \underset{(T,d)}{X} \underset{(d,d_c)}{W_{DKV}}, \qquad \underset{(h,T,d_k)}{K} = \underset{(T,d_c)}{C^{KV}} \underset{(h,d_c,d_k)}{W_{UK}},$$ $$\underset{(h,T,d_v)}{V} = \underset{(T,d_c)}{C^{KV}} \underset{(h,d_c,d_v)}{W_{UV}}.$$

The DeepSeek-V2 paper claims that query compression reduces activation memory during training; DeepSeek-V2-Lite, however, does not compress queries. This gives us a natural comparison: first derive MLA without query compression, then add it back and isolate what changes.

No query compression.

$$\underset{(h,T,d_k)}{Q} = \underset{(T,d)}{X} \underset{(h,d,d_k)}{W_Q}.$$

Query compression.

$$\underset{(T,d_{qc})}{C^Q} = \underset{(T,d)}{X} \underset{(d,d_{qc})}{W_{DQ}}, \qquad \underset{(h,T,d_k)}{Q} = \underset{(T,d_{qc})}{C^Q} \underset{(h,d_{qc},d_k)}{W_{UQ}}.$$

---

RoPE.

In the tensor view, RoPE inserts a position-dependent linear map on the last tensor dimension without changing the tensor shape. Because RoPE acts on $Q$ and $K$, we use their tensor dimension $d_k$ as the rotation dimension. For even $d_k$, the standard rotation matrix is

$$R_t= \operatorname{diag}\!\left( \begin{bmatrix} \cos(t\theta_1) & -\sin(t\theta_1)\\ \sin(t\theta_1) & \cos(t\theta_1) \end{bmatrix}, \ldots, \begin{bmatrix} \cos(t\theta_{d_k/2}) & -\sin(t\theta_{d_k/2})\\ \sin(t\theta_{d_k/2}) & \cos(t\theta_{d_k/2}) \end{bmatrix} \right) \in\mathbb{R}^{d_k\times d_k}.$$

Writing the RoPE base as $b$, we have $\theta_j=b^{-2(j-1)/d_k}$. RoFormer uses $b=10000$; modern language models may use a larger base for longer contexts. Since we use row vectors, stack the transposed position rotations into one tensor: $\underset{(T,d_k,d_k)}{\mathcal R}$, where $\mathcal R_{t,:,:}=R_t^\top$.

RoPE and the subsequent score contraction are then

$$\underset{(h,T,d_k)}{\widetilde Q} = \underset{(h,T,d_k)}{Q} \underset{(T,d_k,d_k)}{\mathcal R}, \qquad \underset{(h,T,d_k)}{\widetilde K} = \underset{(h,T,d_k)}{K} \underset{(T,d_k,d_k)}{\mathcal R}, \qquad \underset{(h,T,T)}{S} = \underset{(h,T,d_k)}{\widetilde Q} \underset{(h,d_k,T)}{\widetilde K^\top}.$$

The two $T$ axes are identified as the same position axis, so only one $T$ remains in the output. The $d_k$ axis of $Q$ or $K$ contracts with the first $d_k$ axis of $\mathcal R$, while the second $d_k$ axis remains; $\mathcal R$ is broadcast across the head axis. Therefore the $(h,T,d_k)$ shapes of $Q,K$ are unchanged. Since $R_t^\top R_u=R_{u-t}$,

$$S_{s,t,u} = Q_{s,t,:}R_t^\top R_uK_{s,u,:}^\top = Q_{s,t,:}R_{u-t}K_{s,u,:}^\top,$$

so each rotation encodes absolute position, while the score depends only on relative position.

The actual elementwise computation. The rotation tensor above is an algebraic representation; an implementation need not construct a dense $d_k\times d_k$ matrix. Let $m=d_k/2$. For any $q\in\mathbb R^{d_k}$, define the adjacent-pair coordinates

$$q_{j,0}:=q_{2j-1}, \qquad q_{j,1}:=q_{2j}, \qquad j=1,\ldots,m.$$

Using the same $\theta_j$ as above, the corresponding phase tensor is

$$\widehat\Theta_{t,2j-1} = \widehat\Theta_{t,2j} = t\theta_j, \qquad j=1,\ldots,m, \qquad \underset{(T,d_k)}{\widehat\Theta}.$$

The implementation then directly computes

$$\underset{(h,T,d_k)}{\widetilde Q} = \underset{(h,T,d_k)}{Q} \odot \cos\underset{(T,d_k)}{\widehat\Theta} + \operatorname{rotate\_half}\!\left( \underset{(h,T,d_k)}{Q} \right) \odot \sin\underset{(T,d_k)}{\widehat\Theta},$$ $$\underset{(h,T,d_k)}{\widetilde K} = \underset{(h,T,d_k)}{K} \odot \cos\underset{(T,d_k)}{\widehat\Theta} + \operatorname{rotate\_half}\!\left( \underset{(h,T,d_k)}{K} \right) \odot \sin\underset{(T,d_k)}{\widehat\Theta}.$$

Here $\odot$ is elementwise multiplication; the cosine and sine tensors broadcast across the head axis. In the pair notation above, rotate_half is defined by

$$\operatorname{rotate\_half}(q)_{j,0}=-q_{j,1}, \qquad \operatorname{rotate\_half}(q)_{j,1}=q_{j,0}.$$

It preserves the shape and removes no axis. This elementwise formulation is exactly equivalent to the rotation-tensor formulation above. We include it in the tensor view, but do not count it separately in the later FLOPs comparison of the main tensor contractions.

4. Computation Order and FLOPs in Prefill, Training, and Decode

We derive each tensor contraction and its FLOPs, comparing MHA with MLA.

4.1 Prefill and Training: The Standard Computation Order

Takeaway. MLA reduces the FLOPs of the linear projections. It does not reduce the $T^2$ core-attention FLOPs, because $QK^\top$ and $PV$ still contract over $d_k$ and $d_v$, respectively.

---

MHA.

For MHA prefill, the complete score expression is

$$S = QK^\top = \underset{(T,d)}{X} \underset{(h,d,d_k)}{W_Q} \underset{(h,d_k,d)}{W_K^\top} \underset{(d,T)}{X^\top}.$$

How to compute $S$ in a cost-effective way? From a FLOPs perspective, the important point is to minimize the coefficient of the dominating $T^2$ term. Since $d_k<d$, the computation order is to first construct $Q=XW_Q$ and $K=XW_K$ to reduce the $d$ feature and keeping $d_k$, then we compute $S=QK^\top$:

$$\begin{aligned} \underset{(T,d)}{X} \underset{(h,d,d_k)}{W_Q} &\rightarrow \underset{(h,T,d_k)}{Q}, &\qquad \mathrm{FLOPs}&=2Thdd_k, \\ \underset{(T,d)}{X} \underset{(h,d,d_k)}{W_K} &\rightarrow \underset{(h,T,d_k)}{K}, &\qquad \mathrm{FLOPs}&=2Thdd_k, \\ \underset{(h,T,d_k)}{Q} \underset{(h,d_k,T)}{K^\top} &\rightarrow \underset{(h,T,T)}{S}, &\qquad \mathrm{FLOPs}&=2hT^2d_k. \end{aligned}$$

The complete value and output chain is:

$$\underset{(h,T,T)}{P} \underset{(T,d)}{X} \underset{(h,d,d_v)}{W_V} \underset{(h,d_v,d)}{W_o} \rightarrow \underset{(T,d)}{Y}.$$

We still focus on the $T^2$ term as our optimization goal. Here the ordering question is which tensor dimension accompanies the quadratic multiplication by $P$. Constructing $V$ first ensures that attention operates on $d_v$, rather than carrying the wider model dimension $d$ through all token pairs.

use $(P(XW_V))W_o$, not $((PX)W_V)W_o$ or $P(X(W_VW_o))$.

The standard order is $XW_V\rightarrow V$, $PV\rightarrow O$, and $OW_o\rightarrow Y$:

$$F_{(P(XW_V))W_o} = \underbrace{2Thdd_v}_{XW_V} +\underbrace{2hT^2d_v}_{PV} +\underbrace{2Thd_vd}_{OW_o}.$$

Computing $PX$ first gives

$$F_{((PX)W_V)W_o} = \underbrace{2hT^2d}_{PX} +\underbrace{2Thdd_v}_{(PX)W_V} +\underbrace{2Thd_vd}_{OW_o}.$$

Precomposing $W_VW_o$ gives

$$F_{P(X(W_VW_o))} = \underbrace{2hd^2d_v}_{W_VW_o} +\underbrace{2hTd^2}_{X(W_VW_o)} +\underbrace{2hT^2d}_{P[X(W_VW_o)]}.$$

Both alternatives attach the quadratic contraction to $d$ rather than $d_v$. Since $d_v\ll d$, the standard order is cheaper: $2hT^2d_v$ instead of $2hT^2d$.

Combining these terms, the dense MHA prefill FLOPs are

$$F_{\mathrm{MHA,prefill}} =\underbrace{2Tdh(2d_k+d_v)}_{XW_Q,\,XW_K,\,XW_V} +\underbrace{2hT^2(d_k+d_v)}_{QK^\top,\,PV} +\underbrace{2Thd_vd}_{OW_o}.$$

---

MLA.

For MLA prefill, first compute the shared KV latent $C^{KV}=XW_{DKV}$. Without query compression, the complete score expression is

$$S = \underset{(T,d)}{X} \underset{(h,d,d_k)}{W_Q} \underset{(h,d_k,d_c)}{W_{UK}^\top} \underset{(d_c,d)}{W_{DKV}^\top} \underset{(d,T)}{X^\top}.$$

With query compression, it becomes

$$S = \underset{(T,d)}{X} \underset{(d,d_{qc})}{W_{DQ}} \underset{(h,d_{qc},d_k)}{W_{UQ}} \underset{(h,d_k,d_c)}{W_{UK}^\top} \underset{(d_c,d)}{W_{DKV}^\top} \underset{(d,T)}{X^\top}.$$

Every way of forming the score has output shape $(h,T,T)$, so the factor $hT^2$ is fixed; we only need to compare the contracted tensor dimension. We therefore want this dimension to be $d_k$ ($d_k=128<d_c=512<d=5120$ in DeepSeek-V2). Thus, we can divide the formula into two parts: $XW_{DQ}W_{UQ}$ and $XW_{DKV}W_{UK}$.

Query-projection order: use $(XW_{DQ})W_{UQ}$, not $X(W_{DQ}W_{UQ})$.

Without query compression, $Q=XW_Q$ is the same as in MHA (FLOPs: $2Thdd_k$). It contains only one contraction, so there is no query-side computation order to choose. The comparison below appears only after $W_Q$ is factorized into $W_{DQ}W_{UQ}$.

With query compression, the factorized order first maps each token to $C^Q$, then expands the head axis through $W_{UQ}$:

$$\begin{aligned} \underset{(T,d)}{X} \underset{(d,d_{qc})}{W_{DQ}} &\rightarrow \underset{(T,d_{qc})}{C^Q}, &\qquad \mathrm{FLOPs}&=2Tdd_{qc}, \\ \underset{(T,d_{qc})}{C^Q} \underset{(h,d_{qc},d_k)}{W_{UQ}} &\rightarrow \underset{(h,T,d_k)}{Q}, &\qquad \mathrm{FLOPs}&=2Thd_{qc}d_k. \end{aligned}$$

If we compose the weights first, the two contractions are

$$\begin{aligned} \underset{(d,d_{qc})}{W_{DQ}} \underset{(h,d_{qc},d_k)}{W_{UQ}} &\rightarrow \underset{(h,d,d_k)}{W_Q^{\mathrm{eff}}}, &\qquad \mathrm{FLOPs}&=2hdd_{qc}d_k, \\ \underset{(T,d)}{X} \underset{(h,d,d_k)}{W_Q^{\mathrm{eff}}} &\rightarrow \underset{(h,T,d_k)}{Q}, &\qquad \mathrm{FLOPs}&=2Thdd_k. \end{aligned}$$

The activation-side costs of the two orders are therefore

$$\begin{aligned} F_{(XW_{DQ})W_{UQ}} &=2Tdd_{qc}+2Thd_{qc}d_k =66.1\text{M}T, \\ F_{X(W_{DQ}W_{UQ})} &=2Thdd_k =167.8\text{M}T \qquad\text{(DeepSeek-V2)}. \end{aligned}$$

where the second expression excludes the additional weight-precomposition cost $2hdd_{qc}d_k$. The factorized order is cheaper when

$$d_{qc}(d+hd_k)<hdd_k.$$

The saving comes from introducing $C^Q:(T,d_{qc})$ between the two contractions, so the per-head contraction uses $d_{qc}$ rather than $d$.

Similarly, for the key and score order, compute $C^{KV}=XW_{DKV}$, reconstruct $K=C^{KV}W_{UK}$, then form $S=QK^\top$.

With the DeepSeek-V2 dimensions, the key-projection order $(XW_{DKV})W_{UK}$ costs $22.0\text{M}T$, while $X(W_{DKV}W_{UK})$ costs $167.8\text{M}T$, excluding the one-time cost of composing the weights.

Value and output order: use $(P(C^{KV}W_{UV}))W_o$, not $((PC^{KV})W_{UV})W_o$ or $P(C^{KV}(W_{UV}W_o))$.

The standard order reconstructs $V$ before $PV$:

$$F_{(P(C^{KV}W_{UV}))W_o} =\underbrace{2Thd_cd_v}_{C^{KV}W_{UV}} +\underbrace{2hT^2d_v}_{PV} +\underbrace{2Thd_vd}_{OW_o}.$$

Computing $PC^{KV}$ first changes the quadratic term to $2hT^2d_c$. Precomposing $W_{UV}W_o$ instead changes it to $2hT^2d$. Since $d_v<d_c<d$, reconstructing $V$ first keeps the quadratic contraction in the smallest tensor dimension, $d_v$.

Without query compression, the dense MLA prefill FLOPs are

$$\begin{aligned} F_{\mathrm{MLA,prefill,noQ}} &=\underbrace{2Thdd_k}_{Q=XW_Q} \\ &\quad+\underbrace{2Tdd_c}_{C^{KV}=XW_{DKV}} +\underbrace{2Thd_c(d_k+d_v)}_{K=C^{KV}W_{UK},\;V=C^{KV}W_{UV}} \\ &\quad+\underbrace{2hT^2(d_k+d_v)}_{S=QK^\top,\;O=PV} +\underbrace{2Thd_vd}_{Y=OW_o}. \end{aligned}$$

With query compression, they are

$$\begin{aligned} F_{\mathrm{MLA,prefill,Qcomp}} &=\underbrace{2Tdd_{qc}}_{C^Q=XW_{DQ}} +\underbrace{2Thd_{qc}d_k}_{Q=C^QW_{UQ}} \\ &\quad+\underbrace{2Tdd_c}_{C^{KV}=XW_{DKV}} +\underbrace{2Thd_c(d_k+d_v)}_{K=C^{KV}W_{UK},\;V=C^{KV}W_{UV}} \\ &\quad+\underbrace{2hT^2(d_k+d_v)}_{S=QK^\top,\;O=PV} +\underbrace{2Thd_vd}_{Y=OW_o}. \end{aligned}$$

---

Comparison.

Prefill FLOPs consist of three parts: the linear projections, the core-attention computation, and the final output projection. The table summarizes the derivation and substitutes the DeepSeek-V2 dimensions; the figure then plots the resulting FLOPs as functions of sequence length $T$, except for the RoPE branch.

4.2 Decode: Latent Caching and Matrix Absorption

Takeaway. In one decode step, the current-query length is $1$ while the cache length is $t$. The cheaper order keeps this $1$ outside, contracts away the other dimensions as early as possible, and expands only afterward. This avoids reconstructing full $K,V$ for all $t$ cached tokens: MLA's FLOPs grow faster with $t$, but it reads much less cache data.

---

MHA.

Before decoding the current position $t$, the historical $K_{<t},V_{<t}$ have already been computed in previous decode steps. The current step only computes the new $q_t,k_t,v_t$; after $k_t,v_t$ are appended, the cache read for this step is $K_{\le t},V_{\le t}$. Across all heads, the current-token projection cost is $2dh(2d_k+d_v)$.

Unlike prefill, there is only one new query, so the quadratic $T^2$ interaction disappears. The growing factor is now $t$: the current query scans $t$ cached keys, and the resulting probabilities combine $t$ cached values.

The score, value, and output contractions are

$$\begin{aligned} \underset{(h,1,d_k)}{q} \underset{(h,d_k,t)}{K_{\le t}^\top} &\rightarrow \underset{(h,1,t)}{\mathrm{score}}, &\qquad \mathrm{FLOPs}&=2htd_k, \\ \underset{(h,1,t)}{p} \underset{(h,t,d_v)}{V_{\le t}} &\rightarrow \underset{(h,1,d_v)}{o}, &\qquad \mathrm{FLOPs}&=2htd_v, \\ \underset{(h,1,d_v)}{o} \underset{(h,d_v,d)}{W_o} &\rightarrow \underset{(1,d)}{y}, &\qquad \mathrm{FLOPs}&=2hd_vd. \end{aligned}$$

Both chains follow the same principle: start from the singleton current-token state and choose the contraction order that keeps every intermediate as small as possible.

The full value-output chain, with every dimension explicit, is

$$\underset{(1,d)}{y} = \underset{(h,1,t)}{p} \underset{(h,t,d_v)}{V_{\le t}} \underset{(h,d_v,d)}{W_o}.$$

The two associations have different costs:

$$F_{(pV_{\le t})W_o} = \underbrace{2htd_v}_{pV_{\le t}:(h,1,d_v)} +\underbrace{2hd_vd}_{(pV_{\le t})W_o:(1,d)},$$ $$F_{p(V_{\le t}W_o)} = \underbrace{2htd_vd}_{V_{\le t}W_o:(h,t,d)} +\underbrace{2htd}_{p(V_{\le t}W_o):(1,d)}.$$

Evaluating from left to right keeps the singleton query axis as an output axis and contracts away the cache-length axis before the output projection. The other order applies $W_o$ to all $t$ cached values, carrying both $t$ and $d$ in the intermediate, which is more costly.

The score path follows the same principle. With dimensions explicit,

$$\underset{(1,d)}{x_t} \underset{(h,d,d_k)}{W_Q} \underset{(h,d_k,t)}{K_{\le t}^{\top}} \rightarrow \underset{(h,1,t)}{\mathrm{score}_t}.$$

Keep the singleton query axis $1$ outside and contract the other feature and cache axes around it, rather than expanding them first.

Therefore, the cached MHA decode FLOPs are

$$F_{\mathrm{MHA,decode,cached}} =\underbrace{2dh(2d_k+d_v)}_{x_tW_Q,\,x_tW_K,\,x_tW_V} +\underbrace{2ht(d_k+d_v)}_{qK_{\le t}^\top,\,pV_{\le t}} +\underbrace{2hd_vd}_{oW_o}.$$

---

MLA.

MLA caches the shared latent rather than per-head $K,V$. At this single decode step, the new latent is

$$\underset{(1,d)}{x_t} \underset{(d,d_c)}{W_{DKV}} \rightarrow \underset{(1,d_c)}{c_t^{KV}}, \qquad \mathrm{FLOPs}=2dd_c,$$

and the cache after appending it is $C^{KV}_{\le t}=[C^{KV}_{<t};c_t^{KV}]$ with shape $(t,d_c)$. We now write the complete contraction chains first, without assuming an association.

Score chain. It is enough to analyze the Q-compressed case. The complete contraction is

$$\underset{(h,1,t)}{\mathrm{score}_t} = \underset{(1,d)}{x_t} \underset{(d,d_{qc})}{W_{DQ}} \underset{(h,d_{qc},d_k)}{W_{UQ}} \underset{(h,d_k,d_c)}{W_{UK}^{\top}} \underset{(d_c,t)}{(C^{KV}_{\le t})^{\top}}.$$

First consider the coefficient of $t$. The cache is the only factor carrying the $t$ axis. If every factor on its left is contracted first, the left intermediate has shape $(h,1,d_c)$, and the final contraction costs

$$2htd_c.$$

This is the smallest possible length-$t$ term: $h$, $t$, and $d_c$ are unavoidable, while the remaining query axis has size $1$. Applying $C^{KV}_{\le t}$ earlier would retain an additional feature axis in a length-$t$ intermediate. Therefore, the cache should be contracted last:

$$\left( \underset{(1,d)}{x_t} \underset{(d,d_{qc})}{W_{DQ}} \underset{(h,d_{qc},d_k)}{W_{UQ}} \underset{(h,d_k,d_c)}{W_{UK}^{\top}} \right) \underset{(d_c,t)}{(C^{KV}_{\le t})^{\top}}.$$

It remains to choose the association of the four-factor prefix

$$\underset{(1,d)}{x_t} \underset{(d,d_{qc})}{W_{DQ}} \underset{(h,d_{qc},d_k)}{W_{UQ}} \underset{(h,d_k,d_c)}{W_{UK}^{\top}}.$$

The table lists the five online binary associations, followed by every distinct offline precomposition of two or three adjacent weights. Excluding the common final cache contraction $2htd_c$, their per-step costs are shown below; the one-time offline composition costs are omitted.

With the DeepSeek-V2 dimensions, the first row remains the minimum even after allowing all offline precompositions. It also has the clearest interpretation: multiplying by the singleton $x_t$ first eliminates $d$, then the same length-$1$ axis is carried through $d_{qc}$, $d_k$, and $d_c$. The chosen order is therefore left to right, with total cost

$$F_{\mathrm{score,Qcomp}} =\underbrace{2dd_{qc}}_{x_tW_{DQ}} +\underbrace{2hd_{qc}d_k}_{c_t^QW_{UQ}} +\underbrace{2hd_kd_c}_{q_tW_{UK}^{\top}} +\underbrace{2htd_c}_{q_{\mathrm{abs}}(C^{KV}_{\le t})^{\top}}.$$

Without Q compression, delete $W_{DQ},W_{UQ}$ and replace them with $W_Q$. The chain has fewer factors, so we leave the same association argument to the reader; the resulting minimum is

$$F_{\mathrm{score,noQ}} =\underbrace{2hdd_k}_{x_tW_Q} +\underbrace{2hd_kd_c}_{q_tW_{UK}^{\top}} +\underbrace{2htd_c}_{q_{\mathrm{abs}}(C^{KV}_{\le t})^{\top}}.$$

Value and output chain. The complete value-side chain is

$$\underset{(1,d)}{y_t} = \underset{(h,1,t)}{p_t} \underset{(t,d_c)}{C^{KV}_{\le t}} \underset{(h,d_c,d_v)}{W_{UV}} \underset{(h,d_v,d)}{W_o}.$$

First consider the coefficient of $t$. Contracting $p_t$ with $C^{KV}_{\le t}$ immediately removes the cache-length axis through the singleton query axis, at a cost of

$$2htd_c.$$

This is the smallest possible length-$t$ term. If $C^{KV}_{\le t}$ is expanded through $W_{UV}$ first, that step alone costs $2htd_cd_v$. Therefore, $p_tC^{KV}_{\le t}$ should be evaluated first, producing $o_c:(h,1,d_c)$.

It remains to associate the three-factor chain

$$\underset{(h,1,d_c)}{o_c} \underset{(h,d_c,d_v)}{W_{UV}} \underset{(h,d_v,d)}{W_o}.$$

Excluding the common first contraction $2htd_c$, the table includes the two online associations and the offline-precomposed variant; its one-time composition cost is omitted:

The first row remains the minimum. Again, the cheaper order keeps the length-$1$ axis outside and evaluates the chain from left to right. The complete value-side cost is

$$F_{\mathrm{value}} =\underbrace{2htd_c}_{p_tC^{KV}_{\le t}} +\underbrace{2hd_cd_v}_{o_cW_{UV}} +\underbrace{2hd_vd}_{oW_o}.$$

Combining the chosen associations, the complete cached-decode costs are

$$\begin{aligned} F_{\mathrm{MLA,decode,cached,noQ}} &=\underbrace{2dd_c+2hdd_k}_{\text{current-token projections}} +\underbrace{2hd_kd_c+4htd_c}_{\text{score absorption and latent reductions}} +\underbrace{2hd_cd_v+2hd_vd}_{\text{value and output projections}}, \\ F_{\mathrm{MLA,decode,cached,Qcomp}} &=\underbrace{2dd_c+2dd_{qc}+2hd_{qc}d_k}_{\text{current-token projections}} +\underbrace{2hd_kd_c+4htd_c}_{\text{score absorption and latent reductions}} +\underbrace{2hd_cd_v+2hd_vd}_{\text{value and output projections}}. \end{aligned}$$

---

Comparison.

In this RoPE-free comparison, MLA's decode FLOPs grow $4\times$ faster with $t$, while its cache is $64\times$ smaller than MHA's, a $98.4\%$ reduction in cached elements.

4.3 Why Compress Q Too?

The two MLA variants above give a direct ablation of Q compression. Since it factorizes only the query projection, its effects are concentrated on query-side activations and projection FLOPs rather than the cache path:

• Training activation memory. We do not know DeepSeek's internal implementation, but we think recomputation and kernel fusion can be used to reduce activation memory.
• Prefill and training FLOPs. The factorized projection is cheaper when $d_{qc}(d+hd_k)<hdd_k$, as the comparison in Section 4.1 shows.
• Decode. Only the current-token query projection changes; the latent cache and the length-$t$ attention term remain the same. Q compression therefore contributes little to decode compared with KV compression.

5. Making MLA Compatible with RoPE

5.1 Why RoPE and MLA Compression Are Incompatible

Takeaway. Without RoPE, $W_{UK}$ can be absorbed into the query, producing one latent query that is reused across all key positions. With RoPE, a key-position-dependent $R_u$ appears before $W_{UK}$, so the same absorbed query can no longer be reused across all keys.

Start with the complete RoPE $QK^\top$ expression. Fix a head $s$. Let $q_{s,t}^C:(1,d_k)$ be its content query at position $t$, let $c_u^{KV}:(1,d_c)$ be the KV latent at position $u$, and let $W_{UK,s}:(d_c,d_k)$. The reconstructed content key is $k_{s,u}^C=c_u^{KV}W_{UK,s}$. Under the row-vector RoPE convention from Section 3, the query and key are right-multiplied by $R_t^\top$ and $R_u^\top$, respectively. The score for this pair of positions is therefore

$$\begin{aligned} S^{\mathrm{RoPE}}_{s,t,u} &= \left(q_{s,t}^C R_t^\top\right) \left((c_u^{KV}W_{UK,s})R_u^\top\right)^\top \\ &= q_{s,t}^C R_t^\top R_u W_{UK,s}^\top (c_u^{KV})^\top. \end{aligned}$$

The only question is how to parenthesize this matrix chain. First remove RoPE and recall why ordinary MLA supports key absorption. Let $T$ be the number of query positions and $U$ the number of key positions. The complete score chain is

$$\underset{(h,T,U)}{S^C} = \underset{(h,T,d_k)}{Q^C} \underset{(h,d_k,d_c)}{W_{UK}^\top} \underset{(d_c,U)}{(C^{KV})^\top}.$$

One order reconstructs all keys from the latent before computing attention:

$$\underset{(h,T,d_k)}{Q^C} \left( \underset{(h,d_k,d_c)}{W_{UK}^\top} \underset{(d_c,U)}{(C^{KV})^\top} \right) = \underset{(h,T,d_k)}{Q^C} \underset{(h,d_k,U)}{(K^C)^\top}.$$

The other order absorbs $W_{UK}$ into the query first:

$$\left( \underset{(h,T,d_k)}{Q^C} \underset{(h,d_k,d_c)}{W_{UK}^\top} \right) \underset{(d_c,U)}{(C^{KV})^\top} = \underset{(h,T,d_c)}{\widehat Q^C} \underset{(d_c,U)}{(C^{KV})^\top}.$$

$W_{UK}$ is a fixed weight, so the absorbed query $\widehat Q^C=Q^CW_{UK}^\top$ is independent of the key position. Each query needs this $(h,T,d_c)$ intermediate only once, after which it can be multiplied by every cached latent. This is key absorption.

Now put RoPE back. The complete score can be evaluated by reconstructing and rotating the key first:

$$\left(q_{s,t}^C R_t^\top\right) \left( R_u \left(W_{UK,s}^\top(c_u^{KV})^\top\right) \right).$$

Or we can try to put every factor before the key latent into the left parenthesis, just as we did above:

$$\underbrace{\left(q_{s,t}^C R_t^\top R_u W_{UK,s}^\top\right)}_{\text{the part we would like to absorb into the query}} (c_u^{KV})^\top.$$

The difference is now visible in the chain itself: the parenthesis contains $R_u$. It changes with key position $u$, so it cannot produce one $\widehat Q^C$ that is shared by all keys. Expanding all query and key positions makes the shape of this intermediate not $(h,T,d_c)$ but

$$(h,T,U,d_c).$$

Here $U$ is the key-position axis: $U=T$ during prefill; for one decode step, $T=1$ and $U=t$. RoPE does not invalidate the matrix multiplication, but it turns the supposed “absorbed query” into a different object for every $(t,u)$ position pair. The key-length axis remains, defeating the purpose of absorption. If the cache stores only $C^{KV}$, decoding must either reconstruct and rotate the historical keys or explicitly form this $(h,1,t,d_c)$ intermediate.

Standard RoPE and position-independent key absorption therefore cannot occupy the same feature path. The elementwise RoPE formulation in Section 3 is exactly equivalent to the rotation-matrix form used here; rewriting the computation elementwise does not remove its dependence on key position $u$, so it cannot recover position-independent absorption. Full MLA separates the two in Section 5.2: an absorbable content branch and a dedicated positional branch for RoPE.

5.2 Full Formulation of MLA: Decoupled Rotary Position Embedding

Takeaway. Full MLA keeps the content channel unrotated and absorbable, while a separate low-dimensional RoPE channel carries position. The cache stores $C^{KV}$ and one shared RoPE key, and the two channel scores are added before softmax.

Section 5.1 showed that one path cannot simultaneously support standard RoPE and position-independent key absorption. Full MLA therefore stops asking the same representation to do both jobs and splits the feature space into two branches: an unrotated content branch that remains fully compressible and absorbable, and a positional branch that is free to apply RoPE without preserving absorption. In short, one branch is designed to compress cleanly, while the other is designed to rotate cleanly.

---

MLA with decoupled RoPE.

Let $d_r$ denote the per-head dimension of the decoupled RoPE branch, written as $d_h^R$ in the DeepSeek-V2 report. We use the same rotation tensor as Section 3, now over $d_r$:

$$\underset{(T,d_r,d_r)}{\mathcal R^{(r)}}, \qquad \mathcal R^{(r)}_{t,:,:}=(R_t^{(r)})^\top.$$

The query path is

$$\underset{(T,d_{qc})}{C^Q} = \underset{(T,d)}{X} \underset{(d,d_{qc})}{W_{DQ}},$$ $$\underset{(h,T,d_k)}{Q^C} = \underset{(T,d_{qc})}{C^Q} \underset{(h,d_{qc},d_k)}{W_{UQ}}, \qquad \underset{(h,T,d_r)}{\overline Q^R} = \underset{(T,d_{qc})}{C^Q} \underset{(h,d_{qc},d_r)}{W_{QR}},$$ $$\underset{(h,T,d_r)}{Q^R} = \underset{(h,T,d_r)}{\overline Q^R} \underset{(T,d_r,d_r)}{\mathcal R^{(r)}}.$$

The query-side RoPE operation can remain unchanged. The key side is different: if the RoPE key were stored in compressed form, every decode step would have to expand and rotate all $t$ historical keys again, creating substantial repeated computation. Full MLA therefore caches the materialized $K^R$, rather than a compressed positional-key latent.

The key-value path is

$$\underset{(T,d_c)}{C^{KV}} = \underset{(T,d)}{X} \underset{(d,d_c)}{W_{DKV}},$$ $$\underset{(h,T,d_k)}{K^C} = \underset{(T,d_c)}{C^{KV}} \underset{(h,d_c,d_k)}{W_{UK}}, \qquad \underset{(h,T,d_v)}{V} = \underset{(T,d_c)}{C^{KV}} \underset{(h,d_c,d_v)}{W_{UV}},$$ $$\underset{(T,d_r)}{\overline K^R} = \underset{(T,d)}{X} \underset{(d,d_r)}{W_{KR}}, \qquad \underset{(T,d_r)}{K^R} = \underset{(T,d_r)}{\overline K^R} \underset{(T,d_r,d_r)}{\mathcal R^{(r)}}.$$

Storing a separate positional key for every head would add $hd_r$ cache elements per token. Full MLA instead uses one $K^R:(T,d_r)$ shared by all $h$ query heads. The RoPE branch is therefore MQA-style: it keeps per-head positional queries $Q^R:(h,T,d_r)$, but uses a single shared positional key head. This keeps the additional cache at $d_r$ elements per token.

The full per-head query and key are the feature-axis concatenations

$$\underset{(h,T,d_k+d_r)}{Q} = \left[ \underset{(h,T,d_k)}{Q^C}; \underset{(h,T,d_r)}{Q^R} \right],$$ $$\underset{(h,T,d_k+d_r)}{K} = \left[ \underset{(h,T,d_k)}{K^C}; \underset{(T,d_r)}{K^R} \text{ broadcast over }h \right].$$

Why does the score become the sum of two terms? Concatenation partitions the last tensor dimension into two disjoint blocks. For head $s$, query position $t$, and key position $u$, the full dot product is

$$\begin{aligned} S_{s,t,u} &= \sum_{\ell=1}^{d_k+d_r}Q_{s,t,\ell}K_{s,u,\ell} \\ &= \sum_{a=1}^{d_k}Q^C_{s,t,a}K^C_{s,u,a} + \sum_{b=1}^{d_r}Q^R_{s,t,b}K^R_{u,b} \\ &=S^C_{s,t,u}+S^R_{s,t,u}. \end{aligned}$$

The same identity can be read as block tensor multiplication:

$$\underset{(h,T,T)}{S} = \left[ \begin{array}{cc} \underset{(h,T,d_k)}{Q^C} & \underset{(h,T,d_r)}{Q^R} \end{array} \right] \left[ \begin{array}{c} \underset{(h,d_k,T)}{(K^C)^\top} \\ \underset{(d_r,T)}{(K^R)^\top} \text{ broadcast over }h \end{array} \right] = \underset{(h,T,T)}{S^C} + \underset{(h,T,T)}{S^R}.$$

There are no cross terms such as $Q^C(K^R)^\top$ or $Q^R(K^C)^\top$: the content and RoPE coordinates occupy different feature blocks, and the dot product pairs coordinates only within the same block. Therefore, the two score channels can be computed separately:

$$\underset{(h,T,T)}{S^C} = \underset{(h,T,d_k)}{Q^C} \underset{(h,d_k,T)}{(K^C)^\top}, \qquad \underset{(h,T,T)}{S^R} = \underset{(h,T,d_r)}{Q^R} \underset{(d_r,T)}{(K^R)^\top},$$ $$\underset{(h,T,T)}{P} = \operatorname{softmax}\!\left( \frac{S^C+S^R}{\sqrt{d_k+d_r}} \right),$$ $$\underset{(h,T,d_v)}{O} = \underset{(h,T,T)}{P} \underset{(h,T,d_v)}{V}, \qquad \underset{(T,d)}{Y} = \underset{(h,T,d_v)}{O} \underset{(h,d_v,d)}{W_o}.$$

For cached decode, content absorption and the positional branch give

$$\underset{(h,1,t)}{\mathrm{score}} = \left( \underset{(h,1,d_k)}{q^C} \underset{(h,d_k,d_c)}{W_{UK}^\top} \right) \underset{(d_c,t)}{(C_{\le t}^{KV})^\top} + \underset{(h,1,d_r)}{q^R} \underset{(d_r,t)}{(K_{\le t}^R)^\top},$$ $$\underset{(h,1,d_v)}{o} = \left( \underset{(h,1,t)}{p} \underset{(t,d_c)}{C_{\le t}^{KV}} \right) \underset{(h,d_c,d_v)}{W_{UV}}.$$

Thus the content key and value remain behind $C^{KV}$, while only the shared rotated key $K^R$ is cached explicitly.

---

Cache supplement.

For DeepSeek-V2, $d_c=512$ and $d_r=64$. The full cache reduction relative to MHA is therefore $32768/576\approx56.9\times$, rather than the $64\times$ obtained from the RoPE-free latent alone.

---

FLOPs supplement.

This supplement follows the accounting in Section 4 and counts the tensor contractions introduced by the decoupled branch. It also counts the explicit addition $S^C+S^R$: one addition for every score element, or $hT^2$ FLOPs in prefill and $ht$ FLOPs in cached decode. The structured elementwise RoPE rotations themselves remain outside the comparison.

For cached decode, the coefficient of $t$ is now $4hd_c+2hd_r+h=278{,}656$ for DeepSeek-V2.

6. How MLA Changes MTP's Speculative-Decoding Gain

Takeaway. MHA speculative verification reuses KV-cache data and can hide the added query computation behind the memory-bound read. MLA reduces that read and moves attention toward the compute-bound regime, increasing the marginal cost of each additional query.

The DeepSeek-V3 Technical Report trains MTP as an auxiliary objective. Its actual configuration uses one MTP module ($D=1$), which predicts one additional token. The module can be discarded for ordinary inference or used to draft a token for speculative decoding. On DeepSeek-V3, the report gives an $85\%$-$90\%$ acceptance rate for this second token and reports $1.8\times$ TPS with speculative decoding.

In MHA, several verification queries reuse the same KV-cache data. Because decode is strongly memory-bound, much of their additional attention computation can be hidden behind that cache read: memory bandwidth is already the bottleneck, while the compute units still have room. MLA changes this balance. It replaces the full per-head KV cache with $C^{KV}$ and the decoupled RoPE key, while absorption adds contractions around $C^{KV}$. The cache read becomes smaller and attention moves toward the compute-bound regime, leaving less room to hide the additional computation. This is why MLA weakens part of MTP's speculative-decoding gain.

Contributions and Acknowledgements

This note was written under NonLinear-1. Place holder contributed to the derivations and writing. We thank Place holder for discussions and feedback.

---

How to cite this article.

If you would like to cite this article, please use:

NonLinear-1. (Jul. 8, 2026). “MLA, dim by dim” [Blog post]. Retrieved from https://nonlinear1.com/en/posts/mla-dim-by-dim

@online{nonlinear1-mla-dim-by-dim,
        title={MLA, dim by dim},
        author={{NonLinear-1}},
        year={2026},
        month={July},
        url={\url{https://nonlinear1.com/en/posts/mla-dim-by-dim}},
}