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 function

broadcast_pull_2stage_kernel

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

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.
• ​result_engine (TensorEngine): Engine of the output tensor.
• ​BLOCK_SIZE (Int): Number of threads per block.

Args:

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