deepseek-ai/DeepEP
GitHub: deepseek-ai/DeepEP
DeepEP 是面向混合专家模型专家并行的高效通信库，提供高吞吐和低延迟的 all-to-all GPU kernel，专为大规模 MoE 模型训练与推理场景优化。
Stars: 9214 | Forks: 1175
# DeepEP
DeepEP 是一个专为混合专家模型和专家并行量身定制的通信库。它提供了高吞吐量和低延迟的 all-to-all GPU kernel（也称为 MoE dispatch 和 combine）。该库还支持低精度运算，包括 FP8。
为了与 [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) 论文中提出的 group-limited gating 算法保持一致，DeepEP 提供了一组针对非对称域带宽转发（例如，从 NVLink 域转发到 RDMA 域）优化的 kernel。这些 kernel 具有高吞吐量，适用于训练和推理 prefilling 任务。此外，它们还支持 SM（流多处理器）数量控制。
对于延迟敏感的推理解码，DeepEP 包含了一组带有纯 RDMA 的低延迟 kernel，以最大程度地减少延迟。该库还引入了一种基于 hook 的通信-计算重叠方法，该方法不占用任何 SM 资源。
注意：该库中的实现可能与 [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) 论文存在一些细微差异。
## 性能
### 使用 NVLink 和 RDMA 转发的普通 kernel
我们在 H800（约 160 GB/s NVLink 最大带宽）上测试了普通 kernel，每个节点连接一张 CX7 InfiniBand 400 Gb/s RDMA 网卡（约 50 GB/s 最大带宽）。我们遵循 DeepSeek-V3/R1 预训练设置（每个 batch 4096 个 token，7168 隐藏层维度，top-4 组，top-8 专家，FP8 dispatch 和 BF16 combine）。
| 类型 | Dispatch #EP | 瓶颈带宽 | Combine #EP | 瓶颈带宽 |
|:---------:|:------------:|:--------------------:|:-----------:|:--------------------:|
| 节点内 | 8 | 153 GB/s (NVLink) | 8 | 158 GB/s (NVLink) |
| 节点间 | 16 | 43 GB/s (RDMA) | 16 | 43 GB/s (RDMA) |
| 节点间 | 32 | 58 GB/s (RDMA) | 32 | 57 GB/s (RDMA) |
| 节点间 | 64 | 51 GB/s (RDMA) | 64 | 50 GB/s (RDMA) |
**新闻 (2025.04.22)**：借助腾讯网络平台部门的优化，性能提升了高达 30%，详情请参见 [#130](https://github.com/deepseek-ai/DeepEP/pull/130)。感谢贡献！
### 使用纯 RDMA 的低延迟 kernel
我们在 H800 上测试了低延迟 kernel，每个节点连接一张 CX7 InfiniBand 400 Gb/s RDMA 网卡（约 50 GB/s 最大带宽）。我们遵循典型的 DeepSeek-V3/R1 生产设置（每个 batch 128 个 token，7168 隐藏层维度，top-8 专家，FP8 dispatch 和 BF16 combine）。
| Dispatch #EP | 延迟 | RDMA 带宽 | Combine #EP | 延迟 | RDMA 带宽 |
|:------------:|:-------:|:--------------:|:-----------:|:-------:|:--------------:|
| 8 | 77 us | 98 GB/s | 8 | 114 us | 127 GB/s |
| 16 | 118 us | 63 GB/s | 16 | 195 us | 74 GB/s |
| 32 | 155 us | 48 GB/s | 32 | 273 us | 53 GB/s |
| 64 | 173 us | 43 GB/s | 64 | 314 us | 46 GB/s |
| 128 | 192 us | 39 GB/s | 128 | 369 us | 39 GB/s |
| 256 | 194 us | 39 GB/s | 256 | 360 us | 40 GB/s |
**新闻 (2025.06.05)**：低延迟 kernel 现已尽可能利用 NVLink，详情请参见 [#173](https://github.com/deepseek-ai/DeepEP/pull/173)。感谢贡献！
## 快速开始
### 环境要求
- Ampere (SM80)、Hopper (SM90) GPU，或其他支持 SM90 PTX ISA 的架构
- Python 3.8 及以上版本
- CUDA 版本
- SM80 GPU 需要 CUDA 11.0 及以上版本
- SM90 GPU 需要 CUDA 12.3 及以上版本
- PyTorch 2.1 及以上版本
- 节点内通信需使用 NVLink
- 节点间通信需使用 RDMA 网络
### 下载并安装 NVSHMEM 依赖
DeepEP 还依赖于 NVSHMEM。请参考我们的 [NVSHMEM 安装指南](third-party/README.md)获取说明。
### 开发
```
# 构建并为 SO 文件创建 symbolic links
NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py build
# 您可以根据自己的平台修改特定的 SO 名称
ln -s build/lib.linux-x86_64-cpython-38/deep_ep_cpp.cpython-38-x86_64-linux-gnu.so
# 运行 test cases
# 注意：您可以修改 `tests/utils.py` 中的 `init_dist` 函数
# 根据您自己的集群设置，并启动到多个节点
python tests/test_intranode.py
python tests/test_internode.py
python tests/test_low_latency.py
```
### 安装
```
NVSHMEM_DIR=/path/to/installed/nvshmem python setup.py install
```
#### 安装环境变量
- `NVSHMEM_DIR`：NVSHMEM 目录的路径，如果未指定，将禁用所有节点间和低延迟功能
- `DISABLE_SM90_FEATURES`：0 或 1，是否禁用 SM90 功能，对于 SM90 设备或 CUDA 11 是必选项
- `TORCH_CUDA_ARCH_LIST`：目标架构列表，例如 `TORCH_CUDA_ARCH_LIST="9.0"`
- `DISABLE_AGGRESSIVE_PTX_INSTRS`：0 或 1，是否禁用激进的 load/store 指令，有关详细信息请参见 [未定义行为 PTX 用法](#undefined-behavior-ptx-usage)
然后，在您的 Python 项目中导入 `deep_ep`，即可开始使用！
## 网络配置
DeepEP 已在 InfiniBand 网络上进行了全面测试。然而，理论上它也兼容 RDMA over Converged Ethernet (RoCE)。
### 流量隔离
InfiniBand 通过 Virtual Lanes (VL) 支持流量隔离。
为了防止不同类型流量之间的干扰，我们建议将工作负载分配到不同的虚拟通道，具体如下：
- 使用普通 kernel 的工作负载
- 使用低延迟 kernel 的工作负载
- 其他工作负载
对于 DeepEP，您可以通过设置 `NVSHMEM_IB_SL` 环境变量来控制虚拟通道分配。
### 自适应路由
自适应路由是 InfiniBand 交换机提供的一种高级路由功能，可以将流量均匀分布在多条路径上。启用自适应路由可以完全消除由路由冲突引起的网络拥塞，但同时也会引入额外的延迟。为了获得最佳性能，我们建议进行以下配置：
- 在网络负载较重的环境中启用自适应路由
- 在网络负载较轻的环境中启用静态路由
### 拥塞控制
拥塞控制已被禁用，因为我们在生产环境中未观察到明显的拥塞现象。
## 接口与示例
### 模型训练或推理 prefilling 的示例用法
普通 kernel 可用于模型训练或推理 prefilling 阶段（不包含反向传播部分），如下面的示例代码所示。
```
import torch
import torch.distributed as dist
from typing import List, Tuple, Optional, Union
from deep_ep import Buffer, EventOverlap
# Communication buffer（将在运行时分配）
_buffer: Optional[Buffer] = None
# 设置要使用的 SM 数量
# 注意：这是一个静态变量
Buffer.set_num_sms(24)
# 您可以在 framework 初始化时调用此函数
def get_buffer(group: dist.ProcessGroup, hidden_bytes: int) -> Buffer:
global _buffer
# NOTES: you may also replace `get_*_config` with your auto-tuned results via all the tests
num_nvl_bytes, num_rdma_bytes = 0, 0
for config in (Buffer.get_dispatch_config(group.size()), Buffer.get_combine_config(group.size())):
num_nvl_bytes = max(config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes)
num_rdma_bytes = max(config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes)
# Allocate a buffer if not existed or not enough buffer size
if _buffer is None or _buffer.group != group or _buffer.num_nvl_bytes < num_nvl_bytes or _buffer.num_rdma_bytes < num_rdma_bytes:
_buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes)
return _buffer
def get_hidden_bytes(x: torch.Tensor) -> int:
t = x[0] if isinstance(x, tuple) else x
return t.size(1) * max(t.element_size(), 2)
def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
topk_idx: torch.Tensor, topk_weights: torch.Tensor,
num_experts: int, previous_event: Optional[EventOverlap] = None) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, List, Tuple, EventOverlap]:
# NOTES: an optional `previous_event` means a CUDA event captured that you want to make it as a dependency
# of the dispatch kernel, it may be useful with communication-computation overlap. For more information, please
# refer to the docs of `Buffer.dispatch`
global _buffer
# Calculate layout before actual dispatch
num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \
_buffer.get_dispatch_layout(topk_idx, num_experts,
previous_event=previous_event, async_finish=True,
allocate_on_comm_stream=previous_event is not None)
# Do MoE dispatch
# NOTES: the CPU will wait for GPU's signal to arrive, so this is not compatible with CUDA graph
# Unless you specify `num_worst_tokens`, but this flag is for intranode only
# For more advanced usages, please refer to the docs of the `dispatch` function
recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event = \
_buffer.dispatch(x, topk_idx=topk_idx, topk_weights=topk_weights,
num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert,
previous_event=previous_event, async_finish=True,
allocate_on_comm_stream=True)
# For event management, please refer to the docs of the `EventOverlap` class
return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event
def dispatch_backward(grad_recv_x: torch.Tensor, grad_recv_topk_weights: torch.Tensor, handle: Tuple) -> \
Tuple[torch.Tensor, torch.Tensor, EventOverlap]:
global _buffer
# The backward process of MoE dispatch is actually a combine
# For more advanced usages, please refer to the docs of the `combine` function
combined_grad_x, combined_grad_recv_topk_weights, event = \
_buffer.combine(grad_recv_x, handle, topk_weights=grad_recv_topk_weights, async_finish=True)
# For event management, please refer to the docs of the `EventOverlap` class
return combined_grad_x, combined_grad_recv_topk_weights, event
def combine_forward(x: torch.Tensor, handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \
Tuple[torch.Tensor, EventOverlap]:
global _buffer
# Do MoE combine
# For more advanced usages, please refer to the docs of the `combine` function
combined_x, _, event = _buffer.combine(x, handle, async_finish=True, previous_event=previous_event,
allocate_on_comm_stream=previous_event is not None)
# For event management, please refer to the docs of the `EventOverlap` class
return combined_x, event
def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
handle: Tuple, previous_event: Optional[EventOverlap] = None) -> \
Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]:
global _buffer
# The backward process of MoE combine is actually a dispatch
# For more advanced usages, please refer to the docs of the `dispatch` function
grad_x, _, _, _, _, event = _buffer.dispatch(grad_combined_x, handle=handle, async_finish=True,
previous_event=previous_event,
allocate_on_comm_stream=previous_event is not None)
# For event management, please refer to the docs of the `EventOverlap` class
return grad_x, event
```
此外，在 dispatch 函数内部，我们可能不知道当前 rank 会接收多少个 token。因此会涉及一个隐式的 CPU 等待 GPU 接收计数信号的过程，如下图所示。
### 推理解码的示例用法
低延迟 kernel 可用于推理解码阶段，如下面的示例代码所示。
```
import torch
import torch.distributed as dist
from typing import Tuple, Optional
from deep_ep import Buffer
# Communication buffer（将在运行时分配）
# 注意：low-latency kernels 没有提供 SM 控制 API
_buffer: Optional[Buffer] = None
# 您可以在 framework 初始化时调用此函数
def get_buffer(group: dist.ProcessGroup, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> Buffer:
# NOTES: the low-latency mode will consume much more space than the normal mode
# So we recommend that `num_max_dispatch_tokens_per_rank` (the actual batch size in the decoding engine) should be less than 256
global _buffer
num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, group.size(), num_experts)
# Allocate a buffer if not existed or not enough buffer size
if _buffer is None or _buffer.group != group or not _buffer.low_latency_mode or _buffer.num_rdma_bytes < num_rdma_bytes:
# NOTES: for the best performance, the QP number **must** be equal to the number of the local experts
assert num_experts % group.size() == 0
_buffer = Buffer(group, 0, num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_experts // group.size())
return _buffer
def low_latency_dispatch(hidden_states: torch.Tensor, topk_idx: torch.Tensor, num_max_dispatch_tokens_per_rank: int, num_experts: int):
global _buffer
# Do MoE dispatch, compatible with CUDA graph (but you may restore some buffer status once you replay)
recv_hidden_states, recv_expert_count, handle, event, hook = \
_buffer.low_latency_dispatch(hidden_states, topk_idx, num_max_dispatch_tokens_per_rank, num_experts,
async_finish=False, return_recv_hook=True)
# NOTES: the actual tensor will not be received only if you call `hook()`,
# it is useful for double-batch overlapping, but **without any SM occupation**
# If you don't want to overlap, please set `return_recv_hook=False`
# Later, you can use our GEMM library to do the computation with this specific format
return recv_hidden_states, recv_expert_count, handle, event, hook
def low_latency_combine(hidden_states: torch.Tensor,
topk_idx: torch.Tensor, topk_weights: torch.Tensor, handle: Tuple):
global _buffer
# Do MoE combine, compatible with CUDA graph (but you may restore some buffer status once you replay)
combined_hidden_states, event_overlap, hook = \
_buffer.low_latency_combine(hidden_states, topk_idx, topk_weights, handle,
async_finish=False, return_recv_hook=True)
# NOTES: the same behavior as described in the dispatch kernel
return combined_hidden_states, event_overlap, hook
```
对于双 micro-batch 重叠，您可以参考下图。通过我们的接收 hook 接口，RDMA 网络流量在后台进行，不会占用计算部分的任何 GPU SM。但请注意，重叠的部分是可以调整的，即 attention/dispatch/MoE/combine 这四个部分可能没有完全相同的执行时间。您可以根据工作负载调整阶段设置。
## 路线图
- [x] AR 支持
- [x] 重构低延迟模式 AR 代码
- [x] A100 支持 (仅限节点内)
- [x] 低延迟 dispatch kernel 支持 BF16
- [x] 节点内低延迟 kernel 支持 NVLink 协议
- [ ] 使用 TMA copy 代替 LD/ST
- [x] 节点内 kernel
- [ ] 节点间 kernel
- [ ] 低延迟 kernel
- [ ] 无 SM kernel 及重构
- [ ] 完全移除未定义行为的 PTX 指令
## 注意事项
#### 更简便的潜在整体设计
当前的 DeepEP 实现使用队列作为通信缓冲区，这节省了内存，但也引入了复杂性和潜在的死锁问题。如果您打算基于 DeepEP 实现自己的版本，可以考虑使用按最大容量分配的固定大小缓冲区，以获得更好的简便性和性能。有关此替代方案的详细讨论，请参见 https://github.com/deepseek-ai/DeepEP/issues/39。
#### 未定义行为 PTX 用法
- 为了追求极致性能，我们发现并使用了一种未定义行为的 PTX 用法：使用只读 PTX `ld.global.nc.L1::no_allocate.L2::256B` 来**读取易失性数据**。PTX 修饰符 `.nc` 表示使用了非一致性缓存。但在 Hopper 架构上，已通过测试验证结合 `.L1::no_allocate` 可以保证正确性，并且性能会好得多。我们猜测原因可能是：非一致性缓存与 L1 统一，且 L1 修饰符不仅仅是一个提示，而是一个强选项，因此 L1 中不会存在脏数据，从而保证了正确性。
- 起初，由于 NVCC 无法自动展开易失性读取 PTX，我们尝试使用 `__ldg`（即 `ld.nc`）。即使与手动展开的易失性读取相比，它的速度也要快得多（这可能是由于额外的编译器优化所致）。然而，这可能导致结果不正确或产生脏数据。在查阅 PTX 文档后，我们发现在 Hopper 架构上 L1 和非一致性缓存是统一的。我们推测 `.L1::no_allocate` 可能会解决这个问题，从而促成了这一发现。
- 如果您发现 kernel 在某些其他平台上无法正常工作，您可以在 `setup.py` 中添加 `DISABLE_AGGRESSIVE_PTX_INSTRS=1` 以禁用此功能，或者提交 issue。
#### 在您的集群上进行自动调优
为了在您的集群上获得更好的性能，我们建议运行所有测试并使用最佳的自动调优配置。默认配置是在 DeepSeek 的内部集群上优化的。
## 许可证
本代码仓库基于 [MIT 许可证](LICENSE) 发布，但引用了 NVSHMEM 的代码（包括 `csrc/kernels/ibgda_device.cuh` 和 `third-party/nvshmem.patch`）除外，这些代码受 [NVSHMEM SLA](https://docs.nvidia.com/nvshmem/api/sla.html) 约束。
## 实验性分支
- [Zero-copy](https://github.com/deepseek-ai/DeepEP/pull/453)
- 移除 PyTorch tensor 和通信缓冲区之间的拷贝，显著降低了普通 kernel 的 SM 使用量
- 此 PR 由 **腾讯网络平台部门** 提交
- [Eager](https://github.com/deepseek-ai/DeepEP/pull/437)
- 使用低延迟协议，消除了由 RDMA 原子操作引入的额外 RTT 延迟
- [Hybrid-EP](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep)
- 使用 TMA 指令实现的新后端，以最小化 SM 占用并支持更大的 NVLink 域
- 针对单 batch 场景的细粒度通信-计算重叠
- 在非 NVLink 环境中支持 PCIe kernel
- 支持 NVFP4 数据类型
- [AntGroup-Opt](https://github.com/deepseek-ai/DeepEP/tree/antgroup-opt)
- 此优化系列由 **蚂蚁集团网络平台部门** 提交
- [Normal-SMFree](https://github.com/deepseek-ai/DeepEP/pull/347) 通过将通信 kernel 执行与 NIC token 传输解耦，消除 RDMA 路径上的 SM 占用，释放 SM 用于计算
- [LL-SBO](https://github.com/deepseek-ai/DeepEP/pull/483) 通过信号机制将 Down GEMM 计算与 Combine Send 通信重叠，以减少端到端延迟
- [LL-Layered](https://github.com/deepseek-ai/DeepEP/pull/500) 使用轨道优化转发和数据合并来优化跨节点 LL 算子通信，以降低延迟
- [Mori-EP](https://github.com/deepseek-ai/DeepEP/tree/mori-ep)
- 由 [MORI](https://github.com/ROCm/mori) 后端支持的 ROCm/AMD GPU 支持（低延迟模式）
## 社区 Fork
- [uccl/uccl-ep](https://github.com/uccl-project/uccl/tree/main/ep) - 支持在异构 GPU（如 Nvidia、AMD）和网卡（如 EFA、Broadcom、CX7）上运行 DeepEP
- [Infrawaves/DeepEP_ibrc_dual-ports_multiQP](https://github.com/Infrawaves/DeepEP_ibrc_dual-ports_multiQP) - 在 IBRC 传输中增加了多 QP 解决方案和双端口网卡支持
- [antgroup/DeepXTrace](
标签：All-to-All, CUDA, DeepSeek, DeepSeek-R1, DeepSeek-V3, DLL 劫持, FP8, GPU内核, GPU通信, HPC, InfiniBand, MoE, NVLink, RDMA, 专家并行, 人工智能, 低延迟, 低精度计算, 凭据扫描, 分布式训练, 大语言模型, 推理加速, 深度学习, 混合专家模型, 熵值分析, 用户模式Hook绕过, 算子优化, 网络性能优化, 逆向工具, 通信库, 预训练, 高吞吐量, 高性能计算