ttnn.experimental.topk_large_indices

ttnn.experimental.topk_large_indices(input_tensor, k, valid_length) → None

Experimental Top-K over the last dimension of a row-major BFLOAT16 tensor. This op is Blackhole-only.

Returns a ROW_MAJOR UINT32 tensor containing sorted descending top-k indices. The output shape matches the input shape except that the last dimension is k.

This op is intended for large row-major rows. Internally it snaps k to the nearest supported LLK size and streams each input row in LLK-sized windows. Input values equal to -inf produce the sentinel index 0xFFFFFFFF when they survive into the final top-k result.

K constraints:

  • k must be in [16, 2048];

  • k must be a multiple of 16;

  • the internal LLK window is snapped to 512, 1024, or 2048 elements.

Input tensor constraints:

  • the input tensor must be allocated on a Blackhole device;

  • rank must be >= 1;

  • all leading dimensions are flattened into independent rows;

  • the flattened leading-dimension row count must fit in uint32_t;

  • the last dimension is the input row length;

  • the flattened row count must be > 0;

  • the last dimension must be >= k and <= 1,073,741,824 elements.

valid_length (optional):

  • restricts the search to the first valid_length columns of each row;

  • the remaining columns are ignored – neither read nor ranked – so an over-allocated row whose tail is stale can be searched without slicing it;

  • must be in [k, last dimension]; defaults to the full last dimension;

  • applied at runtime (no recompile), so a loop growing valid_length reuses one program.

Parameters:

  • input_tensor – device tensor with ROW_MAJOR layout and BFLOAT16 dtype.

  • k – required number of indices to return.

  • valid_length – optional number of leading columns to search (default: full width).