cuda.coop: Cooperative Group Primitives

cuda.coop brings CCCL’s optimized cooperative algorithms to Python GPU kernels. Use it when threads need to work together, such as loading a tile of values or arranging data for the next stage of a computation. These operations run inside a kernel, where you can combine them with your own code and reuse algorithms maintained and tuned for NVIDIA GPUs.

The Numba-CUDA-MLIR integration compiles cooperative calls inside Python GPU kernels. You keep Numba’s kernel syntax and launch conventions and use cuda.coop for the cooperative parts of the work.

Get started

Start with the Numba-CUDA-MLIR Programming Guide for complete kernels and host launch code. Install the extra matching your CUDA major version:

python -m pip install "cuda-coop[numba-cuda-mlir-cu13]"
# Use numba-cuda-mlir-cu12 with CUDA 12.

The guide covers requirements and kernel reuse.

Inside a kernel

A cooperative operation acts on a group of threads, usually a warp or a block. Each thread owns part of the group’s data, held in ThreadData. Together, those parts form a tile.

These calls copy a full tile from one array to another:

# Inside a kernel, with `coop` imported from `cuda`:
block = coop.this_block()
items = coop.ThreadData(items_per_thread)

coop.load(block, source, items, offset=offset)
coop.store(block, destination, items, offset=offset)

The kernel supplies the tile’s offset and a compile-time items_per_thread count. With 128 threads and four items per thread, the block processes 512 values. All threads in the group must reach the cooperative calls, including when only part of the tile is valid.

The programming concepts explain participation, data layouts, operation results, and temporary storage.

Common API and compiler extensions

Use from cuda import coop for the common API. It describes the cooperative part of a kernel independently of its compiler.

The qualified cuda.coop.numba_mlir namespace provides Numba-specific operands and controls. The Numba API comparison explains when to use it.

Import your kernel compiler before cuda.coop to activate its integration automatically. For varying import order, use explicit backend registration.

Explore further

The interactive visualizations show which values each thread owns and how an operation moves them.

Use the API reference for signatures and return values, and configuration for package and compiler settings.

To work on the integration, follow a kernel through the Numba-CUDA-MLIR Developer Guide.