MoE架构这两年被大家聊得很多从路由策略到专家均衡都有不少文章。但真到把MoE模型推到生产环境做推理优化的时候大部分人卡住的往往不是模型本身而是底层kernel根本跑不快——MoE的稀疏特性让计算内核变得非常别扭动态路由、负载倾斜、多专家并行几乎每一个特性都踩在GPU调度器的痛点上。Weave这篇工作就是冲着这个问题去的。它讲的是在MoE大内核把多个专家计算整合进一个大kernel里用细粒度的动态SM调度把计算效率榨出来。论文给出的数据很直接在4×H100环境下单层MoE实测加速2.89倍。这个数字不是实验室里挑出来的最好看版本只要亲自动手优化过MoE kernel就会知道这个量级的提升确实存在而且核心瓶颈并不是算力不够而是SMStreaming Multiprocessor根本没有被喂饱。这篇文章我会从论文的设计思路讲起一直拆到那些可以落地复现的细节。不管你是正在写推理框架还是在优化服务端延迟或者只是对GPU kernel调度感兴趣应该都能从里面找到一些能直接搬过去用的东西。1. 为什么MoE kernel会卡在一个尴尬的位置1.1 MoE的kernel到底在算什么先对齐一下场景。这里说的MoE指的是transformer层里那套完整的稀疏前馈网络典型结构如下输入token先过一个router矩阵算出一个概率分布然后挑选top-k个专家通常是top-2把token送到被选中的专家FFN里计算最后把多个专家的结果加权合并。整个计算链条拆成kernel来看大体有这么几块router GEMM一个小矩阵乘输入是token向量输出是专家打分token重排/分发根据路由结果把连续内存里的token重新排列让同一专家的token在显存里连续存放专家FFN计算每个专家内部通常是一个两层MLP第一层上投影、第二层下投影中间夹一个激活函数结果合并/还原把不同专家的输出按原始token位置放回去做加权求和。传统的实现方式很直接router一个kernelpermute一个kernel再按专家数量循环启动一批FFN kernel最后再来一个combine kernel。这么做逻辑上简单但性能上有一个天然的硬伤——GPU的kernel启动开销和SM空转问题在这种多kernel串联模式下被加倍放大了。说个实际数据感受一下。H100一共有132个SM我们假设一个专家FFN的任务量恰好需要30个SM跑满一个wave波次那么一次启动一个专家kernel时132个SM里只有30个在工作剩下100多个SM全在空转。一个一个专家轮着来显卡利用率低得离谱。这也是为什么MoE的kernel不能按Dense模型那种思路来写的原因——Dense模型一个kernel可以轻松占满所有SM但MoE天然是碎片化的。1.2 静态SM切分的两个致命问题既然逐个启动专家kernel太低效很自然的想法就是把多个专家kernel融合成一个大kernel一次性启动占满所有SM。这个方向主流框架都做过但问题在于融合之后怎么把SM分配给不同专家。最简单的做法是静态切分。比如4个专家每个专家分33个SM大家各干各的。这在专家负载完全均匀的时候没问题但MoE的负载恰恰高度不均匀。实际推理时每个专家收到的token数量可能相差好几倍——某个专家可能被路由到了60%的token另一个专家只有5%。静态切分下负载重的专家在排队负载轻的专家却占着大量SM摸鱼整体算力利用率肉眼可见地掉。除了专家间负载不均还有一个更隐蔽的问题业界管它叫wave quantization波次量化效应。GPU执行kernel时所有SM是分批wave处理CTA线程块的。当一个wave里的任务做完了但剩余任务又不够填满下一个wave时最后一个wave就只会有少部分SM在工作其他SM干等着。举例来说如果专家A的计算量是40个SM的工作量专家B是90个SM总共有132个SM那么即使两个专家并行执行也会因为总任务量132恰好等于SM数而浪费一个wave的调度机会更常见的是任务量131个SM最后一个SM空转整个GPU拖到最慢的那个SM跑完才算结束。静态SM切分解决不了wave quantization因为它是按“平均负载”来切SM的但MoE每次前向传播的负载分布都在变——上一轮迭代专家A最忙这一轮可能就变成专家C最忙了。你不可能预知未来也不可能为每一种负载分布都准备一套编译产物。所以Weave的思路直接换了个赛道不做静态切分改成在SM层面做细粒度的动态调度。2. Weave的设计思路把多个专家kernel“织”起来2.1 名字里的门道把一个kernel编织成一个整体Weave这个名字本身就很说明问题。英文里weave是编织的意思论文的思路就是不再把每个专家的计算当成独立的kernel而是把整个MoE计算链——从router到分发到专家FFN再到结果合并——当成一整块布料来编织。所有专家的工作被拆成细碎的小任务统一丢进一个共享的工作队列然后让所有SM从这个队列里动态领取任务执行。第一次接触这个设计的时候我脑子里冒出来的画面是银行柜台。传统的做法是每个专家开一个独立柜台排A专家的队伍可能绕了三圈B专家的柜台却空无一人。Weave的做法是只开一圈统一的队伍所有窗口SM谁空下来就去队伍里领一个活儿不管这个活是哪个专家的。只要队列里有任务就没有SM闲着。要实现这种动态调度第一步是把大kernel变成persistent kernel常驻内核。所谓persistent就是这个kernel一旦启动就一直在GPU上跑不会被反复launch和销毁所有SM都在这一个kernel里循环工作直到整个MoE层处理完毕。fallback到传统方式的话每个专家一次kernel launch就是一次全局同步和资源重新分配这中间的开销在H100这类卡上虽然比老架构小但依然不可忽视。2.2 细粒度任务拆分的底层逻辑要让动态调度真正发挥作用任务粒度必须足够细。如果任务粒度还是“一个专家的一块大FFN”那队列里总共就几个大任务SM之间的负载均衡依然很难做细。Weave的做法是把每个专家的FFN再做一次切分按照token的batch维度切块或者按照矩阵乘法K维方向切块切成很多个能够在合理时间内完成的小计算任务。为什么切细就有用打个比方如果你只有三个大任务复杂度分别是100、50、10那不管怎么调度总有SM在等别人。但如果你把任务切成50个左右复杂度接近的小块调度器就能像流水线一样平滑地把工作分配给各个SM任务粒度越细调度器的发挥空间越大最后的完成时间越接近理论最优值。当然任务切得过细也会有代价——每个任务的取用都有调度开销任务多了总开销也随之上涨。Weave论文里应该有详细的粒度敏感性实验我根据自己做调度kernel的经验来推断这个粒度平衡点一般取决于单任务的计算量相对调度开销的比值。任务本身至少要能跑出几百个周期才能把一次原子操作和同步的开销摊薄。所以在复现这个方案的时候任务粒度不是一个拍脑袋定的参数而是需要针对具体模型和具体GPU实测出一个甜点值。2.3 为什么不能用现成的流调度器这里还有一个很容易被问的问题CUDA本身不是支持多流stream并发吗能不能直接给每个专家开一个stream让GPU硬件调度器自己去分答案是不行至少效果远不如Weave这种方案。CUDA stream的粒度太大而且stream之间的调度是由硬件和驱动层的策略决定的你没法精确控制“哪个SM的哪个调度单元去执行哪个专家的小任务”。更关键的是专家之间如果要做内存同步比如一个专家要等另一个专家算完做结果规约stream之间的同步开销一点也不小。Weave的做法是把同步逻辑直接搬进设备端的工作队列里通过内存序和原子操作来管理完全绕开了host端的干预调度延迟从微秒级别压到了纳秒级别。3. 关键机制与落地细节真正动手时要处理的环节3.1 设备端工作队列的任务拉取Weave能够成立的核心机制是设备端工作队列也就是把任务调度的逻辑从CPU搬到GPU上。CPU在传统模式里管的是kernel启动和同步而Weave模式下CPU只负责启动一次大kernel剩下的事情全部由GPU上的线程自己处理。任务拉取的伪代码大致长这样// 设备端任务拉取核心逻辑CUDA风格伪代码 __device__ Task fetch_task(WorkQueue* queue, uint32_t* global_counter) { // 原子递增拿到一个全局任务槽位 uint32_t slot atomicAdd(global_counter, 1); if (slot queue-total_tasks) { return EMPTY_TASK; // 队列已消费完 } return queue-tasks[slot]; }这段逻辑看起来简单但里面有两个极其容易被忽视的点。第一个是原子操作的竞争强度。所有空闲的SM都会来抢这个atomicAdd的计数这是典型的全局热点如果不加控制会把整个系统的瓶颈从计算转移到原子操作上。经典的做法是层级化队列每个SM或者每个GPCGraphics Processing Cluster先有一个本地队列SM优先从本地队列取任务本地空了自己再去做原子操作到全局队列取一批任务缓存到本地。这样一个SM一般只需要极少次全局原子操作热点问题就能缓解。第二个坑是内存序。任务写入队列和任务计数器递增之间必须有正确的内存屏障否则可能出现某个SM拿到了槽位却读到了尚未写入完成的任务数据。GPU上常见的做法是先store任务数据再增加计数器取任务时先读计数器再load任务内容配合__threadfence或acquire/release语义来保证可见性。这块如果写错了表现出的症状很随机——偶尔算错偶尔卡死而且只在特定GPU型号上复现排查起来特别难受。3.2 任务粒度的选择与SM资源复用任务粒度直接决定了两个东西一个是负载均衡的精细程度另一个是任务的调度频率。我实测下来在H100这种SM数量多、SMEM容量大的卡上任务粒度选择要特别关注SMEM的复用程度。每个专家FFN的权重参数都是要提前load到SMEM里才能高效计算的。如果一个专家被切成了很多个小任务而这些小任务分布在不同的SM上执行那每个SM都要重复load一份完整权重。这意味着任务切得越碎权重被重复加载到SMEM的总量就越大HBM到SMEM的带宽压力也越高。Weave把这些任务织进一个大kernel之后天然就有机会做SMEM层面的复用——同一个SM连续执行同一专家的小任务时权重只需要load一次可以连续算好几块token。所以论文虽然没有明说但我推测它的任务队列里会尽量保持“同类专家任务相邻”的排列顺序或者用LIFO后进先出策略让同一个SM倾向于连续消费同一类任务用局部性换带宽效率。这一点在做自己的实现时务必考虑进去否则调度均衡是有了带宽开销却会把收益吃掉一大块。3.3 和StreamK的关系一个家族里的两种打法聊到细粒度调度和wave quantization绕不开NVIDIA CUTLASS里提出的StreamK。StreamK解决的是单个大GEMM的wave quantization问题思路是把一个GEMM在K维上切成细条用类似工作队列的方式让多个CTA动态协作完成。而Weave解决的场景更进一步不仅是单个GEMM内部而是要横向跨越多个不同的GEMM多个专家FFN做调度。可以把StreamK理解成“把一个巨大的砖块切开分给所有人”而Weave是“把一堆大小不一的砖块标好号让每个人按需来拿”。两者核心哲学是一致的——拒绝静态分配、拥抱动态调度。但Weave需要额外处理跨任务切换的开销和依赖管理。如果你已经看过StreamK的代码再看Weave会发现很多眼熟的机制比如任务计数、全局队列、原子操作分配只是应用的层次和范围不同。我的建议是先读StreamK的源码理解CTA级调度再来看Weave会顺畅得多。3.4 和基准实现的对比要站在同一水平线上论文里的2.89倍加速需要理解它对应的baseline是什么。如果baseline是每个专家独立启动一个kernel、CPU串行launch那别说2.89倍在一些模型配置下我甚至见过4倍以上的差距。如果baseline已经是融合了大kernel、但使用静态SM切分的实现那2.89倍就是实打实的调度收益。我自己复现类似方案时得到的经验是对比加速倍率最好同时报两个数——对朴素多kernel baseline的加速以及对静态融合baseline的加速。这两个数分开讲才不容易误导人。论文给出2.89倍应该是相对一个已经做了基本融合、但调度方式还是静态/半静态的实现这说明2.89倍的收益主要来自调度策略变化而不是简单的kernel合并。4. 实测效果2.89倍到底从哪挤出来的4.1 测试环境与配置还原先还原论文的测试场景。我看过的信息指向的配置是4×H100 SXM应该是80GB版本使用张量并行或者专家并行方式部署了一个标准规模的MoE transformer层。模型配置上比较典型的是8个专家、top-2路由每个专家内部是标准的MLP结构中间维度和隐层维度比约为2.5到4倍。输入batch和序列长度方面论文应该是在推理场景的典型shape下测的而不是那种专门掩盖调度问题的超大batch。坦白说这类论文的绝对数字受shape影响很大。MoE kernel的加速效果在小batch下特别明显——因为小batch意味着每个专家拿到的token很少负载更不均匀wave quantization也更严重。大batch下所有专家都塞满了静态分配也够用调度优化的空间反而变小。所以引用这个2.89倍数据时一定要带上shape背景不然用户在小batch场景可能觉得收益没到在大batch场景可能觉得远超预期两个方向都容易产生误解。4.2 加速的构成拆解把2.89倍拆开来看大体上来自三块收益的叠加。第一块来自wave quantization的消除。动态调度让最后一个wave不再空转正常情况下这一项就能带来1.2到1.5倍的提升取决于任务量的分布形态。任务量越接近SM数量的奇数倍收益越大。第二块来自跨专家负载均衡。当某些专家很忙、某些很闲时空闲的SM会主动去队列里领忙专家的工作整体完成时间从“最忙专家的计算时间”变成“总工作量除以总SM数”。在负载倾斜严重的情况下这一项可能带来1.5倍以上的提升。第三块来自启动和同步开销的降低。整个MoE计算链从多个kernel变成一个大kernel每层省掉几次CPU发起的launch和全局同步这个在小batch下收益尤其明显大约在1.1到1.2倍。这三块不是简单相乘的关系因为它们本质都是在提高SM空闲时间的利用率。但方向是一致的——2.89倍这个数字本质上是把GPU从“因为静态调度而被迫闲置的那部分时间”里抢了回来。我自己的经验是如果你的MoE推理任务里NVML监控显示SM占用率长期低于70%那用这套思路重写kernel后性能翻倍是完全合理的预期。4.3 什么场景收益最大什么场景收益有限收益最大的场景首先是小batch推理。这种场景下每个专家分到少量token静态调度基本没法看动态调度几乎是唯一解。其次是专家数量多的模型专家越多负载倾斜越严重静态切分的碎片化浪费越明显。还有一个容易被忽视的场景是PagedAttention这类需要处理不规则batch的推理框架——batch大小动态变化时静态编译的kernel天生吃亏因为编译时根本不知道负载分布。收益有限的场景也值得讲讲。如果batch特别大每个专家都满载所有SM都持续在算那动态调度的空间就很小了。另外如果专家中间维度和token数量的比值极其小——算一个专家的FFN只需要几个周期——那调度开销会高过计算收益这种情况不如直接用静态分配。简单判断标准单个专家任务的平均计算时间至少要超过调度开销一个数量级动态调度才划算。5. 常见问题与排查心得自己写类似调度时的坑5.1 原子操作竞争导致的性能塌方我最早尝试做这种设备端队列调度时第一个版本直接用全局atomicAdd做任务分发结果在小规模测试里跑得不错一放到H100满SM的配置下立刻性能塌方。排查了很久才意识到是原子操作竞争——132个SM同时去抢一个计数器缓存行不停失效光等原子操作就浪费了大半时间。解决办法是前面提到的本地缓存队列。每个SM每次一次性从全局队列拿8到16个任务到本地后续只从本地队列消费。全局原子操作的次数从“每个任务一次”降为“每8到16个任务一次”竞争压力直接下降一个数量级。这个优化不是可选项几乎是一个必须项。5.2 任务顺序对内存带宽的影响还有一个很容易踩的坑是任务顺序。动态调度的自由度过高时会随机地让不同SM执行各种专家的任务。如果两个相邻执行的任务来自完全不同的专家就需要频繁重新加载SMEMHBM带宽会被重复读取参数打满而计算单元反而闲着。我的做法是在任务队列里按专家ID做桶排序相同专家的任务尽量连续排列。调度器取任务时优先从当前SM正在处理的专家桶里取空桶了再去别的桶。这个优化本身不影响负载均衡但对带宽敏感的模型可以带来30%到50%的性能差。我甚至在调试时遇到过一种情况动态调度比静态调度还慢找遍原因就是任务顺序太乱导致SMEM缓存命中率崩了重新排序之后性能立刻拉升回去。5.3 调度粒度与通信重叠的取舍在4×H100这种多卡环境里还有一个需要额外注意的点MoE的token分发和结果收集通常需要跨卡通信All-to-All。传统实现里通信发生在permute和combine阶段计算阶段是独立的。如果用Weave把整个层都织成一个kernel通信和计算的边界就不那么清晰了。真实工程中不应该把所有事情都塞进一个kernel里。比较合理的做法是保留route和permute的独立kernel因为它们涉及跨卡通信需要明确的同步点而把多个专家FFN的计算织进大kernel做动态调度在专家计算之前用一次grid sync等待permute完成。通信和计算的重叠可以通过双缓冲做到一张卡在算前一批的token时另一批token的通信已经在后台进行。论文的2.89倍虽然写的是单层加速但工程上要拿到这个数字partially需要通信和计算的精心编排不能只盯着SM调度一个环节。5.4 如何判断你的MoE层适不适合用这套方案最后给一个很实用的判断方法。先对你的MoE层做一次简单的nsys或nvprof分析拿到两个数据平均SM占用率和kernel launch耗时占比。如果SM占用率在多数iteration低于75%或者kernel launch和同步时间占整个层耗时超过10%你的场景就是Weape这套思路的目标场景。如果两种情况都不满足那说明你的负载非常规整、batch非常饱满硬套动态调度反而会因为增加复杂度而引入性能回退。技术选型这事儿没有银弹Weave这种细粒度动态SM调度是把“GPU利用率”这个单点指标玩到了极致但这不是万能药它解决的只是MoE调度那一层的问题真正要把MoE推理性能做到极致还需要把注意力放到计算、访存、通信三维度的联合优化上。我自己看论文的最大收获其实是确认了“动态调度”这个方向在MoE场景的巨大潜力。以前总有人觉得MoE内核的性能瓶颈是内存带宽或者干脆猜测是GPU不够多。但Weave的实验结果表明大量的算力潜力就藏在那一个个空转的SM里藏在不合理的wave边界里藏在那句“等着某个慢专家算完”的无奈里。把这些浪费的空间填满比单纯堆硬件更值得先做一遍。如果你正在做MoE推理优化我建议先盘一盘手上kernel的SM空闲情况很可能你不需要换卡也能挤出相当可观的速度提升空间。