Transformer KV缓存优化（MHA、MQA、GQA、MLA，参考：DeepSeek-V2）

1、说明

本文参考： 论文DeepSeek-V2
我们用我们提出的多头潜在注意力 （MLA） 和 DeepSeekMoE 优化了 Transformer 框架（Vaswani et al.， 2017）中的注意力模块和前馈网络 （FFN）。（1） 在注意力机制的背景下，多头注意力 （MHA） 的键值 （KV） 缓存（Vaswani et al.， 2017）对 LLM 的推理效率构成了重大障碍。已经探索了各种方法来解决这个问题，包括分组查询注意力 （GQA） （Ainslie et al.， 2023） 和多查询注意力 （MQA） （Shazeer， 2019）。但是，这些方法在尝试减少 KV 缓存时通常会牺牲性能。为了实现两全其美的效果，我们引入了 MLA，这是一种配备了低秩键值联合压缩的注意力机制。从经验上讲，MLA 实现了优于 MHA 的性能，同时显著降低了推理过程中的 KV 缓存，从而提高了推理效率。（2） 对于前馈网络 （FFN），我们遵循 DeepSeekMoE 架构 （Dai et al.， 2024），该架构采用细粒度的专家分割和共享的专家隔离，以提高专家专业化的潜力。与 GShard 等传统 MoE 架构相比，DeepSeekMoE 架构表现出巨大的优势（Lepikhin et al.， 2021），使我们能够以经济的成本训练强大的模型。由于我们在训练期间采用专家并行性，因此我们还设计了补充机制来控制通信开销并确保负载平衡。通过结合这两种技术，DeepSeek-V2 同时具有强大的性能（图 1（a））、经济的训练成本和高效的推理吞吐量（图 1（b））。

2、KV缓存优化

2.0 KV缓存（Cache）的核心原理

Transformer 的 KV Cache（Key-Value Cache） 是一种针对自回归生成任务（如文本生成）的推理优化技术，通过缓存中间计算结果减少重复运算，显著提升模型效率。以下从原理、作用、实现等角度详细解析：

---

一、KV Cache 的核心原理

1. 自注意力机制的计算冗余问题
   在 Transformer 的自注意力机制中，每个 token 需计算对应的 Query（Q）、Key（K）、Value（V）矩阵。生成新 token 时，若未缓存 K 和 V，模型需重新计算所有历史 token 的 K 和 V，导致计算复杂度为 O(n²)（n 为序列长度）。

2. KV Cache 的缓存机制
   • 首次推理：处理初始 token 时，完整计算并存储 K 和 V 矩阵。

• 后续推理：生成新 token 时，仅计算当前 token 的 Q，并从缓存中读取历史 K 和 V，避免重复计算。复杂度降至 O(n)。

• 动态更新：将新 token 的 K 和 V 追加至缓存，供后续步骤复用。

---

• 传统Transformer的缺陷
  在自注意力机制中，每个token的输出依赖于对所有历史token的注意力计算：

  1. 查询（Query, Q）、键（Key, K）、值（Value, V） 由输入通过线性层生成：
     $$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$$
  2. 注意力分数计算为：
     $$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$
     在生成任务（如逐个token生成文本）中，若每次生成新token都重新计算所有历史token的K和V，会导致计算量随序列长度呈平方增长（时间复杂度为 $O(n^2)$），效率极低。

• KV缓存的核心思想
  缓存历史层的K和V，避免重复计算：
  在推理阶段（生成任务），当处理第 $i$ 个token时，仅计算当前token的Q，而K和V则复用之前 $i-1$ 个token的缓存结果。具体步骤如下：

  1. 首次计算：处理第1个token时，生成其K₁和V₁，存入缓存。
  2. 后续生成：处理第 $n$ 个token时：
     • 计算当前token的Qₙ。
     • 从缓存中获取前 $n-1$ 个token的K₁₋ₙ₋₁和V₁₋ₙ₋₁，与当前Kₙ、Vₙ拼接后，用于计算注意力。
     • 将新的Kₙ和Vₙ添加到缓存中，供下一个token使用。

  通过这种方式，每次生成新token的时间复杂度降至 $O(n)$（仅新增一次线性计算和拼接操作），大幅提升长序列生成效率。

---

二、KV Cache 的适用场景与限制

1. 适用场景
   • 自回归模型：如 GPT、LLaMA 等生成式模型。

• 实时交互任务：聊天机器人、翻译等对延迟敏感的场景。

2. 限制
   • 因果性要求：仅适用于因果掩码（Causal Mask）模型，如 Decoder-only 架构。

• 显存瓶颈：长序列生成时缓存占用显存，需权衡速度与内存。

---

三、KV Cache 的核心作用

1. 加速推理
   • 通过避免重复计算 K 和 V，推理速度提升 2-3 倍。例如，某聊天机器人的响应时间从 0.5 秒降至 0.2 秒。

• 实验表明，开启 KV Cache 后，生成 1000 个 token 的耗时从指数级增长（如 332 秒）变为线性增长（29 秒）。

2. 降低显存占用
   • 复用缓存减少中间张量的存储需求。例如，某移动设备上的内存占用从 1GB 降至 0.6GB。

• 尽管缓存占用显存（如 Llama3-70B 在 4k token 时需约 10.5GB），但整体效率仍显著优化。

3. 支持长序列生成
   • 长文本生成时，KV Cache 维持对历史上下文的完整理解，确保语义连贯性。

---

2.1 KV缓存优化

我们用我们提出的多头潜在注意力 （MLA） 和 DeepSeekMoE 优化了 Transformer 框架（Vaswani et al.， 2017）中的注意力模块和前馈网络 （FFN）。（1） 在注意力机制的背景下，多头注意力 （MHA） 的键值 （KV） 缓存（Vaswani et al.， 2017）对 LLM 的推理效率构成了重大障碍。已经探索了各种方法来解决这个问题，包括分组查询注意力 （GQA） （Ainslie et al.， 2023） 和多查询注意力 （MQA） （Shazeer， 2019）。但是，这些方法在尝试减少 KV 缓存时通常会牺牲性能。为了实现两全其美的效果，我们引入了 MLA，这是一种配备了低秩键值联合压缩的注意力机制。从经验上讲，MLA 实现了优于 MHA 的性能，同时显著降低了推理过程中的 KV 缓存，从而提高了推理效率。（2） 对于前馈网络 （FFN），我们遵循 DeepSeekMoE 架构 （Dai et al.， 2024），该架构采用细粒度的专家分割和共享的专家隔离，以提高专家专业化的潜力。与 GShard 等传统 MoE 架构相比，DeepSeekMoE 架构表现出巨大的优势（Lepikhin et al.， 2021），使我们能够以经济的成本训练强大的模型。由于我们在训练期间采用专家并行性，因此我们还设计了补充机制来控制通信开销并确保负载平衡。通过结合这两种技术，DeepSeek-V2 同时具有强大的性能（图 1（a））、经济的训练成本和高效的推理吞吐量（图 1（b））。

2.2 性能对比

2.1 多头潜在注意力：提升推理效率
传统Transformer模型通常采用多头注意力（MHA，Vaswani等人，2017），但在生成过程中，其庞大的键值（KV）缓存会成为限制推理效率的瓶颈。为减少KV缓存，研究者提出了多查询注意力（MQA，Shazeer，2019）和分组查询注意力（GQA，Ainslie等人，2023）。这些方法需要的KV缓存规模较小，但性能不及MHA（附录D.1提供了MHA、GQA和MQA的消融实验结果）。
针对DeepSeek-V2，我们设计了一种创新的注意力机制——多头潜在注意力（MLA）。通过低秩键值联合压缩，MLA在实现比MHA更强性能的同时，仅需显著更少的KV缓存。以下将介绍其架构，并在附录D.2中提供MLA与MHA的对比。

我们在表 1 中演示了不同注意力机制之间每个令牌的 KV 缓存的比较。MLA 只需要少量的 KV 缓存，相当于只有 2.25 组的 GQA，但可以达到比 MHA 更强的性能。

表 1 不同注意力机制之间每个令牌的 KV 缓存比较。$n_h$ 表示关注头的数量，$d_h$ 表示每个关注头的维度，l 表示层数，$n_g$ 表示 GQA 中的组数，$d_c$ 和 $d_h^R$ 分别表示 MLA 中解耦查询和 key 的 KV 压缩维度和每头维度。KV 缓存量由元素数量来衡量，与存储精度无关。对于 DeepSeek-V2,$d_c$ 设置为 $4d_h$，$d_h^R$ 设置为 $\frac{d_h}{2}$。因此，它的 KV 缓存等于只有 2.25 组的 GQA，但其性能比 MHA 强。

2.3 架构

图 3 |多头注意力 （MHA）、分组查询注意力 （GQA）、多查询注意力 （MQA） 和多头潜在注意力 （MLA） 的简化图示。MLA 通过将 key 和 values 联合压缩成一个 latent vector，显著降低了推理过程中的 KV 缓存。

2.4多头注意力 （MHA）

预备知识：标准多头注意力
首先介绍标准MHA机制作为背景。设d为嵌入维度，$n_h$为注意力头数， $d_h$为每头维度， $h_t\in {\mathbb{R}}^d$为第t个token在注意力层的输入。标准MHA通过三个矩阵W^Q、WK、$W^V\in {\mathbb{R}}^{d_h{n}_h\times d}$分别生成qt、kt、 $v_t\in {\mathbb{R}}^{d_h{n}_h}$ ：

$n_h$ be the number of attention heads, $d_h$ be the dimension
per head, and $h_t\in {\mathbb{R}}^d$be the attention input of the
𝑡-th token at an attention layer. Standard MHA first produces $q_t$
,kt, $v_t\in {\mathbb{R}}^{d_h{n}_h}$ through three matrices
$w^Q$ , $W^K$ , $W^V\in {\mathbb{R}}^{d_h{n}_h\times d}$ ,
respectively:

$$q_t=W^Q{h}_t,(1)$$

$$k_t=W^K{h}_t,(2)$$

$$v_t=W^V{h}_t,(3)$$

随后，qt、kt、vt被切分为nh个头以进行多头注意力计算：
$$[q_{t,1};q_{t,2};\dots ;q_{t,n_h}]=q_t,(4)$$

$$[k_{t,1};k_{t,2};\dots ;k_{t,n_h}]=k_t,(5)$$

$$[v_{t,1};v_{t,2};\dots ;v_{t,n_h}]=v_t,(6)$$

$$o_{t,i}=\sum _{j=1}^t{\text{Softmax}}_j\left(\frac{q_{t,i}^T{k}_{j,i}}{\sqrt{d_h}}\right)v_{j,i},(7)$$

$$u_t=W^O[o_{t,1};o_{t,2};\dots ;o_{t,n_h}],(8)$$
其中q𝑡,𝑖, k𝑡,𝑖, $v_{t,i}\in {\mathbb{R}}^{d_k}$分别表示第 i 个注意力头的查询、键和值;$W^O\in {\mathbb{R}}^{d\times d_h{n}_h}$ 为输出投影矩阵。推理时，所有键和值需缓存以加速计算，因此MHA每个token需缓存 $2n_h{d}_{h}l$个元素（$l$ 表示层数）。在模型部署中，这种庞大的KV缓存是限制最大批量大小和序列长度的主要瓶颈。

注意： KV，不是超参数$W^K$ , $W^V\in {\mathbb{R}}^{d_h{n}_h\times d}$ ,而是相乘后的KV矩阵：$k_t=W^K{h}_t,(2)、v_t=W^V{h}_t,(3)$

where q𝑡,𝑖, k𝑡,𝑖, $v_{t,i}\in {\mathbb{R}}^{d_k}$ denote the
query, key, and value of the 𝑖-th attention head, respectively;
$W^O\in {\mathbb{R}}^{d\times d_h{n}_h}$ denotes the output projection matrix. During inference, all keys and values need,
to be cached to accelerate inference, so MHA needs to cache $2n_h{d}_{h}l$ elements for each token. In model deployment, this heavy KV cache is a large bottleneck that limits the maximum batch size and sequence length.

2.5 多头潜在注意力 （MLA）

2.5.1 低秩键值联合压缩 （Low-Rank Key-Value ）

MLA的核心是通过键值的低秩联合压缩来减少KV缓存：
$$c_t^{KV}=W^{DKV}{h}_t,(9)$$

$$k_t^C=W^{UK}{c}_t^{KV},(10)$$

$$v_t^C=W^{UV}{c}_t^{KV},(11)$$
其中，$c_t^{KV}\in {\mathbb{R}}^{d_c}$ 是键和值的压缩潜向量；$d_c(\ll d_h{n}_h)$ 表示KV压缩维度；$W^{DKV}\in {\mathbb{R}}^{d_c\times d}$ 是下投影矩阵；$W^{UK}、W^{UV}\in {\mathbb{R}}^{d_h{n}_h\times d_c}$ 分别是键和值的上投影矩阵。推理时，MLA仅需缓存 $c_t^{KV}$，因此其KV缓存仅包含 $d_{c}l$ 个元素（$l$ 表示层数）。此外，由于推理时 $W^{UK}$ 可融入 $W^Q$，$W^{UV}$ 可融入 $W^O$，我们甚至无需显式计算注意力的键和值。图3直观展示了MLA通过键值联合压缩减少KV缓存的过程。

为减少训练时的激活内存，我们还对查询进行低秩压缩（尽管这不会减少KV缓存）：
$$c_t^Q=W^{DQ}{h}_t,(12)$$

$$q_t^C=W^{UQ}{c}_t^Q,(13)$$
其中，$c_t^Q\in {\mathbb{R}}^{d_c^{\prime }}$ 是查询的压缩潜向量；$d_c^{\prime }(\ll d_h{n}_h)$ 表示查询压缩维度；$W^{DQ}\in {\mathbb{R}}^{d_c^{\prime }\times d}、W^{UQ}\in {\mathbb{R}}^{d_h{n}_h\times d_c^{\prime }}$ 分别是查询的下投影和上投影矩阵。

解释：“$W^{UK}$ 可融入 $W^Q$，$W^{UV}$ 可融入 $W^O$”

在MLA的低秩键值联合压缩中，“$W^{UK}$ 可融入 $W^Q$，$W^{UV}$ 可融入 $W^O$”是指通过矩阵运算的结合律，将键/值的上投影矩阵与注意力模块的输入/输出投影矩阵合并，从而在推理时跳过显式计算键和值的步骤。以下结合公式和实例具体解释：

1. 矩阵吸收的数学原理
   假设：

   • $W^Q$ 是标准MHA中生成查询的投影矩阵，形状为 ${\mathbb{R}}^{d_h{n}_h\times d}$
   • $W^{UK}$ 是键的上投影矩阵，形状为 ${\mathbb{R}}^{d_h{n}_h\times d_c}$
   • $W^{DKV}$ 是键值的下投影矩阵，形状为 ${\mathbb{R}}^{d_c\times d}$

   根据公式(9)(10)，键的生成过程为：
   $$k_t^C=W^{UK}{c}_t^{KV}=W^{UK}(W^{DKV}{h}_t)=(W^{UK}{W}^{DKV})h_t$$
   这里，$W^{UK}$ 和 $W^{DKV}$ 可合并为一个新的矩阵 $W^{\prime Q}=W^{UK}{W}^{DKV}$，使得：
   $$k_t^C=W^{\prime Q}{h}_t$$
   此时，生成键的过程被转化为一次矩阵乘法，等价于将 $W^{UK}$ 的操作“吸收”到 $W^Q$ 中（实际是与 $W^Q$ 合并）。同理，值的生成矩阵 $W^{UV}$ 可与输出投影矩阵 $W^O$ 合并，避免显式计算 $v_t^C$。

2. 举例说明（简化维度）
   假设：

   • $d=512$（输入维度），$d_c=128$（KV压缩维度），$d_h{n}_h=1024$（单头维度×头数）
   • $W^{DKV}$ 是 $128\times 512$ 的下投影矩阵，$W^{UK}$ 是 $1024\times 128$ 的键上投影矩阵

   标准MHA流程（需显式计算键）：

   1. $c_t^{KV}=W^{DKV}{h}_t$（128维向量）
   2. $k_t^C=W^{UK}{c}_t^{KV}$（1024维键，需存储）

   MLA优化后（吸收矩阵，跳过键计算）：

   • 合并矩阵：$W^{\prime Q}=W^{UK}{W}^{DKV}$（形状为 $1024\times 512$）
   • 直接计算：$k_t^C=W^{\prime Q}{h}_t$（无需单独存储 $W^{UK}$ 和 $c_t^{KV}$，推理时仅需缓存 $c_t^{KV}$）

---

数学推导与示例
1. 键（Key）的融合
• 原始计算：
$$k_t^C=W^{UK}\cdot c_t^{KV}=W^{UK}\cdot (W^{DKV}\cdot h_t)$$
• 合并权重：

将$W^{UK}\cdot W^{DKV}$合并为新的查询权重矩阵$W^{Q^{\prime }}$：
$$W^{Q^{\prime }}=W^Q\cdot (W^{UK}\cdot W^{DKV})$$
此时，查询计算简化为：
$$q_t^{\prime }=W^{Q^{\prime }}\cdot h_t$$
实际计算注意力时，直接使用$q_t^{\prime }$，无需单独存储$k_t^C$。

2. 值（Value）的融合
• 原始计算：
$$v_t^C=W^{UV}\cdot c_t^{KV}=W^{UV}\cdot (W^{DKV}\cdot h_t)$$
• 合并到输出矩阵：

将$W^{UV}\cdot W^{DKV}$合并到输出矩阵$W^O$中：
$$W^{O^{\prime }}=W^O\cdot (W^{UV}\cdot W^{DKV})$$
最终输出计算为：
$$u_t=W^{O^{\prime }}\cdot \text{Attention}(q_t^{\prime },k_t^{\prime },v_t^{\prime })$$

---

实例说明
假设模型参数为：
• 输入维度$d=512$，压缩维度$d_c=128$，头数$n_h=8$，每头维度$d_h=64$。

• $W^{DKV}\in {\mathbb{R}}^{128\times 512}$，$W^{UK}\in {\mathbb{R}}^{512\times 128}$，$W^Q\in {\mathbb{R}}^{512\times 512}$。

推理优化步骤：

1. 合并键的投影：
   $$W^{Q^{\prime }}=W^Q\cdot (W^{UK}\cdot W^{DKV})\in {\mathbb{R}}^{512\times 512}$$
   实际计算查询时：
   $$q_t^{\prime }=W^{Q^{\prime }}\cdot h_t$$
   原始需要存储$k_t^C\in {\mathbb{R}}^{512}$，优化后直接通过$q_t^{\prime }$隐式包含键信息。

2. 合并值的投影：
   $$W^{O^{\prime }}=W^O\cdot (W^{UV}\cdot W^{DKV})\in {\mathbb{R}}^{512\times 512}$$
   输出计算直接使用合并后的$W^{O^{\prime }}$，无需单独计算值向量$v_t^C$。

---

2.5.2. 解耦旋转位置嵌入（Dcoupled RoPE ）

一、核心问题：传统RoPE与低秩KV压缩的冲突

DeepSeek-V2在设计MLA时遇到一个关键矛盾：

1. RoPE的位置敏感性：RoPE通过旋转矩阵对键（K）和查询（Q）添加位置信息，要求键和查询在计算时携带与位置相关的参数（如不同位置的旋转矩阵 ${\mathbf{R}}_t$）。
2. 矩阵吸收失效：在MLA的低秩压缩中，键的生成需要将上投影矩阵 $W^{UK}$ 融入查询投影矩阵 $W^Q$（见2.1.2节），以跳过显式计算键值对。但RoPE的旋转矩阵会插入在 $W^Q$ 和 $W^{UK}$ 之间，导致：
   $$\text{传统流程：}{k}_t^C=W^{UK}\cdot {\mathbf{R}}_t\cdot W^{DKV}{h}_t$$
   由于矩阵乘法不满足交换律，$W^{UK}$ 无法与 $W^Q$ 合并，必须为每个位置 $t$ 单独计算 $k_t^C$，导致推理时需重新计算所有前缀token的键，计算量随序列长度呈平方级增长（例如128K序列时，每次生成新token需计算128K次键，效率极低）。

在DeepSeek-V2模型中，将旋转位置编码（RoPE）与低秩KV压缩技术结合时，由于矩阵乘法不满足交换律，导致推理过程中必须重新计算所有前缀token的键（Key），从而显著降低效率。以下是该问题的技术细节及实例分析：

---

核心矛盾：RoPE与低秩KV压缩的互斥性

1. 低秩KV压缩的原理
   • 压缩过程：输入向量$h_t$通过低秩矩阵$W^{DKV}\in {\mathbb{R}}^{d_c\times d}$（$d_c$为压缩维度，$d$为原始维度）压缩为潜在向量$c_t^{KV}$：
   $$c_t^{KV}=W^{DKV}\cdot h_t$$
   • 还原键值：通过上投影矩阵$W^{UK}\in {\mathbb{R}}^{d\times d_c}$和$W^{UV}$还原完整键值：
   $$k_t^C=W^{UK}\cdot c_t^{KV},v_t^C=W^{UV}\cdot c_t^{KV}$$
   理想情况下，推理时可通过矩阵结合律将两步合并，例如将$W^{UK}\cdot W^{DKV}$融入$W^Q$，从而避免显式计算$k_t^C$。
   $$W^{Q^{\prime }}=W^Q\cdot (W^{UK}\cdot W^{DKV})\in {\mathbb{R}}^{512\times 512}$$

2. RoPE的位置敏感性
   • RoPE的作用：RoPE通过旋转矩阵$R_t$对查询（Query）和键（Key）引入位置相关的相位偏移：
   $$q_t=R_t\cdot (W^Q\cdot h_t),k_t=R_t\cdot (W^{UK}\cdot c_t^{KV})$$
   • 冲突点：RoPE矩阵$R_t$与低秩矩阵$W^{UK}$的位置耦合，导致矩阵乘法顺序不可交换。

---

数学推导与冲突示例

1. 原始计算流程
   • 键的计算：
   $$k_t^C=R_t\cdot (W^{UK}\cdot c_t^{KV})=R_t\cdot W^{UK}\cdot W^{DKV}\cdot h_t$$
   • 理想合并：若矩阵乘法可交换，可将$W^{UK}\cdot W^{DKV}$提前合并到$W^Q$中：
   $$W^{Q^{\prime }}=W^Q\cdot (W^{UK}\cdot W^{DKV})$$
   此时查询计算简化为$q_t^{\prime }=R_t\cdot W^{Q^{\prime }}\cdot h_t$，无需显式存储$k_t^C$。

2. 实际冲突分析
   • 交换律失效：由于$R_t$与$W^{UK}$的乘法顺序不可交换，即：
   $$R_t\cdot W^{UK}\ne W^{UK}\cdot R_t$$
   导致无法将$W^{UK}$吸收到$W^Q$中。
   • 示例：假设$W^{UK}\in {\mathbb{R}}^{512\times 128}$，$R_t\in {\mathbb{R}}^{512\times 512}$，则：
   $$R_t\cdot W^{UK}\in {\mathbb{R}}^{512\times 128},W^{UK}\cdot R_t\in {\mathbb{R}}^{128\times 512}(\text{维度不匹配})$$
   即使维度匹配（如方阵），数值结果仍不同：
   $$\text{若}{W}^{UK}=\left[\begin{array}{cc}1& 2\\ 3& 4\end{array}\right],R_t=\left[\begin{array}{cc}0& 1\\ 1& 0\end{array}\right],\text{则}{R}_t\cdot W^{UK}=\left[\begin{array}{cc}3& 4\\ 1& 2\end{array}\right]\ne W^{UK}\cdot R_t=\left[\begin{array}{cc}2& 1\\ 4& 3\end{array}\right]$$

---

二、解耦RoPE策略：分离位置信息的“双重表示”

为解决上述问题，DeepSeek-V2提出解耦策略：将查询和键拆分为位置无关的压缩部分（用于高效计算）和位置相关的RoPE部分（用于捕获序列顺序），两者独立处理后拼接使用。

三、关键公式与实例解析

1. 符号定义（简化维度示例）

假设：

• 输入维度 $d=512$，压缩维度 $d_c=128$，查询压缩维度 $d_c^{\prime }=256$
• 注意力头数 $n_h=32$，每头维度 $d_h=16$，解耦位置维度 $d_h^R=8$（每头位置相关子空间维度）
• $c_t^Q$：查询的压缩潜向量（256维，来自2.1.2节的查询压缩）
• $h_t$：输入隐藏层向量（512维）

2. 位置相关查询 $q_t^R$ 的生成（公式14）

$$[q_{t,1}^R;q_{t,2}^R;\dots ;q_{t,32}^R]=q_t^R=\text{RoPE}\left(W^{QR}{c}_t^Q\right)$$

• $W^{QR}$：投影矩阵（形状 $32\times 8\times 256$，将256维压缩查询投影到32头×8维的位置子空间）
• 实例：
  • $c_t^Q=[x_1,x_2,...,x_{256}]$（压缩后的查询）
  • $W^{QR}{c}_t^Q$ 生成32个8维向量（每头一个），经RoPE添加位置信息后，得到每头的位置相关查询 $q_{t,i}^R\in {\mathbb{R}}^8$。

3. 共享位置键 $k_t^R$ 的生成（公式15）

$$k_t^R=\text{RoPE}\left(W^{KR}{h}_t\right)$$

• $W^{KR}$：投影矩阵（形状 $8\times 512$，将512维输入投影到8维位置空间）
• 实例：
  • $h_t$ 经 $W^{KR}$ 投影为8维向量，再经RoPE处理，得到共享的位置键 $k_t^R\in {\mathbb{R}}^8$（所有头共用，减少参数和计算量）。

4. 查询与键的拼接（公式16-17）

$$q_{t,i}=[q_{t,i}^C;q_{t,i}^R],k_{t,i}=[k_{t,i}^C;k_t^R]$$

• 实例：
  • 压缩部分 $q_{t,i}^C\in {\mathbb{R}}^{16}$（每头原始维度），位置部分 $q_{t,i}^R\in {\mathbb{R}}^8$，拼接后查询维度为 $16+8=24$。
  • 键的压缩部分 $k_{t,i}^C\in {\mathbb{R}}^{16}$，位置部分 $k_t^R\in {\mathbb{R}}^8$，拼接后键维度为 $16+8=24$。

5. 注意力计算（公式18）

$$o_{t,i}=\sum _{j=1}^t{\text{Softmax}}_j\left(\frac{q_{t,i}^T{k}_{j,i}}{\sqrt{24}}\right)v_{j,i}^C$$

• 分母 $\sqrt{d_h+d_h^R}=\sqrt{16+8}=\sqrt{24}$ 是拼接后维度的归一化因子，确保注意力分数稳定。

四、核心优势：通过“分离-拼接”实现高效与位置建模的双赢

1. 保留矩阵吸收（高效推理）：

   • 压缩部分 $q_t^C,k_t^C$ 不含位置信息，其投影矩阵 $W^{UK},W^{UV}$ 仍可被吸收到 $W^Q,W^O$ 中（如 $W^Q=W^{UQ}\cdot W^{DQ}$），避免重新计算前缀键。
   • 实例：若序列长度为128K，传统RoPE需缓存 $2\times 32\times 16\times 128K\times l$ 个元素，而解耦策略仅需缓存 $(128+8)\times l\times 128K$ 个元素，缓存量减少93.3%（与文档表1数据一致）。

2. 精准位置建模：

   • 位置相关部分 $q_t^R,k_t^R$ 专门处理序列顺序，通过RoPE捕获长距离依赖，例如：
     • 当token位置从 $t$ 到 $t+1$，$q_t^R$ 和 $q_{t+1}^R$ 的旋转矩阵不同，确保模型区分前后顺序。

3. 计算量控制：

   • 共享键 $k_t^R$ 对所有头生效，无需为每个头单独计算位置键，计算量从 $O(n_h{d}_{h}l)$ 降至 $O((d_c+d_h^R)l)$，支持128K长上下文高效处理（见文档图4，DeepSeek-V2在128K时性能稳定）。

五、与传统RoPE的对比（以1头为例）

步骤	传统RoPE（直接应用）	解耦RoPE（MLA策略）
键生成	$k_t=W^K{h}_t\cdot {\mathbf{R}}_t$（含位置矩阵）	$k_t=[k_t^C;k_t^R]$（$k_t^C$ 无位置，$k_t^R$ 含RoPE）
矩阵吸收	不可吸收（${\mathbf{R}}_t$ 阻断）	可吸收（$k_t^C$ 的 $W^{UK}$ 融入 $W^Q$）
前缀键重计算	每次生成新token需重算所有历史键	仅需计算一次共享 $k_t^R$，$k_t^C$ 可快速生成
KV缓存（1层）	$2\times 16\times 1=32$ 元素	$128（d_c）+8（d_h^R）=136$ 元素（虽略增，但远低于MHA的 $2\times 16\times 32=1024$ 元素）

六、总结：解耦策略如何让MLA“两全其美”

通过将位置信息从压缩的键/查询中分离，解耦RoPE策略实现了三大突破：

1. 技术兼容：在低秩KV压缩的同时，保留RoPE的位置建模能力，避免传统RoPE导致的推理瓶颈。
2. 效率提升：通过矩阵吸收和共享位置键，将键值对的计算复杂度从 $O(n^2)$ 降至 $O(n)$，支持128K长上下文高效生成。
3. 性能保障：实验表明（见文档附录D.2），解耦后的MLA在保持KV缓存大幅减少的同时，性能优于传统MHA，实现了“高效+强性能”的平衡。

这一设计是DeepSeek-V2能在长上下文场景（如128K序列）中保持高吞吐量的关键技术之一，也是其架构创新的核心亮点。

2.5.3.MLA完整计算过程

为了展示MLA的完整计算过程，我们在下面给出其完整公式：

其中，蓝色框选的向量在生成时需要缓存。在推理过程中，原始公式需要从 $c_t^{KV}$ 中恢复 $k_t^c$ 和 $v_t^c$ 用于注意力计算。幸运的是，由于矩阵乘法的结合律，我们可以将 $W^{UK}$ 合并到 $W^{UQ}$ 中，将 $W^{UV}$ 合并到 $w^o$ 中。因此，我们无需为每个查询计算键值对。通过这种优化，我们避免了推理过程中重新计算 $k_r^C$ 和 $v_t^C$ 的计算开销。

数学：矩阵分块计算，多token的矩阵分成单个token相乘

详细解释举例说明：向量与多维矩阵乘法，矩阵分块计算，多token的矩阵如何分成单个token去相乘

向量与多维矩阵乘法

定义：向量与矩阵相乘时，需满足维度匹配。若矩阵为 $A\in {\mathbb{R}}^{m\times n}$，向量为 $v\in {\mathbb{R}}^n$，则乘积 $Av$ 的结果为 ${\mathbb{R}}^m$ 的向量。计算方式为矩阵每行与向量的点积。

示例：

• 设矩阵 $A=\left[\begin{array}{cc}1& 2\\ 3& 4\end{array}\right]$，向量 $v=\left[\begin{array}{c}5\\ 6\end{array}\right]$。
• 计算 $Av$：
  $$Av=\left[\begin{array}{c}1\times 5+2\times 6\\ 3\times 5+4\times 6\end{array}\right]=\left[\begin{array}{c}17\\ 39\end{array}\right].$$

---

矩阵分块计算

定义：将大矩阵划分为小块（子矩阵），通过块间运算简化计算。分块需满足乘法规则：若 $A$ 分块为 $[A_{ij}]$，$B$ 分块为 $[B_{jk}]$，则乘积 $C=AB$ 的块 $C_{ik}={\sum }_j{A}_{ij}{B}_{jk}$。

示例：

• 矩阵 $A=\left[\begin{array}{cc}{A}_{11}& A_{12}\\ A_{21}& A_{22}\end{array}\right]$，$B=\left[\begin{array}{cc}{B}_{11}& B_{12}\\ B_{21}& B_{22}\end{array}\right]$，其中每个子块为 $2\times 2$ 矩阵。
• 计算块 $C_{11}$：
  $$C_{11}=A_{11}{B}_{11}+A_{12}{B}_{21}.$$
• 所有块计算完成后，合并得到 $C$。

---

多Token矩阵分块处理

场景：在自然语言处理中，多个token（如词向量）按行排列成矩阵 $X\in {\mathbb{R}}^{k\times d}$，需计算 $XW$（$W\in {\mathbb{R}}^{d\times h}$）。

分块策略：

1. 按行分块：将 $X$ 拆分为单个token的行向量。

   • 例如，$X=\left[\begin{array}{c}\text{—}{x}_1\text{—}\\ \text{—}{x}_2\text{—}\\ ⋮\\ \text{—}{x}_k\text{—}\end{array}\right]$，每个 $x_i\in {\mathbb{R}}^{1\times d}$。
   • 每个token独立计算：$y_i=x_{i}W$（结果 $y_i\in {\mathbb{R}}^{1\times h}$）。
   • 合并结果：$Y=\left[\begin{array}{c}{y}_1\\ y_2\\ ⋮\\ y_k\end{array}\right]\in {\mathbb{R}}^{k\times h}$。

2. 批量分块：将 $X$ 划分为多个子矩阵，每个子矩阵含多个token。

   • 例如，将 $X$ 分为 $X_1\in {\mathbb{R}}^{k_1\times d}$ 和 $X_2\in {\mathbb{R}}^{k_2\times d}$，则：
     $$Y=\left[\begin{array}{c}{X}_{1}W\\ X_{2}W\end{array}\right].$$
     示例：

• 设 $X=\left[\begin{array}{cc}1& 2\\ 3& 4\\ 5& 6\end{array}\right]$，$W=\left[\begin{array}{cc}0.5& 1\\ 1.5& 2\end{array}\right]$。
• 按行分块：
  • $x_1=[12]\to y_1=[1\times 0.5+2\times 1.5,1\times 1+2\times 2]=[3.5,5]$。
  • 同理计算 $x_2,x_3$，合并得 $Y=\left[\begin{array}{cc}3.5& 5\\ 8.5& 11\\ 13.5& 17\end{array}\right]$。

---

分块计算的优势

1. 并行计算：不同块可分配到多个处理器加速计算。
2. 内存优化：处理大矩阵时，分块减少内存占用。
3. 灵活性：支持动态加载数据，适合分布式系统。

---

总结

• 向量乘矩阵：按行或列的点积计算。
• 矩阵分块：通过子矩阵运算简化复杂计算。
• 多Token处理：按行分块后独立计算，或批量处理子矩阵，灵活平衡效率与资源。