This document is relevant for: Trn2, Trn3

nki.isa.topk#

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

Find the k largest values and their indices in each 16-partition group using GpSIMD Engine.

Each group of 16 partitions operates independently on its own data. Within a group the source is read in a snake layout: elements 0..15 fill partitions 0..15 at free-dim position 0, elements 16..31 fill position 1, and so on. The parameter n gives the number of active BF16 elements in one snake; storage slots past n are ignored. A tile with 128 partitions therefore holds 8 independent top-k problems, solved by one instruction.

For each group, the instruction selects the top-k elements by value from that group’s own data and writes:

val_dst: the k largest values (BF16)
idx_dst: the original 0-based indices of those values (uint32)

k – the number of results – is not a parameter. It is exactly the per-partition free size of val_dst: to ask for the top 64, pass a val_dst whose free dimension is 64.

Results are written in the same snake layout as the input, so result j of a group lands at partition j % 16, free-dim position j // 16, and idx_dst holds indices into that group’s own snake (0..n-1). All k results therefore fit in the first ceil(k / 16) free-dim positions (for k = 64, positions 0..3):

Positions 0 .. ceil(k/16) - 1 are written. The hardware writes each position across all 16 partitions at once, so when k is not a multiple of 16 the last one carries 16*ceil(k/16) - k slots of unspecified data, which may include NaN.
Positions ceil(k/16) .. k - 1 are never written. They exist only because k is read from the free dimension; a narrower tile would simply request a smaller k.

Read only the first k results, using the mapping above.

By default (sorted=False) the k results are written in an unspecified order – the selected set is well defined, but the order within val_dst/idx_dst must not be relied upon. Pass sorted=True to get each group’s results ordered descending by value; val_dst and idx_dst are reordered consistently, so idx_dst[j] always remains the source index of val_dst[j]. The sort is not stable: ties between equal values may resolve in any order. Sorting costs slightly more than the unsorted path and restricts k to at most 8192.

NaN input is undefined behavior: a NaN may appear in the output even though NaN never compares greater than any value. Mask or replace NaN before calling topk. The CPU simulator warns when it sees NaN input.

The source and outputs must be in SBUF. The source must be BF16.

Estimated instruction cost:

4*n + k GpSIMD Engine cycles per instruction. The 16-partition groups execute in parallel, one per GpSIMD core, so activating more of them does not add cycles.

n is the number of active BF16 elements in one 16-partition 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 <= 16384 and k <= n
k <= 8192 when sorted=True
src free dim >= ceil(n / 16)
partition dim must be a multiple of 16
src can have at most three free dimensions (maximum rank 4); val_dst and idx_dst can have at most two free dimensions (maximum rank 3). These ranks include the partition dimension.

Parameters:

val_dst – Output tile for the top-k values (BF16), shape [par_dim, k].
idx_dst – Output tile for the top-k indices (uint32), shape [par_dim, k].
src – Source tile containing the input snake (BF16), shape [par_dim, src_x].
n – Number of active BF16 elements in each 16-partition input snake.
sorted – If True, sort the outputs of each 16-partition group descending by value. Defaults to False, which leaves the output order unspecified.

Example

import nki
import nki.isa as nisa
import nki.language as nl
import numpy as np
import pytest

@nki.jit
def topk_sorted_kernel(in_tensor):
    '''Select the 8 largest of 42 values per 16-partition group, sorted descending.'''
    par_dim = in_tensor.shape[0]
    n = 42  # number of active BF16 elements in the snake
    src_x = (n + 15) // 16  # = 3 columns; the 48 - 42 = 6 trailing slots are ignored
    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, sorted=True)

    # The k results land in the first ceil(k/16) = 1 free-dim position: result j at
    # partition j % 16, position j // 16. Kernel outputs must live in shared_hbm.
    val_out = nl.ndarray((par_dim, k), dtype=nl.bfloat16, buffer=nl.shared_hbm)
    idx_out = nl.ndarray((par_dim, k), dtype=nl.uint32, buffer=nl.shared_hbm)
    nisa.dma_copy(dst=val_out, src=val_dst)
    nisa.dma_copy(dst=idx_out, src=idx_dst)
    return val_out, idx_out

This document is relevant for: Trn2, Trn3

nki.isa.topk

Contents

nki.isa.topk#