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_multimem_kernel

def broadcast_multimem_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, root: Int)

Broadcast kernel using multimem.st for multicast writes.

Root GPU writes to multicast address, data appears on all GPUs. Only root performs the stores; other GPUs just participate in barriers.

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.
root (Int): Rank of the source GPU whose data is broadcast to all GPUs.