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.NkiTensor, packed_source_indices: nl.NkiTensor, send_counts: nl.NkiTensor, recv_capacity: int, replica_group: ReplicaGroup) → tuple[nl.NkiTensor, nl.NkiTensor]#

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

Parameters:

  • hidden_states (nl.NkiTensor) – [N, H]@HBM routed input rows.

  • packed_source_indices (nl.NkiTensor) – [N]@HBM, int32 permutation of source rows in destination-major order. Every entry must be in [0, N), with peer segments defined by send_counts.

  • send_counts (nl.NkiTensor) – [1, EP]@HBM, int32/uint32 row counts.

  • recv_capacity (int) – Static number of rows in the packed receive output. For top-k routing with unique expert IDs and L local experts per destination rank, a dropless bound is global_token_count * min(top_k, L). This reduces to global_token_count when L == 1. Duplicate expert IDs within one token invalidate this bound.

  • replica_group (ReplicaGroup) – EP replica group.

Returns:

[C, H]@HBM. Source-major packed rows followed by deterministic zero padding.

Returns:

[4, EP]@HBM, uint32. Send counts/displacements and receive counts/displacements, all measured in elements.

Return type:

tuple[nl.NkiTensor, nl.NkiTensor]

Dimensions:

  • T: number of local tokens (SP sharded).

  • H: hidden dimension.

  • EP: number of expert-parallel ranks (== replica_group size).

This document is relevant for: Trn2, Trn3