nki.collectives.all_to_all#

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

Perform an all-to-all 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 redistributed in a single collective operation for improved efficiency.

Tensors must reside on HBM. SBUF is not currently supported for all-to-all.

Parameters:

  • srcs – List of input tensors to redistribute

  • dsts – List of output tensors to store results

  • replica_group – ReplicaGroup defining rank groups for the collective

  • collective_dim – Dimension along which input tensors are split and output tensors are concatenated. Currently only 0 is supported.