探秘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
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！