MoE 模型的文章我看得不少,大多数都在讲路由策略、负载均衡 loss、专家并行怎么切分,但专门把矛头指向 GPU 底层 SM(Streaming Multiprocessor)调度的论文,确实少见。最近读的这篇《Weave》让我眼前一亮,标题写得很直接——MoE 大内核里的细粒度动态 SM 调度,而且给了 4×H100 实测 2.89× 层加速的数据。做 LLM 推理优化的人都知道,H100 已经很强了,能在 MoE 层上再抠出接近 3 倍的加速,说明问题不只是算力不够,而是算力根本没用满。文章适合正在折腾 MoE 推理性能、关心 GPU 利用率、或者想理解调度系统底层逻辑的人来读,读完之后你会对“稀疏模型为什么跑不快”有一个完全不同的认识。
1. MoE 推理为什么卡在“调度”而不是“算力”上
1.1 MoE 的算力悖论:稀疏激活反而让 GPU 出现大片空闲
MoE 的核心卖点是稀疏激活,一个 token 进来,router 只挑 top-k 个专家做计算,理论上计算量是稀疏的,同样的参数量能装下更多专家。但真正部署到 GPU 上就会发现,稀疏只是参数层面的稀疏,硬件层的麻烦反而更多了。
问题出在 token 和专家之间的对应关系上。真实场景里 token 的分布极其不均匀,热门专家每天被大量 token 光顾,冷门专家几乎无人问津,这种分布近似 Zipf 分布。在 GPU 上表现为:负责热门专家的 SM 排队排到天荒地老,负责冷门专家的 SM 闲着没事干。MoE 的稀疏性反而制造了新的负载不均,而且是动态的、不可预测的负载不均。
传统的 MoE 推理框架一般怎么做?最粗暴的一种是把专家静态绑定到固定的 SM 子集上,哪个专家被分配到多少 SM 是配置死的。一旦流量分布偏离预期,整个 GPU 的利用率就崩了。简单说,MoE 不是缺算力,而是算力被调度策略困住了。
1.2 静态 SM 分配是浪费的根源
为了说清楚静态分配的浪费,得先理解 MoE 层在 GPU 上到底跑什么。一个典型的 MoE 层包含 router 计算、token 到专家的分发(all-to-all)、专家 FFN 计算、结果回收(all-to-all 回来)这么几个阶段。前两个阶段和最后一个阶段通常是内存带宽瓶颈,中间专家计算阶段是计算瓶颈。
静态分配的问题在专家计算阶段最明显。假设一个 GPU 有 128 个可用 SM,你把 4 个专家各分 32 个 SM 绑死。如果这一批 token 有 70% 都去了专家 A,专家 A 的 32 个 SM 会塞满计算任务,其他专家的 SM 却空着看热闹。你可能会说,那就把 token 重新路由一下呗?
问题就在这,token 应该去哪个专家是模型学出来的语义决定的,强行改路由会伤害模型质量。负载均衡 loss 能在训练阶段缓解这种不平衡,但推理阶段你面对的是一个已经训练好的模型,它的路由偏好已经固化了。况且即使训练的时候加了 aux loss,实际推理时依然会有波动。
1.3 “细粒度”到底细在哪里
既然静态分配不行,干脆动态分配。动态分配本身不是什么新概念,服务器集群里的弹性伸缩就是动态分配,GPU 上的 stream-K 也把 split-K 的小块动态分配给 SM。但过去 MoE 的动态调度粒度太粗,顶多到 GPU 级别或者 GPC(Graphics Processing Cluster)级别,Weave 的思路是把动态调度下放到单个 SM 级别。
细到 SM 级别意味着什么?一个 H100 有 132 个 SM(实际可调度的一般按 128 左右算),4 卡就是 500 多个 SM。如果把专家相关的计算任务拆成很多个小块,让所有 SM 自己过来抢活的干,哪还有空闲和排队的说法?这就是 Weave 的核心直觉——与其让 token 强行匹配专家资源,不如让计算任务主动流到空闲的 SM 上。
2. Weave 的设计拆解:大内核与动态 SM 调度
2.1 MoE 大内核:把一堆 kernel 揉成一个常驻任务
Weave 的第二关键词是“大内核”。常规的 MoE 层实现要 launch 很多个 kernel,router 一个 kernel,每个专家一个 kernel,all-to-all 还要有通信 kernel。每次 kernel launch 都有固定开销,而且 kernel 之间数据要写回全局显存再读出来,白白消耗带宽。
大内核的思路是反着来的:搞一个 persistent kernel,把路由判定、token 分发、专家 FFN、结果回收全部塞进去,kernel 启动一次就一直在 GPU 上运行,直到处理完这一层的所有 token。这样 kernel launch 的开销几乎被清零,中间数据也能尽量留在 SM 的寄存器或者共享内存里,不用频繁回全局内存。
我之前写过不少 CUDA 优化的文章,一直强调 kernel launch overhead 在短小 kernel 场景下有多致命。MoE 层正好符合这个特征,单次专家 FFN 的计算量不大,但调用次数极多,launch 开销占比很高。大内核把这些开销省掉,本身就能带来可观的加速。
2.2 动态调度的机制:任务队列加原子计数器
大内核只是容器,真正聪明的是它内部的调度机制。Weave 的做法可以理解为在每个 GPU 上维护一个全局的任务队列,队列里放的是专家计算任务块。当一个 SM 完成手头的任务,就通过原子操作从队列里抓取下一个任务继续算。
这个机制跟 CPU 调度里的 work stealing 非常像。每个 SM 是一个工作线程,队列是共享任务池,原子计数器保证多个 SM 同时取任务时不会冲突。任务块的粒度很关键,太大不够灵活,太小的话原子操作频次太高,反而把调度开销吃回去。合理的选择是把专家 FFN 的计算按 token 数量切成中等大小的块,让每个块在 SM 上跑几十微秒级别的时间,这样调度开销占比能控制在很低水平。
用生活类比解释就是:以前是餐馆每个服务员固定负责几张桌子,客人多的时候有人忙死有人闲死。现在改成所有服务员站在一个取菜口,谁手里没活了就自己端菜走,整个餐馆的接待效率自然提上去了。
2.3 负载均衡不靠硬掰 token,靠流动的计算任务
做 MoE 的人对负载均衡这个词都很敏感。训练阶段大家都要加负载均衡 loss,还要设 expert capacity factor,token 路由的时候如果某个专家过载,多余 token 直接 drop 掉。这些都是靠静态约束来抑制不均衡。
Weave 对这种思路做了个有意思的颠覆:既然 SM 层面的计算任务是流动的,那 token 爱去哪个专家就去哪个专家,只要专家 FFN 的任务能动态分布到所有 SM 上,过载问题自然就化解了。换句话说,不是修改 token 的分布,而是让计算资源自适应 token 的分布。
这会带来一个连带好处:推理阶段可以放宽甚至去掉对负载均衡 loss 的约束。训练得更自由,模型质量更好,推理时还不怕热门专家打爆某一个 SM 子集。我在自己团队里见过太多因为推理负载不均衡被迫重新训练的案例,如果 Weave 的思路可落地,这类问题会少很多。
2.4 和“负载均衡代码”热词的关系
最近“moe 负载均衡代码”这个词被搜得多,说明大家都被负载不均衡问题折腾得不轻。传统的负载均衡代码集中在 router 端的 Loss 实现上,比如给 router 加辅助损失项,或者设计容量限制函数,这些都是训练期的手段。权重固化之后,推理期的负载均衡只能靠部署层的专家并行和缓存策略来实现。
Weave 换个思路,在 kernel 层面把任务重新调配。底层的原子操作、任务队列、自旋等待这些代码,就是它的“负载均衡代码”。这个思路本质上不需要 model-level 的干预,做的好了对上层模型完全是透明的。
3. 显存问题:MoE 到底要不要全部参数进显存
3.1 先给结论:峰值推理基本要,但可以有取舍
“moe 架构要全部参数进显存吗”这个热搜问得很核心。先说答案:如果追求低延迟,峰值推理时专家参数基本要先放进显存。MoE 虽然激活稀疏,但 token 可以路由到任意专家,你不能提前知道哪个专家不被用到,常见的做法是把全部专家都常驻显存。
但实际操作有取舍。共享专家和路由专家要分开看待,共享专家被所有 token 使用,必须常驻。路由专家数量大、单个占用小,可以按热度做弹性驻留。更深层的现实是,现在 MoE 模型的专家数动不动就是几百个,全部塞进显存需要很高的显存预算。
这就面临一个经典的推理调度权衡——模型放不下的话要么 offload 到 CPU 内存或者 NVMe,要么量化压缩专家权重,要么对冷门专家做按需加载。每种方案都有自己的开销。推理时延最怕的是突然访问到冷门专家,要从硬盘读权重,这个延迟能到几十毫秒级别。
3.2 Weave 这类方案对显存提出了新要求
大内核动态调度看着美好,但需要付出显存代价。Persistent kernel 意味着 GPU 上要预留一块常驻的调度区域,任务描述符队列、token 缓冲、中间计算结果都要占显存。以前静态 kernel 跑完资源就释放了,现在内核常驻,这部分显存不能被其他任务使用。
我做 GPU 显存优化时有个经验,任何新的调度机制都要算一笔账:省了多少显存带宽和 launch 开销,又额外消耗了多少驻留显存。Weave 们这些设计通常面向单层优化,显存增长是局部性的,整体影响不会很大,但如果你在做一个显存预算很紧的多层模型,这些额外开销必须提前规划进去。
另外还有一个细节值得指出,动态调度的任务块如果涉及把大量 token 暂存到显存里的缓冲区,那么这块缓冲区的分配策略就很重要。预分配太低,高峰来了不够用;预分配太高,常态下浪费显存。这个在工程落地上是个比较头疼的 tuning 点。
3.3 不同显存预算下的部署选择对比
把显存配置和调度方式放在一起看,能更清楚地理解 Weave 的适用范围。我整理了一个对照表:
| 部署方案 | 显存占用 | 调度方式 | 推理效率 | 适用场景 |
|---|---|---|---|---|
| 全专家常驻 + 静态 SM 绑定 | 高 | 静态 | 中等,负载不均时浪费明显 | 常规多卡部署 |
| 全专家常驻 + 动态 SM 调度 | 高 | 动态细粒度 | 高,利用率显著提升 | 低延迟、高吞吐在线服务 |
| 部分专家 offload | 中 | 静态/动态混合 | 中低,冷门专家延迟较高 | 显存紧张、容忍长尾延迟 |
| 量化专家 + 常驻 | 低 | 动态 | 中高,需接受精度损失 | 小显存环境 |
从这个表能看出来,Weave 这类动态 SM 调度方案是在“全专家常驻”的前提下才能发挥最大价值的。你占了大量显存,但能换回接近 3 倍的层加速和更高的 GPU 利用率。在 H100 80G 级别的大显存卡上,这个交换非常划算,但在 40G 级别的卡上就得掂量掂量了。
4. 4×H100 实测:2.89× 层加速是怎么来的
4.1 实验环境:4 卡 H100 的 SM 资源池
论文用的环境是 4 张 H100,每一张约 132 个 SM 满血,实际可调度的大概 128 个左右,4 卡加一起就是 500 多个 SM 的动态资源池。卡间互联走 NVLink,带宽在 900GB/s 级别,all-to-all 通信不会被压出明显的瓶颈。
这个环境配置选得很有代表性。现在的 MoE 部署主流方案是 EP(Expert Parallel),把不同的专家放到不同 GPU 上,每个 GPU 处理一部分专家。4 卡是个比较典型的起步配置,往上能扩展到千卡集群,往下也有 2 卡测试的需求,用 4 卡先验证调度机制是合理的。
H100 上的 SM 数量和计算能力都有富余,dynamic scheduling 有了施展空间。如果换到 SM 数量更少的消费级卡,动态调度省下来的时间可能不足以抵消调度开销,效果会大打折扣。
4.2 基准对比:和谁比出 2.89×
论文对比的基线是传统静态绑定的 MoE 层实现,就是最常见的那种方案,每个专家固定占用一组 SM,按批次处理 token。在这种方案里,kernel 按顺序逐个 launch,每次都要经历完整的启动、计算、退出周期。
Weave 的对比目标很有讲究,没有拿一个精心调过的 fused kernel 做基线,而是挑了最主流的静态实现。这样一来加速比里面有相当一部分是“大内核省 launch 开销”贡献的,另一部分是“动态调度省排队空闲”贡献的。行业里做 benchmark 容易犯的毛病就是挑一个最弱的基线来突出自己,所以我看数据的时候一般会留意基线是怎么设置的,这个加速比才可信。
从工程角度看,静态实现确实还是目前大多数生产框架的默认选择,和它对标有实际意义。如果跟一些已经做了 expert 分块优化的实现去比,加速比应该不会这么夸张,大概会缩水到 1.5 到 2 倍之间,这依然很可观。
4.3 2.89× 的数据拆解:延迟到底省在哪
2.89× 的层加速意味着 Weave 处理 MoE 层的平均耗时只有基线的约三分之一。这个收益不是均匀分布的,主要来自三部分。
第一是消除了专家排队。在静态分配下,热门专家所在 SM 的任务队列可能堆积几倍于平均水平的负载,这部分等待时间是纯浪费。Weave 下任务在 500 多个 SM 间流动,理论上任何 SM 都不会长时间空闲,队列等待接近零。
第二是减少了 kernel launch 和数据回显存的开销。传统实现每个专家一个 kernel,启动、结束、写回全局、再读入,每一次都要付出额外的带宽和时间。大内核把这些开销摊薄到整个层处理周期里,几乎可以忽略。
第三是让 SM 利用率更接近理论峰值。我在做性能分析时喜欢看一个指标叫“SM busy ratio”,静态方案在极端不平衡下可能只有 40% 到 60%,Weave 这类动态方案能把它拉到 90% 以上。利用率翻倍,延迟自然就降下来了。
要特别提醒的是,2.89× 是“层加速”而不是“端到端加速”。真实场景下 MoE 层只是整个模型的一部分,前后还有 attention 层、MLP 层、embedding、采样等,端到端收益要看 MoE 层占总延迟的比例。如果 MoE 层占 50%,那端到端大概能提速 40% 左右,这个数字依然很漂亮。
4.4 多卡拓展与千卡部署的思考
顺着 4 卡往上看,千卡级别的 MoE 部署是最近的搜索热点,值得展开说几句。到了千卡规模,每个 GPU 都得承载跨卡 token 分发,所以 all-to-all 通信会成为新的瓶颈。Weave 这种调度本质上是单 GPU 级别的机制,它解决的是单个 GPU 内部 SM 利用率的问题,跨卡的 token 流动还是要靠高速互联。
但有意思的是,如果每张卡内部的 SM 利用率都能被拉满,跨卡通信的传输量是不变的,通信效率自然就成为了下一个亟待解决的焦点。比如说,动态调度如果刚好让某张卡上的专家负载变大,会不会导致 token 在该卡上堆积?这个跨层、跨卡级别的问题 Weave 没直接解决,也给未来的优化方向留了想象空间。
在我看来,将 Weave 应用在千卡集群时最好叠加一层传统的专家并行和 token 路由策略,上层做粗粒度分配,下层用 Weave 做细粒度调度,两个维度配合起来才能把整个集群的效率吃干榨净。
5. 实操落地会踩的坑与排查方法
5.1 调度开销的反噬:原子操作不是免费的
动态调度最大的隐患是原子操作竞争。500 多个 SM 同时去抢一个队列的索引,如果不做优化,原子操作本身的延迟和带宽消耗就能把你辛苦省下的时间全部吐回去。我见过类似的实现,原子竞争严重的时候性能不但没提升,反而比静态方案还慢 10% 到 20%。
几个缓解方案是验证过有效的:一是 SM 每次不是取一个任务块,而是批量取多个任务块存到私有队列里,减少访问全局队列的频率;二是任务块粒度要动态调整,SM 多的时候调小一点,SM 少的时候调大一点;三是对队列数据结构做分区,不同 SM 抢不同分区的任务,减少冲突。
5.2 不同 GPU 世代下的效果差异
我一直强调,H100 上 2.89× 不代表每种 GPU 都有这个收益。H100 的 SM 数量多,计算能力强,任务块在 SM 上的运行时间长,调度开销相对占比小。如果换到 SM 数量较少的卡上,任务块跑得飞快,原子操作和队列访问的开销相对就突出了。
另外不同代际 GPU 的硬件调度能力不同,新卡有更细粒度的硬件调度支持,动态调度可以做得更顺滑,老卡可能只能靠软件模拟。做方案选型之前一定要搞清楚自己的硬件版本,不然精心调好的模块换个环境就失灵了。
5.3 与推理框架集成的现实问题
vLLM、SGLang 这些主流框架目前对 MoE 层的支持都有自己的一套,如果你想把 Weave 的思路套进去,得面对一个现实:框架里的算子实现不是你想换就能换的。Persistent kernel 和 CUDA Graph 在某些情况下是冲突的,因为 CUDA Graph 要求整个执行图预先捕获并静态化,而动态调度天然带有运行时的不确定性。
我在实际项目中踩过这种坑,capture 阶段动态调度器的工作方式跟 graph 捕获机制完全对不上。要么关掉 CUDA Graph 用传统的 kernel launch,损失一部分时间;要么得把调度逻辑改成 graph-compatible 的模式,这会牺牲一些动态灵活性。论文本身可能不会讲这些工程细节,但真正落地的团队大概率都要面对类似取舍。
5.4 常见问题速查表
| 问题现象 | 可能原因 | 排查与对策 |
|---|---|---|
| 动态调度后性能反而下降 | 原子竞争太激烈,任务块粒度太小 | 增大任务块,采用批量取任务或分区队列 |
| SM 利用率始终上不去 | 任务块之间依赖太强,无法并行 | 检查 token 分发的 all-to-all 是否成为瓶颈 |
| 大内核导致显存不足 | persistent kernel 常驻占用 + 缓冲预留过大 | 缩减缓冲预分配,评估动态分配策略 |
| 与 CUDA Graph 集成失败 | 动态调度不符合图捕获的静态要求 | 关掉 graph 或改造调度为可捕获模式 |
| 单卡效果明显但多卡不升 | 跨卡通信成为新瓶颈 | 上层叠加专家并行或优化 all-to-all 拓扑 |
| 部分专家延迟明显偏高 | 冷门专家被 offload 后按需加载 | 结合热度策略做专家驻留管理 |
在这个表之外有一条我认为比什么都重要:动态调度的收益不能只看平均延迟,要看 P99 延迟和尾延迟。MoE 负载分布的长尾问题在静态方案下特别明显,Weave 这类方案对尾延迟的改善往往是平均延迟改善的几倍,排查的时候建议把延迟分布图拉出来看。
还有一个容易忽视的细节,就是多进程或多实例共享 GPU 时的 SM 抢占问题。假设 GPU 上同时跑着两个推理实例,一个用 Weave,一个用静态调度,它们之间会互相干扰,动态调度的效果会大打折扣。生产环境要么统一所有实例的调度策略,要么用 MIG 或整卡独占的方式隔离。
说回我自己读这篇论文的心得。MoE 推理优化做了这么久,大家的目光一直盯着路由策略和显存管理,很少有人走到 CUDA 层去思考 SM 到底在忙什么。Weave 给我最大的启发不是那套任务队列和原子操作,而是它提醒了我一件事:很多时候性能瓶颈不在算法而在于执行模型。你有一个稀疏的路由模型,却用一套静态的、非共享的执行方式去跑它,冲突几乎是必然的。真正的解法是想办法让硬件资源跟上模型的动态性,如果模型的路由是动态的,那你对计算资源的使用方式也应该动态起来。这个思路不仅能用在 MoE 层上,后续处理更大的专家粒度、更复杂的多模态 MoE、甚至端到端的动态执行模型时,都有借鉴价值。正好最近我也在研究把类似机制推广到前缀共享和树解码场景,后续有实测结果了再分享。