!!!DEPRECATION!!!

We deprecated the old NVSHMEM-based kernels in this repo. Please check out our newer kernels in https://github.com/perplexityai/pplx-garden.

Perplexity MoE Kernels

Features:

• ✅ Cuda Graph support
• ✅ Flexible transportation layers: NVLink, IBGDA, IBRC, EFA
• ✅ Overlapping communication and computation

System Requirements

To learn how to set up the system drivers and dependencies, refer to the Install Driver and Dependencies guide.

Installation

cd pplx-kernels
TORCH_CUDA_ARCH_LIST=9.0a+PTX python3 setup.py bdist_wheel
pip install dist/*.whl

Single-node Testing and Benchmarking

Test:

pytest -svx --tb=short tests

Benchmark:

python3 -m tests.bench_all_to_all

Multi-node Testing and Benchmarking

To run an all-to-all test across multiple nodes:

export NODE_RANK= # 0, 1, ..., num_nodes-1
export WORLD_SIZE= # num_nodes * 8
export WORLD_LOCAL_SIZE=8
export MASTER_ADDR= # IP address of rank-0 node
export MASTER_PORT=29500
export NVSHMEM_IB_ENABLE_IBGDA=1

An then run the tests normally:

pytest -svx tests/test_all_to_all.py

Benchmark Results

1 token per GPU:

128 tokens per GPU:

C++ Testing

To build the C++ tests and benchmarks:

cd pplx-kernels
mkdir build-cmake
cd build-cmake

export TORCH_PREFIX_PATH=$(python3 -c 'import torch; print(torch.utils.cmake_prefix_path)')

cmake ../csrc \
    -GNinja \
    -DCMAKE_PREFIX_PATH=$TORCH_PREFIX_PATH \
    -DTORCH_CUDA_ARCH_LIST=9.0a+PTX \
    -DWITH_TESTS=ON \
    -DWITH_BENCHMARKS=ON

ninja test_all_to_all bench_all_to_all

To run the all-to-all tests on one node:

NVSHMEM_REMOTE_TRANSPORT=none mpirun -np 4 ./all_to_all/test_all_to_all

To run the all-to-all benchmarks on one node:

NVSHMEM_REMOTE_TRANSPORT=none mpirun -np 4 ./all_to_all/bench_all_to_all

Citation

If you use this codebase or otherwise find our work valuable, please cite:

@misc{pplx-kernels,
  title={{pplx-kernels}: {Perplexity} {MoE} Kernels},
  author={Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  year={2025},
  publisher = {GitHub},
  howpublished = {\url{https://github.com/perplexityai/pplx-kernels}},
}