Kimi K3 技术报告揭秘 AI Infra 核心原则：KDA、AttnRes、MoonEP 技巧大起底

为什么 Infra 重要？

前面提到，OpenAI 核心工程师翁家翌说大模型训练中工程师主要给 Infra 修 bug，这里的 bug 指什么呢？在大模型团队日常语境中，Infra bug 既包括导致结果错误、训练不稳定或任务中断的传统 bug，也包括虽不影响数学正确性，但会造成 GPU 等待、内存浪费、通信阻塞和吞吐下降的系统性问题。具体而言，要在保证计算逻辑正确的前提下，看 GPU 是充分利用还是有大量闲置算力，即便 GPU 充分利用，还要看是否存在大量不必要的重复计算，以及 GPU 上的计算负载是否均衡、稳定等。满足这些要求需要大量通信在 GPU 之间做负载协调，而通信又可能遇上新瓶颈。这些 bug 最终影响集群吞吐量，反映训练效率。

模型浮点运算利用率 (MFU) 是评估训练效率的标准指标，它是“观测吞吐量”与“理论最大吞吐量”的比值，理论最大吞吐量假设达到峰值浮点运算的 100%。从行业角度看，大模型训练效率仍有很大提升空间。目前公开可验证的万卡、千亿参数、完整 LLM 预训练案例中，字节跳动的 MegaScale 的 55.2% MFU 仍是代表性的最高纪录之一。今年 4 月还出现极端反面案例，据《The Information》报道，马斯克旗下 xAI 坐拥约 55 万张英伟达 GPU，但内部备忘录显示其 MFU 仅 11%，远低于业界 35%-45% 的正常水准，也不及 Meta（43%）、谷歌（46%）等竞争对手。

拥有更高的训练效率意义重大，它可直接转化为更短的训练周期、更低的训练成本、更多可尝试的实验、更大的数据量或模型规模。虽然本文要展现 Infra 设计的重要性，但因缺乏足够数据和参数参考，很难自主判断技术报告中哪些 Infra 技术最重要或贡献最大。为此，本文主要从 Kimi K3 自主提出的架构设计入手，如 KDA、AttnRes、MoonEP，看看其背后的 Infra 技巧体现的原则，让大家从逻辑上感受，选择原则不代表实际重要性排序。

总体设计

先来看看 Kimi K3 的分布式系统设计，其训练采用的并行方法都是比较经典的，如流水线并行、专家并行、ZeRO 分片等。大模型训练的“并行化”简单理解就是，大模型训练过程包含大量串行过程，从前向传播到反向传播，但细节上存在大量矩阵计算，矩阵计算一般可并行化。此外，一张卡的内存很难装下超大参数规模的所有数值。不严谨地说，就像盖房子要从底往上盖（串行过程），但盖房子时可以同时搬很多水泥和钢筋（矩阵计算的并行）。

为此，大模型训练通常从模型参数、训练数据和中间激活值等不同维度进行切分，将计算和存储任务分配到多张 GPU 上，并通过合理的并行策略和调度机制，尽量减少 GPU 等待和空转。例如，把模型的不同神经网络层放到不同 GPU 上依次计算，这叫流水线并行。由于后面的 GPU 要等待前面的 GPU 产生中间结果，流水线中容易出现部分 GPU 空闲的“气泡”。实际训练中通常把一个 batch 切成多个微批次，让不同 GPU 同时处理不同微批次，从而提高流水线利用率。

如果模型的某一层规模很大，单张 GPU 无法容纳，或单卡计算速度不足，还可把这一层中的权重矩阵和矩阵运算拆分到多张 GPU 上共同完成，这叫张量并行。张量并行会在同一层计算过程中频繁交换数据，因此对 GPU 之间的互联带宽和通信延迟要求很高。当通信耗时接近甚至超过计算耗时时，就需要在并行规模、计算效率和通信开销之间进行权衡。

以上只是原理科普，不同的大模型集群有自己的问题。特别是模型越大，越可能出现小模型中不会出现的问题，因为某个原先不被注意的计算进程或内存存储对象会因规模增大，而对 GPU 容量而言变得显著。

技术报告显示：“MoE 层使用在各 EP rank 之间复制的共享专家”，这意味着在专家并行（EP）设置下，每个 GPU 都保存一份参数完全相同的共享专家，以及不同的路由专家，这样可以让共享专家尽可能吸收数据中的共有特征（rank 可理解为一个参与分布式通信的计算进程，实际部署中通常一个 rank 对应一张 GPU，为简单起见本文把 rank 都理解为单个 GPU）。

接下来，看看 Kimi K3 的模型架构核心机制 KDA 和 AttnRes，在匹配的 Infra 设计上有什么巧思，以及 Kimi K3 为优化 MoE 架构的专家并行特有的难题提出的 Infra 方法 MoonEP，体现了什么原则。

KDA：混合 KDA 注意力机制

KDA 全称 Kimi Delta Attention，是一种线性注意力机制。如何理解其“线性”特点呢？标准注意力机制的核心在于注意力的计算，用每个输入 token 计算出对应的 Query、Key、Value 向量（简称 Q、K、V）。借用搜索引擎的比喻，Q、K、V 好比搜索引擎中的搜索查询词、网页索引库、网页内容，标准注意力机制会先将搜索查询词、网页索引库进行配对，即 Q 和 K 先相乘，然后再去匹配网页内容。这样做能精确匹配所有潜在关系，但需要将搜索查询词和网页索引库一一配对，计算量大，数学上是一个随输入长度增长而平方增长的矩阵，这也是制约当前大模型 Scaling Law 的核心因素。在自回归推理中，还需要保存随序列长度线性增长的 KV cache。

KDA 做出的改变是，让网页索引库和网页内容先进行配对，即 K 和 V 先相乘，再用查询搜索词去匹配。K 和 V 相乘好比对可检索的网页内容先进行总结，数学上是一个长和宽固定的矩阵，无论输入长度是多少，这是 KDA 计算效率高的核心机制。在更精细的结构中，KDA 继承 Gated DeltaNet，加入并精细化了总结网页内容即模型记忆库的管控结构，使其可以在推进推理的过程中，合理地忘记部分内容，并记住新内容（所以严谨来说，KDA 实际还包含遗忘门，并不是简单地把所有 K 和 V 相乘后累加）。

KDA 绕开了标准注意力机制的完整记忆矩阵，但也引入了串行依赖：后一个 token 的状态必须等前一个 token 计算完成。所以，这个方法在 Kimi K3 这个参数体量下，会遇到算力瓶颈，它是串行更新的，而 GPU 又偏好大规模、高带宽并行计算，容易导致算力利用率不足和算力浪费。

优化方式基本是把其中串行过程尽可能并行化。于是，月之暗面提出了 FlashKDA，FlashKDA 把 token 进行分块（chunk）计算，chunk 内部可以并行计算，chunk 之间仍需依次传递递归状态。在不同设备（或者 GPU）之间，Kimi K3 采用更进一步的 KDA 上下文并行（KCP）方法，来提升并行度。

对于标准的线性注意力的记忆递归计算，每个输入序列分段的状态转移可以在不知道输入状态的情况下独立求值，随后再精确合成递归链条。通俗来说，不是让每个序列分段计算自己的“历史”，而是先让每个序列分段独立算出自己那段“会怎样改变历史信息”，再把这些结果拼起来，就能还原每一段真正应该接到的前文状态，从而完成“历史演化”的计算。但对于 KDA，这样并不直接可行，因为其状态转移函数具有前面提到的记忆遗忘机制，在数学上导致一些额外的困难。KCP 的做法是，把状态转移中可以独立计算的部分尽可能拆解出来，再分到不同的 GPU 进程中去独立、并行计算，细节比较复杂，这里不再展开。

AttnRes：注意力残差机制

标准残差连接是指，一层神经网络在处理前面层传入的输入时，不是把原来的信息完全替换掉，而是只学习“需要修改或补充的部分”，再把这部分加回原输入，即“输出 = 原输入 + 本层计算结果”。这就像在原稿上做增量修改，而不是每一层都重新写一遍，使原始信息和梯度能沿网络直接传递，从而让很深的模型更容易训练，也减少信息在多层处理中被破坏或遗忘。

但随着神经网络深度增加，会越发稀释靠前的每一层的贡献。注意力残差 AttnRes 就是为了克服这一点而提出的，使其对原输入的继承更加建立在基于语义理解的筛选上。

然而，额外引入的注意力机制在大参数模型中，会带来新的内存负担，所以月之暗面又提出了 Block AttnRes。Block AttnRes 将神经网络层分为多个块，块内使用的还是标准的残差连接，块间使用注意力机制，研究表明，在约 8 个块的情况下，它能够达到 AttnRes 的大部分性能优势。

当然，Block AttnRes 本身引入的内存开销也不可忽视，仍然需要做进一步的优化。注意力机制是在块间进行的，且需要随着前向传播的演进计算多次。所以块表示只需要计算一次，就可以长期保留在 GPU 中，并重复使用。比如按 6 层为一个块，前 6 层的块表示可以被后 6 层在注意力计算使用，也可以被后 7 到 12 层使用。

AttnRes 相比标准的残差连接需要额外计算注意力机制相关的中间结果，如注意力分数等，这些结果在反向传播中也要用到，但如果一直保留到反向传播阶段，会占用大量内存。因此 Kimi K3 选择在前向传播时不在内存中保存这部分中间结果，等反向传播阶段时再重新计算这些数值，属于典型的用额外计算换取显存。

此外，由于训练采用了流水线并行，即模型不同的层放在不同的 GPU 上，而它们在计算注意力残差时需要用到其它块的块表示，这些块表示还保存在其它块所在 GPU 上。由于这些层之间存在串行关联，因此可以把块表示按照层顺序的轨迹，依次传递过去，并进行缓存，而不需要从源头重复传递。比如把 GPU0 的 Block_0 的表示 B_0，传递给 GPU1 的 Block_1 之后，GPU1 将 B_0 缓存，等 Block_1 计算出表示 B_1 后，将 B_0、B_1 一起传递给 GPU2 的 Block_2，GPU2 将 B_0、B_1 缓存。

最后，这些块表示仅限于在一个微批次的训练数据中可重复使用，因此在微批次训练结束后可以释放。综合以上手段，Block AttnRes 的优化方案做到了：不重复计算同一个块；不保存可以重算的中间激活；不重复传输已经缓存的块；不保留已经失效的微批次块；因此显存里只留下“当前时刻确实仍有用的数据”，最终达到 Kimi K3 所声称的“内存占用理论下限”。

MoonEP：解决 MoE 架构难题

在 MoE 架构的大模型中，每一个输入 token 可能激活不同的专家，那么反过来，对于任意一段输入 token，输入每一个专家的 token 数量可能都不同，进而造成激活形状的动态变化。然后，在传统专家并行中，每个专家固定放在某个 GPU 上，而路由器不会保证 token 被平均分给这些专家。于是，有些 GPU 要处理大量 token，有些 GPU 很快就做完并等待。由于整个训练步骤必须等待最忙的 GPU，整体训练吞吐量被最慢的 GPU“拖后腿”。

同时，路由专家每轮收到的 token 数不断变化，使中间张量时大时小，GPU 需要频繁申请和释放不同大小的显存块，留下许多难以复用的不连续空洞，从而浪费显存、增加分配开销。在这种内存碎片化的情况下，虽然零散的剩余显存总量可能不少，但单个新的张量或通信缓冲区需要一块足够大的可用空间，分散的小块无法直接拼成一次有效分配。

为此，Kimi K3 提出了 MoonEP 方法，即直接给每个 GPU 分配相同的任务量，接收 S×K 个 token，其中 S 是序列长度，K 是每个 token 选择的专家数量。而一个 GPU 里不同的专家还是可能分配到不同的任务量，所以需要在 GPU 中配置冗余专家，以应对随时增加的计算负载。冗余专家本身是一个高负载专家（即被分配了较多 token 的专家）的复制，作用是为高负载专家分摊计算任务，从而缩短最忙碌 GPU 的工作时间。

月之暗面证明了，每个 GPU 只需要固定数量的冗余专家（比如专家总数为 100，参与专家并行的 GPU 数量是 20，每个 GPU 最多需要 100/20 = 5 个冗余专家），便在理论上确保可以实现负载均衡。过去的方案一般会设置固定的冗余专家数量，或设置单个 GPU 的 token 上限，容易因为无法满足条件使得训练被迫终止。但要精确应用这个算法会导致计算成本过高，因此实际落地时还是使用了一些近似优化手段，并确保冗余专家数量不超过上限。

写在最后

可以看到，以上三个 Infra 技巧关注不同的方面，KDA 关注串行过程的并行化，AttnRes 关注计算的重复性，MoonEP 关注负载均衡。但把这些技巧放在一起看，是一套贯穿系统的工程思维。凡是串行链条，就寻找可并行拆解的部分，凡是重复计算、存储和通信，就判断能否重算、缓存或增量传输，凡是动态负载，就尽可能把它转化为可预测、可调度的静态形状。

AI Infra 的竞争并非简单堆更多 GPU，而是不断识别那些被默认接受的等待、复制、同步和闲置，再将它们逐一消除。单个优化看似只节省几个百分点，但当模型扩大到万亿参数、训练持续数月时，每个百分点都会被放大成巨额算力、时间和资金。

需要提醒的是，本文关注的这三个点，相对于技术报告的整个 Infra 章节只是冰山一角。Kimi K3 公开了方法，但这不是一套可以直接复制的标准答案。不同芯片、网络、模型结构和业务场景，都会产生不同的最优解。技术报告能帮助后来者少踩一些已知的坑，但真正决定部署效率的，仍然是团队能否在长期运行中发现并修复自己遇到的 Infra bug。总之，模型架构决定了能力的可能性，而 Infra 决定了这种可能性能否以可承受的成本，稳定地变成现实。