nki.collectives.all_gather#

nki.collectives.all_gather(srcs, dsts, replica_group, collective_dim)[source]#

Perform an all-gather on the given replica group and input/output tensors.

The srcs and dsts parameters accept lists of tensors to support coalesced collective communication, which allows multiple tensors to be gathered in a single collective operation for improved efficiency.

Tensors can reside on either HBM or SBUF. However, mixing memory spaces is not supported: all tensors must be on HBM or all must be on SBUF. Coalesced collective communication (multiple tensors) is only supported when tensors are on HBM.

Parameters:

  • srcs – List of input tensors to gather

  • dsts – List of output tensors to store results

  • replica_group – ReplicaGroup defining rank groups for the collective

  • collective_dim – Dimension along which output tensors are concatenated. Currently only 0 is supported for HBM tensors. For SBUF tensors, 0 or 1 is supported as SBUF collectives currently only operate on 2D tensors with a single free dimension.