MoE Kernel动态SM调度：从静态分配到任务池，4卡H100加速2.89×

1. 为什么MoE Kernel在4张H100上依然跑不满先说我遇到的一个典型场景。去年帮朋友调一套MoE模型的离线批量推理服务4张H100按理说单卡132个SM、4卡528个SM处理一个8专家、top-2路由的MoE层应该非常轻松但Nsight Compute一打开SM利用率长期在45%到55%之间晃部分SM的活跃周期只有30%出头。当时第一反应是显存带宽不够Kernel都跑在tensor core上应该不至于。后来一行一行看每张卡上各SM的调度波形问题才真正暴露每个SM手里分到的活严重不均有的SM从头忙到尾有的SM后半程完全空转。这个现象并不是个例。只要用过MoE做训练或推理都会碰到专家冷热不均的问题。刚接触的朋友可能会觉得这是负载均衡loss没调好加大aux loss系数、让router把token摊平不就行了。但实际上路由分布即便在统计上均匀单个batch的瞬时分布依然是剧烈的。Weave这篇论文的核心动作就是把这个问题重新定义了一遍路由层面解决的是token该去哪个专家SM层面解决的是在同一个kernel内部怎么让每个SM都别闲着。两个层面不是一回事。1.1 动态路由带来的专家冷热不均MoE层的前向计算其实不复杂。输入token经过一个gate/router网络每个token按概率选出top-2或top-k个专家然后只激活这几个专家做矩阵乘法。问题就出在“按概率选”这件事上。我一个线上batch是64条请求、每条seq len 512总共32768个token。假设8个专家均匀分布下每个专家应该拿到4096个token。但在真实数据上热门专家拿到的token可能是这个数字的3到4倍冷门专家可能只有几百个token。这个偏斜幅度在不同batch间还会剧烈抖动——上一秒专家3最忙下一秒可能就是专家7。更麻烦的是token到达专家时并不是一个完整的连续矩阵。top-2路由让每个token要被两个专家各算一次而每个专家收到的token在显存里分散在各个位置。所以MoE kernel内部必须要做数据重组把去往同一个专家的token重新拼成一个紧凑的矩阵这个操作叫作“token permutation”或“grouped GEMM前的gather”。一旦做完重组每个专家对应的就是一块独立的GEMM任务。接下来就是调度问题了这块任务交给哪些SM算以什么粒度算任务之间怎么衔接。1.2 大内核模式与静态SM绑定的代价早期MoE实现非常粗暴每个专家单独launch一个GEMM kernel8个专家就launch 8次再加上token重组、缩放、Reduce一层有二三十个kernellaunch开销和kernel间同步损耗大得离谱。后来社区基本都转向了“大内核”方案把整个MoE层揉成一个或几个fused kernel里面用grouped GEMM或者loop内逐个专家计算配合CUDA Graph把kernel launch开销压到极致。大内核带来的新问题就是SM的分配变成了静态的。很多实现的做法是kernel启动时根据每个专家token数量预分配block数量token多的专家分到更多block每个block固定处理某一块数据。听起来合理但执行到后半段就开始露馅——当某个专家的block已经全部执行完原本负责它的SM如果没有从其他专家那里“拿新活”就只能原地等待整个kernel结束。也就是说给定一个batchSM的负载上限是“最忙的那个专家”决定的而不是“所有专家总计算量”决定的。这很像操作系统的处理机调度如果你在启动进程时把CPU核静态分配给线程而不是让线程动态抢任务那么任何一个线程提前结束都会造成核的空闲。Weave干的其实是一件事把SM当成可动态迁移的worker池而不是启动时锁死的资源。1.3 静态SM绑定在什么场景下会空转很多人的直觉是只要block总数大于SM总数就不会有空闲SM。这个直觉在普通GEMM里成立因为GEMM的所有block是同构的哪个block先做完调度器可以把后续block随时补上去。但在MoE的grouped GEMM里不同专家对应的GEMM形状差异巨大而且block和block之间不是无限可互换的——一个block被指派给专家A它就只能在专家A的数据范围内执行。所以在静态方案里即使全局block数量充足也常出现这种情况专家A的block还有20个在排队专家B的block队列已经是0。此时负责“专家B区域”的SM只好提前退休。如果把时间线拉长你会发现每个SM的实际活跃时间只有总kernel时间的一部分这部分“偏斜损失”在路由波动大的batch上尤其明显。有没有办法解决最简单的是在kernel内部加一个global work queue让所有SM执行完当前任务后去队列里抓下一个任务。这个思路听起来谁都能想到但真正做起来有一堆细节Weave的价值就在于把这些细节做成了一套可复用的调度机制。接下来聊聊它的具体设计。2. Weave的调度模型把静态SM图变成运行时任务池Weave的核心设计可以概括成一句话不在启动时决定每个SM做什么而是让每个SM在运行时动态领取自己下一个要做的子任务。在4张H100上这意味着把528个SM组织成一个共享的任务池。每个专家不再是“一堆固定block”而是一个待处理队列每个SM则是一个worker执行完当前子任务就通过原子操作到全局队列里取下一个。2.1 从“每个专家固定一组SM”到“每个SM按需领任务”对比两种模型会清楚很多。传统大内核的调度过程是kernel启动前host端根据每个专家token数算出每个专家分配多少个block然后作为参数传进kernelkernel按block的全局ID天然归属到某一专家的任务。这个方案没有运行时决策好处是简单、可预测坏处就是前面说的静态空转。Weave模型则是kernel启动时只初始化M个任务队列每个队列对应一个专家路由完成后的token矩阵被切成若干子任务挂到对应专家的队列尾部。之后所有SM进入同一个循环从全局索引取出一个待办子任务执行GEMM执行完再取下一个。这里没有“这个SM属于专家A”的概念只有“这个SM现在正在处理专家A的子任务”。从负载均衡角度看这就把静态分区变成了一种动态work stealing。专家A任务多占用SM时间就长专家B任务少它的SM做完后会自然涌向专家A的队列把专家A的积压吃掉。整个过程不需要host参与全部在一个kernel内部完成。2.2 队列结构、原子操作与内存序动态调度的工程基础是队列结构。每个专家队列由一个或多个“子任务描述符”组成描述符里至少要包含数据在重组后buffer中的偏移、token数、对应的专家权重指针。SM每次取任务需要修改一个全局索引这里必须用原子操作保证并发安全。具体到CUDA实现常见的做法是每个专家维护一个32位计数器记录该专家已分配的子任务数。某个SM想取专家i的任务时用atomicAdd(counter[i], 1)得到返回的旧值如果旧值小于该专家任务总数就说明这次领取成功按旧值索引去执行对应子任务如果已经超过总数说明专家i没有任务了SM就去检查下一个可用专家。这里有一个非常关键的工程细节如果每个SM每次只取一个子任务就去抢全局锁原子竞争会直接把调度开销打到不可接受的地步。实测中528个SM同时对一个计数器做atomicAddL2带宽很快就被击穿。Weave的处理思路是让SM成批领取任务——一次抓取4个或8个子任务到本地工作列表处理完再回来领取下一批。这样原子操作的频率下降了10倍以上对L2的冲击小得多。我印象里论文中专门花了篇幅论证这个batch size的选取它本质上是在“调度的公平性”和“原子操作开销”之间找平衡点。2.3 执行循环与任务完成同步有了队列和原子操作剩下就是SM的执行循环了。伪代码大致是这样// 每个SM的入口函数persistent kernel风格 __global__ void weave_moe_kernel(TaskQueue* queues, int numQueues) { int smId get_sm_id(); LocalTask localTasks[FETCH_BATCH]; while (true) { int remain fetch_tasks(localTasks, queues, numQueues); if (remain 0 all_queues_empty(queues)) break; for (int i 0; i remain; i) { exec_gemm(localTasks[i].expert_id, localTasks[i].token_offset, localTasks[i].token_count); } } }关键在退出条件。你不能因为某一个SM发现所有队列暂时是空的就让它退出因为其他SM可能还在执行中之后会产生新的“空位”不会因为任务总量是启动时定好的不存在执行中产生新任务的情况所以只要所有队列都空了就意味着所有GEMM子任务都已被领取当前这批SM可以安全退出。但这里有一个容易忽略的点“被领取”不等于“已执行完”。后续的分布式AllReduce或token交换需要的是整个MoE层的输出都写回显存。所以kernel在真正收尾前必须有一个全局barrier通常用atomicAdd一个完成计数器实现确保所有领取的任务都执行完毕。这个barrier如果写得不好会让所有SM在末尾空等静态方案的“尾部空转”问题就会换个姿势重新出现。我的实际体会是这部分是Weave类似机制里最难调的地方。队列空得太早说明fetch batch太小调度太频繁空得太晚说明batch太大负载均衡效果变差。必须用真实路由分布反复校准。3. 细粒度切分的技术选型调度粒度与tensor core效率的平衡动态调度听起来非常简单但粒度选不好加速比会在另一个维度上损失殆尽。GEMM子任务的粒度要同时满足两个条件足够小让SM之间能灵活均衡足够大让tensor core的矩阵运算效率不被浪费。这二者天然矛盾。3.1 为什么不能直接把整个专家任务当一个子任务如果把一个专家的全部token拼成一个大的GEMM当成一个子任务挂到队列上那动态调度就退化成“先到先服务整个专家”反而比静态分配更糟——一个专家占用大量SM跑完后面任务才能开始完全失去了并行性。如果把粒度切得太细比如把一个专家的token切成单个token级别的子任务每个SM取一个token做GEMV问题更大。H100的tensor core每次矩阵乘的形态是m16n8k16这类小块单个token对应的是m1矩阵乘法的数据复用彻底消失等效成了低效的GEMV吞吐可能掉一个数量级。动态调度的收益再高也补不回来。3.2 合理的子任务形状选择与计算换算实践中一个子任务至少要保证在K维度上有足够的复用同时M维度不要太小。我推荐按“专家内连续token块”来切。假设专家A收到了4096个token隐藏维度hidden4096权重形状是[hidden, expert_dim]也就是[n4096, k4096]这里的K其实是hiddenN是expert输出维度。如果按128个token切成一个子任务每个子任务就是一次m128n?k4096的GEMM。128个token对应m128这个规模已经能让tensor core吃到比较高的利用率同时128个token的粒度又足够细让SM之间的均衡非常灵活。选择切分个数还需要考虑另一个约束任务的领取次数不能太少。如果所有专家总共只有几十个子任务那么动态调度几乎退化成静态分区只有几百个子任务的粒度才是理想的。因为4张H100上528个SM如果一次kernel有800到1500个子任务每个SM平均能领到1到3个批次调度器就有足够的空间去纠正偏斜。我自己的经验公式是子任务总数尽量大于SM总数的2到4倍。太少则均衡能力不足太多则原子操作和调度分支开销偏高。低于这个阈值的时候优先增加切分数高于它的时候优先保证每个子任务m值不小于64。3.3 切分后的数据布局重组buffer该怎么排子任务切分还牵扯一个重要工程问题——token重组的buffer布局。MoE kernel先要按专家把token gather到一起如果后续是动态调度那么这个gather最好提前按“子任务连续、任务之间按专家分组”的方式排布这样SM领取任务后可以直接读一段连续的显存地址不用再额外做索引间接寻址。否则每个子任务都要从描述符里解析出散落的token位置GPU的L2和内存访问模式会变成乱序带宽利用率下滑很明显。论文里应该也提到预排布数据是动态调度能跑出好成绩的前提之一我在自己的实现里验证过乱序索引会让有效带宽掉大概15%到20%。3.4 反向传播怎么办训练场景下MoE层不仅要做前向GEMM还要做反向的weight gradient和activation gradient。反向的计算形状与token分布一一对应因此同样存在专家冷热不均——甚至更严重因为反向的两个GEMM都要处理同样的偏斜分布。好消息是Weave的队列模型天然适用于反向把反向的两个GEMM也切成子任务挂到同一个任务池里SM做完一个就去领下一个。唯一要注意的是反向GEMM的K维度往往是前向的N维度形状不同子任务的最优切分大小也不同最好分别配置不要共用一套参数。我在调优时发现前向和反向用同一套切分参数效果一般单独调反向后整体训练吞吐又提了10%。4. 4×H100实测2.89×实验配置、加速比拆解与适用范围标题里的2.89×是这篇论文最抓眼球的数字。但要理解这个数字必须先搞清三个问题硬件配置是什么基线是什么测的是哪个阶段。否则很容易高估或低估它的价值。4.1 实验配置与基线从论文公开信息来看实测环境是4张H100MoE层按常见的8专家配置路由top-2隐藏维度取4000量级测试的batch规模应该是偏大的——至少是离线批量推理的规模因为batch太小的时候MoE kernel总执行时间太短动态调度本身的原子操作开销会吃掉大部分收益。基线是传统的静态SM分配方案也就是不带动态调度、把block按专家token数静态划分的fused MoE kernel。这一点很重要因为如果基线是per-expert kernel launch那2.89×就没太大说服力——那种实现本身就慢。Weave的基线是已经做过CUDA Graph和大内核优化的方案在这个基础上再提升2.89倍水分就少很多。4.2 加速比从哪里来一个简单的偏斜模型为什么能到2.89×我用一个简化模型说明假设专家数E8静态方案按token数比例分配SM。假设当前batch最忙的专家占据了全部token的35%即 p_max 0.35。均匀情况下每个专家应该只有1/812.5%所以静态方案里最忙专家成为整个kernel的执行瓶颈。理想动态调度下所有SM都在处理堆在队列里的任务整体执行时间约等于总计算量除以SM总数而不是由最忙专家独占决定。此时理论上加速比约为 p_max / (1/E) 0.35 / 0.125 2.8倍。如果这个batch的路由偏斜更极端比如p_max达到0.38理论加速比就到3.04倍。所以2.89×这个数字本质上约等于“负载峰均比”说明它的加速几乎全部来自把SM的总工作量拉平而不是某个魔法级的算子优化。下表给出不同偏斜程度下理论加速上界的参考值E8均匀作为基准最忙专家token占比静态执行相对理想情况下限动态调度理论上界0.151.21.20.252.02.00.302.42.40.352.82.80.403.23.2这里要多说一句理论加速上界不是常态路由偏斜程度随业务数据变化。如果线上batch的token分布一直很均匀动态调度的收益会明显缩水。反过来如果你们的业务有典型的“热门主题集中”现象——比如大量query集中在少数领域——那路由偏斜会非常严重2.89×甚至更高都是可以复现的。4.3 “层加速2.89×”不等于端到端2.89×这个点必须单独讲很多人一看到2.89×就以为整服务快了2.89倍这是误读。“层加速”指的是单个MoE层的kernel执行时间缩短了2.89倍。但一次完整前向里还有attention、FFN、embedding、norm等一系列操作MoE层只占其中一部分。端到端加速比应该这样估算假设MoE层占整体前向时间的50%。单个MoE层时间变为原来的 1/2.89 ≈ 34.6%。那么整体时间变为 0.5 0.5 × 0.346 0.673端到端加速约1.49倍。如果MoE层占比更高比如80%端到端加速约1/(0.20.8×0.346)2.0倍。所以论文报2.89×层加速放到真实模型里可能是1.5到2倍左右的端到端收益这已经是相当可观的数字了。我个人更关心的是这个加速发生在哪个环节。MoE层的kernel时间是token重组、专家计算、AllReduce三块加起来的。动态调度优化的是“专家计算”这一块的GPU利用率不会减少token重组的开销和AllReduce的通信量。所以在分析自己业务的收益空间时先把MoE层总耗时拆开再乘上专家计算占比才能判断Weave类方案的潜力。5. 工程化落地时容易踩的坑如果你看完上面的机制准备在自己的框架里照搬这个方案先别急。动态SM调度在工程上有一连串坑我按踩过之后的教训从深到浅列一下。5.1 CUDA Graph捕获与动态队列不兼容现在的推理引擎普遍用CUDA Graph把整层kernel捕获成一个图减少launch开销。可CUDA Graph的图结构在执行阶段是固定的你在图捕获时不能依赖运行时才能确定的任务队列指针或动态分配的显存。Weave类方案如果要进CUDA Graph所有的队列buffer、子任务描述符、原子计数器必须在捕获前预先分配好并把地址固化在图里。我在实践中的做法是维护一个专用的workspace池把所有描述符预先置为无效值捕获时kernel始终访问这块固定显存。每次执行任务前由host端往这个workspace里重填描述符然后回放CUDA Graph。等于把“动态”限制在数据写入环节图结构本身保持静态。这个技巧不复杂但如果你一开始没预留workspace后面迁移会非常痛苦。5.2 原子操作竞争与L2带宽前面提过fetch batch size问题我再展开一点。H100上原子操作如果全部落在同一个计数器上528个SM的竞争会直接拖垮L2。解决方向有两个一是每个SM本地缓存一部分“配额”批量领取任务减少原子操作调用次数。二是把计数器分片为每个SM维护独立的局部计数器定时合并到全局。分片方式在多任务重叠时更灵活但实现复杂度高。业界做work stealing时通常两种都试。如果你的场景是8个专家以内的中小规模MoE批量领取就够了专家数超过16建议认真考虑分片计数器。我在8专家场景下对比过批量领取让原子操作总量减少了80%以上kernel时间稳定改善。5.3 静态波次效应与persistent kernel的冲突还有一个容易被忽略的问题静态方案虽然不均衡但它至少没有额外的分支分支判断。动态调度的while循环里每个SM都要反复检查队列状态和退出条件这部分指令本身会占用流水线槽位。在MoE层本身非常短、batch很小的情况下这些分支开销可能超过均衡收益结果不升反降。所以落地前先跑一个快速验证拿当前MoE kernel时间乘上SM数看看有没有显著的空闲周期。如果没有明显空转就别上动态调度老老实实用静态分区。我在一个batch16的场景里测过动态调度的开销反而让kernel慢了8%因为总计算量太小原子操作和队列检查的相对占比太高。动态调度适合大batch、高偏斜、kernel时间长的场景它解决的是资源浪费不是性能魔法。5.4 多卡并行时的队列边界最后说多卡。4×H100如果走张量并行或专家并行每个GPU上跑的MoE层实际上是同一个kernel的多副本Weave的队列范围是单卡内的528个SM不会跨卡。跨卡的任务分配属于通信调度问题那是AllReduce/All2All和NCCL的领域不能指望SM级调度去解决。但如果4张卡每张分别处理不同的batch数据数据并行那每张卡上的token分布和偏斜程度都不同动态调度在各卡上独立生效效果是一致的。所以落地时先想清楚并行策略SM级调度解决的是单卡内的SM空闲跨卡负载不均衡需要配合数据并行或专家并行里的其他机制。混为一谈会走弯路。6. 与负载均衡、显存规划的关系以及我的取舍建议最后把Weave放进一个更大的坐标系里看尤其是和两类常规优化手段的关系很多读者会在这里产生误区。6.1 动态SM调度不能替代路由负载均衡第一个误区是“既然SM能动态调度是不是可以让router随便发不用调负载均衡loss了”。不对。动态SM调度只能解决“专家GEMM计算”过程中的GPU资源复用问题但MoE层的其他成本并不会因此消失。一个token如果被路由到某个冷门专家它依然要做一次完整的矩阵运算只是这次运算被哪个SM执行的问题。如果某个batch里70%的token都涌向同一个专家即便动态调度让所有SM都去算这个专家的任务这个专家的计算量本身也会让kernel时间拉长很多。所以router层面的aux loss、负载均衡loss依然重要它控制的是任务总量的分布而Weave解决的是任务分配到资源后的执行效率。我在生产环境里的建议是先用负载均衡loss把专家的长期分布压平再用动态SM调度吸收剩余的单batch波动。两层各干各的不要用动态调度掩盖路由失灵的代价。实际上论文里应该也是在这样的前提下去做对比测试的如果原始路由分布已经很差层加速比会小于2.89×因为最忙专家的总计算量已经超过了SM总量所能消化的范围。6.2 MoE权重是否全部进显存与这个方案无关网上的热搜里总有“MoE架构要全部参数进显存吗”这种问题。这里顺便解释一下MoE的专家参数本来就在推理时必须全部加载到GPU显存不存在只加载一部分的正常玩法。动态SM调度不改变权重占据的显存量它只是改变kernel内部的计算组织方式。真正能省显存的是激活值和中间buffer——动态调度因为减少了任务排队等待理论上只需要更小的临时buffer来缓存中间gather结果但这个收益通常不大做显存规划时不用把它当成重点。如果你的目标是省显存更有效的方向是KV cache压缩、量化专家权重、或者把部分专家offload到CPU。这些和Weave不冲突可以叠加使用但在我们的讨论范围内不要指望动态调度帮你把模型塞进更小的显存。6.3 什么场景值得用什么场景不建议用综合上面的分析我给一个相对清晰的决策清单值得上的场景MoE层耗时占整体前向的40%以上且kernel时间已经从launch开销中摆脱出来进入了纯计算阶段。离线批量推理或大批量训练batch内的token数量足够大专家间偏斜明显。模型专家数在8到64之间任务粒度可调队列管理的复杂度可控。已经确认profiler里存在大量SM空转周期而不是带宽瓶颈。不建议上的场景在线小batch推理比如batch小于16SM本身利用率就上不去动态调度只会增加开销。MoE层总执行时间只有几十到一两百微秒调度分支和原子操作占比过高。你的现状是per-expert多个kernel launch——先老老实实做大内核和CUDA Graph收益会大得多没必要一步跨到动态调度。6.4 落地时的优先顺序如果决定在团队内部尝试我建议按这样的优先顺序推进先把MoE层改成fused大内核接入CUDA Graph确认静态方案下的SM利用率。通过Nsight Compute抓取每个专家的block执行区间量化偏斜损失。只有偏斜损失超过20%才继续。预分配workspace写一个简单的per-expert队列加全局计数器先不做任何精细调参跑通正确性。用真实路由分布数据分别测试fetch batch为2、4、8、16时的性能画出曲线取拐点。前向稳定后把反向GEMM也接入同一套队列注意反向子任务尺寸单独调。再考虑CUDA Graph捕获、多卡数据并行的组合验证。这套流程里每一步都有独立的收益验证点不容易被“看起来很快”带入坑。我自己的项目按这个顺序走下来最终层加速确实落在2.6到3.0倍之间和论文报的数字基本一致。最后提醒一句动态SM调度不是所有MoE性能问题的银弹但它是静态大内核方案之后很值得做的一步。先把静态方案做到位再谈动态化否则只会换来一堆调试成本。