[WIP] Add deep_ep_wrapper to support SGLang or other frameworks

build.sh

@@ -105,7 +105,8 @@ build_efa() {
 
   # EFA requires a custom NCCL.
   cd thirdparty/nccl-sg
-  make src.build -j$(nproc) NVCC_GENCODE="-gencode=arch=compute_80,code=sm_80"
+  # make src.build -j$(nproc) NVCC_GENCODE="-gencode=arch=compute_80,code=sm_80"
+  make src.build -j$(nproc) NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90"
   cd ../..
 
   echo "[container] Copying EFA .so to uccl/lib/"

collective/efa/run_p5en.sh

@@ -10,12 +10,14 @@ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
 # Disable NVLink.
 NV_LINK_DISABLE=0
 MULTI_GROUP=0
-NIC=10.1.0.0/16
+# NIC=10.1.0.0/16
+NIC=172.31.0.0/16
 # Processes/Ranks/GPUs per node.
 PROCS_PER_NODE=8
 
 TEST=${1:-srd}
-NUM_PROCS=${2:-32}
+# NUM_PROCS=${2:-32}
+NUM_PROCS=${2:-16}
 PROG_NAME=${3:-0}
 
 # all_gather_perf  all_reduce_perf  alltoall_perf  broadcast_perf  gather_perf
@@ -32,7 +34,8 @@ else
     exit 1
 fi
 
-CHANNELS=32
+# CHANNELS=32
+CHANNELS=16
 CHANNELS_NET_PEER=1
 
 # UCCL optimal parameters. Yang: for allreduce with nvlink, we need to use larger buffer to catch up with NCCL with larger buffers, and avoid outliers.
@@ -55,7 +58,7 @@ if [ "$TEST" = "srd" ]; then
         >"nccl_test_outputs/output_rank_$rank.log"
     done
 
-    LIBNCCL_PATH="${UCCL_HOME}/thirdparty/nccl/build/lib/libnccl.so"
+    LIBNCCL_PATH="${UCCL_HOME}/thirdparty/nccl-sg/build/lib/libnccl.so"
     PLUGIN_PATH="/opt/amazon/ofi-nccl/lib/x86_64-linux-gnu/libnccl-net.so"
 
     mpirun --bind-to none -np ${NUM_PROCS} -N ${PROCS_PER_NODE} --hostfile $NODEFILE --map-by ppr:8:node \
@@ -75,7 +78,9 @@ if [ "$TEST" = "srd" ]; then
         -x NCCL_NCHANNELS_PER_NET_PEER=${CHANNELS_NET_PEER} \
         -x NCCL_P2P_NET_CHUNKSIZE=${CHUNK_SIZE} \
         -x NCCL_BUFFSIZE=${BUFFSIZE} \
-        ${UCCL_HOME}/thirdparty/nccl-tests/build/${PROG_NAME} \
+        -x UCCL_EFA_DEVICES=rdmap110s0,rdmap112s0,rdmap135s0,rdmap137s0,rdmap160s0,rdmap162s0,rdmap85s0,rdmap87s0,rdmap111s0,rdmap113s0,rdmap136s0,rdmap138s0,rdmap161s0,rdmap163s0,rdmap86s0,rdmap88s0 \
+        -x UCCL_ENA_DEVICES=enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0 \
+        /usr/local/cuda-12.9/efa/test-cuda-12.9/${PROG_NAME} \
         -b 1K -e 1G -f 2 -c 1 -w 5 -n 10 -t 1 -g 1 \
         2>&1 | while read -r line; do
         if [[ "$line" =~ ^\[[0-9]+,([0-9]+)\](.+) ]]; then
@@ -129,7 +134,9 @@ elif [ "$TEST" = "ud" ]; then
         -x NCCL_TOPO_FILE=${UCCL_HOME}/collective/efa/p4d-24xl-topo.xml \
         -x NCCL_PXN_DISABLE=1 \
         -x UCCL_ENGINE_QUIET=1 \
-        ${UCCL_HOME}/thirdparty/nccl-tests/build/${PROG_NAME} \
+        -x UCCL_EFA_DEVICES=rdmap110s0,rdmap112s0,rdmap135s0,rdmap137s0,rdmap160s0,rdmap162s0,rdmap85s0,rdmap87s0,rdmap111s0,rdmap113s0,rdmap136s0,rdmap138s0,rdmap161s0,rdmap163s0,rdmap86s0,rdmap88s0 \
+        -x UCCL_ENA_DEVICES=enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0,enp71s0 \
+        /usr/local/cuda-12.9/efa/test-cuda-12.9/${PROG_NAME} \
         -b 1K -e 1G -f 2 -c 1 -w 5 -n 10 -t 1 -g 1 \
         2>&1 | while read -r line; do
         if [[ "$line" =~ ^\[[0-9]+,([0-9]+)\](.+) ]]; then

docker/Dockerfile.cuda

@@ -1,6 +1,9 @@
-ARG BASE_IMAGE=nvidia/cuda:12.3.2-devel-ubuntu22.04
+# ARG BASE_IMAGE=nvidia/cuda:12.3.2-devel-ubuntu22.04
+# FROM ${BASE_IMAGE}
+# ARG PY_VER=3.13
+ARG BASE_IMAGE=nvidia/cuda:12.9.0-devel-ubuntu22.04
 FROM ${BASE_IMAGE}
-ARG PY_VER=3.13
+ARG PY_VER=3.12
 
 # Non-interactive APT
 ENV DEBIAN_FRONTEND=noninteractive

docker/Dockerfile.efa

@@ -1,6 +1,9 @@
-ARG BASE_IMAGE=nvidia/cuda:12.3.2-devel-ubuntu22.04
+# ARG BASE_IMAGE=nvidia/cuda:12.3.2-devel-ubuntu22.04
+# FROM ${BASE_IMAGE}
+# ARG PY_VER=3.13
+ARG BASE_IMAGE=nvidia/cuda:12.9.0-devel-ubuntu22.04
 FROM ${BASE_IMAGE}
-ARG PY_VER=3.13
+ARG PY_VER=3.12
 
 # Non-interactive APT
 ENV DEBIAN_FRONTEND=noninteractive
@@ -44,7 +47,8 @@ RUN ln -s /usr/lib/x86_64-linux-gnu/libevent_core-2.1.so.7 /usr/lib/x86_64-linux
     ln -s /usr/lib/x86_64-linux-gnu/libhwloc.so.15 /usr/lib/x86_64-linux-gnu/libhwloc15.so
 
 # Install EFA installer (without kernel driver)
-ARG EFA_VER=1.42.0
+# ARG EFA_VER=1.42.0
+ARG EFA_VER=1.43.2
 RUN curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_VER}.tar.gz && \
     tar -xf aws-efa-installer-${EFA_VER}.tar.gz && \
     cd aws-efa-installer && \

ep/deep_ep_wrapper/README.md

@@ -0,0 +1,8 @@
+## DeepEP Wrapper of UCCL-EP
+
+```
+cp ../bench/buffer.py ./  # Change `utils` to `deep_ep.utils`
+cp ../bench/utils.py ./
+
+python setup.py install
+```

ep/deep_ep_wrapper/deep_ep/__init__.py

@@ -0,0 +1,15 @@
+from uccl.ep import Config, EventHandle
+
+from .utils import EventOverlap, check_nvlink_connections, initialize_uccl, destroy_uccl
+from .buffer import Buffer
+import torch.distributed as dist
+
+__all__ = [
+    'Config',
+    'EventHandle',
+    'Buffer',
+    'EventOverlap',
+    'check_nvlink_connections',
+    'initialize_uccl',
+    'destroy_uccl',
+]