DeepSeek V2/V3中的MLA和Matrix Absorption

DeepSeek V3的网络结构基本沿用了DeepSeek V2，采用了MLA和DeepSeekMoE两大特性。本文主要涉及MLA（Multi-Head Latent Attention）。抛开维度变化，DeepSeek V3与V2在MLA结构上差别不大。详细请参见官方论文《DeepSeek-V3 Technical Report》和《DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model》。

关于MLA的介绍网上很多，不多讲了。这里将论文中的示意图，公式与官方代码（modeling_deepseek.py）的对应关系做了标注。

论文中还提到matrix absorption的优化：
Fortunately, due to the associative law of matrix multiplication, we can absorb $W^{UK}$ into $W^{UQ}$, and $W^{UV}$ into $W^O$.

为了方便看出差别，这里按论文中的convention整了下matrix absorption后的示意图和公式，并标注了对应关系。代码由于官方没给出，参考的是SGLang中python/sglang/srt/models/deepseek_v2.py里的DeepseekV2AttentionMLA::forward_absorb。

本质上是利用了矩阵乘的结合率，将MHA转为MQA。其中matrix absorption的部分由于计算只涉及权重参数，因此可以提到初始化时，或者离线做。实现可以参考FlashInfer中tests/test_mla_decode_kernel.py中的DeepseekV2AttentionMatAbsorbDecode。

但注意该优化适用于generation阶段，不适用于prefill阶段。看下优化前后的相关两部分计算量比较：

---

将模型参数代入可发现，generation阶段时可以减少计算量，而prefill阶段时不能。