![]()
MoK:把 token 调度、跨 GPU 通信和专家计算塞进同一个 GPU 内核。
作者丨郑佳美
编辑丨岑 峰
7 月 21 日,NVIDIA 公布了 GB300 NVL72 训练 DeepSeek-V3 的最新成绩:256 张 GPU 上,单卡性能达到 1,648 TFLOPS。
不到两周后,Cursor 开源了 Mixture-of-Kittens,简称 MoK。它没有继续折腾更快的矩阵乘法,而是直接重写了一层 MoE 的执行方式,把 token 调度、跨 GPU 通信和专家计算塞进同一个 GPU 内核。
![]()
这件事有点反直觉。因为 GB300 NVL72 已经把 72 张 GPU 放进同一个 NVLink Domain,整机架 NVLink 总带宽达到 130 TB/s。按这个规格看,GPU 之间传数据应该已经足够快。
可在大规模 MoE 训练里,通信依然会拖住专家计算。
问题出在 MoE 的数据流。Router 每一步都会重新决定 token 去哪个专家,而专家分散在不同 GPU 上。token 要先跨卡送过去,算完再传回来;发送之前还要整理位置,到达之后还要等待数据完整。随着 MXFP8 和 Blackwell Tensor Core 把专家计算压得越来越短,这些原本还能藏在计算后面的等待开始越来越显眼。
Cursor 做 MoK,就是从这里下手。它没有只追求一次 Dispatch 能传多快,而是重新安排 token 怎么到达专家、什么时候开始计算,以及通信和计算怎样同时占用 GPU。
在算力资源被 Blackwell 和 NVLink 堆到极致的今天,开发者发现:硬件跑得再快,也救不了低效的软件编排。
MoK 的开源不仅是一个内核的胜利,它标志着“应用层定义算子”的时代的开始:为了压榨最后 30% 的算力,AI 创业公司正在发起一场底层主权的夺权战争。雷峰网
不过要理解 Cursor 的这套设计为什么有效,先是得看清一层 MoE 到底在什么地方交了“通信税”。
![]()
01
MoE 的「通信税」,远不止把 token 传过去
在普通稠密 FFN 里,token 要经过哪些权重基本是固定的。MoE 加入 Router 之后,每个 token 临时选择几个专家。采用 Expert Parallel 时,专家又被拆到很多 GPU 上,于是一次前向传播至少要发生两轮跨卡通信。雷峰网
![]()
第一轮叫 Dispatch,把 token 送到专家所在的 GPU;专家计算完成后,再通过 Combine 把结果送回 token 原来的位置。训练还有反向传播,对应再执行两轮反方向的通信。
如果只是把一大块连续数据从 A 搬到 B,NVLink 已经足够快。MoE 麻烦在于,Router 每一步给出的数据分布都不同。
某个专家这一步可能收到很多 token,下一步又很少。系统需要先数清楚每个专家有多少 token,再决定这些 token 放在目标 GPU 的什么位置,还要把同一专家的数据尽量排在一起。Grouped GEMM 才能拿到规整的输入,高效喂给 Tensor Core。
![]()
通信结束后也不能立刻开算。目标 GPU 要确认远端写入已经全部完成,数据确实可见。专家负载又不完全均匀,一部分 GPU 可能早早结束,仍然得等最忙的专家收尾。
所以一个“通信”阶段里,其实混着数据搬运、布局生成、同步和负载失衡几件事。130 TB/s 描述的是整机架能够提供的峰值带宽,并不代表每一次动态、零碎的 MoE 通信都能把这些链路同时跑满。
DeepEP 已经把其中的数据搬运做得很快。它本身就是面向 Expert Parallel 的高性能通信库,提供专门的 Dispatch 和 Combine 内核,支持 FP8,并允许控制通信占用的 SM 数量。最新版本甚至可以用明显更少的 SM 保持很高的通信吞吐。
可一层 MoE 仍然要在 Dispatch、Grouped GEMM 和 Combine 之间不断交接。
![]()
这时会碰到一个很实际的矛盾:如果等足够多的 token 到齐再计算,矩阵很大,Tensor Core 跑得舒服,但计算启动得晚;如果 token 到一点就算一点,通信和计算能更早重叠,矩阵又太小,GPU 很多 SM 没有足够的工作。
多个 CUDA Stream 可以让通信和计算并发,却很难始终保证两边拿到恰到好处的 GPU 资源。
MoK 后面的设计,基本都在处理这个“节奏”问题。
![]()
02
怎么让 token 一边传,一边算
MoK 最有意思的改动之一,是把前向 Dispatch 从 Push 改成了 Pull。
传统 Push 很直观:源 GPU 手里有 token,就主动把它写进目标 GPU。麻烦出在目标地址。一张 GPU 会同时收到很多其他 GPU 发来的 token。
每个发送方都得提前知道自己应该写到哪一段,不能互相覆盖;同一个专家的 token 最好还要连续排列,否则后面的 GEMM 又要重新整理。
![]()
随着参与 GPU 越来越多,这个调度过程会越来越重。Pull 换了一种方式:保存专家的目标 GPU 自己去读取需要的 token。它只要知道 token 在哪张源 GPU、位于源数据的什么位置,本地存放在哪儿由自己安排。
![]()
这一下省掉了多张发送方之间对目标地址的协调。收到的数据也可以直接按照本地专家排好。
有意思的是,Pull 并没有少传数据。Cursor 的微基准里,同一个 256 × 256 BF16 数据块,Push 在 NVLink 上移动约 159.6 KB,Pull 反而达到 172.0 KB,因为读取还要额外发送请求。
![]()
Pull 胜在另一件事:MoE 的通信很碎,而且负载不均衡。NVLink 的两个方向拥有独立通道,Pull 能同时利用请求和数据返回两个方向。在专家负载不均衡的测试中,Cursor 测到最高 29% 的 NVLink 利用率提升。
同步上的差距更明显。Push 写完以后,目标 GPU 需要等待其他 GPU 的完成信号,Expert Parallel 规模大时,一个 rank 最多可能涉及另外 71 个 peer。Pull 是本地 GPU 自己发起读取,数据回来以后就能直接使用。
![]()
Cursor 的多节点微基准中,这段 signaling 延迟从 Push 的约 103 微秒降到了 Pull 的 18 微秒。
MoK 也没有所有地方都用 Pull。前向是 Pull Dispatch、Push Combine;反向则使用 Pull Reverse-Combine 和 Push Reverse-Dispatch。Dispatch 阶段需要把很多来源的 token 重新组织成专家输入,Pull 更省协调;Combine 时每个结果应该回到哪个 token 已经很清楚,直接 Push 回去更简单。
改变通信方向以后,MoK 还把通信和专家计算放进同一个 Megakernel。
它把 GPU 的 SM 分成两部分。一部分负责 Dispatch、Combine 和状态管理,另一部分专门执行 Expert FFN。通信侧拿到一批完整 token 后,通过 GPU 本地计数器通知计算侧;计算完成,再通知通信侧把结果传回去。
![]()
这样可以直接决定通信用多少 SM、计算用多少 SM,不用完全交给多个 CUDA Stream 自己竞争。这里最关键的参数叫 minibatch,也就是一次交给专家计算多少 token。
它不能太大。太大意味着第一批计算得等很久。也不能太小。专家 GEMM 最终要拆成大量计算任务分给 SM,如果 token 太少,任务数量不够,很多 SM 会闲着。
![]()
Cursor 用 wave 来判断这个边界。一个完整 wave,可以简单理解成所有计算 SM 都分到了工作。MoK 希望一个 minibatch 至少能够形成两个完整 wave,让 Tensor Core 有足够的任务持续执行。
实际结果很能说明问题。在 Hidden Size 为 7168、专家中间维度为 2048 的 Kimi 2.5 形状上,Cursor 估算 minibatch 至少需要大约 2368 个 token。512 token 时,MoK 前向耗时为 5.981 ms;增加到 2560 token 后,下降到 3.425 ms。再继续增加,速度已经没有明显改善。
![]()
换句话说,通信切得更碎并不会一直变快。真正高效的重叠,需要让通信尽早交出数据,同时又不能把 GEMM 切得太小。
但 MoE 还有一个麻烦:Router 跑完之前,根本不知道每张 GPU 最后会收到多少 token。
如果按照最坏情况准备 buffer,会浪费很多显存。如果先让 GPU 数完 token,再通知 CPU 分配对应空间,GPU 又得停下来等 CPU。
MoK 使用固定大小的 Ring Token Buffer。一块空间先装 Dispatch 进来的 token,专家算完、Combine 把结果传走以后,这块空间立即重新给下一批 token 使用。前一个 macrobatch 的 Combine,还可以和下一个 macrobatch 的 Dispatch 同时进行。
![]()
Ring Buffer 在这里很像一个缓冲层:通信暂时跑得更快,数据就在里面积累;计算消费得更快,就等下一批 token 到达。整个过程靠 GPU 上的状态推进,不需要 CPU 每一轮都进来决定下一步。
Cursor 还把 MXFP8 activation 的量化塞进了 Dispatch、Grouped GEMM 和 SwiGLU 的数据路径里,少掉独立量化 kernel,也少一次中间结果在 HBM 里的来回读写。
Pull、minibatch、SM 分区和 Ring Buffer 放到一起后,MoK 才真正变成了一条连续的 MoE 流水线。
![]()
![]()
03
一整套高度特化的执行方式
Cursor 的 benchmark 测的是完整 MoE 层,包括 Schedule、Dispatch、Expert FFN、Combine 和最后的加权合并,对比对象包括 NCCL + PyTorch、DeepEP、TransformerEngine 以及 HybridEP + Megatron。
在 GB300 NVL72 上,相比每种场景最快的公开基线,MoK 的 MXFP8 前向最高提升 2.37 倍,反向最高提升 1.78 倍;BF16 前向和反向最高分别提升 1.92 倍和 1.58 倍。
![]()
更重要的是端到端训练。Cursor 原来的生产方案已经用了 DeepEP。在 512 张 GB300 GPU 上,换成 MoK 后,单卡吞吐从每秒 760.9 个 token 提高到 1070.2 个,提升约 41%。
![]()
这组结果也要看清边界。Cursor 没有公开完整的逐项消融,所以无法准确说 2.37 倍里面有多少来自 Pull、多少来自 Megakernel、多少来自 Ring Buffer。
![]()
可以单独确认的是 Pull 对 NVLink 利用率和 signaling 延迟的改善,其余收益更多来自整套执行方式组合后的结果。
MoK 对硬件的依赖也很强。它面向 Blackwell 和 NVL72 这样的高速 NVLink Domain。Pull 的远端读取、通信和计算的细粒度交错,都建立在 GPU 之间能够低延迟访问彼此显存的基础上。模型的 Hidden Size、Top-k、专家规模一变,合适的 minibatch 和通信 SM 数量也会跟着变。
![]()
这也是 MoK 最值得看的地方。过去讨论 MoE 优化,很容易盯着两个数字:GEMM 有多少 TFLOPS,All-to-All 有多少 GB/s。到了 GB300 这一代,仅仅把这两个数字继续往上推,已经解释不了全部性能。
token 什么时候到、怎样排成专家需要的布局、攒多少开始计算、通信拿多少 SM、buffer 什么时候释放,这些执行细节开始直接决定训练速度。
一个拥有 130 TB/s NVLink 带宽的机架,最后仍然需要为 MoE 重写 GPU 内核,原因也就在这里:链路已经很快了,接下来要省的是 GPU 等数据的时间。
Cursor重写GPU内核的行为,标志着AI 2.0时代的竞争进入了“全栈主权”的新阶段。
以前,我们笃信“术业有专攻”,做应用就该做应用(Cursor),做底层就该做底层(NVIDIA)。但如今的AI竞争已经到了“去中间商化”的新阶段,Cursor 并不是因为“想做”而去做内核,而是因为“不得不做”。
DeepSeek 开启了工程榨取的新纪元,Cursor 则把这团火烧到了应用层。这种“去中间商化”的趋势正在重塑AI的定价权:以后,决定 AI 公司估值的,不再是它拥有多少 Token,而是它的代码离显存和寄存器到底有多近。
那些无法穿透底层黑盒的 AI 公司,最终将会留在“平庸税”的泥潭里。
参考链接:https://cursor.com/cn/blog/mixture-of-kittens
上车,带你看遍全球 AI 顶会精华
可独家畅览:
专家演讲PPT
大会报告全文
热门论文解读
学术新星访谈
特别声明:以上内容(如有图片或视频亦包括在内)为自媒体平台“网易号”用户上传并发布,本平台仅提供信息存储服务。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.