vLLM 源码完全拆解（八）：MoE 层的前向执行——Router、FusedMoE Kernel 与通信融合

三层架构总览#

FusedMoEFactory 接收约 40 个参数（num_experts、top_k、hidden_size、intermediate_size、use_grouped_topk、scoring_func 等），组装出三个核心对象：

1. Router（FusedMoERouter）：负责从 router_logits 中选出每个 token 的 top_k 个专家，返回 (topk_weights, topk_ids)。通过 router_factory.py 中的 create_fused_moe_router() 根据参数选择具体子类：FusedTopKRouter、GroupedTopKRouter、FusedTopKBiasRouter、ZeroExpertRouter 等。

2. RoutedExperts：管理所有专家权重参数（gate_proj、up_proj、down_proj），持有 quant_method——即 MoE 的实际计算实现。quant_method 可以是 Monolithic（一个 kernel 搞定路由+计算+激活，如 TRT-LLM 的 fused MoE kernel）或 Modular（先 router 选专家，再调 kernel 计算）。

3. MoERunner：编排器，把 Router 和 RoutedExperts 组合成完整的前向流程。它持有 gate（router_logits 的线性层）、shared_experts（可选，共享专家的独立计算）、routed_input_transform / routed_output_transform（Latent MoE 的纬度变换），并负责 dispatch/combine 通信与 All-Reduce 时机。

FusedMoEFactory 的返回类型就是 MoERunner——模型代码拿到它后，在 forward 中调用 runner.forward(hidden_states, router_logits) 即可。

从哪里被调用#

回顾前文，在 GPUModelRunner 的模型执行中，每个 MoE 层被封装成了 FusedMoEFactory 的返回值。在模型的 forward 方法中，对每个 MoE 层：

1
# 伪代码：模型 forward 中的 MoE 层调用
2
router_logits = self.gate(hidden_states)  # gate 计算 router logits
3
hidden_states = self.moe_runner.forward(
4
    hidden_states,
5
    router_logits,
6
)

而 MoERunner.forward() 内部通过一个自定义 PyTorch op（torch.ops.vllm.moe_forward）进入 _forward_impl()，完成从 dispatch 到 combine 的全部工作。这个自定义 op 的设计动机是让 Dynamo / torch.compile 能把 MoE 层当作一个不透明的 kernel 来看待，避免 trace 内部的复杂控制流。

FusedMoEConfig：所有配置的汇聚点#

在工厂函数中，各种分散的参数被汇聚到一个 FusedMoEConfig 对象中：

1
moe_config = FusedMoEConfig(
2
    num_experts=global_num_experts,       # 全局专家数（含冗余）
3
    experts_per_token=top_k,              # 每个 token 激活的专家数
4
    hidden_dim=hidden_size,               # 输入 hidden size
5
    intermediate_size=intermediate_size,  # 专家中间层尺寸
6
    num_local_experts=expert_map_manager.local_num_experts,  # 本地专家数（EP 后）
7
    num_logical_experts=logical_num_experts,  # 逻辑专家数
8
    moe_parallel_config=moe_parallel_config,  # TP/EP/DP 配置
9
    ...
10
)

这个 config 随后被传给 RoutedExperts 和 MoERunner，是整个 MoE 层所有模块决策的单一事实来源。

Router 的子类选择逻辑#

create_fused_moe_router() 的优先级链（按先后顺序）：

1. 如果设置了环境变量 VLLM_MOE_ROUTING_SIMULATION_STRATEGY → RoutingSimulatorRouter（用于模拟路由策略，非生产路径）
2. 如果 zero_expert_type is not None → ZeroExpertRouter（处理 identity zero expert，用于 EPLB 等场景）
3. 如果 use_grouped_topk=True → GroupedTopKRouter（DeepSeek-V3 类型模型的最常见选择）
4. 如果 custom_routing_function is not None → CustomRoutingRouter
5. 如果 e_score_correction_bias is not None 或 hash_indices_table is not None → FusedTopKBiasRouter（DeepSeek-V4 的 hash routing + bias 路径）
6. 如果 num_fused_shared_experts > 0（ROCm AITER）→ AiterSharedRoutedFusedMoERouter
7. 默认 → FusedTopKRouter（标准 softmax top-k，用于 Mixtral 等）

生产环境中最常见的是 3（GroupedTopKRouter）和 7（FusedTopKRouter）。

FusedTopKRouter：标准 softmax Top-K#

这是最简单的路由方式，适用于 Mixtral 8x7B 这类没有专家分组的模型：

1
# FusedTopKRouter._select_experts() 的核心逻辑（伪代码）
2
def _select_experts(self, hidden_states, router_logits):
3
    # 1. 对 router_logits 做 softmax（如果 scoring_func == "softmax"）
4
    if self.scoring_func == "softmax":
5
        routing_weights = router_logits.softmax(dim=-1)
6
    elif self.scoring_func == "sigmoid":
7
        routing_weights = router_logits.sigmoid()
8
    # sqrtsoftplus 等特殊评分函数也在此处理
9

10
    # 2. 取 top_k 个最大值的专家索引和权重
11
    topk_weights, topk_ids = torch.topk(routing_weights, self.top_k, dim=-1)
12

13
    # 3. 可选：重新归一化 top_k 权重，使它们和为 1
14
    if self.renormalize:
15
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
16

17
    return topk_weights, topk_ids

关键点：

• scoring_func 不只是 softmax 和 sigmoid。DeepSeek-V4 使用了 sqrtsoftplus——对 logits 先做 sqrt(softplus(x)) 再 topk，实验表明这种评分函数在专家数极大时更稳定。
• renormalize 控制了 top_k 权重是否重新归一化。大部分模型（DeepSeek-V3、Mixtral）都设为 True。
• routed_scaling_factor 是 DeepSeek-V3 引入的机制——选出的专家权重先乘以这个 factor，论文中设为 1.0 / top_k 的某种变体，实际效果是调整 routed expert 输出与 shared expert（如果有的话）的相对比例。

GroupedTopKRouter：分组的专家选择#

这是 DeepSeek-V3 式 MoE 的核心机制。在 V3 中，256 个专家被分成 8 组，每组 32 个专家。路由不是从全部 256 个中选 top-8，而是先从 8 组中选出 top-4 组（topk_group=4），再从这 4 组中各自选出 top-2 专家（top_k=8 / 4 = 2），合起来就是 8 个专家。

为什么要分组？因为 DeepSeek-V3 的专家之间有”设备亲缘性”（affinity）——同一组的专家放在同一个 GPU 上，这样 All-to-All 通信只需要在组间发生，组内的专家选择不产生额外跨 GPU 通信。分组路由天然对齐了硬件拓扑。

GroupedTopKRouter._select_experts() 的流程：

1. 将 router_logits reshape 为 [num_tokens, num_expert_group, experts_per_group]
2. 在每个 group 内做 scoring（softmax/sigmoid）
3. 对每个 group 求和得到 group scores → 选出 topk_group 个最好的组
4. 在选中的组内，各自做 top-k per group → 总共得到 top_k 个专家
5. 对选出的专家权重做 renormalize

1
# GroupedTopKRouter 的核心逻辑（简化伪代码）
2
def _select_experts(self, hidden_states, router_logits):
3
    # router_logits: [num_tokens, num_experts]
4
    # reshape: [num_tokens, num_expert_group, experts_per_group]
5
    scores = router_logits.view(-1, self.num_expert_group, self.experts_per_group)
6

7
    # 组内 softmax
8
    if self.scoring_func == "softmax":
9
        group_scores = scores.softmax(dim=-1)  # per-group softmax
10

11
    # 每组求和作为组分数
12
    group_total = group_scores.sum(dim=-1)  # [num_tokens, num_expert_group]
13

14
    # 选出 topk_group 个组
15
    _, top_group_idx = torch.topk(group_total, self.topk_group, dim=-1)
16

17
    # 只在选中的组内取 top-k per group
18
    # 这里 mask 掉未选中组的分数，再做 per-group topk
19
    ...
20
    return topk_weights, topk_ids

注意：DeepSeek-V3 的每组内 per-group softmax 然后取组内 top-k，这和先全局 softmax 再取全局 top-k 是不同的——前者保证了组间比较的公平性（每组先内部归一化），后者则是直接比原始 logits 的绝对大小。

topk_weights 和 topk_ids 的语义#

路由完成后，返回两个张量：

• topk_weights：[num_tokens, top_k]，float，每个 token 对每个选中专家的权重。如果 renormalize=True，这 top_k 个权重之和为 1。
• topk_ids：[num_tokens, top_k]，int32/int64，每个 token 选中专家的全局 ID。

这两个张量随后传入 RoutedExperts.forward_modular() 或直接被 Monolithic kernel 消费。

专家权重如何存储#

RoutedExperts 是 torch.nn.Module 的子类，持有所有专家的权重参数。对于标准的 gated MoE（如 SwiGLU 激活）：

• gate_up_proj（或 w13）：[num_local_experts, intermediate_size * 2, hidden_size]——gate 和 up 两个投影合并在一起。之所以合并，是因为它们共享输入，合并后在 kernel 内部只需一次读取 hidden_states。
• down_proj（或 w2）：[num_local_experts, hidden_size, intermediate_size]——从中间维度投影回 hidden_size。

这里 num_local_experts 是经过 Expert Parallelism 切分后，当前 GPU 实际持有的专家数。如果 256 个专家分布在 8 个 GPU 上做 EP，每个 GPU 持有 256/8 = 32 个专家。

MergedColumnParallel 与 RowParallelLinear#

对于张量并行（TP），vLLM 会把专家的权重沿 hidden_size 或 intermediate_size 维度切分：

• gate_up_proj 使用 MergedColumnParallelLinear，沿 output 维度切分（列并行），输入是完整的 hidden_states，输出被 TP group 内各 GPU 分片持有。
• down_proj 使用 RowParallelLinear，沿 input 维度切分（行并行），输入来自上一个 All-Reduce（或 TP group 内的 partial output），输出需要 All-Reduce。

MoE 层中这两个并行策略的交互比较复杂：EP 把专家按数量切分到不同 GPU，TP 把单个专家的权重矩阵按维度切分到同组 GPU。FusedMoEParallelConfig 记录了所有这些并行维度：

1
class FusedMoEParallelConfig:
2
    tp_size: int          # 张量并行度
3
    ep_size: int          # 专家并行度
4
    dp_size: int          # 数据并行度（含弹性 EP 的最大 DP size）
5
    sp_size: int          # 序列并行度
6
    use_ep: bool          # 是否开启 EP
7
    use_all2all_kernels: bool  # 是否使用融合的 all-to-all kernel

量化与 MoE Kernel 选择#

RoutedExperts 的核心属性是 quant_method——决定 MoE 计算以什么精度、什么 kernel 执行。常见选择：

• 未量化（FP16/BF16）：使用 fused_moe Triton kernel（或 CUDA kernel），直接做 FP16 的 batched GEMM。
• FP8 量化：使用 FP8 fused MoE kernel，参考 DeepSeek-V3 的 FP8 混合精度方案。
• INT8 / INT4 量化：使用对应的量化 kernel。
• TRT-LLM 后端：使用 TensorRT-LLM 的 MoE kernel，通常是 Monolithic 模式（路由和计算合并在一个 kernel 中）。

quant_method 还决定了 modular vs monolithic 的执行路径：

• Modular：Router 先出 topk_weights/topk_ids，再传入 forward_modular(x, topk_weights, topk_ids)。
• Monolithic：Router 不单独执行，forward_monolithic(x, router_logits) 在一个 kernel 中完成路由 softmax + top-k + batched GEMM + 激活函数。

Monolithic 的优势是避免 router 输出的显存读写往返，适合 TRT-LLM 这类高度优化的闭源 kernel。Modular 的优势是灵活——可以在 router 后插入 EPLB 重映射、shared expert 融合等逻辑。vLLM 目前两者都支持，但新架构倾向于 Modular。

调用链：forward → _forward_impl#

1
MoERunner.forward(hidden_states, router_logits)
2
  ├─ apply_routed_input_transform()     # Latent MoE: 压缩输入维度
3
  ├─ _maybe_pad_hidden_states()         # 对齐 hidden_dim
4
  ├─ self._forward_entry(...)           # 自定义 PyTorch op → _forward_impl
5
  │   └─ _forward_impl()
6
  │       ├─ _maybe_fuse_gate_weights() # 首次调用时合并 gate 权重
7
  │       ├─ shared_experts.maybe_forward_async()  # 异步启动 shared experts
8
  │       ├─ gate(hidden_states) → router_logits  # （如果 gate 在 runner 内部）
9
  │       ├─ sequence_parallel_context()
10
  │       ├─ _maybe_dispatch()          # EP/DP dispatch
11
  │       ├─ _apply_quant_method()      # 核心：路由 + fused MoE kernel
12
  │       │   ├─ router.select_experts()
13
  │       │   ├─ routed_experts.forward_modular() / forward_monolithic()
14
  │       │   └─ shared_experts.wait()  # 等待异步 shared experts 完成
15
  │       └─ _maybe_combine()           # EP/DP combine
16
  ├─ 截断输出                           # 去掉 padding
17
  ├─ _maybe_reduce_routed_output_before_transform()  # Latent MoE 提前 all-reduce
18
  ├─ _maybe_reduce_shared_expert_output()
19
  ├─ _maybe_apply_routed_scale_to_output()
20
  ├─ apply_routed_output_transform()    # Latent MoE: 恢复输出维度
21
  ├─ shared + routed 输出相加
22
  ├─ _maybe_reduce_final_output()       # 最终 All-Reduce
23
  └─ _maybe_add_zero_expert_output()

这个调用链看起来很长，但大部分步骤在标准 MoE（无 Latent、无 shared expert）下都是 no-op。对于最简情况（Mixtral 8x7B，TP=1，EP=1），实际执行的只有 dispatch → router → forward_modular → combine → final all-reduce。

Gate 计算与 FSE（Fused Shared Expert gate）#

MoERunner 可以选择在内部持有 gate（router logits 的计算层）。如果模型的 gate 权重已经加载到 runner 上，_forward_impl 会直接在这里做 gate 的前向：

1
if self.gate is not None:
2
    if self._fse_fuse_gate:
3
        self._maybe_fuse_gate_weights()
4
        router_logits = dispatch_unquantized_gemm()(
5
            self, hidden_states, self._combined_gate_weight, None
6
        )
7
    else:
8
        router_logits, _ = self.gate(hidden_states)

FSE（Fused Shared Expert）模式下，routed gate 和 shared expert gate 的权重矩阵被拼成一个 [num_experts + num_shared, hidden] 的大矩阵，一次 GEMM 算出所有 logits。TopK kernel 再从中分离出 routed topk 和 shared expert 的激活值（通常用 sigmoid 而非 softmax），省掉一次单独的 shared expert gate 前向。

Dispatch：EP 下的 token 分发#

当 ep_size > 1 时，每个 token 需要被送到持有其选中专家的 GPU 上去计算。传统做法是通过 All-to-All 通信来交换 token——每个 GPU 把自己的 token 中属于其他 GPU 专家的部分发送过去，同时接收其他 GPU 发来的、属于自己专家的 token。

vLLM 的 dispatch 有两种路径：

1. MoE Kernel 内部 dispatch（新路径）：如果 quant_method.supports_internal_mk，则 dispatch/combine 完全在 fused MoE kernel 内部完成。这是 vLLM 新架构的目标状态——kernel 直接通过 NVSHMEM 或类似机制做 RDMA，不经过 host 端显存中转。

2. Naive dispatch（旧路径，过渡期）：get_ep_group().dispatch_router_logits() 负责在 EP group 内分发 hidden_states 和 router_logits。实际上就是 All-to-All 通信。代码在 _maybe_dispatch() 中：

1
def _maybe_dispatch(self, hidden_states, router_logits):
2
    if self.do_naive_dispatch_combine:
3
        result = get_ep_group().dispatch_router_logits(
4
            hidden_states,
5
            router_logits,
6
            self.moe_config.is_sequence_parallel,
7
        )
8
        hidden_states, router_logits = result
9

10
    # PCP (Pipeline Context Parallelism) 的 all_gather
11
    if self.moe_config.pcp_size > 1 and not use_all2all_kernels:
12
        hidden_states = get_pcp_group().all_gather(hidden_states, dim=0)
13
        router_logits = get_pcp_group().all_gather(router_logits, dim=0)
14

15
    return hidden_states, router_logits

do_naive_dispatch_combine 的条件是 dp_size > 1 or is_sequence_parallel（且 kernel 不支持 internal MK）。这意味着在 DP + EP 混合并行和序列并行场景，dispatch 仍然走旧路径。

_apply_quant_method：路由与 kernel 执行#

这是 MoE 层的核心。代码很紧凑：

1
def _apply_quant_method(self, hidden_states, router_logits,
2
                         shared_experts_input, shared_experts_overlapping=False):
3
    # 1. 先执行 shared experts（如果不需要 overlap）
4
    self._maybe_apply_shared_experts(shared_experts_input, SharedExpertsOrder.NO_OVERLAP)
5

6
    # 2. 路由 + 计算
7
    if self.routed_experts.quant_method.is_monolithic:
8
        # Monolithic: router_logits 直接传入 kernel
9
        fused_out = self.routed_experts.forward_monolithic(
10
            x=hidden_states, router_logits=router_logits
11
        )
12
    else:
13
        # Modular: 先 router 选专家，再计算
14
        topk_weights, topk_ids = self.router.select_experts(
15
            hidden_states=hidden_states,
16
            router_logits=router_logits,
17
            topk_indices_dtype=self._quant_method.topk_indices_dtype,
18
        )
19
        fused_out = self.routed_experts.forward_modular(
20
            x=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids,
21
            shared_experts=self._shared_experts,
22
            shared_experts_input=shared_experts_input,
23
        )
24

25
    # 3. 等待异步 shared experts
26
    if shared_experts_overlapping:
27
        self._shared_experts.wait()
28

29
    return shared_output, fused_out

这里的关键设计决策：

• Shared experts 提前执行：在 router 执行之前或同时启动，这样 shared experts 的计算可以和 routed experts 的路由+计算 overlap。如果 shared_experts_overlapping=True，shared experts 在 _forward_impl 的最开始就用单独 CUDA stream 启动了。
• Monolithic vs Modular：取决于 quant_method.is_monolithic，由 kernel 能力决定。TRT-LLM 的 MoE kernel 是 monolithic 的，vLLM 自己的 Triton/CUDA kernel 是 modular 的。

输出处理与 All-Reduce 时机#

MoE 输出处理有三个 All-Reduce 要点（互斥，不会同时触发）：

1. Kernel 内部已 reduce（_fused_output_is_reduced=True）：fused kernel 做完 expert 计算后内部做了 All-Reduce，fused_output 已经是完整结果。此时需要单独对 shared_output 做 All-Reduce（_maybe_reduce_shared_expert_output），然后跳过最终 All-Reduce。

2. Latent MoE 提前 reduce（_maybe_reduce_routed_output_before_transform）：当 MoE 在压缩的 latent 空间做计算（如 NemotronH），输出需要经过可能有非线性的 transform 恢复到完整维度。如果 fused_output 还没 reduce，必须在 transform 之前先 All-Reduce，否则非线性和 partial sum 不交换。

3. 最终 All-Reduce（_maybe_reduce_final_output）：最常见的路径。fused_output 和 shared_output 都还没 reduce，把它们相加后再做一次 All-Reduce。

Shared Experts 的 Overlap 机制#

SharedExperts 类封装了共享专家的异步执行。核心方法是 maybe_forward_async()：在当前 CUDA stream 上启动 shared experts 的前向，返回是否成功启动了 overlap。

如果 overlap 成功，router + routed experts 在主 stream 上执行，shared experts 在另一个 stream 上执行，最后通过 wait() 同步两个 stream 的结果。这种设计让 shared experts 的 compute-heavy 部分（两个 GEMM + 激活）和 router 的 score compute + topk + dispatch 完全重叠。

DeepSeek-V2 的稀疏 MoE 架构（来源：DeepSeekMoE 论文）。图中展示了 Shared Expert（共享专家）和 Routed Expert（路由专家）的分离——前者对所有 token 一律执行，后者只对路由选中的 token 执行。vLLM 的 SharedExperts 异步 overlap 正是利用了这两条路径的天然独立性。

Custom PyTorch Op（moe_forward）#

MoERunner 的 forward 入口被封装为 torch.ops.vllm.moe_forward 自定义 op。这对 torch.compile/Dynamo 非常关键——Dynamo 无法 trace 复杂控制流（条件 dispatch、EPLB 映射、多种 kernel 路径），但面对一个 opaque custom op 时，它可以把它当作单个 kernel 来调度，不影响整个图的编译。

自定义 op 的 fake 实现也简洁：直接返回和输入同 dtype、同 shape 的空张量。这让 Dynamo 在 fake tensor 阶段就能推断整个图的 shape 和 dtype，提前完成编译。

CUDA Graph 中的 MoE 层#

vLLM V1 的 CUDA Graph 对 MoE 有多层支持。FusedMoEConfig 中记录了 max_capture_size——CUDA Graph 能捕获的最大 batch size。当 batch size 在 capture 范围内，整个 MoE 层（包括路由、kernel 计算、通信）被重放而非重新 launch。当 batch size 超出时，退回到 eager 执行。

MoE 层在 CUDA Graph 中的特殊之处在于路由输出的大小不固定——不同 step 的 topk_ids 可能指向不同的专家组合，导致后续 kernel 的 launch config 变化。vLLM 通过 _routing_replay_out 机制处理这个问题：在 graph capture 时预分配一个最大尺寸的 buffer，每次重放时 router 把结果写入这个 buffer，kernel 从 buffer 读取。

通信与计算的重叠#

在 EP 场景中，dispatch All-to-All 和 combine All-to-All 是必须的两次通信。vLLM 的融合 All-to-All kernel（use_all2all_kernels=True 时）尝试把 dispatch 通信和路由计算重叠：

1. Router 先完成 softmax + topk（这是纯计算，不依赖通信）
2. All-to-All dispatch 把 token 送到专家所在 GPU
3. 专家计算（GEMM+激活）在本地 GPU 上执行
4. All-to-All combine 把结果送回原始 GPU

步骤 1 和步骤 2 可以部分重叠——在 topk 完成后立即开始 dispatch，不需要等所有 token 的路由都完成。但这需要精细的 stream 管理和 barrier 插入，当前 vLLM 主要在 DeepEP 后端（DeepSeek 开源的通信库，本站已有DeepEP完全拆解文章）实现这种细粒度重叠。