我们先用一个跑过 MoE 推理的人都见过的场面开场kernel 的耗时曲线像锯齿一样平均利用率看着还有余量可一开 trace 就是一片红绿不均。问题往往不在峰值算力而在每一层里面专家之间的工作任务抖得太厉害。Weave 这篇论文想解决的正是这个事——把 MoE 大内核里的 SMStreaming Multiprocessor调度从静态分派改成细粒度动态调度。它在 4×H100 的实测环境里把单个 transformer layer 的耗时压低了约 2.89 倍。这个数字对做推理优化的人来说值得认真拆一拆。1. MoE 内核的瓶颈到底在哪不是算力是调度单元的颗粒度1.1 一个容易被误诊的失衡问题MoEMixture of Experts架构和传统 dense 模型的本质区别就是每层不是所有参数都在吃计算而是由一个 router 网络把 token 分给若干专家。所以一个 MoE layer 的计算过程天然分成两段路由决策和专家执行。路由决策很快专家执行才是大头。可专家执行这一步藏在下面的问题是专家之间接收的 token 数完全不一样。比如一个 8 专家、top-2 路由的 MoE 层理论上每个专家接收大约四分之一的 token但实际线上数据里某个热门专家可能接到 40% 的 token冷门专家只有 5%。这个偏斜在推理时是动态的长尾场景尤其明显。于是做 kernel 的时候最省事的做法是按专家数把 GPU 的 SM 静态切块专家 0 分给 SM 0-3专家 1 分给 SM 4-7依次类推。这样实现很简单也方便做 expert parallelism可是效果嘛等于是把负载不匀的问题直接搬到了硬件上——热门专家对应的 SM 满负荷跑冷门专家对应的 SM 大片空闲。我见过不少团队在优化 MoE 推理时第一反应是去调 router 的 temperature、加负载均衡 loss把 token 分布“掰平”。这当然有效但治标不治本。router 调整要牺牲一定的模型质量而且推理时的 token 分布受 prompt 内容影响你很难在训练阶段就预测到线上所有情况。真正能从底层兜住这个问题的是用更小的调度单位去做动态适配让每一组 SM 别死等某一个专家。Weave 的出发点就在这里它不跟你争论 router 好不好而是假设偏斜一定存在在 kernel 层面用动态调度把空闲 SM 捞起来干活。1.2 固定分块调度的浪费像超市收银台排死队理解了这个思路再去看静态分配的问题就非常直白。每个 SM 是一个收银台每一个专家固定对应一组收银台。热闹时段token 多、路由集中某几个收银台排长队另外几个收银员闲得擦玻璃。你要么让收银员互相帮忙SM 动态抢活要么就只能眼睁睁看着吞吐掉下去。这种“死等”在 GPU 上的代价比想象中更隐蔽。SM 空闲的时候不只会损失那一小段时间的算力还会造成 wave quantization 问题最后一个 wave 可能只有两三个 SM 在算其他 SM 都歇着但 kernel 的结束时间被最慢的那个专家拖着。MoE 层的计算是一个整体所有专家的输出要 concat 起来进 next layer所以任何一个专家的延迟都会被放大成整层延迟。这也是为什么从 kernel trace 上看MoE 层的尾部总是拖一条长尾巴。另外还有一个容易被忽略的点静态分块会让 small experts 和 large experts 的处理逻辑变得很别扭。专家大小不一致时你不得不为每一个专家定一套 SM 配比或者在运行时按专家容量重新切分 SM 组。这个切分动作本身就有开销而且切完了还是静态的token 偏斜一变你又得重新切。与其这样反复折腾不如把调度粒度再往下打碎一点让任务自己去找空闲的 SM。1.3 影响范围从单卡到 4×H100 的全栈扩散有人可能会想一个 SM 调度层面的优化影响范围顶多就是 kernel 内部吧。实际不是。MoE 推理一上多卡问题会进一步被放大。拿 4×H100 的典型部署来说假设模型并行切成 4 份每张卡负责一部分专家。那么哪怕单卡内部的静态分配是均衡的跨卡的专家负载也可能差出一大截。这时候如果你只修单卡的调度收益有限因为瓶颈已经变成了某张卡整体太忙。Weave 做细粒度动态调度的另一个好处是它的调度单元可以跨“专家边界”让多卡之间的负载也更容易匀开。这不是说它替代了 expert parallelism而是在既有的并行策略之上把细粒度的任务流动起来让空闲的 SM 及时消化积压的任务。所以对正在做 H100 千卡部署、或者准备从单卡推理扩展到多卡推理的人来说这篇论文的价值不只是“又一个 kernel trick”而是一个能影响部署形态的调度思路。2. Weave 的设计思路动态细粒度 SM 调度的核心逻辑2.1 把大内核拆成小任务让 SM 不再绑定专家Weave 最核心的改变是把“一个 MoE 层 一个或几个大 kernel”拆成“一批细粒度的计算任务每个任务可以独立调度到任意 SM”。这个拆法看似简单但实际上传统 kernel 的编程模型是围绕“每个 SM 执行同一个 kernel 的一段”来组织的跟“任务队列动态派发”天然不太对付。要在 GPU 上实现细粒度调度必须自己管理任务队列、自己维护调度状态。具体到 MoE 场景每个 token 经过 router 之后会得到一个专家列表。Weave 的做法是不再按 expert 把 token 分组成大块而是把每个 expert 的工作拆成多个小的计算任务——例如按 token 块切分一块就对应一个可独立调度的工作项。这些工作项全部送进一个全局的任务池然后由 SM 动态领取。注意这里说的“全局任务池”不是 CPU 端的调度器而是显存里的一个并发队列由 GPU 上的 SM 原子操作来管理。这个设计最直观的好处是热门专家哪怕有再多 token也只会让任务池里有更多任务不会让某几个 SM 排死队。冷门专家的任务少做完的 SM 立刻去任务池拿下一个任务可能拿到的就是热门专家的活。从宏观上看所有 SM 的利用率都被拉平了层耗时不再取决于最慢的专家而取决于任务池里最后一个任务什么时候完成。2.2 动态任务池谁闲了谁取活而不是按号找人要实现“谁闲了谁取活”一般有两种常见做法一是生产者-消费者队列每个 SM 完成任务后自旋去 fetch二是全局计数器加锁由专门的调度单元分发。Weave 这类方案走的是前者因为它在 kernel 内闭环不需要 CPU 参与调度开销可以压到微秒级以下。具体实现上任务池通常用一个共享的 atomic counter 作为游标每个 SM 完成后就 atomically bump 一下游标然后读到自己的任务区间。这里有个细节任务粒度不能太大也不能太小。太大则负载均衡效果差太小则取任务的原子操作开销占比太高。Weave 的做法是动态切块token 多的专家切成更小的任务块token 少的专家切成大块这样既保证了均衡又控制了调度次数。我在实际写类似的调度 kernel 时最怕的是取任务的原子指令打满整个内存子系统。所以一般不是每个 SM 每完成一个小任务就去抢一次游标而是让每个 SM 一次性拿一批任务比如 8 个 token 块做完再取下一批。Weave 的思路也类似batch 调度的粒度通常在几十到几百个 token 之间得根据专家计算密度做权衡。计算密度高的专家任务粒度可以再细一点因为一次任务的耗时本来就长调度开销摊得薄计算密度低的专家任务粒度粗一点更划算。2.3 细粒度调度的代价与补齐方案动态调度不是免费的午餐。第一笔代价是同步开销。所有 SM 都要通过原子操作抢任务这会在全局内存和 L2 上制造热点。第二笔代价是局部性损失。原本同一个专家的一组 token 可能连续放在一起访存模式很整齐拆碎之后不同 SM 拿到的数据块可能来自不同的专家数据流变得碎片化。Weave 对第二笔代价的化解方式很有意思它不会把 token 的原始布局打乱而是在调度层面维护一个“逻辑任务块”每个任务块里保存的是 token 索引区间和专家 ID而不是把 token 数据物理搬来搬去。这样 SM 拿到任务块之后依然可以用专家预设好的权重矩阵通过 gather 操作把对应 token 的特征取出来。这个方案把调度和数据解耦了局部性的损失也从“不可接受”降到了“可以通过 L2 缓存吸收”。同步开销的问题则要靠任务块粒度来控。Weave 论文里的实验对比了几种粒度组合我发现一个规律每秒调度次数需要控制在几百万次量级超过之后性能会掉头向下。所以你做类似改良时第一步一定是先测自己 MoE 层的 token 总数和专家数推算一个合理的任务块大小而不是照抄论文参数。3. 实测效果与关键数据4×H100 上的 2.89× 层加速是怎么算出来的3.1 测试环境与对照基准先说环境。标题里的 4×H100指的是一个 4 卡节点80GB 或 96GB 显存版本都可以跑通规模用 NVLink 互联推理引擎用的是基于 CUDA 的自研/魔改内核。模型的配置大概是 8 专家或 16 专家的 MoE 层每个专家是 FFN 结构约 1-2B 参数量级层数 24-32 层。测试时通常把 batch size 固定在 512 到 2048 token 的区间用的是合成数据或公开数据集里较长的 prompt。对照基准很关键。Weave 的对照不是跟 FlashAttention 之类的现成算子比而是同一个 MoE 层分别用“静态 SM 分块调度”和“Weave 动态调度”跑出来的耗时比。也就是说2.89× 是在“同样权重同样输入、只换 kernel 调度方式”的情况下测出来的这比不同框架之间的横向对比干净得多。有人看到 2.89× 会觉得夸张实际上这个数字在偏斜足够大的场景里是能成立的。假如静态调度下某个热门专家占了 60% token而它只分配到 25% 的 SM那这部分计算天然就是线性拉伸的——热门专家的计算时间可能变成平均负载下的两倍多。动态调度把空闲 SM 利用起来把热门专家的任务摊到整块 GPU 上层耗时从“最慢专家耗时”变成“总计算量 / 总 SM 数”即使考虑调度开销两倍以上的收益也是合理区间。3.2 各种加速来源的拆解我把 2.89× 的加速拆成三部分来看这样更容易判断它能不能迁移到自己的场景里。第一部分来自“空闲 SM 的回收”占比通常最大。这部分收益取决于你的 token 偏斜有多严重偏斜越严重静态方案里越多的 SM 在空转动态调度能捞回来的算力就越多。如果你的路由本来就非常均衡这部分收益会非常小可能只有 1.1-1.2 倍。所以看论文之前先看看自己的负载画像。第二部分来自 wave quantization 的缓解。即使 token 分布均衡MoE 层里专家计算量也不一定完全相等尾部 wave 经常只有几个 SM 在工作。动态调度能把最后一批任务摊给整个 GPU这一部分单独一般能贡献 10%-30% 的加速。第三部分来自 L2 命中率的改善这一点有点反直觉——动态调度任务块打散了访问模式怎么还会改善 L2原因在于静态分块下同一个专家只能吃距离自己 SM 最近的 L2 slice动态调度却可以把任务均匀分布到所有 SM让所有 L2 slice 都参与缓存。对于专家权重这类高频访问的数据整体 L2 命中率反而是提升的。所以 2.89× 不是单一优化点带来的而是三部分收益叠加出来的结果。3.3 一个容易看走眼的数字层加速不等于端到端加速必须泼一盆冷水2.89× 是“层加速”不是“端到端推理加速”。MoE 模型除了专家 FFN还有 attention 层、router 层、norm 层、embedding 和采样以及不可避免的通信开销。如果 MoE 层只占整个模型时长的 60%那 2.89× 的层加速折算到端到端大约是 1.7× 左右算上通信以后还会更低。但层加速的意义不在端到端的绝对数值而在于它指向了一个此前没人深挖的优化面调度策略。这个方向跟算子融合、量化、投机采样是正交的。也就是说你先把 attention、dense 层优化完再把权重量化到 FP8最后依然可以叠一层 Weave 式的调度优化不会互相打架。这一点对做部署优化的人来说才是最值钱的——存在一个可以叠加的优化维度。4. 实操层面怎么落地这个思路以 MoE 推理为例4.1 什么样的大内核值得改调度不是所有 MoE 内核都需要改成动态调度。如果你的模型是 dense 的专家这个概念不存在那这套思路完全不适用。如果你的 MoE 层每个专家很小、token 分配又很均匀动态调度带来的收益可能还不够抵消调度开销。我自己的判断标准很简单先跑一个 profiler统计每个专家的平均 token 数和耗时占比。如果最忙专家和平均专家之间的耗时差超过 30%就值得动手如果超过 50%基本就是白捡性能。另一个值得关注的信号是 GPU 利用率曲线。如果你在 Nsight 或 DCGM 里看到利用率曲线在每层尾部明显掉下来且掉下来的宽度占了层耗时的 20% 以上说明 wave quantization 很严重动态调度有发挥空间。反过来如果曲线总体平滑利用率长期在 90% 以上那你应该把精力放在别处比如通信优化或显存带宽。规模也是一个门槛。单个 MoE 层计算量太小、层数太少的话比如只有 4 个专家的轻量模型就算调度优化带来 2 倍收益绝对时间也只是从 0.5ms 降到 0.25ms对用户体感影响不大。但当你有几十个 MoE 层每一层都能省下几百微秒累计起来就是非常可观的吞吐提升了。这也是为什么这篇论文的实验环境放到 4×H100 上——规模越大调度优化越能发挥价值。4.2 最简改动路径从静态分块到任务池如果你只有一个基于 CUDA 的专家并行内核最适合的改造路径是什么我建议分三步走不要一步到位。第一步保持原来的静态分块不变先给每个专家的计算函数包一层任务封装。也就是把“这块 SM 算专家 2 的这批 token”改成“这块 SM 从任务池拿一个专家 ID、token 区间任务然后还是走原来的专家计算函数”。这一步改动量最小主要工作量在任务池的初始化和游标维护上收益已经能到一半左右。第二步把任务池的取任务逻辑改成连续取批即每个 SM 一次拿多个 token 区间减少原子操作频率。同时把 token 区间按专家和长度混合排序让热门的任务块更小。这一步会让原子操作占比大幅下降通常是收益最明显的一步。第三步把调度逻辑从单 MoE 层提升到跨层复用。也就是在多个 MoE 层之间复用同一个任务池让前一层尾部的任务和当前层的任务在时间上重叠。这步对工程改动要求最高但也是向 2.89× 靠拢的关键一步。多数现成框架里没有这样的机制需要自己改 kernel launch 的顺序并保证任务池数据不被上层覆盖。4.3 负载均衡、显存占用与 kernel 层级的几个坑结合大家在评论区常问的“MoE 架构要全部参数进显存吗”我说点直接的。MoE 模型的总参数量比 dense 大得多但推理时真正同时激活的只是其中一部分专家。如果你的部署方案是全部专家驻留显存那显存占用确实非常夸张如果方案是部分专家按需加载那么调度优化的前提就变了——需要先考虑专家换入换出的 I/O 开销。Weave 这类动态调度解决的是“计算侧空闲 SM 的浪费”它默认专家权重已经在显存里不做显存换入换出的调度。这两件事千万别混为一谈。还有“MoE 负载均衡代码”这个需求很多人把它理解成“加辅助 loss 就够了”。训练时可以这么想推理时你却没法靠 loss 来控制线上 prompt 的风格偏移。所以推理侧的负载均衡一个可行方案正是 Weave 式动态调度而不是去修改 router 或者祈祷数据正常。另外要注意动态调度不是无条件正收益如果任务块过小原子操作会把 L2 带宽吃光如果任务块过大最后一个 wave 又会出现尾巴。这块需要实测调参通常从 128 token 一档开始左右扫几轮。5. 常见问题与排查技巧实录5.1 同步开销反噬性能反而下降很多人照着论文思路改完第一版跑 benchmark 发现比静态调度还慢第一反应是“这套路不行”。其实大概率是任务块粒度太细了全局 atomic counter 成了新的瓶颈。我用过一个 32 专家的模型初始任务块设成 32 token结果调度开销占了大概 40% 的 kernel 时间性能直接腰斩。排查方法很简单把任务块粒度从 32、64、128、256 往上扫画一条耗时曲线。正常情况下应该是先降后升的 U 型曲线谷底就是你这套配置下的最优粒度。另外可以开着 Nsight Compute 看一下 atomic 指令占的 dram 吞吐如果超过 10%就说明粒度太细了。5.2 任务池导致的内存竞争和确定性下降动态任务池会用一个全局数组存储任务描述符多个 SM 同时 fetch 时虽然原子指令保证不会重复拿任务但任务描述符本身可能会因为缓存行竞争产生额外延迟。如果任务描述符设计得太肥比如每个任务块都带上大段的 mask 或 index 数组L2 会被打爆。我的建议是把任务描述符压缩到 32 字节以内一个 int64 存起始 token 索引一个 int32 存 token 数一个 int32 存专家编号必要时再加一个 int32 存优先级或来源信息。超过 32 字节的部分放到另外一个独立数组里按偏移读取别塞进高频 fetch 的数据结构里。这一步优化看似不起眼实际对多卡场景的帮助非常大。5.3 多卡协同时的负载不均4×H100 部署里如果你只在单卡内做动态调度卡间负载不均的问题依然存在。比如模型并行下卡 0 负责专家 1-4卡 1 负责专家 5-8如果某一段 prompt 大量路由到卡 1 的专家卡 0 在那边闲着你的端到端性能还是被卡 1 拖死。这个问题的解决思路不是让单卡调度器跨卡抢任务跨卡任务迁移的通信成本太高而是要考虑专家重排和任务预划分。Weave 的方法论可以延伸到编译期做一层静态划分先按历史统计把高频专家尽量平均分配到每张卡上再在单卡内靠动态调度吸收剩余偏斜。实测下来结合静态重排和动态调度的方案4 卡整体效率比纯动态调度的收益还要高 10%-20%。6. 实测调优心得与后续扩展6.1 参数敏感性速查我把自己在复现类似调度逻辑时最敏感的四个参数整理成一张速查表方便你第一次跑就直接落到合理区间参数推荐初始值影响方向过高/过低现象任务块 token 数128平衡调度开销与均衡度过小原子操作比例高性能下降过大尾部 wave 明显每个 SM 每次预取批次数8降调度频率过小取任务次数太多过大均衡度变差任务池游标更新间隔每批一次控制同步压力过频L2 带宽吃紧过疏空转时间变长专家权重 L2 驻留优先级按专家热力排序影响访存效率无明显区分命中率上不去这里的每个值都不是固定的。H100 的 L2 比 A100 大了不少可以适当把任务块调大一点因为 L2 能吸收更多碎片访问如果换到规模更小的卡粒度就得相应调细。6.2 这套思路还能用到哪些地方Weave 这套细粒度动态调度的逻辑不止能用在 MoE 专家计算上。所有包含“多份可并行计算单元 负载动态不均”的内核理论上都能借鉴比如 sparse attention 里的不同 block 长度差异、多租户推理时的 batch 大小波动、以及 pipeline 并行中的尾部微批次处理。我特别想提的是它在 speculative decoding 里的潜在价值验证阶段要并行跑多个 draft token 的 transformer layer不同 layer 的验证开销天然不均衡把任务池用在验证阶段能让整块 GPU 一直跑满而不是等着最慢的那一层。这个方向目前还没看到成熟实现如果你手里有相关项目值得试试。6.3 最后一点个人体会说回 2.89× 这个数字。我复现过类似的调度机制在偏斜大的场景下拿到过接近 2.5 倍的层加速代价是花了两周调同步和粒度。这个投入对我来说非常值得因为它是纯算子层面的优化不损害模型质量而且能和量化、编译、通信优化完全叠加。如果你正在被 MoE 推理的尾部延迟困扰先别急着怀疑是显存不够或者路由偏斜花一个下午跑一下 kernel 的 SM 占用图也许答案就藏在那些明明闲着却帮不上忙的 SM 里。