nccl-extensions (Python)

Python bindings for the nccl-extensions communication libraries.

Package layout

This package installs into the nccl namespace, so the import paths are nccl.ep and nccl.m2n:

import nccl.ep as ep
import nccl.m2n as m2n

It contributes exactly three directories to that namespace, and no nccl/__init__.py:

Install

CUDA_HOME=/usr/local/cuda pip install -e python/

Building requires a CUDA toolkit and a Cython toolchain.

Stage native artifacts before building a distributable wheel:

python/nccl/ep/lib/cu{12,13}/libnccl_ep.so
python/nccl/ep/include/**
python/nccl/m2n/lib/cu{12,13}/libnccl_m2n.so
python/nccl/m2n/include/**

For a complete Linux wheel build, the internal build_assets/build_wheels.sh script builds and stages these artifacts in a temporary project copy before running cibuildwheel. It does not stage artifacts into the source package tree.

Missing shared libraries emit explicit build warnings by default. Set NCCL_EXTENSIONS_REQUIRE_NATIVE_LIBS=1 to turn a missing library into a build error; the production wheel script sets this automatically. With the default value 0, the resulting wheel is not self-contained and needs compatible external libraries at runtime.

The sdist is source-only and excludes native shared libraries. Building a wheel from it must stage the native libraries at the paths above to bundle them, or provide compatible external libraries for runtime loading.

Pick a CUDA-variant extra to pull in the matching runtime stack (they forward to nccl4py's cu12 / cu13 extras, and are mutually exclusive):

pip install -e 'python/[cu13]'

At runtime, the installed cuda.bindings major selects the matching bundled lib/cu12 or lib/cu13 native libraries. NCCL EP compiles kernels at runtime with NCCL_EP_JIT_NVCC, NVCC, or the compiler under CUDA_HOME; that compiler and its headers must match the selected CUDA major.

Do not run Python from inside python/. There is no nccl/__init__.py there, so that directory resolves only as a namespace portion and these modules become invisible. Always go through the editable install.

Regenerating the bindings

Everything under nccl/_extensions/bindings/ is generated and checked in. Do not edit it by hand — re-run build_assets/generate_cython.py after changing a public header, config, or template and commit the result.