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'=KW_{K}^{down}$

$V'=VW_{V}^{down}$

其中，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 日