KimiK3近期终于公布了技术报告，不仅是模型架构，连重要的Infra层面也说了不少细节。

Infra指的是基础设施层——如果把硬件比作手机本体，AI比作应用软件，那AIInfra就是中间的操作系统，没有它硬件就无法发挥作用。

行业普遍认为，国产开源大模型未来的商业化有一个重要的支撑点，那就是大模型厂商自己部署他们的模型，成本可以比第三方部署低好几倍，这里依靠的便是Infra上的技术和经验差距。

OpenAI核心工程师翁家翌也曾在访谈中提到，大模型训练的核心差距在于Infra，工程师的主要工作是给Infra修bug。

也就是说，当模型架构逐渐趋同、公开数据和蒸馏技术降低部分能力复现门槛后，Infra工程能力会成为拉开训练成本、部署效率和运行稳定性的重要差距之一。

那么，KimiK3的技术报告有没有说透Infra的秘密？从商业的角度估计，一定是有所保留的。这意味着一百个公司部署一百个KimiK3，可能有一百种不同的运行效率。

当然，KimiK3好不容易公开了那么多Infra细节，该学的还是要学学。刚好，Infra对很多非从业者来讲是一个很陌生的事儿。

所以在本文中，知危想借KimiK3技术报告中一些有趣的Infra技巧，给大家科普一下AIInfra的核心原则。

为什么Infra重要？

前面我们提到，OpenAI核心工程师翁家翌说，大模型训练过程中工程师的主要工作就是给Infra修bug，那这个bug到底指的是什么呢？

在大模型团队的日常语境中，Infrabug既包括导致结果错误、训练不稳定或任务中断的传统bug，也包括那些虽然不影响数学正确性，却会造成GPU等待、内存浪费、通信阻塞和吞吐下降的系统性问题。

更具体而言，是在保证计算逻辑正确的前提下，看GPU是被充分利用还是有大量闲置算力，就算GPU被充分利用，还要看是否存在大量不必要的重复计算，要看GPU上的计算负载是否均衡、是否稳定，等等。要满足这些要求，还需要大量的通信在GPU之间做负载协调，这时候通信又可能遇上新的瓶颈。

这些bug最终影响的是集群吞吐量，并反映训练效率。

模型浮点运算利用率(MFU)是评估训练效率的标准指标，它是“观测吞吐量”与“理论最大吞吐量”的比值，理论最大吞吐量是假设达到峰值浮点运算的100%。

从行业的角度来看，大模型的训练效率仍然有很大提升空间，目前公开可验证的万卡、千亿参数、完整LLM预训练案例中，字节跳动的MegaScale的55.2%MFU仍是最有代表性的最高纪录之一。今年4月甚至还出现一个极端反面案例，据《TheInformation》报道，马斯克旗下xAI坐拥约55万张英伟达GPU，但内部备忘录显示其MFU仅11%，远低于业界35%-45%的正常水准，也不及Meta（43%）、谷歌（46%）等竞争对手，可以说是尴尬至极。

拥有更高的训练效率的意义是非常大的，它可以直接转化为更短的训练周期、更低的训练成本、更多可尝试的实验、更大的数据量或模型规模。

虽然本文是要展现Infra设计的重要性，但由于没有足够的数据和参数的参考，很难自主判断技术报告中哪些Infra技术是最重要的或贡献最大的。

为此，知危主要还是从KimiK3自主提出的架构设计入手，比如KDA、AttnRes、MoonEP，看看其背后的Infra技巧体现的原则，让大家从逻辑上去感受，选择原则不代表实际重要性排序。

总体设计

先来总体看看KimiK3的分布式系统设计，其训练采用的并行方法都是比较经典的方法，比如流水线并行、专家并行、ZeRO分片等。

大模型训练的“并行化”简单理解就是，大模型训练过程其实包含大量的串行过程，从前向传播到反向传播，但在细节上又存在大量矩阵计算，矩阵计算一般都是可以并行化的。此外，一张卡的内存也很难装下超大参数规模的所有数值。

如果上面这段话你看着头疼，那么不严谨的翻译成人话说就是：盖房子必须从底下一步一步往上盖（串行过程），但是盖房子过程中你可以同时搬好多水泥和钢筋（矩阵计算的并行）。

为此，大模型训练通常会从模型参数、训练数据和中间激活值等不同维度进行切分，将计算和存储任务分配到多张GPU上，并通过合理的并行策略和调度机制，尽量减少GPU等待和空转。

例如，可以把模型的不同神经网络层放到不同GPU上依次计算，这叫流水线并行。由于后面的GPU必须等待前面的GPU产生中间结果，流水线中容易出现部分GPU空闲的“气泡”。实际训练中通常会把一个batch切成多个微批次，让不同GPU同时处理不同微批次，从而提高流水线利用率。

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

以上只是原理上的简单科普，但不同的大模型集群会有自己独有的“烦恼”。

特别是模型越大，越可能出现在小模型中不会出现的问题，因为某个原先不被注意的计算进程或内存存储对象会因为规模的增大，而突然变得对于GPU容量而言显著了。

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

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

KDA

首先是混合KDA注意力机制。

KDA全称KimiDeltaAttention，是一种线性注意力机制，如何理解其“线性”特点呢？

标准注意力机制的核心在于注意力的计算，用每个输入token计算出对应的Query、Key、Value向量（简称Q、K、V）。

借用搜索引擎的比喻，Q、K、V好比搜索引擎中的搜索查询词、网页索引库、网页内容，标准注意力机制会先将搜索查询词、网页索引库进行配对，也就是Q和K先相乘，然后再去匹配网页内容。

这样做的好处是能非常精确匹配所有潜在关系，但需要将搜索查询词和网页索引库进行一一配对，计算量大，数学上是一个随着输入长度的增长而平方增长的矩阵，这也是制约当前大模型ScalingLaw的核心因素。在自回归推理中，还需要保存随序列长度线性增长的KVcache。

KDA做出的改变是，让网页索引库和网页内容先进行配对，也就是K和V先相乘，再用查询搜索词去匹配。K和V相乘好比是对可检索的网页内容先进行总结，数学上是一个长和宽固定的矩阵，无论输入长度是多少，这是KDA计算效率高的核心机制。

在更精细的结构中，KDA继承GatedDeltaNet，加入并精细化了这个总结网页内容也就是模型记忆库的管控结构，使其可以在推进推理的过程中，合理地忘记部分内容，并记住新内容（所以严谨来说，KDA实际还包含遗忘门，并不是简单地把所有K和V相乘后累加）。

KDA绕开了标准注意力机制的完整记忆矩阵，但这也是有代价的，也就是引入了串行依赖：后一个token的状态必须等前一个token计算完成。

所以，这个方法在KimiK3这个参数体量下，也会遇到算力瓶颈，它是串行更新的，而GPU又是天然偏好大规模、高带宽并行计算的，容易导致算力利用率不足和算力浪费。

而优化方式基本就是把其中串行过程尽可能并行化。

于是，月之暗面提出了FlashKDA，FlashKDA把token进行分块（chunk）计算，chunk内部可以并行计算，chunk之间仍需依次传递递归状态。

在不同设备（或者GPU）之间，KimiK3采用更进一步的KDA上下文并行（KCP）方法，来提升并行度。

对于标准的线性注意力的记忆递归计算，每个输入序列分段的状态转移可以在不知道输入状态的情况下独立求值，随后再精确合成递归链条。

通俗来说，不是让每个序列分段计算自己的“历史”，而是先让每个序列分段的独立算出自己那段“会怎样改变历史信息”，再把这些结果拼起来，就能还原每一段真正应该接到的前文状态，从而完成“历史演化”的计算。

但对于KDA，这样并不直接可行，因为其状态转移函数具有前面提到的记忆遗忘机制，在数学上导致一些额外的困难。KCP的做法是，把状态转移中可以独立计算的部分尽可能拆解出来，再分到不同的GPU进程中去独立、并行计算，细节比较复杂，这里就不再展开了。

AttnRes

其次是注意力残差AttnRes机制。

标准残差连接是指，一层神经网络在处理前面层传入的输入时，不是把原来的信息完全替换掉，而是只学习“需要修改或补充的部分”，再把这部分加回原输入，即“输出=原输入+本层计算结果”。

这就像在原稿上做增量修改，而不是每一层都重新写一遍，使原始信息和梯度能沿网络直接传递，从而让很深的模型更容易训练，也减少信息在多层处理中被破坏或遗忘。

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

然而，额外引入的注意力机制在大参数模型中，也会带来新的内存负担，所以月之暗面又提出了BlockAttnRes。

BlockAttnRes将神经网络层分为多个块，块内使用的还是标准的残差连接，块间使用注意力机制，研究表明，在约8个块的情况下，它能够达到AttnRes的大部分性能优势。

当然，BlockAttnRes本身引入的内存开销也是不可忽视的，仍然需要做进一步的优化。

注意力机制是在块间进行的，且需要随着前向传播的演进计算多次。

所以块表示只需要计算一次，就可以长期保留在GPU中，并重复使用。比如按6层为一个块，前6层的块表示可以被后6层在注意力计算使用，也可以被后7到12层使用。

AttnRes相比标准的残差连接需要额外计算注意力机制相关的中间结果，比如注意力分数等，这些结果在反向传播中也要用到，但如果一直保留到反向传播阶段，会占用大量内存。

因此KimiK3选择在前向传播时不在内存中保存这部分中间结果，等反向传播阶段时再重新计算这些数值，属于典型的用额外计算换取显存。

此外，由于训练还采用了流水线并行，也就是模型不同的层放在不同的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缓存。

最后，这些块表示仅限于在一个微批次的训练数据中可重复使用，因此在微批次训练结束后可以释放。

综合以上手段，BlockAttnRes的优化方案做到了：

-不重复计算同一个块；

-不保存可以重算的中间激活；

-不重复传输已经缓存的块；

-不保留已经失效的微批次块；

因此显存里只留下“当前时刻确实仍有用的数据”，最终达到KimiK3所声称的“内存占用理论下限”。

MoonEP

最后，我们看看MoonEP的机制原理。

在MoE架构的大模型中，每一个输入token可能激活不同的专家，那么反过来，对于任意一段输入token，输入每一个专家的token数量可能都是不同的，进而造成激活形状的动态变化。

然后，在传统专家并行中，每个专家固定放在某个GPU上，而路由器不会保证token被平均分给这些专家。

于是，有些GPU要处理大量token，有些GPU很快就做完并等待。由于整个训练步骤必须等待最忙的GPU，整体训练吞吐量被最慢的GPU“拖后腿”。

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

为此，KimiK3提出了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关注负载均衡。

但把这些技巧放在一起看，是一套贯穿系统的工程思维。

凡是串行链条，就寻找可并行拆解的部分，凡是重复计算、存储和通信，就判断能否重算、缓存或增量传输，凡是动态负载，就尽可能把它转化为可预测、可调度的静态形状。

AIInfra的竞争并不是简单堆更多GPU，而是不断识别那些被默认接受的等待、复制、同步和闲置，再将它们逐一消除。

单个优化看起来可能只节省几个百分点，但当模型扩大到万亿参数、训练持续数月时，每个百分点都会被放大成巨额算力、时间和资金。

需要提醒的是，知危关注的这三个点，相对于技术报告的整个Infra章节也只是冰山一角。

KimiK3公开了方法，但这不是一套可以直接复制的标准答案。

不同芯片、网络、模型结构和业务场景，都会产生不同的最优解。技术报告能帮助后来者少踩一些已知的坑，但真正决定部署效率的，仍然是团队能否在长期运行中发现并修复自己遇到的Infrabug。

总之，模型架构决定了能力的可能性，而Infra决定了这种可能性能否以可承受的成本，稳定地变成现实。