Skip to content

Repository files navigation

NCCL Extensions

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.

What's Inside

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.

Getting Started

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.

Contributing

We welcome contributions! See CONTRIBUTING.md to get started.

License

This project is licensed under the Apache License, Version 2.0 — see LICENSE.txt for details.

About

Communication patterns for AI, built on top of NCCL device and host APIs

Resources

Contributing

Security policy

Stars

28 stars

Watchers

2 watching

Forks

Releases

Packages

Contributors

Languages