deepseek-ai/DeepEP
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepEP.git synced 2026-11-05 00:47:32 +00:00
Code Issues 308 Projects Packages Wiki Activity

[PR #521] Add NCCL Communication Backend for DeepEP #607

New issue
Open
opened 2026-10-27 21:04:30 +00:00 by bsakaguchi · 0 comments
bsakaguchi commented 2026-10-27 21:04:30 +00:00
Owner
Copy link

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/521
Author: @aamirshafi
Created: 12/15/2025
Status: 🔄 Open

Base: main ← Head: nccl

📝 Commits (10+)

  • 2527ce7 Support cuMem API in cross process shared memory management (#217)
  • 57d82f6 Enable CUDA Graph for internode dispatch (#438)
  • 7ffa208 Update README.me for Hybrid-EP Experimental Branche (#502)
  • 4cd23f0 Update README for an experimental optimization branch (#512)
  • 10d110a fix: Support multiple PyTorch distributed backends
  • b15f9a0 Added NCCL GIN backend (only able to run init/finalize).
  • 90f30bd adding support for LL dispatch. this is WIP.
  • 079463d Updating the scripts to work with slurm
  • b02cc43 Changing the scripts to run with MPI
  • 6b1b49b Changing the NCCL-GIN backend to enable the use of the high-level API. Removed redundant/dead code.

📊 Changes

22 files changed (+2862 additions, -106 deletions)

View changed files

➕ README-NCCL.md (+168 -0)
📝 csrc/config.hpp (+38 -2)
📝 csrc/deep_ep.cpp (+48 -20)
📝 csrc/deep_ep.hpp (+5 -1)
📝 csrc/kernels/api.cuh (+24 -4)
➕ csrc/kernels/backend_factory.cpp (+113 -0)
➕ csrc/kernels/backend_factory.h (+14 -0)
➕ csrc/kernels/communication_backend.h (+131 -0)
📝 csrc/kernels/configs.cuh (+8 -1)
📝 csrc/kernels/exception.cuh (+1 -0)
📝 csrc/kernels/ibgda_device.cuh (+1 -0)
📝 csrc/kernels/internode.cu (+771 -19)
📝 csrc/kernels/internode_ll.cu (+583 -11)
➕ csrc/kernels/nccl_gin_backend.cu (+461 -0)
➕ csrc/kernels/nccl_gin_backend.h (+101 -0)
📝 csrc/kernels/runtime.cu (+140 -7)
📝 deep_ep/buffer.py (+20 -1)
📝 setup.py (+49 -15)
📝 tests/test_internode.py (+61 -9)
📝 tests/test_intranode.py (+4 -4)

...and 2 more files

📄 Description

Summary

This PR adds NCCL as an additional communication backend for DeepEP, leveraging NCCL's Device API for GPU-initiated
network operations. The integration introduces a CommunicationBackend abstraction that allows users to use NVSHMEM or
NCCL as communication backend. This enables users to select their preferred backend based on deployment requirements.

Build/Runtime Selection of NVSHMEM and NCCL Backends

Both backends are fully supported and can be selected at build/runtime:

Backend Build Runtime Status
NVSHMEM NVSHMEM_DIR=/path/to/nvshmem Default ✅ Existing (unchanged)
NCCL ENABLE_NCCL=1 NCCL_DIR=/path/to/nccl DEEP_EP_BACKEND=nccl ✅ New

Why Support Both NVSHMEM and NCCL Backends:

  • Existing NVSHMEM deployments continue to work without any changes
  • NCCL provides a portable, widely-adopted additional option
  • Users can choose based on their infrastructure and preferences

Unchanged Buffer Interface

The Buffer class interface exported to AI frameworks remains completely unchanged. This means:

  • Existing AI frameworks (e.g., training and inference systems using DeepEP) can utilize NCCL without any code changes
  • The same dispatch(), combine(), low_latency_dispatch(), and low_latency_combine() APIs work transparently with either backend
  • Backend selection is handled via environment variables at build/runtime, not in application code

Communication Backend Interface

We introduce an abstract CommunicationBackend interface that decouples DeepEP kernels from the underlying communication library:

class CommunicationBackend {
public:
    // Initialization/Cleanup
    virtual int init(const std::vector<uint8_t>& root_unique_id, int rank, 
                     int num_ranks, bool low_latency_mode, int qps_per_rank) = 0;
    virtual void finalize() = 0;
    virtual void get_unique_id(void* unique_id) = 0;

    // Synchronization
    virtual void barrier() = 0;

    // Memory Management (RDMA-registered buffers)
    virtual void* alloc(size_t size, size_t alignment) = 0;
    virtual void free(void* ptr) = 0;

    // Query
    virtual int get_rank() const = 0;
    virtual int get_num_ranks() const = 0;
    virtual BackendType get_backend_type() const = 0;
};

Backend Selection:

// Auto-detect from DEEP_EP_BACKEND environment variable
BackendType detect_backend_type();

// Global singleton management
void initialize_backend(BackendType type, ...);
CommunicationBackend* get_backend();
void finalize_backend();

NCCL Backend Integration

The NCCL backend uses the Device API for GPU-initiated network operations (GIN), translating between NVSHMEM's PGAS model and NCCL's window-based model:

Aspect NVSHMEM/IBGDA NCCL
Memory Model Pointer-based addressing Window-based with offset addressing
Data Transfer put_nbi(dst_ptr, src_ptr, count, pe) put(peer, dstWin, dstOff, srcWin, srcOff, bytes)
Synchronization Memory atomics Signal atomics (signal(), readSignal())
Completion Per-QP quiet() Per-context flush()

Key Integration Challenges:

  1. Multi-Communicator Mapping: Each NCCL GIN context wraps a Queue Pair (QP), and NCCL provides 4 GIN contexts per communicator. DeepEP's QP requirements are met using ⌈QPs/4⌉ communicators:

    comm_id = channel_id / 4
ctx_id  = channel_id % 4

  2. Memory Registration: Buffers registered with all communicators; window handles stored in GPU memory for kernel access.

  3. Signal-Based Synchronization: Pre-allocated signal layouts map memory atomics to NCCL signal primitives.

  4. Semantic Preservation: Zero-byte put() with SignalAdd is semantically equivalent to net.signal() but better performing in the current NCCL release.

NCCL GIN Type Selection

For GPU-initiated network operations, NCCL supports multiple backends:

NCCL_GIN_TYPE Backend Description
2 Proxy CPU proxy for network operations
3 GDAKI GPU Direct Async Kernel-Initiated (recommended)

Performance Results

Benchmarked on H100 (900 GB/s NVLink) with 8×400 Gbit/s InfiniBand (~50 GB/s per NIC).

High-Throughput Kernels

4096 tokens, 7168 hidden, top-8 experts, BF16 dispatch, BF16 combine

Backend #EP Dispatch BW Combine BW
NVSHMEM 16 79.7 GB/s 66.4 GB/s
NCCL 16 76.9 GB/s 66.2 GB/s
NVSHMEM 32 62.9 GB/s 62.9 GB/s
NCCL 32 61.7 GB/s 62.3 GB/s
NVSHMEM 64 53.5 GB/s 53.2 GB/s
NCCL 64 52.7 GB/s 52.9 GB/s

Low-Latency Kernels

128 tokens, 7168 hidden, top-8 experts, FP8 dispatch, BF16 combine

Backend #EP Dispatch Latency Dispatch BW Combine Latency Combine BW
NVSHMEM 8 160.7 μs 46.8 GB/s 304.2 μs 47.8 GB/s
NCCL 8 160.8 μs 47.0 GB/s 302.8 μs 47.9 GB/s
NVSHMEM 16 182.3 μs 41.4 GB/s 319.8 μs 45.5 GB/s
NCCL 16 178.6 μs 42.2 GB/s 320.8 μs 45.3 GB/s
NVSHMEM 32 188.7 μs 40.0 GB/s 332.9 μs 43.7 GB/s
NCCL 32 190.1 μs 39.8 GB/s 333.2 μs 43.6 GB/s
NVSHMEM 64 225.1 μs 34.8 GB/s 343.1 μs 42.4 GB/s
NCCL 64 218.9 μs 34.5 GB/s 351.1 μs 41.4 GB/s

Requirements

  • NCCL 2.28.9 (or above) with Device API and GIN support (required)

Tested configurations:

  • CUDA 12.8+
  • Hopper (H100) and Blackwell GPUs

Usage

# Build with NCCL backend
export NCCL_DIR=/path/to/nccl/build
export ENABLE_NCCL=1
python3 setup.py build_ext --inplace && pip install --no-build-isolation .

# Run with NCCL
DEEP_EP_BACKEND=nccl NCCL_GIN_TYPE=3 python3 tests/test_low_latency.py
DEEP_EP_BACKEND=nccl NCCL_GIN_TYPE=3 python3 tests/test_internode.py

Future Work

  1. Reduce Communicator Count: Currently, we initialize multiple NCCL communicators due to the limitation of 4 GIN contexts per communicator. Work is in progress on the NCCL roadmap to increase the number of GIN contexts per communicator, which will simplify the integration.

  2. Unified Backend Interface for NVSHMEM: The current NVSHMEM code does not follow the CommunicationBackend interface in order to stay close to the upstream DeepEP main branch. Migrating NVSHMEM to use this interface should be addressed in consultation with DeepEP maintainers.

References

📄 Paper: GPU-Initiated Networking for NCCL

🔗 NCCL Repository: https://github.com/NVIDIA/nccl

📄 NCCL README: NCCL README

Co-authored with @grtheod.

🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/521 **Author:** [@aamirshafi](https://github.com/aamirshafi) **Created:** 12/15/2025 **Status:** 🔄 Open **Base:** `main` ← **Head:** `nccl` --- ### 📝 Commits (10+) - [`2527ce7`](https://github.com/deepseek-ai/DeepEP/commit/2527ce72c14d4a2bc563beddefa329bba8196ffe) Support cuMem API in cross process shared memory management (#217) - [`57d82f6`](https://github.com/deepseek-ai/DeepEP/commit/57d82f61fb4e1fd562b346c69a8fb9784b83a8aa) Enable CUDA Graph for internode dispatch (#438) - [`7ffa208`](https://github.com/deepseek-ai/DeepEP/commit/7ffa20872f9a8c65670a06990bd35266d91f2073) Update README.me for Hybrid-EP Experimental Branche (#502) - [`4cd23f0`](https://github.com/deepseek-ai/DeepEP/commit/4cd23f0e915f5dc38b31297fdd8caf25a544107d) Update README for an experimental optimization branch (#512) - [`10d110a`](https://github.com/deepseek-ai/DeepEP/commit/10d110a1924c4fbc66086293e971822a9c1cdab5) fix: Support multiple PyTorch distributed backends - [`b15f9a0`](https://github.com/deepseek-ai/DeepEP/commit/b15f9a060843b264f202e5e42e3e03a46fe6d5f8) Added NCCL GIN backend (only able to run init/finalize). - [`90f30bd`](https://github.com/deepseek-ai/DeepEP/commit/90f30bd107db11e57415caa8783d8f5284ac5dbd) adding support for LL dispatch. this is WIP. - [`079463d`](https://github.com/deepseek-ai/DeepEP/commit/079463dc5418b3f56c44642b89ce8d568843512e) Updating the scripts to work with slurm - [`b02cc43`](https://github.com/deepseek-ai/DeepEP/commit/b02cc43933d09f62a579659a8890edfb28e5f983) Changing the scripts to run with MPI - [`6b1b49b`](https://github.com/deepseek-ai/DeepEP/commit/6b1b49ba7e1125a00e3da47f798172fca465b3ce) Changing the NCCL-GIN backend to enable the use of the high-level API. Removed redundant/dead code. ### 📊 Changes **22 files changed** (+2862 additions, -106 deletions) <details> <summary>View changed files</summary> ➕ `README-NCCL.md` (+168 -0) 📝 `csrc/config.hpp` (+38 -2) 📝 `csrc/deep_ep.cpp` (+48 -20) 📝 `csrc/deep_ep.hpp` (+5 -1) 📝 `csrc/kernels/api.cuh` (+24 -4) ➕ `csrc/kernels/backend_factory.cpp` (+113 -0) ➕ `csrc/kernels/backend_factory.h` (+14 -0) ➕ `csrc/kernels/communication_backend.h` (+131 -0) 📝 `csrc/kernels/configs.cuh` (+8 -1) 📝 `csrc/kernels/exception.cuh` (+1 -0) 📝 `csrc/kernels/ibgda_device.cuh` (+1 -0) 📝 `csrc/kernels/internode.cu` (+771 -19) 📝 `csrc/kernels/internode_ll.cu` (+583 -11) ➕ `csrc/kernels/nccl_gin_backend.cu` (+461 -0) ➕ `csrc/kernels/nccl_gin_backend.h` (+101 -0) 📝 `csrc/kernels/runtime.cu` (+140 -7) 📝 `deep_ep/buffer.py` (+20 -1) 📝 `setup.py` (+49 -15) 📝 `tests/test_internode.py` (+61 -9) 📝 `tests/test_intranode.py` (+4 -4) _...and 2 more files_ </details> ### 📄 Description ## Summary This PR adds **NCCL** as an additional communication backend for DeepEP, leveraging NCCL's Device API for GPU-initiated network operations. The integration introduces a `CommunicationBackend` abstraction that allows users to use NVSHMEM or NCCL as communication backend. This enables users to select their preferred backend based on deployment requirements. --- ## Build/Runtime Selection of NVSHMEM and NCCL Backends Both backends are fully supported and can be selected at build/runtime: | Backend | Build | Runtime | Status | |---------|-------|---------|--------| | **NVSHMEM** | `NVSHMEM_DIR=/path/to/nvshmem` | Default | ✅ Existing (unchanged) | | **NCCL** | `ENABLE_NCCL=1 NCCL_DIR=/path/to/nccl` | `DEEP_EP_BACKEND=nccl` | ✅ New | **Why Support Both NVSHMEM and NCCL Backends:** - Existing NVSHMEM deployments continue to work **without any changes** - NCCL provides a **portable, widely-adopted** additional option - Users can choose based on their infrastructure and preferences --- ## Unchanged Buffer Interface The `Buffer` class interface exported to AI frameworks remains **completely unchanged**. This means: - **Existing AI frameworks** (e.g., training and inference systems using DeepEP) can utilize NCCL **without any code changes** - The same `dispatch()`, `combine()`, `low_latency_dispatch()`, and `low_latency_combine()` APIs work transparently with either backend - Backend selection is handled via environment variables at build/runtime, not in application code --- ## Communication Backend Interface We introduce an abstract `CommunicationBackend` interface that decouples DeepEP kernels from the underlying communication library: ```cpp class CommunicationBackend { public: // Initialization/Cleanup virtual int init(const std::vector<uint8_t>& root_unique_id, int rank, int num_ranks, bool low_latency_mode, int qps_per_rank) = 0; virtual void finalize() = 0; virtual void get_unique_id(void* unique_id) = 0; // Synchronization virtual void barrier() = 0; // Memory Management (RDMA-registered buffers) virtual void* alloc(size_t size, size_t alignment) = 0; virtual void free(void* ptr) = 0; // Query virtual int get_rank() const = 0; virtual int get_num_ranks() const = 0; virtual BackendType get_backend_type() const = 0; }; ``` **Backend Selection:** ```cpp // Auto-detect from DEEP_EP_BACKEND environment variable BackendType detect_backend_type(); // Global singleton management void initialize_backend(BackendType type, ...); CommunicationBackend* get_backend(); void finalize_backend(); ``` --- ## NCCL Backend Integration The NCCL backend uses the Device API for GPU-initiated network operations (GIN), translating between NVSHMEM's PGAS model and NCCL's window-based model: | Aspect | NVSHMEM/IBGDA | NCCL | |--------|---------------|------| | **Memory Model** | Pointer-based addressing | Window-based with offset addressing | | **Data Transfer** | `put_nbi(dst_ptr, src_ptr, count, pe)` | `put(peer, dstWin, dstOff, srcWin, srcOff, bytes)` | | **Synchronization** | Memory atomics | Signal atomics (`signal()`, `readSignal()`) | | **Completion** | Per-QP `quiet()` | Per-context `flush()` | **Key Integration Challenges:** 1. **Multi-Communicator Mapping**: Each NCCL GIN context wraps a Queue Pair (QP), and NCCL provides 4 GIN contexts per communicator. DeepEP's QP requirements are met using ⌈QPs/4⌉ communicators: ``` comm_id = channel_id / 4 ctx_id = channel_id % 4 ``` 2. **Memory Registration**: Buffers registered with all communicators; window handles stored in GPU memory for kernel access. 3. **Signal-Based Synchronization**: Pre-allocated signal layouts map memory atomics to NCCL signal primitives. 4. **Semantic Preservation**: Zero-byte `put()` with `SignalAdd` is semantically equivalent to `net.signal()` but better performing in the current NCCL release. --- ## NCCL GIN Type Selection For GPU-initiated network operations, NCCL supports multiple backends: | `NCCL_GIN_TYPE` | Backend | Description | |:---------------:|:-------:|:------------| | `2` | Proxy | CPU proxy for network operations | | `3` | GDAKI | GPU Direct Async Kernel-Initiated (recommended) | --- ## Performance Results Benchmarked on **H100** (900 GB/s NVLink) with **8×400 Gbit/s InfiniBand** (~50 GB/s per NIC). ### High-Throughput Kernels *4096 tokens, 7168 hidden, top-8 experts, BF16 dispatch, BF16 combine* | Backend | #EP | Dispatch BW | Combine BW | |:-------:|:---:|:-----------:|:----------:| | NVSHMEM | 16 | 79.7 GB/s | 66.4 GB/s | | **NCCL** | 16 | 76.9 GB/s | 66.2 GB/s | | NVSHMEM | 32 | 62.9 GB/s | 62.9 GB/s | | **NCCL** | 32 | 61.7 GB/s | 62.3 GB/s | | NVSHMEM | 64 | 53.5 GB/s | 53.2 GB/s | | **NCCL** | 64 | 52.7 GB/s | 52.9 GB/s | ### Low-Latency Kernels *128 tokens, 7168 hidden, top-8 experts, FP8 dispatch, BF16 combine* | Backend | #EP | Dispatch Latency | Dispatch BW | Combine Latency | Combine BW | |:-------:|:---:|:----------------:|:-----------:|:---------------:|:----------:| | NVSHMEM | 8 | 160.7 μs | 46.8 GB/s | 304.2 μs | 47.8 GB/s | | **NCCL** | 8 | 160.8 μs | 47.0 GB/s | 302.8 μs | 47.9 GB/s | | NVSHMEM | 16 | 182.3 μs | 41.4 GB/s | 319.8 μs | 45.5 GB/s | | **NCCL** | 16 | 178.6 μs | 42.2 GB/s | 320.8 μs | 45.3 GB/s | | NVSHMEM | 32 | 188.7 μs | 40.0 GB/s | 332.9 μs | 43.7 GB/s | | **NCCL** | 32 | 190.1 μs | 39.8 GB/s | 333.2 μs | 43.6 GB/s | | NVSHMEM | 64 | 225.1 μs | 34.8 GB/s | 343.1 μs | 42.4 GB/s | | **NCCL** | 64 | 218.9 μs | 34.5 GB/s | 351.1 μs | 41.4 GB/s | --- ## Requirements - **NCCL 2.28.9** (or above) with Device API and GIN support *(required)* **Tested configurations:** - CUDA 12.8+ - Hopper (H100) and Blackwell GPUs --- ## Usage ```bash # Build with NCCL backend export NCCL_DIR=/path/to/nccl/build export ENABLE_NCCL=1 python3 setup.py build_ext --inplace && pip install --no-build-isolation . # Run with NCCL DEEP_EP_BACKEND=nccl NCCL_GIN_TYPE=3 python3 tests/test_low_latency.py DEEP_EP_BACKEND=nccl NCCL_GIN_TYPE=3 python3 tests/test_internode.py ``` --- ## Future Work 1. **Reduce Communicator Count**: Currently, we initialize multiple NCCL communicators due to the limitation of 4 GIN contexts per communicator. Work is in progress on the NCCL roadmap to increase the number of GIN contexts per communicator, which will simplify the integration. 2. **Unified Backend Interface for NVSHMEM**: The current NVSHMEM code does not follow the `CommunicationBackend` interface in order to stay close to the upstream DeepEP main branch. Migrating NVSHMEM to use this interface should be addressed in consultation with DeepEP maintainers. --- ## References 📄 **Paper**: [GPU-Initiated Networking for NCCL](https://arxiv.org/abs/2511.15076) 🔗 **NCCL Repository**: https://github.com/NVIDIA/nccl 📄 **NCCL README**: [NCCL README](https://github.com/aamirshafi/DeepEP/blob/nccl/README-NCCL.md) --- Co-authored with @grtheod. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
hybrid-ep
chore/remove-clang-format-yapf
v2.5
nvDev
try-fix-cu128
mori-ep
antgroup-opt
compat-old-torch
add-linter
eager-rdma
v1.2.1
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepEP#607
No description provided.