☰
MoE负载不均衡怎么办?Weave动态SM调度带来2.89倍加速
2026/10/7 22:57:01 网站建设 项目流程

如果你的 MoE 模型跑起来 GPU 利用率忽高忽低,经常一排 SM 忙成狗、另一排闲着没事干,那你大概率在某个深夜的 profiling 日志里怀疑过人生。这个问题在 MoE 大内核里尤其明显:专家负载天然不均衡,而传统 kernel 里 block 一旦指派给某个 SM 就死守到底,谁也帮不了谁。Weave 这篇论文做的事情,简单说就是把 MoE 大内核里的 SM 从一个“静态工位”改成“动态资源池”,让计算负载在 SM 之间细粒度地迁移,4×H100 环境下实测单层加速 2.89 倍。

这篇内容适合三类人看:一是做 MoE 推理/训练加速的工程师,二是写 CUDA kernel 但总被负载不均衡折磨的性能优化选手,三是对 GPU 调度机制好奇、想理解 SM 层面到底发生了什么的研究者。我会把 Weave 的设计思路、实测收益怎么来的、复现实验要避开哪些坑,以及它适用的边界都拆开讲清楚。

1. 为什么 MoE 会把 SM 调度逼到墙角

1.1 MoE 的负载不均衡从哪来

MoE(Mixture of Experts)的核心是稀疏激活:一个 token 进来,只会被路由到少数几个专家网络,而不是全部专家都参与计算。听上去很省算力,但问题也随之而来——路由结果是动态的,不同 token 分布在不同 batch 里,热门专家和冷门专家的忙闲程度可能差出好几倍。

这就像一个餐厅后厨,炉头(SM)是固定的,但订单(token)是动态来的。有的窗口排队排到门口,有的窗口十分钟没动静。你没法提前知道哪个窗口会忙,因为 token 的分布取决于上一个 token 的 embedding、路由器的 softmax 结果、乃至训练阶段的数据分布。更麻烦的是,为了不让某些专家被无限“挤爆”,系统通常会设置 expert capacity,超出部分会被丢弃或者走残差路径。这就意味着:不仅负载不均衡,还伴随着有效信息被牺牲掉。

行业里现有框架怎么处理这个问题?最常见的是在路由层面做负载均衡 loss,或者在 token 层面重新分配。比如 vLLM、SGLang 这类框架,会在调度阶段尽量让不同 expert 上的 token 数接近。但问题是,这个均衡粒度停在“层”级别、停在 dispatch 之前的 token assignment,真正的 SM 级别不均衡,还是要靠 kernel 内部解决。你在 framework 层面无论怎么路由,只要最后落到 GPU kernel 里还是静态 block 分配,SM 层面的忙闲不均就始终存在。

1.2 SM 分配机制的底层约束

要理解 Weave 为什么有价值,得先明白 NVIDIA GPU 上 kernel 是怎么跑起来的。一个 kernel 启动时,会创建一大堆线程块(thread block),硬件层面的 block scheduler 负责把这些 block 分配到 SM(Streaming Multiprocessor)上。一旦 block 分配到一个 SM,它就会在这个 SM 上执行完,中途不会挪窝。这是 GPU 编程模型的基础。

这就是问题的根源:硬件调度器只管“分配”,不管“平衡”。它按照先来先得的方式填充 SM,但如果某个 block 恰好是重负载 expert 的,那么这个 SM 就要忙很久;旁边一个轻负载 expert 的 block 早跑完了,SM 只能空转等着 kernel 结束。一个 kernel 必须在所有 block 都完成之后才算结束,因此落后的专家拖慢了整个层。大家挤在一起等最慢的那个窗口出菜。

有些优化试图通过把 MoE 层拆成多个小 kernel 来改善——每个 expert 一个 kernel,这样不同 expert 可以跑在不同 stream 上。但代价也很明显:kernel launch overhead 急剧上升、中间结果要反复写显存、不同 expert 的 kernel 之间更没有机会共享 SM 资源。拆得越碎,调度灵活性越大,但搬运开销越高。Weave 的思路是反过来:走“大内核”路线,把 MoE 层整体保持在一个 kernel 里,但在 kernel 内部让 SM 不要死守某个 expert。

1.3 细粒度动态 SM 调度到底在调度什么

既然硬件 scheduler 不做平衡,软件层面能不能干预?能,但干预的方式不是去改写硬件分配逻辑,而是在 kernel 内部重新设计任务的表示方式。Weave 做的事,是把 SM 当作一个可动态再分配的计算池,而不是把每个 block 固定给某个任务。用大白话说:你不是给每个专家固定几个 SM,而是让专家们都去一个公共池里临时借用 SM。

这就涉及到“调度单位”的变化。传统 MoE kernel 的调度单位是一个 expert 对应的整个计算任务,粗颗粒度;Weave 把计算任务切成更小的 tile,比如一个 token 组或一个 block 宽度的数据切片,然后这些 tile 在 SM 之间动态流动。流动的策略不是全局集中式的,而是类似 work-stealing——某个 SM 发现自己空闲了,就去队列里捞一个没做完的 tile 来做。

这个过程其实很像 CPU 领域的“智能核心调度”思路:不是让线程死绑在一个核心上,而是允许核心间迁移,从而消化不平衡。区别在于 CPU 调度发生在操作系统或者运行时层,一次迁移可能几微秒;而 GPU SM 之间的 tile 迁移,必须被控制在几乎无感的时间窗口内,否则调度开销比收益还大。Weave 的核心价值,就是用极低开销的信号机制实现了这种迁移。

2. Weave 的核心设计思路拆解

2.1 第一层:如何感知不平衡

动态调度的第一步,是知道自己不平衡。但 GPU kernel 里收集全局信息是非常昂贵的事。如果你在每个 expert 计算前都做一次全局 barrier 加一次状态统计,那调度还没开始,性能已经被同步开销拖垮了。Weave 的做法,从论文里的实验刻画来看,走的是“局部信号 + 原子协商”路线。

具体来说,每个 expert 的计算队列在显存中维护一个原子计数器,表示自己还有多少待处理的 tile。每当一个 SM 完成当前 tile,它不会傻等,而是去查这些计数器,选一个最需要帮助的队列“偷活”。查询动作本身只有几次原子操作或者 volatile 读,开销极小。这就像外卖骑士送完一单之后,看一眼自己附近的订单热度,再决定去哪接单。

这里有个关键设计问题:信号要多“新鲜”?如果用全局内存里的计数器,每次读取都要过 L2,延迟可能几百周期;但如果你把信号下沉到 L2 里的 per-SM 状态,读开销能到几十周期。Weave 的做法是允许信号有适度延迟——不需要每个周期都精确,只要别让 SM 空转太长时间。这跟 CPU 分支预测器的思路有点像:用未来几个周期内大概率正确的信息做决策,换极高的能效比。

2.2 第二层:如何做动态决策

感知到负载信号之后,接下来是决策。在分布式系统里调度决策通常由中心调度器完成,但 GPU kernel 里不能有中心节点——那会让所有 SM 争抢一份锁,扩展性瞬间崩掉。Weave 的决策机制更接近“局部自由市场”。

每个 SM 在做完手头 tile 后,会根据自己看到的队列长度快照做局部决策:如果本 expert 还有剩余 tile,优先继续做;如果本 expert 已经做完,就去其他队列获取工作。决策只依赖本地 cache 里的信息,不做全局锁。这个机制类似于网约车平台的派单逻辑:乘客等车的位置和价值由平台算好,但司机不在一个地方排队接单,而是在不同区域之间移动寻单。

为了让决策不震荡,Weave 引入了一个偏好因子:SM 优先停留在自己最近的 expert 队列上,只有当自己队列没活且其他队列超过一定阈值时才迁移。这个阈值相当于一个“死区”,避免因为信号波动导致 SM 在多个 expert 之间来回搬。我在实际调度系统里见过太多这类问题——如果对信号反应太灵敏,系统会在两个负载差不多的队列之间抖成筛子,浪费远比收益多。

2.3 第三层:细粒度体现在哪里

Weave 的“细粒度”体现在两个维度:一是调度单位小,二是调度事件频繁。调度单位不再是整层的所有 expert 计算,而是一个 tile 或一个 token 组。这意味着一次 SM 迁移最多迁移一个 tile 的负载,代价上限是可控的。SM 可以在一个 expert 计算某个 token 组时穿插处理另一个 expert 的 tile,从而让每个 SM 始终有事做。

第二层细粒度体现在流水线联动上。一个 MoE 层的计算通常可以分三段:dispatch(把 token 送到对应专家)、compute(专家网络的前向)、reduce(把结果合回主分支)。如果只对 compute 阶段做动态调度,dispatch 和 reduce 阶段还是会成为瓶颈,因为它们在时间上是串行的。Weave 的做法是把三段融合到同一个运行框架里,让 dispatch 产生的 tile 在“生产”阶段就开始给未来的 compute 提供信号,reduce 阶段也遵循同样的队列机制。这样一来,SM 之间流动的不只是计算任务,还包括数据准备的进度信息。

这让我想起做 CPU 智能核心调度时的经验:真正的系统优化不能只调一个点,要把“感知-决策-迁移”当作一个闭环来设计。类似地,Weave 在 SM 之间形成的也不只是一个动态映射表,而是一整套 tile 生命周期循环的调度协议。

3. 4×H100 实测拆解:2.89 倍是怎么来的

3.1 测试环境和测量方法

论文里的 4×H100 实测,对应的应该是 4 卡 H100(可能 PCIe 或 SXM,版本不同对结果有一些影响但趋势一致)。典型被测对象是 MoE backbone 的中间层,比如 8 专家或 16 专家配置。关键指标是“层加速”:把 MoE 层整体从 dispatch 到 reduce 的执行时间,对比 baseline 与 Weave 的耗时比值。2.89 倍就意味着 Weave 把层时间压缩到原来的约 35%。

测量工具上,我建议复现时至少用两层工具交叉验证:第一层是 CUDA event 计时,拿到端到端 wall time;第二层是 nsys timeline,看 MoE 层在不同阶段上的时间分布。有人喜欢直接用 nsight compute 看 kernel 时长,但因为 MoE 层可能不只是一个 kernel,更稳妥的做法是定义一个准确的 CUDA graph 节点边界,把 dispatch、compute、reduce 整体括进一个计时区域。

如果你要自己复现,环境里有几件事必须先确认。GPU 要跑在持续的 clock 频率上,最好用 nvidia-smi 锁定;关闭 ECC 与否会影响显存带宽,进而影响 tile 搬运的开销;如果用了 MIG 或虚拟化,调度行为会显著偏离物理 SM 模型,结果不能直接类比。

3.2 2.89 倍背后的性能账

2.89 倍的层加速听起来很高,但我们可以从利用率角度算一笔账。假设 baseline 时 MoE 层有 16 个专家,每个专家负载方差较大,SM 平均利用率可能只有 50% 左右。也就是说有一半 SM 周期在空转或者等显存。Weave 动态调度让空闲 SM 去接别的专家 tile,理论上能把平均利用率拉到 85% 以上。如果 kernel 是 compute-bound,加速比大约等于利用率之比:85% / 50% ≈ 1.7 倍。如果 baseline 利用率更低,比如只有 30%,那加速比就可以到 2.8 倍。

2.89 倍的实测数字说明测试场景里的 baseline 确实处于极不均衡的状态,也说明 Weave 的调度开销被控制得很低。如果调度开销占掉 10% 的 kernel 时间,理论 3 倍的收益会被压到 2.7 倍,但实测依然有 2.89,说明在 4×H100 这个规模下,调度开销占比应该在 5% 以下。

为什么不是更高的加速?有几个限制因素。第一是 dispatch 和 reduce 阶段的不可压缩开销,这两段虽然也被纳入调度,但它们涉及数据布局的重排,不像 pure matrix multiply 那样容易均衡。第二是尾部效应:即使动态调度,当整批 tile 只剩最后几个时,有一两个 SM 还是要等那最后一块算完,利用率无法到 100%。第三是显存带宽约束:如果一个 tile 的数据需要跨 SM 迁移,带宽会成为新的瓶颈。

3.3 从实验结果反推设计选择

通过 2.89 倍这个数字,我们能反推出 Weave 的几个关键参数区间。比如 tile 大小:如果 tile 太小,迁移开销和原子操作开销会吃掉全部收益;如果 tile 太大,不平衡粒度又太大,SM 空转窗口还是解决不了。以 H100 的规格估算,一个 tile 如果对应 32 到 128 个 token 组的矩阵乘,比较合理。这正好是 warp 级并行度和任务拆分数之间的平衡点。

SM 数量分配也很关键。MoE 层里的专家数量与 SM 数量不是整除关系,因此必然有某些 SM 要负责多个专家 tile。Weave 采取的队列数量一般是专家数乘以一个小系数,比如每个专家两个队列,便于细粒度抢占。这跟操作系统里 run queue 的 per-CPU runqueue 多队列设计有异曲同工之处——队列越多,锁竞争越小,调度就越灵活,但队列管理本身也消耗资源。

4. 实操中的注意事项与踩坑实录

4.1 复现实验的五个关键设置

先列一下我在类似项目里被坑过的地方,给想复现的同学提个醒。

第一,GPU 频率锁定。H100 的 boost 频率会根据功耗和温度浮动,如果你不锁频,跑两次实验的频率不同,测出来的加速比可能波动 20%。用nvidia-smi -lgc 1980这类命令锁在持续频率。

第二,确认 L2 缓存策略。MoE 层里 token 的 dispatch 数据要经过 global memory,如果 L2 策略是 write-back 或 streaming,可能对 tile 迁移时的命中率影响很大。建议用cudaDeviceSetCacheConfig配合实际负载微调。

第三,性能分析器会影响调度。Nsight 在回放模式下会改变某些 kernel 的执行流,存带宽 profile 出来的结果跟真实跑了会有偏差。我建议先用nsys profile -c cuda做一次记录,然后不带 profiler 跑纯计时来验证。

第四,多卡环境里的 P2P 影响。4×H100 如果用了 NVLink,不同卡上的显存访问延迟差异很大。如果你的 MoE 层是张量并行切分过的,要确认只统计了单卡局部执行时长,而不是包含跨卡通信的时间,否则测出来的“层时间”会混合进通信量。

第五,确定 baseline 版本。对比实验里 baseline 一定也必须是优化过的版本,不要拿一个 naive 的 expert-by-expert kernel 来比赚了便宜。最好用当前框架里 SOTA 的 MoE kernel 作为 baseline,比如基座框架自带的 fused MoE 实现,这样 2.89 倍才真实反映增量价值。

4.2 常见问题与排查思路

我在调度类 kernel 的调试中积累了一个排查表,直接贴出来给你参考。

现象可能原因排查/解法
计算结果不稳定,偶发错误原子操作顺序不被保证,tile 被重复消费检查队列计数器的 CAS 逻辑,确认每个 tile 只被抢一次
加速比低于 1.5 倍tile 太大,SM 空转窗口不足以被填平把 tile 减小到 1/4 试一次,同时观察 atomic 操作冲突率
显存占用莫名上升每个 expert 队列的 tile 元数据存储过多改为固定循环缓冲区,信号只保留 cursor 而不保存完整描述符
多卡时加速比严重退化跨卡 NUMA 效应被低估,tile 迁移访问了远端内存强制 tile 分配优先本地内存,只在本地队列空时才跨卡
小 batch 不升反降调度初始化开销占比过高增加最小工作总量阈值,低于阈值退化为静态调度

这些坑里面,第一个原子操作的问题最隐蔽。GPU 上 atomicCAS 和 atomicAdd 的语义跟 CPU 不完全一样,跨 thread block 的原子操作一致性需要你额外显式处理 memory fence。如果你在队列尾部插入和头部消费之间少了一个 fence,你会发现结果在单卡上偶尔对、在多卡上必错,这个问题极其费时间。

4.3 个人经验:什么时候别用动态调度

虽然 Weave 很香,但它不是万灵药。以我的经验,这几类场景不要轻易上动态 SM 调度。

第一,小模型或者低吞吐场景。如果你每个 MoE 层计算时间本身只有 50 微秒,一次 tile 迁移就可能花掉 5 微秒,调度开销占比 10% 起步,基本白做。动态调度需要足够的“剩余利用空间”才能覆盖成本。

第二,层内专家数太少。如果只有一个或两个专家,SM 之间没有足够的队列多样性,work-stealing 无从谈起。至少要有 4 到 8 个专家,调度空间才成立。

第三,如果你已经做了非常激进的 token 级负载均衡,比如把 token 重新打包成固定 token 块,再配合等尺寸专家网络,那 SM 层面的不平衡已经被压得比较低,动态调度的增量收益会缩水。这时候不如直接优化矩阵乘的指令效率或者加 FlashAttention 这一类算子融合。

5. 适用边界与扩展思路

5.1 收益最大的场景特征

什么样的 MoE 层最适合 Weave?我用一个清单来总结:专家数量在 8 到 64 之间,token 路由分布方差大,batch 足够大(至少几千 token),单层计算时间在 200 微秒以上。在这个区间里,动态 SM 调度的收益能稳定保持在 1.8 倍以上。

如果换到 A100 平台,效果会略差一点——A100 的 SM 数量是 108 或更少,H100 是 132 或 144(取决于型号),调度空间的绝对规模变小,但是负载不均衡问题照样存在,加速比一般还有 1.5 到 2.2 倍。L40S 这类偏图形和推理的卡也适用,但需要留意它的 L2 带宽规格和 H100 不同,tile 大小的最优参数会偏移。

与 CPU 智能核心调度做对比能更好理解边界:CPU 上动态迁移负载收益明显是因为核心间中断和 cache 迁移的代价已经很低;GPU 上 tile 迁移其实需要搬运的是寄存器状态和 shared memory 状态,代价不低。Weave 的巧妙之处在于它尽量在 tile 边界做迁移,避免迁移一个正在执行中的计算,而是迁移“还没开始的计算”。这个设计原则是所有想借鉴它的系统都应该采用的。

5.2 与框架、编译器的结合方向

Weave 的调度逻辑其实是与底层硬件指令集强相关的。想把它推广成通用优化,可以考虑和编译期生成代码结合:编译器在生成 MoE kernel 时,能够根据专家数量、token 分布统计静态配置队列数量和 tile 大小。比如 Triton 或者 CUTLASS 能不能生成这种带动态队列的 persistent kernel?我认为未来大概率有人做,但当前最稳的路线还是用 CUTLASS 的 tile scheduler 或者自己写 persistent CUDA kernel。

另一个方向是跟运行时调度层协同。现在大模型推理框架里的调度器主要管 CPU 侧的请求调度,GPU 侧的 kernel 调度被认为“硬件自动完成”。Weave 打破了这种心智,它让 GPU 内部 kernel 层面也能接受运行时传入的“偏好信号”——比如根据请求长度预测负载,提前给某些专家队列更高的优先级权重。这就跟 CPU 智能核心调度里的负载预测形成跨层呼应了。

这里也能看到行业趋势:原来大家把 GPU 当成一个“黑盒执行器”,把调度全部交给 BSP 模型;现在开始往 GPU 内核里做主动的、细粒度的动态调度,把 GPU 当成一个“有内部分布式系统”的资源池来设计。Weave 在 MoE 场景里的验证,大概率只是这条路上的一块铺路石。

最后说点实在的。我一开始接触这类项目时,也怀疑这么复杂的调度逻辑到底能不能落地,毕竟 GPU kernel 的调试难度比 CPU 调度高一个量级。但真正跑通之后,那种“空转 SM 活过来了”的感觉确实很上头。如果你准备在 MoE 层上做类似的细粒度调度工作,我建议先从一个专家块的 persistent kernel 改起,验证 tile 迁移开销模型,再逐步扩展到整层。调度这个东西,设计时永远要记住:开销比算法先落地,迁得动、迁得快,才是王道。

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

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

立即咨询