This document is relevant for: Trn2, Trn3

Build Packed Routing Kernel API Reference#

Build stable destination-major source indices without an [N, EP] HBM table.

ids[N] contains peer IDs and counts[1, EP] their exact histogram. The output is a permutation of range(N), grouped by peer while retaining source order within each peer. A row’s packed position is its peer’s starting offset plus the number of earlier rows assigned to that peer. Each tile scans at most 128 rows in SBUF; per-peer cursors carry the prefix across tiles. Float32 prefix arithmetic is exact up to 2**24 rows. Run with LNC=2, with core zero owning the cursors.

Background#

The build_packed_routing kernel builds stable, destination-major source indices for all-to-all-v token routing without materializing an [N, EP] HBM table. Its output is the packed_source_indices input of permute_a2av and unpermute_a2av.

API Reference#

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

build_packed_routing#

nkilib.experimental.collectives.a2av_train.build_packed_routing(ids, counts)#

Build stable destination-major source indices without an [N, EP] HBM table.

Parameters:

  • ids (nl.NkiTensor) – [N] destination peer ID of each source row, in HBM. 1 <= N <= 2**24.

  • counts (nl.NkiTensor) – [1, EP] exact histogram of ids (rows per peer), in HBM. 2 <= EP <= 128.

Returns:

[N, 1] int32 permutation of range(N) grouped by destination peer, preserving source order within each peer. Requires LNC=2.

Return type:

nl.NkiTensor

This document is relevant for: Trn2, Trn3