DeepEP Benchmark

Build and run DeepEP's Mixture-of-Experts dispatch/combine micro-benchmarks with an EFA (libfabric) NVSHMEM transport on AWS GPU clusters

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 use case builds DeepEP with an EFA (libfabric) NVSHMEM transport and runs its three micro-benchmarks on AWS GPU clusters, on either Slurm or Kubernetes / EKS.

Benchmark	Nodes	Measures
`test_intranode`	1	intra-node dispatch/combine over NVLink
`test_internode`	2	inter-node dispatch/combine (RDMA over EFA)
`test_low_latency`	2	low-latency kernels used on the decode path

Version constraints

This particular version works only with NVSHMEM >= 3.7.0-0 and DeepEP v1 with the NVSHMEM backend (commit 567632d). This targets DeepEP’s legacy NVSHMEM code path (pre-EPv2). EPv2 restructures the kernel sources and switches to the NCCL GIN backend; it is out of scope here.

How the EFA support works

DeepEP’s internode kernels were written for IBGDA (GPU-initiated RDMA over InfiniBand), which EFA does not provide. The image works around this:

NVSHMEM 3.7.0.
setup_deepep_efa.sh — a self-contained script that downloads NVSHMEM, clones DeepEP, and applies an embedded EFA patch. Per the upstream README, this patch replaces IBGDA device calls with NVSHMEM host-proxy QP APIs (nvshmemx_qp_*), maps the dispatch/combine tail updates onto a combined put+signal, and raises NVSHMEM_MAX_TEAMS for NVSHMEM 3.7. See the upstream setup_deepep_efa.sh as the source of truth for the patch behavior.
The build requires PyTorch; PyTorch’s bundled nvidia-nvshmem-cu13 is uninstalled so it does not clash with the OS-installed NVSHMEM 3.7.0 during runtime. (See the pip3 uninstall step in the upstream deepep.Dockerfile.)
The image is built on CUDA 13 (PyTorch cu130). Per the upstream README, this is required for Blackwell (B200/B300): CUDA ≤ 12.9 CUPTI fails on those GPUs (CUPTI_ERROR_INVALID_DEVICE), which breaks the internode/low-latency tuning phase that profiles kernels.

At runtime the NVSHMEM libfabric transport is selected by these environment variables. They are baked into the image as ENV in the upstream deepep.Dockerfile, so every rank in the container inherits them automatically. The launchers re-export the same values explicitly so the runtime transport config stays visible at the launch surface and can be overridden per-job without rebuilding the image.

Variable	Value	Purpose
`FI_PROVIDER`	`efa`	Selects the EFA provider in libfabric.
`NVSHMEM_REMOTE_TRANSPORT`	`libfabric`	Routes NVSHMEM’s inter-node RDMA through the libfabric transport (instead of IBGDA/UCX).
`NVSHMEM_LIBFABRIC_PROVIDER`	`efa`	Tells the NVSHMEM libfabric transport to use the EFA provider.
`NVSHMEM_NETDEVS_POLICY`	`EXTERNAL_SHARING_PCIE_SWITCH_NIC_EXCLUSIVE`	Sets NIC↔PE (processing element) assignment on multi-NIC EFA nodes; per the upstream README, pairs each PE with the NIC under its PCIe switch.

Prerequisites

EFA-enabled GPU nodes.
GPU architecture: the image is built for both Hopper and Blackwell by default, so the same image runs on p5/p5en and p6-b300 nodes. Per the upstream deepep.Dockerfile comments, the default architecture list covers 9.0 for Hopper (H100, sm_90) and 10.0 plus 10.3 for Blackwell (B200/B300, sm_100/sm_103) — that is, Blackwell contributes two compute capabilities, which is why the default TORCH_CUDA_ARCH_LIST has three entries (9.0;10.0;10.3) and the matching NVCC_GENCODE emits three gencode targets. Those build args and their default values are pinned in the upstream deepep.Dockerfile. To build a smaller single-arch image, override those build args (e.g. TORCH_CUDA_ARCH_LIST=9.0 with the matching Hopper-only NVCC_GENCODE, or TORCH_CUDA_ARCH_LIST=10.0 for Blackwell only).

Building the DeepEP image

The deepep.Dockerfile relies on BuildKit features (RUN --mount), so the build must run under BuildKit. Set DOCKER_BUILDKIT=1 explicitly (or use docker buildx build ...). The values in the block below mirror the ARG defaults declared in the upstream deepep.Dockerfile; treat that Dockerfile as the source of truth for the current default values, since the pins there may drift over time.

GDRCOPY_VERSION=v2.5.2
EFA_INSTALLER_VERSION=1.48.0
TORCH_CUDA_ARCH_LIST="9.0;10.0;10.3" # Hopper + Blackwell by default
NVCC_GENCODE="-gencode=arch=compute_90,code=sm_90 -gencode=arch=compute_100,code=sm_100 -gencode=arch=compute_103,code=sm_103"
TAG="efa${EFA_INSTALLER_VERSION}"
DEEPEP_CONTAINER_IMAGE_NAME_TAG="deepep:${TAG}"

DOCKER_BUILDKIT=1 docker build --progress=plain -f ./deepep.Dockerfile \
      --build-arg="GDRCOPY_VERSION=${GDRCOPY_VERSION}" \
      --build-arg="EFA_INSTALLER_VERSION=${EFA_INSTALLER_VERSION}" \
      --build-arg "TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}" \
      --build-arg "NVCC_GENCODE=${NVCC_GENCODE}" \
      -t ${DEEPEP_CONTAINER_IMAGE_NAME_TAG} \
      .

Running the benchmark

The upstream repository provides ready-to-use launchers for both schedulers:

Slurm (Pyxis/Enroot): the slurm/ directory contains test_intranode.sbatch, test_internode.sbatch, and test_low_latency.sbatch, which run the corresponding python3 /DeepEP/tests/test_*.py under srun --mpi=pmix against the built deepep.sqsh image. See the upstream slurm/ directory.
Kubernetes / EKS (MPIJob): the kubernetes/ directory contains test-intranode.yaml, test-internode.yaml, and test-low-latency.yaml (kubeflow.org/v2beta1 MPIJob), which launch the same tests via mpirun -x. See the upstream kubernetes/ directory.

Source

The version pins, build args, the EFA-patch description, the CUPTI/Blackwell note, and the 567632d DeepEP commit reference are drawn from the upstream micro-benchmarks/expert-parallelism/deepep-benchmark/README.md and its accompanying deepep.Dockerfile and setup_deepep_efa.sh.