NVSHMEM4Py

NVSHMEM4Py exposes NVSHMEM host, device, and on-stream interfaces to Python, with interoperability for CUDA Python, Numba, PyTorch, CuPy, and CuTe DSL.

Quick Start

Follow the NVSHMEM4Py installation guide to install the package matching the CUDA major version. This example requires two supported GPUs on one Linux node and mpi4py built for the active MPI installation.

From the repository root:

mpirun -np 2 python nvshmem4py/examples/hello.py

The output includes one line from each PE, in either order:

Hello from PE 0 of 2
Hello from PE 1 of 2

The example selects one GPU per local MPI rank and initializes NVSHMEM from MPI_COMM_WORLD. See the main NVSHMEM quick start for platform requirements and the launching guide for other launchers.

Building Wheels

Wheel targets are part of the NVSHMEM CMake build. See the NVSHMEM4Py CMake configuration for supported options and defaults. For example, to build a Python 3.12 wheel for CUDA 12:

cmake -S . -B build \
  -DNVSHMEM4PY_BUILD_ALL_WHEELS=OFF \
  -DNVSHMEM4PY_PYTHON_VERSIONS=3.12 \
  -DNVSHMEM4PY_CUDA_VERSIONS=12
cmake --build build --target build_nvshmem4py_wheel_cu12_3.12

The wheel is written to build/dist/ and can be installed as described in the NVSHMEM4Py installation guide.

Examples

The examples cover host, device, on-stream, and framework-interoperability workflows. See the main NVSHMEM README for additional documentation and project information.