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_1stage_kernel

def broadcast_pull_1stage_kernel[dtype: DType, layout: TensorLayout, BLOCK_SIZE: Int, ngpus: Int, simd_width: Int = simd_width_of[dtype, get_gpu_target()]()](output: TileTensor[dtype, layout, MutAnyOrigin], input: TileTensor[dtype, layout, ImmutAnyOrigin], rank_sigs: Array[Pointer[Signal, MutAnyOrigin, _safe=False], Int(8)], my_rank: Int)

Single-stage pull broadcast kernel: each GPU reads root's input directly.

All GPUs participate in the start and end barriers; after the start barrier every GPU copies the root's input buffer to its own output buffer using a grid-strided vectorized load/store loop. This one-stage path is preferred for small messages (up to a few MiB) and for 2-GPU configurations where the 2-stage scatter/gather overhead is not justified.

Parameters:

dtype (DType): Element data type of the input and output tensors.
layout (TensorLayout): TensorLayout shared by both tensors.
BLOCK_SIZE (Int): Number of threads per thread block.
ngpus (Int): Number of GPUs participating in the broadcast.
simd_width (Int): Vector width used for memory access (defaults to the device-native SIMD width for dtype).

Args:

output (TileTensor[dtype, layout, MutAnyOrigin]): Output TileTensor for this GPU.
input (TileTensor[dtype, layout, ImmutAnyOrigin]): Input TileTensor (root's data, readable via P2P).
rank_sigs (Array[Pointer[Signal, MutAnyOrigin, _safe=False], Int(8)]): Per-GPU Signal pointers for barrier synchronization.
my_rank (Int): Rank of this GPU in the communicator.