1. 从一次显存告急说起MoE 到底解决了什么问题第一次在生产环境里部署一个千亿参数级别的模型时我盯着监控面板上那条几乎垂直向上的显存曲线心里是有点发怵的。单卡放不下多卡切分之后通信开销又高得离谱推理吞吐量低到没法看。后来把模型换成 MoEMixture of Experts混合专家架构的版本同样的硬件条件下显存占用降下来了吞吐量反而上去了。这个反差让我意识到MoE 不是一个“锦上添花”的学术概念而是实打实影响部署成本和推理性能的工程手段。这篇笔记聚焦的是 MoE 的计算原理和过程也就是它在前向传播时到底做了什么、为什么这么设计、每一步的矩阵运算长什么样。我会把 vLLM 这类推理框架里 MoE 的实现逻辑拆开讲涉及 GEMM、Top-k 路由、SwiGLU 激活这些关键环节。如果你正在做大模型推理部署或者想搞清楚 MoE 为什么能“用更少的算力跑更大的模型”这篇内容应该能帮你把这条链路理顺。我尽量不堆公式而是用“数据怎么流动、张量怎么变形”的方式来讲让即使没系统学过并行计算的人也能跟上。先说结论性的认知MoE 的核心思想是把一个大 FFN前馈网络拆成多个小 FFN每次前向只激活其中一小部分。传统稠密模型里每个 token 都要过完整的 FFN参数量越大算力消耗越大MoE 里每个 token 只被路由到少数几个专家比如 2 个其余专家不参与计算。这样一来模型总参数量可以做得很大但单个 token 的实际计算量FLOPs保持在一个可控水平。这就是所谓的“稀疏激活”。理解这一点之后后面所有的计算细节——路由、分发、GEMM、合并——都是为了在 GPU 上高效地实现这个“稀疏激活”。而难点恰恰在于稀疏是逻辑上的硬件执行是密集的怎么把逻辑稀疏映射到硬件密集计算上是整个 MoE 推理优化的核心矛盾。2. MoE 的整体计算流程拆解2.1 一个 token 在 MoE 层里经历了什么要理解 MoE 的计算过程最直观的方式是跟着一个 token 走一遍。假设我们有一个 MoE 层包含 N 个专家expert每个专家本身是一个 FFN通常结构是Linear - 激活函数 - Linear。此外还有一个门控网络gate / router它是一个小的线性层输出维度等于专家数量 N。一个 token 的隐藏状态hidden state进入 MoE 层后流程大致如下门控计算token 的 hidden state 经过 gate 线性层得到一个长度为 N 的分数向量logits。Top-k 选择从这 N 个分数里选出最高的 k 个常见 k1 或 k2得到被选中的专家索引和对应的权重。权重归一化把这 k 个分数做 softmax或直接归一化得到每个被选中专家的权重系数。专家计算token 被分发到对应的 k 个专家各自过一遍 FFN得到 k 个输出。加权合并把这 k 个输出按权重加权求和得到该 token 在 MoE 层的最终输出。这个过程听起来不复杂但真正在 GPU 上跑的时候问题就来了一个 batch 里有成百上千个 token每个 token 选中的专家可能都不一样。如果按 token 逐个去调用专家GPU 的并行能力完全发挥不出来效率极低。所以实际实现里一定是按专家维度重新组织数据把“哪些 token 该去哪个专家”这件事变成一次批量矩阵运算。2.2 为什么是 Top-k 而不是 Top-1很多人会问既然要稀疏为什么不干脆每个 token 只选一个专家Top-1这样计算量最小。原因在于路由的稳定性和表达能力。Top-1 路由的问题在于门控网络的判断一旦有偏差token 就被送到一个不太合适的专家那里而且没有任何补偿机制。训练时这种硬选择还会导致梯度只流向被选中的专家其他专家得不到更新出现“专家坍缩”——少数专家被反复训练多数专家形同虚设。Top-2或更大的 k相当于给每个 token 一个“备选方案”两个专家的输出加权组合既提升了表达能力也让梯度能更均匀地分配到多个专家上。当然 k 越大计算量越大所以实践中 k2 是最常见的折中。像一些较新的架构也会用 k1 配合更精细的负载均衡策略但那是另一套思路了。2.3 专家并行的物理布局逻辑上我们有 N 个专家物理上这些专家要分布在多张 GPU 上。常见的布局是专家并行Expert Parallelism把 N 个专家平均分到 E 张卡上每张卡负责 N/E 个专家。这样做的好处是每张卡只需要存一部分专家的参数显存压力被分摊了。但代价是通信token 经过路由后可能需要被发送到其他卡上的专家去计算算完再发回来。这个“发送-计算-回收”的过程就是 MoE 推理里通信开销的主要来源。在 vLLM 这类框架里专家并行的通信通常用 all-to-all 或者 all-gather reduce 的方式实现。具体用哪种取决于专家数量和 token 分布的均衡程度。如果 token 分布很均匀all-to-all 效率高如果分布倾斜严重某些卡会成为瓶颈这时候就需要负载均衡策略来兜底。3. 核心计算环节从 GEMM 到 SwiGLU3.1 门控网络与 Top-k 的矩阵视角门控网络本质上就是一个线性层权重矩阵形状是[hidden_dim, num_experts]。假设一个 batch 有 T 个 tokenhidden_dim 为 H那么门控计算就是一次 GEMMgate_logits hidden_states gate_weight # [T, H] [H, N] - [T, N]这一步的计算量不大因为 N专家数通常远小于 H。但它是整个 MoE 的“决策中心”后续所有数据流动都依赖它的输出。得到[T, N]的 logits 后做 Top-k 选择。这里有个实现细节Top-k 在 GPU 上不是简单的排序而是用torch.topk这类算子返回 values 和 indices 两个张量。values 是选中的分数indices 是专家编号。接着对 values 做 softmax 归一化。注意这里是对选中的 k 个分数做 softmax而不是对全部 N 个。这个细节很重要因为它保证了权重之和为 1且只关注被选中的专家。# 伪代码示意 gate_logits hidden_states gate_weight # [T, N] topk_values, topk_indices torch.topk(gate_logits, k2, dim-1) # [T, 2] topk_weights torch.softmax(topk_values, dim-1) # [T, 2]3.2 Token 分发从逻辑稀疏到物理密集拿到topk_indices之后下一步是把 token 送到对应专家那里。如果直接按 token 循环效率会惨不忍睹。实际实现里用的是按专家分组的思路。具体做法是遍历所有 token 的 topk_indices统计每个专家被多少个 token 选中然后根据这个统计结果把 token 重新排列成一个“按专家聚集”的顺序。这个过程叫permutation重排在 vLLM 里通常用moe_permute之类的算子实现。重排之后数据变成这样前一段是专家 0 要处理的 token接着是专家 1 的以此类推。每个专家拿到的 token 数量可能不同这就是所谓的“变长”问题。为了能用 GEMM 批量计算通常会把每个专家的 token 补齐到相同长度padding或者用 grouped GEMM 来处理变长输入。提示padding 会带来无效计算浪费算力。所以更先进的做法是用 grouped GEMM 或者 variable-length GEMM让每个专家只算自己实际拿到的 token。vLLM 在新版本里对这块做了不少优化。3.3 专家内部的 FFN 与 SwiGLU 激活每个专家本身是一个 FFN。传统 FFN 是两层线性加一个 ReLU但现代大模型里更常见的是SwiGLU。SwiGLU 的结构是FFN(x) (Swish(x W_gate) * (x W_up)) W_down也就是说它有三个权重矩阵W_gate、W_up、W_down。输入 x 分别过 W_gate 和 W_up前者经过 Swish 激活也叫 SiLU然后两者逐元素相乘最后过 W_down 降维。为什么用 SwiGLU 而不是 ReLU核心原因是门控机制带来的表达能力强。Swish 本身是平滑的梯度性质比 ReLU 好而gate * up这种乘法结构相当于一个软门控让网络能动态调节信息流。实测下来同样参数量下 SwiGLU 的效果通常优于 ReLU FFN这也是为什么 LLaMA、Qwen 这些模型都用了它。在 MoE 里每个专家都用 SwiGLU意味着每个专家有三个权重矩阵。如果专家数量是 N那总参数量就是 N 倍的单专家参数量。这也是 MoE 模型参数量能做得很大的原因——参数量大但每次只激活少数专家实际计算量可控。3.4 加权合并与输出每个专家算完之后得到的是“按专家聚集”的输出。接下来要把它还原回原来的 token 顺序这个过程叫unpermute逆重排。还原之后每个 token 有 k 个输出来自 k 个专家按之前算好的topk_weights加权求和就得到 MoE 层的最终输出。# 伪代码示意 expert_outputs grouped_gemm(permuted_inputs, expert_weights) # 按专家分组计算 unpermuted unpermute(expert_outputs, reverse_indices) # 还原顺序 final_output (unpermuted * topk_weights.unsqueeze(-1)).sum(dim1) # 加权合并这一步的加权求和看起来简单但在实现上要注意数值精度。因为 k 个专家的输出可能量级差异较大直接相加可能引入误差。有些实现会用 float32 做累加再转回原来的精度。4. 实操中的关键参数与性能调优4.1 专家数量、Top-k 与隐藏维度的权衡MoE 的性能和这几个参数强相关专家数量 N、Top-k 的 k、以及每个专家的隐藏维度。它们之间有一个隐性的平衡关系。专家数量越多模型总参数量越大但每个 token 激活的参数量不变因为 k 固定。这意味着显存占用增加但计算量不增加。所以 N 的选择主要受显存限制。如果显存充足可以适当增大 N 来提升模型容量。k 的选择直接影响计算量。k2 的计算量大约是 k1 的两倍忽略门控和通信开销。所以 k 增大吞吐量会下降。实践中 k2 是效果和性能的常见折中。每个专家的隐藏维度通常比稠密 FFN 小。比如稠密 FFN 的中间维度是 4HMoE 里每个专家的中间维度可能是 H 或 2H。这样即使 N 很大单个专家的计算量也不至于爆炸。4.2 负载均衡MoE 绕不开的工程难题MoE 最让人头疼的问题之一是负载不均衡。理想情况下每个专家被选中的概率差不多这样每张卡的负载均匀通信和计算都能打满。但实际训练出来的模型往往会出现“热门专家”和“冷门专家”——某些专家被大量 token 选中另一些几乎没人用。负载不均衡的后果很直接热门专家所在的卡成为瓶颈其他卡在等它整体吞吐量被拖垮。在推理场景下这个问题尤其明显因为推理时没有训练时的负载均衡损失来约束路由。解决思路主要有两类。一类是在训练阶段引入辅助损失auxiliary loss惩罚不均衡的路由分布让模型学会均匀分配。另一类是在推理阶段做动态调度比如把热门专家的计算拆分到多张卡上或者用 capacity factor 限制每个专家最多处理多少 token超出的 token 直接跳过drop。提示capacity factor 是个双刃剑。设得太小token 被丢弃影响效果设得太大padding 增多浪费算力。一般从 1.25 开始调根据实际负载分布微调。4.3 通信开销的优化思路专家并行带来的通信开销是 MoE 推理性能的关键变量。优化方向主要有几个第一减少通信数据量。token 的 hidden state 维度是 H如果 H 很大传输成本就高。有些实现会在分发前先做降维或者用更紧凑的数据格式。第二重叠通信和计算。在等待其他卡返回结果的同时先算自己能算的部分。这需要框架层面的调度支持vLLM 和 SGLang 在这方面都有持续优化。第三选择合适的并行策略。如果专家数量不多可以考虑把专家放在同一张卡上避免跨卡通信如果专家很多专家并行不可避免那就尽量让 token 分布均匀减少 all-to-all 的等待时间。5. 常见问题排查与避坑记录5.1 推理时显存占用异常有一次部署 MoE 模型发现显存占用比预期高不少。排查下来问题出在专家参数的加载方式上。默认情况下框架可能把所有权重都加载到每张卡上而不是按专家并行切分。这样虽然省了通信但显存直接翻倍。解决办法是确认框架的专家并行配置是否正确。在 vLLM 里需要显式设置专家并行度expert parallel size确保每个专家只加载到对应的卡上。如果配置错了显存占用会明显偏高而且计算也不会走专家并行的路径。5.2 Top-k 路由结果不稳定另一个常见问题是路由结果抖动。同样的输入两次推理选中的专家不一样。这通常是因为门控输出的数值精度问题或者 Top-k 在分数接近时选择了不同的专家。排查思路先看门控输出的 logits 是否有异常值比如 NaN 或极大值再检查 Top-k 的实现是否用了确定性算法。有些 GPU 算子在分数相同时返回的索引不确定需要设置确定性模式。5.3 吞吐量上不去如果 MoE 推理的吞吐量明显低于预期可以从几个方向排查排查项可能原因解决方向负载分布热门专家成为瓶颈检查路由分布考虑 capacity factor通信开销all-to-all 等待时间长优化并行策略重叠通信计算Padding 浪费变长处理不当用 grouped GEMM 替代 padding专家并行度配置不合理调整 expert parallel size批大小太小导致 GPU 利用率低增大 batch或做 continuous batching这张表基本覆盖了我遇到过的大部分性能问题。实际排查时建议先用 profiling 工具定位瓶颈在计算还是通信再针对性优化。5.4 数值精度问题MoE 的加权合并环节对精度比较敏感。如果 k 个专家的输出量级差异大用低精度累加可能出问题。我遇到过用 fp16 累加导致输出偏差的情况改成 fp32 累加后恢复正常。另外SwiGLU 里的 Swish 激活在极端输入下可能产生较大数值需要注意溢出。一般框架会做数值稳定处理但自己实现时要注意。6. 我对 MoE 推理优化的一点个人体会折腾了这么多 MoE 部署的案例我最大的体会是MoE 的性能瓶颈往往不在计算而在数据流动。专家本身的计算量其实不大真正耗时的是 token 的分发、重排、通信和合并。所以优化 MoE 推理重点应该放在减少数据搬运、提升并行度、均衡负载上而不是一味优化单个专家的 GEMM。另一个体会是MoE 的参数配置没有“万能解”。专家数量、Top-k、capacity factor 这些参数需要根据具体的硬件配置、模型结构和业务场景来调。我见过有人直接照搬别人的配置结果性能差了一大截。建议的做法是先用小规模实验摸清负载分布和瓶颈位置再逐步调参。最后分享一个小技巧在调试 MoE 路由时可以把每个专家的被选中次数打印出来画成直方图。如果分布明显倾斜说明负载均衡有问题需要调整辅助损失或 capacity factor。这个简单的可视化往往能快速定位问题比盲目调参高效得多。