For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo module

sync

Signal-based GPU synchronization primitives for multi-GPU collective kernels.

Provides the Signal struct for counter-based cross-GPU thread-block synchronization, peer-to-peer memory access helpers, and the barrier utilities consumed by allreduce, scatter, and other collective operations.

comptime values

MAX_GPUS

comptime MAX_GPUS = 8

Maximum number of GPUs supported in the allreduce implementation.

This constant sets the upper bound for the number of GPUS supported in this algorithm.

MAX_NUM_BLOCKS_UPPER_BOUND

comptime MAX_NUM_BLOCKS_UPPER_BOUND = 1024

Maximum number of thread blocks to use for reduction kernels.

This value has been empirically optimized through grid search across different GPU architectures. While this value is optimal for A100 GPUs, H100 GPUs may benefit from more blocks to fully saturate NVLink bandwidth.

NUM_BARRIER_DOMAINS

comptime NUM_BARRIER_DOMAINS = MAX_GPUS

Number of disjoint barrier counter banks held by a Signal buffer.

_multi_gpu_barrier keys its self_counter/peer_counter slots by the in-block thread index, not the global device rank. When subgroup collectives (e.g. a group_size-of-4 TP collective) and full-world collectives share one Signal buffer, they advance the same counter slots at different rates and the generation counters desync -- a subsequent full-world barrier then spins forever waiting for a generation a peer will never publish.

Giving each collective domain its own counter bank makes the barrier histories disjoint so domains can never poison each other. MAX_GPUS banks is the loosest sufficient bound: there can be at most one distinct device-group per GPU.

Callers map by group WIDTH (0 if group_size == num_devices else group_size), so DIFFERENT ops of the same width share a bank. Sharing is sound only while every rank in the domain issues the same barrier sequence: every collective must pair start with end (a start-only one lets a fast rank overwrite a shard its peers are still P2P-reading), and any conditional skip must be group-uniform. Sibling groups are safe regardless -- their rank_sigs are disjoint. The cost is negligible next to the embedded Lamport region.

Structs

• ​Signal: A synchronization primitive for coordinating GPU thread blocks across multiple devices.

Functions

• ​circular_add: Addition modulo n, assuming 0 <= x < n and 0 <= y < n.
• ​enable_p2p: Enable peer-to-peer memory access between all GPU pairs if supported.
• ​group_end: Marks the end of a collective operation group.
• ​group_start: Marks the start of a collective operation group.
• ​init_signal_buffer: Initializes a freshly allocated signal buffer for any comm collective.
• ​is_p2p_enabled: Checks whether P2P access is available between GPUs.