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

allgather

def allgather[dtype: DType, in_layout: TensorLayout, in_origin: ImmOrigin, out_layout: TensorLayout, out_origin: MutOrigin, //, ngpus: Int](inputs: Array[TileTensor[dtype, in_layout, in_origin], ngpus], outputs: Array[TileTensor[dtype, out_layout, out_origin], (ngpus * ngpus)], list_of_ctx: List[DeviceContext])

Performs an allgather across all GPUs via the vendor CCL library.

Each GPU contributes its local inputs[i] chunk; after the call every output slot outputs[dev * ngpus + src] on device dev holds a copy of the input from device src. Uses ncclAllGather (NVIDIA) or rcclAllGather (AMD) internally, wrapped inside an NCCL group for correct pipelining.

Parameters:

dtype (DType): Element data type of all input and output tensors.
in_layout (TensorLayout): TensorLayout of each per-GPU input.
in_origin (ImmOrigin): Origin tag for the input tensors.
out_layout (TensorLayout): TensorLayout of each per-GPU output slot.
out_origin (MutOrigin): Mutable origin tag for the output tensors.
ngpus (Int): Number of participating GPUs.

Args:

inputs (Array[TileTensor[dtype, in_layout, in_origin], ngpus]): Per-GPU input TileTensors; must all have the same element count.
outputs (Array[TileTensor[dtype, out_layout, out_origin], (ngpus * ngpus)]): Flat array of ngpus * ngpus output TileTensors. Slot i * ngpus + j on device i receives device j's input.
list_of_ctx (List[DeviceContext]): Device context for each GPU; length must equal ngpus.

Raises:

If ngpus < 1 or ngpus > MAX_GPUS. If len(list_of_ctx) != ngpus. If any input element count differs from inputs[0]'s count. If the CCL collective call fails.