DeepSeek V3 Multi-head Latent Attention (MLA) — pydata: Huiming's learning notes

MLA plot

在上一篇blog中，我们讨论了DeepSeek-V3 采用了 MLA (Multi-Head Latent Attention) 架构，其核心目的是通过低秩联合压缩 Key 和 Value 来减少 KV 缓存。特别是在推理的时候，通过低秩压缩，从而在generation（inference）的时候KV catch可以减少内存使用，增加速度。其架构设计的核心目标是用推理时的 KV 缓存节省，来换取训练和推理时投影矩阵（降维/升维）带来的额外计算。

虽然在配置文件中查询头 (Query Heads, num_attention_heads) 键值头 (Key/Value Heads, num_key_value_heads) $n_h = 128$ 且 $n_{kv} = 128$（参数上是 MHA），但在其DeepSeek MLA 的内部实现中，它采用了 MQA 的思想，即每个 Key/Value 的latent vector $\mathbf{c}_t^{KV}$ 会在内核级别共享给所有查询头，从而实现 KV 缓存的显著减少。

因此，虽然从配置文件参数上看 $n_h = n_{kv}=128$，但从实际的内存优化策略上看，该模型采用了类似 MQA/GQA 的思想来处理 Key/Value 缓存。

下面我们根据原始的paper，仔细分解MLA的结构。首先，我们引入一些向量维度的记号：

记号

$$ \begin{aligned} & d = \text{embedding的长度 = 7168} \\ & n_h = \text{计算attention的head的个数 = num_attention_heads = 128} \text{，也就是 num of Q 的个数} \\ & n_{kv} = \text{计算attention的 K,V 的head的个数 = num_kv_heads = 128} \\[1ex] & d_c = \text{压缩向量} \mathbf{c}_t^{KV} \text{的长度 = kv_lora_rank = 512} \\ & d_c^{\prime} = \text{压缩向量} \mathbf{c}_t^{Q} \text{的长度 = q_lora_rank = 1536} \\[1ex] & d_h = \text{计算attention的head的NoPE的长度 = qk_nope_head_dim = 128} \\ & d_h^R = \text{计算attention的head的RoPE的长度 = qk_rope_head_dim} = 64 \end{aligned} $$

1. 低秩压缩

MLA 引入了低秩压缩并且缓存这个KV 压缩来减少额外的投影矩阵带来的运算。

将 Key/Value 通过投影矩阵 $\mathbf{W}^{DKV}$ 投影到一个低维latent vector space $\mathbf{C}_t^{KV} \in \mathbb{R}^{d_c}$, $(d_c = 512 \ll d_h \cdot n_h = 128 \cdot 128)$，并缓存它们。$d_c = \text{kv_lora_rank} = 512$;
在Forward的时候，通过上投影矩阵恢复 $\mathbf{k}_t^C$ 和 $\mathbf{v}_t^C$。

压缩

$$ \mathbf{c}_t^{KV} = \mathbf{W}^{DKV} h_t, \qquad \mathbf{c}_t^Q = \mathbf{W}^{DQ} h_t $$

解压

$$ \begin{aligned} & [\mathbf{k}_{t,1}^{C}, \mathbf{k}_{t,2}^{C}, ..., \mathbf{k}_{t,n_h}^{C}] = \mathbf{k}_t^C = \mathbf{W}^{UK} c_t^{KV}, \\[1ex] & [\mathbf{v}_{t,1}^{C}, \mathbf{v}_{t,2}^{C}, ..., \mathbf{v}_{t,n_h}^{C}] = \mathbf{v}_t^C = \mathbf{W}^{UV} c_t^{KV}, \\[1ex] & [\mathbf{q}_{t,1}^{C}, \mathbf{q}_{t,2}^{C}, ..., \mathbf{q}_{t,n_h}^{C}] = \mathbf{q}_t^C = \mathbf{W}^{UQ} c_t^Q \end{aligned} $$

其中

$\mathbf{W}^{DKV} \in \mathbb{R}^{d_c \times d} = \mathbb{R}^{512 \times 7168}$ 是降维投影矩阵。
$\mathbf{W}^{DQ} \in \mathbb{R}^{d_c^{\prime} \times d} = \mathbb{R}^{ 1536 \times 7168}$ 是降维投影矩阵。
$\mathbf{W}^{UK} \in \mathbb{R}^{d_{h_k} \cdot n_h \times d_c}$
- Paper中的设置，$d_{h_k} = d_{h} = 128$
- $\mathbf{W}^{UK}$ 是 K 的 NoPE 部分的上投影，它的作用是将 $d_c$ 维度的向量投影到 K 的 NoPE 部分 ($d_{h}$维度乘以 $n_h$个head)。
- 对 K（这里只有NoPE部分），$d_{h_k}$ = $\text{qk_nope_head_dim}$ = $128$ (每头维度)
- $\mathbf{W}^{UK} \in \mathbb{R}^{d_{h_k} \cdot n_h \times d_c} = \mathbb{R}^{128 \cdot 128 \times 512}$是每个head的 K 的升维矩阵。
$\mathbf{W}^{UV} \in \mathbb{R}^{d_{h_v} \cdot n_h \times d_c}$
- Paper中的设置，$d_{h_v} = d_{h} = 128$
- $\mathbf{W}^{UV}$ 是 V 的 NoPE 部分上投影。同样是从 $d_c$ 维度的向量投影到 V 的 NoPE 部分 ($d_{h}$维度乘以 $n_h$个head)。
- $\mathbf{W}^{UV} \in \mathbb{R}^{d_{h_v} \cdot n_h \times d_c} = \mathbb{R}^{128 \cdot 128 \times 512}$是每个head的 V 的升维矩阵。
如果总的加起来的话，$\mathbf{W}^{UVK} \text{ 的维度是 } = [(128+128)*128, 512] = (32768, 512)$
$\mathbf{W}^{UQ} \in \mathbb{R}^{d_{h_q} \cdot n_h \times d_c^{\prime}}$
- Paper中的设置，$d_{h_q} = d_{h} = 128$
- $\mathbf{W}^{UQ}$ 是 Q 的上投影，每个 Q NoPE Head 维度 $d_{h_q}$ = $\text{qk_nope_head_dim}$ = $128$.
- $\mathbf{W}^{UQ} \in \mathbb{R}^{d_{h_q} \cdot n_h \times d_c^{\prime}} = \mathbb{R}^{128 \cdot 128 \times 1536}$

Notes：因为KQ要做内积，所以$d_{h_k} == d_{h_q}$，但是不一定非要等于 $d_{h_v}$。尽管在本文中它们都相等，在好多paper里面它们也都相等。但理论上他们可以不等。

从这儿可以看到，与标准的multihead attention（MHA）直接用kqv计算相比，这里面增加了降维和升维的过程，也就是增加了额外的运算。

2. Training时的计算过程

DeepSeek-V3 的 KV 压缩机制 $\left(\mathbf{h} \text{ to } \mathbf{c}^{KV} \text{ to } \mathbf{k}, \mathbf{v} \right)$ 在training的时候总 FLOPs 其实是增加了（多了两个投影 $\mathbf{W}^{DKV}$ 和 $\mathbf{W}^{UK}, \mathbf{W}^{UV}$），但它在推理时大幅减少了内存带宽需求和 KV 缓存大小，从而实现了更快的生成速度和更长的上下文窗口。

恢复 K V:

$$ \mathbf{k}_t^{C} = \mathbf{W}^{UK} \mathbf{c}_t^{KV}, \qquad \mathbf{v}_t^{C} = \mathbf{W}^{UV} \mathbf{c}_t^{KV} $$

恢复 Q:

$$ \mathbf{q}_t^{C} = \mathbf{W}^{UQ} \mathbf{c}_t^{Q} $$

Attention

$$ \mathbf{o}_t = [\mathbf{o}_{t,1};\ \mathbf{o}_{t,2};\ \ldots;\ \mathbf{o}_{t,n_h}] $$

对每个head，

$$ \mathbf{o}_{t,i} = \sum_{j=1}^{t} \mathrm{Softmax} \left( \frac{\left(\mathbf{q}_{t,i}^{C} \right)^\top \mathbf{k}_{j,i}^{C}}{\sqrt{d_h}} \right) \mathbf{v}_{j,i}^{C} $$

得到经过 K Q weighted average 过的 V 的值。

其中

$$ \begin{aligned} & \mathbf{q}_{t,i}^C = \mathbf{W}_i^{UQ} \mathbf{c}_{t,i}^Q \in \mathbb{R}^{d_h} = \mathbb{R}^{128}, \qquad \mathbf{W}_i^{UQ} \in \mathbb{R}^{d_h \times d_c^{\prime}} = \mathbb{R}^{128 \times 1536} \\[1ex] & \mathbf{k}_{t,i}^C = \mathbf{W}_i^{UK} \mathbf{c}_{t,i}^{KV} \in \mathbb{R}^{d_h} = \mathbb{R}^{128} , \qquad \mathbf{W}_i^{UK} \in \mathbb{R}^{{d}_h \times {d}_c} = \mathbb{R}^{128 \times 512} \\[1ex] & \mathbf{v}_{t,i}^C = \mathbf{W}_i^{UV} \mathbf{c}_{t,i}^{KV} \in \mathbb{R}^{d_h} = \mathbb{R}^{128}, \qquad \mathbf{W}_i^{UV} \in \mathbb{R}^{{d}_h \times {d}_c} = \mathbb{R}^{128 \times 512} \\[1ex] & \mathbf{c}_{t,i}^{KV} = \mathbf{W}_i^{DKV} \mathbf{h}_{t,i} \in \mathbb{R}^{{d}_c} = \mathbb{R}^{512}, \qquad \mathbf{c}_{t,i}^{Q} = \mathbf{W}_i^{DQ} \mathbf{h}_{t,i} \in \mathbb{R}^{{d}_c^{\prime}} = \mathbb{R}^{1536} \\[1ex] & \mathbf{W}_i^{DKV} \in \mathbb{R}^{d_c \times d} = \mathbb{R}^{512 \times 7168}, \qquad \mathbf{W}_i^{DQ} \in \mathbb{R}^{d_c^{\prime} \times d} = \mathbb{R}^{ 1536 \times 7168} \end{aligned} $$

3. Inference时的计算过程

在推理阶段（自回归生成），有前面$n$个token生成第$n+1$个token，每次只处理一个新 token，但是在计算Attention的时候，需要与前面所有的token做内积，这个时候如果每次都从前面token的embedding来计算，会增加很多的计算量。实际当中为了减少计算和内存的消耗，是使用的之前的所有 token 的 KV cache。KV catch的简单原理参见KV catch

压缩

对 token $t$，计算：

$$ \mathbf{c}_t^{KV} = \mathbf{W}^{DKV} \mathbf{h}_t $$

把 $\mathbf{c}_t^{KV}$ 存入缓存。这里面会极大的节省内存： 1. $\mathbf{c}_t^{KV}$ 被所有的 head（在deepseek-v3 是128个head）共享。所以对token $t$，只要cache一个512维的向量就可以了 2. $\mathbf{c}_t^{KV}$ 是$d_c=512$维度，而如果向原始的MHA那样保存 key $\mathbf{k_t}$ 和value $\mathbf{v_t}$，假如它们都和hidden vector或者embedding vector 一样的size (i.e, $d_h = 7168$)，那么一个head一个token就差 7168/512倍，对$n_h=128$个head，上千个token或更多 (ds-v3 最多支持上下文的token个数 $\text{max_position_embeddings}$ = $163840$)，节省的内存会更多.

解压

对所有的历史 token $j \le t$ ，因为它们的压缩矩阵 $\mathbf{c}_j^{KV}$ 都被缓存了。所以，直接从缓存的 $\mathbf{c}_j^{KV}$ 恢复：

$$ \mathbf{k}_j^{C} = \mathbf{W}^{UK} \mathbf{c}_j^{KV}, \qquad \mathbf{v}_j^{C} = \mathbf{W}^{UV} \mathbf{c}_j^{KV} $$

对当前的查询 Q，直接从 token $t$ 的embedding vector计算（降维升维）：

$$ \mathbf{q}_t^{C} = \mathbf{W}^{UQ} \mathbf{c}_t^{Q} = \mathbf{W}^{UQ} \mathbf{W}^{DQ} \mathbf{h}_t $$

Attention

$$ \mathbf{o}_{t,i} = \sum_{j=1}^{t} \mathrm{Softmax}\left( \frac{\left(\mathbf{q}_{t,i}^{C}\right)^\top \mathbf{k}_{j,i}^{C}}{\sqrt{d_h}} \right) \mathbf{v}_{j,i}^{C} $$

得到经过 K Q weighted average 过的 V 的值。

到目前为止，所有的计算都有一个问题：打乱他们的顺序，$\mathbf{K^TQ}V$ 值不变。正如 MHA 所做的，MLA 也需要引入位置参数。

4. RoPE

背景知识：传统的transformer模型中，RoPE 的核心思想是将绝对位置信息编码到 Q (Query) 和 K (Key) 向量的点积中，而不需要像最初的绝对位置编码那样将位置向量直接加到词嵌入上。在 Transformer 模型中，由于自注意力机制是排列不变性的（即打乱词语顺序不影响计算结果），所以必须引入位置编码来告诉模型每个词语在序列中的相对或绝对位置。RoPE 通过将位置信息编码为旋转矩阵的形式，作用于 Query 和 Key 向量。它必须满足以下几个条件： 1. 对于序列中的第 $m$ 个词，其对应的 Query 向量 $\mathbf{q}_m = \mathbf{W}_Q \mathbf{h}_m$ 和第$n$个词 Key 向量 $\mathbf{k}_n=\mathbf{W}_K \mathbf{h}_n$，它们同维度而且同基($\mathbf{h}_i$) 2. 它们会被一个旋转矩阵作用得到 $\mathbf{q}_m' = \mathbf{R}_m(\mathbf{q}_m)$ 和 $\mathbf{k}_n' = \mathbf{R}_n(\mathbf{k}_n)$ 用来做attention的计算。 3. 注意力 Attention = $\left<\mathbf{q}_m'^{T}, \mathbf{k}_n'\right>$ = $\left< (\mathbf{R}_m(\mathbf{q}_m)^T, \mathbf{R}_n(\mathbf{k}_n) \right>$ = $\left< \mathbf{q}_m^T, \mathbf{R}_{m - n} (\mathbf{k}_n) \right>$. 即只跟两个token的距离 $m-n$ 有关，而跟绝对位置 $m$ $n$本身无关。

4.1. RoPE 与 KV 压缩的冲突

但是在MLA里面，下面几个原因导致 RoPE 不能直接作用在 K V 向量上，导致原来的RoPE性质不再成立。 1. MLA里面，$\mathbf{q}$ 和 $\mathbf{k}$ 不是来自同样的基。它们来自不同长度的压缩矩阵$\mathbf{c}_t^{Q}$, $\mathbf{c}_t^{KV}$ 2. 升维不是等距嵌入。 $\left< \mathbf{q}_m, \mathbf{k}_n \right>$ = $\left< \mathbf{q}_m^C, \mathbf{k}_n^C\right>$ = $\left<\mathbf{W}^{UQ} \mathbf{c}_m^{Q}, \mathbf{W}^{UK} \mathbf{c}_n^{KV} \right>$. 由于$\mathbf{W}^{UQ}$和$\mathbf{W}^{UK}$ 不相互正交，所以前面的内积不等于$\left<\mathbf{c}_m^{Q}, \mathbf{c}_n^{KV} \right>$, 它们也没法做内积 3. 如果在升维以后直接强制做RoPE， $\left< R_m(\mathbf{q}_m), R_n(\mathbf{k}_n) \right>$ = $\left< R_m(\mathbf{W}^{UQ} \mathbf{c}_m^{Q}), R_n(\mathbf{W}^{UK} \mathbf{c}_n^{KV}) \right>$ = $(\mathbf{c}_m^{Q})^T (\mathbf{W}^{UQ})^T R_{m-n}(\mathbf{W}^{UK} \mathbf{c}_n^{KV})$, 而 $(\mathbf{W}^{UQ})^T R_{m-n}(\mathbf{W}^{UK}) \neq R_{m-n}$.

这些原因导致在降维/升维的情况下，RoPE不再可行。为了解决这个问题，MLA采用了 RoPE 和 NoPE 分离的办法。RoPE单独的有一个自己的向量，然后与NoPE向量concat在一起。

解决办法：Decouple RoPE

1. Key 部分

$\mathbf{k}_{t}^R$ (RoPE 部分，注意这一部分跟head无关，所有的 head 都用同样的 $\mathbf{k}_{t}^R$ ) 1. $\mathbf{k}_{t}^R = \text{RoPE}\left(\mathbf{W}^{KR} \mathbf{h}_t \right)$ 2. 输入：矩阵 $\mathbf{W}^{KR}$ 作用于 $\mathbf{h}_t$ (Input Hidden $\mathbf{h}_t$)。 3. 作用： $\mathbf{W}^{KR}$ 是一个独立的投影矩阵（维度 $64 \times 7168$），专门用于将 $\mathbf{h}_t$ 投影到 $\mathbf{k_t}^R$ 的 64 维空间。

$\mathbf{k}_{t,i}^C$ (NoPE 部分，对head i) 1. 不应用 RoPE 的部分 $\mathbf{k}_{t,i}^C$ 才是来自低秩潜在向量 $\mathbf{c}_{t}^{KV}$ 的。 2. 路径： $\mathbf{h}_t \to \mathbf{W}^{DKV} \to \mathbf{c}_{t}^{KV} \to \mathbf{W}^{UK} \to \mathbf{k}_{t,i}^C$。

2. Query 部分

$\mathbf{q}_{t,i}^R$ (RoPE 部分) 1. $\mathbf{q}_{t,i}^R = \text{RoPE}(\mathbf{W}^{QR} \mathbf{c}_{t}^Q)$ 2. 输入：矩阵 $\mathbf{W}^{QR}$ 作用于低秩压缩向量$\mathbf{c}_{t}^Q$。 3. 作用： $\mathbf{W}^{QR}$ 是一个独立的投影矩阵（维度 $64 \times 1536$），用于将 $\mathbf{c}_{t}^Q$ 投影到 $\mathbf{q}^R$ 的 64 维空间。

$\mathbf{q}_{t,i}^C$ (NoPE 部分) 1. 不应用 RoPE 的部分 $\mathbf{q}_{t,i}^C$ 同样来自低秩潜在向量 $\mathbf{c}_{t}^{Q}$ 的。 2. 路径： $ \mathbf{c}_{t}^{Q} \to \mathbf{W}^{UQ} \to \mathbf{q}_{t,i}^C$。

4.2. MLA + RoPE 在训练时的计算

1. 压缩

压缩部分跟上面提到的没有 RoPE 一样，都是从embedding vector（hidden status）压缩到一个低维的空间。

$$ \mathbf{c}_t^{KV} = \mathbf{W}^{DKV} \mathbf{h}_t $$

$$ \mathbf{c}_t^Q = \mathbf{W}^{DQ} \mathbf{h}_t $$

2. 解压

从压缩空间 (latenct vector) 恢复 Key, Value, Query 向量:

$$ \mathbf{k}_t^C = \mathbf{W}^{UK} \mathbf{c}_t^{KV}, \\[1ex] \mathbf{v}_t^C = \mathbf{W}^{UV} \mathbf{c}_t^{KV} , \\[1ex] \mathbf{q}_t^C = \mathbf{W}^{UQ} \mathbf{c}_t^Q $$

3. RoPE 部分

生成一组新的 Key Query 向量，它们不是用来和上面的向量相加做RoPE，而是做完RoPE以后跟上面的向量 concat:

$$ \mathbf{k}_t^R = \text{RoPE}\left(\mathbf{W}^{KR} \mathbf{h}_t \right) \\[1ex] \mathbf{q}_t^R = \text{RoPE}\left(\mathbf{W}^{QR} \mathbf{c}_t^Q \right) $$

4. concat

把 K Q 的 RoPE 部分和 NoPE 拼接起来，新的向量维度为 $d_h + d_h^R = 128 + 64$.

$$ \mathbf{q}_{t, i} = \left[\mathbf{q}_{t, i}^C; \mathbf{q}_{t, i}^R\right], \quad \mathbf{k}_{t, i} = \left[\mathbf{k}_{t, i}^C; \mathbf{k}_{t}^R \right] $$

5. 计算注意力

第 $t$ 个 token 的 Q 跟所有的历史 K 做内积，对 head $i$

$$ \mathbf{o}_{t,i} = \sum_{j=1}^t \text{Softmax}\left(\frac{(\mathbf{q}_{t,i})^{\top} \mathbf{k}_{j,i}}{\sqrt{d_h^C + d_h^R}}\right) \mathbf{v}_{j,i}^C $$

对所有的 head

$$ \mathbf{o}_t = [\mathbf{o}_{t,1}, \mathbf{o}_{t,2}, \dots, \mathbf{o}_{t,n_h}] $$

各个向量和矩阵的维度

$$ \mathbf{q}_{t,i} = \left[\mathbf{W}^{UQ}_i \mathbf{c}_{t,i}^Q; \text{ RoPE}(\mathbf{W}^{QR}_i \mathbf{c}_{t,i}^Q) \right] \in \mathbb{R}^{d_h + d_h^R} = \mathbb{R}^{128 + 64}, $$

$$ \quad \mathbf{W}^{UQ}_i \in \mathbb{R}^{d_h \times d_c^{\prime}} = \mathbb{R}^{128 \times 1536} $$

$$ \mathbf{W}^{QR}_i \in \mathbb{R}^{d_h^R \times d_c^{\prime}} = \mathbb{R}^{64 \times 1536} $$

$$ \mathbf{k}_{t,i} = \left[\mathbf{W}^{UK}_i \mathbf{c}_{t,i}^{KV}; \text{RoPE}(\mathbf{W}^{KR}_i \mathbf{h}_t) \right] \in \mathbb{R}^{d_h + d_h^R} = \mathbb{R}^{128 + 64}, $$

$$ \quad \mathbf{W}^{UK}_i \in \mathbb{R}^{d_h \times d_c} = \mathbb{R}^{128 \times 512} $$

$$ \mathbf{W}^{KR}_i \in \mathbb{R}^{d_h^R \times d} = \mathbb{R}^{64 \times 7168} $$

$$ \mathbf{v}_{t,i}^C = \mathbf{W}^{UV}_i \mathbf{c}_{t,i}^{KV} \in \mathbb{R}^{d_h} = \mathbb{R}^{128}, \quad \mathbf{W}^{UV}_i \in \mathbb{R}^{d_h \times d_c} = \mathbb{R}^{128 \times 512} $$

$$ \mathbf{c}_{t,i}^{KV} = \mathbf{W}^{DKV}_i \mathbf{h}_{t,i} \in \mathbb{R}^{d_c} = \mathbb{R}^{512}, \quad \mathbf{c}_{t,i}^Q = \mathbf{W}^{DQ}_i \mathbf{h}_{t,i} \in \mathbb{R}^{d_c^{\prime}} = \mathbb{R}^{1536} $$

$$ \mathbf{W}^{DKV}_i \in \mathbb{R}^{d_c \times d} \mathbb{R}^{512 \times 7168}, \quad \mathbf{W}^{DQ}_i \in \mathbb{R}^{d_c^{\prime} \times d} \mathbb{R}^{1536 \times 7168} $$

4.3. MLA + RoPE 在推理时的计算

在自回归推理中，模型预测出 token $t$ 得到对应的 $\mathbf{h}_t$ 以后，token $t$ 会被压缩被缓存，加上 token $t$ 之前的所有 token 的缓存，然后利用缓存的 KV 进行注意力计算，预测下一个 token $t+1$。

1. 压缩并缓存

路径 $\text{token t} \to \mathbf{h}_t \to \mathbf{c}_t^{KV} \text{ and } \mathbf{k}_t^R$. 注意：对 $j < t$ 的情况，所有的 $\mathbf{c}_j^{KV}$ 都已经被Cache了。

从 token $t$ 的 embedding vector 到压缩向量:
$$ \mathbf{c}_t^{KV} = \mathbf{W}^{DKV} \mathbf{h}_t $$
从token $t$ 的 embedding vector 到新的 RoPE K:
$$ \mathbf{k}_t^R = \text{RoPE}(\mathbf{W}^{KR} \mathbf{h}_t) $$
Cache $\mathbf{c}_t^{KV}$ 和 $\mathbf{k}_t^R$ 。这样的话从 $j$ 从 $ 1 \to t$ 的 $\mathbf{c}_j^{KV}$ 和 $\mathbf{k}_j^R$ 都已经 Cache 了。

2. 解压

历史 Key/Value ( $j \le k$ ) 直接从缓存恢复:

$$ \mathbf{k}_j^C = \mathbf{W}^{UK} \mathbf{c}_j^{KV}, \quad \mathbf{v}_j^C = \mathbf{W}^{UV} \mathbf{c}_j^{KV}, \quad \forall j \le k $$

RoPE K 的向量 $\mathbf{k}_j^R$ $(\forall j \le k)$ 直接从缓存读取。
当前 $\text{token t}$ 的 Query:
$$ \mathbf{q}_t^C = \mathbf{W}^{UQ} \mathbf{W}^{DQ} \mathbf{h}_t, \quad \mathbf{q}_t^R = \text{RoPE}(\mathbf{W}^{QR} \mathbf{W}^{DQ} \mathbf{h}_t) $$

3. Concat

把 NoPE 和 RoPE 拼接起来：

$$ \mathbf{q}_{t,i} = [\mathbf{q}_{t,i}^C; \mathbf{q}_{t,i}^R], \quad \mathbf{k}_{j,i} = [\mathbf{k}_{j,i}^C; \mathbf{k}_{j,i}^R] $$

4. 注意力公式:

$$ \mathbf{o}_{t,i} = \sum_{j=1}^t \text{Softmax}\left(\frac{\mathbf{q}_{t,i}^{\top} \mathbf{k}_{j,i}}{\sqrt{d_h + d_h^R}}\right) \mathbf{v}_{j,i}^C $$

在MLA Inference里面，只需要缓存 $\mathbf{c}_t^{KV} \in \mathbb{R}^{512}$ 和 $\mathbf{k}_t^R \in \mathbb{R}^{64}$，并且他们会被所有的head共享，这样比Cache原始的$\mathbf{h}_t \in \mathbb{R}^{7168}$ 或者 cache K V 减少了很多内存。

DeepSeek-V3 的设计目标不是减少 FLOPs（浮点运算次数），而是减少内存带宽和显存占用，这在 LLM 的生成/推理阶段尤为关键。

最后总结一下 MLA 部分的超参：

组件	参数名	数值
注意力头数	$n_h$	128
每个头的维度	$d_h$	128
KV 压缩维度	$d_c$ (kv_lora_rank)	512
Query 压缩维度	$d_{c,q}$ (q_lora_rank)	1536
解耦 RoPE 维度	$d_h^R$ (rotary_dim)	64
KV 上采样维度	$d_h^C$ (nop_dim)	128

组件	参数名	数值
注意力头数	\(n_h\)	128
每个头的维度	\(d_h\)	128
KV 压缩维度	\(d_c\) (kv_lora_rank)	512
Query 压缩维度	\(d_{c,q}\) (q_lora_rank)	1536
解耦 RoPE 维度	\(d_h^R\) (rotary_dim)	64
KV 上采样维度	\(d_h^C\) (nop_dim)	128