1. 先搞清楚 MoK 到底是什么以及它解决了什么核心问题如果你最近在关注大模型训练尤其是混合专家模型那你可能已经看到了 Cursor 开源的 Mixture-of-Kittens。这个名字听起来有点“萌”但它要解决的问题非常硬核如何高效、确定性地在像 GB300 NVL72 这样拥有海量 GPU 的超级计算机集群上训练一个 MoE 模型。这里有几个关键点需要拆开看。首先MoE 是 Mixture of Experts 的缩写它不是一个新的概念但在大模型时代被重新重视。简单来说MoE 模型不像传统的“稠密”模型那样每一层都用所有的参数处理每一个输入。相反它有一个“路由”机制对于每个输入只激活一小部分“专家”网络进行计算。这就像你有一个庞大的专家团队每次只请几位最相关的专家来解答问题而不是让所有人同时发言。这样做的好处是在模型总参数量巨大的情况下实际计算量可以大幅降低从而提升训练和推理速度。但 MoE 训练有个老大难问题负载不均衡和通信开销。因为输入是动态路由的你无法保证每个 GPU 上的专家被均匀地激活。可能有些 GPU 上的专家忙得要死有些却闲着。在单机多卡或者小规模集群上这个问题还能通过各种调度策略缓解。但当 GPU 数量膨胀到 GB300 NVL72 这种级别可以理解为由 72 个顶级 GPU 组成的超大规模计算节点时负载不均衡和随之而来的 GPU 间数据交换通信就会成为性能瓶颈甚至导致训练过程不稳定、不可复现即非确定性。Cursor 开源的 Mixture-of-Kittens 就是为了解决这个问题而生的。它不是一个完整的训练框架而是一个“Megakernel”。你可以把它理解为一个高度优化、针对特定硬件和任务定制的“超级计算内核”。它的目标很明确在 GB300 NVL72 这样的特定硬件架构上实现 MoE 层前向传播和反向传播的确定性、高性能执行。所以这篇文章的核心不是教你从零开始训练一个 MoE 模型而是带你理解当你的计算规模大到需要关注每一个时钟周期和每一字节的通信时像 MoK 这样的底层核心里到底发生了什么。这对于从事大规模分布式训练、框架底层优化或者单纯想理解前沿技术如何落地的工程师来说价值非常大。2. 从“原理可行”到“生产稳定”MoE 训练的挑战与 MoK 的切入点在深入 MoK 之前我们必须先理解为什么常规的 MoE 实现在超大规模集群上会“失灵”。很多论文和开源代码展示了 MoE 的原理在小规模比如 8 卡、16 卡上也能跑出漂亮的数据。但一旦规模上去问题就暴露了。2.1 常规 MoE 训练的三大痛点动态路由带来的不确定性路由决策哪个输入由哪个专家处理如果在不同 GPU 上、甚至不同运行批次间有细微差异就会导致后续的梯度计算和参数更新路径完全不同。这在科学研究中是无法接受的因为你无法复现实验结果。在生产中这会导致模型训练波动大难以调试。专家负载的严重不均衡即使路由算法设计得再精妙在超大规模分布式环境下由于输入数据的随机性和通信延迟几乎不可能保证 72 个 GPU 上的专家工作量完全均等。这会导致“木桶效应”整个集群的速度取决于最慢的那个 GPU其他 GPU 大量时间在空闲等待。爆炸式的通信开销MoE 层需要将输入数据根据路由结果从一批 GPU 发送到另一批 GPU 上进行专家计算然后再将结果收集回来。这个“All-to-All”或“All-to-Many”的通信模式在 GPU 数量巨大时通信量会呈指数级增长成为绝对的主导开销。2.2 MoK 的设计哲学确定性、融合与硬件协同Mixture-of-Kittens 的“Megakernel”设计正是直击这些痛点。它的思路不是在上层的调度算法上修修补补而是深入到最底层的计算单元。确定性优先MoK 的核心目标之一是保证每次运行的计算结果完全一致。这意味着从路由逻辑、数据分发、专家计算到结果聚合整个流程必须是无二义性的、可重复的。它通常会采用确定性的算法来处理路由和通信避免任何可能引入随机性的操作。计算与通信的极致融合传统做法是先完成路由然后发起通信等数据到位后再进行计算。MoK 作为一个“大内核”会尝试将通信和计算重叠起来甚至将多个步骤融合在一个内核中执行减少 GPU 核心的启动开销和内存访问延迟。这就是“Kernel Fusion”的思想在分布式场景下的应用。为特定硬件量身定制标题中提到的“GB300 NVL72”不是泛指的集群而是一个具体的、高性能的硬件配置通常指基于 NVIDIA Grace Hopper 超级芯片和 NVLink 高速互联的系统。MoK 会充分利用这种架构的特性比如NVLink 高速互联优化 GPU 间点对点通信的数据路径。共享内存模型在 Grace CPU 和 Hopper GPU 之间可能存在的高带宽内存MoK 可以设计特定的数据驻留策略。Tensor Core 利用确保专家网络的计算部分能最大程度调用 GPU 的 Tensor Core 进行矩阵运算。简单来说MoK 是把 MoE 层在超算集群上的“该怎么算”和“该怎么传数据”这两个最棘手的问题打包成一个高度优化的、确定性的超级程序块Megakernel。上层框架比如 PyTorch只需要调用这个内核而不必关心底下复杂的分布式协同细节。3. 环境与思想准备你离运行 MoK 有多远看到这里你可能摩拳擦掌想试试。但我们必须现实一点Mixture-of-Kittens 不是一个面向个人开发者或小团队的即插即用工具。3.1 硬性环境要求硬件架构它的首要目标平台是GB300 NVL72或类似规模的、基于 NVIDIA Hopper 架构且通过 NVLink 高速互联的超大规模 GPU 集群。你手头的单台 RTX 4090 游戏 PC或者一个小型 A100/H100 服务器很可能无法直接运行或者无法体现出其价值。因为它深度绑定了特定的大规模互联拓扑和内存层次结构。软件栈你需要有对应的底层驱动、CUDA 版本、以及可能特定的通信库如 NCCL版本支持。MoK 作为底层内核通常需要编译安装对系统环境要求苛刻。框架集成MoK 本身是一个内核你需要将它集成到像 PyTorch 这样的深度学习框架中。这需要你具备修改框架源码或编写自定义 C/CUDA 扩展的能力。3.2 你可以从 MoK 中学到什么即使你没有 GB300 集群研究 MoK 的设计和代码也极具价值理解确定性训练的实现你可以学习它如何通过算法设计如确定性的排序、哈希来消除分布式计算中的随机性。这个思想可以应用到你自己项目的分布式训练中。学习内核融合与优化技巧看看它是如何将多个操作路由、数据搬运、矩阵乘、激活函数融合到一个内核里减少全局内存访问和内核启动开销的。这些优化技巧在小规模 GPU 编程中同样适用。掌握面向硬件的编程思想MoK 是硬件感知编程的典范。研究它如何根据 NVLink 拓扑设计通信、如何安排数据在 HBM 和 SRAM 中的布局能极大提升你对高性能计算的理解。所以对于大多数读者我建议把本文和 MoK 项目当作一个“高级案例”来学习而不是一个“马上要部署的工具”。它的价值在于展示了大模型训练技术的前沿和工程极限。4. 核心流程拆解一个 Megakernel 内部是如何工作的虽然我们不能直接运行但我们可以深入原理看看 MoK 这个“黑盒”里大概经历了哪些步骤。这对于理解分布式 MoE 训练至关重要。假设我们有一个简单的 MoE 层分布在 4 个 GPU 上为了简化实际是72个。每个 GPU 上有若干个专家。一批输入数据也被分布在这 4 个 GPU 上。4.1 步骤一确定性的本地路由与全局协调本地路由计算每个 GPU 独立地对自己持有的那部分输入数据计算路由权重例如通过一个门控网络。这里的关键是路由计算本身必须是确定性的。不能使用任何会产生随机种子的操作。全局信息交换每个 GPU 需要告诉其他 GPU“我有哪些数据Token要发送给你”。MoK 需要高效地完成这次全局通信形成一张“发送-接收”映射表。为了实现确定性这个过程通常基于一个所有 GPU 达成共识的规则例如对所有 Token 按某种唯一 ID如全局索引进行排序后再决定归属。4.2 步骤二高效的数据置换Permutation这是通信密集型阶段。根据上一步的映射表GPU 之间需要交换数据。在经典实现中这通常是一个AllToAll操作。MoK 的优化可能体现在通信与计算重叠在发送/接收数据的同时GPU 可能已经开始处理已经到位的部分数据。利用 NVLink 拓扑不是所有 GPU 对之间都有直接的 NVLink 连接。MoK 会规划数据转发的路径使得通信尽可能发生在高速链路上避免经过慢速的 PCIe 或网络。打包Packing与小批量聚合为了减少通信次数可能会将多个发往同一目标 GPU 的小数据包打包成一个大的数据包再发送。4.3 步骤三专家计算数据到达正确的 GPU即该专家所在的 GPU后开始执行专家网络的前向计算。这里 MoK 的优化点在于内核融合专家网络可能由线性层、激活函数等组成。MoK 可能会将这些操作融合成一个内核避免中间结果写回全局内存。利用 Tensor Core确保矩阵乘运算以最优的方式调用 Tensor Core。负载均衡虽然路由阶段尽力均衡了但专家计算量仍有差异。MoK 可能在内核内部采用更细粒度的并行如更精细的线程块划分来消化这种不均衡。4.4 步骤四反向聚合与梯度同步前向传播的逆过程。计算得到的专家输出需要按原路返回给对应的原始 GPU。梯度也需要沿着相似的路径进行同步和聚合。MoK 的“Megakernel”可能将前向的步骤 2、3、4 以及反向的相应步骤融合或紧密耦合在一起形成一个从输入到输出梯度的高性能流水线。整个过程中通信和计算之间的空隙被压缩到最小。5. 关键参数与配置如果我要设计这样一个系统该关注什么如果你有志于从事相关领域的开发或者需要评估类似的技术以下是你需要关注的核心维度。这些也是 MoK 这类项目必须做出设计和权衡的地方。5.1 模型与算法参数参数含义影响与考量专家数量MoE 层中专家网络的总数。数量越多模型容量越大但路由和通信复杂度呈平方级增长。需要与 GPU 数量匹配。激活专家数 (k)每个输入 Token 实际经过的专家数量通常是 top-1, top-2。k 越大精度可能更高但计算和通信成本也翻倍。是精度与效率的关键权衡点。专家容量因子为每个专家分配的缓冲区大小用于处理可能超过其“公平份额”的输入。设置过小会导致溢出某些专家太忙输入被丢弃影响精度设置过大会浪费显存和计算资源。路由算法决定输入分配给哪个专家的机制如线性门控、可学习路由。算法本身的复杂性、确定性和负载均衡能力。MoK 要求路由算法必须是确定性的。5.2 分布式与硬件参数参数含义影响与考量GPU 拓扑GPU 之间的连接方式NVLink, PCIe, InfiniBand。这是 MoK 优化的核心依据。它决定了数据交换的最优路径。MoK 需要感知 NVL72 这样的具体拓扑。数据并行维度除了 MoE 的专家并行是否同时采用了数据并行DP、张量并行TP、流水线并行PP。MoK 主要解决专家并行EP的问题。在实际训练中EP 常与 DP/TP/PP 结合形成复杂的 4D 并行这会给通信带来更大挑战。批次大小与序列长度训练的全局批次大小和输入序列长度。直接影响每次通信的数据量。数据量越大通信开销占比可能相对降低但对显存压力越大。5.3 性能与稳定性参数参数含义影响与考量确定性种子确保整个分布式系统随机性一致的种子。对于可复现性至关重要。必须贯穿路由、数据重排、Dropout如果使用等所有环节。通信-计算重叠策略如何将数据通信和 GPU 计算同时进行。MoK 性能提升的关键。策略的好坏直接决定了 GPU 利用率。溢出处理策略当某个专家的输入超过其容量时如何处理。直接丢弃会影响精度复杂的负载均衡或重路由算法会增加复杂性和不确定性。MoK 需要清晰的定义。注意对于使用 MoK 的最终用户研究员很多底层参数可能是封装好的。但当你遇到性能瓶颈或收敛问题时理解这些底层参数能帮你更好地定位问题是与框架开发者沟通的基础。6. 排查与调试当大规模 MoE 训练出问题时从哪里入手假设你在一个大规模集群上运行集成了 MoK 的 MoE 训练任务出现了问题如速度不达预期、显存溢出、训练不稳定可以遵循以下排查路径。这个思路同样适用于其他分布式训练场景。6.1 第一步确认单卡与单机健康度不要一上来就怀疑分布式逻辑。首先排除基础问题。单卡算力用标准的矩阵乘基准测试如torch.cuda.amp.GemmBenchmark检查单块 GPU 的 FP16/BF16/TF32 算力是否正常。单机通信在单台服务器内的多块 GPU 之间运行 NCCL 测试如nccl-tests检查 NVLink 带宽是否达到预期。显存与内存检查是否有其他进程占用资源。确保没有内存泄漏。6.2 第二步缩小规模进行确定性验证这是调试分布式问题的黄金法则。最小可复现样例将模型规模专家数、层数、数据规模批次大小、序列长度和 GPU 数量都降到最小例如 2 个专家2 个 GPU。开启确定性模式在框架中设置torch.use_deterministic_algorithms(True)和torch.backends.cudnn.deterministic True并固定所有随机种子。运行并记录运行一个训练 step记录下损失值、关键中间变量的统计量均值、方差。对比结果多次运行检查结果是否完全一致。如果不一致问题就出在确定性上可能是 MoK 的集成有漏洞或者你的数据加载、初始化等环节引入了随机性。6.3 第三步性能剖析与瓶颈定位当训练能跑通但速度慢时使用性能分析工具。PyTorch Profiler这是最直接的武器。它会生成一个时间线清晰地展示出MoK 内核的执行时间看看它本身是快是慢。CPU 与 GPU 的间隙如果 GPU 有大段空闲说明它在等待可能是数据加载也可能是通信。通信操作耗时重点看all_to_all,all_gather等通信原语的耗时。MoK 内部可能封装了这些操作但 profiler 通常能捕捉到。NCCL 调试设置环境变量NCCL_DEBUGINFO或NCCL_DEBUGWARN运行训练观察 NCCL 的日志输出看是否有通信错误或警告。GPU 利用率使用nvidia-smi dmon或nvtop实时监控 GPU 的 SM流处理器利用率和显存带宽利用率。理想的 MoK 执行期间利用率应持续在高位。6.4 第四步问题归因与策略调整根据 profiling 结果问题通常归于以下几类计算瓶颈如果 MoK 内核本身执行时间很长且 GPU 利用率高可能是专家网络计算量过大或者内核融合/优化不够。考虑是否要减少专家容量、降低模型维度。通信瓶颈如果通信操作耗时占比极高GPU 大量时间在等待。检查数据量是否批次过大或激活的专家数k过多导致通信数据膨胀检查拓扑在跨机训练时网络InfiniBand带宽是否成为瓶颈MoK 的设计可能更优化机内 NVLink跨机通信需要额外评估。重叠不足分析时间线看通信是否能与计算更好地重叠。MoK 应该在这方面做了优化但如果你的使用方式不对比如频繁同步可能破坏了这种重叠。负载不均衡如果某些 GPU 的 kernel 执行时间明显长于其他 GPU。检查路由路由算法是否导致某些专家过载可以输出每个专家的负载统计信息。检查容量因子专家容量因子是否设置过小导致溢出和丢弃破坏了负载均衡我个人的经验是大规模训练的问题十有八九出在通信和负载均衡上。计算瓶颈相对容易通过缩放模型规模来预估而通信和动态负载的复杂性往往超出预期。MoK 的价值正是通过一个精心设计的 Megakernel系统性地缓解这两个核心痛点。7. 总结与展望MoK 对普通开发者的启示Cursor 开源 Mixture-of-Kittens其意义远不止于公布一个针对 GB300 NVL72 的优化内核。它更像一个技术宣言展示了当大模型训练进入“硬核工程”深水区后我们需要怎样的思维方式。对于没有超算集群的绝大多数开发者我们可以从中汲取以下几点确定性是可调试性的基石无论项目大小尽力让你的训练过程可复现。固定随机种子、使用确定性算法、记录完整的配置和环境这些习惯在排查诡异问题时能救命。MoK 将确定性作为首要目标值得我们学习。理解硬件是突破性能瓶颈的关键你不一定需要为 NVLink 拓扑写内核但你需要了解你的 GPU 内存带宽、PCIe 带宽、CPU-GPU 数据传输成本。优化数据加载、减少不必要的 CPU-GPU 拷贝、合理使用混合精度这些都是在“吃透硬件”思想下的具体实践。通信是分布式训练的“阿喀琉斯之踵”在你的多卡训练任务中试着用 Profiler 看看all_reduce花了多少时间。思考你的数据并行、模型并行策略是否引入了过多的同步点。MoK 对通信-计算重叠的极致追求提醒我们时刻关注 GPU 的“空闲等待时间”。负载均衡是动态系统的永恒课题MoE 是负载均衡的极端案例。在你的系统中是否有类似“热点”问题比如微服务中某个实例压力过大或者数据处理流水线中某个环节成为瓶颈动态、自适应的负载均衡策略是构建稳健系统的核心。最后回到 MoK 项目本身。它目前是一个高度特化的、面向顶尖硬件的解决方案。它的直接应用范围很窄但它的设计思想、代码实现和性能数据为整个行业设立了一个标杆。随着更多厂商推出大规模集成 GPU 的服务器这类深度硬件协同的 Megakernel 可能会成为大模型训练基础设施中的标准组件。对于学习者我建议将 MoK 的代码仓库作为一个“高级阅读材料”。不必急于运行它而是去阅读它的设计文档、核心算法描述甚至尝试理解其中关键的 CUDA 内核代码片段。这个过程本身就是向大规模AI系统工程的深处迈进了一步。