编者按： 模型在生成长序列文本时，键值缓存占用的显存是否让 GPU 不堪重负？如何在保持模型性能的同时有效降低大语言模型推理阶段的内存占用？
每一次模型推理，大量重复计算和庞大的内存占用不仅大幅增加了运营成本，还严重限制了模型的批处理能力和响应速度。多查询注意力机制（MQA）和分组查询注意力机制（GQA）虽能部分缓解问题，但往往以牺牲模型性能为代价。
DeepSeek 团队在模型优化过程中开发的多头潜在注意力机制（MLA）为资源受限场景提供了可行性方案。通过巧妙的低维潜在向量压缩和解耦式 RoPE 设计， MLA 能在显存效率与建模能力间实现更优平衡。
作者 | Shirley Li
编译 | 岳扬
本文是「DeepSeek-V3 技术解析」系列文章的第一篇，我们将深入解读深度求索公司最新开源的 DeepSeek-V3 模型1,2。
在本系列文章中，我们计划涵盖两大主题：
- DeepSeek-V3 的主要架构创新，包括：多头潜在注意力（Multi-head Latent Attention，MLA）3、DeepSeekMoE4、无辅助损失的负载均衡技术（auxiliary-loss-free load balancing）5、多词元预测训练方法（multi-token prediction training）。
- DeepSeek-V3 的训练过程，包括：预训练阶段、微调阶段、强化学习对齐阶段。
本文重点解析首次在 DeepSeek-V2 中提出、并延续应用于 DeepSeek-V3 的多头潜在注意力（MLA）架构。
内容大纲：
- 技术背景：从标准多头注意力（MHA）机制出发，解析推理阶段 Key-Value 缓存的必要性，探讨 MQA（多查询注意力机制）和 GQA（分组查询注意力机制）的优化思路，并阐述"旋转"位置编码（RoPE）的工作原理。
- 多头潜在注意力机制（MLA） ：深入解析 MLA 的创新动机，阐释解耦式 RoPE 的设计原理，并展示其性能优势。
- 参考文献
01 技术背景
为了便于读者更好地理解 MLA 架构并保持本文的完整性，在深入探讨 MLA 的技术细节之前，我们将首先回顾几个相关技术概念。
1.1 Decoder-only Transformers 中的多头注意力机制
有一点需要特别说明，MLA 架构的研发初衷是加快自回归文本生成的推理效率，因此本文讨论的 MHA（多头注意力机制）特指 Decoder-only Transformer 架构。
图 1 对比了三种可用于解码的 Transformer 架构：(a) 展示的是《Attention is All You Need》论文提出的编码器-解码器架构。其解码器部分随后被文献 6 简化，形成如 (b) 所示的 decoder-only Transformer 架构，该架构被 GPT 等生成模型广泛采用8。
(c) 则是当前大语言模型更常采用的优化版本，对输入而不是输出进行归一化。并采用 RMS Norm 替代 LayerNorm，本文将以 (c) 所示架构作为 baseline 架构展开讨论。
图 1. Transformer 架构。(a) 文献 6 中提出的编码器-解码器结构；(b) 文献 7 提出的 Decoder-only Transformer，并被 GPT 8 采用；(c)是(b) 的优化版本，在注意力层前加入 RMS Norm 3。
在这种情况下， MHA 计算大致遵循文献 6 中的流程，如下图所示：
图 2. 缩放点积注意力与多头注意力架构对比（图源文献6）
假设模型包含 n_h 个注意力头，每个注意力头维度为 d_h，则拼接后总维度为 n_h·d_h。
对于具有 l 层的模型，若设某层第 t 个 token 的输入表示为维度 d 的向量 h_t，则需要通过线性映射矩阵（linear mapping matrices）将 h_t 的维度从 d 映射到（h_n - d_h）。
更形式化地表述如下（公式引用自文献3）：
其中，W^Q、W^K 和 W^V 是线性映射矩阵：
完成映射后，q_t、k_t 和 v_t 将被分割为 n_h 个注意力头进行缩放点积注意力计算：
最终通过投影矩阵 W^O 将维度从 (h_n - d_h) 反向映射到 d：
需要强调的是，上述公式 (1)-(8) 描述的是单 token 处理流程。在推理过程中，每个新生成的 token 都需要重复这一计算流程，这产生了大量重复计算 ------ 这正是催生键值缓存（Key-Value Cache）技术的关键动因。
1.2 Key-Value Cache
正如其名称所示，键值缓存（Key-Value Cache）是一种通过缓存和复用先前计算的键值对（Keys and Values）来加速自回归过程的技术，避免了在每个解码步骤重新计算这些参数的需求。
需要注意的是，键值缓存通常仅在推理阶段使用，因为在训练阶段，我们仍然需要并行处理整个输入序列。
键值缓存通常以滚动缓冲区（rolling buffer）的形式实现。在每个解码步骤中，仅计算新的查询向量 Q（Query），而缓存中存储的 K（Keys）和 V（Values）会被复用，注意力机制将通过新计算的 Q 与复用的 K、V 进行运算。与此同时，新生成的 token 对应的 K 和 V 也会被追加到缓存中以供后续使用。
然而，键值缓存带来的加速是以内存消耗为代价的。由于键值缓存的大小通常与批处理量大小（batch size）× 序列长度（sequence length）× 隐藏层维度（hidden size）× 注意力头数（number of heads）成正比，当出现更大批处理量或更长序列时，极易形成内存瓶颈。
这种限制进一步催生了两项旨在突破该瓶颈的技术：多查询注意力机制（Multi-Query Attention）和分组查询注意力机制（Grouped-Query Attention）。
1.3 多查询注意力机制（MQA） vs 分组查询注意力机制（GQA）
下图展示了原始的多头注意力机制（MHA）、分组查询注意力机制（GQA）10 和多查询注意力机制（MQA）9 的对比。
图 3. MHA 6、GQA 10 与 MQA 9（图源文献10）
MQA 的核心思想是让所有查询头（query heads）共享一个单独的键头（key head）和值头（value head）。 这种做法能显著降低内存占用，但也会对注意力的计算精度产生负面影响。
GQA 可视为 MHA 与 MQA 的一种折中方案。 在这种方法中，一组查询头仅共享一对键头和值头，而非所有查询头共享同一对。尽管如此，相较于原始的 MHA，其效果仍会有所逊色。
在后续章节中，我们将探讨 MLA 如何在内存效率与建模精度之间实现平衡。
1.4 旋转位置编码（RoPE）
最后需要提及的相关背景知识是旋转位置编码（RoPE）11。该方法通过在多头注意力机制中对查询向量（Query）和键向量（Key）施加基于正弦函数的旋转操作，将位置信息直接编码到注意力计算中。
具体而言，RoPE 会对每个词元（token）的查询向量和键向量应用一个与位置相关的旋转矩阵。该矩阵以正弦和余弦函数为基础，但以一种独特的方式应用它们来实现旋转。
为理解其"位置相关（position-dependent）"特性，我们以一个 4 维的嵌入向量（x₁, x₂, x₃, x₄）为例。
要应用 RoPE，首先需要将连续的维度两两配对：
- (x₁, x₂) → 位置1
- (x₃, x₄) → 位置2
然后，对其中的每一对应用旋转矩阵：
图 4. 应用到词元对的旋转矩阵示意图。由原文作者供图。
其中旋转角度 θ = θ(p) = p ⋅ θ₀（θ₀为基频参数）。在此示例中，(x₁, x₂) 会被旋转 θ₀ 角度，而 (x₃, x₄) 则被旋转 2 · θ₀ 角度。
因此，我们称这种旋转矩阵为"位置相关（position-dependent）"：在每个位置（或每对），将应用不同的旋转矩阵，其中旋转角度由位置决定。
RoPE 因其在长序列编码中的高效性被现代大语言模型广泛采用。然而，从上述公式可见，该方法对查询向量和键向量的位置敏感性，使其在部分场景下与 MLA 机制不兼容。
02 Multi-head Latent Attention
最后我们终于可以进入 MLA 部分的讲解。在这一章节中，我们将首先阐述 MLA 的核心设计思想，然后深入探讨其为何需要修改 RoPE，最后展示 MLA 的具体算法实现及其性能表现。
2.1 MLA：核心设计思想
MLA 的核心设计思想是将注意力机制的输入 h_t 压缩为一个低维的潜在向量（维度 d_c，且 d_c 远小于原始维度（h_n·d_h））。 当需要计算注意力时，我们可以将这个潜在向量重新映射回高维空间，以还原键向量和值向量。通过这种方式，只需存储潜在向量即可，从而实现显存占用的大幅降低。
这个过程可以通过以下公式更正式地进行描述。其中 c^{KV}_t 表示潜在向量，W^{DKV} 是压缩矩阵（上标 D 代表"下投影"，即降维操作），负责将 h_t 的维度从（h_n·d_h）压缩到d_c；而 W^{UK} 和 W^{UV} 则是上投影矩阵，负责将共享的潜在向量映射回高维空间。
类似地，我们也可以将查询向量映射到一个潜在的低维向量，然后再将其映射回原始的高维空间：
2.2 为何需要解耦式 RoPE
正如前文所述，RoPE 是训练生成模型处理长序列的常用位置编码方案。如果我们直接应用上述 MLA 策略，将会与 RoPE 不兼容。
为了更清晰地理解这一点，来看看使用公式 (7) 计算注意力时的情形：当我们对 q 进行转置并与 k 相乘时，矩阵 W^Q 和 W^{UK} 会出现在中间环节，它们的组合等效于从将 d_c（输入向量的维度）映射到目标维度 d。
在文献 3 中，作者将此现象描述为 W^{UK} 可以被 W^Q "吸收"，因此我们无需在缓存中存储 W^{UK}，从而进一步降低了内存使用量。
然而，当我们将图 (4) 中的旋转矩阵考虑在内时，情况就不是这样了 ------ RoPE 会在 W^{UK} 左侧施加一个旋转矩阵，而该旋转矩阵最终会夹在转置后的 W^Q 和 W^{UK} 之间。
正如技术背景部分所述，这个旋转矩阵是位置相关（position-dependent）的，即每个位置对应的旋转矩阵都不同。因此，W^{UK} 无法再被 W^Q 吸收。
为解决这一矛盾，作者提出了"解耦式RoPE"方案：通过引入额外的查询向量和一个共享的键向量，并仅在 RoPE 过程中使用这些新增向量，同时保持原始键向量与旋转矩阵的隔离。
完整的 MLA 流程可总结如下（公式编号沿用文献 3 附录 C 的编排）：
图 5. MLA 处理流程（作者根据文献 3 中的公式所编辑的示意图）
其中：
- 公式 (37)-(40) 描述查询向量的处理流程
- 公式 (41)-(42) 描述键向量的处理流程
- 公式 (43)-(44) 说明如何将新增的共享键向量用于 RoPE，需注意公式 (42) 的输出不参与 RoPE 计算
- 公式 (45) 描述值向量的处理流程
在此过程中，仅需缓存蓝色变量。该流程可通过下图更直观地展示：
图 6. MLA 流程图（来源：文献3）
2.3 MLA 的性能表现
下表对比了 MHA、GQA、MQA 与 MLA 在 KV 缓存（每个 token）所需元素数量及建模能力上的差异，证明 MLA 确实能在显存效率与建模能力间实现更优平衡。
有趣的是，MLA 的建模能力甚至超越了原始 MHA 架构。
文献 3 中的表 1
下表展示了 7B 参数规模下 MHA、GQA 与 MQA 的性能表现，其中 MHA 显著优于 MQA 和 GQA。
文献 3 中的表 8
文献 3 的作者还对 MHA 与 MLA 进行了对比分析，结果如下表所示。数据显示 MLA 在整体指标上表现更优。
文献 3 中的表 9
03 References
1 DeepSeek（www.deepseek.com/）
2 DeepSeek-V3 Technical Report（github.com/deepseek-ai...
3 DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model（arxiv.org/abs/2405.04...
4 DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models（arxiv.org/abs/2401.06...
5 Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts（arxiv.org/abs/2408.15...
6 Attention Is All You Need（arxiv.org/abs/1706.03...
7 Generating Wikipedia by Summarizing Long Sequences（arxiv.org/pdf/1801.10...
8 Improving Language Understanding by Generative Pre-Training（cdn.openai.com/research-co...
9 Fast Transformer Decoding: One Write-Head is All You Need（arxiv.org/pdf/1911.02...
10 GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints（arxiv.org/abs/2305.13...
11 RoFormer: Enhanced Transformer with Rotary Position Embedding（arxiv.org/abs/2104.09...
Thanks for reading!
Hope you have enjoyed and learned new things from this blog!
About the author
Shirley Li
I am a Machine Learning Engineer working on building multi-modality models to solve real-world problems.
END
本期互动内容 🍻
❓你最看重的是模型性能还是推理效率？为什么？欢迎在评论区谈谈你的看法！
原文链接：