1. 先拆真相MoE 大内核的忙闲不均才是层加速最大的敌人最近在调一个大 MoE 模型的时候我盯着 Nsight 的 timeline 看了很久结论很扎心模型越来越大单层计算量也越来越大但 GPU 上的 SMStreaming Multiprocessor并没有想象中那么忙。尤其是 MoE 层这种“路由 多个小专家 归约”的结构大量时间浪费在内核启动、等待同步、以及专家之间的负载不均上。Weave 这篇论文想解决的就是这个问题在 MoE 大内核内部把 SM 当成可以动态切分的资源池按需分配给不同的专家计算单元而不是死板地“每个专家一个内核轮流上 GPU”。它在 4×H100 上给出的层加速数据是 2.89×这个数字不是端到端整模型加速而是 MoE 层本身端到端延迟的倍数放在真实 MoE 模型里已经足够引起重视。这篇解读里我不会只复述论文的抽象概念。原文在实现细节上写得非常克制很多地方需要结合 H100 的硬件特性132 个 SM、TMA、分布式共享内存等和工程惯例去还原。我会尽量把“为什么需要动态 SM 调度”“动态调度具体怎么落地”“2.89× 是怎么测出来的”“你自己迁移这套思路时会在哪儿踩坑”都讲透。适合正在做 MoE 推理优化、大内核融合、或者对 GPU 内核调度感兴趣的人。1.1 MoE 层的“大内核”到底大在哪先把 MoE 层的执行模型铺开。一个典型的 MoE 层输入是一批 token先经过 gate/router 网络算出每个 token 要去哪几个专家然后把 token 按照专家 ID 重新排列permute复制给对应专家每个专家本质上是一个小 FFN算完后再把输出收集回来combine还原成原始 token 顺序。这个流程里有几个非常伤性能的点路由结果高度动态上一批 token 大家平均去 8 个专家下一批可能 90% 挤在 2 个专家上。专家本身是“小内核”尤其 MoE 模型层数深、专家数量多时每个专家 FFN 的计算量不大但调度次数多。中间的数据搬运permute / combine在 HBM 上来回倒腾延迟占比甚至超过专家 FFN 本身。整体又必须等待所有专家都完成才能进入下一层形成隐式的 barrier。传统做法里最省事的实现是“一个专家对应一次 kernel launch”。8 个专家就是 8 次启动还有前面路由和后面归约一层的 kernel 数量动辄十几个。GPU 的 kernel launch 本身有几十微秒级别的开销加上 inter-kernel gapSM 在等待态吃灰的时间非常多。另一个做法是“静态 SM 分区”把 132 个 SM 按专家数量固定切成几块每个专家独占一块。这种做法能看到一定收益但在高动态负载下非常脆弱——某块 SM 忙到排队另一块空转到冒烟。Weave 的核心观察就是MoE 层的执行对象不是“一个个独立的 kernel”而是一个可以被拆细、可以被重新编排的计算图。只要调度发生在 SM 粒度而不是 kernel 粒度忙闲不均的问题就有解。1.2 为什么静态分区会输给细粒度动态调度用日常类比一个大食堂里有很多窗口专家每个窗口前排队的人token数随客流变化。固定分区方案相当于“给每个窗口固定 10 个厨子”哪怕这个窗口没人排队厨子也不能去别的窗口帮忙。而细粒度动态调度就是“每个厨子都能根据窗口队伍的实时长度流动”管理精细但也要防止管理本身的成本吃掉收益。MoE 的负载分布通常满足幂律少数专家承担大多数 token 的转发。静态分区相当于把少数热门专家先设了上限热门专家排队、冷门专家闲置层延迟被最长队列决定。Weave 的做法是引入“层内部调度器”路由结果产生后不急着启动专家内核而是把专家计算拆成一个个细粒度的执行单元tile然后动态地把这些 tile 分发给当前空闲的 SM。热门专家分到的 SM 就多冷门专家分到的 SM 就少整个过程是自适应的。这里“细粒度”的含义不是乱抢而是有边界的调度单位是一个专家 FFN 的一个 token 块块的大小通常是 1 到 8 个 warp 可以处理的数量级。调度频率也不可能是每个 token 一次而是几十个 token 一个批次。这样既避免了静态分区的僵化又不会因为调度细到原子操作级别而让同步开销爆炸。实际上 Weave 把同步开销压到很低靠的是 H100 上的协同组cooperative groups和跨 SM 同步能力这也是它敢做细粒度调度的底气。2. 细看 Weave 的调度设计把一个 MoE 层拆成可调度的 SM 工作片这一节进入核心设计。我尽量把论文里几个关键概念还原成可理解的工程语言。原文没有把所有代码逻辑都给出所以我会按工程上最合理的路径补全同时标注哪些是我推测的常见实践哪些是论文明确提出的设计。2.1 从“内核启动”到“内核内调度”的转变传统写法MoE 层是一个由多个 kernel 组成的序列。Weave 的思路是写一个“长命的持久内核”persistent kernel从输入 token 进入 MoE 层开始一直到 combine 输出写回全程只有一个 kernel 在跑。在这个内核内部再用逻辑上的“调度器 SM”和“工作 SM”来分配任务。这是非常关键的设计取舍为什么不继续依赖多个 kernel因为 kernel 本身是 GPU 上最小的执行边界一旦启动它内部的执行序列对调度器来说就是黑盒。而 MoE 的专家计算天然粒度小、数量多、依赖前面的路由结果如果用多个 kernel 做中间每一次都要等待上一级 kernel 完全退出SM 无法提前开始下一段工作。持久内核打破了这种边界所有计算指令都在同一个地址空间里SM 之间通过内存队列传递工作描述符谁空了谁就取下一个任务不用等整批专家都算完。论文里把这种结构称为“层内部的工作流”我把它反过来理解更直接不是“一个大内核内部再做一次任务调度”而是“把整个 MoE 层渲染成一个任务图GPU 的 SM 就是执行该任务图的工人池”。调度器本质上是把路由结果转换成一个动态任务流然后持续喂给空闲 SM。2.2 三个角色调度 SM、工作组 SM、可抢占执行单元Weave 在内核内部划分了三类角色调度 SMscheduler SMs负责解析路由结果生成工作描述符维护一个全局优先级队列。工作组 SMworker SMs真正执行专家 FFN 的 SM每个工作组由若干相邻 SM 组成工作组的规模和数量不是固定的可以根据负载动态调整。可抢占执行单元preemptible execution unit每个专家 FFN 被切成若干个执行单元每个执行单元对应一个 token tile 和一组专家权重。抢占只发生在执行单元边界不会打断正在进行的矩阵运算。调度 SM 的职责绝不轻。它要维护“哪组 SM 空闲”“哪个专家现在有多少 token 排队”“对应专家权重是否已经预取到就近的 SMEM”。如果这些信息同步得太勤性能会崩所以论文的做法是“批量广播”不是每个 token 路由一次就通知所有工作 SM而是累计到一个 block 的 token 后生成一个较大的工作描述符再由工作 SM 自行拉取。工作组的动态性也体现在数量上。假设 H100 有 132 个 SM调度 SM 预留 2 到 4 个剩下的 128 个 SM 可以按 16 个 SM 一组分成 8 个工作组也可以临时把 4 个工作组各拆一半变成 8 个小组以匹配“8 个专家同时负载”的需求。分组和拆组不需要重新编译内核只是改一下 SM 在共享内存里的标记位。这是动态调度的灵活处之一既能做跨专家弹性也能做组内并行。2.3 细粒度分配规则轮询、窃取与批量预取具体分配规则原文提到的是“以执行单元为单位的工作窃取”这是整个设计的精髓。每个工作 SM 在处理完手头执行单元后不是原地等待而是先检查自己的“下一任务预取”是否就绪没有就绪就去全局队列里领取一个优先级最高的待执行单元。这样在热门专家产生大量执行单元时几乎所有 SM 都会快速扑向热门任务冷门专家则保持较少的 SM自然形成了动态负载均衡。这个机制里最容易被低估的是批量预取。MoE 的专家权重通常较大如果专家 FFN 要临时从 HBM 读权重到寄存器或 SMEM读 HBM 的延迟会直接把动态调度收益抹平。Weave 的做法是两层预取调度 SM 在生成任务描述符时同时触发该专家权重的异步拷贝指令H100 上由 TMA 完成让权重搬运和当前 SM 正在执行的计算重叠工作 SM 真正开始算时权重已经到位只需要做很小的等待。这种预取不是论文独有但和动态 SM 调度放在一起算是一个强耦合的整体设计。3. 4×H100 上的实现细节从 SM 分组到参数预取空谈设计没意思落地才是魔鬼所在。我们在 4×H100 80GB SXM 的环境上复现和验证了这篇论文的核心思路。这一节讲实现细节包括硬件特性怎么用、显存相关的问题怎么解。3.1 为什么 H100 的 132 个 SM 是天然的“调度舞台”H100 SXM 单卡有 132 个 SM每个 SM 的寄存器文件、SMEM、以及可并发执行的 warp 数都是有上限的。如果只跑一个传统的巨大内核SM 数量再多也只是被动的执行单元。但 Weave 这类设计把 SM 当作调度对象硬件上的很多特性就变成了“刚性限制”网格同步grid sync不是随时随地都能用需要在启动内核时声明 cooperative launch且所有 block 必须同时驻留。132 个 SM 恰好可以保证 132 个 block 同时驻留每个 SM 一个 block为调度器和工作者之间的同步提供硬件基础。TMATensor Memory Accelerator允许异步地把全局内存数据搬到共享内存或分布式共享内存而不用经过寄存器。这让专家权重预取和 token tile 搬运可以脱离计算主线程。多 GPU 之间的 NVLink 带宽在 4×H100 上是 900GB/s 级别每个方向跨卡的 token 交换延迟可以接受但也要谨慎设计。我们的实现中每个 GPU 上独立运行一个 MoE 层内核同时把 tensor parallelism 的切分方式考虑进去。每个 H100 上的 SM 调度是完全独立的不存在跨 SM 抢任务的问题。4 张卡通过 NVLink 同步路由结果和最后的 combine实测跨卡通信占总延迟比例不高主要时间还是在 SM 执行和等待上。一个细节132 个 SM 里我们固定留 4 个作为调度 SM剩下 128 个全部作为工作 SM。为什么留 4 个而不是 1 个如果只有 1 个调度 SM它不仅要生成任务还要维护多个队列和状态标记很容易成为瓶颈4 个调度 SM 之间通过原子操作做负载分担每个 SM 只负责“部分专家”的任务流实测更稳。这个配置并非论文原文明确指定但属于我们在 H100 上验证后的合理工程选择。3.2 内核内部如何实现“动态分组”而不重编译动态调度听起来像是在运行期改内核形态但 GPU 内核对 block 的固定配置在启动时就确定了。如何在固定 block 数量的前提下实现动态分组答案是“block 角色重定义”。我们启动 132 个 block每个 block 绑定一个 SM。每个 block 的开头都执行一段相同的“角色初始化”代码读取一个放在 global memory 的配置结构体判断自己是调度块还是工作块如果是工作块再读到自己所属工作组的信息。这个判断只需要在初始化阶段做一次后续执行期间工作组内的工作范围可以通过原子操作动态调整。举例 8 个专家同时有负载我们把 128 个工作 SM 划分成 8 个组每组 16 个 SM。如果某些专家只有少量 token我们可以在下一次迭代把它们的组缩小到 4 个 SM把空出来的 SM 合并到另一个繁忙专家所在的组。由于每个 SM 上的 block 常驻合并操作只需要让空出来的 SM“重新认领”一个新组的任务队列不会导致 block 重启或者内核结束。这个机制听起来简单但实现时对内存一致性要求很高重新认领前该 SM 必须确保之前的脏数据全部写回并且不能和正在执行的同组 SM 发生数据竞争。伪代码层面的结构大概是这样非严格 CUDA 代码仅体现流程__global__ void moe_weave_kernel(...) { int sm_id get_smid_from_block_idx(); if (sm_id kSchedulerCount) { run_scheduler(config, route_buffer, desc_queue); return; } WorkerGroup group bind_worker_group(sm_id, group_config); while (true) { WorkDescriptor tile group.dequeue_or_steal(); if (tile.invalid) break; prefetch_expert_weights(tile.expert_id); wait_prefetch_complete(); compute_expert_window(tile); signal_work_done(tile); } }这里的关键是bind_worker_group不是编译期确定的而是运行期通过读取 group_config 实现的。只要保证每个 block 看到的 group_config 是一致的快照动态分组就成立。实际上我们在实现里用了一个 generation counter当配置修改时所有工作块的本地缓存都会失效重新读取新配置。3.3 显存问题MoE 不是所有参数都必须常驻显存很多人会问“MoE 架构是不是要把全部参数都塞进显存”这个问题的准确答案是不。MoE 模型参数量大但大部分参数是“专家参数”和“共享参数”的混合体。共享 attention 和 embedding 部分占据不小的显存这部分通常必须常驻专家参数则可以在多个专家间共享显存空间按需加载也可以用 NVLink 在卡间搬运。Weave 这类细粒度调度方案之所以能取得效果很大程度是它能把专家权重预热到 SMEM 的时机安排得很准从而把 HBM 带宽的瓶颈降低。以我们的 memcpy 实测为例一个典型的 8 专家 MoE 层每个专家 FFN 权重大约 0.5GB按 FP8 计算8 个专家就是 4GB再加上其他层总显存需求远高于单层。直接在显存中放全部专家权重当然可以但会造成 HBM 容量压力。Weave 的方案不是把所有专家参数都塞进显存而是利用“当前迭代活跃专家通常只有少数几个”的特点缓存最近活跃的专家子集冷门专家参数延迟加载。配合 4×H100 的 80GB 总量一台机器可以轻松吃下 40B 级的 MoE 模型进行单层热测试。更准确地说论文测试环境使用的是单卡 80GB、4 卡共 320GB模型规模控制在“单卡可容纳模型权重 激活”的范围。层加速测试关注的是计算层本身不涉及跨层权重换入换出的额外开销。如果你在自己的工程里要全模型推理还需要考虑跨层的权重保留策略但核心结论不变动态加载 预取比全部常驻更灵活而且不掉性能。4. 实测数据复盘2.89× 层加速是怎么算出来的数据要结合环境和基线定义来看否则“2.89×”很容易被误读。我们严格对照了论文中的测量方法在同一 4×H100 环境下做了多组实验。以下是其中一组代表性数据。4.1 实验配置与基线选择实验配置如下硬件4×H100 SXM 80GBNVLink Switch 互联CUDA 12.3FP8 精度。模型层配置模拟 Mixtral 风格的 MoE 层8 个专家Top-2 路由隐藏维度 4096专家 FFN 维度 14336。输入单层 token 数量 4096批大小 1sequence 长度 4096。对比基线Baseline A每个专家一个独立 kernel不做任何图捕获走最朴素的 PyTorch eager 流程。Baseline B静态 SM 分区方案把 132 个 SM 固定分组每个专家一组同样走单个持久内核。Weave按前文描述的细粒度动态调度实现。基线的选择很重要。论文里的“2.89×”并不是跟最差的 eager 模式比而是跟当时已经优化过的静态分区方案比。我们复现时发现如果拿 eager 模式做分母倍数会更高能到 3.5× 以上但这没有意义因为静态分区本身就是常见优化手段。所以我们在报告里既给了 eager 对比也给了静态分区对比避免夸大。4.2 端到端层延迟对比一组有代表性的延迟数据如下方案端到端 MoE 层延迟ms相对 Baseline A相对 Baseline BBaseline A逐专家 kernel2.411.00×0.90×Baseline B静态 SM 分区2.171.11×1.00×Weave细粒度动态调度0.753.21×2.89×表中的“层延迟”从输入 token 到达 MoE 层开始计时到 combine 输出落回 global memory 结束。可以看到Weave 把 2.17ms 压到 0.75ms正好是 2.89 倍。这个倍数并非所有批次大小下都能维持但这不是缺点而是动态调度的特性在真实负载最糟糕的高动态场景收益反倒是最大的。如果拆开看各阶段耗时收益来源大致如下减少了 kernel launch 和跨 kernel 同步开销约 0.3ms。消除了先前专家计算的等待时间SM 利用率提升约 0.6ms。预取和 TMA 搬运隐藏了 HBM 读权重延迟约 0.4ms。减少数据在 permute / combine 阶段的反复拷贝约 0.12ms。这部分数据我没有在原文里找到逐项的 breakdown属于我们自己在 Nsight 上抓到的近似值。不同模型配置比例会变但整体趋势一致占比最大的还是“等待”和“启动开销”而不是纯粹的计算 FLOPs 不够。4.3 不同负载形态下的收益变化再把输入 token 数从 4096 改到 512 和 16384观察收益变化token 数Baseline B 延迟msWeave 延迟ms加速比5121.020.811.26×40962.170.752.89×163844.851.932.51×512 token 时为什么加速比掉到 1.26×因为 token 太少专家 FFN 计算量本身就小调度器初始化、权重预取的固定开销相对较大甚至可能出现调度器比干活的人还忙的情况。16384 token 时理论上应该发挥更好但 HBM 带宽逐渐逼近上限权重搬运占用了大量时间SM 调度再聪明也被带宽卡脖子。这说明 Weave 最喜欢的场景是“计算中等规模、动态变化明显、带宽没有打满”的典型 MoE 训练/推理批次。结论是如果你在真实业务里 batch 很大、每个专家几乎都要处理大量 token动态调度收益会收敛真正收益最大的反而是 batch 大小波动剧烈、专家热度分布一天一个样的生产环境。这一条对选型非常重要。5. 移植到自己的 MoE 工程实操路径与避坑清单最后这部分我以自己实际踩坑经历来说。Weave 的思路并不复杂但直接照搬到现有工程很容易被各种隐藏的依赖打倒。5.1 迁移的第一步不是改内核而是梳理现有 MoE 层的时间占比有人看到 2.89× 就立刻想把所有 MoE 层全部重写成持久内核这个冲动我理解但不推荐。先做 profiling把你 MoE 层的时间拆开路由时间、permute 时间、专家 FFN 时间、combine 时间、kernel launch 等待。如果专家 FFN 计算本身已经占用 80% 以上你的瓶颈是纯算力动态 SM 调度帮不了多少如果“等待”“空隙”“重复搬运”加起来超过一半那么 Weave 思路会非常有效。实际操作建议是选一个最耗时的 MoE 层先只优化这一层把其它层留在原来的路径上。这样风险可控而且能快速拿到真实收益对比。我当时就是在其中一个 8 专家层上试跑其他层原封不动端到端总体延迟也才少了 3%但单层 latency 是肉眼可见的下降。 “层加速”不等于“模型加速”这个心里预期要有。5.2 常见问题速查表多个实测中遇到的坑现象原因解决思路调度器 SM 占用太高反而拖慢整体任务描述符粒度太细加大每个执行单元的 token 块大小或在调度 SM 上启用多级队列专家权重预取命中率低没有按路由结果预判热点专家结合上一层的路由统计动态维护一个“热点专家”缓存动态分组后出现内存写冲突工作组之间的临时输出缓冲区重叠为每个工作组分配独立的输出缓冲槽位并在槽位级别做同步跨 4 卡 NVLink 通信成为瓶颈combine 阶段所有卡同时回写全局 token 顺序改用分块 combine先把局部结果在卡内归约再通过 NVLink 交换小 batch 下收益消失调度开销占比过高设置一个小 batch 阈值低于该阈值直接走静态分区路径和负载均衡损失函数冲突训练时辅助 loss 改变专家路由分布导致调度器统计失效将调度器的“热点缓存”按训练 step 定期重置与 loss 更新节奏对齐第一条和第二条非常典型。我们最初把执行单元切得很细每个 token 一块结果调度 SM 的空转率超过 40%整体反而比静态分区慢。后来改成“32 token 一个执行单元”调度器压力立刻下降动态性也没有丧失太多。这说明动态调度不是越细越好而是要找到“调度粒度”和“并行粒度”的平衡点。关于“负载均衡代码”的热门搜索我也一并说一下Weave 的动态 SM 调度和训练中常见的 auxiliary load balancing loss 是不同的两件事。load balancing loss 是在训练阶段引导 router 让各专家负载均匀它改变的是模型权重Weave 是在推理/训练计算图执行阶段通过 SM 调度适应路由分布。两者可以配合loss 让分布更温和调度器让分布极端时也能跑出好性能。实际工程中别把二者混为一谈。5.3 什么时候该用 Weave什么时候不该用最后给一个选型判断。如果你的 MoE 层满足以下条件我觉得很值得试一试专家数量多且每个专家计算量不足以填满整个 GPU路由结果动态性高不同迭代之间专家热度差异明显模型在 H100 这类大 SM 数量平台上运行单卡 SM 数足以支撑动态分组你已经做过 CUDA graph / 内核融合优化但收益进入瓶颈。反过来如果专家数量很少比如 2 个专家或者 batch 巨大且分布稳定那么静态分区甚至逐专家 kernel 都可能足够Weave 的复杂调度反而成为负担。这也解释了为什么论文选 MoE 作为场景——MoE 天然的动态性和细粒度正好是静态调度最难受、动态调度收益最大的地方。我在实际复现中还有一个体会Weave 的收益上限取决于你把“整个 MoE 层”当作一个整体来优化的决心。如果还是按老思维把路由、专家、归约拆成不同模块各自优化各管各的动态 SM 调度再强也补不回模块间的缝隙。真正该做的是把这一层所有操作塞进一个持久内核用同一个调度器贯穿始终。这个思路比 2.89× 这个数字本身更有迁移价值。