This document is relevant for: Trn2, Trn3

nki.collectives.all_gather_v_prepare

nki.collectives.all_gather_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_gather_v().

A plain all_gather_v() couples its metadata setup with the execution of the ag-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.

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_gather_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_gather_v(); see it for their meaning and current restrictions.

Current restrictions:

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

Returns:

The NkiInstruction to pass to the paired all_gather_v_trigger().

This document is relevant for: Trn2, Trn3