UCCL-EP Benchmark

Micro-benchmark for UCCL expert-parallelism dispatch/combine on NVLink and RDMA over EFA, runnable on Slurm and Amazon EKS

Overview

The UCCL-EP benchmark measures expert-parallelism (EP) dispatch and combine performance for Mixture-of-Experts (MoE) communication using UCCL. It runs both intranode (over NVLink) and internode (RDMA over EFA) tests, plus a low-latency decode-path test, and can be launched on Slurm or on Amazon EKS.

Container image

Build the image and, for Slurm, import it with enroot:

docker build -t uccl-ep -f uccl-ep.Dockerfile .

enroot import -o ./uccl-ep.sqsh dockerd://uccl-ep

Run on Slurm

Submit the provided sbatch scripts. The intranode test exercises NVLink, while the internode test exercises RDMA over EFA, and the low-latency test measures the decode-path:

sbatch test_intranode.sbatch

sbatch test_internode.sbatch

sbatch test_low_latency.sbatch

Representative dispatch/combine bandwidths reported by the benchmark:

Run on Amazon EKS

On EKS the benchmarks run as MPIJobs, the Kubernetes equivalents of the Slurm sbatch scripts. Each job launches one MPI rank per GPU (slotsPerWorker: 8); the launcher maps the OMPI per-rank variables to the standard torch.distributed env contract and passes the EFA/NCCL environment through mpirun -x.

Prerequisites

You need an EKS cluster with EFA-enabled GPU nodes (e.g. p6-b300.48xlarge), the NVIDIA device plugin, and the AWS EFA device plugin installed:

aws eks update-kubeconfig --name <EKS_CLUSTER_NAME>
kubectl config current-context

The MPIJob CRD is provided by the standalone Kubeflow MPI Operator:

kubectl apply -f https://raw.githubusercontent.com/kubeflow/mpi-operator/v0.6.0/deploy/v2beta1/mpi-operator.yaml
kubectl get crd mpijobs.kubeflow.org      # confirm the CRD exists

Build & push the container image

The image is built for both Hopper (sm_90) and Blackwell (sm_100/sm_103) with PTX fallback, so one image runs on p5/p5en and p6-b300 without a rebuild. Push it to ECR:

export AWS_REGION=$(aws ec2 describe-availability-zones --output text --query 'AvailabilityZones[0].[RegionName]')
export ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
export REGISTRY=${ACCOUNT}.dkr.ecr.${AWS_REGION}.amazonaws.com
export IMAGE_URI=${REGISTRY}/uccl-ep:efa1.48.0-uccl0dc87eb-cu13

aws ecr create-repository --repository-name uccl-ep --region ${AWS_REGION} 2>/dev/null || true
aws ecr get-login-password --region ${AWS_REGION} | docker login --username AWS --password-stdin ${REGISTRY}

cd ..   # build context is the uccl-ep-benchmark/ directory
docker build --progress=plain -f ./uccl-ep.Dockerfile -t ${IMAGE_URI} .
docker push ${IMAGE_URI}

Configure & launch

Copy the committed template to env_vars (gitignored), edit it with your image URI, instance type, and per-node device counts, then envsubst the manifest into kubectl. Set NP = NUM_NODES * GPU_PER_NODE (the total rank count):

cp env_vars.example env_vars   # then edit env_vars
source env_vars

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

# Internode (RDMA over EFA) test  — e.g. NUM_NODES=8, NP=64
envsubst '$IMAGE_URI $INSTANCE_TYPE $GPU_PER_NODE $EFA_PER_NODE $NUM_NODES $NP' \
  < test-internode.yaml | kubectl apply -f -

# Low-latency (decode-path) test  — e.g. NUM_NODES=8, NP=64
envsubst '$IMAGE_URI $INSTANCE_TYPE $GPU_PER_NODE $EFA_PER_NODE $NUM_NODES $NP' \
  < test-low-latency.yaml | kubectl apply -f -

The EP shape is baked into each manifest to a DeepSeek-V3-class configuration: --num-tokens 4096 (128 for low-latency), --hidden 7168, --num-topk 8, --num-experts 256. num-experts=256 divides evenly across 64 ranks (8 nodes × 8 GPU) and matches the DeepEP NVSHMEM benchmark, so the dispatch/combine bandwidths are comparable. (The upstream Slurm scripts use --num-experts=288, which is not divisible by 64, so the Kubernetes manifests override it to 256.)

Monitor

kubectl get mpijob
watch kubectl get pods -o wide
kubectl logs -f $(kubectl get pods | grep uccl-ep.*launcher | cut -d ' ' -f 1)

Clean up

MPIJob names are fixed, so delete a run before re-applying it:

envsubst '$IMAGE_URI $INSTANCE_TYPE $GPU_PER_NODE $EFA_PER_NODE $NUM_NODES $NP' \
  < test-internode.yaml | kubectl delete -f -

Notes for p6-b300.48xlarge (B300 / Blackwell)

• EFA count: EFA_PER_NODE=16 (a p6-b300.48xlarge exposes 16 EFA NICs); p5.48xlarge exposes 32 — adjust per instance type.
• Tolerations: the manifests tolerate the nvidia.com/gpu and capacity-reservation taints; add matching tolerations if your GPU nodes carry additional ones.
• GPU profiling: the internode and low_latency tests profile their kernels during the tuning phase, so the worker securityContext grants SYS_ADMIN (nodes load the driver with NVreg_RestrictProfilingToAdminUsers=1). Remove it if your nodes already permit non-admin profiling.