Perplexity Kernels Benchmark

Benchmark expert-parallelism communication kernels (pplx-kernels) that accelerate Mixture-of-Experts communication using NVSHMEM

View on GitHub

Overview

This benchmark exercises Perplexity’s expert-parallelism kernels (pplx-kernels), which accelerate Mixture-of-Experts (MoE) communication using NVSHMEM. Per the upstream README, it is updated to NVSHMEM 3.4.5-0 and the Aug 6, 2025 pplx-kernels commit.

The container image built here also bundles NCCL and the AWS EFA installer as dependencies, and the NVSHMEM build enables NCCL and libfabric (EFA) support.

Clone NVSHMEM

For the NVSHMEM release pinned upstream:

git clone https://github.com/NVIDIA/nvshmem.git && cd ./nvshmem && git checkout df2814155acfba6227534dd81a8bf338da9e55f2 && cd ..

Alternatively, to use the devel branch:

git clone https://github.com/NVIDIA/nvshmem.git && cd ./nvshmem && git checkout devel && cd ..

Build the container image

Set the build variables that pin the dependency versions. The exact values below are the ones pinned in the upstream benchmark; treat the upstream pplx-kernels.Dockerfile and README.md as the source of truth if they drift:

GDRCOPY_VERSION=v2.5.1
EFA_INSTALLER_VERSION=1.43.2
NCCL_VERSION=v2.27.7-1
NCCL_TESTS_VERSION=v2.16.9
NVSHMEM_VERSION=3.4.5-0
TAG="efa${EFA_INSTALLER_VERSION}-nccl${NCCL_VERSION}-tests${NCCL_TESTS_VERSION}-nvshmem${NVSHMEM_VERSION}"
PPLX_CONTAINER_IMAGE_NAME_TAG="pplx-kernels:${TAG}"

Build the image with pplx-kernels.Dockerfile:

docker build --progress=plain -f ./pplx-kernels.Dockerfile \
       --build-arg="EFA_INSTALLER_VERSION=${EFA_INSTALLER_VERSION}" \
       --build-arg="NCCL_VERSION=${NCCL_VERSION}" \
       --build-arg="NCCL_TESTS_VERSION=${NCCL_TESTS_VERSION}" \
       --build-arg="NVSHMEM_VERSION=${NVSHMEM_VERSION}" \
       -t ${PPLX_CONTAINER_IMAGE_NAME_TAG} \
       .

Convert the Docker image to a squash file for use with enroot/Pyxis:

enroot import -o ./pplx-kernels.sqsh dockerd://${PPLX_CONTAINER_IMAGE_NAME_TAG}

Run the benchmark on Slurm

The benchmark is launched through the provided pplx-kernels.sbatch script (in the upstream benchmark directory), which configures the EFA and NCCL environment variables and runs the all-to-all benchmark via srun using the pplx-kernels.sqsh container. The upstream script requests 2 nodes with 8 tasks per node; consult the script for its exact node and task counts:

sbatch pplx-kernels.sbatch

Check the logs

tail -f -n +0 slurm-XXX.out

Core dump

To capture core dumps for debugging:

  1. Run ulimit -c unlimited and verify that srun -N <num of nodes> bash -c "ulimit -c" prints unlimited once per node.

  2. Configure the core pattern on each node:

    srun -N <num of nodes> sudo bash -c "mkdir -p /tmp/coredump && echo '/tmp/coredump/core.%e.%p' > /proc/sys/kernel/core_pattern"

  3. Submit the job again with sbatch pplx-kernels.sbatch.

Caution: step 2 writes to /proc/sys/kernel/core_pattern with sudo on every node, modifying a host-wide kernel setting cluster-wide. Only run it on nodes you control (e.g. your own HyperPod/Slurm cluster), not on shared infrastructure.

Source

The NVSHMEM clone commit, the Docker build args, the sbatch pplx-kernels.sbatch / log / core-dump steps, and the version pins are transcribed from the upstream micro-benchmarks/expert-parallelism/pplx-kernels-benchmark/README.md. The 2-node, 8-tasks-per-node topology comes from pplx-kernels.sbatch. pplx-kernels.Dockerfile and pplx-kernels.sbatch live in that upstream benchmark directory.