MoE 少算 FLOPs,延迟反而涨了:慢在搬 token 的路上
Avi Chawla 给过一道面试题:dense FFN 换成 top-2 MoE,profiler 显示每 token FLOPs 下降,端到端延迟却上升。答案藏在 router 之后的 permutation、dispatch、all-to-all 和 combine 里。小 batch decode 会把这些成本放大。
一道技术面试题:把 Transformer 里的 dense feed-forward layer 换成 top-2 MoE。profiler 说每 token 执行的 FLOPs 更少,端到端推理延迟却上升。为什么?
Avi Chawla 在 MoE inference engineering 一文里拆过这条路径。先把两个概念分开看。
dense 和 MoE 的差别
标准 Transformer 层里,attention 先混合 token 位置的信息,随后 FFN 对每个 token 独立做变换。所有 token 走同一个 FFN,权重相同。把 FFN 换成 MoE 后,一层里有多个 expert,每个 expert 是一个独立 FFN。router 读取每个 token 的 hidden state,给可用 expert 打分,选出 top-2,并给出两个 routing weight。只有被选中的 expert 处理这个 token。模型里存着很多 expert 权重,每 token 只执行一小部分。
所以 profiler 看到每 token FLOPs 下降,合理。问题在端到端。MoE 省的是专家计算,搬运和调度没有省。
router 之后还有一串动作
top-2 路由意味着一个 token 产生两条 assignment。一个 batch 有 N 个 token,就有 2N 条 assignment。运行时按 expert 分组:E1 收到哪些 token,E2 收到哪些 token,形成大小不同的矩阵。这个分组叫 dispatch。token 行按 expert 排序叫 permutation。每个 expert 跑自己的矩阵,通常用 grouped GEMM 把多个专家矩阵合到一次 kernel launch。结果再用 inverse mapping 还原 token 顺序,乘 routing weight,相加,叫 combine 和 unpermutation。
prefill 一次进来很多 token,expert 矩阵行数多,GPU 好用。decode 一次每个请求只加一个 token,并发不高时,某个 expert 可能只收到一两行。kernel launch、读权重、搬 token 行的成本会冒出来。grouped GEMM 能减少 launch,不能给没有 assignment 的 expert 造出工作。固定形状 padding 还会把空行算进去。profile 只看 FLOPs,会漏掉这些段。
跨 GPU 时,激活要出门
expert 多到单卡放不下,就要 expert parallelism。不同 expert 放在不同 GPU。token 在 GPU 0 上,选中了 GPU 1 上的 expert,激活向量和路由 metadata 就要通过网络过去,算完再回来。一个 batch 里不同 token 选不同 expert,每块 GPU 可能同时向多块 GPU 发激活,也从多块 GPU 收激活。这个模式叫 all-to-all。
素材里的视频把一次 batch 的路径画出来了。T1 选本卡 E1,激活留在 GPU 0。T2 选同服务器另一块 GPU 上的 E4,走 NVLink。T3 选另一台服务器上的 E7,走 InfiniBand 或 Ethernet。每层 MoE 都重复一次这样的交换。跨服务器链路比机内链路贵,延迟会累积。
最慢的 GPU 决定这一步
路由很难把工作切得完全均匀。一个热门 expert 收到半个 batch 的 assignment,它所在 GPU 就要算更大的矩阵。其他 GPU 早算完也要等。MoE 层的 step latency 由最慢的参与 GPU 决定。平均利用率会掩盖尾巴。要按 serving step 记录 assignment 数、dispatch 时间、expert 执行时间、combine 时间、GPU idle 时间,还要按通信域记录字节和时间。机内 NVLink 和跨机 InfiniBand 混在一个总计数器里,看不出问题在哪。
placement 可以把常被一起选中的 expert 放近一点,把热门 expert 放到没那么忙的 GPU。replication 可以给一个逻辑 expert 建多个物理副本,router 选择不变,运行时把 assignment 发给可用副本。这两种办法都占 GPU 内存,也占控制带宽。流量尖峰可能在 rebalance 完成前就结束。
内存按总参数算
MoE 还有一个容易混的账。Qwen3-30B-A3B 的模型卡写 30.5B total parameters 和 3.3B activated parameters,每 token 选 128 个 routed experts 里的 8 个,共 48 层。3.3B 描述的是一个 token 的执行路径,部署要存的是 30.5B 总参数、attention、embedding、临时 buffer 和 KV cache。按每个参数 2 字节存权重,30.5B 需要约 61 GB,这还没算 KV cache 和通信 buffer。稀疏激活降低专家计算,不降低权重存储。激活参数和常驻参数要分开测。
优化从测量开始
能做的优化有几类。把常被选中的 expert 放进同一个 NVLink 域,减少远程流量。让路由更均衡,避免一块 GPU 拖住整层。把远程激活传输和本地 expert 计算 overlap。用 grouped GEMM 减少 kernel launch。量化权重或激活,减少内存和网络字节,但要测数值误差。改 top-k 会改变执行的 expert 集合,属于模型行为变化,要用同一批请求做质量对比。placement 和 kernel 改动改的是运行时,量化改数值,top-k 改计算集合,这几类实验要分开。
端到端延迟不按 FLOPs 结算。它按最慢的那条链路结算。MoE 把计算切稀疏了,也把 token 推上了路。profiler 里少掉的 FLOPs,可能变成 dispatch、all-to-all 和等待。
参考:Avi Chawla 的 MoE inference engineering 文章 https://x.com/i/article/2100852768403861504 ,原推 https://x.com/DailyDoseOfDS_/status/2118999411254767695 。
发布时间: 2026-11-07 23:27