For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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 scope its own counter bank makes the barrier histories
disjoint so scopes can never poison each other. MAX_GPUS banks is the loosest
sufficient bound: there can be at most one distinct device-group per GPU, so a
1 + group_start // group_size scope mapping never exceeds this. The cost is
negligible -- the counters are tiny 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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!