nki.collectives.all_reduce#

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

Perform an all-reduce 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 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

  • dsts – List of output tensors to store results

  • replica_group – ReplicaGroup defining rank groups for the collective

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