如果你关注大模型训练最近一定被各种 MoEMixture of Experts架构刷屏了。从 Google 的 Switch Transformer 到 Mistral AI 的 Mixtral再到国内各大厂商的“专家混合”模型MoE 似乎成了降低大模型训练与推理成本的“银弹”。但当你真正想动手训练一个 MoE 模型或者想理解其底层效率时往往会发现理论很美好工程实现却是一团乱麻。专家路由的动态性、GPU 显存的碎片化、通信开销的不可预测性让 MoE 训练的效率常常远低于理论峰值。这不仅仅是学术问题而是真金白银的成本问题。训练一个千亿参数模型动辄需要数百万美元的计算资源。效率提升 10%可能就意味着节省数十万美元。今天我们要讨论的Cursor 开源的 Mixture-of-Kittens (MoK)正是瞄准了这个工程深水区。它不是一个新模型架构而是一个面向 NVIDIA GB300 NVL72 这类超大规模 GPU 集群的、确定性的 MoE 训练 Megakernel。简单来说MoK 想解决一个核心矛盾如何让 MoE 训练在超大规模集群上像训练稠密Dense模型一样稳定、高效且可预测它通过一个精心设计的“巨型内核”Megakernel将 MoE 层中复杂的专家选择、数据分发、计算与通信融合成一个高度优化的单一 GPU 核函数执行单元从而实现了前所未有的确定性与性能。本文将为你深入拆解 MoK 的核心思想、技术原理并提供一个基于 PyTorch 的简化概念实现帮助你理解“确定性 MoE 训练”究竟意味着什么以及它如何改变我们构建大模型的方式。无论你是算法研究员、机器学习工程师还是对底层系统优化感兴趣的高阶开发者这篇文章都将为你提供扎实的技术洞察和可操作的代码参考。1. 为什么 MoE 训练如此“不确定”且低效在深入 MoK 之前我们必须先理解传统 MoE 训练的痛点。MoE 的核心思想是“分而治之”一个输入样本只激活整个网络中的一小部分“专家”即子网络从而大幅减少每次前向传播的计算量。1.1 传统 MoE 的经典流程与瓶颈一个标准的 MoE 层通常包含以下步骤门控网络Gating Network计算输入样本应该分配给哪些专家。专家选择Expert Selection根据门控分数选择 top-k 个专家通常 k1 或 2。数据分发Data Dispatching将输入数据根据选择结果路由到对应的专家设备上。专家计算Expert Computation每个专家在其所在的设备上独立处理分配到的数据。结果聚合Result Aggregation将各个专家的计算结果汇总作为该层的输出。这个过程在单卡或小规模集群上问题不大但在 GB300 NVL7272 个 GPU 通过 NVLink 高速互联的巨型机架这样的规模下瓶颈立刻显现动态负载不均衡每个批次的样本激活的专家组合是随机的。可能导致某些 GPU专家过载而其他 GPU 闲置形成“木桶效应”。通信开销爆炸样本需要根据路由结果在 GPU 间频繁搬运。这种 All-to-All 式的通信模式其开销高度依赖于路由结果难以预测和优化。内核启动开销大上述每个步骤门控、分发、计算、聚合都可能涉及多次 GPU 内核启动和同步产生了大量细粒度的开销。显存碎片化由于每个专家处理的数据量动态变化为其分配的显存缓冲区可能利用率低下。这些因素共同导致了MoE 训练的不确定性同样规模的集群和模型不同批次的训练时间可能波动很大给资源调度、成本估算和调试带来了巨大困难。1.2 MoK 的破局思路Megakernel 与确定性Cursor 的 MoK 提出了一个根本性的解决方案将整个 MoE 层的前向传播甚至反向传播融合进一个单一的、庞大的 GPU 内核Megakernel中。这个思路的精妙之处在于消除内核启动开销所有操作在一个内核内完成避免了多次内核启动和同步的延迟。实现确定性调度在 Megakernel 内部可以预先规划好所有样本的路由、计算和通信顺序使得无论输入数据如何执行路径和耗时都是固定的。最大化硬件利用率通过对计算和通信的精细交织Interleaving可以几乎打满 GPU 的计算单元和 NVLink 的带宽逼近硬件的理论峰值。优化显存访问在单个内核内管理数据生命周期可以减少全局显存的访问次数充分利用共享内存和寄存器。简单类比传统的 MoE 训练像是一个混乱的十字路口每个车辆数据都要临时决定走哪条路专家导致交通灯频繁切换内核启动和拥堵负载不均。而 MoK 则像是一个高度智能的中央调度系统在车辆出发前就规划好了所有路线和通行时刻表让整个车流平稳、确定地通过。2. Mixture-of-Kittens (MoK) 核心架构解析MoK 的设计目标是在超大规模集群上提供确定性的、接近硬件峰值的 MoE 训练性能。其架构核心围绕“Megakernel”展开。2.1 系统层级视图一个完整的 MoK 训练系统可能包含以下层次分布式运行时层负责在 GB300 NVL72 这样的多机多卡环境下进行初始化和资源管理。Megakernel 调度层这是 MoK 的核心。它将一个 MoE 层的所有操作门控、路由、计算、通信描述为一个静态的数据流图并编译成可在所有 GPU 上协同执行的单一内核。确定性通信原语定制化的 All-to-All 通信操作其通信量、通信伙伴和时序在编译期即可确定避免了运行时动态决策的开销。专家计算内核库针对不同专家类型如 FFN、Attention高度优化的计算内核能够无缝嵌入到 Megakernel 的数据流中。2.2 Megakernel 的关键技术编译时静态规划MoK 的“确定性”源于编译时Compile-time的静态分析。在模型定义好后、训练开始前MoK 的编译器会进行如下分析专家布局分析确定每个专家被放置在哪个 GPU 上。MoK 可能采用均衡的、固定的布局策略。通信模式推导根据模型结构如专家数量、top-k值和批次大小推导出理论上可能发生的所有通信模式。注意虽然样本路由是动态的但所有可能的通信链路集合是静态的。内核融合与调度将门控计算、数据打包、通信发送/接收、专家计算、结果解包与聚合等操作按照数据依赖关系融合成一个宏大的计算图。并为这个图中的每个操作分配具体的执行时间片和硬件资源如流多处理器 SM。生成确定性执行计划最终输出一个“执行计划”这个计划精确规定了每个 GPU 在每一个时钟周期或更粗粒度的时间步应该执行什么操作。只要输入批次大小固定这个计划就固定执行时间也就确定。2.3 与现有框架如 Megatron-LM, DeepSpeed的对比现有的分布式训练框架如 Megatron-LM 的 Tensor Parallelism Pipeline Parallelism DeepSpeed 的 ZeRO主要解决的是稠密模型的切分与通信问题。它们处理 MoE 时通常是在其并行策略之上“嫁接”一个 MoE 层其路由和通信是动态的、框架感知度不高的。MoK 则是一种MoE-First的设计。它从第一性原理出发为 MoE 这个特定范式重新设计了整个执行引擎。可以认为MoK 是比上述框架更底层的“内核级”优化。未来MoK 完全可以作为 Megatron 或 DeepSpeed 中 MoE 层的后端实现为其提供确定性的高性能支撑。3. 环境准备与概念验证由于 MoK 是面向 GB300 NVL72 级集群的系统级项目完全复现其环境对绝大多数开发者不现实。但我们可以搭建一个概念验证环境在单机多卡甚至模拟环境下理解其“确定性调度”和“内核融合”的核心思想。我们将使用PyTorch和CUDA 编程通过 PyTorch 的 Custom Ops来模拟一个简化版的 MoE 层并尝试将部分操作融合。3.1 基础环境操作系统: Ubuntu 20.04 LTS 或更高版本Python: 3.8PyTorch: 2.0 (需要支持torch.compile和torch.distributed)CUDA Toolkit: 11.8NVIDIA GPU: 至少 2 张具有 NVLink 的 GPU如 V100, A100, H100以获得更好的多卡通信体验。仅学习原理的话单卡也可。NCCL: 正确安装的高版本 NCCL 库用于分布式通信。3.2 项目结构初始化创建一个新的项目目录并初始化虚拟环境。mkdir mixture-of-kittens-demo cd mixture-of-kittens-demo python -m venv venv source venv/bin/activate # Linux/Mac # venv\Scripts\activate # Windows pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install ninja # 用于加速CUDA扩展编译4. 实现一个简化的确定性 MoE 层PyTorch 版我们将实现一个名为DeterministicMoE的模块。它不会实现真正的 Megakernel但会体现“确定性”的核心预先计算好所有可能的路由和通信计划并在每次前向传播中复用这个计划。4.1 核心数据结构路由表与通信计划# deterministic_moe.py import torch import torch.nn as nn import torch.nn.functional as F from typing import List, Optional import numpy as np class DeterministicMoE(nn.Module): 一个简化的确定性MoE层概念实现。 核心思想在初始化时根据批次大小和专家布局预先计算一个确定性的“路由与通信计划”。 在前向传播时严格按计划执行避免动态决策。 def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_experts: int, top_k: int 2, capacity_factor: float 1.0, num_local_experts: int -1, expert_group: Optional[torch.distributed.ProcessGroup] None): super().__init__() self.input_dim input_dim self.hidden_dim hidden_dim self.output_dim output_dim self.num_experts num_experts self.top_k top_k self.capacity_factor capacity_factor self.expert_group expert_group # 1. 创建专家网络。每个专家是一个简单的FFN。 # 在实际MoK中每个专家可能分布在不同的GPU上。 self.experts nn.ModuleList([ nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, output_dim) ) for _ in range(num_experts) ]) # 2. 门控网络 (Gating Network) self.gate nn.Linear(input_dim, num_experts, biasFalse) # 3. 关键路由计划缓存初始为空在第一个批次后生成 self.register_buffer(routing_plan, None) # 缓存路由索引 self.register_buffer(comm_send_counts, None) # 发送数据量 self.register_buffer(comm_recv_counts, None) # 接收数据量 self.plan_batch_size -1 # 当前计划对应的批次大小 # 4. 用于模拟分布式环境的rank和world_size self.rank 0 self.world_size 1 if torch.distributed.is_initialized() and expert_group is not None: self.rank torch.distributed.get_rank(groupexpert_group) self.world_size torch.distributed.get_world_size(groupexpert_group) # 简化假设专家均匀分布在各个rank上 self.num_local_experts num_local_experts if num_local_experts 0 else (num_experts // self.world_size) self.local_expert_indices list(range(self.rank * self.num_local_experts, (self.rank 1) * self.num_local_experts)) else: self.num_local_experts num_experts self.local_expert_indices list(range(num_experts)) def _build_deterministic_plan(self, batch_size: int): 构建确定性路由与通信计划。 这是一个简化的模拟我们假设门控权重是固定的或变化很小 因此可以基于初始门控权重计算一个‘典型’路由并固定下来。 真实MoK会在编译期进行更复杂的静态分析。 print(f[Rank {self.rank}] Building deterministic plan for batch_size{batch_size}) # 为简单起见我们使用随机但固定的门控权重来生成一个“代表性”路由 # 在实际中这可能基于对训练数据分布的统计分析 torch.manual_seed(42) # 固定随机种子以确保确定性 dummy_input torch.randn(batch_size, self.input_dim, deviceself.gate.weight.device) with torch.no_grad(): logits self.gate(dummy_input) # [batch_size, num_experts] scores F.softmax(logits, dim-1) # Top-k 专家选择 topk_vals, topk_indices torch.topk(scores, kself.top_k, dim-1) # [batch_size, top_k] # --- 核心生成确定性的路由索引映射 --- # 我们将每个样本分配给top-k专家并考虑专家容量限制 expert_capacity int(self.capacity_factor * batch_size / self.num_experts) # 初始化路由表: [num_experts, expert_capacity, 2] - (batch_idx, score) routing_table -torch.ones((self.num_experts, expert_capacity, 2), devicescores.device) # 一个简单的确定性分配算法例如按分数排序后顺序填充 # 注意这是一个简化模拟真实MoK的分配算法会更复杂以平衡负载。 for batch_idx in range(batch_size): for k in range(self.top_k): expert_idx topk_indices[batch_idx, k].item() score topk_vals[batch_idx, k].item() # 找到该专家第一个空位 cap_idx (routing_table[expert_idx, :, 0] 0).nonzero(as_tupleTrue)[0] if cap_idx.numel() 0 and cap_idx[0] expert_capacity: routing_table[expert_idx, cap_idx[0], 0] batch_idx routing_table[expert_idx, cap_idx[0], 1] score # 根据路由表生成每个rank需要发送/接收的数据索引 # 这里简化我们只计算本rank拥有的专家需要处理哪些样本 send_indices [] recv_indices [] # ... (详细的索引计算逻辑涉及分布式通信规划) # 由于篇幅此处省略复杂的跨设备索引计算仅示意 # 真实实现需要利用 torch.distributed 进行 all-to-all 的规划。 # 缓存计划 self.routing_plan { topk_indices: topk_indices, # 缓存下来后续直接使用 topk_vals: topk_vals, routing_table: routing_table, expert_capacity: expert_capacity, } self.plan_batch_size batch_size print(f[Rank {self.rank}] Plan built. Expert capacity: {expert_capacity}) def forward(self, x): 确定性的前向传播。 如果批次大小与缓存计划不符则重新构建计划。 然后严格按照计划进行路由和计算。 batch_size x.shape[0] if self.routing_plan is None or batch_size ! self.plan_batch_size: self._build_deterministic_plan(batch_size) plan self.routing_plan topk_indices plan[topk_indices] # 直接使用预计算的路由 topk_vals plan[topk_vals] routing_table plan[routing_table] expert_capacity plan[expert_capacity] # 1. 根据路由表将输入数据分发到对应的专家缓冲区 # 创建专家输入缓冲区: [num_local_experts, expert_capacity, input_dim] expert_inputs torch.zeros((self.num_local_experts, expert_capacity, self.input_dim), devicex.device, dtypex.dtype) expert_input_counts torch.zeros(self.num_local_experts, devicex.device, dtypetorch.long) # 填充缓冲区 (简化版仅处理本地专家) for local_exp_idx, global_exp_idx in enumerate(self.local_expert_indices): mask routing_table[global_exp_idx, :, 0] 0 # 该专家需要处理的样本 batch_indices routing_table[global_exp_idx, mask, 0].long() if batch_indices.numel() 0: count batch_indices.shape[0] expert_inputs[local_exp_idx, :count] x[batch_indices] expert_input_counts[local_exp_idx] count # 2. 本地专家计算 expert_outputs [] for local_exp_idx in range(self.num_local_experts): count expert_input_counts[local_exp_idx] if count 0: inp expert_inputs[local_exp_idx, :count] # [count, input_dim] out self.experts[self.local_expert_indices[local_exp_idx]](inp) # [count, output_dim] expert_outputs.append((local_exp_idx, out)) else: expert_outputs.append((local_exp_idx, None)) # 3. 将计算结果根据路由表写回最终输出张量 final_output torch.zeros((batch_size, self.output_dim), devicex.device, dtypex.dtype) # 这里简化了加权求和的过程... # 真实情况需要根据topk_vals进行加权并处理多个专家对同一样本的贡献。 # 4. 模拟一个确定性的All-to-All通信此处仅为示意实际需调用NCCL if self.world_size 1: # 在真实MoK中通信的缓冲区大小、通信伙伴在计划中是确定的。 # 这里我们只是模拟一个同步点强调其确定性。 torch.distributed.barrier(groupself.expert_group) # 为了演示我们暂时只返回一个简单结果 # 实际应返回聚合了所有专家结果的 final_output return final_output.mean(dim0, keepdimTrue).expand(batch_size, -1) # 临时返回均值这个简化实现的核心在于_build_deterministic_plan函数。它在第一个批次或批次大小变化时根据固定的随机种子生成一个确定性的路由计划topk_indices,routing_table并在后续所有批次中复用这个计划。这就模拟了 MoK “编译时规划运行时执行”的确定性思想。5. 编写测试脚本验证确定性让我们编写一个测试脚本来验证在相同输入下我们的DeterministicMoE层是否真的产生了确定性的输出即使内部有随机初始化的门控网络。# test_deterministic.py import torch import torch.distributed as dist import os from deterministic_moe import DeterministicMoE def setup_distributed(rank, world_size): 初始化分布式环境用于多卡测试 os.environ[MASTER_ADDR] localhost os.environ[MASTER_PORT] 29500 dist.init_process_group(nccl, rankrank, world_sizeworld_size) torch.cuda.set_device(rank) def test_determinism_single_gpu(): 单GPU确定性测试 print( Single GPU Determinism Test ) torch.manual_seed(1234) model DeterministicMoE(input_dim16, hidden_dim32, output_dim16, num_experts4, top_k2).cuda() model.train() # 创建固定输入 dummy_input torch.randn(8, 16, devicecuda) # 第一次前向传播 (会构建计划) with torch.no_grad(): output1 model(dummy_input) print(fOutput 1 mean: {output1.mean().item():.6f}) # 重置模型状态但计划已缓存 # 改变门控网络的权重模拟训练中的更新 with torch.no_grad(): model.gate.weight.data 0.01 * torch.randn_like(model.gate.weight) # 第二次前向传播 (使用缓存的计划尽管权重变了) with torch.no_grad(): output2 model(dummy_input) # 注意这里使用的还是旧的、基于旧权重的计划 print(fOutput 2 mean (with same plan, diff weights): {output2.mean().item():.6f}) print(fAre outputs identical? {torch.allclose(output1, output2, rtol1e-5)}) # 预期输出相同因为路由计划是固定的。 # 第三次前向传播但输入批次大小变化触发重新规划 dummy_input_new torch.randn(12, 16, devicecuda) with torch.no_grad(): output3 model(dummy_input_new) print(fOutput 3 mean (new batch size): {output3.mean().item():.6f}) # 此时会重新构建计划输出基于新的但依然是确定性的计划。 def run_multi_gpu_test(rank, world_size): 多GPU分布式测试函数 setup_distributed(rank, world_size) torch.manual_seed(1234 rank) # 不同rank不同种子但模型同步后应一致 # 创建一个进程组用于MoE专家并行 expert_group dist.new_group(rankslist(range(world_size))) model DeterministicMoE(input_dim16, hidden_dim32, output_dim16, num_experts8, top_k2, expert_groupexpert_group).cuda() # 注意需要将模型参数在不同rank间同步以确保一致性此处简化 # dist.broadcast(model.gate.weight, src0, groupexpert_group) # for exp in model.experts: # for p in exp.parameters(): # dist.broadcast(p, src(p专家所属的rank), groupexpert_group) dummy_input torch.randn(8, 16, devicefcuda:{rank}) # 确保所有rank输入相同仅测试用 if rank 0: input_to_bcast dummy_input else: input_to_bcast torch.empty_like(dummy_input) dist.broadcast(input_to_bcast, src0, groupexpert_group) dummy_input input_to_bcast with torch.no_grad(): output model(dummy_input) print(f[Rank {rank}] Output shape: {output.shape}, mean: {output.mean().item():.6f}) dist.destroy_process_group(expert_group) dist.destroy_process_group() if __name__ __main__: # 测试1单卡确定性 test_determinism_single_gpu() # 测试2多卡分布式需要至少2张GPU world_size torch.cuda.device_count() if world_size 2: print(f\n Multi-GPU Distributed Test (World Size: {world_size}) ) import torch.multiprocessing as mp mp.spawn(run_multi_gpu_test, args(world_size,), nprocsworld_size, joinTrue) else: print(\nSkipping multi-GPU test (requires 2 GPUs).)运行测试脚本python test_deterministic.py预期观察结果在单卡测试中output1和output2应该完全一致尽管门控网络的权重在两次前向之间发生了变化。这是因为我们复用了第一次前向时生成的确定性路由计划。这直观地展示了“确定性”的含义执行路径不随权重微小变化而改变。当批次大小从 8 变为 12 时会触发重新构建计划 (_build_deterministic_plan)因此output3是基于新计划的输出。在多卡测试中如果正确实现了参数同步所有 rank 对于相同输入应计算出相同或可聚合为相同的输出。6. 从概念到现实MoK 的真正挑战与价值我们的简化实现仅仅揭示了确定性 MoE 的冰山一角。真正的 MoK 在 GB300 NVL72 上面临的挑战和带来的价值远超于此。6.1 核心工程挑战巨型内核Megakernel编译如何将包含条件分支、循环、设备间通信的复杂数据流编译成一个高效的、占用大量寄存器且能并发执行数万个线程的单一 GPU 内核这需要极其先进的编译器技术如 Triton、CUDA Graph 与静态调度的结合。通信与计算重叠在 Megakernel 内部需要精细调度使得 GPU 在等待 NVLink 数据时能立刻切换到其他就绪的计算任务完全隐藏通信延迟。这需要对硬件流水线有深刻理解。负载均衡的静态保证如何在编译时就设计路由策略确保任何输入分布下各个专家的负载都是均衡的这可能需要引入“辅助负载”或动态容量因子调整的静态版本。与现有生态集成如何让 MoK 与 PyTorch、JAX 等主流框架的 Autograd 系统兼容如何支持复杂的模型架构如 MoE 与 Attention 交错6.2 带来的核心价值可预测的训练时间对于云服务商和大型实验室确定性的训练时间意味着更精确的资源预留和成本核算避免了因性能波动导致的资源浪费或任务延期。极致的硬件利用率通过编译期的全局优化MoK 可以理论上逼近 GB300 NVL72 集群的算力与通信带宽的峰值将每美元的计算效率提升到新的高度。降低调试难度确定性的执行使得 bug 复现和性能 profiling 变得简单。如果一次训练迭代慢了可以精确地定位到是哪个静态环节出了问题。赋能新的模型探索当训练效率不再是瓶颈时研究人员可以更自由地探索更庞大、更复杂的 MoE 架构例如更多专家、更动态的路由策略而不必过分担心工程实现的复杂性。7. 常见问题与排查思路问题现象可能原因排查方式解决方案训练结果不稳定即使种子固定1. 路由计划未正确缓存或复用。2. 分布式环境下参数未同步。3. 数据加载中存在非确定性操作。1. 检查plan_batch_size是否与当前批次匹配。2. 检查分布式初始化代码确保expert_group正确。3. 检查 DataLoader 的worker_init_fn和generator。1. 确保_build_deterministic_plan只在批次大小变化时调用。2. 在训练开始时广播所有模型的初始参数。3. 为 DataLoader 设置固定的随机种子。多卡测试时程序挂起或报错1. NCCL 通信死锁。2. 各 rank 执行流不一致例如条件分支不同。3. 显存不足。1. 使用NCCL_DEBUGINFO环境变量运行。2. 在每个通信操作前后添加打印注意同步。3. 使用nvidia-smi监控显存。1. 确保所有 rank 都按相同顺序调用相同的通信原语。2. 确定性计划必须保证所有 rank 的逻辑完全一致。3. 减少专家容量因子或批次大小。性能未提升甚至下降1. 模拟的确定性计划生成开销过大。2. 专家计算过于简单无法掩盖内核融合带来的收益。3. 通信模拟如dist.barrier成为瓶颈。1. 使用 PyTorch Profiler 分析耗时。2. 对比关闭确定性计划即动态路由时的性能。1. 计划生成应尽可能简单或离线进行。2. 增加专家网络的复杂度如更大的 hidden_dim。3. 真实部署需使用高度优化的 NCCL All-to-All。_build_deterministic_plan中的expert_capacity计算导致溢出或利用率低capacity_factor设置不合理。打印每个专家的实际负载与expert_capacity的比值。调整capacity_factor。通常 1.0如1.1~1.5以容纳负载波动但过大会浪费显存。8. 最佳实践与工程建议如果你想在自己的研究或项目中借鉴 MoK 的思想可以参考以下建议从分析瓶颈开始不要盲目追求“确定性”。先用 Profiler如 PyTorch Profiler, NSight Systems分析你的 MoE 模型确认瓶颈到底是在动态路由、内核启动、还是通信上。分层实现可以先在PyTorch CUDA Graphs的层面实现一个“准确定性”版本。CUDA Graphs 可以捕获一次迭代的计算图并重复执行能消除内核启动开销是实现 Megakernel 思想的良好起点。关注编译器技术深入学习Triton这类 GPU 编程语言和编译器。它允许你以更高的抽象级别编写可融合的内核是构建复杂 Megakernel 的利器。小规模验证先在 2-4 张 GPU 的小集群上验证确定性调度和通信重叠的逻辑正确性再扩展到大规模集群。与现有框架结合考虑将你的优化实现为PyTorch 的 Custom Autograd Function或JAX 的 custom_vjp这样可以嵌入到现有的模型训练流水线中利用成熟的优化器、分布式数据并行等组件。设计灵活的接口即使底层是确定性的也应向上层提供灵活的配置接口如允许用户选择不同的路由算法如Top-2 Gating,Noisy Top-K Gating和负载均衡损失函数。9. 总结Cursor 开源的 Mixture-of-Kittens (MoK) 项目将 MoE 训练从“动态的、难以预测的”领域推向了一个“确定性的、可极致优化”的新阶段。其核心Megakernel思想——通过编译时静态规划将整个 MoE 层的计算与通信融合为一个单一的、确定性的执行单元——为解决超大规模 AI 训练的效率与成本问题提供了极具潜力的方向。本文通过一个简化的 PyTorch 实现揭示了确定性 MoE 的核心在于将动态决策提前到编译期或初始化阶段生成一个固定的执行计划并复用。虽然真正的 MoK 涉及极其复杂的编译器、运行时和硬件调度技术但其底层逻辑是清晰且可理解的。对于广大开发者和研究者而言MoK 的价值不仅在于其代码本身更在于它展示了一种系统级优化的范式当算法MoE遇到规模瓶颈时从硬件和编译器的角度进行跨层协同设计往往能带来数量级的提升。在 AI 基础设施越来越成为竞争焦点的今天深入理解像 MoK 这样的项目将帮助你站在技术演进的最前沿。你可以从本文提供的概念代码出发逐步深入 CUDA 编程、分布式通信和编译器领域探索属于你自己的“确定性优化”之路。