NCCL Extensions is a repository of communication patterns for AI use cases, built on top of NCCL device and host APIs. It speeds up tensor communication for workloads like MoE token shuffle and reinforcement learning weight rollout.
This is an evolving space, and the content here is under constant development and subject to change. We will continue exploring it and welcome your contributions.
nccl_ep/ — Expert Parallelism
Optimized dispatch and combine primitives for Mixture-of-Experts (MoE) token routing, built on NCCL's Device API (LSA and GIN operations).
nccl_m2n/ — Mesh-to-Mesh Rollout
Experimental library for resharding a tensor between two disjoint groups of GPU processes (e.g. trainer and inference ranks) in a single, zero-copy call, built on NCCL's window API.
python/ — Python bindings
Python package (nccl-extensions) providing Pythonic wrappers for nccl_ep
and nccl_m2n as nccl.ep and nccl.m2n. See
python/README.md for details.
This repo vendors NCCL as a git submodule. Clone with:
git clone --recursive <repo-url>(or git submodule update --init --recursive after a normal clone). See each
subproject's README for build instructions.
We welcome contributions! See CONTRIBUTING.md to get started.
This project is licensed under the Apache License, Version 2.0 — see LICENSE.txt for details.