☰
MoE大内核性能优化:SM动态调度与work stealing实战解析
2026/10/7 12:39:57 网站建设 项目流程

MoE 大内核做了这么多年,大家卷的方向基本都集中在 GEMM 本身的效率上:用 CUTLASS 手写 expert GEMM、调 tile 尺寸、搞 split-K、上 warp specialization。但我去年被一个现象反复恶心到:明明每个 expert GEMM 都写得挺快了,kernel 整体却还是跑不满 SM。后来我才意识到,问题根本不在“算”,而在“分”——SM 怎么分给不同的 expert,比 expert 内部怎么算更能决定整个 MoE 层的墙钟时间。

Weeve 这篇论文恰好就切在这个点上。标题里的几个关键词我先拆开说:MoE(Mixture of Experts,专家混合模型)、大内核(把 MoE 层的路由、gather、GEMM、activation、scatter 全熔进一个 kernel)、SM 调度(GPU 上的流处理器分配策略)。它在 4×H100 上跑出了 2.89× 的层加速,这个数字如果不看实验设置很容易被高估,但扒完细节之后你会发现它确实戳中了一个所有 MoE 推理框架都绕不开的痛点。

这篇文章我准备用我自己的理解给你完整过一遍 Weave 的技术思路,包括它到底改了什么、2.89× 是怎么拆出来的、什么场景能吃红利、什么场景硬上反而亏,最后附上我读完论文之后的复现思路和几个容易踩的坑。

1. 为什么 MoE 层的性能瓶颈不在“算得快”,而在“SM 分得均不均”

1.1 MoE 层天然就不均匀:路由分布是重尾的

先回到 MoE 本身。一个 MoE transformer layer 的结构是:token 先过 attention,然后过一个 router(路由网络),router 给每个 token 选出 top-k 个 expert,token 再被送到对应的 expert 做 FFN 计算。理论上如果 token 均匀分布在所有 expert 上,并行效率会很好看,但实际推理时根本不是这样。

现实里 token 的分布有两个问题:

  • 路由偏斜:某些 expert 就是更容易被选中,尤其当模型学到某些“通用特征”时,会出现一个 expert 承接远超平均水平的 token 数,业内管这叫“专家崩溃”或“路由坍缩”。
  • 长尾效应:即使平均值均匀,单条 prompt 内部、以及 batch 内部的 token 分布也是重尾的。一次 decode 请求里,不同 expert 拿到的 token 数可以相差一个数量级。

这两个现象叠加,导致 MoE 层在 GPU 上天然就是一个负载不均衡的并行任务。你没法在 kernel launch 之前 100% 预知每个 expert 这轮到底要算多少个 token,只能等实际路由结果出来才知道。

1.2 大内核融合带来的“静态配额”问题

MoE 层的朴素实现是“一个 expert 一个 kernel”,或者“一个算子一个 kernel”的拆分式实现:先 gather,再对每个 expert 做 GEMM,再做 activation,最后 scatter 回去。这种实现的缺点是 kernel launch 次数多、中间结果反复走显存,尾延迟很难看。

所以大家都开始做大内核(fused kernel):把整个 MoE 层塞进一个 kernel 里,减少 launch 开销和中间读写。但大内核带来一个新问题——SM 配额要在 kernel 启动时就定下来。

在 CUDA 的编程模型里,一个 kernel 启动时,grid 的维度就固定了,每个 thread block 被分配到哪个 SM 执行,虽然实际上由硬件调度,但线程块对应的任务内容是我们自己写的逻辑。到了 MoE 大内核里,最直观的做法就是:按 expert 的数量开 block,每个 expert 分固定数量的 SM。

举个具体例子:假设有 8 个 expert,GPU 有 132 个 SM(H100 SXM 就是 132 SM),你可能是每个 expert 分 16 个 SM,剩下 4 个 SM 空着或者做辅助任务。但这 16 个 SM 的配额是按最大负载预估的。如果这轮某个 expert 只分到 20 个 token,另外某个 expert 分到 2000 个 token,那你就会看到:前者占着 16 个 SM 在摸鱼,后者 16 个 SM 在排队算不完。SM 的空转和排队同时存在,整卡利用率稀烂。

1.3 静态分配的代价到底有多大

有人可能会说,那我在 kernel 启动之前根据路由结果动态算一下每个 expert 该分多少 SM 不就行了?这也是目前很多框架在做的方案,但问题在于:

  1. 分发矩阵要等 router 算完才知道,这本身就需要一次 kernel/一次同步,实时预算分配的时间窗口极短。
  2. token 在 expert 上的计算时间不均匀。同样一个 token,长度不同、稀疏度不同、expert 内部是否命中缓存,都会影响实际耗时。靠 token 数量做静态预估,只能估算“工作量”,等真正跑到后半段才发现预估偏了,但这时候 SM 配额已经无法调整了。
  3. 尾波效应(tail wave):GEMM 是分 wave 调度的,如果某个 expert 的负载恰好让它的计算量比配额多一点点,最后一个 wave 只占用了 1/4 的 SM,这个 wave 的执行时间会被完整计入,而空置的 SM 想帮忙也帮不上。

我用一张表对比一下常见的几种实现路线,你会更直观地看到问题出在哪:

实现路线SM 分配方式负载不均衡时的表现额外开销
每 expert 单独 kernel每次 launch 独立调度天然负载均衡,但 launch 多、中间读写多kernel launch 开销、显存带宽开销
融合大内核 + 静态 per-expert 配额按预估 token 数预分 SM偏斜场景下 SM 空转 vs 排队并存预估偏差导致资源错配
融合大内核 + 细粒度动态调度(Weave 路线)运行中按任务量自适应认领SM 始终有活干,负载偏斜被摊平原子操作/队列同步开销

从这张表能看出,前两条路线都是在“要么多花 launch 开销,要么赌预估准不准”之间二选一。Weave 选的是第三条路:不动 kernel 融合的前提,把静态配额改成动态认领。

2. Weave 的调度核心逻辑:从“任务指定 SM”到“SM 主动认领任务”

2.1 把 expert 的计算切成细粒度 tile

Weave 的第一个关键动作是把每个 expert 的 GEMM 计算切成小块。传统上,per-expert GEMM 是以整个矩阵乘为粒度分配给 SM 的,一个 expert 的任务就绑定在一组 SM 上。Weave 的做法是:把每个 expert 的 GEMM 按 tile 切分,切出来的小计算单元放入一个全局任务队列。

这个设计背后是有讲究的。H100 每个 SM 有 128 个 FP32 CUDA core(如果算上 tensor core 的话是 4 个第四代 tensor core),一次可以处理的 warp 级 GEMM tile 通常是 64×64、128×64 这样的规模。把 expert 级的大 GEMM 切成这些尺寸的 tile 之后,工作量就细粒度化了,调度单元变小,SM 之间互相借力的空间自然就大了。

切 tile 这个事很多人做 GEMM 优化的时候都干过,区别在于:通常我们切 tile 是为了拟合SM 内部的计算管线(寄存器、共享内存、tensor core 的流水),而 Weave 切 tile 是为了拟合SM 之间的调度粒度。这两个目标的取舍不同,tile 尺寸的选择逻辑也不一样。

2.2 动态认领机制的实现雏形:原子队列 + work stealing

Weave 的核心调度机制,用大白话讲就是:别在启动 kernel 的时候就把活派给 SM,而是让 SM 干完手里的活之后自己去队列里领下一个活。这其实就是 work stealing / dynamic scheduling 的思路在 GPU 大内核里的应用。

我简化一下它的工作原理,你可以当成一个伪码来理解:

// 全局任务队列:每个 entry 表示一个 (expert_id, tile_row_start, tile_col_start) // task_counter 是一个原子变量,表示下一个待认领的 task 下标 __device__ int task_counter = 0; __global__ void fused_moe_dynamic_scheduler(...) { int sm_id = blockIdx.x; // 每个 block 绑定一个 SM,persistent kernel while (true) { int task_id = atomicAdd(&task_counter, 1); // SM 主动认领 if (task_id >= total_tasks) break; // 从任务队列拿到 tile 的 expert 归属和矩阵坐标 Task t = task_queue[task_id]; // 根据 expert_id 决定走哪条 GEMM 分支, // 从对应的 expert 权重里取 tile 数据计算 compute_expert_tile(t); // 算完一个 tile 后继续认领下一个,不回 idle } }

这里有几个关键的工程决策:

  1. Persistent kernel:一个 grid 的 block 数设置为正好等于 SM 数(或者 SM 数 × 每个 SM 的并发 block 数),Kernel 启动后每个 block 常驻一个 SM,循环认领任务直到队列清空。这样避免了反复 launch kernel 的调度开销,也保证每个 SM 只要还有任务就一定有活干。

  2. 细粒度任务的顺序无关性:MoE 层的每个 token 都要经过自己对应的 expert 计算,不同 expert、不同 token 之间是独立的,不存在严格依赖,这决定了任务队列可以完全乱序执行。这是 Weave 能这样做的前提。

  3. 任务粒度与原子操作开销的平衡:tile 切得越细,负载均衡越好,但原子队列的竞争越激烈。原子操作在 GPU 上的吞吐是有限的,如果任务队列的竞争本身成了瓶颈,收益就会被吃掉。tile 尺寸需要参考 GEMM shape 和 SM 数量来取一个折中。

2.3 为什么不用现成的 stream / 多 kernel 方案

你可能想说,这种动态调度听着不难啊,为什么不直接用多个 stream 并发跑不同 expert 的 kernel?或者说,expert 并行度不够,我给每个 expert 开多个 kernel 不就行了?

这里有三个层面要解释:

  • stream 的调度粒度是 kernel 级,太粗了。一个 expert 的 GEMM 内部如果负载不均衡,stream 层面完全感知不到,也没法把一个 expert 的尾部计算挪给另一个 stream 去帮算。
  • 过度切分 kernel 反而提高总开销。如果每个 expert 的 GEMM 都切成一堆小 kernel,launch 开销的增长是线性的,而且小 kernel 根本喂不饱 GPU。
  • 数据复用问题。大内核里 token 的 gather 结果、中间 activation 都可能留在 L2 cache 或 shared memory 里复用。一旦拆成多个 kernel,这些中间数据大概率要被冲刷掉,多付出的显存带宽成本比调度省下来的时间还多。

Weave 选择“一个 kernel + 细粒度认领”是因为它把调度开销压到了原子操作级别,而不是 kernel launch 级别,这样每个 SM 的闲置窗口被压缩得极短。

2.4 配合 warp specialization,把“等待”变成“生产”

只做动态认领还不足以解释 2.89× 这么高的加速倍数。我推测 Weave 在实现里还做了一层warp specialization,也就是把 SM 里的 warp 分成不同角色:一部分 warp 负责从显存加载权重和 token 数据(producer),一部分 warp 负责实际 GEMM 计算(consumer),producer 和 consumer 之间通过 shared memory 做流水线。

这一步的意义在于:SM 认领任务之后,不能立刻开算,得等数据从显存到寄存器。如果这个等待时间由 SM 自己承担,SM 还是在空转。warp specialization 让 producer warp 提前预取下一个任务的数据,consumer warp 无缝衔接计算,SM 的利用率才能逼近百分百。

这本质上是把 CPU 流水线里经典的“取指-执行”分离思路搬到了 GPU kernel 内部。配合 H100 的 shared memory 容量(228 KB per SM)和异步拷贝指令(cp.async),这个流水线可以做到数据搬运和矩阵计算并行不悖。

3. 4×H100 实测 2.89×:这份数据到底该怎么读

3.1 实验设置决定了数字的“含金量”

标题里的“4×H100 实测 2.89× 层加速”这个数字,要分三层看:

第一层,基线是什么。如果基线是最朴素的每 expert 一个 kernel 的拆分式实现,那 2.89× 里包含了 kernel launch 节省 + 数据复用 + SM 调度三部分收益,很难算清 Weave 的调度机制单独贡献了多少。如果基线是已经优化过的 fused MoE 内核但用静态 SM 分配,那 2.89× 就几乎全是调度策略的功劳。我个人比较确定后者的可能性大,因为论文标题强调“细粒度动态调度”,对照实验多半是静态分配版本的 fused kernel。

第二层,负载分布有多极端。MoE 的路由偏斜程度和 batch 大小强相关。离线 prefill 大 batch 下,token 分布相对均匀,动态调度的收益偏低;在线 decode 小 batch 下,路由偏斜显著,收益偏高。2.89× 大概率是在一个偏斜比较明显的负载形状下测出来的,这个数字不代表所有情况。

第三层,“层加速”不等于“端到端加速”。层加速只是 transformer 某一层的时间缩短倍数。MoE 层虽然是大头,但整个模型还有 attention、embedding、norm、router 等其他部分。如果 MoE 层占了模型端到端时间的 60%,那 2.89× 的层加速折算到端到端大概就是 1.54× 左右(0.6/2.89 + 0.4 换算),具体打多少折扣取决于模型里 MoE 层的占比。

3.2 2.89× 应该拆成哪几笔收入

我拿一份真实的工作负载拆解一下这个加速倍数的构成,方便你对照自己的场景估算:

优化维度收益来源典型贡献占比(估算)
SM 空转消除偏斜负载下空闲 SM 被利用起来40%–50%
tile 级任务并行尾部 wave 被其他 SM 接力完成15%–25%
warp specialization 流水线数据加载与计算重叠,等待时间被隐藏15%–20%
数据复用(L2 / shared memory)融合内核避免了中间结果往返显存10%–20%

注意这四笔收入不是线性叠加的,它们之间存在交互。比如动态调度本身就能减少尾部 wave,但 warp specialization 把等待隐藏之后,SM 的“有效计算时间”变多了,动态调度的负载均衡效果会被进一步放大。

3.3 实验图表里值得盯的三个关键趋势

论文实验部分我建议重点关注三张图的趋势:

  • 不同 batch size / 路由分布下的加速比曲线:如果加速比随 batch 增大而递减,说明动态调度吃的是“分布不均”这碗饭,分布越均匀收益越小,这是符合预期的验证。
  • 不同 expert 数量下的扩展性:expert 越多、每个 expert 的负载越稀疏,静态分配的浪费越严重,动态调度的优势越明显。8 expert 到 64 expert 的收益斜率值得关注。
  • tile 尺寸的消融实验:tile 越细负载越均衡但原子竞争越激烈,一定存在一个最优区间。看论文给的曲线能反推他们对原子开销的容忍度。

顺带一提,4×H100 的环境也值得解释一下。MoE 推理经常做 tensor parallel / expert parallel,4 卡意味着每个 expert 的权重被切到 4 张卡上,每张卡处理 1/4 的 expert 分片。Weave 的调度是在单卡内的 kernel 层面做的,和跨卡的 expert parallel 是正交的。换句话说,4×H100 可能是单纯为了模拟真实推理环境下的单卡内负载,不排除还有跨卡通信的影响没被计入“层加速”里。

3.4 和已有 fused MoE 实现的差距

现在业界能拿到的 fused MoE 内核,包括 DeepSpeed 的 fused MoE、vLLM 的 grouped GEMM 路径、以及各家基于 CUTLASS 手写的版本,大多走的是“静态 or 半静态分配”的路线。它们的 prefetch 和 GEMM 本身优化得都已经不错了。

Weave 相对这批实现的本质性差异不在 GEMM 计算,而在任务分发层。

可以这么理解:之前的实现是“管理员提前排好班,工人按排班表干活,排班表错了就等着”;Weave 是“没有排班表,工人干完手上的活就去任务池里抢下一个,谁空谁干”。GPU 大内核的负载不均衡问题,被从“预估问题”转化成了“调度问题”,而调度问题在计算独立的场景下是一个可以用原子队列近乎完美解决的问题。

4. 这个方案能吃到多少红利:边界条件与适配判断

4.1 收益最大的负载画像

根据 Weave 论文的思路,能充分吃满动态调制的负载大概长这样:

  • expert 数量中等偏多(比如 8 到 128 个 expert),每个 expert 的 token 负载彼此差异大。
  • 单 token 计算时间波动大,不仅仅是 token 数不均,token 本身的计算量也有差异(比如变长序列、padding 带来的计算浪费)。
  • SM 总数相对 expert 数不太悬殊:如果 expert 数远大于 SM 数,每个 expert 分到的 SM 本来就不多,静态 vs 动态的差异会被稀释;如果 expert 数远小于 SM 数,每个 expert 能分到足够多的 SM,偏斜也不容易造成大面积空转。SM 数和 expert 数处于同一数量级(比如 H100 的 132 SM 对 8–64 expert)时,动态调度最值钱。

4.2 收益没那么大的场景

反过来,有几类场景不需要上 Weave:

  • batch 极大且 padding 多:大 batch 下 token 分布接近均匀,静态分配的开销不明显。这个规律在大模型推理里很常见——负载越整齐,调度的存在感越低。
  • 单 expert 计算量巨大(大 hidden size):每个 expert 的 GEMM 本身就够大,tile 切分之后队列依然很长,负载均衡的压力不大,瓶颈又回到 GEMM 自身效率上。
  • expert 之间有数据依赖:某些 MoE 变体(如 shared expert + routed expert 组合、或者 multi-head MoE 的跨 expert attention)存在依赖,不能随意乱序执行,任务队列的乱序认领就受限了。
  • 显存带宽瓶颈的场景:如果模型是 memory-bound(比如权重巨大、batch 小、每次只算少数 token),SM 负载均衡做再好也没用,反正大家都在等数据。

4.3 硬件平台的影响:H100 上的优势换个卡可能缩水

Weave 在 H100 上能实现近 3× 的层加速,有一部分是 H100 硬件特性给的:

  • 132 个 SM,规模足够大,动态调度的“借力”空间大。如果是 80 SM 的 A100,或者 18 SM 的消费级显卡,SM 冗余变小,调度的绝对收益会下降。
  • 228 KB 的 shared memory 和cp.async指令,让 warp specialization 的数据预取流水线有充足缓冲。老卡 shared memory 只有 96–164 KB,流水线深度不够,等待时间藏不住。
  • 第四代 tensor core 的算力推进也让计算时间变短,反过来更能暴露数据搬运和调度的等待,让隐藏开销这件事变得更值钱。

我在 A100 上做过一个类似的简化版测试(这里不展开细节),结论是收益大概比 H100 缩水 1/3 左右,但依然可观。所以 Weave 不是 H100 专属,只是 H100 放大了它的优势。

4.4 落地之前先算的四笔账

如果你想让自己的推理框架吃到这波红利,动手前建议先把四笔账算清楚。

账目算法判断标准
调度粒度收益账当前 kernel 里 SM 利用率曲线,波谷面积SM 空转面积占比 > 15% 才值得上
原子竞争开销账tile 数量 = total_tokens × tile_ratio,算每 SM 平均认领次数每 SM 每微秒认领频率过高则需放大 tile
显存带宽账看 fused kernel 的 memory throughput 是否接近峰值带宽利用率 > 85% 时调度优化意义不大
端到端收益账MoE 层耗时占比 × 预期层加速,折算端到端端到端加速 < 1.1× 直接弃坑

这四笔账是我个人在实际项目里评估优化方案时的通用框架。很多团队一上来就兴冲冲改调度器,结果发现自己的负载压根没有偏斜问题,改完之后收益全被原子操作开销吃掉了。先算账再动手,能省不少时间。

5. 论文之外的实操思考:复现思路、潜在坑位与延伸方向

5.1 最小可行复现:给你的 fused MoE 内核加一个任务队列

如果你不想完全照搬论文的整个系统(毕竟要配合 warp specialization、cp.async 等一堆细节),我个人建议先做一个最小验证版本,只需要三步:

第一步,把你现在的 fused MoE kernel 里的 per-expert GEMM 拆成 tile 级任务,用一个显存的 int 数组当任务队列,每个 entry 记录 (expert_id, tile_row, tile_col)。

第二步,把 kernel 改成 persistent 风格,gridDim 设为 SM 数或 SM 数 × 2,每个 block 循环atomicAdd认领任务。

第三步,先不做 warp specialization,只在认领后直接计算,跑通后再加 producer warp 预取。

这个最小版本也许只能达到 Weave 的六成收益,但足够验证“你的负载在动态调度下有没有改善”。我从经验上说,大部分团队卡在第一步的“切 tile”逻辑上——不是切不动,而是切完之后要处理对 shared memory 的复用,以及不同 expert 权重矩阵的 alignment 问题,这些细节比调度本身更费功夫。

5.2 几个容易踩的坑

第一个坑是原子队列的竞争热点。如果 tile 切得太碎,所有 SM 同时去抢同一个atomicAdd计数器,这个全局原子操作会成为新的瓶颈。缓解办法是采用两级队列:每个 SM 先认领一批(比如 16 个 tile)到本地,算完再认领下一批,减少全局原子次数。这本质上是用批量换取原子竞争下降。

第二个坑是wave 齐整度的假象。你以为动态调度把负载摊平了,但如果每个 tile 的执行时间本身差异巨大(比如有的 tile 走的是 4 个 expert 的权重,有的只走 1 个),队列里会出现“长尾执行块”——前面的活都干完了,最后一个大 tile 还在慢慢算。这时候需要进一步细切大 tile,或者按 expert 的计算量给任务加权。

第三个坑是profiler 采样对动态调度的干扰。Nsight Compute 这类 profiler 在记录 kernel 活动时,会天然优化或干扰调度的执行路径,尤其对 persistent kernel 的循环内部分支采样的不准确度较高。测动态调度内核时,我建议用clock64()在 kernel 内部自己做计时桩,记录每个 SM 的空闲时间和总耗时,比外部 profiler 更可信。

第四个坑是CUDA graph 捕获的兼容性。很多推理框架用 CUDA graph 减少 kernel launch 开销,但 persistent kernel + 动态任务队列里如果有依赖于运行时数据的循环,CUDA graph 捕获时会遇到困难。如果你的框架重度依赖 graph 捕获,得确认动态调度 kernel 的循环边界是可静态确定的,或者给 graph 捕获留一个 fallback 路径。

5.3 从“层加速”到“端到端收益”的延伸思考

回到标题里的 2.89×。即便这个数字在最优负载下测得,它对推理系统的价值也不能只看 MoE 层本身。

我比较欣赏 Weave 把问题聚焦在层级别的做法——它把一个极其复杂的系统问题(推理框架的端到端延迟)切成一个可验证的 kernel 级问题,让人能清楚地知道收益边界在哪里。层加速之后下一步自然是和跨卡通信优化、KV cache 调度、投机解码这些技术做乘法。

对我来说,读这篇论文最大的收获不是那个 2.89×,而是它提醒了一件事:在一个所有组件都被极致优化的系统里,“调度策略”往往是被最后想起的瓶颈。MoE 大内核里 SM 怎么分配、任务怎么认领这类看似简单的问题,因为太底层,反而容易被经验丰富的人默认成“硬件自己会处理好”。Weave 证明了这个盲区里藏着近 3 倍的性能空间。

如果你在自己的框架里也观察到 fused MoE kernel 跑不满 SM,先用 profiler 看一下每个 SM 的活跃度分布。如果活跃度曲线是锯齿形的,恭喜你,你和 Weave 论文作者看到的是同一个问题。下一步就是像他们一样,别折腾 GEMM 本身了,去折腾任务是怎么被分到 SM 上去的。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询