Mixture-of-Kittens: our open-source MoE megakernel for NVL72s
5 hours ago
- Mixture-of-Kittens (MoK) 是一个开源的 MoE 训练超级内核，用于 NVL72，它将所有 MoE 通信和计算融合到单个确定性内核中。
- MoK 解决了 MoE 层瓶颈问题，该瓶颈可能消耗超过一半的训练时间，并在生产环境中跨多个 NVL72 机架将端到端每秒 tokens 数提高了 1.41 倍。
- 关键技术包括基于拉取的调度和基于推送的合并，以优化 NVLink 带宽并减少信令开销；可调节的小批量大小以实现计算与通信的重叠；以及使用专用 SM 进行通信和计算的 SM 间重叠。
- 环形令牌缓冲区（宏批次）消除了 CPU-GPU 同步，反向环形顺序在反向传播过程中最小化前向重放。
- MoK 支持 BF16 和 MXFP8 精度，包括融合到调度和 GEMM 中的 MXFP8 量化，并高效计算路由权重梯度。
- 基准测试显示，与最快的基线（DeepEP、NCCL 等）相比，MoK 在 MXFP8 前向中速度提升高达 2.37 倍，在 MXFP8 后向中速度提升 1.78 倍。
- 该超级内核使用集群启动控制（CLC）进行 RDMA 重叠，并使用软件分区实现确定性 SM 分配。
- MoK 在 GitHub 上完全开源，并设计为灵活且易于修改，以支持其他平台。