大模型KV Cache节省神器MLA学习笔记（包含推理时的矩阵吸收分析）

文章预览

首先，本文回顾了MHA的计算方式以及KV Cache的原理，然后深入到了DeepSeek V2的MLA的原理介绍，同时对MLA节省的KV Cache比例做了详细的计算解读。接着，带着对原理的理解理清了HuggingFace MLA的全部实现，每行代码都去对应了完整公式中的具体行并且对每个操作前后的Tensor Shape变化也进行了解析。我们可以看到目前的官方实现在存储KV Cache的时候并不是存储隐向量，而是把隐向量都解压缩变成了标准的MHA的KV Cache，实际上是完全不能节省显存的。接着，就继续学习了一下清华大学的ZHANG Mingxing组实现的MLA矩阵吸收的工程实现，在这一节也详细分析了原理包括 和 分别如何吸收到 和 中，分析了实现了矩阵吸收的每行代码的原理以及操作发生前后相关Tensor的维度变化。接着，对矩阵吸收代码实现里的矩阵乘法的性质进行分析，可以看到MLA在大多数阶段都是计算 ………………………………