This document is relevant for: Trn2, Trn3

Chunked Top-K Kernel API Reference#

Top-k over inp[BxS, vocab] for a vocab wider than nisa.topk’s ceiling.

Background#

The chunked_topk kernel computes top-k over a vocabulary wider than the 65,535-element limit of nisa.topk on a single core. It splits the vocabulary into chunks, takes a top-k of each, and reduces the candidates with a second top-k, with no collective.

API Reference#

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

chunked_topk#

nkilib.experimental.topk.chunked_topk(inp: nl.NkiTensor, config: ChunkedTopkConfig) → Tuple[nl.NkiTensor, nl.NkiTensor]#

Top-k over inp[BxS, vocab] for a vocab wider than nisa.topk’s ceiling.

Parameters:

  • inp (nl.NkiTensor) – [BxS, vocab] bfloat16 tensor in HBM – or, when config.interleaved_batch_groups > 1 (jit-in-jit callers only), a LIST of batch-group slab tensors in the slabbed receive layout, one per producing exchange, so stage-1 tiles data-depend per slab.

  • config (ChunkedTopkConfig) – ChunkedTopkConfig. config.chunks == 1 delegates to gpsimd_topk.

Dimensions:

  • BxS: number of logical rows (tokens)

  • vocab: full reduction length; may exceed 65,535

This document is relevant for: Trn2, Trn3