ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Weave:MoE大内核的动态SM调度,2.89×层加速如何实现?

Weave:MoE大内核的动态SM调度,2.89×层加速如何实现? 如果你盯过 Nsight 里的 SM 占用率曲线大概率见过这个让人血压升高的场景MoE 模型跑起来之后路由的负载均衡 loss 已经训得很平了各专家收到的 token 数看起来也差不多可 GPU 里就是有一片 SM 在摸鱼。你加专家并行、调 capacity factor、换 all-to-all 通信实现折腾一圈下来发现最缺的不是专家分了多均匀而是就算分发到了 GPUSM 也没把这活排满。这篇论文解读想说的就是这件事。标题里的 Weave做的就是 MoE 大内核内部的细粒度动态 SM 调度官方实验环境是 4×H100在 MoE 层上拿到了 2.89× 的层加速。下面我会按问题链路 → 切分思想 → 核心同步机制 → 实验数据拆解 → 工程落地 → 适用边界的顺序把我读这篇论文时最在意的几个点全部展开。1. 先看问题MoE 层里的负载均衡到底卡在哪一层1.1 MoE 层执行的完整链路远不止路由分均匀一个 MoE transformer 层在执行推理时大致是这条链路token 进入 router算出每个 token 对每个专家的分值取 top-k把 token 映射到选中的专家做 dispatch也就是把 token 的激活值搬运到对应专家所在的位置每个专家跑自己的 FFN 计算通常就是两个 or 三个 MatMulgate、up、down或者 fused 版本combine把多个专家结果加权求和输出给下一层。很多人讲到 MoE 负载均衡时注意力都放在第 1 步的 router loss 上。也就是让每个专家在平均意义上收到的 token 数差不多。但 GPU 上的真实开销在第 4 步那个 FFN 才是整个 MoE 层里最重的大内核。路由分发只决定了 token 去了哪个专家它完全决定不了这些 token 到了 expert 之后GPU 内部的 SM 怎么排班。这就是标题里大内核的含义MoE 层的 FFN 计算会融合成一个巨大的 persistent kernel几十个专家在这个 kernel 里连续执行所有专家共享同一批 GPU 资源。在这个大内核里任务之间天然存在忙闲不均而 GPU 硬件调度器本身又不会聪明到去重新安排这些任务。1.2 三种不均衡路由不均衡、token 量不均衡、SM 调度不均衡我自己的习惯是把它拆成三个层级来看因为这三层的解决手段完全不同。第一层是路由不均衡。训练时加的 load balancing loss、DeepSeek 那类 bias 项都是为了压制这一层的问题。但训练压平不代表推理时均衡因为推理 batch 是动态到达的请求长短不一瞬时分布抖得厉害。第二层是 token 量不均衡。就算路由总分均匀decode 阶段每个请求的 token 形态是完全动态的。一个用户的请求长度可能是 1024另一个可能只有 8。按专家的 token 数去切块块与块之间的计算密度差异很大。第三层是 SM 调度不均衡这是最容易被忽略、也是 Weave 主要解决的问题。你有一个 132 SM 的 H100专家计算被切成很多个 block交给 GPU 调度器分配到 SM 上执行。block 不是无限小的每个 block 固定占据一个 SM 直到算完。如果最后只剩 3 个 block 在跑那 129 个 SM 全在空转。这个现象业内一般叫尾部效应或者 wave quantization。1.3 为什么静态分块解决不了 SM 空闲一个很常见的反面教材是我按 token 数把专家任务切成 64 token 一块想着只要任务粒度够细SM 总能被填满。但问题在于任务切得太细计算效率会掉。具体算一下。假设一个专家收到 1000 个 tokenFFN 的中间维度是 14336。如果我把任务切成每块 64 token那一个专家就有 1000 / 64 ≈ 16 个矩阵乘块每块的 M64、Khidden_dim、N14336。M64 对小 GEMM 来说不算差但如果这批任务里还有大量 M8、M16 的块那就是另一回事了——decode 阶段经常如此。再进一步如果任务粒度只有 8 token单块计算量大概在几十微秒级这部分时间可能还没 kernel launch 或者调度开销大。静态切分唯一的调节旋钮是块大小块大了尾部效应严重块小了计算效率和 L2 局部性崩盘。你永远找不到一个既能让 SM 排满、又不牺牲 GEMM 效率的固定块大小因为负载本身是动态的。Weave 的出现就是冲着这个静态切块调参调到怀疑人生的困境来的。2. Weave 的核心思路把大内核变成可以被偷走的小任务2.1 从一个 block 负责一个专家到一个 block 消费任意任务传统 MoE kernel 的逻辑是一个 block 绑定一个计算区间你负责这个专家、这批 token你从头算到尾算完退出或者继续处理下一下预先分配给你的活。这种模式的隐含假设是所有计算量可以被预先估计好而且各 SM 的工作量天然一致。Weave 把这个假设直接掀了。它不再把一个 SM 绑定到特定任务上而是把整个大内核拆成一个动态任务池。任务被做成很小的粒度放进一个所有 SM 都能看到的公共结构里。每个 SM 先取自己本地队列的任务本地没活了就去别人那里偷。这就是标题里细粒度动态 SM 调度最核心的含义调度的对象不再是块而是微任务调度时机不再是编译期或者 kernel 启动前而是 SM 运行时的实时决策。2.2 任务切分规则M 维和 K 维都要切那么一个任务到底切成多大论文里用的不是简单按 token 切而是二维切分。一个专家 FFN 计算里M 维是 token 数K 维是 hidden 维度N 维是 FFN 中间维度。Weave 把 M 维按 token group 切把 K 维按 chunk 切两者组合成一个任务单元。一个任务可以看成五元组expert_id、token_group 起始位置、token_group 大小、k_chunk 起始位置、k_chunk 大小。任务执行时只需取对应专家的权重矩阵中的 k_chunk 范围和对应的 token 激活做一次小 GEMM把部分和写到结果缓冲里。这样切有一个很直接的好处M 维和 K 维的粒度可以分开调。M 维度决定 token 级并行的细度K 维度决定权重局部性。K 维切得大权重被重复加载的次数少K 维切得小任务粒度更细动态调度的空间更大。2.3 任务粒度不是越小越好算清调度开销这笔账任务粒度具体取多少是 Weave 设计里的一个核心权衡。任务太小比如单个任务只有 16×16 的矩阵乘那主流开销反而是取任务、放任务的原子操作和内存 fence真正的计算少得可怜。任务太大又会重新出现尾部效应因为最后一个 wave 还是塞不满。我按论文实验的常见范围估算过在 H100 上如果任务相当于是 64 token × 64 hidden chunk 的规模单任务执行时间在亚微秒级别到几微秒之间。此时做一次原子操作取任务开销占比可能在 1% 到 5% 上下是划算的。如果任务缩小到 16×16单任务执行时间可能只有几百纳秒原子操作占比就会高到不可接受。所以 Weave 的粒度选择是一个工程上必须反复调的参数不是越细越好。理想粒度是单任务执行时间在微秒量级、且一次调度开销占比小于几个百分点。2.4 和现有方案的对比grouped GEMM 和静态 EP 差在哪现在推理引擎里最常见的两种 MoE kernel 实现一个是 grouped GEMM一个是静态专家并行。grouped GEMM 把多个专家的矩阵乘打包成一个大的 kernel每个 group 对应一个专家block 按照 group 顺序依次处理。静态专家并行则是把整个服务器内的专家固定分到不同的 GPU每张 GPU 只跑自己分到的那些专家。这两种方案有一个共同问题GPU 内部的 block 调度是启动时定死的。grouped GEMM 启动时每个 group 有固定数量的 block静态 EP 启动时每个专家在 GPU 上也有固定数量的 block。一旦某个专家慢下来整个 GPU 的 SM 就得等它。Weave 是在这两者之上加了一个运行时调度层让空闲 SM 主动去找活而不是排队等活。3. 核心机制深挖Per-SM 任务队列、窃取策略与数据局部性3.1 双阶段 persistent kerneldispatch 与 compute 分离Weave 典型的 kernel 结构是双阶段的 persistent kernel。第一阶段是 dispatch一个很小的 kernel 或者 kernel 的前几个 block负责扫描路由结果把所有专家任务切成微任务填入任务队列。第二阶段是 computepersistent worker block 一直存活每个 block 绑定一个 SM不断从自己的本地队列取任务执行。为什么要保持 persistent因为 MoE 层如果按传统方式反复 launch 几十个专家 kernel每次 launch 的空档足够让所有 SM 歇菜。persistent kernel 让 worker block 一直都在任务之间没有 launch 间隙调度变成了纯软件逻辑。阶段之间用 global memory 里的队列做交接。dispatch 阶段把任务按专家进行分组写入compute 阶段通过队列指针消费。这两个阶段可以重叠dispatch 填充任务的速度远大于 compute 消费的速度差距就是队列冗余度。3.2 每 SM 本地队列的数据结构Weave 的队列不是一个大全局队列而是每个 SM 一个队列。为什么因为如果所有任务都塞进一个全局队列所有 SM 都用同一个原子计数器去抢任务那高并发下原子操作会成为新的瓶颈等于把尾部效应换成了并发冲突。每 SM 一个队列worker 优先从自己的队列取任务。这种 local-first 的设计同时解决两个问题一是减少全局原子操作竞争二是保住了数据局部性——同一个 SM 的任务尽量来自同一个专家、同一段 token权重和激活在 L2 里的复用率会更高。队列在 global memory 里用 head 和 tail 指针管理配合 atomicCAS 做并发控制。head 指针是本 SM 自己 pop 用的tail 指针是 dispatch 阶段 push 用的。偷取者会从别人的 tail 一侧取走任务尽量减少和被偷者本地 pop 的冲突。3.3 窃取触发条件与 victim 选择任务切得再细总有一些 SM 会先干完。当一个 SM 发现自己本地队列空了就会开始窃取。窃取的目标选择我用过的系统里最有效的策略是随机 victim而不是轮询。轮询看起来公平但在多 SM 同时窃取的场景下容易形成共振——所有空闲 SM 同时跑去偷同一个 victim。随机能让冲突概率摊开。偷取时不是随便拿一个任务就走而是尽量从 victim 队列的尾部取那种计算量大的任务。这样对 victim 的影响最小同时偷取者一次能拿到足够的活避免偷了个芝麻大小的任务、下次又得继续偷。还有一个很关键的判据如果 victim 队列里的任务太小偷取者会选择直接跳过宁可再随机选一个 victim。因为偷取开销是固定的偷一个执行时间比偷取开销还短的任务纯属负收益。3.4 数据局部性为什么动态调度还能不牺牲 L2 命中动态调度最怕的一件事是任务到处乱飞每个 SM 都要去 HBM 重新拉一遍专家权重把 MoE 本来就高的访存压力再推高一层。Weave 缓解这个问题的思路是任务亲和性。具体来说dispatch 阶段在把任务塞进队列时会尽量让同一个专家、相邻 token group 的任务落在相邻的几个 SM 队列里。这样当一个 SM 处理完一个任务下一个任务大概率还是同一个专家的专家权重在 L2 里还是热的。即使发生窃取偷取者拿到的任务也大概率来自同一个专家的相邻 chunk相当于把 L2 的复用单位从一个 SM 扩展到了两个 SM。另外H100 的 L2 有 residency control 特性可以把指定地址范围的专家权重标记为优先留在 L2。论文里用到了类似机制把高频专家的权重锁在 L2 里进一步降低 HBM 带宽压力。这一层优化虽然没有动态调度那么炫但对最终加速比的影响一点都不小。3.5 代码级骨架一套简化版队列窃取逻辑这一块我用伪代码给一个最小可跑思路。真正工程实现会复杂不少但核心逻辑就这几步struct Task { int expert_id; int token_start; int chunk_start; int chunk_size; }; __device__ int queue_head[MAX_SM]; __device__ int queue_tail[MAX_SM]; __device__ Task queue[MAX_SM][MAX_TASKS_PER_SM]; __device__ bool pop_local(int sm_id, Task* t) { int h queue_head[sm_id]; if (h queue_tail[sm_id]) return false; // 本地 pop用 atomicAdd 推进 head 指针 int slot atomicAdd(queue_head[sm_id], 1); if (slot queue_tail[sm_id]) { *t queue[sm_id][slot % MAX_TASKS_PER_SM]; return true; } return false; } __device__ bool steal(int from_sm, Task* t) { // 从 victim 的 tail 侧取任务用 atomicCAS 锁住 tail int t_old atomicAdd(queue_tail[from_sm], -1); if (t_old queue_head[from_sm]) { // 队列为空把 tail 恢复 atomicAdd(queue_tail[from_sm], 1); return false; } *t queue[from_sm][(t_old - 1) % MAX_TASKS_PER_SM]; return true; }这个简化版本里有两个细节容易踩坑。第一个是队列容量如果 dispatch 阶段产生的任务量瞬间超过队列容量push 端的 tail 指针会覆盖还没消费的数据。工程上要么把任务空间预分配得足够大并加上积压计数要么在 push 端做阻塞等待。第二个是内存可见性一个 SM 往队列里写数据另一个 SM 偷取时读数据中间必须有__threadfence()保证写入对所有 SM 可见否则你偷到的任务可能是一个 state 还没写完的半成品。4. 4×H100 实测解读2.89× 层加速到底来自哪里4.1 实验设置和评测口径论文的测试环境是 4 张 H100 SXM 80GB通过 NVLink 连接。模型是标准的 MoE LLM专家规模覆盖 8 专家到几百专家的范围测试场景同时包含 prefill 和 decode。这一点很重要因为2.89× 层加速这个数不是所有场景下都能拿到它是特定负载下的峰值收益。对照基线是静态调度的 persistent kernel也就是把专家计算按固定块大小切好、block 按预设映射分配到 SM 的那种传统实现。两边的算子本身用的都是高优化度的融合 kernel所以加速比主要反映的是调度策略的差距而不是基础 GEMM 实现谁写得更好。评测口径是单 MoE layer 的端到端时间router、dispatch、FFN、combine 全部算进去。这也就是标题里层加速的含义。论文报告的 2.89× 是 4×H100 上 decode 混合负载场景的层加速峰值。4.2 不同负载下的收益差异为什么 decode 收益远大于 prefill我把论文里不同场景的收益逻辑拆开看会发现一个明显规律任务越碎、越不均衡Weave 的收益越大。decode 场景下batch 里的每个请求 token 数不同生成进度完全异步路由结果每步都变。此时传统静态 kernel 的尾部效应最严重可能 GPU 上 132 个 SM 里有一半在空等。Weave 的动态窃取能让空闲 SM 立刻投入剩余任务SM 利用率提升非常显著。2.89× 这个峰值就是在 decode、混合 token 长度、专家负载波动大的场景下测出来的。prefill 场景尤其是长序列 prefill任务天然均匀一批 token 一次性进来每个专家的 token 数相对稳定任务块大而齐整。此时静态调度和动态调度的差距会缩小到 1.1~1.4 倍左右因为尾部效应本来就不严重动态调度省下的时间有一大半被调度原子操作抵消了。4.3 数字拆解2.89× 是怎么算出来的这个 2.89× 的背后我理解基本可以拆成两部分SM 利用率的提升和 L2 局部性的改善。第一部分也是占大头的一块是 SM 利用率。如果基线实现里 SM 的平均利用率只有 35%而 Weave 把它拉到 90% 上下单看利用率就是接近 2.6 倍的提升。剩下一部分来自局部性优化权重不反复回 HBM带宽压力下降计算单元的等待时间变短。这里有个容易误读的点2.89× 不是把 MoE 层从零优化到满分而是对比一个不算差、但确实有调度缺陷的基线。如果基线本身的持久性 kernel 已经写得很优化比如 block 粒度已经调过一轮那收益可能掉到 1.5× 左右。峰值加速这个表述本身就意味着不是常态。4.4 从层加速到端到端加速怎么估算整体收益做工程的人最关心的是层加速 2.89×整个模型的推理时延能提升多少这不能直接画等号要看 MoE 层在整个模型推理时间里的占比。假设一个模型里 MoE 层占了总时间的 p那么端到端加速比按 Amdahl 定律算总加速比 1 / ((1 - p) p / 2.89)如果 p 是 70%换算出来端到端大概 1.85 倍如果模型 MoE 层占比只有 50%端到端大概 1.49 倍。很多 MoE 模型里 FFN 相关计算确实能占到 60%-80%所以如果你在 4×H100 上部署这类模型层加速 2.89× 带来的实际吞吐提升很可能落在 1.5-2 倍这个区间。5. 工程落地在推理引擎里复现 Weave 的关键要点与坑5.1 集成位置替换 MoE Kernel不是改路由很多人在自己的推理框架里想引入 Weave 式调度下意识会先动路由逻辑、调负载均衡 loss方向是完全错的。Weave 是在执行层做优化输入依然是路由结果输出依然是专家计算结果只是它把给专家计算排班的职责从 GPU 硬件调度器手里接管过来改成了软件动态调度。在 vLLM、TensorRT-LLM 这类引擎里对应的工作是把原有的 grouped GEMM MoE kernel 替换成带动态调度逻辑的 persistent kernel。路由部分、所有对通信的优化、显存管理都可以保持原样。这也是 Weave 这类方案能落地的前提之一它不要求模型训练方式重新调整只替换推理实现。5.2 关键实现要点队列容量、block 数和原子调优第一个要确定的参数是 persistent block 数量。最简单合理的设置是 block 数等于单 GPU 的 SM 数或者 SM 数 × 2。H100 单卡 132 SM4 卡就是 528 个 SM。4×H100 环境下worker block 数直接取 528 即可。第二个是队列容量。任务总量最坏情况可以估算专家数 × token 组数 × K chunk 数。实际运行时不是所有任务都同时存在dispatch 是流式填充的所以队列深度不需要等于任务总量通常按一个专家最多同时活跃的任务数来预留具体数值要用 profiling 数据迭代调。第三个是原子操作调优。全局队列的原子计数器如果只用一个头指针所有 SM 并发 pop极端情况下每秒几亿次原子操作会拖慢整体速度。Per-SM 队列之所以比全局队列快本质就是把原子操作从全 GPU 竞争一个计数器变成了每个 SM 竞争自己的计数器冲突概率下降一个数量级。5.3 坑一跨 SM 队列的可见性问题这个坑我印象极深。GPU 的全局内存虽然所有 SM 都能访问但它有缓存层级L1 是每 SM 私有。一个 SM 往任务队列写数据另一个 SM 通过 L1 读到旧值是可能发生的事。在写队列数据结构时push 端写完任务字段后必须做__threadfence()把写入刷到对全局可见的级别再推进 tail 指针pop 端读到 tail 推进后还要再读一次任务字段确保读到的不是缓存旧值。如果省略这个 fence大概率出现偷取者拿到一个 expert_id 对、token 数据却不对的任务计算出来是错的。调试这种问题特别恶心因为它不是稳定复现而是频率低、随机触发。我的建议是写队列逻辑时第一版就严格把 fence 放在该放的位置别图省事。5.4 坑二高争用下偷取反而会变慢动态窃取不是免费的。在负载极高、所有 SM 都在忙、本地队列几乎不空的情况下窃取逻辑引入的原子操作就是纯开销。这个时候一个设计得不够聪明的 stealing 实现可能比静态调度还慢。解决办法是加阈值控制。当全局积压任务数低、空闲 SM 数量高时才启动盗窃逻辑。或者更细一点每个 SM 在连续 N 次本地 pop 都成功的情况下根本不去检查其他队列只有本地 pop 失败才触发窃取。这个 N 值可以做成运行时可调的用真实负载 profiling 后确定。5.5 与 CUDAGraph 和 vLLM 的兼容性推理引擎现在普遍用 CUDAGraph 来减少 kernel launch 开销。persistent kernel 动态任务队列这个形态天然适合 CUDAGraph因为 kernel 是长驻的graph 里只需要 capture 那一次 launch。任务队列在 graph 重放之间需要重置这个重置要小心不能在 graph capture 时把队列清空要在前一次计算完全结束后、下一次 dispatch 开始前完成 reset。vLLM 里接入时我的建议是先做独立 kernel 验证再进模型层集成。先写一个单独的 CUDA kernel把 MoE FFN 部分替换成 persistent worker用合成数据验证正确性和加速比再挂到完整的模型跑。直接改 MoE layer 再上模型调试成本会成倍增加。6. 边界与启发显存、千卡部署和不需要 Weave的场景6.1 先回应那个常见问题MoE 架构要全部参数进显存吗答案是要。而且这是理解 Weave 适用场景的前提。MoE 的稀疏性只体现在计算上一个 token 只激活 top-k 个专家所以计算量相对 dense 模型可以大幅减少。但权重存储上是另一回事全部专家权重必须常驻 GPU 显存因为你无法预知下一个 batch 的路由结果会击中哪些专家。任何专家都有可能在下一轮被激活提前换出去会付出灾难性的 HBM 回载代价。所以 MoE 部署的显存压力本质就在权重全驻留。Weave 不碰这个层面不压缩权重、不减少显存占用它做的是把已经驻留在显存里的权重用得更充分。如果你在规划千卡集群先算清楚的是总权重能不能塞进显卡Weave 解决的是塞进去之后怎么跑得更快。6.2 千卡部署下SM 级调度在一个节点内有效很大规模的部署里单是 4×H100 一个节点就开始出现新的瓶颈跨节点 all-to-all 通信。MoE 专家并行下token 经常要跨节点搬运NVLink 和 InfiniBand 的带宽再大也架不住反复搬。这时候要认清 Weave 的边界它优化的是token 到了 GPU 之后SM 怎么更快算完这一段。如果推理时间里有 40% 都耗在跨节点通信上SM 调度优化得再漂亮端到端收益也会被通信吃掉大半。我的建议是先做一轮 profiling把单层时间拆成通信时间和计算时间如果通信占比已经超过 40%优先去优化通信重叠和路由策略而不是继续压 kernel。等通信降下来之后把 Weave 这类 SM 级调度再叠加上去才能吃到那笔收益。6.3 什么时候不要用动态调度从工程角度讲至少有三种情况不适合硬上 Weave第一种专家数很少比如经典的 8 专家 Mixtral 规模。专家少的时候任务分布相对整齐静态持久性 kernel 调好块大小通常已经接近利用率上限动态调度的收益空间有限反而引入不必要的复杂度。第二种长序列大规模 prefill。负载越大、越均匀静态调度越接近最优动态调度能补的窟窿越小。第三种任务粒度过小且 K 维权重复用价值低的任务。如果矩阵本身很小原子调度开销在总时间里占比过高动态不如静态。判断标准我一般用一个粗暴的经验值先用 Nsight 看 SM 占用率。如果平均占用率已经超过 85%动态调度大概率帮不上忙如果占用率在 50% 上下波动那就值得认真考虑。6.4 这篇论文给我的最大启发读 Weave 让我重新意识到一件事GPU 调度优化远没到天花板。过去几年大家拼命优化算子和通信把 GEMM 效率推到极致但调度层一个简单的 work stealing 设计就能在大内核里再挖出接近 3 倍的单层性能。这个空间的入口就是让 SM 不再傻等去找活干。对于想自己试水的人我最后给一个具体动作拿你自己线上模型的某一个 MoE 层先用静态 persistent kernel 跑一遍记录 SM 占用率和 tail gap如果你发现占用率波谷明显、尾部有大量空闲再按照文章里讲的队列和窃取逻辑做一个最小实现优先对比 decode 场景的层延迟。这一套下来你就能快速判断 Weave 式调度值不值得引入到自己的框架里。
返回列表