For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

topk_row_bounds_kernel

def topk_row_bounds_kernel[IROLayoutType: TensorLayout, iro_origin: ImmOrigin, cache_lengths_layout: TensorLayout, use_causal_mask: Bool](row_bounds: Pointer[Int32, MutAnyOrigin], input_row_offsets: TileTensor[DType.uint32, IROLayoutType, iro_origin], cache_lengths: TileTensor[DType.uint32, cache_lengths_layout, ImmutAnyOrigin], total_seq_len: Int32, max_num_keys: Int32)

Compute each token row's live-key count for the bounded top-k.

Writes row_bounds[token] = min(num_keys, max_num_keys) with num_keys from the shared indexer_key_bound helper (causal: cache_len + local_seq_idx + 1; non-causal: cache_len + seq_len).

This is exactly the range the scorers write for that row (they compute the same helper's bound), so a top-k clamped to it reads only written score slots. max_num_keys may be a capture-time upper bound far above the batch's real lengths; the clamp keeps every bound within the row stride.

Parameters:

• ​IROLayoutType (TensorLayout): Layout of the input_row_offsets tensor.
• ​iro_origin (ImmOrigin): Origin of the input_row_offsets tensor.
• ​cache_lengths_layout (TensorLayout): Layout of the cache_lengths tensor.
• ​use_causal_mask (Bool): Whether each token is restricted to keys up to its own position.

Args:

• ​row_bounds (Pointer[Int32, MutAnyOrigin]): Output buffer of shape [total_seq_len].
• ​input_row_offsets (TileTensor[DType.uint32, IROLayoutType, iro_origin]): Ragged row offsets per batch, length batch_size + 1.
• ​cache_lengths (TileTensor[DType.uint32, cache_lengths_layout, ImmutAnyOrigin]): Per-batch cached-prefix length.
• ​total_seq_len (Int32): Number of token rows.
• ​max_num_keys (Int32): Row stride of the scores buffer (upper bound on any row's key count).