This document is relevant for: Trn2, Trn3

Permute A2AV Kernel API Reference#

Permute tokens by destination EP rank and dispatch via all-to-all-v.

Gathers tokens per destination EP rank into a packed send buffer at cumsum(send_counts) row offsets, builds the (4, EP) uint32 metadata tensor, then exchanges via ncc.all_to_all_v. The receive-side layout is determined by the runtime (packed cumsum of recv_counts).

Background#

The permute_a2av kernel is the MoE training dispatch step: it permutes tokens by their destination expert-parallel (EP) rank and exchanges them across ranks via ncc.all_to_all_v.

API Reference#

Source code for this kernel API can be found at: permute_a2av.py

permute_a2av#

nkilib.experimental.collectives.a2av_train.permute_a2av(hidden_states: nl.ndarray, send_indices: nl.ndarray, send_counts: nl.ndarray, replica_group: ReplicaGroup) → tuple[nl.ndarray, nl.ndarray]#

Permute tokens by destination EP rank and dispatch via all-to-all-v.

Parameters:

hidden_states (nl.ndarray) – [T, H]@HBM. Input tokens (bf16/fp16/fp32).
send_indices (nl.ndarray) – [T, EP]@HBM, int32. Source-row indices per destination rank. Entries with value = T are skipped (via oob_mode.skip), so slots beyond send_counts[d] for destination d are no-ops.
send_counts (nl.ndarray) – [1, EP]@HBM, int32/uint32. Tokens sent per dest rank.
replica_group (ReplicaGroup) – EP replica group.

Returns:

[EP*T, H]@HBM. Received tokens packed by the runtime at cumsum(recv_counts) offsets.

Return type:

nl.ndarray

Returns:

[4, EP]@HBM, uint32. After the collective, row 2 contains the runtime-populated recv_counts * H.