vllm源码剖析14-vLLM 分布式推理-专家并行EP

文章目录

    • [一 MOE + EP 概述](#一 MOE + EP 概述)
      • [1.1 MOE 出现的背景](#1.1 MOE 出现的背景)
      • [1.2 MOE 的基本计算流程](#1.2 MOE 的基本计算流程)
      • [1.3 EP 概述](#1.3 EP 概述)
      • [1.4 MOE + EP 的基本计算过程](#1.4 MOE + EP 的基本计算过程)
      • [1.5 MoE + EP 推理场景实例](#1.5 MoE + EP 推理场景实例)
        • [1.5.1 推理流程简单拆解](#1.5.1 推理流程简单拆解)
    • [二 AlltoAll 通信原理](#二 AlltoAll 通信原理)
      • [2.1 Scatter & ReduceScatter](#2.1 Scatter & ReduceScatter)
      • [2.2 ReduceScatter](#2.2 ReduceScatter)
      • [2.3 Alltoall](#2.3 Alltoall)
        • [1. 准备阶段:切分](#1. 准备阶段:切分)
        • [2. 传输阶段:路由与交换](#2. 传输阶段:路由与交换)
        • [3. 结果阶段:拼接 (Concat)](#3. 结果阶段:拼接 (Concat))
        • [4. 代码示例](#4. 代码示例)
    • [三 MoE 层架构和推理流程](#三 MoE 层架构和推理流程)
    • [四 Dispatch-Compute-Combine 流程](#四 Dispatch-Compute-Combine 流程)
    • [五 EP 之专家计算内核实现分析](#五 EP 之专家计算内核实现分析)
    • [六 EP 之完整函数调用链总结](#六 EP 之完整函数调用链总结)
    • 参考资料

一 MOE + EP 概述

1.1 MOE 出现的背景

大模型的算力与显存瓶颈,随着大模型规模越来越大,参数量动辄上百亿,但我们的显存和算力却跟不上。尤其是在 Transformer 的 Decoder 层中,除了 Attention 之外,参数量最多、计算最重的部分就是 MLP(多层感知机)层

一个典型的 Llama 系列模型的 MLP 层通常包含两个大规模矩阵乘法,计算流程如下:

  1. MLP 层计算流程
  • 门控和上投影: Z g , Z u = X ⋅ W g a t e _ u p Z_g, Z_u = X \cdot W_{gate\up} Zg,Zu=X⋅Wgate_up + b(其中 W g a t e _ u p W{gate\_up} Wgate_up 实际包含 gate_proj 和 up_proj 两组权重)
  • 激活并逐元素相乘: H = SiLU ( Z g ) ⊙ Z u H = \text{SiLU}(Z_g) \odot Z_u H=SiLU(Zg)⊙Zu
  • 矩阵乘法2: Y = H ⋅ W 2 + b 2 Y = H \cdot W_2 + b_2 Y=H⋅W2+b2(其中 W 2 ∈ R h × d W_2 \in \mathbb{R}^{h \times d} W2∈Rh×d,输出维度回到 d,这里的 W 2 W_2 W2 实际是 down 线性层权重)
  1. 参数量计算(以 d model = 4096 d_{\text{model}}=4096 dmodel=4096 为例)
  • 隐藏层尺寸: h i d d e n _ s i z e = 4 × 4096 = 16384 hidden\_size = 4 \times 4096 = 16384 hidden_size=4×4096=16384
  • 单个 MLP 层参数量: 2 × 4096 × 16384 + 16384 × 4096 ≈ 201 M parameters 2 \times 4096 \times 16384 + 16384 \times 4096 \approx 201\text{M parameters} 2×4096×16384+16384×4096≈201M parameters(FP16 精度下约占 384MB 存储,不含 bias)

有人研究了 gpt-like 模型(opt)的参数分布,如下图所示:

可以看出,随着模型变大,MLP参数量占比越来越大,接近 2/3,也就是说优化 MLP模块的计算量刻不容缓。更重要的是,MLP 是密集计算(dense computation): 每一个输入 token 都要和全部参数做完整计算,没有任何跳过或稀疏性。这就导致:模型越大,推理越慢,显存压力越大,扩容变得极其困难。

为了解决这个问题,研究人员提出了 MOE(Mixture of Experts) 架构------它的核心思想对于整个模型来说不是减少参数,而是把 MLP 模块拆成多个专家 Experts(小 MLP),每个 token 只激活其中少数几个(比如 8 个)专家,从而在不显著增加每 token 计算量的前提下,大幅提升模型总参数量和表达能力。

1.2 MOE 的基本计算流程

  1. MOE 的输入: hidden_states 张量,也就是每个 token 的隐状态表示,而不是离散 token id。
  2. 专家池
  • 包含 E E E个独立的 MLP 专家(例如 E = 8 E=8 E=8),每个专家有专属权重 W e ( 1 ) , W e ( 2 ) W_e^{(1)}, W_e^{(2)} We(1),We(2),结构与普通MLP类似;
  • 总参数量 ≈ E × (单个专家参数量 ) \approx E\times(单个专家参数量) ≈E×(单个专家参数量),远大于普通 dense MLP。
  1. 门控网络(Gating Network)
  • 这是一个轻量级网络(通常为线性层 + softmax),计算式为: g = TopK ( softmax ( W g ⋅ x ) , k ) g = \text{TopK}(\text{softmax}(W_g \cdot x), k) g=TopK(softmax(Wg⋅x),k);
  • 其中 W g ∈ R E × d W_g \in \mathbb{R}^{E \times d} Wg∈RE×d,输出每个专家的被选中概率,再选出 Top-K 个专家(如 K = 2 K=2 K=2),也就是说这里的 g 表示 MoE 路由器输出的路由结果,也可以理解成:这个 token 被分配到哪些专家topk_ids,以及对应的专家权重是多少topk_weights。
  • 门控网络自身参数量极小,可忽略不计。
  1. 稀疏激活与加权融合
  • 将输入 tokens 张量送入被激活的 K K K个专家,分别计算输出;
  • 最终输出是这些专家结果的加权和,权重来自路由得分

MoE 的核心是通过在前向传播时激活少量专家,这样可以在保持大参数量(提升模型能力)的同时,控制每次推理的实际计算量。举个例子:在 Mixtral-8x7B 中,有 8 8 8 个专家,但每个 token 只激活其中 2 2 2 个。所以 总参数量 ≈ 47B,但 每个 token 的实际计算量 ≈ 12B dense 模型。对于这类模型,常见现象是"总参数量"和"每个 token 实际参与计算的参数量"会明显不同。

如果不考虑使用 EP 并行策略优化计算效率,moe 层推理的朴素实现应该是下述这样:

python 复制代码
batch_outputs = []
for token in batch_tokens:
    topk_weights, topk_ids = gate(token)
    expert_outputs = [get_expert(expert_id)(token) for expert_id in topk_ids]
    output = sum(
        weight * expert_output
        for weight, expert_output in zip(topk_weights, expert_outputs)
    )
    batch_outputs.append(output)

1.3 EP 概述

专家并行(Expert Parallelism, EP) 是一种专为 MoE 层设计的分布式策略,其核心思想是将 MoE 层中的海量专家分布至不同的计算设备(如 GPU)上。EP 的主要优势在于:

  1. 提高专家计算的并行度。不同 GPU 可以并行处理被路由到本地专家的 token,从而提升 MoE 层整体吞吐。实际收益取决于 token 路由分布、专家负载、通信后端、跨节点拓扑和 kernel 效率,不能简单等价为 batch size 线性增大。
  2. 降低单卡专家权重的显存压力。启用 EP 后,每个 rank 通常只需要加载本地专家权重;
    EP 的实现需要两次 All-to-All 通信,涵盖 Dispatch(分发)Combine(聚合) 两个阶段:
  • Dispatch:根据 Router/Gate 的路由结果,将 token、top-k expert id、top-k weight 等信息送到能够执行目标专家的设备。
  • Local Computation:每个设备对本地专家接收到的 token 做专家 MLP 计算。
  • Combine:专家计算完成后,将结果按原始 token 位置回收并聚合,交还给后续模型层。

虽然 EP 有各种好处,但是,EP 也给推理系统的通信与负载均衡带来了复杂性:

  • 跨节点通信开销:EP 会引入跨 GPU、跨节点的数据传输。为提高吞吐,系统需要尽量减少通信量,并在可能时让计算与通信重叠。
  • 负载均衡挑战:
    • EP 常与 DP、TP 等并行策略组合使用。不同 rank 或不同 DP 组之间的 token 数量、专家命中数量可能不均衡。
    • 专家本身也可能出现负载不均。vLLM 支持 EPLB(Expert Parallelism Load Balancing),用于在启用 EP 的 MoE 模型中维护专家负载状态,并可配合冗余物理专家等机制缓解热点专家问题。
  • 总结
    在 MoE 推理中,"DP + EP" 通常是比 "纯 TP" 或 "DP + TP + EP" 更经典且高效的策略,DeepSeek 后期推荐的 decode 阶段推理的并行配置也是 DP+EP 并行(Attention 还是完整的权重)。原因在于:MoE 模型架构已经解决了最大的显存和算力问题(即大参数量的 MLP 被分割成了多个小专家),而 Attention 层的参数量通常单卡能撑住,如果在 Decode(小 Batch、低计算密度)场景下强行引入 TP(张量并行),不仅无法带来显著性能提升,反而会引入额外 AllReduce/AllGather 通信开销及工程复杂度,因此,MOE业界普遍倾向于 DP + EP
    不过,为了兼容更大的 MoE 模型和更复杂的系统需求,vLLM 也支持 TP + DP + EP 等混合并行。此时 Attention 等非 MoE 层仍可以使用 TP,而 MoE 层内部在启用 EP 后会把专家维度映射到 EP ranks:每个设备完整持有一部分专家,MoE 层本身不再把单个专家继续按 TP 切分。
    但是 vLLM 事实上支持更为复杂的混合并行策略,DP、TP、PP、EP、SP,其中 SP 是 Megatron 引入的针对 Norm 层的序列并行。经典并行策略关系:EP SIZE = TP SIZE × DP SIZE。

vLLM 针对 EP + MoE 系统的设计思想是遵循经典的 Dispatch-Compute-Combine 模式:

  1. 分发 (Dispatch):通信阶段。每个 GPU 先通过本地 gate/router 得到 token 的 topk_ids 和 topk_weights。系统随后根据 expert id 和当前 EP rank 的映射关系,将 token 或路由相关信息送到能够执行对应专家的设备。
    这个过程需要一次全局的数据交换,也就是说,每个 GPU 对自己的输入tokens 跑 gate,得到 topk_indices(专家 id)和 topk_weights,随后通过专家 id 能算出这些专家在哪个 GPU 上(EP rank)。
    在单卡内部,先根据目标 GPU rank把 tokens 分组、重排,例如 GPU0 上:去 GPU0 的 tokens去 GPU1 的 tokens去 GPU2 的 tokens去 GPU3 的 tokens,此时经过一次All-to-All之后,所有该由本机 experts 处理的 tokens,都被聚拢到这台 GPU 上了。
  2. 本地计算 (Local Computation):计算阶段。当目标专家所需的 token 到达本地,GPU 就可以独立执行本地专家的 MLP 计算,启用 EP 时,每个设备通常完整持有一部分专家;
  3. 回收 (Combine):通信阶段。专家计算完成后,结果需要回到原始 token 所属的位置,才能继续后续模型层计算。概念上,这相当于把各 GPU 上的专家输出再交换回来;在 vLLM 中,具体表现为 combine 抽象,底层可能是 reduce-scatter、All2All 或由 MoE kernel 内部完成的 combine/reduce 逻辑。

1.4 MOE + EP 的基本计算过程

在 Mixture of Experts (MoE) 架构中,随着模型规模的扩大,专家(Expert)的参数量巨大,单个 GPU 显存往往无法容纳所有专家权重。因此,需要采用专家并行 (Expert Parallelism, EP) 策略,将不同的专家分布存储在不同的 GPU (节点)上。但是,当 Token 经过门控网络(Router)路由时,其选中的 Top-K 专家可能位于当前 Token 所在的 GPU(本地),也可能位于其他 GPU(远程)。这种数据与计算资源的不对齐,导致必须引入跨设备的通信操作(All-to-All),将 Token 发送到其目标专家所在的设备进行计算,然后再传回。

EP + MOE 的前向推理过程可以划分为如下五个标准阶段:

  1. 输入:Token Embeddings,假设有 N N N个token,每个维度为 d d d,输入张量为 X ∈ R N × d X \in \mathbb{R}^{N \times d} X∈RN×d
  2. 门控网络(Gating / Router):门控网络决定了每个 Token 由哪些专家处理。
  • 计算流程通过轻量线性层 + softmax + Top-K实现,得分计算式: s c o r e s = softmax ( X W g ) ,其中 W g ∈ R d × E scores = \text{softmax}(XW_g) ,其中 W_g \in \mathbb{R}^{d \times E} scores=softmax(XWg),其中Wg∈Rd×E
  • 核心操作:为每个token选出Top-K个专家
  • 输出内容:
    • 每个 token 被分配到的专家索引(topk_ids)
    • 用于加权融合的门控权重(topk_weights)
  1. 数据重排与分发
    因为在开启 EP 后,专家权重分布式存储在不同 GPU 上,而 Token 可能被路由至任意专家,所以需根据路由结果重新组织数据,将 Token 搬运至目标专家处。
  • 本地分组:每个 GPU 根据本地 Token 的目标专家索引,对其进行分组。
  • 全局交换:通过 All-to-All 集合通信,各 GPU 将分组后的 Token 发送至对应专家所在的 GPU。完成后,每个 GPU 均持有全部需要由本地专家处理的所有 Tokens。如果未启用 EP all2all kernels,某些部署会退化为其他 dispatch/combine 方式,例如 AllGather + ReduceScatter。
  1. 专家并行计算
  • 每个 GPU 对收到的 token 调用本地专家进行前向计算。
  • 典型形式可以写成: Y i = E x p e r t i ( X i ) Y_i = Expert_i(X_i) Yi=Experti(Xi),其中 X i X_i Xi 表示送到第 i i i个本地专家的 token 子集。
  1. 结果收集与聚合(再次 All-to-All)
    专家计算完成后,每个的 tokens 的输出结果分散在各 GPU 上,需按原始 Token 顺序进行收集。
  • 专家计算完成后,输出会按路由映射回收,并恢复成原始 token 顺序。
  • vLLM 通过 moe_unpermute 按 topk_weights 对多个专家的输出进行加权融合,得到每个 token 的最终表示 Y = ∑ k = 1 K g k ⋅ Expert e k ( x ) Y = \sum_{k=1}^{K} g_k \cdot \text{Expert}_{e_k}(x) Y=∑k=1Kgk⋅Expertek(x)
  • 最终输出 Y ∈ R N × d Y \in \mathbb{R}^{N \times d} Y∈RN×d,其顺序与输入保持一致,可继续传递至下一层网络。
    EP + MOE 的前向推理过程的伪代码实现如下所示:
shell 复制代码
输入: hidden_states (num_tokens, hidden_dim), router_logits (num_tokens, num_experts)

1. [路由选择] 根据 router_logits 选择 top-k 专家
2. [Dispatch] 将 token 分发到对应的 GPU(专家并行)
3. [Permute] 重排 token,按专家分组
4. [Compute] 各 GPU 上的专家并行计算(融合 GEMM)
5. [Unpermute] 恢复 token 顺序并加权合并
6. [Combine] 收集各 GPU 的结果并合并

输出: final_hidden_states (num_tokens, hidden_dim)

1.5 MoE + EP 推理场景实例

为了快速了解 MoE 层 EP 执行的流程,这里以 Qwen3-30B-A3B 推理为例,先来看下模型配置和结构信息,Qwen3-30B-A3B 和 Qwen3-235B-A22B 模型架构一模一样,其中 MOE layer:

  • 专家数量: 128 128 128 个,激活的专家数量: 8 8 8个。
  • 系统会对选中专家的权重进行重新归一化:routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
    模型结构信息如下所示:
python 复制代码
Qwen3-30B-A3B model structure: Qwen3MoeModel(
  (embed_tokens): Embedding(151936, 2048)
  (layers): ModuleList(
    (0-47): 48 x Qwen3MoeDecoderLayer(
      (self_attn): Qwen3MoeAttention(
        (q_proj): Linear(in_features=2048, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2048, out_features=512, bias=False)
        (v_proj): Linear(in_features=2048, out_features=512, bias=False)
        (o_proj): Linear(in_features=4096, out_features=2048, bias=False)
        (q_norm): Qwen3MoeRMSNorm((128,), eps=1e-06)
        (k_norm): Qwen3MoeRMSNorm((128,), eps=1e-06)
      )
      (mlp): Qwen3MoeSparseMoeBlock(
        (gate): Linear(in_features=2048, out_features=128, bias=False)
        (experts): ModuleList(
          (0-127): 128 x Qwen3MoeMLP(
            (gate_proj): Linear(in_features=2048, out_features=768, bias=False)
            (up_proj): Linear(in_features=2048, out_features=768, bias=False)
            (down_proj): Linear(in_features=768, out_features=2048, bias=False)
            (act_fn): SiLU()
          )
        )
      )
      (input_layernorm): Qwen3MoeRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): Qwen3MoeRMSNorm((2048,), eps=1e-06)
    )
  )
  (norm): Qwen3MoeRMSNorm((2048,), eps=1e-06)
  (rotary_emb): Qwen3MoeRotaryEmbedding()
)

8 卡的 2080ti(单张显卡 11g 显存) 硬件,并行的模式是 tp =4, dp = 2, vllm 启动后的 nvidia-smi 终端显示如下:

这里开启 tp 的原因主要还是单卡显存太小了,如果不在权重上做切分没有办法加载起 30b 参数数量的模型。

shell 复制代码
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A           45336      C   VLLM::Worker_DP0_TP0_EP0              10236MiB |
|    1   N/A  N/A           45339      C   VLLM::Worker_DP0_TP1_EP1              10236MiB |
|    2   N/A  N/A           45341      C   VLLM::Worker_DP0_TP2_EP2              10236MiB |
|    3   N/A  N/A           45343      C   VLLM::Worker_DP0_TP3_EP3              10236MiB |
|    4   N/A  N/A           45337      C   VLLM::Worker_DP1_TP0_EP4              10234MiB |
|    5   N/A  N/A           45338      C   VLLM::Worker_DP1_TP1_EP5              10234MiB |
|    6   N/A  N/A           45340      C   VLLM::Worker_DP1_TP2_EP6              10234MiB |
|    7   N/A  N/A           45342      C   VLLM::Worker_DP1_TP3_EP7              10234MiB |
+-----------------------------------------------------------------------------------------+

需要注意的是,在 vLLM 的 MoE 层里,开启 EP 后,专家并行的大小并不是只看 TP 或只看 DP,而是会把 DP * PCP * TP 展平为 EP 维度。上面的 tp=4, dp=2 示例中,MoE 层对应的 EP size 是 2 * 4 = 8,所以进程名里会出现 EP0 到 EP7。在默认 expert_placement_strategy="linear" 且没有 EPLB 冗余专家的情况下,128 个专家会按连续编号分配到 8 个 EP rank 上,每个 rank 约持有 16 个专家。

启动程序的 vscode 配置如下:

json 复制代码
{
    "version": "0.2.0",
    "configurations": [
        {
            "name": "vLLM: API Server Debug",
            "type": "debugpy",
            "request": "launch",
            "module": "vllm.entrypoints.openai.api_server",
            "args": [
                "--model",
                "/model/HuggingFace/Qwen/Qwen3-30B-A3B-Instruct-2507",
                "--dtype",
                "float16",
                "--max-model-len",
                "1024",
                "--gpu-memory-utilization",
                "0.9",
                "--max-num-batched-tokens",
                "4096",
                "--max-num-seqs",
                "8",
                "--port",
                "13333",
                "--tensor-parallel-size",
                "4",
                "--data-parallel-size",
                "2",
                "--enable-expert-parallel",
                "--enforce-eager"
            ],
            "env": {
                // 开启 vLLM 调试日志
                "VLLM_LOGGING_LEVEL": "DEBUG",
                "PYTHONPATH": "${workspaceFolder}:${env:PYTHONPATH}",
                // 国内 Hugging Face 源配置
                "HF_ENDPOINT": "https://hf-mirror.com",
                "HF_HOME": "${workspaceFolder}/.cache/huggingface",
            },
            "console": "integratedTerminal",
            "justMyCode": false,
            "cwd": "${workspaceFolder}",
            "stopOnEntry": false
        }
    ]
}

下面通过一个简单实例来理解 moe + ep 推理流程。推理场景设定如下:

  • 硬件环境:一个配备 4 张 GPU 的服务器 (Rank 0, 1, 2, 3)(简单场景方便描述)
  • 模型:Qwen3-30B-A3B。
    • 总专家数 (num_experts): 128
    • 激活专家数 num_experts_per_tok: 8。处理每一个 Token 时,门控网络会根据输入特征计算专家分数,只选择分数最高的 8 个专家参与计算,其余 120 个专家对当前 Token 不执行专家 MLP 计算。
  • 专家分布 (EP 核心):128 个专家被平均分配到 4 张 GPU 上。
    • GPU 0 (Rank 0): 持有专家 E0, E1, ..., E31
    • GPU 1 (Rank 1): 持有专家 E32, E33, ..., E63
    • GPU 2 (Rank 2): 持有专家 E64, E65, ..., E95
    • GPU 3 (Rank 3): 持有专家 E96, E97, ..., E127
  • Prefill 阶段场景推理:一个包含 4 个请求的批次,每个请求 512 个令牌,总共有 2048 个令牌 需要被处理。EP 必须和 DP 结合,由于数据并行 (DP) 的存在, tokens 被均匀分布在 4 个 GPU 上:
    • GPU 0: 持有 512 个令牌 (T0...T511)
    • GPU 1: 持有 512 个令牌 (T512...T1023)
    • GPU 2: 持有 512 个令牌 (T1024...T1535)
    • GPU 3: 持有 512 个令牌 (T1536...T2047)

这里之所以以 prefill 阶段来总结 ep 推理流程,是因为其不用考虑 decode 阶段的专家负载不均衡问题,即每个专家接收到的 tokens 数量不一致的问题,而 prefill 阶段是推理第一步,一般来讲经过 dp 后, tokens 会被均匀的分布在 GPU 上。

在每一层 MoE 中,每个 rank 会先基于本地 hidden states 计算 router logits。随后,Gate 为每个 Token 选出 Top-8 专家。由于某个 Token 选中的专家可能不在当前 GPU 上,例如 GPU 0 上的某个 Token 可能选中了 E37、E81、E103 等远端专家,EP 就需要把对应的 hidden states 和 routing 信息发送到持有这些专家的 rank 上执行专家 MLP。

1.5.1 推理流程简单拆解

步骤 1: 本地门控与路由决策 (Local Gating)

计算流进入 MoE 层 (Qwen3MoeSparseMoeBlock),步骤 1 是在每个 GPU 内部独立完成的, 通过 FusedMoE.select_experts 函数实现。

  • Gate 层:每个 GPU 将其本地的 512 个 token 输入到共享的门控网络 (gate)。对于每个 token,门控网络会输出一个长度为 128 的 logits 向量,代表该 token 与 128 个物理专家的匹配分数。
  • Top-k 选择:vLLM 对这个 logits 向量执行 torch.topk(... , k=topk) 操作,为每个 token 选出分数最高的 8 个专家索引 (topk_indices) 及其对应的权重 (topk_weights)。
python 复制代码
class FusedMoE(CustomOp):
    def select_experts(...):
        topk_weights, topk_ids, token_expert_indices = fused_topk(
            hidden_states=hidden_states,
            gating_output=router_logits,
            topk=top_k,
            renormalize=renormalize,
            indices_type=indices_type,
        )

此时,GPU 0(为例)内存中现在有一个形状为 512, 8 的 topk_indices 张量和一个 512, 8 的 topk_weights 张量。举例理解 topk_indices,topk_indices0可能为 5, 18, 33, 59, 2, 9, 12, 17,意味着 T0 token,也就是输入序列中的第一个 token 需要被专家 E5、E18、E33、E59、E2、E9、E12、E17 处理

步骤 2: 第一次 All-to-All 令牌分发 (Dispatch)

准备工作就绪后,第一次全局通信开始。所有 GPU 同时调用 get_ep_group().dispatch。但需要注意,Naive 模式(默认/兼容模式)与理想的 All-to-All 不同: 它实际上执行的是 All-Gather,即简单粗暴地将所有输入 token 广播给所有 EP 节点,让每个节点获得全量数据自行筛选。只有在使用 PPLX 或 DeepEP 等优化内核时,才会实现真正的有选择性发送(All-to-All),即只将 token 发送给其目标专家所在的 GPU。Dispatch(All-to-All)通信执行示例:

发送阶段:

  • GPU 0 发送:需要发送到: GPU0(150个 tokens), GPU1(120个), GPU2(140个), GPU3(102个) = 总共 512 个 tokens
  • GPU 1 发送:需要发送到: GPU0(130个 tokens), GPU1(110个), GPU2(150个), GPU3(122个) = 总共 512 个 tokens
  • GPU 2 发送:需要发送到: GPU0(140个 tokens), GPU1(130个), GPU2(100个), GPU3(142个) = 总共 512 个 tokens
  • GPU 3 发送:需要发送到: GPU0(120个 tokens), GPU1(140个), GPU2(130个), GPU3(122个) = 总共 512 个 tokens
    接收阶段,每个 GPU 接收到的 tokens 数量(接收总数不一致):
  • GPU 0 接收:150 + 130 + 140 + 120 = 540 个 tokens
  • GPU 1 接收:120 + 110 + 130 + 140 = 500 个 tokens
  • GPU 2 接收:140 + 150 + 100 + 130 = 520 个 tokens
  • GPU 3 接收:102 + 122 + 142 + 122 = 488 个 tokens

步骤 3: 本地专家计算 (Local Expert Computation)

dispatch 之后,每个 rank 只负责计算本地持有的专家。在 Qwen3 MoE 中,本地专家数量由物理专家总数和 EP size 决定:

复制代码
n_local_physical_experts = n_physical_experts // ep_size

每个 rank 只保存并计算自己的本地专家权重。收到或收集到 token 后,vLLM 会根据 topk_ids 和 expert map 判断哪些 token 需要当前 rank 的专家处理。本地专家计算的具体 kernel 取决于 MoE backend、量化方式和硬件平台。

当前源码中可以看到多种路径,例如:

  • Triton fused MoE kernel。
  • CUTLASS MoE kernel。
  • DeepGEMM grouped GEMM。
  • FlashInfer MoE kernel。
  • CPU grouped GEMM 路径。

vLLM 会尽量使用 fused MoE 或 grouped/batched GEMM 类 kernel,把多个专家的小 GEMM 组织成更适合 GPU 执行的计算形式,而不是简单地为每个专家单独发射一个普通 GEMM。以 CUTLASS MoE 路径为例,源码中会先做 token 重排,使同一专家要处理的 token 尽量连续,然后构造每个专家对应的 GEMM problem size,再调用 ops.cutlass_moe_mm() 完成专家计算。相关逻辑包括:

复制代码
moe_permute(...)
ops.get_cutlass_moe_mm_problem_sizes_from_expert_offsets(...)
ops.cutlass_moe_mm(...)
moe_unpermute(...)

Grouped/Batched GEMM 的关键理解可以这样讲:

  • 每个专家对应一个 GEMM 子问题。
  • 不同专家收到的 token 数可能不同,因此每个子问题的 M 维可能不同。
  • K 维通常对应 hidden size,N 维对应专家中间层或输出维度。
  • grouped/batched kernel 可以把多个专家的 GEMM 组织在同一次调度或同一类 kernel 流程里,减少大量小 GEMM 带来的 kernel launch 开销,并改善 GPU 调度效率。
    这类 kernel 能减少重复调度开销,并提高多专家小矩阵乘法的执行效率;具体的数据复用和访存行为取决于 kernel 实现、矩阵形状、量化格式和硬件平台。

步骤 4: 第二次 All-to-All - 结果回收 (Combine)

专家计算完成后,计算结果(修改后的 hidden_states)需要被送回它们各自的原先节点---即 token 最初所在的 GPU。这通过调用 get_ep_group().combine 通信算子完成。这是第一次 All-to-All 的逆向过程:每个 GPU 将处理后的 tokens 按照它们原始来源 GPU 进行分组并发送回去。

Combine(All-to-All)通信执行示例:

发送阶段:

  • GPU 0 发送:需要发送到: GPU0(150个 tokens), GPU1(130个), GPU2(140个), GPU3(120个) = 总共 540 个 tokens(对应步骤2中GPU 0接收到的540个tokens)
  • GPU 1 发送:需要发送到: GPU0(120个 tokens), GPU1(110个), GPU2(130个), GPU3(140个) = 总共 500 个 tokens(对应步骤2中GPU 1接收到的500个tokens)
  • GPU 2 发送:需要发送到: GPU0(140个 tokens), GPU1(150个), GPU2(100个), GPU3(130个) = 总共 520 个 tokens(对应步骤2中GPU 2接收到的520个tokens)
  • GPU 3 发送:需要发送到: GPU0(102个 tokens), GPU1(122个), GPU2(142个), GPU3(122个) = 总共 488 个 tokens(对应步骤2中GPU 3接收到的488个tokens)
    接收阶段,每个 GPU 接收到的 tokens 数量(接收总数恢复为 512,与初始状态一致):
  • GPU 0 接收:150 + 120 + 140 + 102 = 512 个 tokens
  • GPU 1 接收:130 + 110 + 150 + 122 = 512 个 tokens
  • GPU 2 接收:140 + 130 + 100 + 142 = 512 个 tokens
  • GPU 3 接收:120 + 140 + 130 + 122 = 512 个 tokens
    步骤 5: 最终合并 - 加权求和 (Weighted Sum)
  • 操作:每个 GPU 使用步骤 1 中计算出的 topk_weights,对接收到的 8 份专家输出进行加权求和。
  • 输出:得到 MoE 层的最终输出 final_hidden_states。
python 复制代码
final_hidden_state[token] =
    sum(topk_weights[token, i] * expert_output[token, topk_ids[token, i]])

为了方便理解,上述步骤总结,没有描述 tokens 重排 + 逆向重排(恢复原始顺序)的过程,下章代码分析会给出。至此,MoE 层 + EP 的前向推理过程大致总结完成,计算流将带着 final_hidden_states 张量进入模型的下一个层。

二 AlltoAll 通信原理

2.1 Scatter & ReduceScatter

2.2 ReduceScatter

2.3 Alltoall

1. 准备阶段:切分
2. 传输阶段:路由与交换
3. 结果阶段:拼接 (Concat)
4. 代码示例

三 MoE 层架构和推理流程

四 Dispatch-Compute-Combine 流程

五 EP 之专家计算内核实现分析

六 EP 之完整函数调用链总结

参考资料

相关推荐
七夜zippoe3 小时前
深入解析CANN仓库中的HCCL分布式通信库
pytorch·分布式·cann·hccl·通信库
星恒随风3 小时前
C++ STL 栈详解:stack 的使用、经典题目与简单模拟实现
开发语言·数据结构·c++·笔记·学习
富士康质检员张全蛋3 小时前
Kafka的操作 消费者组 消费位置查看
分布式·kafka
运维行者_3 小时前
如何查看每个IP的带宽使用情况?NetFlow 技术实战指南
开发语言·网络·分布式·后端·架构·带宽
玖玥拾5 小时前
Unity 3D 笔记(十一)UI 框架进阶:栈弹窗交互、BasePanel 基类虚方法、DoTween 界面动画
笔记·3d·unity
FellAveal5 小时前
【Go语言入门学习笔记】Part13.结构体与接口
笔记·学习·golang
学计算机的计算基5 小时前
回溯算法下篇:四道经典题讲透约束剪枝、原地标记、预计算与状态压缩
java·笔记·算法
雷工笔记5 小时前
Ubuntu迁移记录
笔记·ubuntu
Token炼金师5 小时前
框架的擂台:LangChain、LlamaIndex、Dify、AutoGen 与 LangGraph —— 应用框架选型五局
人工智能·深度学习·llm