Mixture-of-Kittens：我们为NVL72s提供的开源MoE巨型内核（25分钟阅读）

# Mixture-of-Kittens：我们为 NVL72 打造的开源 MoE 巨型内核 来源：https://cursor.com/blog/mixture-of-kittens 今天，我们开源了 Mixture-of-Kittens（MoK）(https://github.com/cursor/mixture-of-kittens)，这是我们为 NVL72 打造的、用于生产环境的 MoE 训练巨型内核。 随着我们扩展 Composer (https://cursor.com/blog/composer-2)（我们的智能体编码模型）的训练和推理规模，混合专家（MoE）层始终是最大的瓶颈。根据工作负载和训练配置的不同，它可能消耗端到端训练时间的一半以上。 MoK 通过将所有 MoE 通信和计算融合到单个完全确定性的内核中来解决这一瓶颈。它如今已支撑着数万张 GPU 上的 Composer 训练。 你可以尝试 MoK 并在 GitHub (https://github.com/cursor/mixture-of-kittens) 上探索相关代码。我们期待你的反馈和贡献。 MoK 源自此前多次加速 MoE 层的尝试。在过去一年里，我们编写了自己的 MXFP8 (https://cursor.com/blog/kernels) 和 NVFP4 训练内核，并为 MoE 推理开发了“warp decode”(https://cursor.com/blog/warp-decode) 方法。 但这些技术只优化了该层的计算部分，并假设 GPU 间通信会被单独处理。在我们的生产工作负载中，通信已成为限制因素。这促使我们从第一性原理重新设计整个 MoE 层，并将通信直接构建到内核中。 此外，我们迁移到 GB300 NVL72 也在两个重要方面改变了问题。首先，NVL72 是一个位于单一 NVLink 域内的多节点机架，能够在全部 72 张 GPU 之间实现快速、细粒度的计算与通信重叠。 其次，集成的 Grace CPU（GB300 中的“G”）相对于 GPU 而言往往较慢。我们发现 GPU 流很容易追上 CPU 端的工作，导致 GPU 在这段时间内完全空闲。因此，我们必须激进地最小化 CPU 工作和 CPU-GPU 同步。 我们针对这一系列挑战的解决方案便是 Mixture-of-Kittens（MoK），这是一个基于第一性原理为 NVL72 构建的高度优化的 MoE 训练巨型内核。MoK 将所有 MoE 通信和计算融合到单个内核中，完全确定性，并在公开可用的实现中达到了最先进的性能。 在我们跨多个 NVL72 机架的生产训练栈中，MoK 将端到端令牌/秒提升了 1.41 倍。 本文其余部分将解释 MoK 背后的关键思想，包括我们如何选择正确的通信方向、如何结构化计算与通信的重叠，以及如何通过环形令牌缓冲区消除 CPU-GPU 同步。我们还会介绍巨型内核设计、确定性、MXFP8 支持以及其他一些实现细节。 ## 在 MoE 中重叠计算和通信 MoK 针对 DeepSeek-V3 (DSV3) 风格的 MoE 层，这类层被广泛应用于包括 GLM、Qwen、Kimi（最高到 K2.7）和 DSV 本身在内的众多开放权重模型。这些层将一个共享专家与许多路由专家（通常数百个）结合在一起。 对于进入 MoE 层的每个令牌，路由器投影会选择 top-k 个路由专家，并为每个专家分配一个路由器权重。然后，每个被选中的专家执行标准的 feed-forward 网络计算，包括 up 和 gate 投影、SwiGLU 激活以及 down 投影。该层使用路由器权重组合共享专家和路由专家的输出。 我们使用以下记号： - \(D\) = 模型维度 - \(D_{expert}\) = 专家中间维度 - \(\mathcal{E}\) = 路由 top-k 专家集合 - \(x\) = 输入令牌 - \(w\) = 路由器权重 - \(W_{up}, W_{gate}, W_{down}\) = 专家权重，\(W_{up}, W_{gate} \in \mathbb{R}^{D_{expert} \times D}\)，\(W_{down} \in \mathbb{R}^{D \times D_{expert}}\) MoE 层的计算如下： \[ x' = W_{shared}(x) + \sum_{e \in \mathcal{E}} w_e \cdot W_{down}^e \cdot \text{SwiGLU}(W_{up}^e (x), W_{gate}^e (x)) \] 其中 \(\text{SwiGLU}(u, g) = \text{SiLU}(u) \odot g\)。 使用专家并行（EP）时，我们切分路由专家并将其权重分布到许多 GPU（或 rank）上，我们将共同持有所有专家权重的 rank 数量称为 EP 度。例如，有 256 个路由专家且 EP 度为 64 时，每个 rank 持有 4 个路由专家，外加共享专家。因此，根据路由器投影的结果，令牌必须在 MoE 层之前和之后跨 GPU 传输。 最简单的分布式 MoE 实现会将每个令牌发送到持有其被分配专家的 rank（dispatch all-to-all），运行 FFN，将结果返回到每个令牌的原始 rank（combine all-to-all），并对专家输出进行加权求和。但由于通信可能和计算本身一样耗时，顺序执行二者效率很低。 标准的补救措施是通过流水线化将 dispatch/combine^1 (https://cursor.com/blog/mixture-of-kittens#fn-1) 通信与每个专家的 FFN 重叠：传输一批令牌，在其上计算 FFN，同时重叠传输下一批令牌，如此重复。MoK 是该方案的一种变体，并带有一组新颖的、针对特定目标的技术，使其比现有基线更快。 MoE 前向计算与通信的重叠 ## 选择正确的通信方向 在跨 GPU 发送令牌时，可以选择基于推送（push）的机制，即拥有令牌的 GPU 主动将令牌存储到远程目标 GPU；也可以选择基于拉取（pull）的机制，即需要令牌的 GPU 从远程源 GPU 加载令牌。现有方法通常依赖基于推送的通信在 GPU 之间散射和收集令牌（例如 DeepEP）。 常见的观点是推送方式能更好地饱和 GPU 间链路，因为它涉及更少的协议通信，因此成为默认选择。然而，我们的观察是每种机制各有取舍，为每个通信算子选择正确的方式对于最大化性能至关重要，原因有以下三点。 ### 调度 为了尽可能快地 dispatch 和 combine 令牌，必须满足以下条件： 1. 机架中连接 72 张 GPU 的所有 NVLink 通道必须保持饱和。我们不能容忍令牌仅通过部分（源 → 目标）通道传输的时间段，因此必须选择令牌，使得在任意时刻每个源 rank 的发送都能均匀地分布到所有目标 rank。 2. 发送到某个 rank 的令牌应按照该 rank 的本地专家顺序到达。如果到达是无序的，则专家分组的 GEMM 需要等待更长时间才能获得完整令牌块，然后才能开始张量核心矩阵乘法。 3. 应避免本地拷贝。发往某个专家的令牌应直接落在连续内存中，这样分组 GEMM 无需在本地重新排序即可开始。 4. 满足上述三个条件的开销必须保持在最低水平。我们希望将大部分时间用于实际发送令牌，而很少用于调度或搜索要发送的令牌。 对于基于推送的 dispatch，我们需要生成一个调度表，其列为 `{src_index, dst_rank, dst_index}`，其中 `src_index` 是本地传入激活缓冲区中的索引，`dst_index` 是令牌必须落在目标 rank 内存中的位置。该表的行索引将决定令牌在 NVLink 上的发送顺序。我们希望这个表的行按 `dst_rank` 循环轮转，以便每条相连的通道保持忙碌。 在针对给定 `dst_rank` 的行中，`dst_index` 值必须在所有源 GPU 之间均匀交错，并且还必须遵循该目标 rank 上本地专家的递增顺序。构建此表涉及多次排序，并且每个 rank 的调度必须考虑其他所有 rank，因为不能有两个源 rank 写入相同的 `dst_index`。 | 调度索引 | 源索引 | 目标 rank | 目标索引 | |----------|--------|-----------|----------| | 0 | 17 | 0 | 10 | | 1 | 02 | 9 | 29 | | 2 | 03 | 3 | 30 | | ... | ... | ... | ... | 基于推送的调度：连续条目必须在目标 rank 之间交错以实现全网络饱和。 | 调度索引 | 源 rank | 源索引 | |----------|---------|--------| | 0 | 0 | 15 | | 1 | 2 | 21 | | 2 | 3 | 33 | | ... | ... | ... | 基于拉取的调度：连续条目必须在源 rank 之间交错以实现全网络饱和。 对于基于拉取的 dispatch，调度表简化为两列 `{src_rank, src_index}`，且表的行索引直接对应本地目标缓冲区中的令牌索引。 理论上，表中的行数和目标缓冲区中的行数应匹配（这将需要较大的内存分配；将在后面的章节中详述）。这里不需要排序。我们遍历路由器投影结果，每当找到一个应落在当前 rank 上的令牌，就将其源 rank 和索引写入调度表。算法如下： **输入：** 路由张量 \(R\)，其中 \(R[t][r]\) 是源 rank \(s\) 上令牌 \(t\) 的第 \(r\) 条路由被分配到的专家的全局索引；本地 rank \(l\)；每个 rank 的专家数 \(E_{local}\) **输出：** 每个本地专家的令牌计数 \(counts\)、每个本地专家和源 rank 的令牌计数 \(counts_src\)、区域偏移量 \(offsets\)、调度表 \(S\) 初始化 \(counts\)、\(counts_src\)、\(offsets\) 为 0；\(S\) 为空 对于每个令牌 \(t\)（从 0 到 \(T-1\)）： 对于每条路由 \(r\)（从 0 到 \(R-1\)）： 如果 \(R[t][r]\) 位于本地 rank 的专家范围内： 令 \(e = \text{local\_expert\_index}(R[t][r])\) 增加 \(counts[e]\) 增加 \(counts_src[e][s]\)（其中 \(s\) 是源 rank） 将 \((s, t)\) 追加到 \(S\) 在计算完所有计数后，计算前缀和以得到偏移量，并据此确定每个专家和源 rank 区域在目标缓冲区中的位置。 在实践中，我们实现该算法的调度内核耗时不到 MoE 总运行时的 3%，并且完全在设备端运行，无需任何 CPU-GPU 通信。 我们还可以将这一调度原样复用于 combine，只需选择基于推送的 combine，将 `{src_rank, src_index}` 读作 `{dst_rank, dst_index}` 即可。事实上，我们可以构建一次调度，并在前向和反向的四种通信操作中复用它，只需选择： - 前向 dispatch：基于拉取 - 前向 combine：基于推送 - 反向 reverse-combine：基于拉取 - 反向 reverse-dispatch：基于推送 在最坏情况下，调度表也只有几 MB，因此我们可以将其保留以供复用，而不会有任何内存压力。另一个好处是这完全消除了 GPU 间多通道信号传递，详见下文。 ### NVLink 带宽利用率 与任何网络系统一样，通过 NVLink 传输的所有用户数据（有效载荷）都会附带额外的协议元数据，其中包含源和目标信息、确认信息等，具体取决于网络协议。虽然 NVLink 通信协议的细节未公开，但我们可以观察链路上发送的数据并仔细推理。 我们发现，基于推送的 NVLink 通信移动的总字节数更少（即协议元数据加有效载荷），并且几乎所有数据都在一个方向上发送，因此在所有通道完全繁忙时能达到更高的带宽利用率。 另一方面，基于拉取的传输移动的总字节数更多，但协议元数据在链路的两个方向上的分布更加分散。拉取方先在一个方向上发送元数据，然后在另一个方向上接收更多元数据以及有效载荷。 我们可以通过编写一个简单的跨 GPU 传输内核并用 NCU 进行分析来验证这一点。例如，在我们的微基准测试中，通过 NVLink 发送一个 256x256 BF16 块（131,072 字节），我们观察到以下情况： | 方向 | 总计 | 总接收 | 总发送 | 协议接收 | 协议发送 | 有效载荷接收 | 有效载荷发送 | |----------|------------|------------|------------|----------|----------|--------------|--------------| | 推送 | 159.6 KB | 2.9 KB (1.84%) | 155.6 KB (99.16%) | 2.7 KB | 24.6 KB | 0 KB | 131.1 KB | | 拉取 | 172.0 KB | 147.5 KB (85.71%) | 24.6 KB (14.29%) | 16.4 KB | 24.6 KB | 131.1 KB | 0 KB | 在上表中，RX 指接收方向，TX 指发送方向。我们可以看到，推送方式总共减少约 12.4 KB 的字节数，理想情况下能带来更好的 NVLink 整体饱和。然而，NVLink 为每个方向提供独立通道。要达到 NVLink 规范中宣传的带宽（GB300 上第五代 NVLink 为 1.8 TB/s），流量必须在两个方向上都充分流动。而对于 MoE dispatch/combine 这样的工作负载，每次传输往往只有几千字节且不平衡，会在链路上留下执行气泡。 根据经验，在专家不均衡情况下分发令牌时，基于拉取的通信比基于推送的通信提供最高 29% 的 NVLink 带宽利用率提升，因此是更好的选择。 ### 信号传递 信号传递是必需的，因为 GPU 必须知道 dispatch/combine 通信何时完成。GPU 只有在收到完成信号后才能开始处理接收到的令牌。最小化这种信号传递开销非常重要，而选择正确的通信方向会产生显著影响。 对于基于推送的 dispatch 或基于拉取的 combine，完成信号必须跨 GPU 传递。在专家分组的 GEMM 开始之前，一个 rank 必须等待来自多达 71 个对端的信号，并通过发出内存栅栏在整个机架上刷新内存。此外，在信号传输期间，数据已经位于目标位置。 对于基于拉取的 dispatch 和基于推送的 combine，跨 GPU 信号传递消失了。每个 rank 发出一个 load，等待数据到达，然后就可以立即开始使用它。不需要多节点同步。此外，只有单一实体执行信号传递，因此开销不会随专家并行度增长。 我们发现这种信号传递开销相当大。在我们的多节点微基准测试中，基于推送的 dispatch 信号传递比基于拉取的 dispatch 信号传递延迟高出约 5.8 倍，分别为 103 μs 和 18 μs，而这种成本在巨型内核中会迅速累积。 基于这些结果，MoK 使用基于拉取的前向 dispatch、基于推送的前向 combine、基于拉取的反向 reverse-combine 和基于推送的反向 reverse-dispatch。 基于推送与基于拉取的 dispatch 对比 ## 恰当结构化计算-通信重叠 ### 计算-通信粒度 计算-通信粒度是计算任务和通信任务相互通知完成的最小单位。它带来一个有趣的权衡。 在一个极端，通信可以非常细粒度。我们只发送足以让张量核心启动一条完整矩阵乘加指令（例如在 Blackwell 上为 256）的令牌，并持续将每次传输与矩阵乘法重叠。Comet (https://arxiv.org/abs/2502.19811) 是这种方法的一个很好例子。 在另一个极端，通信可以是粗粒度的。我们一次发送数千或数万个令牌，并等待完整批次到达后才开始张量核心操作。DeepEP (https://github.com/deepseek-ai/DeepEP) 是这里的一个很好例子。 我们的观察是，两个极端都不是正确选择。最佳点位于中间，并取决于工作负载。 如果通信过于细粒度，张量核心永远无法完全饱和。Nvidia 张量核心是高度流水线化、带宽优化的加速器，我们希望持续喂饱它们，而不是每隔几条 MMA 指令就遇到一次屏障等待。 如果通信过于粗粒度，张量核心会等待太长时间才能等到第一批令牌到达，以及最终输出组合完成。 细粒度与粗粒度的计算-通信重叠 我们称每轮发送的令牌集合为一个**小批次（minibatch）**。MoK 将小批次大小设为可调参数，以便在训练前轻松扫描可能的配置。