Module nccl_backend 

NCCL-based communication backend for expert parallelism.

Uses TCP bootstrap: rank 0 generates a unique ID and sends it to all other ranks via a TCP listener. Then all ranks call ncclCommInitRank with the shared ID.

Optimizations for 2-rank EP with small messages (4 KB):

• Pre-registers buffers with NCCL (ncclCommRegister) to cache IB memory registration
• Uses paired ncclSend/ncclRecv + local BF16 add instead of ncclAllReduce

Health monitoring and recovery:

• Checks ncclCommGetAsyncError after each collective
• Detects broadcast timeouts (>30s) via stream sync + wall-clock check
• Aborts dead communicators via ncclCommAbort and reconnects

Safety contract for the unsafe { ... } calls below

All unsafe blocks in this file wrap a single FFI call into either NCCL (nccl*) or the CUDA Driver API (cu*). The invariants are uniform:

• NCCL handles: NcclComm instances are constructed via nccl::comm_init_rank after a successful TCP bootstrap and are Drop-cleaned via ncclCommDestroy. They are never aliased across threads without a Mutex guarding the comm.
• CUDA buffers passed to NCCL come from a prior cuMemAlloc_v2 on the same device that owns the comm; size in bytes matches the allocation.
• Streams referenced via u64 are owned by the caller and outlive the in-flight collective.
• extern "C" ABI: matches the NCCL 2.20+ headers and the cuMemAlloc_v2/cuLaunchKernel/etc. shapes declared just below.

Per-site // SAFETY: comments are omitted because the contract is identical for every call. Deviations get a per-site comment.

Structs

NcclBackend

NCCL communication backend for multi-GPU / multi-node EP.

Constants

ALL_REDUCE_DTYPE_BYTES

Element width, in bytes, of the dtype the 2-rank send/recv all-reduce moves.

Functions

required_recv_bytes

Bytes required for the 2-rank all-reduce receive buffer.