Add DeepEP V2 benchmark (NCCL GIN)

micro-benchmarks/expert-parallelism/deepep-v2-benchmark/README.md

@@ -0,0 +1,125 @@
+# DeepEP V2 Benchmark (NCCL GIN / EFA-GDA)
+
+[DeepEP](https://github.com/deepseek-ai/DeepEP) is a communication library for
+Mixture-of-Experts (MoE) **expert parallelism** — its dispatch and combine kernels perform the
+GPU all-to-all that routes tokens to experts and gathers the results back.
+
+This directory provides **[`setup_deepep_gin.sh`](./setup_deepep_gin.sh)**, which installs the
+**DeepEP V2** Python package — the version whose internode transport is **NCCL GIN
+(GPU-Initiated Networking)** — into your container or environment (for example a vLLM image
+that already ships a DeepEP): it uninstalls any existing `deep_ep`, validates that the NCCL it
+is pointed at is GIN-capable (`include/nccl_device.h`, errors out otherwise), and builds +
+installs DeepEP against your torch. No arguments needed for the default flow; see `--help`.
+
+Everything else here is reference material for using the script and testing that the
+installation works:
+
+- **[`deepep.Dockerfile`](./deepep.Dockerfile)** — a reference image showing the full stack the
+  script needs (EFA userspace, a GIN-capable NCCL, aws-ofi-nccl with the EFA-GDA backend,
+  torch), built from a bare CUDA base; it can also be used directly for testing.
+- **[`slurm/`](./slurm/)** — launchers that run DeepEP's own `tests/elastic/test_ep.py` to
+  validate the installation:
+
+| Benchmark | Nodes | Measures |
+|-----------|-------|----------|
+| `test-intranode` | 1 | intra-node dispatch/combine over NVLink |
+| `test-internode` | 2 | inter-node dispatch/combine (GPU-initiated RDMA over EFA) |
+
+## ⚠️ Version constraints
+
+> - **NCCL >= 2.31** — the GIN device API (`nccl_device.h`) DeepEP V2 links against.
+>   The Dockerfile builds NCCL from source and fails the build if the header is missing.
+> - **EFA installer >= 1.50** — EFA-GDA requires **libfabric >= 2.5**; installers up to
+>   1.49.0 ship libfabric 2.4 and will not work (the image build fails loudly on a
+>   too-old installer).
+> - **aws-ofi-nccl built against that libfabric** — required for the **EFA-GDA** GIN backend
+>   (GIN type 5) this benchmark measures. DeepEP's kernels can also run over the CPU-proxy
+>   GIN backend, which has no such floor.
+> - A host EFA kernel driver **>= 3.3.0**, the first release with the completion-counter API
+>   (`efadv_create_comp_cntr`). Check with `modinfo efa | grep ^version`; stock AMIs may ship
+>   older and need a driver upgrade. The container ships only the userspace stack.
+> - The **gdrcopy kernel module (`gdrdrv`) loaded on compute nodes** — the launchers
+>   bind-mount `/dev/gdrdrv`, and the plugin's GIN initialization opens a gdr handle.
+>   Check with `lsmod | grep gdrdrv`.
+> - **DeepEP V2** from [amazon-contributing/DeepEP](https://github.com/amazon-contributing/DeepEP) `main`.
+
+## How the EFA support works
+
+DeepEP V2's internode kernels communicate through **NCCL GIN**: the dispatch/combine kernels
+call NCCL's device-side communication API, and a GIN backend carries the traffic. Backends come
+in two kinds: GPU-initiated ones, where the NIC work queues are mapped into GPU memory and the
+kernels post RDMA themselves, and the CPU-proxy backend (available on EFA as well as IB/RoCE),
+where the GPU hands work to a proxy thread. On EFA, the **aws-ofi-nccl** plugin provides the
+GPU-initiated backend, **EFA-GDA** (GIN type 5). The launchers pin `NCCL_GIN_TYPE=5` to select
+it; left unset, NCCL selects the CPU-proxy backend.
+
+| Variable | Value | Purpose |
+|----------|-------|---------|
+| `NCCL_GIN_TYPE` | `5` | Pins the EFA-GDA GIN backend; unset falls back to the CPU-proxy backend. |
+| `EP_NCCL_ROOT_DIR` | `$NCCL_HOME` | Points DeepEP's build/runtime checks at the GIN-capable NCCL. |
+| `LD_PRELOAD` | `$NCCL_HOME/lib/libnccl.so.2` | Guarantees the GIN-capable NCCL wins over any other on the loader path. |
+
+## Prerequisites
+
+- EFA-enabled GPU nodes with the EFA kernel driver installed (section below).
+- **GPU architecture:** the image builds for **Hopper (`sm_90`) and Blackwell (`sm_100`)** by
+  default, so one image runs on `p5`/`p5en` and `p6`. Earlier DeepEP revisions required
+  **CUDA >= 13.1** on p6/Blackwell (ptxas 13.0.88 rejects the 32-bit `st.bulk` size operand);
+  DeepEP main carries the fix ([amazon-contributing/DeepEP#3](https://github.com/amazon-contributing/DeepEP/pull/3)),
+  so CUDA 13.0 works as well. Override
+  `--build-arg TORCH_CUDA_ARCH_LIST` / `--build-arg NVCC_GENCODE` for a single-arch image.
+- Docker with BuildKit, and enroot + pyxis on the cluster for the Slurm flow.
+
+## Using the script in your own container
+
+```bash
+# inside a container that already has torch (and possibly an older deep_ep):
+./setup_deepep_gin.sh --nccl-root /path/to/gin-capable/nccl
+```
+
+The script refuses to run if the NCCL at `--nccl-root` lacks the GIN device API, so it cannot
+silently produce a DeepEP that falls back to a slower path. Runtime still requires the
+aws-ofi-nccl GIN plugin built with EFA-GDA support.
+
+## Building the reference image
+
+### External use (public NCCL release + official EFA installer)
+
+```bash
+DOCKER_BUILDKIT=1 docker build --progress=plain -f ./deepep.Dockerfile \
+  -t deepep-v2:gin .
+```
+
+Override `NCCL_REF` or `EFA_INSTALLER_VERSION` for different versions:
+
+```bash
+DOCKER_BUILDKIT=1 docker build --progress=plain -f ./deepep.Dockerfile \
+  --build-arg NCCL_REF=v2.31.2-1 \
+  --build-arg EFA_INSTALLER_VERSION=1.50.0 \
+  -t deepep-v2:gin .
+```
+
+## Run (Slurm + enroot)
+
+### 1. Convert the Docker image to a squashfs container
+
+```bash
+enroot import -o ~/deepep-v2.sqsh dockerd://deepep-v2:gin
+```
+
+### 2. Submit the Slurm jobs
+
+```bash
+SQSH=~/deepep-v2.sqsh sbatch slurm/test-intranode.sbatch   # 1 node, NVLink
+SQSH=~/deepep-v2.sqsh sbatch slurm/test-internode.sbatch   # 2 nodes, EFA-GDA
+```
+
+Each job runs `tests/elastic/test_ep.py`, which checks dispatch/combine correctness and prints
+per-rank bandwidth lines like:
+
+```
+* EP:   3/16 | dispatch: 81 GB/s (SO), 263 GB/s (SU), ... | copy: 3068 GB/s, ...
+```
+
+`SO` is scale-out (inter-node over EFA), `SU` scale-up (NVLink). `RUN_EXIT=0` on every rank
+means the correctness checks passed.

micro-benchmarks/expert-parallelism/deepep-v2-benchmark/deepep.Dockerfile

@@ -0,0 +1,195 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+#
+# DeepEP V2 image for AWS EFA clusters.
+# See README.md for build and run instructions.
+
+# CUDA 13.1 is the default. Earlier DeepEP revisions needed >= 13.1 on
+# p6/Blackwell (ptxas 13.0.88 rejects the 32-bit st.bulk size operand); DeepEP
+# main carries the 64-bit fix (amazon-contributing/DeepEP#3), so 13.0 works too.
+ARG CUDA_VERSION=13.1.2
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
+
+ARG TORCH_VERSION=2.11.0
+# Torch wheel index. Kept explicit rather than derived from nvcc: pytorch.org has
+# no cu131 index, and cu130 wheels run on a 13.1 toolkit (same CUDA major).
+ARG TORCH_CUDA_INDEX=cu130
+ARG GDRCOPY_VERSION=v2.5.2
+
+# EFA installer: provides OpenMPI, the EFA userspace runtime, and libfabric.
+# EFA-GDA requires libfabric >= 2.5, so installer >= 1.50 is required (1.49.0 and
+# older ship libfabric 2.4; the build fails loudly on a too-old installer).
+# EFA_INSTALLER_TARBALL overrides the download with a local tarball staged next
+# to this Dockerfile.
+ARG EFA_INSTALLER_VERSION=1.50.0
+ARG EFA_INSTALLER_TARBALL=
+
+# NCCL: defaults to the public NVIDIA release. For internal use, override with
+# the upstream fork + staging branch (see build instructions in README.md).
+ARG NCCL_REPO=https://github.com/NVIDIA/nccl.git
+ARG NCCL_REF=v2.31.2-1
+
+# aws-ofi-nccl: built from source at a pinned tag so the plugin version is
+# controlled by this file rather than by whatever the EFA installer bundles.
+# EFA-GDA is auto-detected by configure from
+# the libfabric >= 2.5 GDA ops + hardware counters + CUDA; the build gate below
+# fails the image if the detection did not fire.
+ARG AWS_OFI_NCCL_REPO=https://github.com/aws/aws-ofi-nccl.git
+ARG AWS_OFI_NCCL_REF=v1.21.1
+
+# DeepEP V2. The repo is pinned to amazon-contributing/DeepEP inside
+# setup_deepep_gin.sh; only the ref is overridable.
+ARG DEEPEP_REF=main
+
+# CUDA architecture(s), semicolon-separated:
+# 9.0 = Hopper (H100/H200, sm_90); 10.0/10.3 = Blackwell (B200/B300, sm_100/sm_103).
+# Defaults to Hopper + Blackwell so one image runs on p5/p5en and p6.
+ARG TORCH_CUDA_ARCH_LIST="9.0;10.0"
+ARG NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90 -gencode=arch=compute_100,code=sm_100"
+ARG CUDA_HOME="/usr/local/cuda"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+## Remove conflicting distro packages
+RUN apt-get update -y \
+    && apt-get remove -y --allow-change-held-packages \
+    ibverbs-utils \
+    libibverbs-dev \
+    libibverbs1 \
+    libmlx5-1 \
+    libnccl2 \
+    libnccl-dev || true
+RUN rm -rf /opt/hpcx /usr/local/mpi /etc/ld.so.conf.d/hpcx.conf && ldconfig
+
+## Toolchain
+RUN apt-get update -y && apt-get install -y \
+    apt-utils autoconf automake build-essential check cmake curl debhelper \
+    devscripts git gcc gdb libsubunit-dev libtool ninja-build meson pandoc \
+    pkg-config vim wget \
+    python3.10-dev python3.10-venv python3-distutils \
+    cython3 \
+    libnl-3-dev libnl-route-3-dev libudev-dev libsystemd-dev \
+    libhwloc-dev \
+    && rm -rf /var/lib/apt/lists/*
+RUN apt-get purge -y cuda-compat-* || true
+
+RUN curl https://bootstrap.pypa.io/get-pip.py -o /tmp/get-pip.py \
+    && python3 /tmp/get-pip.py \
+    && pip3 install awscli nvidia-ml-py ninja Cython
+
+## GDRCopy
+# Required by aws-ofi-nccl: the plugin's GIN initialization opens a gdr handle
+# regardless of backend, so building or running without GDRCopy leaves the GIN
+# backends unable to initialize (verified: DeepEP then fails with "NCCL GIN is
+# unavailable"). Built from source because no distro package ships it. The
+# matching gdrdrv kernel module must be loaded on the host (see README).
+ARG GDRCOPY_PREFIX="/opt/gdrcopy"
+RUN git clone -b ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \
+    && cd /tmp/gdrcopy \
+    && make prefix="${GDRCOPY_PREFIX}" install \
+    && rm -rf /tmp/gdrcopy
+ENV LD_LIBRARY_PATH="${GDRCOPY_PREFIX}/lib:$LD_LIBRARY_PATH"
+ENV LIBRARY_PATH="${GDRCOPY_PREFIX}/lib:$LIBRARY_PATH"
+ENV PATH="${GDRCOPY_PREFIX}/bin:$PATH"
+
+## EFA installer (OpenMPI + EFA runtime; local tarball wins when set)
+RUN --mount=type=bind,target=/ctx \
+    apt-get update -y \
+    && cd /tmp \
+    && if [ -n "${EFA_INSTALLER_TARBALL}" ]; then \
+         echo "=== EFA installer: LOCAL tarball /ctx/${EFA_INSTALLER_TARBALL} ==="; \
+         if [ ! -f "/ctx/${EFA_INSTALLER_TARBALL}" ]; then \
+           echo "ERROR: ${EFA_INSTALLER_TARBALL} not found in build context" >&2; exit 1; \
+         fi; \
+         tar -xf "/ctx/${EFA_INSTALLER_TARBALL}"; \
+       else \
+         echo "=== EFA installer: downloading v${EFA_INSTALLER_VERSION} ==="; \
+         curl -fsSL -O "https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz" \
+         && tar -xf "aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz"; \
+       fi \
+    && cd aws-efa-installer \
+    && ./efa_installer.sh --disable-ngc -y --skip-kmod --skip-limit-conf --no-verify \
+    && ldconfig \
+    && FI_VER=$(/opt/amazon/efa/bin/fi_info --version | grep -oP "libfabric: \K[0-9]+\.[0-9]+") \
+    && { [ "$(printf "%s\n2.5\n" "$FI_VER" | sort -V | head -1)" = "2.5" ] \
+         || { echo "ERROR: installer libfabric $FI_VER < 2.5; EFA-GDA requires >= 2.5 (use installer >= 1.50)" >&2; exit 1; }; } \
+    && rm -rf /tmp/aws-efa-installer* /var/lib/apt/lists/*
+ENV LD_LIBRARY_PATH=/opt/amazon/openmpi/lib:$LD_LIBRARY_PATH
+ENV PATH=/opt/amazon/openmpi/bin:/opt/amazon/efa/bin:$PATH
+ENV EFA_PREFIX=/opt/amazon/efa
+
+## NCCL
+ENV NCCL_HOME=/opt/nccl/build
+RUN git clone ${NCCL_REPO} /opt/nccl \
+    && cd /opt/nccl \
+    && git checkout ${NCCL_REF} \
+    && make -j"$(nproc)" src.build CUDA_HOME="${CUDA_HOME}" NVCC_GENCODE="${NVCC_GENCODE}" \
+    && test -f ${NCCL_HOME}/include/nccl_device.h \
+       || (echo "ERROR: nccl_device.h missing -- this NCCL is not GIN-capable" >&2 && exit 1)
+ENV LD_LIBRARY_PATH="${NCCL_HOME}/lib:${LD_LIBRARY_PATH}"
+
+## aws-ofi-nccl. EFA-GDA support is auto-detected by configure from the libfabric
+## above; the ncclGinPlugin_v14 export check makes the detection authoritative
+## (v11/v13 op-tables are exported by every build; the v14 table is EFA-GDA-only).
+ENV OFI_HOME=/opt/aws-ofi-nccl
+RUN set -e; \
+    git clone --recursive ${AWS_OFI_NCCL_REPO} /tmp/aws-ofi-nccl \
+    && cd /tmp/aws-ofi-nccl \
+    && git checkout ${AWS_OFI_NCCL_REF} \
+    && git submodule sync && git submodule update --init --recursive \
+    && ./autogen.sh \
+    && { ./configure \
+        --prefix=${OFI_HOME} \
+        --with-mpi=/opt/amazon/openmpi \
+        --with-libfabric=${EFA_PREFIX} \
+        --with-nccl=${NCCL_HOME} \
+        --with-cuda=${CUDA_HOME} \
+        --with-gdrcopy=${GDRCOPY_PREFIX} \
+        --enable-platform-aws \
+      || { tail -n 160 config.log 2>/dev/null; exit 1; }; } \
+    && make -j"$(nproc)" \
+    && make install \
+    && if nm -D "${OFI_HOME}/lib/libnccl-net-ofi.so" 2>/dev/null | grep -qw ncclGinPlugin_v14; then \
+         echo "OK: EFA-GDA (ncclGinPlugin_v14) present"; \
+       else \
+         echo "ERROR: no ncclGinPlugin_v14 export -- configure did not detect EFA-GDA (libfabric >= 2.5 with hw counters + CUDA required)" >&2; exit 1; \
+       fi \
+    && echo "${OFI_HOME}/lib" > /etc/ld.so.conf.d/aws-ofi-nccl.conf && ldconfig \
+    && rm -rf /tmp/aws-ofi-nccl
+ENV LD_LIBRARY_PATH="${OFI_HOME}/lib:${LD_LIBRARY_PATH}"
+ENV NCCL_NET_PLUGIN="${OFI_HOME}/lib/libnccl-net-ofi.so"
+ENV NCCL_GIN_PLUGIN="${OFI_HOME}/lib/libnccl-net-ofi.so"
+
+# Keep NCCL's bootstrap off the docker/loopback/veth interfaces.
+ENV NCCL_SOCKET_IFNAME=^docker,lo,veth
+
+## PyTorch.
+# The torch wheel drags in pip's stock NCCL (nvidia-nccl-cu*), which predates GIN.
+# We remove it rather than keep both: DeepEP refuses to import when it detects two
+# NCCL runtimes, and with both on the loader path which one wins is load-order
+# luck. The GIN capability itself is enforced above at NCCL build time
+# (nccl_device.h check errors the build), so this cannot silently downgrade: the
+# image ends up with exactly one NCCL, the GIN-capable one it was built and
+# tested with.
+RUN pip3 install torch==${TORCH_VERSION} numpy --index-url https://download.pytorch.org/whl/${TORCH_CUDA_INDEX} \
+    && pip3 uninstall -y nvidia-nccl-cu13 nvidia-nccl-cu12 nvidia-nccl 2>/dev/null || true
+
+## No NVSHMEM: DeepEP V2's transport is NCCL GIN, so the NVSHMEM backend is not
+## built (the DeepEP fork's build only links NVSHMEM when NVSHMEM_DIR points at an
+## install). To additionally build the legacy NVSHMEM backend, install the
+## libnvshmem3-*-cuda-13 packages and set NVSHMEM_DIR before running the setup script.
+
+## DeepEP V2.
+ARG DEEPEP_PREFIX="/opt/amazon/deepep"
+RUN --mount=type=bind,source=setup_deepep_gin.sh,target=/tmp/setup_deepep_gin.sh \
+    set -e; \
+    TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \
+    /tmp/setup_deepep_gin.sh \
+        --deepep-ref "${DEEPEP_REF}" \
+        --deepep-prefix "${DEEPEP_PREFIX}" \
+        --nccl-root "${NCCL_HOME}"
+
+## Runtime env
+ENV NVIDIA_GDRCOPY=enabled
+ENV NCCL_OFI_RDMA_GDR_FLUSH_DISABLE=0
+WORKDIR /root