nki.collectives#

NKI Collectives#

Collective operations for multi-rank communication.

`all_reduce`	Perform an all-reduce on the given replica group and input/output tensors.
`all_gather`	Perform an all-gather on the given replica group and input/output tensors.
`reduce_scatter`	Perform a reduce-scatter on the given replica group and input/output tensors.
`all_to_all`	Perform an all-to-all on the given replica group and input/output tensors.
`collective_permute`	Send and receive data between ranks based on explicitly defined source-target pairs.
`collective_permute_implicit`	Send and receive data between ranks in a ring, where sources and destinations are implicitly determined by the ring structure during runtime.
`collective_permute_implicit_reduce`	Perform an implicit collective permute with reduction in a ring, where sources and destinations are implicitly determined by the ring structure during runtime.
`collective_permute_implicit_current_processing_rank_id`	Returns the rank ID of the data to be processed in the current ring iteration.
`rank_id`	Get the rank ID of the current rank.

Defines a group of ranks that participate in a collective operation.