文章摘要
近期高性能AI硬件算力与互联带宽提升，但大规模混合专家模型（MoE）训练仍受通信瓶颈影响。不久前相关团队开源的Mixture-of-Kittens（MoK）重写执行逻辑带来新优化路径。它将前向Dispatch改为Pull模式，整合通信与计算，用固定大小的Ring Token Buffer解决显存问题。基准测试显示其性能提升显著，标志AI训练进入“全栈主权”新阶段。

近期高性能AI硬件的算力与互联带宽已达到全新高度，但大规模混合专家模型（MoE）的训练中，通信瓶颈依然会显著拖慢整体效率。不久前，相关团队开源了Mixture-of-Kittens（简称MoK），并未继续优化矩阵乘法速度，而是直接重写了MoE的执行逻辑，将token调度、跨GPU通信与专家计算整合进同一个GPU内核，为MoE训练带来了全新的优化路径。

7月下旬，相关硬件厂商公布了新一代训练集群的性能数据：在256张加速卡组成的集群上，单卡算力达到1648 TFLOPS。该集群将72张加速卡整合进同一高速互联域，整机总互联带宽达到130 TB/s，按常理GPU间的数据传输已足够高效，但实际MoE训练中，通信环节依然会拖累专家计算的整体效率。

MoE的核心特点是通过路由模块动态为每个token分配不同的专家模型，而专家模型通常会被拆分到多张加速卡上。一次前向传播至少需要两轮跨卡通信：第一轮将token发送到对应专家所在的GPU（Dispatch阶段），第二轮将专家计算后的结果送回原始token的位置（Combine阶段）；反向传播则需要再执行两轮反向通信操作。

若仅传输连续大块数据，高速互联链路可轻松满足需求，但MoE的token分配是动态的：每一轮路由都会生成不同的token分布，每个专家在不同迭代中收到的token数量差异极大。系统需要先统计每个专家的token数量，再规划目标GPU的存储位置，将同一专家的token尽量连续排列，才能让后续的分组矩阵乘法（Grouped GEMM）高效运行。

通信完成后，目标GPU需确认所有远端数据写入完成才能开始计算。同时专家负载通常并不均衡，部分GPU会提前完成任务，需等待最慢的专家收尾。因此看似简单的“通信”阶段，实际包含数据搬运、布局规划、同步等待与负载失衡补偿等多个环节。虽然整机互联带宽峰值可达130 TB/s，但动态且零碎的MoE通信无法完全跑满所有链路。

现有优化方案已在数据搬运层面做了大量优化，比如专门针对专家并行的高性能通信库，支持FP8精度并可控制通信占用的流多处理器（SM）数量，甚至可通过更少的SM保持较高通信吞吐。但单层MoE仍需在Dispatch、Grouped GEMM和Combine间不断切换交接，依然存在实际效率损耗。

这里还存在实际矛盾：若等待足够多token到齐再计算，矩阵规模足够大可让张量核心高效运行，但计算启动延迟会增加；若token到一点就开始计算，虽可让通信与计算更早重叠，但矩阵规模过小，会导致大量SM资源闲置。使用多个CUDA流可让通信和计算并发，但很难始终让两边获得恰到好处的GPU资源。MoK的核心设计正是围绕解决这个“节奏”问题展开。

MoK最核心的改动之一，是将前向Dispatch阶段从传统Push模式改为Pull模式。传统Push模式直观易懂：源GPU持有token后，主动将数据写入目标GPU。但该模式的问题在于目标地址协调：一张GPU可能同时收到多张其他GPU的token，每个发送源需提前知晓应写入目标GPU的哪段内存，避免互相覆盖；同时同一专家的token最好连续排列，否则后续矩阵乘法还需重新整理数据。当参与通信的GPU数量增多时，这种调度协调的开销会越来越大。

Pull模式则换了思路：由保存专家的目标GPU主动读取需要的token。目标GPU只需知晓token所在的源GPU位置与数据偏移量，本地存储布局可自行规划。这种方式省去了多个发送源间对目标地址的协调工作，收到的数据可直接按本地专家需求完成排列。

有意思的是，Pull模式并未减少传输的数据总量。根据相关团队的微基准测试，同一个256×256的BF16数据块，Push模式在高速互联上传输约159.6 KB，而Pull模式反而达到172.0 KB，因为读取操作需要额外发送请求。但Pull模式的优势体现在其他方面：MoE的通信通常零碎且负载不均衡，高速互联的两个方向拥有独立通道，Pull模式可同时利用请求发送与数据返回两个通道。在专家负载不均衡的测试场景中，相关团队测到最高29%的互联带宽利用率提升。

同步层面的改进更加明显。使用Push模式时，目标GPU需等待所有其他GPU的完成信号，当专家并行规模较大时，一个节点最多需等待71个对等节点的信号。而Pull模式由本地GPU主动发起读取，数据返回后可直接使用，无需额外同步等待。根据多节点微基准测试，这种同步信号延迟从Push模式的约103微秒降到了Pull模式的18微秒。

MoK并未在所有环节都使用Pull模式：前向传播采用Pull Dispatch与Push Combine的组合，因为Dispatch阶段需将多个来源的token重新组织为专家输入，Pull模式更节省协调开销；而Combine阶段每个结果应回到哪个token已非常明确，直接Push回去的方式更简单。反向传播则采用Pull Reverse-Combine与Push Reverse-Dispatch的对应组合。

在改变通信方向的基础上，MoK还将通信与专家计算整合进同一个巨型内核（Megakernel）。它将GPU的SM划分为两部分：一部分负责Dispatch、Combine与状态管理，另一部分专门执行专家的前馈网络计算。通信侧拿到一批完整token后，通过GPU本地计数器通知计算侧；计算完成后，再通知通信侧将结果传送回去。这种方式可直接分配通信与计算各自使用的SM数量，无需完全依赖多个CUDA流自行竞争资源。

这里最关键的参数是minibatch，即一次交给专家计算的token数量。该参数不能过大，否则第一批计算需等待很久才能启动；也不能过小，否则专家的矩阵乘法会被拆分为太少的计算任务，大量SM会处于闲置状态。相关团队用“wave”衡量这个边界，一个完整的wave可理解为所有计算SM都分配到了工作任务。MoK希望一个minibatch至少能形成两个完整的wave，让张量核心有足够任务持续运行。

根据相关测试，在隐藏层维度为7168、专家中间维度为2048的模型结构下，minibatch至少需要约2368个token才能达到最优效率。当minibatch为512个token时，MoK的前向耗时为5.981 ms；增加到2560个token后，耗时下降到3.425 ms；继续增加minibatch大小，速度无明显改善。这说明通信拆分过细并不会持续提升效率，真正高效的重叠需要让通信尽早交付数据，同时又不能将矩阵乘法切得过小。

MoE还有一个额外麻烦：在路由模块完成计算前，系统无法得知每张GPU最终会收到多少token。若按最坏情况准备显存缓冲区，会浪费大量内存；若先让GPU统计完token数量，再通知CPU分配对应空间，GPU又需停下来等待CPU操作。MoK使用固定大小的Ring Token Buffer解决这个问题：一块内存空间先存储Dispatch阶段送来的token，当专家计算完成、Combine阶段将结果传送走后，这块空间会立即被下一批token复用。前一个宏批次（macrobatch）的Combine操作，还可与下一个宏批次的Dispatch操作同时进行。

Ring Buffer就像一个缓冲层：若通信速度暂时更快，数据会在缓冲区中积累；若计算消费速度更快，则等待下一批token到达。整个过程由GPU上的状态机推进，无需CPU每一轮都介入决定下一步操作。此外，MoK还将MXFP8激活量化嵌入到Dispatch、Grouped GEMM和SwiGLU的数据路径中，省去独立的量化内核，减少了中间结果在高带宽内存中的来回读写。

将Pull调度、minibatch控制、SM分区和Ring Buffer整合后，MoK真正构建出了一条连续的MoE训练流水线。

相关团队的基准测试覆盖了完整的MoE层，包括调度、Dispatch、专家FFN、Combine和最后的加权合并，对比对象包括NCCL+PyTorch、DeepEP、TransformerEngine以及HybridEP+Megatron等现有方案。

在新一代高速互联集群上，相比各场景下最快的公开基线方案，MoK的MXFP8精度前向传播最高提升2.37倍，反向传播最高提升1.78倍；BF16精度的前向和反向传播最高分别提升1.92倍和1.58倍。

更具实际意义的是端到端训练性能：相关团队原本的生产方案已使用DeepEP，在512张新一代加速卡组成的集群上，替换为MoK后，单卡吞吐从每秒760.9个token提升到1070.2个，整体提升约41%。

需要注意的是，相关团队未公开完整的逐项消融测试数据，因此无法准确量化2.37倍的性能提升中，有多少来自Pull模式、多少来自巨型内核、多少来自Ring Buffer。可以确认的是Pull模式对互联带宽利用率和同步延迟的改善，其余性能收益更多来自整套执行方式组合后的协同效果。

同时MoK对硬件依赖较强，它面向类似新一代高速互联域的硬件环境，Pull模式的远端读取、通信和计算的细粒度交错，都建立在GPU之间能够低延迟访问彼此显存的基础上。模型的隐藏层维度、Top-K选择、专家规模发生变化时，合适的minibatch和通信SM数量也需要相应调整。

这也是MoK最值得关注的地方：过去讨论MoE优化时，人们很容易只关注两个数字——矩阵乘法的TFLOPS性能和All-to-All通信的GB/s带宽。但到了新一代硬件时代，仅继续提升这两个数字已无法完全解释训练效率的差距。

token何时到达、如何排列为专家需要的布局、积累多少token才开始计算、通信操作占用多少SM、缓冲区何时释放，这些执行细节开始直接决定训练的整体速度。即便拥有130 TB/s的互联带宽，依然需要为MoE重写GPU内核，核心原因就在于：链路速度已足够快，接下来需要节省的是GPU等待数据的时间。

相关团队重写GPU内核的行为，标志着AI训练的竞争进入了“全栈主权”的新阶段。过去人们普遍认为“术业有专攻”，做应用的团队专注应用层，做底层的团队专注硬件和内核开发。但如今的AI竞争已进入“去中间商化”的新阶段，应用团队并非出于兴趣开发底层内核，而是出于必须：只有深入控制底层执行细节，才能榨干硬件的最后一丝性能。

相关团队开启了工程优化的新纪元，将性能压榨的战火蔓延到了应用层。这种“去中间商化”的趋势正在重塑AI行业的定价权：未来决定AI公司估值的，不再是它拥有多少token，而是它的代码离显存和寄存器到底有多近。那些无法穿透底层黑盒的AI公司，最终会陷入“平庸税”的泥潭，无法获得最优的训练效率和成本控制。

以上内容不代表本平台立场，仅供读者参考