!!!已弃用!!!

我们已在此仓库中弃用了基于 NVSHMEM 的旧内核。请查看我们在 https://github.com/perplexityai/pplx-garden 中的较新内核。

Perplexity MoE 内核

特性：

• ✅ 支持 CUDA 图
• ✅ 灵活的传输层：NVLink、IBGDA、IBRC、EFA
• ✅ 通信与计算重叠

系统要求

有关如何设置系统驱动程序和依赖项，请参阅 安装驱动程序和依赖项 指南。

安装

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

单节点测试和基准测试

测试：

pytest -svx --tb=short tests

基准：

python3 -m tests.bench_all_to_all

多节点测试和基准测试

要在多个节点之间运行全对全测试：

export NODE_RANK= # 0, 1, ..., num_nodes-1
export WORLD_SIZE= # num_nodes * 8
export WORLD_LOCAL_SIZE=8
export MASTER_ADDR= # rank-0 节点的 IP 地址
export MASTER_PORT=29500
export NVSHMEM_IB_ENABLE_IBGDA=1

然后正常运行测试：

pytest -svx tests/test_all_to_all.py

基准测试结果

每 GPU 1 个 token：

每 GPU 1 个 token	EP128	EP64	EP32	EP16	EP8
NVLINK 分发	x	x	x	x	41.6μs ± 1.3μs
IBGDA 分发	125.9μs ± 0.6μs	121.0μs ± 0.2μs	115.7μs ± 1.4μs	102.7μs ± 8.7μs	x
IBRC 分发	488.4μs ± 51.0μs	525.0μs ± 9.4μs	421.2μs ± 35.5μs	290.5μs ± 4.7μs	x
NVLINK 聚合	x	x	x	x	41.7μs ± 3.0μs
IBGDA 聚合	63.2μs ± 8.3μs	58.6μs ± 1.0μs	55.4μs ± 0.8μs	62.7μs ± 0.7μs	x
IBRC 聚合	786.8μs ± 149.8μs	400.0μs ± 47.9μs	122.1μs ± 38.2μs	85.9μs ± 5.3μs	x
Torch 全对全	132.0μs ± 25.9μs	101.6μs ± 15.7μs	95.7μs ± 14.3μs	109.7μs ± 3.1μs	24.4μs ± 16.3μs
NVLINK NVSHMEM 全对全	x	x	x	x	59.9μs ± 30.7μs
IBGDA NVSHMEM 全对全	132.4μs ± 73.3μs	95.3μs ± 23.5μs	77.3μs ± 23.0μs	71.7μs ± 14.6μs	x
IBRC NVSHMEM 全对全	258.8μs ± 145.3μs	98.9μs ± 57.1μs	63.2μs ± 20.3μs	55.4μs ± 12.6μs	x

每 GPU 128 个 token：

每 GPU 128 个 token	EP128	EP64	EP32	EP16	EP8
DeepEP 分发	192μs	186μs	182μs	173μs	163μs
NVLINK 分发	x	x	x	x	83.6μs ± 1.0μs
IBGDA 分发	307.7μs ± 3.0μs	317.4μs ± 1.5μs	427.6μs ± 1.4μs	622.4μs ± 1.7μs	x
IBRC 分发	2038.5μs ± 77.0μs	1669.3μs ± 64.0μs	973.5μs ± 37.9μs	687.1μs ± 12.9μs	x
DeepEP 聚合	369μs	353μs	350μs	329μs	318μs
NVLINK 聚合	x	x	x	x	102.3μs ± 0.6μs
IBGDA 聚合	593.9μs ± 6.6μs	529.9μs ± 6.7μs	481.4μs ± 3.6μs	668.1μs ± 3.4μs	x
IBRC 聚合	1184.8μs ± 79.7μs	1058.5μs ± 49.6μs	916.5μs ± 45.1μs	633.4μs ± 14.0μs	x
Torch 全对全	4972.0μs ± 135.8μs	5418.1μs ± 241.4μs	4225.9μs ± 69.5μs	3213.9μs ± 19.7μs	699.9μs ± 2.2μs
NVLINK NVSHMEM 全对全	x	x	x	x	6585.3μs ± 2.4μs
IBGDA NVSHMEM 全对全	6180.1μs ± 344.7μs	6916.3μs ± 315.4μs	4603.4μs ± 133.1μs	3444.8μs ± 15.3μs	x
IBRC NVSHMEM 全对全	6378.5μs ± 375.9μs	6625.1μs ± 371.3μs	4371.3μs ± 148.8μs	3410.1μs ± 20.2μs	x

C++ 测试

要构建 C++ 测试和基准：

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

要在单个节点上运行全对全测试：

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

要在单个节点上运行全对全基准：

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

引用

如果您使用此代码库或以其他方式认为我们的工作有价值，请引用：

@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}},
}

pplx-kernels 快速上手指南

⚠️ 重要提示：本仓库中基于旧版 NVSHMEM 的内核已弃用。如需使用最新功能，请迁移至新仓库：pplx-garden。以下内容仅针对当前仓库的历史版本或特定兼容需求。

pplx-kernels 是 Perplexity AI 开源的高性能 MoE（混合专家模型）内核库，支持 CUDA Graph、多种通信后端（NVLink, IBGDA, IBRC, EFA）以及计算与通信重叠优化。

环境准备

在开始之前，请确保您的系统满足以下要求：

• GPU 架构：需要支持 CUDA Compute Capability 9.0a (Hopper 架构，如 H100) 或更高。
• 驱动程序与依赖：必须正确安装 NVIDIA 驱动、CUDA Toolkit 以及相关的通信库（如 NVSHMEM, InfiniBand verbs 等）。
• 配置指南：详细的驱动安装和依赖配置步骤，请参考官方文档 Install Driver and Dependencies。

安装步骤

请克隆仓库并执行以下命令进行编译和安装。注意 TORCH_CUDA_ARCH_LIST 需根据您的显卡架构调整（默认示例为 9.0a+PTX）。

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

💡 国内加速建议：如果 pip install 下载缓慢，可指定国内镜像源： pip install dist/*.whl -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

1. 单节点测试与基准测试

安装完成后，您可以立即运行单元测试或基准测试来验证安装是否成功。

运行测试：

pytest -svx --tb=short tests

运行基准测试（All-to-All）：

python3 -m tests.bench_all_to_all

2. 多节点测试（简要）

若需在多节点环境下运行 All-to-All 测试，需先设置环境变量（以 8 卡节点为例）：

export NODE_RANK=0                  # 当前节点排名 (0, 1, ...)
export WORLD_SIZE=8                 # 总进程数 (节点数 * 8)
export WORLD_LOCAL_SIZE=8           # 单节点进程数
export MASTER_ADDR=192.168.1.100    # Rank-0 节点的 IP 地址
export MASTER_PORT=29500
export NVSHMEM_IB_ENABLE_IBGDA=1    # 启用 IBGDA

设置完成后，正常运行测试脚本：

pytest -svx tests/test_all_to_all.py

3. C++ 构建与测试（可选）

如果您需要构建 C++ 层面的测试或基准工具：

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

运行单节点 C++ 测试：

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