上一篇 MoE 笔记把架构家族理了一遍——稀疏、稠密、Soft MoE，gate 在前还是在后，决定了算力形态和适用场景。
但那篇有一段伪代码一直让我不太舒服：
1for prob, idx in zip(topk_probs.unbind(1), topk_idx.unbind(1)):
2 out += prob.unsqueeze(1) * self.experts[idx](x)逻辑上没问题。但想象一下真实场景：一个 batch 有几千个 token，每个 token 选 2 个 expert，就是上万个 token-expert 计算任务。用 Python for loop 逐个调 expert 的 forward，GPU 利用率会低得离谱。
那实际做的时候，token 到底是怎么被"分发"给各个 expert 的？分发之后，几十个 expert 各自不同的矩阵乘法又是怎么在 GPU 上高效执行的？
这篇想把单个 Sparse MoE Layer 的执行流从头走一遍。
FFN：每个 token 都过同一套参数#
标准 Decoder-only Transformer Block（Pre-Norm）的两步计算，在之前的文章里拆过了：
$$ H = X + \operatorname{Attention}(\operatorname{Norm}_1(X)) $$$$ Y = H + \operatorname{FFN}(\operatorname{Norm}_2(H)) $$Attention 做 token 之间的信息路由，FFN 做逐 token 的非线性变换。关键词是"逐 token"——所有 token 共享完全相同的 FFN 权重矩阵，互不交互。
如果想扩大模型容量，最直接的做法就是把 FFN 的 intermediate dimension 搞大。但代价是线性的：参数翻倍，每个 token 的计算量也翻倍。
那篇 MoE 笔记已经提过，MoE 的核心想法是把这两者部分解耦：准备 \(E\) 个独立的 FFN（Expert），但每个 token 只用其中 \(K\) 个。设每个 Expert 的参数量为 \(P\)：
$$ P_{\text{total,experts}} = E \cdot P $$$$ P_{\text{active,experts}} = K \cdot P $$这里只计 Expert 层的参数，Attention、Embedding、Router 等共享组件不包含在内。
MoE 通过让不同 token 激活不同的参数子集，实现总参数规模与单 token 计算量的部分解耦。
“部分解耦"不是免费的。Router、Token Dispatch、Expert 并行执行都有额外开销，总 FLOPs 不会简单地等于 Dense baseline 的 \(K/E\) 倍。参数量、FLOPs、显存占用是三个不同的概念。
和 MHA 是不是类似的#
我第一次看 MoE 结构时，脑子里跳出来的联想就是 Multi-Head Attention：MHA 有多个 Head，MoE 有多个 Expert，都是"多个子网络并行 → 合并结果"的形态。
但仔细一想，差别其实很大。
MHA 的每个 Head 都参与计算。8 个 Head 就是 8 个全算，然后 concat 拼起来。在标准配置下（\(d_{\text{model}} = 512\)，8 Head，\(d_{\text{head}} = 64\)），总计算量大致等于一个 512 维 single-head 的水平——那篇文章拆过这个数。增加 Head 数量（保持 \(d_{\text{model}}\) 不变）不增加参数量和计算量，只是把子空间切得更细。
Sparse MoE 则完全不同。8 个 Expert、Top-2 意味着每个 token 只用 2 个 Expert，剩下 6 个不参与。增加 Expert 数量会线性增加总参数量，但不增加单 token 的计算量（\(K\) 不变时）。
还有一点容易误解：不要把 Expert 想象成"Expert 1 专门处理数学，Expert 2 专门处理语言"这样的清晰语义分工。Expert 的功能分化是 Router 在训练中自然学到的，程度因模型和数据而异，不是架构本身保证的。
Router 的选择过程#
用一组具体数字来走一遍。
1Hidden Size = 4096
2Number of Experts (E) = 8
3Top-K = 2输入 token 的表示 \(x \in \mathbb{R}^{4096}\)。Router 就是一个线性层：
$$ s = x W_g, \quad W_g \in \mathbb{R}^{4096 \times 8} $$输出 8 个 score，对应 8 个 Expert。假设算出来是：
1s = [1.2, -0.3, 0.8, 2.1, -1.0, 0.5, 0.1, -0.7]Top-2 选出得分最高的两个：Expert 3（score 2.1）和 Expert 0（score 1.2）。
对选中的 score 做 Softmax，得到路由权重：
$$ p_3 = \frac{e^{2.1}}{e^{2.1} + e^{1.2}} \approx 0.71, \quad p_0 = \frac{e^{1.2}}{e^{2.1} + e^{1.2}} \approx 0.29 $$这个 token 的 MoE 输出：
$$ \operatorname{MoE}(x) = \sum_{i \in \operatorname{TopK}(s)} p_i \cdot E_i(x) $$几个容易忽略的点：
- 不同 token 通常会选到不同的 Expert 组合——同一个 batch 里，有的 token 走 Expert 0+3，有的走 Expert 1+5。
- 同一个 token 在不同 Transformer Layer 可以选不同的 Expert。每层有自己独立的 Router 参数和 Expert 参数。
- 每个 Expert 的输出维度和输入维度相同（都是 Hidden Size），所以加权求和后 shape 不变。
关于 Softmax 的位置：上面的例子是"先 Top-K，再对选中 score 做 Softmax”。实际实现中也有"先对全部 8 个 score 做 Softmax，再取 Top-K，然后可能再 renormalize"的做法。不同论文和框架的选择不同。
Expert 里面是什么#
典型 Transformer MoE 的 Expert 就是一个独立的 FFN。以 SwiGLU 为例：
$$ E_i(x) = W_{\text{down},i}\left(\operatorname{SiLU}(x W_{\text{gate},i}) \odot (x W_{\text{up},i})\right) $$三个投影矩阵：\(W_{\text{gate},i}\) 通过 SiLU 激活产生门控信号，\(W_{\text{up},i}\) 将输入映射到 intermediate dimension，\(\odot\) 是逐元素乘法，\(W_{\text{down},i}\) 映射回 Hidden Size。
这里特别要区分两个"Gate"。MoE 的 Router（有时也叫 Gating Network）负责选哪些 Expert。SwiGLU 里的 Gate Projection 是 FFN 内部的门控计算。两者都叫 Gate，但不是一回事。
用 Mixtral 8×7B 核对一下真实参数：
| 参数 | 值 |
|---|---|
| Hidden Size | 4096 |
| Intermediate Size | 14336 |
| Number of Experts | 8 |
| Top-K | 2 |
每个 Expert 有 3 个矩阵，单个 Expert 的参数量约 \(3 \times 4096 \times 14336 \approx 1.76 \times 10^8\)，8 个 Expert 合计约 14 亿参数。
Mixtral 8×7B 这个名字容易让人误以为有八套完整的 7B Transformer。实际上只有 FFN 层被替换成了 8 个 Expert，Attention 层是所有 Expert 共享的。八个 Expert 的 FFN 参数加上共享的 Attention、Embedding 等组件，总参数量约 46.7B，每个 token 只激活 2 个 Expert，活跃参数量约 12.9B。
梯度怎么过 top-k 这一刀#
Top-K 的选择是离散的——选了 Expert 0 和 Expert 3，没选 Expert 1。这个"选/不选"的决定怎么做梯度反传？
类比 Dropout 有助于理解。两者的共同点：forward 时只有部分路径被激活，backward 时梯度只沿激活过的路径传回去。没被选中的 Expert 不会从该 token 获得梯度。
区别在于：Dropout 的 mask 是随机的，和输入无关；MoE Router 的选择取决于输入 \(x\) 和可学习的 Router 参数 \(W_g\)。确定性 Router 下相同输入和参数得到相同路由。
具体看梯度路径。假设选了 Expert 0 和 Expert 2：
$$ y = p_0 \cdot E_0(x) + p_2 \cdot E_2(x) $$Expert 0 和 Expert 2 的参数通过各自的 \(E_i(x)\) 获得梯度。Router 参数 \(W_g\) 通过路由权重 \(p_0, p_2\) 参与计算图——这些权重是 Softmax 产生的连续值，可以正常传梯度。
Top-K 操作产生的是离散索引，不可导。但选中之后的路由权重 \(p_0, p_2\) 是可导的，梯度通过它们回传到 Router。
有个微妙之处。如果是 hard top-1（只选一个 Expert，输出权重恒为 1），那么 Router 的 Softmax 权重变成常数，任务损失对 Router 几乎没有梯度信号。这也是很多 MoE 论文要加辅助损失（Auxiliary Loss）的原因之一。辅助损失的具体设计不在本文展开。
Token Dispatch：从 token-major 到 expert-major#
前面说的都是单个 token 的视角。现在切换到 batch 级别。
4 个 Token、3 个 Expert、Top-2。原始输入 \(X\) 的 shape 为 \([4, 4096]\)。Router 给出的路由结果：
| Token | 选中 Expert | 路由权重 |
|---|---|---|
| T0 | E0, E2 | 0.7, 0.3 |
| T1 | E1, E2 | 0.4, 0.6 |
| T2 | E0, E1 | 0.2, 0.8 |
| T3 | E0, E2 | 0.9, 0.1 |
问题来了。Expert 0 需要处理 T0、T2、T3。Expert 1 需要处理 T1、T2。Expert 2 需要处理 T0、T1、T3。
但在原始 Tensor 里，这 4 个 token 按索引 0-3 连续排列。Expert 0 的三个 token 散落在第 0、2、3 行——内存中不连续。而且不同 Expert 有不同的权重矩阵，不能混着算。
如果不做重排，就只能用 for loop 逐个 Expert 地 gather 数据再做矩阵乘。这正是开头那段伪代码的问题。
Token Dispatch 做的事情是：把输入 Tensor 从 token-major 重排成 expert-major，让每个 Expert 要处理的 token 在内存中连续排列。
按 Expert 分组后的布局：
1Expert 0: T0, T2, T3 → 3 个 token
2Expert 1: T1, T2 → 2 个 token
3Expert 2: T0, T1, T3 → 3 个 token拼成一个连续的 Dispatch Buffer：
$$ X_{\text{dispatch}}.\text{shape} = [8, 4096] $$这个 8 是怎么来的——原始只有 4 个 token，但 Top-2 产生了 \(4 \times 2 = 8\) 个 token-expert 计算任务。同一个 token（比如 T0）在 Dispatch Buffer 中出现了两次，分别给 Expert 0 和 Expert 2。重排的是整个 hidden state 所在的行，不是对特征维度排序。
记录重排信息：
1permuted_token_ids = [0, 2, 3, 1, 2, 0, 1, 3]
2tokens_per_expert = [3, 2, 3]用 tokens_per_expert 做 split，每个 Expert 得到连续的子矩阵：
1E0 的输入: X_dispatch[0:3] → shape [3, 4096]
2E1 的输入: X_dispatch[3:5] → shape [2, 4096]
3E2 的输入: X_dispatch[5:8] → shape [3, 4096]每个 Expert 拿到的就是一个规整的 \([M_i, 4096]\) 矩阵，可以直接做 GEMM。
Token Dispatch 有两层目的：
- 按 Expert 的独立参数组织计算任务——每个 Expert 有自己的权重矩阵，需要知道"哪些 token 归我"。
- 为 GEMM 准备连续的数据布局——GPU 上的矩阵乘法要求输入在内存中连续。把同一个 Expert 的 token 排在一起，每个 Expert 就能直接做一次 GEMM，不需要额外的 scatter/gather。
为什么不能一个一个 GEMM 地算#
Dispatch 解决了数据布局问题，但还有一个问题没解决：怎么高效地执行这些 GEMM。
假设 64 个 Expert，每个 Expert 平均收到 8 个 token，需要执行 64 个独立的 GEMM（以 SwiGLU 的一个投影为例）：
$$ [8, 4096] \times [4096, 14336] $$逐个 Expert 启动 Kernel 有两个问题。
Kernel Launch 开销。64 次独立 Launch，每次涉及 CPU-GPU 命令提交和 Kernel 调度。当 Kernel 本身计算量很小时（M=8 就是一个很小的矩阵乘），Launch 开销在总时间中的占比会变得显著。
GPU 利用率不足。单个 \([8, 4096] \times [4096, 14336]\) 的 GEMM，M 维度只有 8。以常见的 128×128 Tile 来说，M 方向只能切出 1 个 Tile。GPU 上大量 SM 可能只有少数几个有活干——算术强度那篇讨论过，小 GEMM 很容易掉进 memory-bound 区间。
Grouped GEMM 的做法是把这 64 个 GEMM 打包，用同一个 Kernel 统一调度。
和普通 GEMM 的区别：每个子 GEMM 使用不同的权重矩阵，M 维度也可以不同（每个 Expert 收到的 token 数量不一定一样）。
和 Batched GEMM 的区别：Batched GEMM 通常假设所有子 GEMM 有相同的 M/N/K 尺寸。MoE 场景下各 Expert 的 token 数量不同，Batched GEMM 不直接适用。
和 CUDA Graph 的区别：CUDA Graph 优化一组 CUDA 操作的提交开销，把多次 Launch 打包成一次提交。但 CUDA Graph 不会把 64 个 GEMM 融合进同一个 Kernel。Grouped GEMM 在 Kernel 内部统一调度多个 GEMM 的 Tile，不只是减少 Launch 次数，还能改善 Tile 在 SM 之间的分配。 两者可以组合使用。
用一个简化的性能模型感受差异：
$$ T_{\text{separate}} \approx N \cdot L + \sum_i T_i $$其中 \(L\) 是每次独立 Launch 的固定开销，\(N\) 是 GEMM 数量，\(T_i\) 是第 \(i\) 个 GEMM 的纯计算时间。Grouped GEMM 的 Launch 次数大幅减少，同时 Tile 调度更紧凑。
这个模型是粗略的，不能当作精确计时公式。实际的 Grouped GEMM 实现也不一定只有一个 Kernel Launch——具体取决于 API 和实现。
当每个 Expert 收到的 token 很少时（\(T_i\) 小，\(L\) 占比高），Grouped GEMM 的优势最明显。反过来，如果每个 Expert 收到几千个 token，单个 GEMM 本身就足够大，Launch 开销可以忽略，SM 也能被充分利用，Grouped GEMM 的相对优势就减弱了。不能说它在所有矩阵形状下都更快。
NVIDIA 的 CUTLASS 库和 Transformer Engine 都提供了 Grouped GEMM 实现。Grouped GEMM 本身是通用的 GPU 计算技术，不是 MoE 专属。
Token Combine：把结果合回去#
Expert 计算完成后，Dispatch Buffer 里的 8 个结果需要按原始 token 索引合并回去。
沿用前面的例子。Expert 计算后，Dispatch Buffer 中每个位置的结果和归属：
1位置 0: E0(T0) → Token 0, 权重 0.7
2位置 1: E0(T2) → Token 2, 权重 0.2
3位置 2: E0(T3) → Token 3, 权重 0.9
4位置 3: E1(T1) → Token 1, 权重 0.4
5位置 4: E1(T2) → Token 2, 权重 0.8
6位置 5: E2(T0) → Token 0, 权重 0.3
7位置 6: E2(T1) → Token 1, 权重 0.6
8位置 7: E2(T3) → Token 3, 权重 0.1按 token 加权聚合：
$$ Y_{T0} = 0.7 \cdot E_0(T0) + 0.3 \cdot E_2(T0) $$$$ Y_{T1} = 0.4 \cdot E_1(T1) + 0.6 \cdot E_2(T1) $$$$ Y_{T2} = 0.2 \cdot E_0(T2) + 0.8 \cdot E_1(T2) $$$$ Y_{T3} = 0.9 \cdot E_0(T3) + 0.1 \cdot E_2(T3) $$最终 \(Y\) 的 shape 为 \([4, 4096]\)。Token 数量没有变——Dispatch 阶段被复制给多个 Expert 处理的 token，在 Combine 阶段又被加权合并了回来。
用 PyTorch 的基本操作表示这个逻辑：
1# expert_output: [8, 4096] -- Dispatch Buffer 中每个位置的 Expert 输出
2# token_ids: [8] -- 每个位置属于哪个原始 token
3# weights: [8] -- 每个位置对应的路由权重
4
5token_ids = torch.tensor([0, 2, 3, 1, 2, 0, 1, 3])
6weights = torch.tensor([0.7, 0.2, 0.9, 0.4, 0.8, 0.3, 0.6, 0.1])
7
8weighted_output = expert_output * weights.unsqueeze(1) # [8, 4096]
9
10output = torch.zeros(4, 4096)
11output.index_add_(0, token_ids, weighted_output) # [4, 4096]index_add_ 的语义：对于每个位置 \(j\)，把 weighted_output[j] 累加到 output[token_ids[j]]。同一个 token 从不同 Expert 收到的加权结果就被加在了一起。
这段代码是教学用的逻辑等价代码。真实训练代码中的 weights 必须是从 Router Softmax 计算图中流过来的 Tensor，不能用硬编码常量替代——否则 Router 参数收不到梯度。
完整的数据流从 Router 到最终输出：
如果 Expert 不在同一张 GPU 上#
到目前为止，所有讨论都假设全部 Expert 在同一张 GPU 上。
但真实 LLM 的 Expert 数量可能很多（DeepSeek-V3 有 256 个），单卡放不下。一个自然的做法是把 Expert 分散到多张 GPU 上——比如 8 个 Expert、4 张 GPU，每卡放 2 个。
这时候 Token Dispatch 就不再是本地的内存重排。当某个 token 选中的 Expert 不在本 GPU 上时，hidden state 需要通过网络发到那张 GPU，计算完再发回来。通信原语那篇介绍过的 All-to-All 正好适合这个场景：每张 GPU 需要向其他 GPU 发送不同数量的 token，同时接收来自其他 GPU 的 token。
MoE 是模型架构，Expert Parallelism 是分布式执行策略。 就像 Tensor Parallelism 把矩阵乘法的计算拆到多卡一样，Expert Parallelism 把不同 Expert 的计算分配到不同 GPU。
留一个问题给下一篇：
如果 Router 可以把任意 token 路由到任意 Expert，而 Expert 被分散在不同 GPU 上，那跨 GPU 的 Token Dispatch 应该怎么组织，才能减少通信量并提高整体吞吐？
参考资料#
- Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (Shazeer et al., 2017)
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (Fedus et al., 2021)
- Mixtral of Experts (Jiang et al., 2024)
- NVIDIA CUTLASS: Grouped GEMM
- NVIDIA Transformer Engine Documentation