MLA 架构

注：本文为 “MLA 架构” 相关文章合辑。
未整理去重。

---

DeepSeek 的 MLA 架构

原创 老彭坚持 产物经理修炼之道 2025 年 01 月 28 日 10:15 江西
DeepSeek 的 MLA（Multi-head Latent Attention，多头潜在留意力）架构
是一种优化后的留意力机制，旨在减少推理过程中的键值（KV）缓存，同时保持与标准多头留意力（MHA）相称的性能。
以下是 MLA 架构的工作原理：
低秩联合压缩

核心头脑：MLA 通过对留意力机制中的键（Key）和值（Value）举行低秩联合压缩，减少推理时必要缓存的 KV 对数量。
详细实现：设嵌入维度为 d model，留意力头数为 h，每个头的维度为 d k，则对于给定留意力层中第 i 个 token 的留意力输入，MLA 将键和值压缩为低秩潜在向量：
                                              K                            ′                                  =                         K                                   W                            K                                       d                               o                               w                               n                                                 K'=KW_{K}^{down}                  K′=KWKdown​
                                              V                            ′                                  =                         V                                   W                            V                                       d                               o                               w                               n                                                 V'=VW_{V}^{down}                  V′=VWVdown​
此中，W Kdown 和 W Vdown 是下投影矩阵，用于将原始的键和值映射到低秩空间。
优势：通过这种低秩压缩，MLA 显著减少了 KV 缓存的内存占用，从而在推理过程中减少了内存利用，提拔了推理效率。
查询的低秩压缩

优化目的：除了对键和值举行压缩，MLA 还对留意力查询（Query）举行低秩压缩，以减少训练过程中的激活内存。
实现方式：查询的低秩压缩通过雷同的投影操作实现，进一步减少了模型在训练阶段的内存需求。
与标准多头留意力的对比

性能保持：尽管 MLA 通过低秩压缩减少了 KV 缓存和激活内存，但它仍旧能够保持与标准多头留意力（MHA）相称的性能。
推理效率提拔：在推理过程中，MLA 只必要缓存压缩后的键和值，这显著减少了内存占用，使得模型能够处理更长的上下文长度。
旋转位置嵌入（RoPE）

位置信息处理：MLA 架构还结合了旋转位置嵌入（RoPE），有效处理了长序列中的位置依靠标题。
作用：RoPE 通过旋转操作将位置信息嵌入到键和查询中，使得模型能够更好地捕捉长距离依靠关系。
综上所述，MLA 架构通过低秩压缩技术减少了 KV 缓存和激活内存，同时结合 RoPE 处理长序列位置信息，显著提拔了模型的推理效率，同时保持了与标准多头留意力相称的性能。

---

DeepSeek 惊艳背后的架构创新

原创 非子爱 子非 AI 2025 年 01 月 23 日 21:07 美国
DeepSeek v3 震撼发布，仅用十分之一计算量就到达比肩 Llama 3 405B 的性能！其法门在于三大架构创新：多头潜留意力 (MLA) 大幅降低长文本推理成本，混淆专家模型 (MoE) 创新办理了路由崩溃困难，多令牌预测显著提拔推理速率。DeepSeek 团队对 Transformer 的深刻明白和精妙计划，为 AI 领域树立了新的标杆。
小计算量，大智慧：DeepSeek v3 的惊艳表态

你是否曾经由于大模型推理成本过高而望而却步？你是否曾经由于处理长文本而感到力不从心？现在，DeepSeek v3 来了！它以革命性的架构创新，将长文本推理成本暴降，算力需求狂砍 90%！
DeepSeek 近期发布的 DeepSeek v3 模型，在开源权重模型中，以其杰出的基准测试性能脱颖而出，可与当前开始进的模型相媲美。更令人惊叹的是，DeepSeek v3 仅用了约 280 万 H800 小时的训练硬件时间，就实现了这一领先性能。这相称于约 4e24 FLOP 的计算量（假设 MFU，即模型 FLOP 利用率为 40%），与性能相近的 Llama 3 405B 相比，训练计算量足足减少了约十倍！
这一突破性的进展，不仅彰显了 DeepSeek 团队强盛的技术实力，也为 AI 领域的发展带来了新的启示：通过奇妙的架构计划，可以大幅提拔模型的效率和性能，降低 AI 应用的门槛。DeepSeek v3 究竟是怎样做到的？让我们一探究竟！
架构揭秘：DeepSeek v3 的三大创新利器

图 1：DeepSeek v3 架构概览图，展示了其两大核心改进：DeepSeekMoE 和多头潜留意力 (MLA)。图中未显示多令牌预测部门。
DeepSeek v3 之所以能够以小博大，关键在于其三大架构创新：多头潜留意力 (MLA)、混淆专家模型 (MoE) 的改进以及多令牌预测。这三大创新分别针对 Transformer 架构中的不同瓶颈，实现了性能和效率的双重提拔。
1. 多头潜留意力 (MLA)：突破长文本推理的性能瓶颈

• 什么是 KV 缓存？它为什么紧张？
  

想象一下，你在读一本很长的小说，为了明白后面的情节，你必要记住前面的人物关系和事件发展。Transformer 模型也是一样，在举行推理时，为了明白当前输入与历史信息之间的关系，必要访问所有历史信息。为了避免重复计算，模型会将历史信息中的关键信息（键和值向量）存储起来，这就是所谓的 KV 缓存。
KV 缓存的巨细直接影响了模型的推理速率和内存斲丧，尤其是在处理长文本时，KV 缓存的开销会变得非常巨大。以 GPT-3 为例，它有 96 个留意力头，每个留意力头的维度是 128，并且有 96 个 Transformer 块。这意味着对于每个令牌，我们必要一个 2.36M 参数的 KV 缓存，或者在每个 KV 缓存参数精度为 2 字节时为 4.7 MB。可以想象，当处理长文本时，这个缓存会变得多么庞大！

• 传统方法的局限：分组查询留意力等方法怎样捐躯模型质量换取性能？
  

为了减少 KV 缓存的开销，业界提出了多种方法，例如分组查询留意力 (Grouped-Query Attention, GQA) 和多查询留意力 (Multi-Query Attention, MQA)。这些方法通过将多个查询头分组共享同一个键和值头，从而减少 KV 缓存的巨细。例如，Llama 3 70B 和 Mistral Large 2 等模型就接纳了 GQA，将 KV 缓存巨细减少了大约一个数量级。然而，这些方法本质上是一种 “妥协”，它们通过捐躯模型质量来换取性能提拔。就好比为了减轻行李负担，你不得不丢掉一些大概有用的物品。

• MLA 的突破：怎样在不捐躯质量的前提下，大幅减少 KV 缓存巨细？
  

DeepSeek 提出的多头潜留意力 (MLA) 则另辟蹊径，它在不捐躯模型质量的前提下，大幅减少了 KV 缓存的巨细。MLA 的核心头脑是将键和值向量的计算分解成两个步骤，并在推理过程中只缓存中心的 “潜向量”，而不是完备的键和值向量。

• MLA 的核心原理：低秩矩阵与潜向量的奇妙应用
  

图 2：MLA 原理图。MLA 将键和值向量的计算分解成两个步骤，引入了潜变量。
如果把计算键和值向量比作从原材料加工成产物，传统的方法是直接举行加工。而 MLA 则引入了一个 “中心产物”—— 潜向量。在传统的 Transformer 中，键和值向量是通过将残差流向量与一个形状为 (留意力头数 * 留意力头维度) x (模型维度) 的矩阵相乘得到的。MLA 则将这个矩阵分解成两个较小的矩阵：一个形状为 (潜变量维度) x (模型维度) 的矩阵，和一个形状为 (留意力头数 * 留意力头维度) x (潜变量维度) 的矩阵。
在推理过程中，只必要缓存形状为 (潜变量维度) x (序列长度) 的潜向量，而不是完备的键和值向量。这就好比，我们只必要存储 “中心产物”，而不必要存储最终产物，大大节省了存储空间。通过减小潜变量维度，就可以大幅减少 KV 缓存的巨细。

• 低秩压缩的优势：信息共享与差异化表达的完美结合
  

MLA 之所以有效，是由于它利用了不同留意力头之间信息重叠的特性。多个留意力头必要的信息每每有许多相似之处，低秩压缩则提取出这些共同的信息，并允许不同的留意力头以不同的方式利用这些信息。这就好比，多个厨师可以用相同的食材做出不同的菜肴。这不仅减少了冗余，甚至大概对训练产生有益的正则化效果。
DeepSeek 团队的 “好咀嚼” 在 MLA 这里体现得淋漓尽致。他们没有盲目寻求新技术，而是深入明白了留意力机制的本质，奇妙地利用了低秩压缩和潜向量，在不捐躯模型质量的前提下，大幅提拔了长文本推理的效率。
2. 混淆专家模型 (MoE) 的进化：告别路由崩溃，拥抱高效训练

• MoE 的基本原理：怎样通过专家分工提拔模型能力？
  

混淆专家模型 (Mixture of Experts, MoE) 是一种通过将 Transformer 的前馈块替换为多个 “专家” 网络来提拔模型能力的方法。MoE 模型引入了一个路由机制，该机制根据当前的输入，将每个令牌动态地路由到少数几个专家举行处理。这种 “专家分工” 的模式，使得模型可以用更少的计算量，学习到更多的知识。

• 路由崩溃困难：MoE 训练中常见的 “拦路虎”
  

MoE 模型在训练过程中常常会遇到 “路由崩溃” 的标题，即模型总是倾向于将所有令牌都路由到少数几个专家，导致其他专家无法得到充分的训练。这就好比一个团队里，总是少数几个人承担了所有工作，而其他人则无事可做，长此以往，团队的团体效率就会降落。
这是由于梯度降落算法总是倾向于沿着当前最优的方向前进，一旦少数几个专家体现出优势，就会形成正反馈循环，导致其他专家逐渐被 “荒凉”。

• DeepSeek v3 的办理方案：多管齐下，攻克路由崩溃
  

DeepSeek v3 针对路由崩溃标题，提出了两种创新的办理方案：无辅助损失的负载均衡和共享专家机制。

• 告别辅助损失：利用动态偏置项实现负载均衡
  

传统的办理方法是在训练损失中添加一个辅助损失项，用于权衡专家路由的均衡程度。但 DeepSeek v3 的研究表明，这种 “胡萝卜加大棒” 式的强制均衡会侵害模型的性能。他们提出了一种更优雅的办理方案：利用动态偏置项。
详细来说，每个专家都有一个可学习的偏置项，该偏置项会被加到专家亲和力上。在训练过程中，如果某个专家被激活的次数过少，则会逐渐增长其偏置项，使其更容易被选中。这就好比，给每个专家设定一个 “底子人气值”，如果某个专家的人气过低，就轻微提拔一下，使其有更多的机会被选中。这种方法避免了辅助损失带来的性能损失，同时又能有效地办理路由崩溃标题。

• 共享专家机制：打破均衡路由的束缚，实现更灵活的知识分配
  

DeepSeek v3 进一步提出了共享专家机制，将专家分为共享专家和路由专家。共享专家始终被激活，负责处理通用知识，而路由专家则根据必要举行选择，负责处理专业知识。这种机制打破了强制均衡路由的束缚，使得 MoE 模型可以更有效地利用不同专家的专业知识，进步了模型的团体性能。
这就好比一个团队里，既有负责日常工作的 “常驻成员”(共享专家)，也有根据项目必要随时调配的 “特聘专家”(路由专家)。如许，既能包管日常工作的稳定举行，也能灵活应对各种专业使命。

• 共享专家的优势：通用知识与专业知识的公道分配
  

共享专家机制的优势在于，它可以更公道地分配通用知识和专业知识。例如，一个语言模型必要频仍利用语法和常用词汇等通用知识，而对特定领域的专业知识则访问较少。共享专家可以专门负责处理这些通用知识，而路由专家则可以根据必要处理不同的专业知识。
3. 多令牌预测：加快推理，效率翻倍

• 多令牌预测的机制：一次前向传递，预测多个令牌
  

图 3：多令牌预测机制图。DeepSeek v3 在预测下一个令牌之后，还会将最终的残差流向量输入到另一个 Transformer 块中，并利用该块的输出预测第二个令牌。
DeepSeek v3 还引入了多令牌预测机制，即在每个前向传递中预测多个令牌，而不是像传统的 Transformer 那样只预测下一个令牌。详细来说，DeepSeek v3 在预测下一个令牌之后，还会将最终的残差流向量输入到另一个 Transformer 块中，并利用该块的输出预测第二个令牌。在训练过程中，DeepSeek v3 会将这两个令牌的预测损失加权求和，作为最终的训练损失。

• 多令牌预测的优势：提拔训练效率，实现推测性解码
  

多令牌预测机制有两个主要的优势：一是提拔了训练效率，由于模型可以在一次前向传递中学习到更多信息；二是实现了推测性解码，从而可以显著进步推理速率。

• 推测性解码：怎样利用多令牌预测加快推理过程？
  

推测性解码是一种利用小模型生成多个候选令牌，然后用大模型举行验证的推理加快方法。DeepSeek v3 的多令牌预测机制天然支持推测性解码。在推理过程中，DeepSeek v3 可以一次生成两个令牌，然后用自身举行验证。根据 DeepSeek v3 的技术报告，第二个令牌的接受率在 85% 到 90% 之间，这意味着推理速率几乎可以翻倍！
这就好比，你写文章时，不是一个字一个字地写，而是一次写两个字，然后查抄一下这两个字是否通顺，如果通顺就保留，不通顺就修改。如许，你的写作速率就会大大进步！
创新背后的思索：DeepSeek 团队的 “好咀嚼”

DeepSeek v3 的这些创新并非简单的暴力搜刮，而是基于对 Transformer 架构的深刻明白。DeepSeek 团队展现了他们对 AI 研究的 “好咀嚼”：他们没有盲目寻求新技术，而是深入明白了 Transformer 架构的本质，奇妙地利用了已有的技术本领，办理了实际标题。
例如，在 MLA 中，他们利用了低秩压缩和潜向量；在 MoE 的改进中，他们利用了动态偏置项和共享专家机制；在多令牌预测中，他们利用了推测性解码的头脑。这些方法都不是全新的，但 DeepSeek 团队将它们奇妙地结合在一起，发挥出了巨大的威力。
未来展望：计算优先级的探索

DeepSeek v3 的乐成，为未来的 AI 研究指明了方向。一个值得探索的方向是计算优先级的探索。目前，Transformer 模型对每个令牌的计算量都是相同的，这大概存在效率标题。例如，模型在处理一个简单的令牌和一个复杂的令牌时，利用的计算量是相同的。
未来的改进方向大概是根据预测的难易程度，动态调整计算量。例如，可以利用基于留意力分数的动态路由，将更多的计算资源分配给更紧张的令牌；或者可以利用早期退出机制，在模型对某个令牌的预测已经足够自信时，提前结束计算。这就好比，考试的时间，我们应该把更多的时间花在困难上，而不是在简单的标题上浪费时间。

---

【100 个 AI 核心概念】MLA

原创 道上 AI 道上 2025 年 01 月 17 日 19:23 浙江
一、留意力机制

在先容多头潜在留意力之前，必要先了解下留意力机制。在天然语言处理（NLP）和计算机视觉等领域，处理信息时通常会遇到大量的数据，比如，在翻译一段话时，模型必要明白每个单词的意思，以及它们之间的关系。传统的模型在处理信息时，每每是以固定的方式来考虑所有输入。这种方法有时会导致模型忽视紧张的信息。
留意力机制正是为了办理这个标题，它的核心头脑是让模型在处理信息时，能够动态地选择关注的部门。通过计算不同部门的紧张性，模型可以更好地明白数据。例如，在翻译句子时，模型可以更关注当前翻译的单词和源语言中相关的单词。
二、留意力机制的工作原理

留意力机制的工作原理可以用一个简单的例子来说明。假设我们有一句话：“小猫在阳光下睡觉。” 当我们想要翻译这句话时，模型必要关注 “猫” 和 “睡觉” 这两个词。留意力机制通过计算每个词的紧张性，来决定模型在生成翻译时应该关注哪些词。
留意力机制通常会生成一个 “留意力权重”，这个权重体现了每个输入部门的紧张性。权重越高，体现这个部门对当前使命越紧张。模型根据这些权重来加权输入信息，从而生成更准确的输出。
三、多头留意力的概念

在了解了留意力机制后，我们再来看看多头留意力。多头留意力是对单一留意力机制的扩展。它的核心头脑是同时利用多个留意力机制，来捕捉输入信息的不同方面。
假设我们在翻译一段复杂的句子。单一的留意力机制大概只能关注到句子中的某一个方面，比如语法布局。而多头留意力则可以同时关注到多个方面，比如语法、语义和上下文等。如许，模型能够更全面地明白输入信息，从而生成更准确的输出。
四、多头潜在留意力（MLA）的定义

多头潜在留意力（MLA，Multi-Head Latent Attention）是一种将多头留意力机制与潜在空间概念结合的先进模型架构。为了更好地明白 MLA，我们必要先明确几个关键概念：多头留意力、潜在空间及其在深度学习中的应用。
1. 多头留意力机制

多头留意力机制是 Transformer 架构的核心组件之一。其基本头脑是，通过多个 “留意力头” 并行计算输入序列中不同部门之间的相关性，从而使模型能够关注输入的不同特性。详细来说，每个留意力头会学习到不同的权重矩阵，从而对输入数据的不同方面举行建模。这种机制的优势在于，它允许模型在处理信息时考虑多种上下文信息，从而进步了模型的表达能力和鲁棒性。
在标准的留意力机制中，给定一个输入序列，模型管帐算查询（Query）、键（Key）和值（Value）之间的相似度。每个留意力头通过不同的线性变更生成这些查询、键和值，并通过加权求和的方式得到输出。最终，所有留意力头的输出会被拼接并经过线性变更，形成最终的输出结果。
2. 潜在空间

潜在空间（Latent Space）是指在呆板学习模型中，数据通过某种映射被转化为一个高维空间的体现。在这个空间中，数据的特性和布局可以被更好地捕捉和明白。潜在空间的构建通常依靠于降维技术或嵌入方法，如主身分分析（PCA）、自编码器、词嵌入（如 Word2Vec 和 GloVe）等。
在潜在空间中，数据的相似性可以通过距离度量来体现。相似的样本在潜在空间中会被映射到彼此靠近的位置，而不同的样本则会被映射到较远的位置。这种体现方式使得模型能够更有效地捕捉到数据的内涵布局和特性。
3. MLA 的工作原理

多头潜在留意力（MLA）结合了以上两个概念，首先将输入数据映射到潜在空间，然后利用多个留意力头来计算输入的不同部门在潜在空间中的紧张性。详细步骤如下：

• 1. 输入映射：输入数据首先通过嵌入层（如词嵌入）被映射到一个高维的潜在空间。这个过程可以看作是对输入数据的特性提取。
  
• 2. 留意力计算：在潜在空间中，模型为每个留意力头计算留意力权重。每个头利用独立的权重矩阵来处理输入数据，捕捉不同的特性和上下文信息。
  
• 3. 加权求和：计算出的留意力权重会被应用到输入数据上，通过加权求和的方式生成一个更加精确的体现。
  
• 4. 结果合并：最后，所有留意力头的输出结果会被合并，形成最终的输出。这种方式不仅进步了模型的表达能力，也加强了其对复杂模式的捕捉能力。
  

通过这种方式，MLA 能够有效地处理高维数据，并提取出有用的特性，从而在各种使命中体现精彩。
五、MLA 的详细实现

实现多头潜在留意力（MLA）通常涉及多个步骤，从数据预处理到模型训练。以下是 MLA 实现的详细过程：
1. 数据预处理

在构建 MLA 模型之前，首先必要对输入数据举行预处理。这一步骤的目的是将原始数据转换为模型可以明白的格式。预处理包罗以下几个步骤：

• 分词：将文本数据分解为单词或子词单元。这一步骤通常利用天然语言处理（NLP）中的分词工具，如 NLTK 或 spaCy。
  
• 去除停用词：停用词是指在文本中频仍出现但对语义贡献较小的词，如 “的”、“是”、“在” 等。去除停用词可以帮助模型更好地聚焦于有意义的词汇。
  
• 词嵌入：将分词后的文本转换为向量体现。常用的词嵌入技术包罗 Word2Vec、GloVe 和 FastText 等。这些技术能够将词汇映射到一个高维空间中，使得相似的词在向量空间中距离较近。
  

2. 映射到潜在空间

在完成数据预处理后，输入数据将被嵌入到一个高维的潜在空间中。

• 初始化嵌入矩阵：根据词汇表的巨细和嵌入维度初始化词嵌入矩阵。每个词汇会对应一个向量，这些向量会在训练过程中不停更新。
  
• 输入数据转换：将处理后的文本数据转换为嵌入向量。每个词汇会根据其在嵌入矩阵中的位置被替换为对应的向量体现。
  

3. 留意力权重计算

在潜在空间中，模型会为每个留意力头计算留意力权重。步骤如下：

• 生成查询、键和值：每个留意力头会通过独立的线性变更生成查询（Q）、键（K）和值（V）。这些矩阵的维度通常与嵌入维度相同。
  
• 计算留意力分数：通过点积计算查询与键之间的相似度，得到留意力分数。然后，通过 Softmax 函数将这些分数归一化为概率分布，以体现各个输入部门的紧张性。
  
• 加权求和：将计算出的留意力权重应用到值（V）上，通过加权求和的方式生成每个留意力头的输出。这一步骤确保了模型能够根据上下文动态调整关注的输入部门。
  

4. 结果合并

最后，所有留意力头的输出结果会被合并。

• 拼接输出：将每个留意力头的输出在特性维度上举行拼接，形成一个更大的向量体现。
  
• 线性变更：对拼接后的结果举行线性变更，通常通过一个全毗连层将其映射到目的输出的维度。这一步骤可以帮助模型进一步整合信息，并生成最终的输出。
  

5. 模型训练

在 MLA 模型构建完成后，接下来是模型的训练过程。

• 损失函数定义：根据详细使命（如分类、回归等）定义损失函数。常见的损失函数包罗交织熵损失、均方误差等。
  
• 优化算法选择：选择合适的优化算法（如 Adam、SGD 等）来更新模型参数。优化算法的选择会影响模型的收敛速率和最终性能。
  
• 迭代训练：通过多次迭代（Epochs）训练模型。在每次迭代中，利用训练数据计算损失并更新模型参数。训练过程中通常会利用验证集来监控模型性能，防止过拟合。
  

6. 模型评估与应用

在模型训练完成后，必要对模型举行评估。评估的步骤包罗：

• 测试集评估：利用未见过的测试集对模型举行评估，计算模型在各种指标上的体现，如准确率、召回率、F1-score 等。
  
• 超参数调优：根据评估结果调整模型的超参数（如学习率、批量巨细等），以进一步进步模型性能。
  
• 实际应用：将训练好的 MLA 模型应用于实际使命中，如文天职类、感情分析、呆板翻译等。
  

六、MLA 在实际应用中的优势

首先，它能够捕捉输入数据的多种特性。通过利用多个留意力头，模型可以同时关注到不同的信息。这种多样性使得模型在处理复杂使命时体现更佳。
其次，MLA 能够进步模型的灵活性。在处理不同类型的数据时，模型可以根据必要调整留意力头的数量和权重。这种灵活性使得 MLA 能够适应各种应用场景，如呆板翻译、文本生成和感情分析等。
最后，MLA 还能够进步模型的性能。通过在潜在空间中举行计算，模型能够更好地明白数据的布局和关系。这种明白有助于进步模型的预测准确性和生成质量。
七、总结一下

多头潜在留意力的优势在于它能够同时关注输入数据的多个方面。这种能力使得模型在处理复杂使命时体现更佳。通过结合多个留意力头，模型能够更全面地明白数据，从而生成更准确的输出。
此外，MLA 的灵活性和适应性使得它在各种应用场景中都能发挥紧张作用。无论是在呆板翻译、文本生成照旧感情分析中，MLA 都能提供更好的性能和结果。

---

再读 MLA，还有多少细节是你不知道的

原创 猛猿 大猿搬砖简记 2025 年 01 月 20 日 19:11 北京
关于 MLA，我想先简单记载下我了解它的心路历程：

• 我第一次了解 MLA，是在它刚出来的档口**。在我读过它的原生实践后，我发现它既不省 KV cache，还徒增了计算量。这个标题让我困扰了很久，当时网上对 MLA 细节的讨论很少，所以我觉得大概是我没弄懂，所以它就被我放到一边去了。别人问起我 MLA，我都答复不知道，由于我确实没想明确。
  
• 直到近来 dpsk v3 的风刮起来，再加上实际体验后它的推理速率确实很快（当然我知道不止于 MLA），所以又有了重新认识它的爱好。而幸运的是，此时网上已经有关于 MLA 的诸多讨论了，如许就能和我的明白相互校验。在这一次的重新认识中，我发现我之前的认知其实没有错，原生 MLA 确实存在上述的 2 个标题，并且开源社区已有了各种对原生 MLA 的优化方法，这里我主要参考了：
  
• blog：https://zhuanlan.zhihu.com/p/700214123
  
• 代码：https://github.com/madsys-dev/deepseekv2-profile/tree/main
  
• 章明星老师关于 MLA 优化的思绪，本文里先容的 MLA 优化思绪也来自于此，我保举各人都看看这块，写得非常具有启发性 (链接见上，编辑出标题了 sry)
  
• 知乎上一篇对 MLA 计算量的分析，在这个分析思绪的启发下，本文按照自己的明白重新做了 MLA 各种优化方法的计算量和 KV cache 分析：https://zhuanlan.zhihu.com/p/714761319
  

我个人认为，明白 MLA 的难点之一，是它算法计划颇为 “绕”，不管是用数学公式，或者是用流程图，似乎都很难一下找到它计划的核心头脑。所以本文第一部门，将会抛开所有复杂的计算细节，根据我自己的明白，抽象出 MLA 的计划方案，基于此再来谈计算细节和各种优化。
一、MLA 的基本头脑

1.1 MLA，MQA 与 GQA

我们先来快速复习一下 decoder 架构的 MHA 的运作流程，如下图：

这里 head_num = 4，图中刻画了 head0 的运算过程，包罗 attn_weights = Matmul (q, k) 以及 attn_sv = Matmul (attn_weights, v)，图中虚线灰框体现在 head0 上的结果是由包罗其在内的若干前置 tokens 计算结果汇总而来。为了表达轻便，这里省去了诸如 softmax，的计算过程。图中被红色虚线框圈起来的部门，就是各人熟知的将被用在推理阶段的 KV cache。
KV cache 的存在，本来是为了避免在推理阶段对前置序列的重复计算的。但是，随着前置序列的长度变长（我们记为 kv_len），必要读取的 KV cache 也将越来越大，数据的传输成本增长，这就使得 attn 计算逐渐酿成 memory bound。我们采取了一些计谋来缓解 KV cache 过大的标题，此中 2 种就是各人熟知的 MQA 和 GQA。
MQA 和 GQA 的运作方式如下：

• 在 MQA 的情况下，一个 token 所有的 heads 都共享同一个 k 和 v。如许在降低 param weights 巨细的同时，还让本来必要生存 num_heads 份的 kv cache 降低为只需生存 1 份。
  
• 但是，MQA 大概造成模型效果上的损失，毕竟原来对于 1 个 token，它的每个 head 都有各自的 k、v 信息的，现在却被压缩为一份。所以 GQA 作为一种折衷的方案出现了，即将 1 个 token 的 head 分成 num_group 组，每个 group 内共享同一个 k，v 信息，使得信息压缩不像 GQA 那样严重。
  

但是，不管是 MQA 照旧 GQA，对于 1 个 token 来说，总是存在 heads 上 k、v 信息被压缩的情况。那么是否有一种办法，能在尽量不压缩 head 上 k，v 信息的情况下，节省 kv cache，进步团体推理速率呢？ 那么接下来，我们就来大致看一下 MLA 的计划头脑。
1.2 MLA 的团体计划头脑

在本节中，我们会以 K cache 为例，抽象出 MLA 的核心优化头脑。V cache 的优化头脑也是同理，但不在这节赘述，而是合并到后文对于 MLA 的细节解说中（参见 2.3 节 A_CC）。
现在先让我们回到 MHA 上（图 1.1），来思索一个标题：为什么对于一个 token，我们必要生存它所有 heads 上的 K 值作为 K cache 呢？

主要缘故原由我们在上文解释过：这是由于每个 k_head 附带有不同的信息，它将用这份独有的信息和对应的 q_head 举行 attn 的计算，用公式体现即为，这里的是合并了所有 head 对应的 param weight 后的表达。
我们现在的总目的是节省 K cache，当你再次端详上面这幅图时，一个 idea 在你的头脑中出现：

• 当前我要存的 K cache 是 4 个 k_head（图中深绿色框），但如果我能从这 4 个 k_head 中抽取出 1 份共有的信息，然后在做 attn 计算时，每个 head 都用这 1 份共有的信息做计算，那么我也只需存这 1 份共有信息作为 K cache 了。如许我就把 K cache 从原来 num_heads = 4 酿成 num_heads = 1，这不就能节省 K cache 了吗？
  
• 但是等等，现在共有的 k_head 信息是抽取出来了，那么相异的 k_head 信息呢？（简单来说，就是由不同 head 部门学习到的相异信息）。我们当然是渴望 k_head 间相异的信息也能保留下来，那么该把它们保留至哪里呢？当你回顾 attn_weights 的计算公式时，一个想法在你脑中闪现：q 部门不是也有 heads 吗！我可以把每个 k_head 独有的信息转移到对应的 q_head 上吗！写成公式解释就是：
  
• 原来，括号体现运算次序，即先各自算 2 个括号内的，再做 * 计算
  
• 现在 ，同理括号体现运算次序。
  
• 也就是说，这里我们通过矩阵乘法的交换律，奇妙地把 1 个 token 上 k_heads 独有的信息转移到了对应的 q_head 上来，如许 1 个 token 上 k_heads 间共享的相同信息就能被我们当作 K cache 存储下来。
  

（在这里，你可以抽象地把明白成是 4 个 k_heads 共享的信息，但最终 K cache 的形式还会在这底子上有所变化。我知道此时你脑海中一定有许多迷惑。但我们先不要纠结细节的标题，由于在后文会展示全部细节，这里我们要做的是从宏观上明白 MLA 计划的核心头脑。）
现在我们更详细地画出上面这套 “信息转移” 方案的详细流程：

⚠️：再次说明，在本部门，我们偏重于抽象出 MLA 的优化思绪，各人在阅读上面这幅图时，请不要带入任何详细的细节（例如矩阵尺寸）等去做计算，这部门细节我们会在下文详细先容。
我们来详细看这幅图：

• 对于每个 token 的 k_heads，我们必要抽取出它们的相异信息，而这个相异信息本质上是由维护的。观测到所有 tokens 都共享 1 个，所以我们对于 q_heads，我们只需做 1 次对于的吸取，就能统一获取所有 tokens 的所有 k_heads 上的相异信息。
  
• 对于每个 tokens 的 k_heads，我们还必要抽取出它们的相同信息，而这个相同信息应该是每个 tokens 的所有 k_heads 共享一份，同时不在不同 tokens 间共享。那么我们天然而然想到，可以学习一个 linear 参数矩阵，从原始 token 中提取出这份共有信息，以此作为我们的 K cache。而不管是从 “信息提取” 照旧从 “进一步节省 K cache 巨细” 的角度来说，似乎这个 linear 参数参数矩阵如果能把压缩到一个更低维的空间，会收获更精密的信息表达和更小的存储量，这也是图中 compress_k 的由来。
  
• 最后，我们利用压缩后了共有信息的 compress_k，和吸取了相异信息的 q_head 做计算，得到 attn_weights**。
  

对 v cache 的优化也是同理，这里额外提几点：

• 究竟上，当我们考虑到 v cache 优化时，上图中的 compress_k 其实应该被明白成 compress_kv，也就是它是 1 个 token 所有 k_heads 和 v_heads 的共有信息。
  
• 可以和作吸取，我们在后文会讲这块细节。
  
• 总结起来，我们应该尽量直接利用 compress_kv，而不要将其用或者举行展开，由于 q 已经吸取过这两者的信息了。**
  

好，到此为止，我们抽象出了 MLA 的团体优化思绪，从中你可以发现：

• 固然从形式上来说，MLA 和 MQA/GQA 很像，似乎都是通过压缩 k/v_heads 的数量来节省 KV cache 巨细的。但 MLA 是压缩 num_heads，不压缩信息（把信息转移到了 q_heads 上）；而 MQA/GQA 则在一定程度上对信息做了压缩。
  
• 从这一点上看，我个人认为 MLA 其实更应该明白成是 MHA 的变种优化，而并非 MQA/GQA 的变种优化。**
  

二、MLA 的运作流程

2.1 CD (CacheDecompressed, dpsk MLA 的原生实现）

现在我们可以来看 MLA 的运作细节了。

本图参考自 huggingface 上 dpsk v2 的 MLA 原生实现：
https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/main/modeling_deepseek.py#L682
关于 MLA 的更多细节，例如 “为什么要区分 nope 和 rope 维度”，“为什么要做 low rank 压缩” 等，已经有前人的分析珠玉在前，因此本文不做重点阐述，更多细节请参考：

• 缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA - 科学空间|Scientific Spaces By 苏剑林 | 2024-05-13
  https://kexue.fm/archives/10091
  
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
  https://arxiv.org/pdf/2405.04434
  

这里假设 q_len = 1，kv_len = 1024，nope 体现非 pe 部门的 head_dim，rope 体现 pe 部门的 head_dim。别的维度已标注在图中。此中红色体现 param_weights，此中：

• q_b_proj：是 q 计算中的升维矩阵，它包罗了两部门，分别体现对 q 的 nope/rope 部门的计算。
  
• kv_a_proj_with_mqa：是对原始 hidden_states 的压缩矩阵，它包罗了两部门，分别用于计算 compress_kv（即抽取 k_heads 和 v_heads 的共同信息），以及计算 k_pe 的部门。
  
• kv_b_proj：它包罗了两部门，分别体现对 k_nope 和 v 部门的计算。
  

以上符号体现皆服从 dpsk 原始论文，下标体现 Down 降维，体现 Up 升维，体现做 Rope（诸如就体现和 K 的 rope 相关）。
好，现在关于这个 MLA 的原生实现，我们来讨论几个有意思的点：
（1）在 MLA 中，每个 head_dim 的尺寸更大了。观察到原始 hidden_size = 5120，如果按照 num_heads = 128 来看的话，正常来说一个 head_dim = 40 (5120/128=40)。但是在 MLA 中，一个 head_dim = 128，宏大于 40。也就说 MLA 其实是用比一样平常 MHA 更大的 head_dim（或者也大概是 num_heads）来做 attn 计算的，然后在最终的矩阵中映射回原来的 hidden_size。对此我个人给出一些简单推测：如果推理阶段 KV cache 造成的 memory bound 的标题已经得到办理的话，那么训练时我就能少一点后顾之忧，然后通过提拔模型的复杂度来取得与 MHA 比肩或更好的效果（训练阶段还有别的优化方式）。如许当我回到推理阶段时，我的团体计算强度就上去了（每读 1 次，算的次数更多了）只要没有到达 compute bound 的界限，如许的提拔就是有利益的。
（2）原生 MLA 的计算最终展开成了 MHA 的计算。这一点可以参见图中 q（蓝色），k（绿色），v（黄色），它们最终都酿成了标准 MHA 的计算。从理论上来说，这一点也不奇怪，由于我们在第一部门说过 MLA 就是 MHA 的变种，只是它在 MHA 的底子上做了信息从 k/v_head 向 q_head 的转移。嗯？！！但是等等，从上图这个原生 MLA 上来看，固然产出了 compress_kv，但是好像并没有做什么信息转移，也就是粗糙来看目前的计算流程照旧而不是转移后的：

• 是的，如果你有这个迷惑，那么恭喜你发现了原生 MLA 的标题，也就是它没有做任何的信息转移**。
  
• 同时，原生 MLA 生存的 KV cache 并不是图中绘制的 compress_kv，而是图中已经成形的完备的 k（绿色）和 v（黄色），这一点在上面的代码中可以看见。
  
• 再有，考虑到这里 head_dim = 128（宏大于同 num_heads 数量下的标准 head_dim=40），所以原生 MLA 增长算力所付出的代价是，KV cache 显存反而增长了。
  

基于这些，我们管原生 MLA 的实现方式为 CD（CacheDecompressed），即存储的 KV cache 是没有经过任何压缩的。为什么 dpsk 放出来的原生 MLA 会如许呢？这一点我一直没有想通，这也是为什么我在 MLA 刚出来那阵，看完它的实践就决定先停息探索的缘故原由。当时没有实际的业务需求，自己也没动力去细想，以及考虑到 MLA 算法的复杂性，我还以为是我明白错了。但是随着时间推移，后续开源社区有一系列对 MLA 的优化实现，直到近期再次捡起来后，才使我对 MLA 有了更多的了解。目前来看，这个原生 MLA 似乎以提供 “MLA 的概念” 为主，而详细的优化实践方式照旧要看个人。我们立刻就来看后一些做过 “信息转移 / 吸取” 的优化方法，不过在此之前，我们先对原生 MLA 的计算量和 KV cache 做一个分析。
(公众号编辑表格太难了，这里我直接从我笔记截图了，各人可以点开放大看)

我们对这张表格做一些说明：

• 这张表格形貌的是推理阶段的 Attn 部门相关的计算量 (MFLOPs, Million FLOPs) 和 KV Cache 巨细 (per layer per token, kB)。后面这类型的表格都是同理。**
  
• 首先，再次声明，在 CD 方法下，我们存储的是图中 MHA 形态的 K，V 值（绿色，黄色框）。所以 KV 相关的计算量都是从这两个框之后开始的，各人可以对着图看。
  
• 表格的前 3 行，是每个做 query 的 token 必须要实行的操作。此中，对于 kv_a_proj_with_mqa 这个操作，是当前做 query 的 token 也会产出自己的 KV 值。
  
• MFLOPs/per_token**：这个指标更准确的说是 MFLOPs/per_layer/per_token，权衡的是单 token 的计算量。我们知道 1 个 token 有分 q 部门和 kv 部门，那么表格中除以 1 的就可以明白为是 q 部门相关的计算量，除以 1024 的可以明白成是 kv 部门相关的计算量。通过这种方式，可以帮助我们更好明白计算量是怎样在 q 和 kv 之间转移的，我们在后文会更好感受这个 “转移”
  
• 我们接下来会对比除 CD 外的 MLA 优化方法，这里提前预告下，在这些比力中，你会发现 CD 的计算量是最小的（不难明白，由于它相称于是从流程图最靠近下方的部门开始计算的），但是它的 KV cache 是最大的**。留意，这里我们说的计算量小，是 CD 相比于别的 MLA 优化方法，而不是 CD 相比于 head_dim=40 的标准 MHA。
  
• 最后总结一下，在 CD 中，单 token KV cache = 80 KB ，单 token KV 相关的计算量为 0.08 MFLOPs。**
  

2.2 CC (CacheCompressed）

好，在进入各人从第一部门开始就心心念念的 “k/v_head 信息向 q 转移（或者明白成被 q 吸取）”这个优化先容前，我们先先容基于原生实践和这个优化的一个中心态：CC (CacheCompressed）。在这个中心态中，我们终于是以 compress_kv 为 kv cache 了，但是我们没做任何吸取。之所以要先容这个中心态，是方便各人更好感受 “吸取” 的利益。
我们直接对着 2.1 的图，列出 CC 表格：

不难发现，在这个中心态 CC 优化的 MLA 下：

• 单 token KV cache = 1.13 KB ，相比 CD 有了显著降低。
  
• 单 token 的 kv 计算量 = 33.55 + 0.05 + 0.03 = 33.63 MFLOPs。主要犯罪怀疑人就在 kv_b_proj 上。简单来说，在没有做吸取 / 转移前，一个矩阵必要作用在 kv_len = 1024 条数据上，但是现在它只必要被 q_len=1 条数据算 1 次就好了，即我们把属于 kv 的计算量转移到了 q 上。
  

2.3 A_CC（AbsorbCacheCompressed）

现在，终于来到我们心心念念的涉及吸取的优化了：

• 单 token KV cache = 1.13 KB
  
• 单 token 的 KV 计算量 = 0.15 + 0.13 = 0.25 MFLOPs
  
• 到达了节省 KV cache 的同时，维持单 token KV 计算量不变的需求。
  

这里解释下为什么 A_CC 相比于 CC，总计算量降低了许多，但单 token 计算量却没有变化：

• 这是由于单 token 计算量分成作用在 q 和作用在 kv 上的。而 q 对应的 seq_len = 1，kv 对应的 seq_len=1024
  
• A_CC 相比于 CC，把原来属于单 kv 的计算量转移到 q 上了，而 q 的 seq_len=1，对总计算量的影响本来就少。
  

2.4 A_CC_ME

最后，这个优化其实就是在 A_CC 的底子上，在计算 attn_weights 的时间，把 nope 和 rope 的部门拆开算，然后再求和。如许做是为了避开无用的数据拷贝和广播（可以看代码，你会发现 A_CC 为了做数据拼接，是先初始化一个一个拼接好的空张量，再往里塞数据，如许就是 2 倍的显存开销。而避开拼接各自算各自的，可以直接复用已有的数据），实际测起来这种方法性能是最好的。

三、MLA 可以用 TP 吗

现在，回来看一个常常被讨论的标题：MLA 可以做 TP 吗？由于看样子，对于每一个 token 来说，它所有的 num_heads 上的 kv 信息已经被压缩成 compress_kv 了，好像是不能再切分了？
这里先说结论：MLA 可以做 TP，但是它大概必要一些定制化的 TP 方式，而不是直接套用惯常 decoder 模型的 TP 方式。
为了解答这个标题，我们这里再贴出 2.1 中的流程图：

我们偏重关注流程图中红色部门（也就是 param_weights），各人追念一下之前的先容：尽管 compress_kv 已经被抽取成雷同单头的形式了（1 个 token 只有 1 个，且不区分 heads），但是它能如许做的缘故原由是由于 kv_heads 上的信息转移去 q_heads 了，对了！q 照旧有 heads 的！
我们首先来看一下，dpsk 官方是怎样在上面这张流程图中做 TP 切分的，详细代码可以参见这里：
https://github.com/deepseek-ai/DeepSeek-V3/blob/ee4c4ea32bfd89e197616b80e713466954c51c75/inference/model.py#L409
从图里来说：

• 在 TP 的情况下，一个 tp_group 内的每个 tp_rank 都共享相同的 compress_kv。（这里有一个优化方案，那就是可以通过雷同 cache pool 的方式，只维护一份 compress_kv，然后每个 tp_rank 都从上面去取，就能避免数据的重复存储了）
  
• 在 dpsk 官方给出的 tp 示例中，图中所有带 num_heads 的红色矩阵是按照 num_heads 做 TP 分别的，这一点不难明白，由于 q 依然是按照 heads 分别的，而 kv 的信息又是被吸取到各个 q_heads 上的
  
• 而对于 q_a_proj 和 kv_a_proj_with_mqa，dpsk 官方是不切的，即各个 tp_rank 上重复维护同样的 param_weights。但这不意味着这两个 param_weights 不能再切分，例如实践上，你可以用 col/row 对他们做切分，只要包管在输入输出的位置做好正确的通讯即可。详细的方式则要看实际应用场景。
  

修改于 2025 年 01 月 20 日

---

via:

• DeepSeek 的 MLA 架构
  https://mp.weixin.qq.com/s/Kgo1Tn-B9LWcXrwGbvvc3w
  
• DeepSeek 惊艳背后的架构创新
  https://mp.weixin.qq.com/s/fIvUFHbEM1v4nwPYe3fxmA
  
• 【100 个 AI 核心概念】MLA
  https://mp.weixin.qq.com/s/Duc5ADHIGrUjHdmO4ylffg
  
• 再读 MLA，还有多少细节是你不知道的
  https://mp.weixin.qq.com/s/E7NwwMYw14FRT6OKzuVXFA
  

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。