ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Weave:MoE大内核细粒度动态SM调度,4×H100实现2.89×层加速

Weave:MoE大内核细粒度动态SM调度,4×H100实现2.89×层加速 MoE 模型的文章我看得不少大多数都在讲路由策略、负载均衡 loss、专家并行怎么切分但专门把矛头指向 GPU 底层 SMStreaming 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 级别或者 GPCGraphics 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 很多个 kernelrouter 一个 kernel每个专家一个 kernelall-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 factortoken 路由的时候如果某个专家过载多余 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 部署主流方案是 EPExpert 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、甚至端到端的动态执行模型时都有借鉴价值。正好最近我也在研究把类似机制推广到前缀共享和树解码场景后续有实测结果了再分享。
返回列表