高性能专家并行（EP）内核解剖学解读

背景

大型语言模型（LLM）的体量日益庞大，这导致运行它们需要消耗大量的 GPU 资源。理想情况下，LLM 推理应当是“极度并行”（embarrassingly parallel）的，即我们可以简单地在不同的 GPU 上独立计算互不相关的事物。然而现实是，为了在 LLM 推理中利用大量的 GPU，我们必须让 GPU 之间进行通信。

实现多 GPU 协作有多种方式，包括张量并行（Tensor Parallelism）、流水线并行（Pipeline Parallelism）、上下文并行（Context Parallelism）以及专家并行（Expert Parallelism, EP）等。虽然每种并行方式都有其适用场景，但对于混合专家模型（MoE）而言，在 MoE 层进行大规模服务时，“宽专家并行”（wide Expert Parallelism, wideEP）是王者。

例如，vLLM 曾发布过一篇关于 DeepSeek 大规模服务的博客，展示了在 H200 集群上使用 wideEP 和数据并行注意力机制，实现了每 GPU 2.2k tokens/s 的生产级吞吐量。

与其他并行方式不同，张量并行、流水线并行等所需的通信模式由架构固定：谁发送、谁接收以及发送多少数据，在前向传播开始之前就已确定，并且在每一步中都是相同的。因此，这些通信可以作为标准的集合操作（collectives）运行。

专家并行（EP）则截然不同。 哪些 token 需要到达哪些 GPU，是由路由器（router）根据数据在运行时动态决定的，且在每一个 MoE 层都是全新的。此外，token 有明确的来源：我们假设采用 DeepSeek 服务中的“数据并行注意力”架构，即每个 token 仅存在于集群中的某一个 rank（即某一块 GPU）上。专家（experts）分散在这些相同的 rank 上，因此，一个 token 及其被路由到的专家通常不在同一个位置。

这种动态性使得 EP 通信内核的设计变得复杂且关键。现代此类内核的形态由 DeepSeek 的 DeepEP 库确立。本文将深入剖析 DeepEP 风格的 dispatch（分发）和 combine（合并）内核，首先探讨高吞吐量的形态，随后讨论低延迟的形态。

核心内容

任务设定与数据流

为了具体化讨论，我们设定如下环境：

• 硬件拓扑：8 块 GPU 分布在 2 个节点上，通过 RDMA 连接，每个数据并行 rank 拥有一块 GPU。
• 注意力机制：注意力在每个 GPU 上针对一批 token 运行，batch size 在不同 GPU 间可能不同。
• 专家并行配置：每个 GPU 拥有 2 个专家，每个 token 路由到 2 个专家（2-out-of-16）。

在每个 rank 进入 EP 层时，我们拥有一个形状为 [batch_size, hidden_size] 的激活张量。路由层会在本地运行，为每个 token 提供专家分配信息。具体来说，路由器为每个 token 生成一个长度为 16 的 logits 张量，我们取其中最大的 2 个值的索引，得到一个形状为 [batch_size, 2] 的张量。例如，如果 token $i$ 被路由到专家 3 和专家 7，那么该行的值就是 [3, 7]。

因此，在 EP 层入口处，每个 rank 持有两样东西：

1. 它产生的激活行（activations）。
2. 经过本地路由后，每个激活行对应的 top-2 专家分配。

并非所有专家都位于本地。有些专家位于邻近的 NVLink 对等节点上，有些则位于只能通过 RDMA 访问的远端节点上。专家并行内核的目标是将激活数据发送到它们需要去的地方，在到达后运行专家 GEMM（通用矩阵乘法），然后将结果带回原处。

吞吐量与延迟的权衡

通信优化通常分为两个方向：吞吐量（Throughput） 和 延迟（Latency）。这对应于推理的两个阶段：

• Prefill 阶段：带来大型、计算密集型的 batch，有足够的其他工作来隐藏通信开销。
• Decode 阶段：几乎没有其他工作可做，传输本身成为等待的主要瓶颈。

下文首先讨论针对吞吐量优化的标准 EP 形态。

高吞吐量：先询问，再发送（Ask, then send）

1. Dispatch（分发）阶段的目标

Dispatch 的核心目的是为分组 GEMM（Grouped GEMM）准备数据。分组 GEMM 是指一次内核启动中，针对每个专家运行独立的矩阵乘法，每个乘法处理的行数不同。路由产生的正是这种“不规则性”（raggedness）：每个专家组的行数取决于它抽到了多少 token。

在路由之前，这些 token 分散在集群的每个 rank 上。因此，在每个 rank 上，dispatch 必须将绑定到本地专家的 token 收集到一个单一的密集缓冲区中，以便分组 GEMM 可以一次性消费。

2. 核心难点：未知的大小

主要困难在于我们事先不知道这个缓冲区的大小。

• 哪个 token 去哪个专家是由路由器在运行时决定的。
• 分布是不均匀的（lumpy）且每一步都在变化。
• 一个专家现在可能抽到 1000 个 token，下次可能一个都没有。
• 直到其他所有 rank 都运行完路由器之前，我们都不知道本地专家会收到多少 token。

因此，本地激活缓冲区的大小，以及任何特定 token 在缓冲区中的位置，在事前都是未知的。

3. 解决方案：固定矩形 vs. 精确分配

应对未知大小有两种主要策略：

• 策略 A：固定矩形（Fixed Rectangle）

  • 方法：为最坏情况预留空间，分配一个带有填充槽位的固定矩形缓冲区。
  • 缺点：最坏情况不具备可扩展性。如果所有对等节点同时将整个 batch 路由到同一个专家，每个填充槽位都必须大到足以容纳所有人。在 prefill 的大 batch 规模下，这意味着大量的 HBM（高带宽内存）被用于存储“空位”，而非数据。HBM 是宝贵资源，应释放出来用于 KV cache，以保持序列的活跃。
  • 适用场景：低延迟场景（下文详述）。

• 策略 B：精确分配（Exact Allocation）—— 吞吐量路径的选择

  • 方法：先找出真实的计数，然后分配确切需要的空间。
  • 优势：当 batch 较大且 GEMM 计算受限时，学习真实计数所需的额外通信开销可以被计算任务隐藏。推理栈（如 SGLang）通过双 batch 重叠技术制造隐藏空间：将步骤拆分为两个微批次（microbatches），当一个微批次的 token 在传输线上时，另一个微批次的 GEMM 保持 SM（流式多处理器）忙碌。
  • 实现：吞吐量路径分配一个不规则缓冲区（ragged buffer），其大小仅基于实际接收到的 token 数量。

4. 协调传递（Coordination Pass）

为了实现精确分配，必须在任何激活数据传输之前获取计数信息。

• 每个 rank 已经知道从自己的路由中向每个对等节点发送了多少 token。
• 如果所有 rank 交换这些数字，每个 rank 就可以计算出它将接收多少 token。
• 通信效率：协调传递的字节数很少，每个对等节点只需交换几个整数，而不是兆字节级的隐藏状态。
• 通信路径：计数交换镜像了物理拓扑：计数首先通过 RDMA 在节点之间交叉，然后在节点内的 GPU 之间通过 NVLink 交叉，沿途汇聚。因此，协调成本与真实分发所需的两次跳转相同。

5. 缓冲区布局与发送

一旦 rank 知道了从每个来源接收多少 token，缓冲区的安全布局便通过**前缀和（prefix sum）**自然得出：

• 第一个来源的 token 从索引 0 开始。
• 下一个来源的 token 从上一个来源结束的位置开始，依此类推。

计数信息提供了我们缺失的两样东西：缓冲区的大小，以及每个数据块在其中的位置。

6. 实际发送机制

布局确定后，即可发送激活