DeepEP README

The project's README from github.com/deepseek-ai/DeepEP. As one markdown file: README.md.

DeepEP

DeepEP (DeepEveryParallel) is a high-performance communication library for modern machine learning training and inference. The library currently focuses on expert parallelism (EP) — providing high-throughput and low-latency all-to-all GPU kernels (MoE dispatch and combine) with low-precision support including FP8 — while also offering experimental primitives for pipeline parallelism (PP), context parallelism (CP), and remote memory access (Engram), all designed for zero or minimal SM occupation. All kernels are compiled at runtime via a lightweight Just-In-Time (JIT) module, requiring no CUDA compilation during installation.

Despite its lightweight design, DeepEP's performance matches or exceeds hardware bandwidth limits across various configurations.

News

New features

Notes

Still on-going features

For the legacy V1 documentation (NVSHMEM-based), see docs/legacy.md.

Performance

Following V3's configuration, we tested with 8K tokens per batch, 7168 hidden dimensions, top 8 experts, FP8 dispatching, and BF16 combining, and obtained the following results:

Arch NIC type Topo Dispatch Bottleneck Bandwidth Combine Bottleneck Bandwidth #SMs
SM90 CX7 EP 8 x 2 90 GB/s (RDMA) 81 GB/s (RDMA) 12
SM90 CX7 EP 8 x 4 61 GB/s (RDMA) 61 GB/s (RDMA) 6
SM100 CX7 EP 8 x 2 90 GB/s (RDMA) 91 GB/s (RDMA) 12
SM100 N/A EP 8 726 GB/s (NVLink) 740 GB/s (NVLink) 64 (Max perf)
SM100 N/A EP 8 643 GB/s (NVLink) 675 GB/s (NVLink) 24 (Min #SM)

Notes: the results are logical bandwidth. For example, under the EP 8 x 2 case, 90 GB/s actually contains local rank traffic.

Comparing with V1, V2 achieves up to 1.3x peak performance, while saving up to 4x SM count.

We omit results for larger EP configurations for the time being, but encourage interested users to benchmark them directly. Based on our internal experience, we expect the kernel to continue saturating hardware bandwidth at scale.

For V1 performance data, see docs/legacy.md.

Quick start

Requirements

Install NCCL dependency

We recommend using pip to install NCCL so that DeepEP can automatically locate it within the Python environment. You can install it using the following command:

pip install "nvidia-nccl-cu13>=2.30.4" --no-deps

Install NVSHMEM dependency

DeepEP also depends on NVSHMEM to provide support for legacy methods. Please refer to our NVSHMEM Installation Guide for instructions.

Development

# Build and make symbolic links for SO files
python setup.py build
# You may modify the specific SO names according to your own platform
ln -s build/lib.linux-x86_64-cpython-38/deep_ep_cpp.cpython-38-x86_64-linux-gnu.so

# Run test cases
# NOTES: you may modify the `init_dist` function in `tests/utils/envs.py`
# according to your own cluster settings, and launch into multiple nodes
python tests/elastic/test_ep.py
python tests/elastic/test_agrs.py
python tests/elastic/test_engram.py
python tests/elastic/test_pp.py

Installation

python setup.py install

Then, import deep_ep in your Python project, and enjoy!

Interfaces and examples

Buffer initialization

In V2, all EP operations — high-throughput and low-latency — are unified under a single ElasticBuffer interface. The buffer can be initialized by specifying MoE settings directly, and the optimal SM and QP counts are calculated analytically.

import torch
import torch.distributed as dist
from typing import Optional

from deep_ep import ElasticBuffer

# Communication buffer (will allocate at runtime)
_buffer: Optional[ElasticBuffer] = None

# Number of SMs to use for communication kernels (will be set at buffer creation)
_num_comm_sms: int = 0


def get_buffer(group: dist.ProcessGroup,
               num_max_tokens_per_rank: int,
               hidden: int,
               num_topk: int,
               num_experts: int,
               use_fp8_dispatch: bool = False) -> ElasticBuffer:
    """Initialize or retrieve the ElasticBuffer for EP communication."""
    global _buffer, _num_comm_sms

    # Check if we can reuse the existing buffer
    required_bytes = ElasticBuffer.get_buffer_size_hint(
        group, num_max_tokens_per_rank, hidden,
        num_topk=num_topk, use_fp8_dispatch=use_fp8_dispatch,
    )
    if _buffer is not None and _buffer.group == group and _buffer.num_bytes >= required_bytes:
        return _buffer

    # Allocate a new buffer with MoE settings
    # NOTES: V2 buffer size consumption is larger than V1
    _buffer = ElasticBuffer(
        group,
        num_max_tokens_per_rank=num_max_tokens_per_rank,
        hidden=hidden,
        num_topk=num_topk,
        use_fp8_dispatch=use_fp8_dispatch,
    )

    # V2 analytically calculates the optimal SM count — no more auto-tuning needed
    # You may also specify `num_sms` manually in dispatch/combine calls to override
    _num_comm_sms = _buffer.get_theoretical_num_sms(num_experts, num_topk)

    return _buffer

Example use in model training or inference prefilling

V2 unifies the dispatch and combine APIs into a single ElasticBuffer interface. The example below shows how to use them for training (with backward passes) or inference prefilling.

import torch
import torch.distributed as dist
from typing import Tuple, Union

from deep_ep import ElasticBuffer, EPHandle, EventOverlap


def dispatch_forward(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                     topk_idx: torch.Tensor, topk_weights: torch.Tensor,
                     num_experts: int,
                     num_max_tokens_per_rank: int,
                     expert_alignment: int = 1) -> \
        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
              torch.Tensor, torch.Tensor, EPHandle, EventOverlap]:
    """
    MoE dispatch: route tokens to the corresponding experts across all ranks.
    Supports both BF16 and FP8 (x as a tuple of [data, scale_factors]) inputs.
    """
    global _buffer, _num_comm_sms

    recv_x, recv_topk_idx, recv_topk_weights, handle, event = _buffer.dispatch(
        x,
        topk_idx=topk_idx,
        topk_weights=topk_weights,
        num_experts=num_experts,
        num_max_tokens_per_rank=num_max_tokens_per_rank,
        expert_alignment=expert_alignment,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    # `handle` contains routing metadata for the subsequent combine call
    # `handle.num_recv_tokens_per_expert_list` provides per-expert token counts for GEMM
    # Use `event.current_stream_wait()` to synchronize the compute stream before using results
    return recv_x, recv_topk_idx, recv_topk_weights, handle, event


def dispatch_backward(grad_recv_x: torch.Tensor,
                      grad_recv_topk_weights: torch.Tensor,
                      handle: EPHandle) -> Tuple[torch.Tensor, torch.Tensor, EventOverlap]:
    """The backward pass of MoE dispatch is actually a combine."""
    global _buffer, _num_comm_sms

    combined_grad_x, combined_grad_topk_weights, event = _buffer.combine(
        grad_recv_x,
        handle=handle,
        topk_weights=grad_recv_topk_weights,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    return combined_grad_x, combined_grad_topk_weights, event


def combine_forward(x: torch.Tensor,
                    handle: EPHandle) -> Tuple[torch.Tensor, EventOverlap]:
    """MoE combine: reduce expert outputs back to their original ranks."""
    global _buffer, _num_comm_sms

    combined_x, _, event = _buffer.combine(
        x,
        handle=handle,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    return combined_x, event


def combine_backward(grad_combined_x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                     handle: EPHandle) -> \
        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], EventOverlap]:
    """The backward pass of MoE combine is actually a dispatch."""
    global _buffer, _num_comm_sms

    grad_x, _, _, _, event = _buffer.dispatch(
        grad_combined_x,
        handle=handle,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    return grad_x, event

For communication-computation overlap, use the EventOverlap interface to manage dependencies between the communication stream and the compute stream:

# After dispatch, overlap computation while communication is in-flight
recv_x, recv_topk_idx, recv_topk_weights, handle, event = dispatch_forward(...)

# ... do some independent computation here ...

# Wait for communication to finish before using results
event.current_stream_wait()

# Now safe to use recv_x, recv_topk_idx, recv_topk_weights

Example use in inference decoding

For inference decoding, the same ElasticBuffer is used. The handle-caching pattern allows reusing routing metadata across iterations when the gating decisions remain unchanged, avoiding redundant CPU synchronization.

import torch
from typing import Tuple, Optional, Union

from deep_ep import ElasticBuffer, EPHandle, EventOverlap


def decode_dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                    topk_idx: torch.Tensor, topk_weights: torch.Tensor,
                    num_experts: int,
                    num_max_tokens_per_rank: int,
                    cached_handle: Optional[EPHandle] = None) -> \
        Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
              torch.Tensor, torch.Tensor, EPHandle, EventOverlap]:
    """
    MoE dispatch for inference decoding.
    If `cached_handle` is provided, the layout is reused without CPU synchronization.
    """
    global _buffer, _num_comm_sms

    if cached_handle is not None:
        # Reuse cached handle: skip layout recomputation and CPU sync
        recv_x, _, _, handle, event = _buffer.dispatch(
            x,
            handle=cached_handle,
            num_sms=_num_comm_sms,
            async_with_compute_stream=True,
        )
        return recv_x, cached_handle.topk_idx, None, handle, event

    recv_x, recv_topk_idx, recv_topk_weights, handle, event = _buffer.dispatch(
        x,
        topk_idx=topk_idx,
        topk_weights=topk_weights,
        num_experts=num_experts,
        num_max_tokens_per_rank=num_max_tokens_per_rank,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    return recv_x, recv_topk_idx, recv_topk_weights, handle, event


def decode_combine(x: torch.Tensor,
                   handle: EPHandle) -> Tuple[torch.Tensor, EventOverlap]:
    """MoE combine for inference decoding."""
    global _buffer, _num_comm_sms

    combined_x, _, event = _buffer.combine(
        x,
        handle=handle,
        num_sms=_num_comm_sms,
        async_with_compute_stream=True,
    )

    return combined_x, event

Environment variables

The library provides some environment variables, which may be useful:

Some environment variables are persistent: they are captured at build time and baked into the installed package as default values. At import time, these defaults are applied automatically unless overridden by current environment variables. The persistent variables are: EP_JIT_CACHE_DIR, EP_JIT_PRINT_COMPILER_COMMAND, EP_NUM_TOPK_IDX_BITS, EP_NCCL_ROOT_DIR.

For additional details, please refer to the test code or review the corresponding Python documentation.

Network configurations

DeepEP is fully tested with InfiniBand networks. However, it is theoretically compatible with RDMA over Converged Ethernet (RoCE) as well.

Traffic isolation

Traffic isolation is supported by InfiniBand through Virtual Lanes (VL).

To prevent interference between different types of traffic, we recommend segregating workloads across different virtual lanes as follows:

For DeepEP V2, you can control the virtual lane assignment by setting the sl_idx argument or the EP_OVERRIDE_RDMA_SL environment variable.

Adaptive routing

Adaptive routing is an advanced routing feature provided by InfiniBand switches that can evenly distribute traffic across multiple paths. Even though adaptive routing introduces additional latency, we still recommend enabling it under all network load conditions.

Congestion control

Congestion control is disabled because it hurts maximum bandwidth. If congestion is unavoidable in some scenarios, we recommend assigning those workloads to low-priority virtual lanes.

PCI atomic mode

If the hardware supports it, we recommend using the following command to set the NIC's PCI_ATOMIC_MODE to improve RDMA atomic operation performance:

sudo mlxconfig -y -d mlx5_$i set PCI_ATOMIC_MODE=4

Experimental branches

Community forks

Acknowledgement

DeepEP V2 is built on top of the NCCL Gin backend. Thanks to @sjeaugey, @pakmarkthub, @sb17v, @xiaofanl-nvidia, and the NCCL team for their support!

License

This code repository is released under the MIT License.

Citation

@misc{deepep2025,
      title={DeepEP: an efficient expert-parallel communication library},
      author={Chenggang Zhao and Shangyan Zhou and Liyue Zhang and Chengqi Deng and Zhean Xu and Yuxuan Liu and Kuai Yu and Jiashi Li and Liang Zhao},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/DeepEP}},
}

DeepEP (deepseek-ai/DeepEP)

Searching Database…

10.3K ·GitHub
DeepEP
这是什么？DeepEP 是一个面向 MoE（混合专家）模型的高性能通信库，提供快速的全对全 GPU 内核，用于机器学习训练和推理中的专家并行分发（dispatch）与合并（combine）操作。
这是训练框架，还是可以接入我现有 PyTorch 环境的库？它是一个库，而不是完整的框架：你可以把它导入到 PyTorch 的训练或推理代码中，以加速专家并行通信。内核通过轻量的 JIT 模块在运行时编译，而不是提前进行 CUDA 编译。
我该如何安装并开始使用？先用 pip 安装 NCCL 依赖，让 DeepEP 能自动找到它，然后在 Python 项目中直接 import deepep 即可。内核会在首次使用时进行 JIT 编译，因此不需要单独的 CUDA 构建步骤。
我需要什么样的 GPU 和软件版本？它需要 Hopper（SM90）GPU，或其他兼容 SM90-PTX 的硬件，还需要 Python 3.8+、CUDA 12.3+、PyTorch 2.10+ 和 NCCL 2.30.4+。节点内通信需要 NVLink，节点间通信需要 RDMA 网络。
它能在 Hopper 之外的 GPU 上运行吗？V1 的所有功能在 V2 中仍然都支持吗？官方只支持 Hopper（SM90）GPU；更广泛的 GPU 支持（例如 AMD）仅通过社区分支提供。另外，V2 不再支持 V1 中的 0-SM 低延迟 RDMA EP 模式，而 Engram、PP、CP 被标记为实验性功能。
V1 和 V2 有什么区别？V2 是完全重写的版本，所使用的流式多处理器（SM）数量减少了好几倍；通信后端从 NVSHMEM 切换为 NCCL Gin，并将高吞吐和低延迟的 API 统一为单一的 ElasticBuffer 接口，支持扩展到 EP2048。
什么是 NCCL Gin 后端？这是 DeepEP V2 用来替代 NVSHMEM 的新型轻量后端。它还允许 DeepEP 复用应用已有的 NCCL 通信器，而不必另外创建新的通信器。
Engram、PP 和 CP 功能是什么？它们是与专家并行一同构建的实验性零 SM 原语：Engram 用于远程内存访问（通过 RDMA），PP 用于流水线并行（通过 RDMA），CP 用于上下文并行（通过 Copy Engine）。
我能在非 NVIDIA 的 GPU 上运行 DeepEP 吗？官方不支持，但社区分支对其做了扩展：ROCm/mori 将后端移植到了 AMD GPU 上，uccl/uccl-ep 则支持在异构 GPU 和网卡之间运行 DeepEP。
What are similar projects?DeepGEMM, DeepEP-Ascend, TensorRT-LLM, DeepGEMM-Ascend, taskflow, tilelang, DeepSpeed, ColossalAI
Which topics is it in?hpc, distributed-computing, header-only, jit, gpu-programming, mixture-of-experts, moe, high-performance, gpu-acceleration, distributed-systems, inference, nvidia
Creator
deepseek-ai
Size
3.8 MB
Stars
10.3K
Updated
CommitOct 15Releasev1.2.1Sep 16
Age
1y
Commits · 90 Days