从零理解 MoE 算子：路由、专家并行与 Token 重排

MoE（Mixture of Experts，混合专家模型）已经成为大模型扩展参数规模的重要技术路线。 本文从普通 FFN 出发，逐步解释 Router、Top-K 路由、专家计算、Expert Parallel、AllToAll 通信、Token Permute / Unpermute，以及 MoE 中常见的 Capacity、负载均衡和通信优化问题。

1. 为什么需要 MoE

Transformer Block 通常由 Attention 和 FFN 两个主要部分组成。为了理解 MoE，可以先忽略 Attention，只看 FFN。

一个常见的 SwiGLU FFN 可以写成：

FFN(x) = down_proj(silu(gate_proj(x)) * up_proj(x))

其中：

普通 FFN 的特点是：所有 token 都经过同一组参数。想增加模型容量，通常需要增大 FFN，这会直接增加每个 token 的计算量。

MoE 提供了另一种扩展方式：

准备多个 FFN 专家，但每个 token 只选择少数几个专家执行。

这类设计称为稀疏 MoE。

这里的“稀疏”指的是计算路径稀疏，而不是权重矩阵本身稀疏。

2. 普通 FFN 到 MoE FFN

普通 FFN 可以理解为：

所有 token -> 同一个 FFN

MoE FFN 可以理解为：

不同 token -> 不同专家 FFN

假设有 E=8 个专家，每个专家都是一个独立 FFN。对于每个 token，只选择得分最高的 top_k=2 个专家执行。

token x
-> Router
-> Top-K experts
-> Expert FFN
-> 按路由权重加权求和

例如：

总专家数 E = 64
每个 token 选择 top_k = 2

模型拥有 64 组专家参数，但每个 token 实际只经过其中 2 组。

3. Router：每个 Token 应该去哪里

Router 也常被称为 Gate。它通常是一个轻量级线性层：

router_logits = x @ router_weight.T

如果：

x : [T, H]
router_weight : [E, H]

那么：

router_logits : [T, E]

router_logits[t, e] 表示第 t 个 token 交给第 e 个专家处理的倾向。

随后，对每个 token 选择得分最高的 K 个专家：

topk_scores, topk_indices = topk(router_logits, k=top_k)
topk_weights = softmax(topk_scores)

得到：

topk_indices : [T, K]
topk_weights : [T, K]

例如：

4. Top-K 为什么会让 Token 变多

当 top_k > 1 时，一个原始 token 会被发送给多个专家。

逻辑上可以理解为：

原始 token 数 = T
进入专家计算的路由 token 数 = T * K

例如：

T = 1024
K = 2

专家侧总共需要处理：

1024 * 2 = 2048 个路由 token

不过需要注意：

Top-K 复制首先是一个逻辑概念。

在实际实现中，不一定一开始就真的复制完整的激活张量。很多系统会先保存：

expert_ids
source_token_ids
topk_weights

真正的数据复制或打包，可能发生在后面的 pack、dispatch 或 permute 阶段。

5. 单设备 MoE 的完整流程

如果所有专家都位于同一设备，MoE 的流程可以概括为：

Router
-> Top-K 选择专家
-> 按专家整理 token
-> 执行专家 FFN
-> 按原 token 位置和路由权重合并结果

更具体地说：

输入 tokens [T, H]
-> Router 得到 topk_indices 和 topk_weights
-> 构造路由 token [T * K, H]
-> 按专家编号重排 token
-> 专家 FFN 计算
-> 根据 source_token_ids 和 topk_weights 合并
-> 输出 tokens [T, H]

6. 为什么需要按专家分组

假设路由结果是：

原始顺序是：

[A(exp1), B(exp0), C(exp1), D(exp0)]

为了批量执行专家计算，需要整理为：

[B(exp0), D(exp0), A(exp1), C(exp1)]

这就是 Token Permute：

按照专家编号重新排列 token，使同一个专家的 token 连续存放。

它不改变 token 的内容，只改变 token 的排列顺序。

7. GroupedMatMul：批量执行多个专家

每个专家都有一套独立权重。如果逐 token、逐专家调用小矩阵乘，硬件利用率通常很差。

例如有 64 个专家，如果每个专家单独调用一次矩阵乘，可能产生大量小 GEMM 和大量 kernel launch，整体效率很低。

更常见的做法是：

将同一专家的 token 连续排列，
然后用 GroupedMatMul 一次性执行多个专家的矩阵乘。

可以理解为：

Expert 0: X0 @ W0
Expert 1: X1 @ W1
Expert 2: X2 @ W2
…

GroupedMatMul 需要知道每个专家对应的 token 范围。常见做法是传入累计 token 数：

每个专家 token 数: [3, 2, 4]
group_list: [3, 5, 9]

它表示：

tokens [0, 3) -> expert 0
tokens [3, 5) -> expert 1
tokens [5, 9) -> expert 2

因此，group_list 本质上是专家 token 分段边界。

8. 专家内部仍然是 FFN

每个专家本质上仍然是一个 SwiGLU FFN：

Expert_e(x) = W2_e(silu(W1_gate_e(x)) * W1_up_e(x))

实际执行通常拆为：

GMM1
-> SwiGLU
-> GMM2

其中：

GMM1: 升维
输入 routed_tokens : [T_routed, H]
权重 up_weight : [E, H, 2I]
输出 gate_up : [T_routed, 2I]

SwiGLU: 激活
输入 gate_up : [T_routed, 2I]
过程 gate, up = split(gate_up)
hidden = silu(gate) * up
输出 hidden : [T_routed, I]

GMM2: 降维
输入 hidden : [T_routed, I]
权重 down_weight : [E, I, H]
输出 expert_out : [T_routed, H]

可以概括为：

按专家排列的 tokens
-> GMM1
-> SwiGLU
-> GMM2
-> 按专家排列的输出

9. 为什么需要 Expert Parallel

当专家数量很多时，单个设备可能无法容纳所有专家权重。

Expert Parallel，简称 EP，就是将不同专家分布到不同设备或 Rank 上。Rank 可以简单理解为“分布式系统中的一个进程编号”，一般情况下，一张GPU或者NPU就是一个Rank。

假设：

总专家数 E = 4
EP 并行度 = 2

可以分配为：

Rank 0: Expert 0, Expert 1
Rank 1: Expert 2, Expert 3

Router 仍可能将 Rank 0 上的 token 路由给 Expert 2。由于 Expert 2 位于 Rank 1，token 必须跨设备发送。

这就是 MoE 中 AllToAll 通信的来源。

10. AllToAll Dispatch：将 Token 发送给专家

在 EP 场景中，每个 Rank 都可能需要向多个 Rank 发送 token，同时也会从多个 Rank 接收 token。

Dispatch 阶段的目标是：

将每个路由 token 发送给持有目标专家的 Rank。

例如：

Rank 0 上的 token 可能要发送给 Rank 0 和 Rank 1
Rank 1 上的 token 也可能要发送给 Rank 0 和 Rank 1

AllToAll Dispatch 之后，每个 Rank 得到的是：

本 Rank 持有的专家需要处理的 token

也就是说，Dispatch 解决的是：

token 应该去哪台设备

而不是：

本地 token 是否已经按专家排好

11. Rank-Major 与 Expert-Major

Dispatch 之后，本地接收 buffer 可能按照来源 Rank 排列。

例如当前 Rank 持有两个本地专家：

Expert 0
Expert 1

它从两个 Rank 收到了 token：

Rank 0 -> Expert 0: [A, B]
Rank 0 -> Expert 1: [C]
Rank 1 -> Expert 0: [D]
Rank 1 -> Expert 1: [E, F]

Dispatch 后可能得到 rank-major 布局：

[A(E0), B(E0), C(E1), D(E0), E(E1), F(E1)]

也就是：

[R0-E0 | R0-E1 | R1-E0 | R1-E1]

这种布局称为 rank-major：

先按来源 Rank 分组，再按专家分组。

但是 GMM 更希望同一个专家的 token 连续：

[A(E0), B(E0), D(E0), C(E1), E(E1), F(E1)]

也就是：

[R0-E0 | R1-E0 | R0-E1 | R1-E1]

这种布局称为 expert-major：

先按专家分组，再按来源 Rank 分组。

12. Permute：从 Rank-Major 到 Expert-Major

Dispatch 解决了跨 Rank 发送问题，但本地为了高效执行 GMM，通常还需要一次 Permute。

rank-major
-> permute
-> expert-major

例如：

rank-major:
[A(E0), B(E0), C(E1), D(E0), E(E1), F(E1)]

Permute 后：

expert-major:
[A(E0), B(E0), D(E0), C(E1), E(E1), F(E1)]

对应的 permute 索引为：

[0, 1, 3, 2, 4, 5]

每个专家 token 数为：

[3, 3]

GMM 使用的累计 group_list 为：

[3, 6]

Permute 可以简化理解为：

permuted_tokens = tokens[permute_indices]

13. Unpermute：从 Expert-Major 恢复 Rank-Major

本地专家计算完成后，输出仍然是 expert-major：

[A', B', D', C', E', F']

但 Combine AllToAll 通常希望按照来源 Rank 连续发送回去，因此需要恢复为 rank-major：

[A', B', C', D', E', F']

这就是 Unpermute。

可以简化理解为：

restored = zeros_like_original_buffer()
restored[permute_indices] = expert_output

它是 Permute 的逆操作。

14. Dispatch 和 Permute 的边界不是固定的

需要特别注意：不同框架中，Dispatch 和 Permute 的边界并不完全一样。

有些实现会把流程拆得很清楚：

Dispatch = AllToAll 通信
Permute = 本地 token 重排

也就是：

AllToAll Dispatch
-> rank-major
-> Permute
-> expert-major

但也有些实现会把多个步骤融合在一起：

Dispatch = pack + permute + alltoall

甚至 Dispatch 输出就已经是适合专家计算的 expert-major 布局。

因此，阅读源码时不要只看函数名字，而要看这个函数的输入输出布局。

判断标准是：

输出 token 是否已经按专家连续排列？

如果已经按专家连续排列，那么后面可能不再需要单独的 Permute。

如果只是完成了跨 Rank 通信，那么后面通常还需要本地 Permute。

15. AllToAll Combine：将专家结果返回来源 Rank

专家计算完成后，每个 token 的专家输出需要回到原始 token 所在的 Rank。

Combine 阶段的目标是：

将专家输出发送回 token 的来源 Rank。

整体可以理解为 Dispatch 的反向通信：

Dispatch:
原始 Rank -> 专家所在 Rank

Combine:
专家所在 Rank -> 原始 Rank

Combine 完成后，每个 Rank 拿回自己原始 token 对应的专家输出。

如果 top_k = 2，每个原始 token 会对应两个专家输出。最后还需要根据路由权重做加权合并。

16. Top-K 专家输出如何合并

一个原始 token 可能被多个专家处理。

例如：

token_0 -> expert_1, weight = 0.7
token_0 -> expert_3, weight = 0.3

两个专家分别输出：

out_1
out_3

最终输出为：

output[token_0] = 0.7 * out_1 + 0.3 * out_3

一般形式为：

output[token] = sum(weight[token, k] * expert_output[token, k])

这里必须保证：

expert_output[token, k]

能够正确映射回原始 token 和对应的 top-k 位置。

因此实现中通常需要保存：

source_token_ids
topk_indices
topk_weights

或者保存等价的反向映射。

17. EP MoE 正向流程总览

现在可以把 EP MoE 的完整正向流程串起来：

Router
-> Top-K 选择
-> 构造路由 token
-> AllToAll Dispatch
-> Permute
-> GMM1
-> SwiGLU
-> GMM2
-> Unpermute
-> AllToAll Combine
-> 按原 token 合并 Top-K 专家输出

更具体地说：

输入 tokens [T, H]
-> Router 得到 [T, E] logits
-> Top-K 得到 [T, K] expert_ids 和 weights
-> 逻辑上形成 [T * K, H] routed_tokens
-> Dispatch 发送到专家所在 Rank
-> Permute 整理成本地 expert-major
-> GMM1 升维
-> SwiGLU 激活
-> GMM2 降维
-> Unpermute 恢复 rank-major
-> Combine 发送回来源 Rank
-> Weighted Merge 合并 Top-K 输出
-> 输出 tokens [T, H]

18. Capacity：每个专家最多处理多少 Token

Router 不会天然保证每个专家收到相同数量的 token。

可能出现：

Expert 0: 100 tokens
Expert 1: 900 tokens
Expert 2: 0 tokens
Expert 3: 300 tokens

如果不加限制，某些专家会收到过多 token，导致：

计算长尾
通信不均衡
显存占用不稳定
GMM 规模波动很大

因此很多 MoE 系统会引入 Capacity，也就是每个专家最多能接收多少 token。

一个常见形式是：

expert_capacity = ceil(tokens_per_rank * top_k / num_experts * capacity_factor)

例如：

tokens_per_rank = 1024
top_k = 2
num_experts = 8
capacity_factor = 1.25

平均每个专家接收：

1024 * 2 / 8 = 256

容量为：

256 * 1.25 = 320

也就是说，每个专家最多处理大约 320 个 token。

超过容量的 token 可能会被：

drop
reroute
保留残差路径
或者使用其他策略处理

Capacity 的核心作用是：

限制专家负载上界，让计算和通信更稳定。

19. Load Balancing：为什么需要负载均衡

如果 Router 完全自由选择专家，它可能会逐渐偏向少数专家。

极端情况下：

Expert 0: 很多 token
Expert 1: 很少 token
Expert 2: 几乎不用
Expert 3: 几乎不用

这会带来两个问题。

第一是性能问题：

部分专家过载，部分专家空闲。

第二是训练问题：

部分专家得不到充分训练。

因此 MoE 训练中通常会加入负载均衡机制。常见做法包括：

辅助负载均衡损失
容量限制
路由噪声
专家复制
更受约束的路由策略

20. Auxiliary Loss：辅助负载均衡损失

Auxiliary Loss，简称 Aux Loss，是 MoE 中常见的负载均衡手段。

它的目标是：

鼓励 Router 尽量均匀地使用不同专家。

直观地说，如果某个 batch 中大多数 token 都去了 Expert 0，而其他专家几乎没有 token，那么 Aux Loss 会变大。

训练时总损失变成：

total_loss = task_loss + aux_loss_weight * aux_loss

其中：

task_loss

是原本的语言模型损失，

aux_loss

是负载均衡损失。

Aux Loss 不直接提高单个 token 的预测能力，但它能改善专家使用分布，让训练更稳定。

不过 Aux Loss 也有代价：

过强的负载均衡约束可能影响模型自由选择最合适专家。

因此实际训练中需要调节它的权重。

21. Group Limited Routing：减少跨节点通信

最基础的 MoE 路由是：

从所有专家中选择 Top-K 专家。

但在大规模分布式训练中，这可能导致大量跨节点通信。

例如某个 token 在 Rank 0 上，但 Router 选择的专家都在远端 Rank 上，那么这个 token 必须跨设备发送。

为了降低通信开销，一些现代 MoE 系统会使用更受限制的路由策略，例如 Group Limited Routing。

它的大致思想是：

先把专家划分成若干组
Router 先选择少数几个专家组
再在这些组内选择具体专家

这样可以减少 token 被路由到任意远端专家的概率，从而降低 AllToAll 通信压力。

可以理解为：

普通 Top-K Routing:
token 可以从所有专家里自由选择

Group Limited Routing:
token 先选专家组，再在组内选专家

它牺牲了一部分路由自由度，换取更好的通信局部性和系统效率。

22. 反向传播如何理解

MoE 的反向传播可以从两个原则理解：

1. 通信方向与正向相反
2. 重排操作的反向是对应的逆重排

正向流程可以简化为：

Dispatch
-> Permute
-> Experts
-> Unpermute
-> Combine

反向传播大致对应：

Combine 的反向
-> Unpermute 的反向
-> Experts 的反向
-> Permute 的反向
-> Dispatch 的反向

其中：

Permute 的反向 ~= Unpermute
Unpermute 的反向 ~= Permute

专家内部 FFN 的反向包括：

GMM2 backward
-> SwiGLU backward
-> GMM1 backward

需要计算：

输入梯度 dX
升维投影权重梯度 dW1
降维投影权重梯度 dW2

对于 SwiGLU 专家：

hidden = silu(gate) * up

反向时需要分别计算：

d_gate
d_up

然后再通过 GMM1 的反向得到输入梯度和权重梯度。

23. MoE 为什么难优化

MoE 的性能问题不是单一矩阵乘问题，而是通信、计算和数据重排交织在一起。

23.1 AllToAll 通信开销

EP 引入两次主要通信：

Dispatch: 将 token 发送给目标专家
Combine : 将专家结果返回 token 来源 Rank

EP 规模增大后，通信延迟、通信带宽和网络拥塞都会影响性能。

对于大规模 MoE 来说，AllToAll 往往是最关键的瓶颈之一。

23.2 Token 重排开销

Permute 和 Unpermute 会读写大量激活值：

[T_routed, H]

它们不涉及复杂数学运算，但会消耗显著内存带宽。

尤其当：

T 很大
H 很大
top_k > 1

时，重排开销可能非常明显。

23.3 负载不均衡

Router 可能让不同专家收到不同数量的 token：

Expert 0: 100 tokens
Expert 1: 900 tokens
Expert 2: 0 tokens
Expert 3: 300 tokens

这会导致：

部分专家成为长尾
部分设备工作量更大
某些专家收到零 token
GMM 的矩阵规模波动
通信量不均衡

负载不均衡不仅影响性能，也会影响训练稳定性。

23.4 小算子和中间结果开销

MoE 流程中存在多个步骤：

Dispatch
-> Permute
-> GMM1
-> SwiGLU
-> GMM2
-> Unpermute
-> Combine

如果每一步都独立执行，可能产生额外的调度开销和中间结果读写。

因此 MoE 优化经常围绕以下方向展开：

算子融合
通信与计算重叠
减少中间激活读写
Token 重排优化
GroupedMatMul 优化
更好的负载均衡

24. 通信与计算重叠

朴素实现可能是：

等待 Dispatch 完成
-> 再开始专家计算
-> 等待专家计算完成
-> 再执行 Combine

这样通信和计算是串行的。

更高效的实现会尝试：

一边接收 token
一边对已经到达的 token 执行专家计算

也就是通信与计算重叠。

理想情况下：

通信开销被部分隐藏在计算开销后面。

例如：

Rank 先收到 Expert 0 的 token
就可以先执行 Expert 0 的 GMM
同时继续接收其他 Expert 的 token

这样可以减少总耗时。

不过通信与计算重叠会让实现复杂很多，因为需要处理：

buffer 管理
异步通信
执行顺序
专家 token 到达顺序
流同步
反向传播中的依赖关系

25. 算子融合与 Megakernel 思路

MoE 中很多步骤本身并不复杂，但每一步都会读写中间结果。

例如：

GMM1 输出 gate_up
SwiGLU 读取 gate_up，写出 hidden
GMM2 再读取 hidden

如果完全分开执行，中间激活会多次写入和读取显存。

优化方向之一是算子融合：

减少中间结果落显存
减少 kernel launch
提高数据局部性

更激进的方向是 Megakernel，即把多个阶段融合到一个更大的 kernel 或更紧密的执行单元中。

不过 Megakernel 不一定总是更好。它可能带来：

寄存器压力增加
代码复杂度上升
并行粒度变差
调度灵活性下降

因此实际优化需要结合硬件特点、token 分布、专家规模和通信模式综合判断。

26. 实现时容易忽略的边界条件

MoE 的基础流程看起来不复杂，但工程实现需要处理很多边界条件。

26.1 零 Token 专家

某个专家可能没有收到任何 token：

tokens_per_expert = [4, 0, 7, 2]

实现中不能假设每个专家都有输入。

26.2 某个 Rank 接收零 Token

在极端不均衡情况下，某个 Rank 可能完全没有收到 token。

通信、内存分配和反向传播仍然要保持正确。

26.3 Token 数无法整除批处理粒度

专家收到的 token 数通常不规则。

最后一个处理块可能不足完整大小，需要处理尾部 token。

26.4 Top-K 合并必须正确

一个原始 token 可能由多个专家共同处理。

输出需要按照路由权重正确合并：

output[token] = sum(weight[token, k] * expert_output[token, k])

如果 source token 映射错误，结果会直接错误。

26.5 Permute 映射必须可逆

Unpermute 需要知道每个 expert-major token 原本位于 rank-major buffer 的哪个位置。

常见方案有两类：

保存 Permute 索引
根据路由计数和偏移重新构造映射

无论采用哪种方案，都必须保证映射可逆。

26.6 Capacity 溢出

如果某个专家超过 capacity，需要明确处理策略：

丢弃
截断
重新路由
走残差路径
或者使用其他补偿机制

这部分必须和训练逻辑保持一致。

26.7 不同精度下的数值稳定性

MoE 中 Router、Softmax、Top-K、专家 FFN 可能使用不同精度。

常见组合包括：

FP32
FP16
BF16
FP8

Router 的 logits 和 softmax 对数值稳定性比较敏感，实际实现中经常需要特别处理。

27. 用伪代码串起整个过程

下面的伪代码省略了很多工程细节，但可以帮助建立整体认识：

def moe_forward(x, router_weight, expert_weights, top_k, ep_group):
# x: [T, H]

# 1. Router
logits = linear(x, router_weight) # [T, E]

# 2. Top-K routing
scores, expert_ids = topk(logits, k=top_k) # [T, K]
weights = softmax(scores, dim=–1) # [T, K]

# 3. Expand routed tokens logically
routed_x, source_token_ids = expand_by_topk(
x,
expert_ids
) # [T * K, H]

# 4. Send routed tokens to ranks that own target experts
received_x, received_expert_ids = all_to_all_dispatch(
routed_x,
expert_ids,
group=ep_group
)

# 5. rank-major -> expert-major
permuted_x, permute_indices, group_list = permute_by_expert(
received_x,
received_expert_ids
)

# 6. Expert FFN
gate_up = grouped_matmul(
permuted_x,
expert_weights.up,
group_list
) # [T_routed, 2I]

hidden = swiglu(gate_up) # [T_routed, I]

expert_out = grouped_matmul(
hidden,
expert_weights.down,
group_list
) # [T_routed, H]

# 7. expert-major -> rank-major
rank_major_out = unpermute(
expert_out,
permute_indices
)

# 8. Send expert results back to source ranks
returned_out = all_to_all_combine(
rank_major_out,
group=ep_group
)

# 9. Merge top-k expert results for each original token
y = weighted_merge(
returned_out,
source_token_ids,
weights
) # [T, H]

return y

28. 常见术语速查

29. 总结

理解 MoE，可以抓住四条主线。

第一条是数学结构：

Router
-> Top-K 专家
-> 加权合并

第二条是专家计算：

按专家分组
-> GMM1
-> SwiGLU
-> GMM2

第三条是分布式数据流：

Dispatch
-> Permute
-> Experts
-> Unpermute
-> Combine

第四条是系统优化：

Capacity
-> Load Balancing
-> AllToAll 优化
-> 通信计算重叠
-> Token 重排优化
-> 算子融合

其中最容易混淆的是：

Dispatch 负责跨 Rank 发送 token
Permute 负责将本地 token 整理为 GMM 需要的 Expert-Major 布局
Unpermute 负责恢复适合 Combine 的 Rank-Major 布局
Combine 负责将结果返回原始 Rank

但在具体实现中，Dispatch 和 Permute 的边界并不固定。有些系统会把 pack、permute、alltoall 融合进一个 dispatch 接口。因此阅读源码时，最重要的不是函数名字，而是看清楚：

输入是什么布局？
输出是什么布局？
token 是否已经按专家连续排列？
是否还需要恢复来源 Rank 的顺序？

只要抓住这些数据布局变化，MoE 的大部分工程实现都会变得更容易理解。