DeepEP V2 Benchmark (NCCL GIN / EFA-GDA)

Install and validate DeepEP V2 MoE expert-parallelism dispatch/combine kernels over NCCL GIN with the EFA-GDA backend on Slurm and EKS.

View on GitHub

Overview

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 benchmark installs and validates DeepEP V2 — the version whose internode transport is NCCL GIN (GPU-Initiated Networking) — and measures dispatch/combine bandwidth over both NVLink (intra-node) and EFA (inter-node). It provides:

  • setup_deepep_gin.sh — installs the DeepEP V2 Python package into your container or environment. It uninstalls any existing deep_ep, validates that the NCCL it is pointed at is GIN-capable (include/nccl_device.h, erroring out otherwise), and builds + installs DeepEP against your torch.
  • deepep.Dockerfile — a reference image showing the full stack (the EFA installer, which bundles the EFA-GDA-capable aws-ofi-nccl plugin, a GIN-capable NCCL, and torch), built from a bare CUDA base.
  • slurm/ and kubernetes/ launchers that run DeepEP’s own tests/elastic/test_ep.py to validate the installation.
BenchmarkNodesMeasures
test-intranode1intra-node dispatch/combine over NVLink
test-internode2inter-node dispatch/combine (GPU-initiated RDMA over EFA)

Upstream repositories

Two GitHub repos are involved, and it is worth keeping them straight:

  • The benchmark itself (this page, the setup script, the Dockerfile, and the Slurm/Kubernetes launchers) lives in awslabs/awsome-distributed-ai under micro-benchmarks/expert-parallelism/deepep-v2-benchmark. That is the source of truth for this page.
  • The DeepEP V2 library the benchmark builds and installs comes from the amazon-contributing/DeepEP fork. The clone URL is fixed to this fork; the setup script and the Dockerfile default the checked-out ref to the fork’s main branch (the --deepep-ref option / DEEPEP_REF build arg both default to main), so the default flow tracks a moving branch rather than a pinned commit. No arguments are needed for the default flow (run the script with --help for options). You can point the build at any ref of that fork with --deepep-ref (a commit, branch, tag, or PR ref) or build from an existing source tree with --deepep-src.

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, 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; if left unset, NCCL selects the CPU-proxy backend when running on EFA.

VariableValuePurpose
NCCL_GIN_TYPE5Pins the EFA-GDA GIN backend; unset falls back to the CPU-proxy backend.

The Slurm and Kubernetes launchers set NCCL_GIN_TYPE=5 (alongside NCCL_SYM_GIN_KERNELS_ENABLE=0) so runs use the EFA-GDA backend rather than silently taking the CPU-proxy path.

Version constraints

The following version floors and requirements are transcribed from the upstream benchmark README:

  • NCCL >= 2.31 — provides the GIN device API (nccl_device.h) DeepEP V2 links against. The Dockerfile installs NVIDIA’s precompiled NCCL from PyPI as a Python package. The exact NCCL version used by the reference image is pinned by the deepep.Dockerfile (NCCL_VERSION build arg).
  • EFA installer >= 1.50 — provides everything EFA-GDA needs in the container (the userspace runtime and the aws-ofi-nccl plugin with the EFA-GDA GIN backend, type 5). Per the upstream README, installers up to 1.49.0 will not work and the image build fails loudly on a too-old installer. DeepEP’s kernels can also run over the CPU-proxy GIN backend, which has no such floor. The exact installer version used by the reference image is pinned by the same Dockerfile (EFA_INSTALLER_VERSION build arg).
  • Host EFA kernel driver >= 3.3.0 — the first release with the completion-counter API (efadv_create_comp_cntr). Run the EFA installer on the node itself (not in the container) to install/upgrade the kernel driver; the provided Dockerfile ships only the userspace stack. Check with modinfo efa | grep ^version — stock AMIs may ship older.
  • 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, defaulting to the fork’s main branch.

Prerequisites

  • EFA-enabled GPU nodes with the EFA kernel driver installed.
  • GPU architecture: by default the Dockerfile’s TORCH_CUDA_ARCH_LIST / NVCC_GENCODE build args target Hopper (sm_90) and Blackwell (sm_100, sm_103), so one image runs on p5/p5en and p6 per the upstream README. (These are the Dockerfile defaults and can drift, so treat the deepep.Dockerfile build args as the source of truth.) 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), so CUDA 13.0 works. 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

Inside a container that already has torch (and possibly an older deep_ep), point the script at a GIN-capable NCCL install:

# 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 (include/nccl_device.h), so it cannot silently produce a DeepEP that falls back to a slower path. Runtime still requires the EFA-GDA-capable aws-ofi-nccl GIN plugin, which is bundled with the EFA installer (>= 1.50) run inside the container.

Building the reference image

Build the image externally (public NCCL release + official EFA installer) with BuildKit:

DOCKER_BUILDKIT=1 docker build --progress=plain -f ./deepep.Dockerfile \
  -t deepep-v2:gin .

Override NCCL_VERSION or EFA_INSTALLER_VERSION for different versions (the default pins live in the Dockerfile):

DOCKER_BUILDKIT=1 docker build --progress=plain -f ./deepep.Dockerfile \
  --build-arg NCCL_VERSION=2.31.2 \
  --build-arg EFA_INSTALLER_VERSION=1.50.0 \
  -t deepep-v2:gin .

Run on Slurm (enroot + pyxis)

Convert the Docker image to a squashfs container, then submit the jobs:

enroot import -o ~/deepep-v2.sqsh dockerd://deepep-v2:gin

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. The upstream README shows an example line (numbers vary per run):

* 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 is scale-up (NVLink). RUN_EXIT=0 on every rank means the correctness checks passed.

Run on EKS (MPIJob)

On Amazon EKS the same tests run as Kubeflow MPIJobs, the Kubernetes equivalents of the Slurm scripts. Each job launches one MPI process per node (slotsPerWorker: 1), and that process spawns 8 local GPU ranks via torch.multiprocessing.

Build and push the image to ECR, then substitute your image URI, instance type, and per-node device counts into the manifests before applying them:

cd kubernetes
cp env_vars.example env_vars   # then edit env_vars
source env_vars

# Single-node intranode (NVLink) test
envsubst '$IMAGE_URI $INSTANCE_TYPE $GPU_PER_NODE $EFA_PER_NODE $NUM_NODES' \
  < test-intranode.yaml | kubectl apply -f -

# Two-node internode (RDMA over EFA) test — set NUM_NODES=2
envsubst '$IMAGE_URI $INSTANCE_TYPE $GPU_PER_NODE $EFA_PER_NODE $NUM_NODES' \
  < test-internode.yaml | kubectl apply -f -

The manifests require the standalone Kubeflow MPI Operator (which provides the MPIJob CRD), the NVIDIA device plugin, and the AWS EFA device plugin. To expose the gdrdrv kernel driver, install the NVIDIA device plugin with gdrcopyEnabled=true. Watch progress with:

kubectl get mpijob
watch kubectl get pods -o wide
kubectl logs -l training.kubeflow.org/job-name=deepep-v2-internode -f