Fused MoE V2 优化速览
这篇文档把 Ling-2.6-1T MoE serving 优化拆成三层： 先看普通 MoE / EPMoE 的数据排布，再看 Fused MoE V1 为什么要把 scatter、expert FFN、gather 放进同一个 Pallas kernel，最后看 Fused MoE V2 怎样把 token、weight、output accumulator 的生命周期改成一条 TPU 上更容易调度的流水线。
核心结论可以先记住一句：V2 的收益不来自更少的 MoE FLOPs，而来自把通信、HBM 读写、VMEM 常驻、VPU scale / activation、MXU dot 放到同一个 kernel 时间线里， 让能被隐藏的成本尽量离开 critical path。
1. MoE 层在做什么
MoE 层可以看成 router + routed expert FFN + weighted sum：
Ling-2.6-1T 的典型形状：
| 项 | 数值 |
|---|---|
| routed experts | 256 |
| top-k | 8 |
| hidden | 8192 |
| intermediate | 2048 |
| EP size | 32 |
| 每个 device 持有的 local experts | 8 |
| 16384-token prefill 下 local tokens | 512 |
EP=32 时，一个 token 选中的 8 个 expert 很可能分布在不同 device 上。 所以每层 MoE 不只是 FFN 计算，还天然包含一次 token scatter 和一次 output return / gather。优化的重点会落到两个问题：
- token 怎么按 expert 重排、跨 chip 发送、再放回原 token 顺序。
- expert FFN 怎么在不等通信、不等 HBM 的情况下持续喂给 MXU。
2. EPMoE baseline：permute、GMM、unpermute
普通 EPMoE 的代码结构贴近 MoE 的数学定义：
_permute -> _dispatch -> _gmm_compute -> _unpermute -> psum / combine它先根据 router 结果把 token/top-k slots 按 expert 聚到一起，让 GMM 可以按 expert group 连续计算；GMM 完成后，再把 expert-sorted output 放回原来的 token/top-k slots，最后按 routing weight 做加权求和。
permute / unpermute 到底在排什么
用一个 top_k=2 的小例子说明。原始 flatten 后是 token-major：
| 原 slot | token | top-k slot | selected expert |
|---|---|---|---|
| 0 | T0 | k0 | E2 |
| 1 | T0 | k1 | E0 |
| 2 | T1 | k0 | E1 |
| 3 | T1 | k1 | E2 |
| 4 | T2 | k0 | E0 |
| 5 | T2 | k1 | E1 |
| 6 | T3 | k0 | E2 |
| 7 | T3 | k1 | E1 |
_permute 之后变成 expert-major：
E0 group: slot1, slot4
E1 group: slot2, slot5, slot7
E2 group: slot0, slot3, slot6GMM 在这个顺序里连续处理每个 expert group。_unpermute 再用 inverse permutation 把输出放回 slot0..slot7，然后 reshape 成 [tokens, top_k, hidden]，按 top-k routing weight 做 weighted sum。
实现上，_permute 不一定物化完整的 sorted token tensor。它主要生成 token_indices、sorted_selected_experts 和 group_sizes，真正读取 token 可以放在 GMM 里做，例如 inputs_2d[token_indices]。这样可以减少一段大的 [T * top_k, hidden] 中间物化。
EPMoE 的优点是表达清楚、通用性强；瓶颈是通信、GMM、unpermute、psum 被多个 较粗粒度 op 串起来，kernel 很难同时看到 ICI、HBM、VMEM、VPU 和 MXU 的完整 时间线。
3. TPU v7 MXU mental model
TPU 的数据流可以简化成：
HBM -> VMEM -> MXU / VPU -> VMEM -> HBMMXU 是 systolic array。矩阵乘 C[M,N] = A[M,K] x B[K,N] 进入 MXU 时， 可以把 A 看成从一侧流入的 activation / token tile，把 B 看成从另一侧流入的 weight tile，C 是持续累加的 output accumulator。
Google Cloud TPU 架构文档里的硬件事实可以用来建立直觉： TPU v6e 和 TPU7x 预览版的 MXU 是 256 x 256 systolic array，v6e 之前是 128 x 128；MXU 提供 TensorCore 里的主要计算能力；乘法输入是 bfloat16， 累加是 FP32。
让 MXU 更忙的四个条件
| 条件 | 含义 |
|---|---|
| tile 要够大 | BT/BK/BN 要摊薄 systolic fill/drain、vmatmul setup、scale 和 reorder 的固定成本。 |
| 复用侧要常驻 | 常驻的不一定永远是 A；谁复用次数高、放得下、能减少关键路径等待，就优先留在 VMEM。 |
| C 要常驻 | output accumulator 留在 VMEM，K / bf 循环里只累加，最后一次性写回 HBM。 |
| 流动侧要预取 | 当前 tile 在 MXU 上计算时，下一块 streamed tile 用 DMA 进另一个 VMEM slot。 |
对 Fused MoE V2 来说，常驻侧通常是 token tile 和 output accumulator； weight tile 更大、随 expert/projection 变化更快，所以用 double buffer 连续流入。
好的 tile 不只是 K 对齐
对 256 x 256 MXU 来说，K 最好是 256 的整数倍，但这只是必要条件之一。 M/N 决定阵列有多少行列可以同时工作，K 决定这段工作可以持续多久，总 FLOPs 决定固定成本能不能被摊薄。
| tile 形状 | 利用率直觉 | 原因 |
|---|---|---|
[256,1024] x [1024,256] | 好 | M/N 覆盖 256 x 256 阵列，K 有 4 个 256-wide wave。 |
[512,1024] x [1024,512] | 好 | 可以分成多个 256 x 256 output tiles，每个 tile 都有足够 K 流水。 |
[32,8192] x [8192,256] | 差 | K 很长，但 M 只有 32，阵列行方向吃不满；MoE decode 常见。 |
[256,129] x [129,256] | 差 | M/N 看起来够大，但 K 太短且不对齐，padding/setup 占比高。 |
TPU v7 microbenchmark 的规律也支持这个判断：独立 GEMM 在几十 GFLOPs 以下很容易被 fixed cost 主导；超过约 0.1T FLOPs 后利用率才明显健康；同时要有足够的 K_tiles_256 和足够大的 M/N output face。Fused MoE 的单 expert tile 通常不是 天然接近 dense GEMM peak 的形状，因此 V2 的重点是 pipeline scheduling 和 overlap， 不是宣称每个小 expert GEMM 都能打满峰值。
4. Fused MoE V1：先拿到调度权
V1 的第一步不是改变 MoE 数学结构，而是改变执行边界：把 scatter、expert FFN、 gather 收进同一个 Pallas kernel。这样 kernel 才能同时安排 ICI 通信、HBM 读写、 VMEM staging、MXU dot 和 VPU activation。
V1 gather 是什么
这里的 gather 不是单纯的本地数组 gather。scatter 把每个 token/top-k slot 送到 持有目标 expert 的 device；expert 计算完成后，gather 是这条路径的回程：
- 把每个 routed expert 的输出从 expert-local scratch 送回原 token 所在的位置。
- 按原始 token/top-k slot 顺序做本地重排。
- 根据 routing weight 做 top-k 加权聚合，形成
T x hidden输出。
因此它既包含跨 device 的 return / all-to-all-like 数据移动，也包含本地 gather 和 weighted sum。名字叫 gather，是因为从 token 视角看，它在收回自己的 top-k expert outputs；从网络视角看，它是 scatter 的反向数据搬运。
V1 block config 控制什么
V1 的 config 同时给外层 token 生命周期和内层 FFN 生命周期定颗粒度：
| 旋钮 | 控制对象 | 影响 |
|---|---|---|
bt | local token tile | 一次让多少 local tokens 进入 fused kernel。太小会让 scatter/gather 和 metadata 过碎，太大推高 scratch 上界。 |
bts / btc | expert 内 token staging / compute M 维 | bts 是 HBM -> VMEM token staging 粒度；btc 是实际 dot 的 token subtile。 |
bf / bfc | FFN intermediate tile | bf 越大，intermediate wave 越厚、wave 数越少；代价是 W1/W3/W2 buffer 和 accumulator 变大。 |
bd1 | FFN1 hidden input slice | V1 为了 VMEM 放得下而切 hidden=8192 的第一刀。 |
bd2 | FFN2 hidden output slice | 第二刀；每个 bd2 slice 都可能触发 FFN2 partial output 的 load/add/store。 |
Ling V1 tuned config 在 16384-token prefill、EP=32 下是：
bt=128, bf=1024, bd1=2048, bd2=2048,
bts=128, btc=32, bfc=1024, bd1c=2048, bd2c=2048, bse=1024把它展开成循环次数：
- local tokens = 512，
bt=128，所以每个 device 有 4 个外层bttiles。 - intermediate = 2048，
bf=1024，所以每个有 token 的 expert 有 2 个bfwaves。 - hidden = 8192，
bd1=bd2=2048，所以 hidden 被切成 4 片。 - 每个
bfwave 里，FFN1 对 4 个bd1slices 做 gate/up 两组 dot； FFN2 对 4 个bd2slices 做 down projection。
FFN1 和 FFN2 的理论 dot 量并不一样：同一个 bf wave 里，FFN1 要分别算 W1 gate 和 W3 up，所以是 4 个 hidden slices x 2 组 dot；FFN2 是 4 个 hidden output slices x 1 组 dot。实际耗时还会受 activation、scale、DMA、HBM staging 影响。
V1 的 b_acc_vmem
b_acc_vmem 保存一个 bf wave 内的 gate/up intermediate accumulator。V1 会按 bd1 扫 hidden input：每扫一片 hidden，就做一次 partial dot，并把结果累加到 b_acc_vmem。这和 tensor parallel 切 weight 后需要对 partial result 做累加的 直觉类似，只是这里的切分发生在 expert kernel 内部的 hidden 维。
在 Ling V1 tuned config 下，b_acc_vmem 的主尺寸约为：
2 x a2a_max_tokens x bf x fp32
= 2 x (128 x 32) x 1024 x 4B
= 32 MiB这里的 2 是 gate/up 两个 accumulator，a2a_max_tokens=bt*EP=4096 是一个 bt tile 在最坏路由下可能落到任意 expert/device 的 routed slots 上界。 V1 的 scatter 和 expert FFN 是两阶段组织：scatter 先把这个 bt tile 可能产生的 routed tokens 放进 A2A scratch，expert_ffn 再从 scratch 里按 expert 处理。 因此 V1 的 accumulator / scratch layout 需要覆盖整个 bt tile 可能产生的 routed token 空间，而不是像 V2 那样只拿一个较小的 bts resident tile 周转。
V1 的限制
V1 已经拿到了 kernel 内调度权，但 expert 内部仍然围绕 bf -> bd1/bd2 周转：
bd1/bd2 控制住了 VMEM 占用，却把 FFN 拆成多轮小 GEMM、小 DMA 和小 VPU work。 更关键的是，FFN2 output 需要跨多个 bf waves 累加，V1 的 partial output accumulator 仍然在 HBM，后续 wave 要 read-add-store。
5. 为什么 V1 不能直接 full-hidden
直觉上，把 bd1/bd2 放大到 full hidden 可以减少 slice 数。但在 V1 的 buffer layout 里，这不是一个独立旋钮：hidden slice 变大时，token staging、W1/W3/W2 double buffer、FFN2 staging buffer 会一起变大。即使某个形状勉强放得下，V1 的 loop order 仍然是 bf -> bd1/bd2，FFN2 accumulator 仍然在 HBM；这只是更大的 V1， 不是 V2。
以 Ling prefill 形状估算主要 VMEM scratch：
| 方案 | 主要 buffer | 估算 |
|---|---|---|
| V1 tuned | W1/W3/W2 double buffer、token stage、FFN2 stage、b_acc_vmem、gather/top-k scratch | 约 55 MiB |
| V1 假设 full hidden | bd1/bd2 从 2048 放大到 8192，多个 hidden-sliced buffer 同时放大 | 约 99 MiB |
| V2 tuned | full-hidden weight streaming、b_x_vmem、gate/up intermediate、b_y_acc_vmem、writeback/gather scratch | 约 46 MiB |
V2 的关键不是简单把 V1 的 hidden slice 调大，而是重新定义哪些东西常驻、哪些东西 流过来。
6. Fused MoE V2：token/output 常驻，weight 流水
V2 的 loop order 可以概括为：
for each bts token tile:
load tokens once into b_x_vmem
initialize / keep b_y_acc_vmem
for each bf intermediate tile:
prefetch W1/W3/W2 into double-buffered VMEM slots
W1/W3 dot -> gate/up accumulator
activation / scale
W2 dot -> accumulate into b_y_acc_vmem
write output once也就是源码注释里的 bts loop is OUTER, bf loop is INNER。
V2 的 VMEM 角色更清晰：
| VMEM buffer | 作用 |
|---|---|
b_x_vmem | 当前 bts token tile，跨多个 bf weight tiles 复用。 |
b_y_acc_vmem | f32 final output accumulator，跨所有 bf waves 累加。 |
b_gate_acc_vmem / b_up_acc_vmem | 当前 bf wave 的 gate/up intermediate，仍然存在但只按 bts x bf 周转。 |
b_w1_x2_vmem / b_w3_x2_vmem / b_w2_x2_vmem | 权重双缓冲；当前 tile 计算时，下一 tile 通过 DMA 进入另一个 slot。 |
b_x_scale_vmem | activation quant 打开时的 per-token scale。 |
V2 会不会多读 weight
会。V2 选择 token/output 常驻，weight 连续流入；当一个 expert 的 routed tokens 超过 bts 时，同一组 expert weights 可能被不同 bts tiles 反复读取。以 local routed tokens 约 512、bts=160 为例，一个 expert 可能需要约 4 个 bts tiles。
这个 trade-off 的目标不是让 weight bytes 的理论值最小，而是让 weight DMA 更大块、 更规则、更容易被 double buffering 藏在 MXU compute 后面。V1 可能读更少的 weight slice，但它的关键路径被小 slice staging、b_acc_vmem 大上界和 FFN2 HBM read-add-store 切碎。V2 接受更多可预取的 weight streaming，换来更长的 resident compute window。
V2 不切 hidden 的好处
V2 没有 bd1/bd2，所以每个 bf tile 直接读取 full hidden 方向的权重。 这带来三个效果：
- FFN1/FFN2 不再被 hidden slices 切成很多小 dot，MXU 的 K 维流水更完整。
- final output accumulator 在 VMEM 中跨
bf累加，避免 FFN2 partial output 的 HBM read-add-store。 - token tile 进入 VMEM 后跨多个
bftiles 复用，DMA/MXU/VPU 的 overlap 窗口更长。
V2 block config 主旋钮
| 旋钮 | 控制对象 | 说明 |
|---|---|---|
bt | 外层 token tile | routing、scatter、gather 的单位。 |
bts / btc | resident token tile / compute subtile | bts 控制 token 常驻工作集，btc 控制 MXU token M 维 subtile。 |
bf | intermediate / weight wave | W1/W3 的 N，W2 的 K；也是 weight streaming 的主节奏。 |
bse | shared expert intermediate tile | V2 把 shared expert 放进 kernel 时使用。 |
Ling full shared-expert、act-quant ON 的 V2 tuned config 使用：
bt=128, bf=512, btc=80, bse=512, bts=160bts=160 让 token/output 常驻工作集保持在 VMEM ceiling 内；bf=512 降低单个 intermediate wave 的 accumulator / weight buffer 压力，为 double buffering、 scatter/gather 和 shared expert 留出空间。
7. Activation quant、direct scaled dot、shared expert
Activation quant
V2 可以先把 bf16 token quantize 成 fp8 token，并把每个 token 的 scale 带进 VMEM。scatter 阶段如果打开 activation quant，就搬 fp8 token 而不是 bf16 token， ICI bytes 直接降低。代价是 VPU 需要做 quantize、scale、pack/reorder；如果 e2e 瓶颈已经转移到别处，act-quant 的端到端收益可能接近噪声。
Per-channel FP8 和 direct scaled dot
如果 scale 依赖 K 维，例如 per-block scale，大 GEMM 会被拆成多个 K-block dot， 每段 dot 后都要乘 scale 再累加。Ling 的 MoE 权重是 per-channel FP8，scale 只依赖 输出 channel，所以可以先做完整 fp8 dot，再在输出上乘 scale。这样 MXU 看到的是 更完整的 K 维流水。
| scale 类型 | MXU 看到的形态 |
|---|---|
| per-block | K 维被切成多段 dot，VPU scale 反复插入。 |
| per-channel | scale 不依赖 K，可以 hoist 到 dot 之后，dot 更完整。 |
In-kernel shared expert
Ling 每层还有一个 shared expert。如果它在 kernel 外单独跑，就会在 routed MoE 之外新增一段 dense MLP critical path。V2 把 shared expert 放到 routed expert loop 之前，利用 scatter recv / metadata 的时间窗，并复用 routed path 的 token、 weight 和 output VMEM buffers。
V1 在 Ling linear path 里 shared expert 是 external dense MLP；V2 full path 则把 shared expert 放进 kernel，因此 V1 的 tuned config 不能直接套给 V2。
8. 性能数字和剩余瓶颈
原始分享材料里的关键数字：
| 场景 | V1 | V2 | 解释 |
|---|---|---|---|
| MoE layer prefill，16384 tokens | 约 5.16 ms | 约 2.44 ms | 约 53% latency 降低，主要来自 VMEM 常驻、大 tile、通信 overlap、act quant 和 shared expert overlap。 |
| MoE layer decode，512 tokens | 约 0.249 ms | 约 0.211 ms | 约 15% 提升；decode 的 per-expert M 太小，更接近 weight HBM 读 lower bound。 |
| 端到端 prefill throughput | baseline | 约 25% 提升 | MoE 不再是唯一瓶颈，GLA / Lightning Linear Attention prefill kernel 成为主要剩余成本。 |
V2 不是无限优化，剩余约束仍然明确：
- 通信侧：32-way all-to-all 受 topology、multi-hop forwarding、link contention 限制。
- 计算侧：fp8 pack/unpack、lane reorder、scale broadcast 会占用 VPU。
- VMEM 侧：约 64 MiB/core 的 VMEM ceiling 限制
bts继续变大，也限制更多 prefetch buffer 同时常驻。
9. 代码入口
以下路径用于把上面的数据流对应回 sglang-jax 代码：
| 主题 | 代码位置 | 关注点 |
|---|---|---|
| EPMoE baseline | python/sgl_jax/srt/layers/moe.py | _permute -> _gmm_compute -> _unpermute -> psum。 |
| Fused V1 kernel | python/sgl_jax/srt/kernels/fused_moe/v1/kernel.py | bt/bf/bd1/bd2 的 hidden-sliced 生命周期。 |
| Fused V2 kernel | python/sgl_jax/srt/kernels/fused_moe/v2/kernel.py | bts outer、bf inner、b_x_vmem、b_y_acc_vmem、weight double buffer。 |
| Fused MoE wrapper | python/sgl_jax/srt/layers/fused_moe.py | V1/V2 backend 选择、direct scaled dot、activation quant 参数。 |
| V2 tuned config | python/sgl_jax/srt/kernels/fused_moe/v2/tuned_block_configs.py | Ling shape 下的 bt/bf/bts/btc/bse。 |
10. 总结
EPMoE 的优点是数学表达清晰；V1 的关键进步是把 MoE 的执行边界收进一个 Pallas kernel，拿到通信和 expert compute 的调度权；V2 的关键进步是改变数据 生命周期，让 token tile 和 output accumulator 常驻 VMEM，让权重以更大、更规则的 tile 连续流入，并用 double buffering 把 HBM DMA 藏在 MXU compute 后面。
因此，Fused MoE V2 的本质不是“少几个 op”，而是把 MoE 从多个 op 的组合改造成 一条可调度的数据流生产线。通信、HBM、VPU、MXU 都在同一个 kernel 时间线上排队， 优化目标是让能隐藏的成本尽量不出现在 critical path 上。