ttnn.experimental.topk_large_indices

ttnn.experimental.topk_large_indices(input_tensor, k, valid_length, valid_length_tensor, valid_length_offset, subdevice_id, sub_core_grids, valid_end_tensor) → 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.

Core selection:

  • subdevice_id selects a TENSIX subdevice from the active manager;

  • sub_core_grids optionally restricts work to a fully contained CoreRangeSet;

  • the resolved grid is structural and participates in the program-cache key;

  • kernels and circular buffers are created only on the resolved grid.

Parameters:

  • input_tensordevice tensor with ROW_MAJOR layout and BFLOAT16 dtype.

  • krequired number of indices to return.

  • valid_lengthoptional number of leading columns to search (default: full width).

  • valid_length_tensoroptional 1-element UINT32 row-major DRAM tensor read on-device. Use for dynamic valid lengths during trace replay. Mutually exclusive with valid_length.

  • valid_length_offsetconstant added to valid_length_tensor[0] (default 0).

  • subdevice_idoptional active subdevice containing the top-k workers.

  • sub_core_gridsoptional core restriction within subdevice_id.

  • valid_end_tensoroptional 1-element UINT32 row-major DRAM tensor holding the REAL token end; caps the derived search length at ceil32(valid_end). Requires valid_length_tensor. Without it a padded chunk ranks columns that were never written. Pass the SAME bound to ring_indexer_score_dsa’s valid_end_tensor – a looser score with a tighter top-k drops real keys, the reverse ranks a stale tail.