MoE 大内核做了这么多年大家卷的方向基本都集中在 GEMM 本身的效率上用 CUTLASS 手写 expert GEMM、调 tile 尺寸、搞 split-K、上 warp specialization。但我去年被一个现象反复恶心到明明每个 expert GEMM 都写得挺快了kernel 整体却还是跑不满 SM。后来我才意识到问题根本不在“算”而在“分”——SM 怎么分给不同的 expert比 expert 内部怎么算更能决定整个 MoE 层的墙钟时间。Weeve 这篇论文恰好就切在这个点上。标题里的几个关键词我先拆开说MoEMixture of Experts专家混合模型、大内核把 MoE 层的路由、gather、GEMM、activation、scatter 全熔进一个 kernel、SM 调度GPU 上的流处理器分配策略。它在 4×H100 上跑出了 2.89× 的层加速这个数字如果不看实验设置很容易被高估但扒完细节之后你会发现它确实戳中了一个所有 MoE 推理框架都绕不开的痛点。这篇文章我准备用我自己的理解给你完整过一遍 Weave 的技术思路包括它到底改了什么、2.89× 是怎么拆出来的、什么场景能吃红利、什么场景硬上反而亏最后附上我读完论文之后的复现思路和几个容易踩的坑。1. 为什么 MoE 层的性能瓶颈不在“算得快”而在“SM 分得均不均”1.1 MoE 层天然就不均匀路由分布是重尾的先回到 MoE 本身。一个 MoE transformer layer 的结构是token 先过 attention然后过一个 router路由网络router 给每个 token 选出 top-k 个 experttoken 再被送到对应的 expert 做 FFN 计算。理论上如果 token 均匀分布在所有 expert 上并行效率会很好看但实际推理时根本不是这样。现实里 token 的分布有两个问题路由偏斜某些 expert 就是更容易被选中尤其当模型学到某些“通用特征”时会出现一个 expert 承接远超平均水平的 token 数业内管这叫“专家崩溃”或“路由坍缩”。长尾效应即使平均值均匀单条 prompt 内部、以及 batch 内部的 token 分布也是重尾的。一次 decode 请求里不同 expert 拿到的 token 数可以相差一个数量级。这两个现象叠加导致 MoE 层在 GPU 上天然就是一个负载不均衡的并行任务。你没法在 kernel launch 之前 100% 预知每个 expert 这轮到底要算多少个 token只能等实际路由结果出来才知道。1.2 大内核融合带来的“静态配额”问题MoE 层的朴素实现是“一个 expert 一个 kernel”或者“一个算子一个 kernel”的拆分式实现先 gather再对每个 expert 做 GEMM再做 activation最后 scatter 回去。这种实现的缺点是 kernel launch 次数多、中间结果反复走显存尾延迟很难看。所以大家都开始做大内核fused kernel把整个 MoE 层塞进一个 kernel 里减少 launch 开销和中间读写。但大内核带来一个新问题——SM 配额要在 kernel 启动时就定下来。在 CUDA 的编程模型里一个 kernel 启动时grid 的维度就固定了每个 thread block 被分配到哪个 SM 执行虽然实际上由硬件调度但线程块对应的任务内容是我们自己写的逻辑。到了 MoE 大内核里最直观的做法就是按 expert 的数量开 block每个 expert 分固定数量的 SM。举个具体例子假设有 8 个 expertGPU 有 132 个 SMH100 SXM 就是 132 SM你可能是每个 expert 分 16 个 SM剩下 4 个 SM 空着或者做辅助任务。但这 16 个 SM 的配额是按最大负载预估的。如果这轮某个 expert 只分到 20 个 token另外某个 expert 分到 2000 个 token那你就会看到前者占着 16 个 SM 在摸鱼后者 16 个 SM 在排队算不完。SM 的空转和排队同时存在整卡利用率稀烂。1.3 静态分配的代价到底有多大有人可能会说那我在 kernel 启动之前根据路由结果动态算一下每个 expert 该分多少 SM 不就行了这也是目前很多框架在做的方案但问题在于分发矩阵要等 router 算完才知道这本身就需要一次 kernel/一次同步实时预算分配的时间窗口极短。token 在 expert 上的计算时间不均匀。同样一个 token长度不同、稀疏度不同、expert 内部是否命中缓存都会影响实际耗时。靠 token 数量做静态预估只能估算“工作量”等真正跑到后半段才发现预估偏了但这时候 SM 配额已经无法调整了。尾波效应tail waveGEMM 是分 wave 调度的如果某个 expert 的负载恰好让它的计算量比配额多一点点最后一个 wave 只占用了 1/4 的 SM这个 wave 的执行时间会被完整计入而空置的 SM 想帮忙也帮不上。我用一张表对比一下常见的几种实现路线你会更直观地看到问题出在哪实现路线SM 分配方式负载不均衡时的表现额外开销每 expert 单独 kernel每次 launch 独立调度天然负载均衡但 launch 多、中间读写多kernel launch 开销、显存带宽开销融合大内核 静态 per-expert 配额按预估 token 数预分 SM偏斜场景下 SM 空转 vs 排队并存预估偏差导致资源错配融合大内核 细粒度动态调度Weave 路线运行中按任务量自适应认领SM 始终有活干负载偏斜被摊平原子操作/队列同步开销从这张表能看出前两条路线都是在“要么多花 launch 开销要么赌预估准不准”之间二选一。Weave 选的是第三条路不动 kernel 融合的前提把静态配额改成动态认领。2. Weave 的调度核心逻辑从“任务指定 SM”到“SM 主动认领任务”2.1 把 expert 的计算切成细粒度 tileWeave 的第一个关键动作是把每个 expert 的 GEMM 计算切成小块。传统上per-expert GEMM 是以整个矩阵乘为粒度分配给 SM 的一个 expert 的任务就绑定在一组 SM 上。Weave 的做法是把每个 expert 的 GEMM 按 tile 切分切出来的小计算单元放入一个全局任务队列。这个设计背后是有讲究的。H100 每个 SM 有 128 个 FP32 CUDA core如果算上 tensor core 的话是 4 个第四代 tensor core一次可以处理的 warp 级 GEMM tile 通常是 64×64、128×64 这样的规模。把 expert 级的大 GEMM 切成这些尺寸的 tile 之后工作量就细粒度化了调度单元变小SM 之间互相借力的空间自然就大了。切 tile 这个事很多人做 GEMM 优化的时候都干过区别在于通常我们切 tile 是为了拟合SM 内部的计算管线寄存器、共享内存、tensor core 的流水而 Weave 切 tile 是为了拟合SM 之间的调度粒度。这两个目标的取舍不同tile 尺寸的选择逻辑也不一样。2.2 动态认领机制的实现雏形原子队列 work stealingWeave 的核心调度机制用大白话讲就是别在启动 kernel 的时候就把活派给 SM而是让 SM 干完手里的活之后自己去队列里领下一个活。这其实就是 work stealing / dynamic scheduling 的思路在 GPU 大内核里的应用。我简化一下它的工作原理你可以当成一个伪码来理解// 全局任务队列每个 entry 表示一个 (expert_id, tile_row_start, tile_col_start) // task_counter 是一个原子变量表示下一个待认领的 task 下标 __device__ int task_counter 0; __global__ void fused_moe_dynamic_scheduler(...) { int sm_id blockIdx.x; // 每个 block 绑定一个 SMpersistent kernel while (true) { int task_id atomicAdd(task_counter, 1); // SM 主动认领 if (task_id total_tasks) break; // 从任务队列拿到 tile 的 expert 归属和矩阵坐标 Task t task_queue[task_id]; // 根据 expert_id 决定走哪条 GEMM 分支 // 从对应的 expert 权重里取 tile 数据计算 compute_expert_tile(t); // 算完一个 tile 后继续认领下一个不回 idle } }这里有几个关键的工程决策Persistent kernel一个 grid 的 block 数设置为正好等于 SM 数或者 SM 数 × 每个 SM 的并发 block 数Kernel 启动后每个 block 常驻一个 SM循环认领任务直到队列清空。这样避免了反复 launch kernel 的调度开销也保证每个 SM 只要还有任务就一定有活干。细粒度任务的顺序无关性MoE 层的每个 token 都要经过自己对应的 expert 计算不同 expert、不同 token 之间是独立的不存在严格依赖这决定了任务队列可以完全乱序执行。这是 Weave 能这样做的前提。任务粒度与原子操作开销的平衡tile 切得越细负载均衡越好但原子队列的竞争越激烈。原子操作在 GPU 上的吞吐是有限的如果任务队列的竞争本身成了瓶颈收益就会被吃掉。tile 尺寸需要参考 GEMM shape 和 SM 数量来取一个折中。2.3 为什么不用现成的 stream / 多 kernel 方案你可能想说这种动态调度听着不难啊为什么不直接用多个 stream 并发跑不同 expert 的 kernel或者说expert 并行度不够我给每个 expert 开多个 kernel 不就行了这里有三个层面要解释:stream 的调度粒度是 kernel 级太粗了。一个 expert 的 GEMM 内部如果负载不均衡stream 层面完全感知不到也没法把一个 expert 的尾部计算挪给另一个 stream 去帮算。过度切分 kernel 反而提高总开销。如果每个 expert 的 GEMM 都切成一堆小 kernellaunch 开销的增长是线性的而且小 kernel 根本喂不饱 GPU。数据复用问题。大内核里 token 的 gather 结果、中间 activation 都可能留在 L2 cache 或 shared memory 里复用。一旦拆成多个 kernel这些中间数据大概率要被冲刷掉多付出的显存带宽成本比调度省下来的时间还多。Weave 选择“一个 kernel 细粒度认领”是因为它把调度开销压到了原子操作级别而不是 kernel launch 级别这样每个 SM 的闲置窗口被压缩得极短。2.4 配合 warp specialization把“等待”变成“生产”只做动态认领还不足以解释 2.89× 这么高的加速倍数。我推测 Weave 在实现里还做了一层warp specialization也就是把 SM 里的 warp 分成不同角色一部分 warp 负责从显存加载权重和 token 数据producer一部分 warp 负责实际 GEMM 计算consumerproducer 和 consumer 之间通过 shared memory 做流水线。这一步的意义在于SM 认领任务之后不能立刻开算得等数据从显存到寄存器。如果这个等待时间由 SM 自己承担SM 还是在空转。warp specialization 让 producer warp 提前预取下一个任务的数据consumer warp 无缝衔接计算SM 的利用率才能逼近百分百。这本质上是把 CPU 流水线里经典的“取指-执行”分离思路搬到了 GPU kernel 内部。配合 H100 的 shared memory 容量228 KB per SM和异步拷贝指令cp.async这个流水线可以做到数据搬运和矩阵计算并行不悖。3. 4×H100 实测 2.89×这份数据到底该怎么读3.1 实验设置决定了数字的“含金量”标题里的“4×H100 实测 2.89× 层加速”这个数字要分三层看第一层基线是什么。如果基线是最朴素的每 expert 一个 kernel 的拆分式实现那 2.89× 里包含了 kernel launch 节省 数据复用 SM 调度三部分收益很难算清 Weave 的调度机制单独贡献了多少。如果基线是已经优化过的 fused MoE 内核但用静态 SM 分配那 2.89× 就几乎全是调度策略的功劳。我个人比较确定后者的可能性大因为论文标题强调“细粒度动态调度”对照实验多半是静态分配版本的 fused kernel。第二层负载分布有多极端。MoE 的路由偏斜程度和 batch 大小强相关。离线 prefill 大 batch 下token 分布相对均匀动态调度的收益偏低在线 decode 小 batch 下路由偏斜显著收益偏高。2.89× 大概率是在一个偏斜比较明显的负载形状下测出来的这个数字不代表所有情况。第三层“层加速”不等于“端到端加速”。层加速只是 transformer 某一层的时间缩短倍数。MoE 层虽然是大头但整个模型还有 attention、embedding、norm、router 等其他部分。如果 MoE 层占了模型端到端时间的 60%那 2.89× 的层加速折算到端到端大概就是 1.54× 左右0.6/2.89 0.4 换算具体打多少折扣取决于模型里 MoE 层的占比。3.2 2.89× 应该拆成哪几笔收入我拿一份真实的工作负载拆解一下这个加速倍数的构成方便你对照自己的场景估算优化维度收益来源典型贡献占比估算SM 空转消除偏斜负载下空闲 SM 被利用起来40%–50%tile 级任务并行尾部 wave 被其他 SM 接力完成15%–25%warp specialization 流水线数据加载与计算重叠等待时间被隐藏15%–20%数据复用L2 / shared memory融合内核避免了中间结果往返显存10%–20%注意这四笔收入不是线性叠加的它们之间存在交互。比如动态调度本身就能减少尾部 wave但 warp specialization 把等待隐藏之后SM 的“有效计算时间”变多了动态调度的负载均衡效果会被进一步放大。3.3 实验图表里值得盯的三个关键趋势论文实验部分我建议重点关注三张图的趋势不同 batch size / 路由分布下的加速比曲线如果加速比随 batch 增大而递减说明动态调度吃的是“分布不均”这碗饭分布越均匀收益越小这是符合预期的验证。不同 expert 数量下的扩展性expert 越多、每个 expert 的负载越稀疏静态分配的浪费越严重动态调度的优势越明显。8 expert 到 64 expert 的收益斜率值得关注。tile 尺寸的消融实验tile 越细负载越均衡但原子竞争越激烈一定存在一个最优区间。看论文给的曲线能反推他们对原子开销的容忍度。顺带一提4×H100 的环境也值得解释一下。MoE 推理经常做 tensor parallel / expert parallel4 卡意味着每个 expert 的权重被切到 4 张卡上每张卡处理 1/4 的 expert 分片。Weave 的调度是在单卡内的 kernel 层面做的和跨卡的 expert parallel 是正交的。换句话说4×H100 可能是单纯为了模拟真实推理环境下的单卡内负载不排除还有跨卡通信的影响没被计入“层加速”里。3.4 和已有 fused MoE 实现的差距现在业界能拿到的 fused MoE 内核包括 DeepSpeed 的 fused MoE、vLLM 的 grouped GEMM 路径、以及各家基于 CUTLASS 手写的版本大多走的是“静态 or 半静态分配”的路线。它们的 prefetch 和 GEMM 本身优化得都已经不错了。Weave 相对这批实现的本质性差异不在 GEMM 计算而在任务分发层。可以这么理解之前的实现是“管理员提前排好班工人按排班表干活排班表错了就等着”Weave 是“没有排班表工人干完手上的活就去任务池里抢下一个谁空谁干”。GPU 大内核的负载不均衡问题被从“预估问题”转化成了“调度问题”而调度问题在计算独立的场景下是一个可以用原子队列近乎完美解决的问题。4. 这个方案能吃到多少红利边界条件与适配判断4.1 收益最大的负载画像根据 Weave 论文的思路能充分吃满动态调制的负载大概长这样expert 数量中等偏多比如 8 到 128 个 expert每个 expert 的 token 负载彼此差异大。单 token 计算时间波动大不仅仅是 token 数不均token 本身的计算量也有差异比如变长序列、padding 带来的计算浪费。SM 总数相对 expert 数不太悬殊如果 expert 数远大于 SM 数每个 expert 分到的 SM 本来就不多静态 vs 动态的差异会被稀释如果 expert 数远小于 SM 数每个 expert 能分到足够多的 SM偏斜也不容易造成大面积空转。SM 数和 expert 数处于同一数量级比如 H100 的 132 SM 对 8–64 expert时动态调度最值钱。4.2 收益没那么大的场景反过来有几类场景不需要上 Weavebatch 极大且 padding 多大 batch 下 token 分布接近均匀静态分配的开销不明显。这个规律在大模型推理里很常见——负载越整齐调度的存在感越低。单 expert 计算量巨大大 hidden size每个 expert 的 GEMM 本身就够大tile 切分之后队列依然很长负载均衡的压力不大瓶颈又回到 GEMM 自身效率上。expert 之间有数据依赖某些 MoE 变体如 shared expert routed expert 组合、或者 multi-head MoE 的跨 expert attention存在依赖不能随意乱序执行任务队列的乱序认领就受限了。显存带宽瓶颈的场景如果模型是 memory-bound比如权重巨大、batch 小、每次只算少数 tokenSM 负载均衡做再好也没用反正大家都在等数据。4.3 硬件平台的影响H100 上的优势换个卡可能缩水Weave 在 H100 上能实现近 3× 的层加速有一部分是 H100 硬件特性给的132 个 SM规模足够大动态调度的“借力”空间大。如果是 80 SM 的 A100或者 18 SM 的消费级显卡SM 冗余变小调度的绝对收益会下降。228 KB 的 shared memory 和cp.async指令让 warp specialization 的数据预取流水线有充足缓冲。老卡 shared memory 只有 96–164 KB流水线深度不够等待时间藏不住。第四代 tensor core 的算力推进也让计算时间变短反过来更能暴露数据搬运和调度的等待让隐藏开销这件事变得更值钱。我在 A100 上做过一个类似的简化版测试这里不展开细节结论是收益大概比 H100 缩水 1/3 左右但依然可观。所以 Weave 不是 H100 专属只是 H100 放大了它的优势。4.4 落地之前先算的四笔账如果你想让自己的推理框架吃到这波红利动手前建议先把四笔账算清楚。账目算法判断标准调度粒度收益账当前 kernel 里 SM 利用率曲线波谷面积SM 空转面积占比 15% 才值得上原子竞争开销账tile 数量 total_tokens × tile_ratio算每 SM 平均认领次数每 SM 每微秒认领频率过高则需放大 tile显存带宽账看 fused kernel 的 memory throughput 是否接近峰值带宽利用率 85% 时调度优化意义不大端到端收益账MoE 层耗时占比 × 预期层加速折算端到端端到端加速 1.1× 直接弃坑这四笔账是我个人在实际项目里评估优化方案时的通用框架。很多团队一上来就兴冲冲改调度器结果发现自己的负载压根没有偏斜问题改完之后收益全被原子操作开销吃掉了。先算账再动手能省不少时间。5. 论文之外的实操思考复现思路、潜在坑位与延伸方向5.1 最小可行复现给你的 fused MoE 内核加一个任务队列如果你不想完全照搬论文的整个系统毕竟要配合 warp specialization、cp.async 等一堆细节我个人建议先做一个最小验证版本只需要三步第一步把你现在的 fused MoE kernel 里的 per-expert GEMM 拆成 tile 级任务用一个显存的 int 数组当任务队列每个 entry 记录 (expert_id, tile_row, tile_col)。第二步把 kernel 改成 persistent 风格gridDim 设为 SM 数或 SM 数 × 2每个 block 循环atomicAdd认领任务。第三步先不做 warp specialization只在认领后直接计算跑通后再加 producer warp 预取。这个最小版本也许只能达到 Weave 的六成收益但足够验证“你的负载在动态调度下有没有改善”。我从经验上说大部分团队卡在第一步的“切 tile”逻辑上——不是切不动而是切完之后要处理对 shared memory 的复用以及不同 expert 权重矩阵的 alignment 问题这些细节比调度本身更费功夫。5.2 几个容易踩的坑第一个坑是原子队列的竞争热点。如果 tile 切得太碎所有 SM 同时去抢同一个atomicAdd计数器这个全局原子操作会成为新的瓶颈。缓解办法是采用两级队列每个 SM 先认领一批比如 16 个 tile到本地算完再认领下一批减少全局原子次数。这本质上是用批量换取原子竞争下降。第二个坑是wave 齐整度的假象。你以为动态调度把负载摊平了但如果每个 tile 的执行时间本身差异巨大比如有的 tile 走的是 4 个 expert 的权重有的只走 1 个队列里会出现“长尾执行块”——前面的活都干完了最后一个大 tile 还在慢慢算。这时候需要进一步细切大 tile或者按 expert 的计算量给任务加权。第三个坑是profiler 采样对动态调度的干扰。Nsight Compute 这类 profiler 在记录 kernel 活动时会天然优化或干扰调度的执行路径尤其对 persistent kernel 的循环内部分支采样的不准确度较高。测动态调度内核时我建议用clock64()在 kernel 内部自己做计时桩记录每个 SM 的空闲时间和总耗时比外部 profiler 更可信。第四个坑是CUDA graph 捕获的兼容性。很多推理框架用 CUDA graph 减少 kernel launch 开销但 persistent kernel 动态任务队列里如果有依赖于运行时数据的循环CUDA graph 捕获时会遇到困难。如果你的框架重度依赖 graph 捕获得确认动态调度 kernel 的循环边界是可静态确定的或者给 graph 捕获留一个 fallback 路径。5.3 从“层加速”到“端到端收益”的延伸思考回到标题里的 2.89×。即便这个数字在最优负载下测得它对推理系统的价值也不能只看 MoE 层本身。我比较欣赏 Weave 把问题聚焦在层级别的做法——它把一个极其复杂的系统问题推理框架的端到端延迟切成一个可验证的 kernel 级问题让人能清楚地知道收益边界在哪里。层加速之后下一步自然是和跨卡通信优化、KV cache 调度、投机解码这些技术做乘法。对我来说读这篇论文最大的收获不是那个 2.89×而是它提醒了一件事在一个所有组件都被极致优化的系统里“调度策略”往往是被最后想起的瓶颈。MoE 大内核里 SM 怎么分配、任务怎么认领这类看似简单的问题因为太底层反而容易被经验丰富的人默认成“硬件自己会处理好”。Weave 证明了这个盲区里藏着近 3 倍的性能空间。如果你在自己的框架里也观察到 fused MoE kernel 跑不满 SM先用 profiler 看一下每个 SM 的活跃度分布。如果活跃度曲线是锯齿形的恭喜你你和 Weave 论文作者看到的是同一个问题。下一步就是像他们一样别折腾 GEMM 本身了去折腾任务是怎么被分到 SM 上去的。