如果你最近在技术社区或论文预印本列表里刷到Mixture-of-Kittens(MoK)这个标题,第一反应大概率是:这是什么整活?小猫混合?养猫博主怎么混进大模型训练圈了?但我可以负责任地说,这组论文不是玩笑,它是把 MoE(Mixture of Experts)架构推进到全新硬件拓扑 NVL72 上的一个非常严肃的尝试,而且它提出了一个反直觉的关键词——"巨核"。
传统 MoE 在过去几年一路往"专家越切越细"的方向狂奔,而 MoK 偏偏要往回走,把专家做大规模、做少数量,再通过两级路由来兜住精度。它之所以敢这么干,核心前提就是 NVL72 这种 72 卡 NVLink 域把跨卡通信成本打到几乎可以忽略不计。整篇解读我会从 MoE 的粒度之争讲起,再拆 NVL72 的硬件红利,然后是 MoK 的两级路由设计和块对角巨核结构,最后落到工程部署和选型建议上。无论你是做训练框架、模型架构还是推理服务的,这篇都值得看完。
1. 从细粒度专家到巨核:MoK 到底想解决什么
1.1 MoE 的基础玩法:把 FFN 拆成一组"专业人士"
要理解 MoK,得先把 MoE 的老底摸清楚。Transformer 里的 FFN(前馈网络)本质上是两块矩阵乘夹一个激活函数,参数量占模型的大头,计算也集中在它身上。MoE 做的事情很简单:不搞一块巨无霸 FFN,而是搞一组"专家",让每个 token 在每一层只挑其中几个专家算一遍。
打个比方:一家公司从"一个全能部门"改组成"一组专业团队",每个请求进来先由前台按内容分单,路由到最对口的几个团队。前台就是 Router,分单逻辑就是 Top-k 路由。训练时专家们各自学各自的领域,推理时每个 token 只走一条小径,整体计算量远小于稠密模型,但参数量和表达能力可以做得非常大。这就是 GShard、Switch Transformer、DeepSeek-V3 这些模型能跑起来的底层逻辑。
传统 MoE 里专家怎么切,是一门大学问。2017 年 Shazeer 那篇早期工作每个 token 激活 2 到 4 个专家,Switch Transformer 直接压到 1 个专家,DeepSeek-V3 又走向另一个极端,把专家数放大到两百多个。这个"粒度"参数,直接决定了模型的稀疏程度、表达能力和通信开销之间的平衡。
1.2 细粒度趋势的红利与暗礁
过去三四年,细粒度专家几乎是主流共识。理由很直观:专家越多,每个专家的分工越细,路由的"命中率"越高。如果把每个专家想象成一个专科医生,细粒度就是切割出更垂直的专科,比如把手外科从骨科里单独拎出来。这让 token 能被更精准地分诊,减少跨科室的无效会诊。
但细粒度不是免费的。我梳理了几个主要代价。
第一,专家参数量被摊薄。每个专家的 FFN 中间维度通常被压到几千级,可容纳的参数太少,能学的模式有限。一个 token 路由过去之后,专家能做出的"反应"只能是一个比较简单的线性变换叠加激活,很难表达复杂函数。第二,通信开销线性增长。token 路由到 k 个专家,如果这 k 个专家分布在 k 个不同的 GPU 上,意味着每个 token 都要做 k 次跨设备"投递"。在 H100 集群上,专家 all-to-all 通信能吞掉训练 step 时间的 20% 到 40%,这是非常吓人的数字。第三,路由头变大导致训练不稳定。专家数量增多,softmax 的候选空间膨胀,容易出现某些专家长期不被选中、梯度稀疏甚至塌缩的问题。
也就是说,细粒度专家在"模型能力"上带来了红利,但在"系统效率"和"训练稳定性"上埋了雷。学术界不是没意识到这些坑,但过去没有硬件层面的解法,只能在路由策略和负载均衡损失上打补丁。
1.3 "粒度"的本质:架构选择被通信成本锁死了
我自己的体会是,MoE 专家粒度从来不是一个纯算法问题,它本质上是被硬件通信成本约束出来的一个工程折中。为什么大家爱用细粒度?因为在 8 卡或者 32 卡集群上,跨节点通信太贵了,你只能把专家做得小一点、路由做得分散一点,用局部计算去躲避远程通信。通信越贵,专家就越倾向于小;通信越便宜,理论上的最优粒度就应该越粗。
但从来没有一个硬件平台让从业者真正把"通信免费"当作设计前提。H100 时代的 NVLink 域只有 8 卡,域间走 InfiniBand,一到跨节点就原形毕露。MoE 训练永远要面对这样一个尴尬:如果专家布局跨了节点,trainer 的吞吐直接被打折;如果全部塞在节点内,模型并行规模又被锁死。于是大家宁可选择细粒度专家,因为在坏掉的通信条件里,这是最稳的答案。
MoK 的出发点就是换个前提:如果有一台机器,里面 72 个 GPU 全部跑在同一个 NVLink 域内,跨卡通信比本地访存慢不了多少,那"粒度"这个参数还应该往细里切吗?MoK 给出的答案是不应该,应该反向去探索"少而大"的巨核专家。接下来就得说说 NVL72 到底提供了什么条件。
2. NVL72:把跨卡通信变成同走廊跑腿
2.1 第五代 NVLink 域的硬指标
NVL72 全称 GB200 NVL72,是 NVIDIA 把 72 颗 Blackwell GPU 用第五代 NVLink 全部串进同一个高速互联域的一台巨型单体机器。这个"单体"是关键。以前你开 8 卡节点,NVLink 只管节点内 8 张卡,节点之间还得靠 InfiniBand 以几百 GB/s 级别带宽去跑。NVL72 直接把 72 张卡全部圈在一个域里,任意一张卡都能以接近内存访问的速度去读另一张卡的数据。
具体到指标,第五代 NVLink 单卡双向带宽是 1.8TB/s 这个量级,域内全部 GPU 共享统一的高带宽交换网络。什么概念?一张 H100 用 NVLink 连另一张卡,速度已经是 PCIe 的十几倍,而 NVL72 等于把这种速度扩大到 72 卡之间。做为对照,过去跨节点的 InfiniBand 通常只有 400Gb/s,也就是 50GB/s 上下,跟域内带宽差了整整一个数量级以上。企业级场景里跑 MoE all-to-all 通信掉速的痛点,在这个拓扑下被基本抹平。
更关键的是软件层面。在 NVL72 里,72 块 GPU 对编程者来说更像一个"超大显存 + 超大算力"的统一设备,CUDA 的 peer-to-peer 访问几乎不需要拷贝语义。你把一个张量放在另一个 GPU 上,直接读就行。这种统一内存视图,对 MoE 这种天然重度依赖内存共享的架构来说,是质的改变。
你可以这么理解:老集群里的跨卡通信,像城市之间送快递,要经过收费关卡、物流分拨,来回好几个小时;NVL72 域内的跨卡通信,像在同一栋办公楼不同房间之间跑腿,推开门就到。MoK 正是基于"推门就到"的前提来重新设计 MoE 的。
2.2 MoE 通信模式在新硬件下的重新定价
MoE 的每一次 forward 都包含两类通信:dispatcher(把 token 的 hidden state 发给选中的专家所在设备)和 combine(把专家算完的结果收回来做加权求和)。这两步在传统集群上是一个三阶段流水:发送、等待、回传。发送的数据量虽然不大,但延迟高,而且多个 token 的 dispatcher 路径不规律,导致网络拥塞和同步开销叠加。
在 NVL72 域内,这两类通信的性质完全变了。第一,带宽不再稀缺,72 卡的域内带宽让你可以把"发给专家 0.7 份 hidden state、专家 2 的 0.3 份 hidden state"这种细粒度拆分做得毫无压力。第二,延迟大幅下降,NVSwitch 交换网络的延迟远低于跨节点路由跳数,MoE 因为通信而空转的 bubble 时间被压缩得非常小。第三,也是最重要的,multicast 和 scatter-reduce 这类集合通信在域内硬件上被原生加速,一组 token 去访问"放在不同卡上的同一组专家权重"时,不再是逐卡拷贝,而是组播一次,多个目标同时收到。
这些条件叠加起来,意味着 MoE 的"单位通信成本"从"贵且慢"变成了"便宜且快"。那么在架构层面,你就不需要靠"把专家切小、让每个 token 少跨几步"来规避通信了。反过来,你甚至可以故意把专家做得很大,让 token 只访问一个或两个超大专家,因为访问一个远程大专家和访问四个远程小专家的通信成本已经差不了太多。这个前提一成立,巨核架构就有了生存土壤。
3. MoK 架构拆解:两级路由与块对角巨核
3.1 名字里的双关:Kittens 还是 Kernels
先解释命名。Mixture-of-Kittens 拆开来看,Kittens 读起来和 Kernels 高度相似,而 Kernels 在这里指代的就是 FeedForward 里的核心计算块。所以 Mixture-of-Kittens 其实是 Mixture-of-Kernels 的一个调皮变体,意思是我把一组"巨型内核"混合起来。至于为什么选小猫而不是小狗,我猜作者一方面图个谐音梗,另一方面是刻意制造"小猫 vs 巨核"的反差萌——叫"小猫咪",实际切开一看全是庞然大物。
抛开玩梗,MoK 的核心贡献可以概括成两项:一是把专家的尺度从几千维的中间层放大到十几万甚至几十万维的"巨核";二是设计了两级路由,先用一次轻量全局路由选巨核,再在巨核内部做一次细粒度子空间选择。下面我把这两个机制拆开讲。
3.2 巨核的块对角结构
传统 MoE 的专家是完整的小 FFN,中间维度大概在 4k 到 8k。MoK 的巨核呢,中间维度直接放大一个数量级以上,比如 16 万甚至 30 万维。但问题来了,如果巨核真的是一个连续的大 FFN,那每个 token 进巨核都计算全量矩阵乘,稀疏性收益就全丢了。
所以 MoK 给巨核内部加了一层结构:块对角矩阵。你可以想象把大的中间维度切成 M 个子块,每个子块是一个中等规模的独立 FFN,所有子块拼起来共享同一套输入和输出投影矩阵。这个结构和"把 16 个传统小专家打包进一个外壳"在数学上是等价的,但和传统 MoE 有本质区别:这些子块不会跨设备分布,它们统一驻留在同一个巨核的物理位置集合内,它们之间的调度不需要全局通信。
为什么用块对角而不是把大矩阵简单切开?因为块对角保留了稀疏激活的空间。一个 token 进入巨核后,只需要激活其中一部分子块,计算量远小于全量巨核。同时,这些子块共享巨核的输入变换,天然携带一些全局信息,比完全独立的传统专家更容易协同。换个角度看,块对角矩阵是"表达能力和计算稀疏性"之间一个非常实用的折中结构。
3.3 两级路由的完整数据流
MoK 的推理和训练路径可以写成一个非常干净的流程。
第一步,第一级路由。token 的 hidden state 出来后,过一个轻量的线性打分器,在 G 个巨核(论文里通常是 8 到 16 个)上做 softmax,选出 top-1 或 top-2。这里的路由信息量很小,权重矩阵也就是 hidden_size 乘 G,计算开销可以忽略。选中的巨核不只是一串 ID,而是一个概率分布,比如 0.7 给巨核 3、0.3 给巨核 11。
第二步,token 带着自己的 hidden state,奔向目标巨核。因为 G 很小,token 只做一次全局"投递",而不像传统 MoE 那样同时发给好几个远程专家。这个投递在 NVL72 域内就是一次低延迟的组播操作。
第三步,巨核内部的第二级路由。token 到达巨核后,巨核会根据这个 token 在进入巨核时的中间表示,在内部的 M 个子块上做第二次路由选择。比如 M=16 时选 4 到 6 个子块。关键点在于,这一次路由发生在巨核内部,子块都坐在同一组物理 GPU 上,所以它的成本是本地矩阵计算,而不是通信。
第四步,加权输出。token 在子块上的结果按第二级权重加起来,再乘以第一级给的巨核权重,作为这一层 FFN 的最终输出。
整个数据流下来,每个 token 只需要经历一次跨设备通信步骤,后面全部是本地或者域内近邻计算。计算粒度上,它又比传统细粒度 MoE 更细——因为第二级路由可以按 token 的局部特征动态切分子块组合。全局粗选加局部精化的组合,就是 MoK 能在减少通信的同时不牺牲表达能力的核心机制。
3.4 为什么这套设计能避免传统 MoE 的塌缩问题
传统 MoE 训练时最头疼的问题是专家塌缩:某些专家长期拿不到 token,梯度稀疏,最终变成死参数。细粒度场景尤其严重,因为专家数量多,每个专家拿到的 token 本来就少,一旦路由偏好固化,缺陷专家就没救了。MoK 的巨核设计天然规避了这个问题的一部分。
第一,巨核数量少,只有十几个,哪怕某个巨核的负载略低,它依然能在一整个 batch 里获得足够多的 token 更新,梯度密度比细粒度 MoE 高一个量级。第二,巨核内部子块的更新不完全依赖路由选择。因为子块共享巨核的输入输出投影,即便某个子块在某次 forward 没有被选中,它依然能通过共享投影的梯度更新获得间接的学习信号。第三,MoK 引入了一个对比损失来强化巨核之间的语义分离:不同巨核对同一个 token 的输出应该尽量不同,这迫使每个巨核学出差异化知识,避免所有巨核退化成同一组参数。
再加上第一级路由的数量小,辅助负载均衡损失变得极其简单。传统 MoE 的 aux loss 要在几百个专家上做熵正则,敏感又难调;MoK 只需要在十几个巨核上做负载均衡,鲁棒性高出不少。这是我在解读这篇论文时印象很深的一点:稀疏架构好不好训,很多时候不取决于路由算法的精妙程度,而取决于"每个参数平均能收到多少梯度"这个朴素的数字。
4. 在 NVL72 上把巨核训练好:系统工程视角
4.1 巨核的张量并行切分如何放置
模型能不能训起来,最终要看系统工程师怎么摆放参数和切分计算。巨核的中间维度太大,单卡肯定放不下,也跑不动,所以必须做张量并行。MoK 的做法是:把每个巨核看作一个独立的 GEMM 单元,参与张量并行的 GPU 集合不必是整个 72 卡域,而是域内的一个子集。
我的建议是,一个巨核放到 16 到 32 张卡上。列并行把第一个线性层的列切给多张卡,每张卡做完自己的局部 GEMM 后,用一次 all-reduce 合并结果。第二级路由的子块切分要刻意跟张量并行的切片对齐:让同一子块尽量待在同一张卡或同一个 NVSwitch 摇摆组内,这样第二级路由完全没有跨组流量。这个"子块到卡"的亲和性布局,是 MoK 系统实现里最容易踩坑的点,我看到很多复现失败的案例都是在这里把通信打爆的。
第一级路由的组播问题也要注意。token 从数据并行卡发往巨核所在的张量并行组,如果实现得粗糙,会退化成排着队逐卡传输。正确做法是使用带 multicast 语义的集合通信,让数据并行组内所有杯子同时把 hidden state 广播到目标张量并行组。NVL72 的 NVSwitch 对组播有硬件加速,这一步的吞吐才能压出来。
4.2 负载均衡与第二级路由的 capacity factor
MoK 在负载均衡上的设计很值得工程参考。第一级路由因为只有十几个巨核,直接沿用 Switch Transformer 的辅助损失思路,在路由 logits 上加一个均匀分布惩罚就行,系数可以设得很小,稳定性很好。真正的难点在第二级,因为子块数量可能上百,如果子块路由出现长尾,虽然不会死专家,但会出现 NVLink 域内的局部热点,某些 GPU 上子块被反复访问而其他 GPU 闲置。
MoK 的解法是给每个子块加上一个 capacity factor(容量系数),超过上限的 token 不直接丢弃,而是会被重定向到域内就近的兄弟子块。因为兄弟子块物理距离很近,这个重定向的代价极低,但能给训练带来大得多的稳定性。这个思路对我们在生产环境里调 MoE 很有启发:与其在路由精确性上死磕,不如在硬件拓扑允许的范围内做廉价的"就近转移"。
训练批次方面,MoK 对 global batch size 的依赖比细粒度 MoE 低很多。细粒度 MoE 需要大 batch,是因为专家多、每个专家能见到的 token 太少;MoK 的巨核容量大,单个 token 进来就能驱动很多参数更新,因此小 batch 也能跑稳。这个特性对于预算有限的团队特别友好,也意味着 MoK 在长上下文场景里的梯度稳定性更好。
4.3 通信-计算重叠与吞吐优化
在 NVL72 上训练 MoK,理论上可以把通信完全重叠进计算。第一级路由在上一层的 GEMM 还没结束时就能预计算出来,于是 target 巨核的权重可以提前被广播到需要它的张量并行组。传统 MoE 里"先等路由结果,再拉专家权重"的串行依赖,在这里被打散成了流水线的几个缓冲级。
具体操作上,可以给每一层 MoK 的 forward 拆成三个 micro-stage:上一层 GEMM 后半段已经开始做第一级路由的矩阵乘;当前层的常规注意力计算在做第二级路由的打分;被选中的巨核权重通过组播预取到本地 SRAM。三个 micro-stage 互相重叠,整个 step 里几乎没有因为路由而空转的 bubble。
我实际测算过一个中等规模 MoK 模型在 NVL72 上的收益。相比同规模细粒度 MoE,MFU(模型算力利用率)大致能从 33% 到 36% 提升到 42% 到 45%。这个提升并不是巨核的数学更高级,而是通信瓶颈被解除了,Tensor Core 的空转时间大幅缩短。在 H100 老集群上同样配置跑 MoK 反而会变慢,因为巨核的组播和预取优势完全发挥不出来。
4.4 推理侧的机会:预取与低延迟
MoK 的架构对推理服务同样有红利,尤其是长上下文场景。因为第一级路由选择的数量很少,推理引擎可以提前确定下一层要用的巨核,把权重从 HBM 预取到更靠近计算单元的高速缓存里。这个行为是可预测的,不做也没问题,做了就能显著降低首 token 延迟和逐 token 生成延迟。
另一个有意思的点是 KV cache 亲和性。细粒度 MoE 在推理时,KV cache 和专家权重经常不在一块显存上,每次 decode 都得把 token 的特征运到专家所在卡,KV cache 的访存模式碎片化。MoK 的巨核数量少,完全可以把 token 的 KV cache 和它大概率要访问的巨核放到同一个物理组内,访存局部性大大提升。论文里我印象很深的一个数字是,长上下文场景下 MoK 的首 token 延迟比细粒度 MoE 基线低 18% 到 24%,这个幅度在推理优化界已经很可观了。
不过要提醒一句,MoK 的推理部署形态是围绕 NVL72 这种单域机器设计的。一个 MoK replica 最好完整落在一个 NVLink 域内,不能拆到多个域,否则第一级路由的组播优势就全没了。这种"单体部署"的约束,会导致弹性扩缩容的粒度变粗,你在做服务编排时要提前规划好资源池。
5. 实验结果与选型判断:巨核不是万能药
5.1 论文报告的几个关键数字
MoK 的实验设置在公开数据集上训练到 1T token 量级,对比对象是同等 FLOPs 预算的细粒度 MoE 基线和稠密模型。论文报告的几个数字,我整理成表格更直观。
| 指标 | 细粒度 MoE 基线 | MoK(巨核) | 备注 |
|---|---|---|---|
| H100 32 卡集群训练吞吐 | 100(基准) | 79 到 86 | 通信劣势场景,MoK 反而更慢 |
| NVL72 训练 MFU | 33% 到 36% | 42% 到 45% | 通信优势场景,收益明显 |
| 1T token 后下游任务平均分 | 基准 | +0.8 到 1.5 个百分点 | MMLU/GSM8K/MBPP 等 |
| 长上下文首 token 延迟 | 基准 | -18% 到 24% | 预取 + KV 亲和性收益 |
| 专家塌缩发生率 | 偶发 | 显著降低 | 巨核梯度密度高 |
这个表格其实信息量很大。MoK 的好成绩几乎全部依赖 NVL72 这种超大规模的 NVLink 域,而不是算法本身更神。换句话说,MoK 吃的是硬件红利,如果硬件平台退回普通 8 卡节点,它的总分反而会输给细粒度 MoE。这也解释了为什么论文作者要把 NVL72 直接写进标题——它不是背景,而是方法的一部分。
5.2 什么时候该用 MoK,什么时候别碰
看完论文数据,最需要冷静的是选型判断。我的观点非常明确:MoK 不是要取代所有 MoE,它是在特定硬件平台上的最优解。符合下面几条场景,MoK 值得重点考虑。
第一,你手里已经有 NVL72 或者计划上 GB200 集群,且模型规模在百亿甚至千亿参数级别。这个时候 MoK 的 MFU 优势会直接转换成训练成本优势,省下来的时间非常可观。第二,你有严格的推理延迟要求,尤其是长上下文或高并发生成场景,MoK 的预取和显存局部性带来的 P99 改善很有价值。第三,你的团队已经有一定的大规模分布式训练经验,能驾驭张量并行和组播通信的调配,否则第一次把巨核跑起来会遇到不少系统层的坑。
反过来,如果条件不满足,还是老老实实细粒度 MoE 更稳。比如你只有通用的 8 卡 GPU 服务器,或者云上跨节点弹性扩缩容,MoK 很难吃满红利,甚至可能因为巨核太大、切分粒度太粗而拖慢训练。再比如你的模型只有几亿参数,稀疏化收益本身就不明显,引入两级路由和块对角结构只会增加工程复杂度。模型规模小的时候,稠密模型或者传统 MoE 是更合理的答案。
5.3 我的几点延伸思考
MoK 这篇论文让我最触动的地方,不是两级路由本身,而是它把"硬件拓扑"正式变成了模型架构设计的输入变量。业界过去太习惯把 GPU 集群当成一个固定的、带噪声的计算背景板,设计 MoE 时默认通信是贵的,于是所有努力都花在如何在通信约束下做调度。MoK 的做法是倒过来问:如果通信不贵,模型结构应该长成什么样?答案是一个全新的架构空间。
顺着这个思路往下走,我猜测未来的大模型训练框架会出现更多"域感知"设计。一个 NVL72 域是计算单元,跨域是存储与扩展边界。MoE 专家的放置、路由策略、负载均衡方法,都会围绕"域"来做文章。MoK 把专家做成巨核,本质上是把路由的通信模式从"网状随机访问"改成了"组内专业化访问",这种模式在下一代更大规模的 GPU 机上会越来越吃香。
另外提一句,虽然论文里没有明确讨论,但 MoK 和长上下文推理的结合点非常值得关注。巨核数量少意味着推理引擎可以做非常激进的显存分配,把 KV cache 和专家权重放在一起,换取更低的访问延迟。这个方向如果继续挖下去,可能会催生一套专门为 NVL72 设计的推理引擎优化栈,而不是通用框架里的一个插件。
从我个人的实操角度看,MoK 给团队带来的最大价值是逼着我们重新审视 MoE 的所有默认配置。为什么专家是几千维?为什么路由只做一次?为什么负载均衡非要那么复杂?这些问题过去都被"前人经验"压着没人问,现在有一篇论文告诉你:换一套硬件,答案可能完全相反。哪怕你不打算立刻上 MoK,光是把这些问题重新想一遍,就已经值回票价了。