引言

分组查询注意力（GQA）作为一种针对 MHA 的计算效率改进方案，消融研究（例如原始 GQA 论文以及 Llama 2 论文中的实验）表明，在大语言模型的建模性能方面，GQA 与标准 MHA 的表现大致相当。

现在，多头潜在注意力（MLA）——它被用于 DeepSeek V2、V3 和 R1 ——提供了一种不同的节省内存策略，并且与 KV 缓存配合得尤为良好。与 GQA 通过共享 key 和 value 头不同，MLA 会在将 key 和 value 存入 KV 缓存之前，将它们压缩到一个更低维度的空间中。

在推理阶段，这些被压缩的张量会在使用之前被投影回原始维度，如下图所示。这会增加一次额外的矩阵乘法，但能够显著降低内存使用。

（顺带一提，query 在训练阶段也会被压缩，但在推理阶段不会。）

另外，正如前面提到的，MLA 并非在 DeepSeek V3 中首次出现，因为它的前身 DeepSeek V2 也使用了（甚至是首次引入了）这一机制。此外，V2 论文中包含了一些有趣的消融实验，或许可以解释 DeepSeek 团队为何选择 MLA 而不是 GQA（见下图）。

如上图所示，GQA 的性能似乎劣于 MHA，而 MLA 的建模性能则优于 MHA，这很可能是 DeepSeek 团队选择 MLA 而非 GQA 的原因。（如果能看到 MLA 与 GQA 在“每 token 的 KV 缓存”方面的节省对比，将会非常有意思！）

总结这一部分，在进入下一个架构组件之前，可以说 MLA 是一种巧妙的技巧，它在减少 KV 缓存内存使用的同时，甚至在建模性能上还略微优于 MHA。

MLA 的内存节省

内存节省主要体现在 KV 存储上。我们可以使用如下公式来计算 KV 存储大小：


     
     bytes ≈ batch_size × seqlen × n_layers × latent_dim × bytes_per_elem

相比之下，MHA 的 KV 缓存内存计算方式如下：


     
     bytes ≈ batch_size × seqlen × n_layers × embed_dim × 2 (K,V) × bytes_per_elem

这意味着，在 MLA 中，我们将 “embed_dim × 2（K,V）” 降低为 “latent_dim”，因为我们只存储了压缩后的潜在表示，而不是完整的 key 和 value 向量，如前面的示意图所示。

你可以使用本目录中的 memory_estimator_mla.py 脚本，将其应用到不同的模型配置中，以查看使用 MLA 相对于 MHA 能节省多少内存：


     
     ➜ uv run memory_estimator_mla.py \
  --context_length 8192 \
  --emb_dim 2048 \
  --n_heads 24 \
  --n_layers 48 \
  --n_kv_groups 4 \
  --batch_size 1 \
  --dtype bf16 \
  --latent_dim 1024

==== Config ====
context_length   : 8192
emb_dim          : 2048
n_heads          : 24
n_layers         : 48
n_kv_groups      : 4
latent_dim       : 1024
batch_size       : 1
dtype            : bf16 (2 Bytes/elem)
head_dim         : 86
GQA n_kv_heads   : 6

==== KV-cache totals across all layers ====
MHA total KV cache  : 3.25 GB
GQA total KV cache  : 0.81 GB
MLA total KV cache  : 0.81 GB
Ratio (MHA / GQA)   : 4.00x
Savings (GQA vs MHA): 75.00%
Ratio (MHA / MLA)   : 4.03x
Savings (MLA vs MHA): 75.19%

请注意，上述压缩（--emb_dim 2048 -> latent_dim 1024）是为了实现与 GQA 类似的节省效果。在实践中，压缩比例是一个需要仔细研究的超参数，因为如果将 latent_dim 设得过小，可能会对建模性能产生负面影响（这与在 GQA 中选择过多的 n_kv_groups 类似）。

下图进一步展示了在不同 latent_dim 值下，随着上下文长度变化，使用 MLA 相对于 MHA 的节省情况：

你可以通过运行 uv run plot_memory_estimates_mla.py 来复现该图。

MLA 代码示例

本目录中的 gpt_with_kv_mha.py 和 gpt_with_kv_mla.py 脚本，提供了在 GPT 模型实现中对比 MHA 与 MLA 内存使用情况的实践示例。

在这里，MLA 的代码实现受到了以下实现的启发：

https://huggingface.co/bird-of-paradise/deepseek-mla

需要注意的是，MLA 也可以与 GQA 结合使用，但为了简化起见，这里并未这样做。（目前我也并不知道有哪一个知名的大模型在实际中采用了这种组合。）

还需要注意的是，该模型并未进行训练，因此生成的文本并没有实际意义。不过，你可以将它作为第 5–7 章中标准 GPT 模型的直接替换，并对其进行训练。

最后，该实现使用了在另一个附加章节中介绍的 KV 缓存机制，因此内存节省效果会更加明显。


     
     uv run gpt_with_kv_mha.py \
--max_new_tokens 32768 \
--n_heads 24 \
--n_layers 12 \
--emb_dim 768

...

Time: 453.81 sec
72 tokens/sec
Max memory allocated: 1.54 GB


     
     uv run gpt_with_kv_mla.py \
--max_new_tokens 32768 \
--n_heads 24 \
--n_layers 12 \
--emb_dim 768 \
--latent_dim 192 # (768×2)/192 = 8× compression

...

Time: 487.21 sec
67 tokens/sec
Max memory allocated: 0.68 GB

之所以没有看到像前面图表中那样显著的节省效果，主要有两个原因：

1. 我使用了一个较小的配置，以便模型能够在合理的时间内完成生成。
2. 更重要的是，这里观察的是整个模型，而不仅仅是注意力机制；模型中的全连接层占用了大部分内存（但这是另一个需要单独分析的话题）。

本杂志是一个个人热情驱动的项目，你的支持有助于它持续发展。

如果你愿意支持我的工作，可以考虑我的书 《Build a Large Language Model (From Scratch)》，或它的续作 《Build a Reasoning Model (From Scratch)》。（我相信你会从中收获颇多；它们以你在其他地方很难看到的深度，系统讲解了 LLM 的工作原理。）

感谢你的阅读，也感谢你对独立研究的支持。