探秘Transformer系列之（28）--- DeepSeek MLA
探秘Transformer系列之（28）--- DeepSeek MLA目录
[*]探秘Transformer系列之（28）--- DeepSeek MLA
[*]0x00 概述
[*]0x01 原理
[*]1.1 问题
[*]1.2 当前状况
[*]1.3 改进思路
[*]1.3.1 增强信息压缩能力
[*]思路
[*]方案
[*]问题
[*]1.3.2 丰富信息表达
[*]思路
[*]方案
[*]1.3.2 解决位置编码冲突
[*]1.4 架构图 & 流程
[*]1.5 代码
[*]0x02 核心要点
[*]2.1 低秩KV联合压缩
[*]2.1.1 低秩分解
[*]2.1.2 思路
[*]2.1.3 向下投影
[*]2.1.4 向上投影
[*]null
[*]2.1.5 完整过程
[*]2.2 权重吸收
[*]2.2.1 当前状态
[*]2.2.2 权重吸收
[*]2.2.3 推导
[*]KQ合并
[*]VO合并
[*]结合
[*]2.2.4 讨论
[*]训练
[*]MHA
[*]不合并
[*]2.3 解耦RoPE
[*]2.3.1 RoPE背景
[*]2.3.2 问题
[*]无法直接应用到低秩压缩
[*]与权重吸收不兼容
[*]2.3.3 解决方案
[*]2.3.5 和权重吸收结合
[*]2.4 资源占用
[*]2.4.1 参数量
[*]2.4.2 内存占用
[*]2.4.3 计算量
[*]2.4.4 信息转移
[*]2.5 并行
[*]0x03 计算过程
[*]3.1 公式
[*]3.2 原始流程
[*]3.3 吸收
[*]3.3.1 过程
[*]3.3.2 吸收结果
[*]3.3.3 MQA形式
[*]0x04 代码
[*]4.1 配置
[*]4.2 定义
[*]4.3 操作Q
[*]4.3.1 变量定义
[*]4.3.2 变量操作
[*]4.4 操作KV
[*]4.4.1 变量定义
[*]4.4.2 变量操作
[*]4.5 注意力操作
[*]4.5.1 变量定义
[*]4.5.2 变量操作
[*]4.6 前向传播
[*]4.7 V3 代码
[*]0x05 优化代码
[*]5.1 压缩优化
[*]5.2 权重吸收
[*]5.2.1 absorbed_cache_compressed.py
[*]\(W^{UK}\)
[*]\(W^{UV}\)
[*]5.2.2 Move Elision
[*]5.2.3 Materializing Projection Matrices
[*]5.3 融合算子
[*]5.4 矩阵乘的重排序（增补@2025-04-19）
[*]0x06 转换
[*]6.1 GQA
[*]6.1.1 思路
[*]6.1.2 方案
[*]6.2 MHA
[*]6.2.1 partial-RoPE
[*]MHA
[*]拆解
[*]6.2.2 低秩近似
[*]0xFF 参考
0x00 概述
MLA（Multi-head Latent Attention / 多头潜在注意力）的基本思想是将注意力输入\(h_t\) 压缩成一个低维的潜在向量 \(c^{KV}_t\) ，维度为 \(d_c\)，且 \(d_c\) 远小于原始的维度（\(h_nd_h\)）。在需要计算注意力时，可将这个潜在向量 \(c^{KV}_t\) 映射回高维空间。因此，只需要存储潜在向量 \(c^{KV}_t\) ，就可以显著减少内存的占用。
这个过程可以通过以下公式更正式地进行描述。其中 \(c^{KV}_t\) 表示潜在向量；\(W^{DKV}\) 是压缩矩阵（上标 D 代表"下投影"，即降维操作），负责将 \(h_t\) 的维度从（\(h_n·d_h\)）压缩到\(d_c\)；\(W^{UK}\)和 \(W^{UV}\) 是上投影矩阵，负责将共享的潜在向量 \(c^{KV}_t\) 映射回高维空间。只需要存储这个潜在向量 \(c^{KV}_t\) ，就能获得对应不同文本特征的Key和Value，而不需要对每个文本特征都存储对应的Key和Value。
类似地，我们也可以将查询向量映射到一个潜在的低维向量，然后再将其映射回原始的高维空间。而且，MLA又结合了权重吸收技术，减少了计算开销。
注：
[*]全部文章列表在这里，估计最终在35篇左右，后续每发一篇文章，会修改此文章列表。cnblogs 探秘Transformer系列之文章列表
[*]本系列是对论文、博客和代码的学习和解读，借鉴了很多网上朋友的文章，在此表示感谢，并且会在参考中列出。因为本系列参考文章实在太多，可能有漏给出处的现象。如果原作者或者其他朋友发现，还请指出，我在参考文献中进行增补。
0x01 原理
1.1 问题
标准Transformer的一大障碍就是KV Cache的空间占用问题：多头注意力机制需要为每个注意力头单独存储历史生成的Key和Value向量（即KV缓存）。随着序列长度增加，KV缓存的存储需求呈指数级增长，导致内存占用急剧上升。而GPU的显存空间往往非常有限，较大的KV Cache会导致同时处理的request数量变少，也即batch size较小；为了减少KV缓存需求，研究人员提出了像Multi-Query Attention（MQA）和Group-Query Attention（GQA）这些方法。这些方法虽然降低了缓存要求，可模型的性能也受到影响。MQA或GQA算子在计算注意力的过程中，所有KV Cache中的数据读取后都仅参与一次或几次计算，导致该算子的MFU极低，并且由于每个request有自己的KV Cache，这一问题无法通过提高batch size的方式解决。
因此，如何减少推理过程的KV Cache，从而实现在更少的设备上推理更长的Context，或者在相同的Context长度下增大batch size，实现更快的推理速度或者更大的吞吐总量，最终降低推理成本。是一个关键问题。
1.2 当前状况
我们首先总结下当前各种方案的情况来看看有没有可以改进的空间。下图给出了MHA、GQA、MQA 以及 MLA 做法。
图上从左到右依次是 MHA、GQA、MQA 以及 MLA 。图中有阴影的长条表示会缓存到显存的结果。MHA、GQA、MQA 都需要将 KVCache 缓存到显存。几种方案特点如下。
<ul>MHA：MHA KVCache 在注意力头这个维度和 Q 矩阵一样，属于“一对一”。MHA把一个注意力计算拆成多个注意力头，每个注意力头使用独立的Q、K、V进行计算，需要把K、V都存储下来，KV Cache中每个token需要缓存的参数量为\(2
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！
页:
[1]