nki.collectives.collective_permute#

nki.collectives.collective_permute(srcs, dsts, source_target_pairs)[source]#

Send and receive data between ranks based on explicitly defined source-target pairs.

Each pair (source, target) specifies that data from the source rank should be sent to the target rank. This gives you full control over the communication pattern (e.g., pairwise swaps, arbitrary shuffles).

Prefer collective_permute_implicit() when the communication follows a ring topology, as the hardware can optimize that pattern.

Tensors must reside on HBM. SBUF is not currently supported for collective_permute.

Coalesced collective communication (multiple tensors) is not currently supported; each list parameter must contain exactly one tensor.

Parameters:

  • srcs – List of source tensors to send

  • dsts – List of destination tensors to receive into

  • source_target_pairs – List of (source, target) rank ID pairs