先说我读完标题的第一反应2.89× 这个数在 MoE 大内核优化里确实属于看着就肉疼的量级但真正让我感兴趣的其实不是这个数字而是标题里那句细粒度动态 SM 调度。我过去很长时间都在跟 MoE 的 kernel 较劲天天看 GPU 算力监控曲线——整体利用率七成显存带宽也吃不满但大内核就是快不了。后来想明白了GPU 并不是缺算力是算力被绑在了一条不均匀的执行流上大量 SM 在等干活。Weave 这篇论文解决的就是如何让空闲的 SM 自己找活干这个老大难问题。这篇文章适合三类人看做 LLM 推理和训练优化的自己手写 CUDA kernel 的还有搞 ML 系统的。我会把论文的核心机制、4×H100 上的实测拆解、以及它能直接搬到工程里的思路一次讲透。1. MoE 大内核的空闲困境2.89× 层加速是闲出来的先说一个经常被误解的点深度学习里说的大内核Large Kernel不是卷积里那个 7×7 大卷积核而是指把一串本来要逐个启动的算子融合成一个巨大的 CUDA kernel一口气在 GPU 上执行完。在 MoE 场景里这个大内核通常包含了路由计算、top-k 专家选择、各个专家网络的矩阵乘、结果的加权合并、残差和归一化这些步骤。为什么要这么干因为每启动一个 kernel都有一次 CPU 到 GPU 的 launch 开销几百纳秒到几微秒不等多 kernel 之间还要把中间结果写回显存再读回来速度和带宽都浪费掉了。把多个算子合并成一个内核相当于后厨不再每道菜单独开一次火而是由一个厨师团队按一道完整工序把所有菜在同一个锅里连续处理完端到端的等待自然就短了。但是 MoE 的大内核有个特殊的难处路由结果是运行时才确定的。门控网络要根据当前输入的 token 动态算出每个 token 应该被派给哪几个专家这带来两个问题——一是不能在编译期静态规划好每个 SM 的负载二是不管你怎么规划专家之间的 token 数量天然不均衡。这个不均衡直接催生了一个 MoE 大内核里臭名昭著的现象kernel 尾部的大量 SM 空转。1.1 尾巴效应排队论里的最后一个波次理解这个问题的关键是波次wave的概念。GPU 硬件以线程块block为单位派发任务一批能同时派发的 block 数量是有限的由所有 SM 的驻留槽位总和决定。举个例子假设一块 GPU 上有 32 个 SM 槽位某次任务里有 60 个 block那第一波 32 个 block 被填满第二波只能填 28 个——最后那半波里就有 4 个槽位空着没活干。这些空着的 SM 虽然功耗不大但也不能干别的就只能干等下一个 kernel 被启动。这个最后半波不齐整的现象在 GPU 计算里专门有个名字叫尾巴效应tail effect。在 MoE 大内核里尾巴效应被三个因素放大得非常严重。第一专家路由的负载是长尾分布。热门专家分到的 token 多冷门专家分到的少冷门专家对应的子任务 block 数量就少最容易形成零头。第二不同专家子任务的矩阵形状差异很大。即便 token 数量相同专家并行切分方式也可能让一部分专家子任务更快完成先完成的 SM 空下来后面的专家还在算。第三MoE 解码阶段还有投机解码、动态批等机制token 数是动态变化的你没法像稠密模型那样把 shape 在编译期焊死。结果就是每个 wave 的尾部都有一部分 SM 提前进入空闲等待。所以这里有个关键结论MoE 大内核快不快不只是看算得有多快更要看 SM 有没有被及时填满。2.89× 层加速本质上是把那些本该空转的 SM 时间给抢回来了。1.2 多 kernel 拆分救不了粒度太粗换汤不换药可能有人会问既然大内核尾部空转那我不做大内核了还是回到多 kernel 拆分的方案每个专家单独 launch空出来的 SM 不就能在下一次 launch 时被重新派活吗这个方向我在工程里确实试过思路是通的但收益非常有限。原因是 kernel 级切换的粒度太粗了。你想让一个 SM 干完活后立刻去执行另一个专家的计算得先把这个 kernel 整体结束让 CPU 重新发起一个新的 kernel中间隔着一整套 launch 机制的开销还要等待前面的数据依赖全部满足。为了在一个 kernel 结束时能快速切换人们通常会叠加多个 CUDA stream 来自然并发让空余的 SM 在下一次 launch 中尽量有活干。但 stream 并发依然受限于 kernel 边界而且管理依赖关系非常麻烦一旦搞错就是 data race每两三个 kernel 就要插一个 event 同步点开销一点都不小。更致命的是多 kernel 方案要求把中间结果写回全局内存再读出来专家权重也要重新从显存加载到寄存器、共享内存这些访存和状态恢复成本会吃掉很大一部分优化空间。所以真正要做的不是kernel 级拆分而是在同一个大 kernel 内部以 SM 为粒度做动态再分配。这就把问题从要不要融合推向了融合之后怎么调度——Weave 走的正是后一条路。1.3 大内核里的精细化需求任务解耦顺序不再重要Weave 这套方案的一个前提是把一个大内核里的工作拆成很多个细粒度子任务而且这些子任务之间不能按照谁先谁后死板地排队。理想情况下一个大 kernel 内的任务流像是一个共享的取餐台每个子任务是一个写好了菜品的打包餐盒哪个 SM 空下来了就去取餐台上拿一个能够立即开始处理的餐盒继续做而不需要在意这个餐盒原本属于流水线的第几步。这听起来简单但实际工程量很大。因为你一旦允许 SM 自由领取任务就需要一套机制来回答三个问题哪些任务是已经被做完的哪些任务是当前 SM 能合法领取的哪些任务虽然排在后边但因为数据依赖还不够格开跑这三点就是我接下来要拆的 Weave 核心设计前两点靠任务队列解决最后一点靠细粒度依赖管理解决。2. Weave 的破局思路让空闲的 SM 自己找活干Weave 的核心贡献可以概括成一句话在 MoE 大内核内部维护一个精细的任务池让每个 SM 在执行完当前任务后能够直接从任务池里领取下一个可执行任务而不是傻等一轮全局 wave 结束。这就把硬件的批量派发变成了软件的按需取取。注意这跟 CPU 上的 work stealing 是同一个思想但 GPU 端把它做出来难度高得多因为 GPU 上有成千上万个并发执行单元同步成本极高任务粒度又必须控制得恰到好处搞不好调度开销比收益还大。2.1 核心机制细粒度任务队列Weave 在每个大内核里引入了一个全局的任务描述符队列队列里的每一项代表一小段计算这段计算的执行函数相同只是输入数据、权重和参数不同。每个 block 一进入 kernel第一件事就是通过原子操作从队列头部取一个任务号然后执行对应的计算段。做完之后立刻再取下一个任务直到队列为空。这段逻辑在 CUDA 里用伪代码写出来其实很少// 简化版MoE 大内核中的自调度任务队列 struct TaskQueue { int head; // 下一个待分配的任务编号 int num_tasks; // 任务总数 }; __device__ int fetch_task(TaskQueue* queue) { int task atomicAdd(queue-head, 1); return (task queue-num_tasks) ? task : -1; } __global__ void moe_layer_kernel(TaskQueue queue, float* input, float* output) { while (true) { int task fetch_task(queue); if (task -1) break; // 根据 task 编号分发到对应的专家计算、合并等子任务 dispatch_task(task, input, output); } }核心就一个原子操作atomicAdd但它把整个调度模型从编译器静态排布变成运行时动态领取。这里要强调一下你不用去关心这个 block 最终落在哪个 SM 上硬件自然会把它派发给一个空闲的 SM。你要做的只是保证队列里的任务足够多、足够细让任何 SM 任何时候都有任务可以取。这个设计跟多 kernel 方案的本质区别在于它没有 kernel 边界。一个 block 算完一个子任务后它所在的 SM 状态寄存器、共享内存、L1 缓存里的权重全部保留直接执行下一个子任务完全不用等 CPU 重新 launch也基本不需要把中间结果写回显存再读回来。这省下来的开销比 kernel 级切换低好几个数量级。2.2 SM 级调度的工程真相软件排队硬件分配很多初次接触 Weave 的人会有一个疑问CUDA 里不是不能指定某个 block 在某个 SM 上执行吗那SM 级调度是怎么做到的这里要说一个容易被误解的工程事实确实在通用 CUDA 模型下程序员无法直接控制 SM 分配SM 由硬件调度器自动管理。但硬件调度器在做分配时遵循一个简单的原则只要有一个 SM 的空槽位并且 grid 里还有未执行的 block就去那里塞一个 block 进来。所以如果你能把任务池里的任务全部转换成尚未执行的 block/任务段的形式暴露给硬件调度器那硬件就会自动把空闲 SM 填满。Weave 的实际作用并不是直接命令某个 SM 干活而是保证空闲 SM 抬头的瞬间总能从任务池里拿到合法的下一个任务让硬件调度器总能找到往空槽里填的东西。这就是为什么标题里用的词是细粒度动态 SM 调度而不是直接管理 SM。粒度越细硬件在块与块之间的选择自由度就越大。一个大的专家子任务如果算一个任务那么一个 SM 跑完这个专家之后就只能干等但如果把这个专家内部的矩阵乘法再拆成若干行块、列块级别的子任务那么完成一个子任务后就能立刻领取下一个子任务空闲窗口就被填平了。2.3 依赖与同步从按层同步到按数据就绪顺序同步动态任务队列解决了SM 有活干的问题但没有解决哪些活现在能合法干的问题。如果一个任务依赖的任务还没算完SM 把它取走了就是数据竞争结果必然错误。传统的融合 kernel 靠的是程序顺序和全局同步屏障来保证这一点但全局 barrier 恰恰是动态调度的大敌——所有 SM 必须等到最慢的那个完成才能继续又把空闲问题带回来了。Weave 的处理方式是把同步粒度从全局降为数据依赖。每个任务描述符上带一个依赖计数ref count只有当它依赖的所有上游任务完成之后这个任务才会进入可领取池。上游任务完成时通过原子操作递减下游任务的依赖计数计到零就把它标记为就绪。这样每个任务都像是交通系统里的信号灯绿灯不是由固定的时间表决定而是由前序车辆是否真正通过路口决定。这个设计有一个很重要的工程特点同步是从数据流推导出来的而不是从代码行号推导出来的。只要上层把任务的 producer-consumer 关系描述好kernel 内部就能按照数据就绪顺序自动执行不需要全局 barrier。对于 MoE 大内核这种天然存在大量并行分支的场景按数据就绪顺序的做法比按代码顺序 反复同步的做法不知道高到哪里去了。2.4 为什么不做 CUDA Dynamic Parallelism可能还有人会问CUDA 本身有动态并行Dynamic ParallelismCDP允许核函数内部再启动新的子核函数听起来也能实现SM 空闲时干新活为什么要自己造一套任务队列我之所以特别想讲这一点是因为我自己在做类似优化时最初也试图依赖 CDP实测下来确实不划算。CDP 的开销来源于三方面第一它仍然要走完整的 kernel launch 路径即使发起方在设备端仍然要经过驱动层的调度程序开销是微秒级的第二子 kernel 的启动通常伴随着栈空间分配和上下文管理比普通全局内存读写贵得多第三嵌套的 grid 会让 Nsight 的分析和依赖管理变得极其复杂调试体验非常痛苦。Weave 这套自管理任务队列避开了 CDP 的三个痛点它不需要启动新的 grid只是当前 grid 内的 block 去队列里取一个新任务它不需要新的栈和上下文因为任务执行函数还是同一个 kernel 内的代码它的依赖状态全部存在显存里通过原子操作维护调试时可以用常规工具直接看数据。一句话总结CDP 是在 GPU 内部重新启动一个小 CPU而任务队列是本来就是同一批工人在共享工厂里按订单干活。前者重后者轻MoE 大内核的高频场景显然适合后者。3. 4×H100 实测2.89× 层加速从哪来、可信度怎么评估论文给出的 2.89× 是层加速layer-wise speedup不是端到端推理加速这两个概念的差别非常大。我在下面按工程测试的习惯把硬件环境、基准设置和加速拆解一条条讲清楚。3.1 测试环境和负载设定4×H100 是验证 MoE 大内核优化最常见的配置因为 H100 的单卡 SM 数量多SXM 版本 132 个 SM4 卡之间通过 NVLink 互连既可以做张量并行也可以做专家并行。Weave 的验证场景通常会选用一个标准的 MoE Transformer 层包含门控路由、top-2 专家选择、多个专家 FFN、结果合并和残差归一。测法上论文一般不会只看层总耗时这一个数字还要用 Nsight/CUPTI 统计 kernel 内的波次分布和各阶段耗时这样才能把动态调度省了多少空转量化出来。我自己在复现同类实验时也是按这个流程走的先跑一个不开启动态调度的 baseline 大内核再跑开启 Weave 调度的大内核注意只改调度策略不改算子和权重保证对照的纯度。3.2 kernel 耗时拆解2.89× 是怎么凑出来的从论文的表述方式和标题信息来看2.89× 是 Weave 相对普通 MoE 大内核 static 版本的层耗时倍差。为了把加速拆开看我用一组归一化耗时做示意数值不代表论文原始数据仅用于构成分析实现方案归一化单层耗时相对静态大内核版本主要耗时构成逐专家多 kernel baseline1.000.44×kernel launch 开销 中间结果读写静态融合大内核0.441.00×尾部 SM 空转约 40%无 launch 开销Weave 细粒度动态调度0.152.89×尾部空转显著下降调度开销约 8%从这个拆解能看出三个信息第一静态融合本身已经把多 kernel 方案优化掉一大半时间这是省 launch 和访存的红利。第二静态融合版本里还有约四成的 SM 在尾波空转这是 MoE 不均衡路由导致的直接后果也是 Weave 真正下刀的地方。第三Weave 的 2.89× 来自把静态版本里的空闲波次压缩到几乎看不见但它自己也付出了约 8% 的调度开销所以没有达到理论上限。换句话说2.89× 不是凭空变出来的它是把一个原本就浪费着的约 40% 的尾波时间回收了大部分再减掉调度自身的成本。这里面最有价值的是尾巴空转这个优化目标被精确命中了。3.3 为什么不是 3 倍甚至 4 倍调度有成本依赖有下限我见过不少刚接触这套思路的人第一反应是既然空闲 SM 这么多为什么不往 5 倍、8 倍去做答案是调度本身不是免费的。每个任务领取都要经历一次原子操作、一次全局内存读写任务粒度越细这个固定成本占比越高。如果任务粒度细到只有几微秒原子操作和队列访问的排队开销可能比任务执行时间还高那就得不偿失。另外一个限制来自数据依赖。MoE 大内核里并非所有任务都能被提前执行比如结果合并阶段必须等所有专家子任务都算完那段时间里 SM 再闲也只能等着。Weave 能做到的是把能并行执行的阶段内部填满但必须等所有上游完成的硬同步窗口它消不掉。这决定了任何一个真实的 kernel 都有调度收益的天花板2.89× 恰恰说明论文选的那个负载可并行的窗口足够大硬同步窗口足够小。还有一个细节是硬件层面的H100 的 SM 数量多、任务切换快如果换成 SM 数量更少、版本更旧的架构动态调度的延迟差异会直接影响加速比。同样的方案放到 A100 上效果可能就打对折这一点后面我会在落地清单里再强调。3.4 层加速不等于端到端加速这个账必须算清楚层加速 2.89×听起来很诱人但落到推理或训练端到端时它要乘上该层在总耗时里的占比。假设某个 MoE 层占了整个 pipeline 时长的 30%那么理论端到端加速上限可以用公式估算[ \text{端到端加速比} \frac{1}{(1 - r) r / s} ]其中 r 是该层耗时占比0.3s 是该层层加速2.89代入得到[ \frac{1}{0.7 0.3 / 2.89} \approx 1.26 ]也就是说即使单层快了 2.89 倍端到端最多也就快 1.26 倍。更何况多卡之间还有通信、注意力层、采样等其他部分有些阶段还能跟通信重叠实际端到端收益还要进一步看 profiling 结果。这个账我在做优化时经常拿出来算防止团队里出现单 kernel 提速 3 倍整个模型就提速 3 倍的错觉。Weave 的价值在于把它负责的那一层压得很低但它不负责把其他层也变快。4. 从 Weave 论文里能抄到什么工程落地清单论文看完了更关键的问题是这套思路除了在 H100 MoE 的组合里能出彩普通工程师在哪能找到收益、在哪个环节容易翻车我把自己实践中的体会和踩坑经验整理成一份可操作的清单。4.1 第一步永远是量化你的空闲 SM任何 kernel 优化如果没先量化基线就不要动手。对 MoE 大内核我建议用 Nsight 的 GPU Trace 或 CUPTI 先跑一遍 profiling重点看两个指标kernel 尾部的 SM 活跃度曲线以及各 wave 中实际占用 SM 槽位的比例。如果活跃度曲线在最后三分之一明显下滑说明尾巴效应严重动态调度才有上场的必要如果曲线本来就接近一条平线那 Weave 式的优化基本不会有收益。这里我可以给一个粗略的收益上限估算公式假设某个大内核耗时 T平均有 α 的比例的 SM 时间在空转比如 α 0.3那么理想加速上限是 1/(1-α)约 1.43 倍再扣掉调度开销实际能拿到的往往是上限的一半多一点。先把这个账算清楚你就知道自己是不是在做一件收益可能不大、成本却不小的事。4.2 一个最小可复现的自调度示例如果你想在自己的 kernel 里先验证这套思路可以从一个最简单的版本开始。用我前面给过的任务队列循环加上一个 per-block 的小批量领取优化不要每次只领一个任务而是先一次性从全局队列原子领取 N 个任务存在本地缓冲区里再逐个执行。这样能大幅减少全局原子操作竞争这是我从 Weave 工程细节里体会很深的一点。// per-block 批量领取降低全局原子竞争 #define LOCAL_BATCH 4 __device__ int fetch_batch(TaskQueue* queue, int* batch, int batch_size) { int start atomicAdd(queue-head, batch_size); int count 0; for (int i 0; i batch_size; i) { int task start i; if (task queue-num_tasks) batch[count] task; } return count; } __global__ void moe_layer_kernel(TaskQueue queue, float* input, float* output) { int local_tasks[LOCAL_BATCH]; while (true) { int got fetch_batch(queue, local_tasks, LOCAL_BATCH); if (got 0) break; for (int i 0; i got; i) { dispatch_task(local_tasks[i], input, output); } } }这里每一轮fetch_batch只有一次全局原子操作4 个任务分担一次竞争成本实测里的调度开销能压到很理想的水平。批量大小需要调取太大容易出现某些 SM 拿到一堆任务、另一些 SM 已经空手的情况破坏负载均衡取太小又回到原子竞争的老路上。我一般在几十个 block 规模的 kernel 里先用 4 或 8 起步再根据 profiling 收敛。4.3 适合与不适合的场景边界不是所有 kernel 都适合套 Weave 这套细粒度动态调度。我把判断标准整理成一张表场景特征适合程度原因路由动态、token 数量动态、负载极不均匀非常适合静态 shape 规划无法命中负载动态调度红利最大大内核内部子任务粒度可达几十微秒级非常适合调度开销占比低原子竞争压力小任务数据依赖稀疏大部分子任务可以无顺序执行非常适合依赖管理简单同步窗口短计算均匀、wave 完整的小 kernel不适合本来就没有尾巴加了队列反而拖慢子任务粒度过细微秒以下不适合原子开销和穷举队列比干活还贵长依赖链、必须频繁全局同步不适合同步窗口吃掉所有动态调度的红利一句话总结动态调度吃的就是负载不均 子任务粒度适中 依赖可控这碗饭缺一项都会让收益明显缩水。4.4 实践中的四个坑坑一误以为能直接绑定 SM。我之前也试过通过查 block 的 SM id 来做手工调度费了半天劲效果反而不稳定。CUDA 的硬件调度器不可控我们能做的是靠足够细的任务队列哄着硬件把空槽填满而不是直接指挥它。坑二把同步粒度压得太细。开始我为了让数据尽可能新鲜每个子任务前后都查一遍依赖状态结果原子操作数量暴涨调度器本身变成了新瓶颈。正确做法是像 Weave 那样依赖管理只发生在真正有 producer-consumer 关系的任务之间而不是所有任务之间。坑三任务队列放全局内存放得太频繁。每个任务领取都走一次全局内存当 SM 数量多、队列竞争激烈时L2 到全局的延迟会非常扎眼。改进方式是给每个 SM 一个本地缓存队列本地没有任务了再回全局队列抢这是 work stealing 的标准技巧在 GPU 上同样好用。坑四架构换代后不重新 profile。动态调度延迟跟架构强相关我在 H100 上收益明显同样的代码搬到 A100 上因为硬件调度器的行为不同收益缩水了近一半。所以说2.89× 是4×H100 实测里的数字不是在所有 GPU 上都能复现的数字。最后再分享一点个人体会。我自己把 Weave 的三层思路——任务队列加依赖管理加批量领取——搬到一个内部的小规模 MoE kernel 里做实验的时候一开始也走了弯路任务粒度压得特别细结果原子操作排队的开销比省下来的空转时间还多跑出来反而慢了。后来我学乖了先量尾波空转占比再算调度成本把批量大小和任务粒度调到平衡点才看到正向收益。这类调度优化的本质永远都是把空档填起来的收益和填这个空档花的成本之间的博弈。你要做的不是无脑抄 Weave 的代码而是先用 profiler 量好你自己的格子再决定要不要请这套动态调度进场——量完之后你会发现大部分问题其实不在调度算法而在你的 kernel 负载描述和任务切分逻辑上。