DeepEP 学习笔记：MoE 场景下的专家并行通信

这篇文章记录我对 DeepEP 的个人理解。DeepEP 本质上是 All-to-All 通信原语的一种高性能实现，专门服务于 MoE（Mixture of Experts）场景下的 dispatch 与 combine。

文章覆盖三个主题：

• DeepEP v1：最经典的实现，打基础。
• DeepEP v2：direct / hybrid 模式与去重思想。
• MoonEP：用通信换计算的在线搬专家思路。

Table of contents

前置知识与一个物流类比

MoE 里的几个基础概念——dispatch、combine、router（gate）、topk——先快速对齐。与其干巴巴地讲公式，不如用快递物流网络来类比：

DeepEP v1 支持三种模式：

1. intranode：单机多卡，走 NVLink。
2. internode：多机多卡，机内走 NVLink，机间走 RDMA。
3. low-latency：纯 RDMA（IBGDA，绕过 CPU），适用于推理 decode 阶段追求低延迟。

训练或 prefill 阶段一般用普通模式（intranode / internode）。

DeepEP v1

dispatch 前：全局布局与通知

先以 intranode 为例。假设 $R=4$ 个 rank，专家总数 $E=8$，$\text{EP}=4$（每个 rank 上 2 个专家），$\text{topk}=2$。

dispatch 前要先调用 get_dispatch_layout，让每个 rank 获得本次传输的全局信息（layout）：

• num_tokens_per_expert：每个工厂会收到多少个包裹。
• num_tokens_per_rank：每个城市会收到多少个包裹。
• is_token_in_rank：每个包裹会发往哪些城市（把 top-k 摊平后才能快速查表）。

先统计这些信息，才能知道每个专家那边要准备多大的 buffer 来接收，否则空间不够就尴尬了。

接下来是 notify_dispatch：

1. 每个 rank 把自己的 num_tokens_per_rank 汇总，获得每个城市的视角：自己总共会进多少包裹。
2. 所有 rank 之间做一次 barrier 同步，确保大家对自己的接收量达成共识。
3. 计算 channel_prefix_matrix。

dispatch 过程

queue 是核心结构：每个 (channel, 目标 rank) 都有一个 queue。把 channel 想象成一个分拣点，就要给输入的包裹建 $R$ 条分拣传送带。因为有很多分拣点一起分担，所以每个分拣点只负责某个区间范围内的包裹的发送。

一个具体的 queue 需要记录：

• 本 channel 的搬运范围；
• queue 的头和尾；
• 要传输的具体内容（x、src_idx、topk_idx、topk_weights 等元信息）。

sender 负责把范围内包裹发送到目标城市的分拣传送带上。流程大概是：

1. 检查目标 queue 是否有空位。
2. 一次遍历一个 chunk（有限数量的 tokens）。
3. 对每个 token，查看 is_token_in_rank 是否命中当前目标 rank；命中就放入队列，并更新 send_head。

receiver 则计算本 channel 应该把数据写到 recv_x 的哪个 offset，拷贝数据和相关元信息，然后把队头弹出。

recv_x 就是接收货物的地方，后续交给 expert 做 GEMM。

combine 过程

1. 先用 notify_combine 清空 queue。
2. combine 的 sender 根据 handle 中保存的反向路由信息，把 expert 输出传回原始 rank 的 queue。
3. combine 的 receiver 查每个原始 token 在 dispatch 时被发到哪个 slot，等所有数据到齐后做 reduce。

多机转发：NVLink + RDMA

多机好比跨国快递：国家与国家之间传输更慢、更不方便，但同样需要传送带机制。

dispatch kernel 中涉及 5 个角色：

Buffer 与内存视角

Buffer 好比物流公司总部：不亲自搬运，但管理所有资源。每个 rank 有一个 Buffer。

初始化时：

1. 知道自己属于哪个国家（RDMA rank）、哪个城市（NVL rank）。
2. 分配 NVLink buffer（仓库）。
3. 拿到 buffer 的 IPC handle（仓库钥匙）。
4. 装好对讲机（barrier signal 记录的地方）。
5. 分配 workspace。
6. 初始化计数器为 -1，表示尚未收到数据。

初始化完后调用 sync()：

• 各 rank 交换 IPC handle，从而可以直接访问其他 rank 的 NVLink buffer。
• 把指针表拷到 GPU。
• 如果是多机场景，还要初始化 RDMA 全局地址空间（NVSHMEM）。

Buffer 对外暴露的核心方法包括：get_dispatch_layout、intranode_dispatch、intranode_combine。

三个方法概览

• get_dispatch_layout

  1. 切到 comm stream，等待前一个事件完成。
  2. 分配统计表，launch get_dispatch_layout kernel。
  3. 如果 async 就返回 event handle；否则让 compute stream 等待本次 layout 完成。
  4. 返回结果 + event。

• intranode_dispatch

  1. 切到 comm stream，等待前一个事件完成。
  2. 分配 metadata tensor，launch notify_dispatch kernel。
  3. CPU 同步等待 GPU 把接收计数写好。
  4. 分配接收 tensor，launch dispatch kernel。
  5. 返回结果 + handle + event。handle 就是“发货单”，后面 combine 会用它。

• intranode_combine

  1. 读取 handle，切到 comm stream，等待前一个事件完成。
  2. launch notify_combine kernel，清空 queue head/tail。
  3. launch combine kernel。
  4. 返回 combined_x + event。这里不需要 CPU 同步，因为 handle 已经告诉了输出大小。

internode 相比 intranode 返回更多 tensor：每个 rank 要维护国家间 channel 分配、全球 channel 分配、从各国累计收多少、从全球累计收多少、RDMA 发货登记表、NVLink 发货登记表等。多机时的发货单更厚，因为一个 token 要经过两次运输。

low-latency 模式下，两个 RDMA buffer 做 ping-pong，用异或翻转使用 buffer，不需要每次分配，也不做 CPU 同步。

内存布局

buffer_ptrs 有 8 个，指向一段连续内存：

buffer_ptrs[nvl_rank]
│
├─ [0, num_nvl_bytes)                        : ① NVL data buffer（真正放数据）
├─ [num_nvl_bytes, +32)                      : ② barrier_signal（8 个 int）
├─ [+32, +32+64)                             : ③ buffer_ptrs_gpu（8 个 void*）
└─ [+32+64, +32+64+64)                       : ④ barrier_signal_ptrs_gpu（8 个 int*）buffer_ptrs

num_nvl_bytes 在 dispatch 时被组织成 (num_channels * num_ranks) 个 queue，按顺序保存相应变量。为什么要这么组织？因为避免多个 channel 抢同一把锁——每个 channel 对自己的 queue 读写，互不竞争。

RDMA buffer 也类似。每个 (channel_id, dst_rdma_rank) 对应一段空间：

rdma_buffer_ptr
│
├─ rdma_channel_data[channel_id][dst_rdma_rank]
│    大小: num_max_rdma_chunked_recv_tokens * num_bytes_per_token
│    作用: 实际 token 数据（含 x, scales, SourceMeta, topk_idx, topk_weights）
│
├─ rdma_channel_meta[channel_id][dst_rdma_rank]
│    大小: (8 * 2 + 2) * sizeof(int)
│    作用: 每个 dst_nvl_rank 的 start/end offset + RDMA start/end
│
├─ rdma_channel_head[channel_id][dst_rdma_rank]
│    大小: sizeof(uint64_t)
│
└─ rdma_channel_tail[channel_id][dst_rdma_rank]
     大小: sizeof(uint64_t)rdma_buffer_ptr

这些缓冲区是对称的：发送端和接收端用同样的索引规则定位同一块内存。

内存生命周期

进程启动
   │
   ▼
Buffer 构造函数
   ├── 分配 NVLink buffer（含 data + barrier + ptr tables）
   ├── 分配 workspace
   ├── 分配 CPU-mapped counters
   └── 此时 rdma_buffer_ptr 还未分配
   │
   ▼
sync()
   ├── 交换 IPC handles → 打开其他 rank 的 NVLink buffer
   ├── 把 ptr tables 拷到 GPU
   └── 若 num_rdma_bytes > 0: nvshmem::init() + 分配 RDMA buffer
   │
   ▼
每次 dispatch/combine
   ├── 在 comm_stream 上 launch kernel
   ├── kernel 读写 NVLink/RDMA buffer
   ├── 临时分配 PyTorch tensor
   └── CPU 可能轮询 moe_recv_counter
   │
   ▼
destroy()
   ├── barrier 同步
   ├── 关闭 remote IPC
   ├── nvshmem::free(rdma_buffer_ptr)
   ├── cudaFree(workspace)
   └── cudaFreeHost(moe_recv_counter)DeepEP

DeepEP v2

V2 相比 V1 有很大升级，哪怕把 V1 看熟了，V2 也要重新学。

Buffer 进化为 ElasticBuffer，对外只暴露 dispatch 和 combine 两个方法。两种模式：

• direct 模式
• hybrid 模式

direct 模式

direct 模式把所有 EP rank 展平成一个逻辑通信域：有 NVLink 就直接写远端地址；没有就走 Gin RDMA（放到本地 send buffer，再发起 RDMA put）。相比 v1 的多机路径，不再有“先节点内聚合 → 再跨节点 → 再节点内转发”的层级逻辑。

dispatch kernel 里设置两种 warps：

• notify warps：替代 v1 的 get_dispatch_layout，从 topk_idx 即时计算 is_token_in_rank。
• dispatch warps：执行实际搬运。

epilogue 负责把通信 buffer 里的数据整理到 tensor 中。同样会保存 handle，供 combine 做反向路由。

v2 还在通信层面做了去重：对当前 token，不再以“每个要去的 expert”为单位发送，而是以“要去的 rank”为单位，发送一份完整 activation。同一 destination rank 只发一份。

combine 使用 handle，把 rank-level partial result 写回 source rank。为什么叫 partial result？因为一个 source token 可能从同一个 rank 的多个 expert 收到结果，先在这个 rank 内部做一次 partial reduce，再把结果发去做 combine。epilogue 按不同 destination rank 做最终归约。

hybrid 模式

hybrid 模式回到多机多卡的层级模式：

1. 先按“目标节点”去重，跨机走 RDMA（scaleout）。
2. 再在目标节点内按“目标 GPU”去重，走 NVLink（scaleup）。

其中 scaleup 指单机内有多少个 rank，scaleout 指有多少台机器。

dispatch 有三类 warps：

• notify warps（同上）
• scaleout warps（先按目标节点去重跨 RDMA）
• forward warps（再在目标节点内按目标 GPU 去重走 NVLink）

为此引入三类 buffer：scaleout_send_buffer、scaleout_recv_buffer、scaleup_buffer。

combine 在跨节点场景下：先在节点内本地规约，再通过 RDMA 发送，最后 source 端做最终归约。

MoonEP

MoonEP 的特点是完美的计算负载均衡。它通过在线搬专家的方式，减少计算上的瓶颈——本质上是用通信换计算。适用于通信快、计算慢的场景，能有效缓解“计算慢”这块短板。

小结

参考资料

1. DeepSeek DeepEP 仓库
2. DeepSeek 官方技术报告（DeepEP 相关章节）