This document is relevant for: Trn2, Trn3

nki.isa.topk#

nki.isa.topk(val_dst: NkiTensor, idx_dst: NkiTensor, src: NkiTensor, n, name=None)[source]#

Find the K largest values and their indices from a source tile using GpSIMD Engine.

Each partition operates independently. The source tile is interpreted as a 16-partition snake layout: elements 0..15 fill partitions 0..15 of free-dim column 0, elements 16..31 fill column 1, etc. The parameter n specifies the total number of active BF16 elements in this snake (which may be less than src_x * 16).

For each partition, the instruction selects the top-K elements by value from the source data and writes:

  • val_dst: the K largest values in ascending order (BF16)

  • idx_dst: the original 0-based indices of those values (uint32)

The source and outputs must be in SBUF. The source must be BF16. K is derived from the free dimension of val_dst.

Estimated instruction cost:

4*n + K GpSIMD Engine cycles per partition, where:

  • n is the number of BF16 elements in the 16P input snake

  • K is the number of top elements to select (= val_dst free dim)

Constraints:

  • Source must be BF16 in SBUF.

  • val_dst must be BF16 in SBUF with shape [par_dim, k].

  • idx_dst must be uint32 in SBUF with shape [par_dim, k].

  • 8 <= n < 65536

  • 1 <= K < 32768 and K <= n

  • src free dim >= ceil(n / 16)

  • partition dim must be a multiple of 16

Parameters:

  • val_dst – Output tile for top-K values (BF16), shape [par_dim, k].

  • idx_dst – Output tile for top-K indices (uint32), shape [par_dim, k].

  • src – Source tile containing the input snake (BF16), shape [par_dim, src_x].

  • n – Number of BF16 elements in the 16P input snake.

Example

def topk_kernel(in_tensor):
    par_dim = in_tensor.shape[0]
    n = 42  # actual number of BF16 elements in the snake
    src_x = (n + 15) // 16  # = 3
    k = 8

    src = nl.ndarray((par_dim, src_x), dtype=nl.bfloat16, buffer=nl.sbuf)
    nisa.dma_copy(dst=src, src=in_tensor)

    val_dst = nl.ndarray((par_dim, k), dtype=nl.bfloat16, buffer=nl.sbuf)
    idx_dst = nl.ndarray((par_dim, k), dtype=nl.uint32, buffer=nl.sbuf)
    nisa.topk(val_dst=val_dst, idx_dst=idx_dst, src=src, n=n)

    return val_dst, idx_dst

This document is relevant for: Trn2, Trn3