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 function

broadcast_pull_2stage_kernel

def broadcast_pull_2stage_kernel[dtype: DType, OutputLayout: TensorLayout, ngpus: Int, *, BLOCK_SIZE: Int](result: TileTensor[dtype, OutputLayout, MutAnyOrigin], root_input_ptr: Pointer[Scalar[dtype], ImmutAnyOrigin, _safe=False], rank_sigs: Array[Pointer[Signal, MutAnyOrigin, _safe=False], Int(8)], num_elements: Int, my_rank: Int, root: Int)

Two-stage broadcast: scatter from root, then allgather among all GPUs.

Stage 1 (Scatter): Root's data is split into ngpus chunks. Each GPU reads its assigned chunk directly from root's input buffer and writes it to its signal payload. Non-root GPUs also write to their result buffer. Root copies all N elements from source to dest (local operation).

Stage 2 (Allgather): Non-root GPUs gather the remaining chunks from all other GPUs' signal payloads (including root's). Root skips this stage since it already has all data.

Parameters:

dtype (DType): Data dtype of tensor elements.
OutputLayout (TensorLayout): Layout of the output TileTensor.
ngpus (Int): Number of GPUs participating.
BLOCK_SIZE (Int): Number of threads per block.

Args:

result (TileTensor[dtype, OutputLayout, MutAnyOrigin]): Output TileTensor for broadcast result.
root_input_ptr (Pointer[Scalar[dtype], ImmutAnyOrigin, _safe=False]): Pointer to root's input data (all GPUs read from this).
rank_sigs (Array[Pointer[Signal, MutAnyOrigin, _safe=False], Int(8)]): Signal pointers for synchronization. IMPORTANT: Signal pointers have trailing buffers for communication.
num_elements (Int): Number of elements to broadcast.
my_rank (Int): Current GPU rank.
root (Int): Root GPU rank (source of broadcast).