Add NCCL Communication Backend for DeepEP

[aamirshafi]

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.

[alpha-baby]

good job!

It seems that the version of NCCL v2.29.1 has not been released, but how did you get the code of this version?

[polarstormx]

good job!

It seems that the version of NCCL v2.29.1 has not been released, but how did you get the code of this version?

Because he works at NVIDIA lol

[aamirshafi]

@alpha-baby - Thanks for letting us know. Updated the NCCL version to 2.28.9.

@polarstormx - You got that right 💚