Mojo function

allgather

allgather[type: DType, rank: Int, ngpus: Int, //](input_buffers: InlineArray[NDBuffer[type, rank, MutableAnyOrigin], ngpus], output_buffers: InlineArray[NDBuffer[type, rank, MutableAnyOrigin], (ngpus * ngpus)], ctxs: List[DeviceContext])

Performs all-gather across GPUs with variadic output.

Each device receives individual copies of all input buffers.

Parameters:

• ​type (DType): DType - The data type of tensor elements.
• ​rank (Int): Int - Number of dimensions in input tensors.
• ​ngpus (Int): Int - Number of GPUs participating in all-gather.

Args:

• ​input_buffers (InlineArray[NDBuffer[type, rank, MutableAnyOrigin], ngpus]): Input buffers from each GPU.
• ​output_buffers (InlineArray[NDBuffer[type, rank, MutableAnyOrigin], (ngpus * ngpus)]): Flat array of ngpus * ngpus output buffers. Layout: output_buffers[device_idx * ngpus + input_idx] contains device_idx's copy of input_idx's data.
• ​ctxs (List[DeviceContext]): List of device contexts for participating GPUs.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit