MoE混合专家模型计算原理与推理部署实战解析

发布时间：2026/11/3 16:20:14
MoE混合专家模型计算原理与推理部署实战解析

MoE（Mixture of Experts，混合专家）这两年从学术论文里的一个"稀疏激活"技巧，变成了大模型推理部署绕不开的架构选择。DeepSeek、Qwen、Mixtral 这些模型把 MoE 推到台前之后，很多做 AI Infra 的人发现：光会调 vLLM 的接口已经不够用了，你得真正搞明白一次前向传播里，token 是怎么被路由到不同专家、专家又是怎么算的、显存和算力到底花在哪。这篇笔记就围绕 MoE 的计算原理和过程展开，把路由、门控、专家计算、负载均衡这几件事从数学到工程实现讲透，适合正在做推理优化、想搞懂 MoE 底层计算流程的工程师参考。

1. 从稠密到稀疏：MoE 到底在解决什么问题

1.1 稠密模型的算力困境

先想清楚一件事：为什么要有 MoE。传统的 Transformer 是稠密（Dense）结构，一个 FFN 层里所有的参数，对每一个 token 都要参与计算。假设一个 FFN 的中间维度是 4 倍隐藏维度，隐藏维度 4096，那这个 FFN 就有大约 4096 × 16384 × 2 ≈ 1.3 亿参数。模型层数一多，参数量轻松上百亿。

问题在于，推理时每个 token 都要过一遍全部参数。参数量越大，单 token 的计算量（FLOPs）就越大，延迟越高。你想让模型更"聪明"，就得加参数；但加了参数，推理成本就线性上涨。这就是稠密模型的死结：模型容量和推理成本被死死绑在一起。

MoE 的思路很直接：既然不是每个 token 都需要用到全部知识，那能不能让不同的 token 走不同的参数子集？把一个大 FFN 拆成 N 个小 FFN（每个叫一个"专家"Expert），每个 token 只激活其中 top-k 个专家。这样总参数量可以做得很大（模型容量大），但每个 token 实际参与计算的参数量很小（计算量小）。

1.2 稀疏激活带来的账本变化

举个具体的账。假设有 8 个专家，每个专家参数量 1 亿，总共 8 亿参数。如果 top-2 路由，每个 token 只激活 2 个专家，实际计算量相当于 2 亿参数的稠密模型。也就是说，你用 8 亿的显存占用，换来了接近 2 亿稠密模型的计算量。这就是稀疏激活的核心红利。

但这里有个容易被忽略的点：显存占用并不会因为稀疏激活而减少。所有专家的参数都得放在显存里待命，因为你不知道下一个 token 会路由到哪个专家。所以 MoE 省的是算力（FLOPs），不是显存。这一点在做部署规划时特别关键，很多人第一次接触 MoE 会误以为"只激活 2 个专家，那显存只要装 2 个专家就行"，这是错的。

注意：MoE 的显存占用约等于全部专家参数之和，加上路由网络和共享层。做显存估算时，按总参数量算，不要按激活参数量算。

1.3 MoE 在 Transformer 里的位置

标准做法是把 Transformer 里的 FFN 层替换成 MoE 层，Attention 层保持不变。也就是说，每个 Transformer Block 的结构变成：LayerNorm → Self-Attention → LayerNorm → MoE-FFN。MoE-FFN 内部包含一个门控网络（Router/Gate）和 N 个专家 FFN。

为什么替换 FFN 而不是 Attention？因为 FFN 占据了 Transformer 里绝大部分参数量（通常 2/3 以上），而且 FFN 是逐 token 独立计算的，天然适合做"按 token 分流"。Attention 涉及 token 之间的交互，拆开会让信息流动变复杂，工程上不划算。

2. 门控与路由：token 是怎么被分配到专家的

2.1 门控网络的计算过程

门控网络（Gate）本质上是一个小的线性层。假设隐藏维度是 d，专家数量是 N，那门控就是一个 d × N 的权重矩阵 W_g。对一个 token 的隐藏状态 h（维度 d），门控输出是：

logits = h @ W_g # 形状 [N]

这 N 个 logits 就是该 token 对每个专家的"偏好分数"。接下来做 softmax 归一化，得到概率分布：

gates = softmax(logits) # 形状 [N]，所有元素和为 1

然后取 top-k 个最大的概率值，对应的专家就是被激活的专家。假设 k=2，选出的两个专家是 e1 和 e2，对应的门控值是 g1 和 g2。

2.2 Top-k 选择与权重归一化

这里有个细节：top-k 选出来的 g1、g2 加起来不一定等于 1（因为 softmax 是对全部 N 个专家做的）。所以通常要对选中的 k 个门控值再做一次归一化：

g1' = g1 / (g1 + g2) g2' = g2 / (g1 + g2)

最终该 token 的 MoE 层输出是：

output = g1' * Expert_e1(h) + g2' * Expert_e2(h)

这个加权求和是 MoE 计算的核心。每个被选中的专家独立对 h 做一次 FFN 计算，然后按归一化后的门控值加权合并。

2.3 路由的两种粒度：Token Choice 与 Expert Choice

上面讲的是 Token Choice 路由：每个 token 主动选 top-k 个专家。这是最主流的做法，Mixtral、DeepSeek 都用这个。但它有个问题：如果很多 token 都挤到同一个专家，那个专家就会过载，其他专家闲着，负载不均。

另一种是 Expert Choice 路由：每个专家主动选 top-k 个 token。这样能保证每个专家处理的 token 数量完全一致，负载天然均衡。但缺点是 token 可能被"漏掉"（没有专家选它），而且实现上更复杂，主流推理框架支持得少。

实际工程里，Token Choice + 负载均衡损失（Load Balancing Loss）是更常见的组合。负载均衡损失在训练时加到总损失里，惩罚专家负载不均的情况，让门控网络学会把 token 均匀分出去。

2.4 一个容易踩的坑：路由的数值稳定性

门控的 logits 如果数值范围太大，softmax 之后会出现极端分布（某个专家概率接近 1，其他接近 0），导致路由坍缩——所有 token 都往一两个专家跑。训练时靠负载均衡损失缓解，推理时则要注意门控权重的精度。

我在实际调试时遇到过一个问题：某些 MoE 模型在 FP16 下推理，门控的 top-k 选择结果和 FP32 下不一致，导致输出有细微偏差。原因是 FP16 的精度不足以区分两个非常接近的 logits。解决办法是把门控网络的计算强制用 FP32，专家计算保持 FP16/BF16。这个细节在 vLLM 这类框架里通常有专门的配置项控制，值得留意。

3. 专家计算：GEMM 视角下的 MoE 前向过程

3.1 单个专家的 FFN 结构

每个专家就是一个标准的 FFN，通常是两层线性变换加一个激活函数。以 SwiGLU 为例（现在主流 MoE 模型基本都用 SwiGLU）：

# 输入 h: [d] gate = h @ W_gate # [d] -> [d_ff] up = h @ W_up # [d] -> [d_ff] act = SiLU(gate) * up # SwiGLU 激活 out = act @ W_down # [d_ff] -> [d]

SwiGLU 相比传统 ReLU FFN 多了一个门控分支（gate 和 up 两个投影），所以参数量是普通 FFN 的 1.5 倍左右，但效果更好。这也是为什么 MoE 模型的专家虽然"小"，但单个专家的参数量并不算少。

3.2 从单 token 到批量：GEMM 的形态变化

单个 token 过专家，是一堆矩阵向量乘法（GEMV）。但推理时是批量处理的，一批可能有几百上千个 token。这时候计算就变成了 GEMM（通用矩阵乘法）。

关键问题是：不同 token 被路由到不同专家，怎么组织成高效的 GEMM？

最朴素的做法是 for 循环：对每个专家，收集路由到它的所有 token，拼成一个矩阵，做一次 GEMM。假设 8 个专家，就要做 8 次 GEMM。每次 GEMM 的 M 维度（token 数）取决于有多少 token 路由到这个专家。

这里就体现出负载均衡的重要性了。如果负载均匀，8 个专家的 GEMM 规模差不多，GPU 利用率高；如果负载不均，某个专家的 GEMM 很大，其他很小，GPU 就会在等待中浪费。

3.3 专家并行的两种切分方式

当专家数量多到单卡放不下时，就要做专家并行（Expert Parallelism）。有两种切分思路：

切分方式做法优点缺点
专家维度切分不同专家放不同 GPU实现简单，专家独立需要 All-to-All 通信
张量维度切分单个专家内部再切通信模式简单专家数受限于单卡

主流做法是专家维度切分：8 个专家分到 4 张卡，每张卡放 2 个专家。前向时，每张卡先算本地专家能处理的 token，然后把不属于本地专家的 token 通过 All-to-All 通信发到对应卡上。这个 All-to-All 是 MoE 并行的性能瓶颈，通信量正比于 token 数和隐藏维度。

3.4 计算量的定量估算

我们来算一笔账，看看 MoE 的计算量到底怎么估。假设：

  • 隐藏维度 d = 4096
  • 专家 FFN 中间维度 d_ff = 14336（约 3.5 倍）
  • 专家数 N = 8
  • top-k = 2
  • 序列长度 S = 4096，batch = 1

单个 token 过单个专家的 FLOPs（SwiGLU，含 gate/up/down 三个投影）：

gate: 2 * d * d_ff = 2 * 4096 * 14336 ≈ 1.17e8 up: 2 * d * d_ff ≈ 1.17e8 down: 2 * d_ff * d ≈ 1.17e8 合计约 3.5e8 FLOPs

每个 token 激活 2 个专家，所以单 token MoE 层 FLOPs ≈ 7e8。4096 个 token 就是约 2.9e12 FLOPs，也就是 2.9 TFLOPs。如果换成稠密模型，8 个专家全激活，就是 4 倍的量，约 11.6 TFLOPs。稀疏激活省了 4 倍算力，这就是 MoE 的价值。

4. 负载均衡：MoE 训练和推理都绕不开的难题

4.1 负载不均是怎么产生的

门控网络是学出来的，它天然倾向于把 token 送给"表现好"的专家。训练初期，某些专家碰巧初始化得好一点，就会吸引更多 token，然后被训练得更充分，进一步吸引更多 token。这是个正反馈，最后导致赢者通吃——少数专家承担了大部分 token，其他专家几乎不被激活。

这个现象叫路由坍缩（Routing Collapse）。后果很严重：模型实际容量远小于理论容量（因为大部分专家没被训练好），而且推理时负载不均导致 GPU 利用率低。

4.2 负载均衡损失的数学形式

训练时最常用的手段是加一个辅助损失（Auxiliary Loss）。以经典的 Switch Transformer 为例，定义：

  • f_i：路由到专家 i 的 token 比例
  • P_i：专家 i 收到的平均门控概率

负载均衡损失为：

L_balance = N * sum(f_i * P_i)

当负载完全均匀时，f_i = 1/N，P_i = 1/N，L_balance = N * N * (1/N) * (1/N) = 1。当负载极度不均时，这个值会远大于 1。所以最小化这个损失就是在鼓励均匀路由。

这个损失通常乘以一个系数（比如 0.01）加到主损失上。系数太大会影响模型效果（门控被迫均匀，可能牺牲质量），太小起不到均衡作用，需要调。

4.3 推理时的负载均衡：容量因子

训练靠损失，推理靠什么？推理时没法改门控权重，只能靠容量因子（Capacity Factor）做硬限制。每个专家能处理的 token 数上限是：

capacity = capacity_factor * (total_tokens / num_experts)

如果某个专家收到的 token 超过 capacity，多出来的 token 就被"丢弃"（drop），直接走残差连接（不经过专家）。capacity_factor 通常设 1.0 到 1.25。

这里有个权衡：capacity_factor 太小，丢的 token 多，模型质量下降；太大，显存和算力浪费（要预留处理空间）。实际部署时，这个值要根据实际负载分布调。

提示：推理时如果发现输出质量明显下降，先检查是不是 capacity_factor 设太小导致大量 token 被 drop。可以临时调大看看效果。

4.4 无辅助损失的负载均衡

DeepSeek 系列模型提出了一种不需要辅助损失的负载均衡方法：给每个专家加一个可学习的偏置项 b_i，路由时用 logits + b_i 做 top-k 选择。如果某个专家过载，就减小它的 b_i；如果某个专家闲置，就增大它的 b_i。这个偏置只影响路由选择，不参与梯度回传（用 detach），所以不会干扰主损失。

这个方法的巧妙之处在于：它把负载均衡从"损失函数层面"下移到了"路由决策层面"，避免了辅助损失对模型质量的干扰。实测下来，这种方法的负载均衡效果和辅助损失相当，但模型效果更好。

5. 在 vLLM 里观察 MoE 的实际执行

5.1 vLLM 对 MoE 的支持方式

vLLM 从 0.4 版本开始支持 MoE 模型，核心是把专家计算做成了 fused MoE kernel。这个 kernel 把"按专家分组 token → 批量 GEMM → 加权合并"这一整套流程融合在一起，减少 kernel launch 开销和中间张量的显存读写。

具体来说，vLLM 的 fused MoE 会：

  1. 对每个 token 算门控，选 top-k 专家
  2. 统计每个专家收到的 token 数，做前缀和（prefix sum）确定每个专家在缓冲区里的偏移
  3. 把所有 token 按专家顺序重排（permute）到连续缓冲区
  4. 对每个专家做一次 GEMM（可以用 grouped GEMM 进一步优化）
  5. 把结果按原顺序还原（unpermute），加权合并

这套流程里，permute 和 unpermute 是额外的开销，但换来的是 GEMM 的高效执行，整体是划算的。

5.2 观察 MoE 执行的关键指标

部署 MoE 模型时，有几个指标值得盯：

  • 专家负载分布：每个专家处理的 token 数。如果分布很偏，说明负载均衡有问题。
  • Token drop 率：被 capacity 限制丢弃的 token 比例。理想情况接近 0。
  • All-to-All 通信耗时：多卡专家并行时，这个通信占比可能很高。
  • GEMM 利用率：专家 GEMM 的实际算力利用率，负载不均时会很低。

在 vLLM 里可以通过日志和 profiling 工具观察这些。我一般会先跑一个固定长度的请求，看专家负载的直方图，如果某个专家明显偏高，就要考虑调 capacity_factor 或者检查模型本身的门控权重。

5.3 一个实际的调优案例

之前部署一个 8 专家的 MoE 模型，发现吞吐上不去。排查下来是专家负载严重不均：专家 0 和专家 3 承担了 60% 以上的 token，其他专家基本闲置。GPU 利用率只有 30% 左右。

尝试了几个方向：

  1. 调大 capacity_factor 到 1.5，让过载专家能多处理一些，但显存占用上去了，吞吐没明显改善。
  2. 检查输入数据的分布，发现某类请求特别容易触发固定几个专家。做了请求层面的负载分散后，情况好转。
  3. 最终发现是模型本身的门控权重就偏，训练时负载均衡损失系数设小了。这个在推理侧没法根治，只能靠 capacity 和调度缓解。

这个案例说明：MoE 的负载问题很多时候是训练阶段埋下的，推理侧只能缓解不能根治。选模型时，如果要做高吞吐部署，最好先看看它的专家负载分布是否健康。

6. MoE 计算流程的完整梳理与常见误区

6.1 一次前向的完整链路

把前面讲的串起来，一个 token 在 MoE 层的完整计算链路是：

  1. 输入隐藏状态 h（来自上一层 Attention 的输出）
  2. 门控网络算 logits = h @ W_g，softmax 得概率分布
  3. 选 top-k 专家，对选中的门控值归一化
  4. 对每个选中的专家，独立计算 FFN（SwiGLU）
  5. 按归一化门控值加权求和，得到 MoE 层输出
  6. 输出加上残差，进入下一层

批量执行时，步骤 3-5 会做 token 重排和分组 GEMM，这是工程实现的核心。

6.2 几个高频误区

误区一：MoE 省显存。前面强调过，MoE 省的是算力不是显存。所有专家参数都要常驻显存。

误区二：激活参数少所以推理快。激活参数少确实降低 FLOPs，但 MoE 有额外的路由开销、token 重排开销、通信开销。实际加速比往往低于理论值，尤其在 batch 小的时候，路由开销占比高，可能还不如稠密模型。

误区三：专家越多越好。专家多了，单个专家分到的 token 少，训练不充分；而且路由空间变大，负载均衡更难。专家数要和训练数据量、算力匹配。

误区四：top-k 越大越好。top-k 大，激活参数多，计算量上去了，稀疏性红利就没了。一般 k=2 是性价比比较高的选择，k=1 太激进（路由错误没有补偿），k>2 收益递减。

6.3 和其他并行策略的组合

MoE 的专家并行（EP）通常要和张量并行（TP）、流水线并行（PP）、数据并行（DP）组合使用。一个常见的组合是：TP 处理 Attention 和专家内部的张量切分，EP 处理专家维度的切分，DP 做数据并行。

组合时要注意通信开销的叠加。EP 的 All-to-All 和 TP 的 All-Reduce 如果都在关键路径上，延迟会很难看。实践中会把 EP 的通信和计算做 overlap，用异步通信隐藏一部分延迟。

7. 一些实操层面的经验

做 MoE 推理优化这段时间，有几个体会比较深。

第一，先搞清楚模型的专家负载特性再谈优化。不同 MoE 模型的负载分布差异很大，有的很均匀，有的偏得离谱。优化手段（capacity_factor、调度策略）要根据实际分布来定，不能照搬。

第二，小 batch 场景下 MoE 未必划算。batch 小的时候，每个专家分到的 token 少，GEMM 的 M 维度小，GPU 利用率低，路由和重排的开销占比反而高。这种场景下稠密模型可能更稳。MoE 的优势在大 batch、高吞吐场景才明显。

第三，门控的精度别省。前面提过 FP16 门控可能导致路由不一致，这个坑很隐蔽，输出看起来正常但和预期有偏差。门控用 FP32，专家用 BF16，是比较稳妥的组合。

第四，关注 All-to-All 通信。多卡 EP 部署时，All-to-All 往往是瓶颈。可以观察通信和计算的时间占比，如果通信占比超过 30%，就要考虑优化通信（比如用更高效的集合通信库，或者调整专家到卡的映射减少跨卡流量）。

第五，capacity_factor 要动态调。固定值在不同负载下表现不一样。负载轻的时候可以调小省显存，负载重的时候调大防丢 token。有些框架支持动态调整，值得用起来。

MoE 的计算原理本身不复杂，难的是工程实现里的负载均衡、通信优化、精度控制这些细节。把原理搞透，再看框架的实现，很多设计选择就顺理成章了。