This document is relevant for: Trn2, Trn3

nki.collectives.all_to_all_v_prepare

nki.collectives.all_to_all_v_prepare(srcs: List, dsts: List, replica_group: ReplicaGroup, metadata_tensor, recv_counts_known: bool = False, has_rdispls: bool = False, priority: Optional[int] = None, name: Optional[str] = None) → NkiInstruction

Begin a split all_to_all_v().

A plain all_to_all_v() couples its metadata setup with the execution of the a2a-v collective, so the setup cannot start until the instruction that fills srcs has completed. Splitting it into prepare + trigger lets the setup run ahead and overlap with unrelated compute, hiding its latency. This helps when the per-rank counts are known before the data is — e.g. in MoE, where send_counts follow the router/top-k while srcs are only filled after a reshuffle.

prepare reads metadata_tensor, which must already hold valid data. It does not read the contents of srcs/dsts, so it may be issued before they are filled.

Each prepare must be followed by exactly one all_to_all_v_trigger(), with no other prepare in between. Prepares share a single slot, so an intervening prepare of any kind is rejected.

All arguments match all_to_all_v(); see it for their meaning. The trigger reuses them from this call, so all_to_all_v_trigger() takes only the instruction this call returns.

Current restrictions:

NeuronSwitch-based instances (the Trn3 Ultraserver) do not currently support the prepare/trigger split. Use plain all_to_all_v() on those instance types.

Returns:

The NkiInstruction to pass to the paired all_to_all_v_trigger().

This document is relevant for: Trn2, Trn3