This document is relevant for: Trn2, Trn3

GPSIMD Top-K Kernel API Reference#

Top-k over the last dimension using the GpSIMD nisa.topk instruction.

Background#

The gpsimd_topk kernel computes top-k over the last dimension using the GpSIMD nisa.topk instruction. The companion create_gpsimd_topk_config builds a GpsimdTopkConfig from an input shape (2D or 3D) and parameters.

API Reference#

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

create_gpsimd_topk_config#

nkilib.experimental.topk.create_gpsimd_topk_config(inp_shape: Tuple, inp_dtype: np.dtype, k: int, sorted: bool = True, num_programs: int = 2) → GpsimdTopkConfig#

Build a GpsimdTopkConfig from an input shape (2D or 3D) and parameters.

gpsimd_topk#

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

Top-k over the last dimension using the GpSIMD nisa.topk instruction.

Parameters:

  • inp (nl.NkiTensor) – [BxS, V] bfloat16 input tensor in HBM.

  • config (GpsimdTopkConfig) – GpsimdTopkConfig describing the problem and sharding.

Notes:

  • Each vocab row is loaded into its 16-partition nisa.topk snake with a BLOCKED / contiguous DMA: partition p reads the contiguous HBM run inp[row, p*n_cols:(p+1)*n_cols] (free-stride 1), so snake position s = p + 16*c holds vocab index p*n_cols + c. The 16-partition snake LAYOUT is mandated by nisa.topk, but the ORDER of the placement is a free bijection (any order yields the same top-k value set). The blocked order is chosen so the load is contiguous; the alternative “snake position i == vocab index i” fill gives an identity remap but forces a transpose-on-load (.ap [[1,16],[16,n]], free-stride 16 -> non-contiguous per partition), which is avoided here. The returned snake-position indices are remapped back to vocab space on-chip before the index store.

  • 8 rows (groups of 16 partitions) are processed per nisa.topk call.

  • The hardware nisa.topk output order is not relied upon: the K values + paired snake-position indices are de-snaked and then sorted descending on-chip; the snake->vocab index remap is applied to the final k indices.

  • Two-phase structure: phase 1 runs nisa.topk per 8-row tile and de-snakes the K (value, index) pairs into per-row HBM buffers; phase 2 runs the descending sort ONCE over up to 128 rows (one row per partition) instead of once per 8-row tile. max8 / nc_match_replace8 / nc_n_gather are per-partition free-dim ops, so widening the sort from 8 to up to 128 partitions is free Vector-engine parallelism and removes the redundant per-tile sort passes (the measured HW bottleneck).

  • config.sorted gates only the phase-2 descending sort. When False the sort is skipped and the K results are compacted to [0, k) in arbitrary order; the value set and value<->index pairing are unchanged.

Dimensions:

  • BxS: number of rows (flattened batch*sequence)

  • V: vocab size (reduction dimension), 8 <= V < 65536

This document is relevant for: Trn2, Trn3