nki.collectives.reduce_scatter#

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

Perform a reduce-scatter 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 reduced and scattered 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 reduce and scatter

  • 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. Currently only 0 is supported.

  • op – The reduction operation to perform (nl.add, nl.minimum, or nl.maximum)