This document is relevant for: Trn2, Trn3

Unpermute A2AV Kernel API Reference#

All-to-all-v combine and unpermute to original token order.

Accepts fixed-stride expert output (source s at rows [s*T, s*T + recv_counts[s])), re-permutes into a packed send buffer using cumsum(recv_counts) offsets, builds the (4, EP) metadata, exchanges via ncc.all_to_all_v, then scatter-adds received rows to original token positions via send_indices.

Background#

The unpermute_a2av kernel is the MoE training combine step: it exchanges expert output via ncc.all_to_all_v and unpermutes tokens back to their original order. It supports top-k > 1 because the scatter accumulates across all EP contributions.

API Reference#

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

unpermute_a2av#

nkilib.experimental.collectives.a2av_train.unpermute_a2av(output: nl.NkiTensor, packed_source_indices: nl.NkiTensor, send_counts: nl.NkiTensor, recv_counts: nl.NkiTensor, replica_group: ReplicaGroup) → nl.NkiTensor#

All-to-all-v combine and unpermute to original token order.

Parameters:

  • output (nl.NkiTensor) – [C, H]@HBM packed expert output.

  • packed_source_indices (nl.NkiTensor) – [N]@HBM, int32 permutation mapping packed return rows to original routed-row indices, in the same destination-major order used for dispatch.

  • send_counts (nl.NkiTensor) – [1, EP]@HBM original dispatch send counts.

  • recv_counts (nl.NkiTensor) – [1, EP]@HBM original dispatch receive counts.

  • replica_group (ReplicaGroup) – EP replica group.

Returns:

[N, H]@HBM in original routed-row order.

Return type:

nl.NkiTensor

Notes:

  • Supports top-k > 1 because the scatter accumulates across all EP contributions.

  • With LNC=2, EP is partitioned across cores and partial results are combined via nisa.sendrecv before core 0 writes the final output.

Dimensions:

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

  • H: hidden dimension.

This document is relevant for: Trn2, Trn3