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

fill_invalid_topk_kernel

def fill_invalid_topk_kernel[IROLayoutType: TensorLayout, iro_origin: ImmOrigin, cache_lengths_layout: TensorLayout, IROEngine: TensorEngine, CacheLengthsEngine: TensorEngine, use_causal_mask: Bool, kpool: Int = Int(1)](output_indices: Pointer[Int32, MutAnyOrigin], topk_indices: Pointer[Int32, MutAnyOrigin], input_row_offsets: TileTensor[.uint32, IROLayoutType, iro_origin, Engine=IROEngine], cache_lengths: TileTensor[.uint32, cache_lengths_layout, ImmutAnyOrigin, Engine=CacheLengthsEngine], total_seq_len: Int32, top_k: Int32, effective_k: Int32)

Scatter the compact top-k indices into the top_k-strided output, with invalid positions set to -1.

topk_gpu wrote valid indices to positions [0, effective_k) of the compact topk_indices buffer (row stride effective_k). This kernel copies them into output_indices (row stride top_k), setting positions that should be -1:

• Positions [effective_k, top_k): no computed value.
• Positions where k_idx >= num_keys for that token.
• Positions where the index VALUE >= num_keys (topk selected an invalid key).

Output shape: [total_seq_len, top_k].

With causal masking, each token can only see keys up to its position: num_keys = cache_len + local_seq_idx + 1 Without causal masking, each token can see all keys in the batch: num_keys = cache_len + seq_len

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.
• ​IROEngine (TensorEngine): Engine policy of the input_row_offsets tensor.
• ​CacheLengthsEngine (TensorEngine): Engine policy of the cache_lengths tensor.
• ​use_causal_mask (Bool): Whether each token is restricted to keys up to its own position.
• ​kpool (Int): Tokens per pooled cache row. 1 scores one row per token; k > 1 scores one pooled key per k consecutive tokens, so every candidate count and the caller's top_k are pool-granular.

Args:

• ​output_indices (Pointer[Int32, MutAnyOrigin]): Output buffer of shape [total_seq_len, top_k] with invalid positions set to -1.
• ​topk_indices (Pointer[Int32, MutAnyOrigin]): Compact top-k index buffer of shape [total_seq_len, effective_k] produced by topk_gpu.
• ​input_row_offsets (TileTensor[.uint32, IROLayoutType, iro_origin, Engine=IROEngine]): Ragged row offsets per batch, length batch_size + 1.
• ​cache_lengths (TileTensor[.uint32, cache_lengths_layout, ImmutAnyOrigin, Engine=CacheLengthsEngine]): Per-batch cached-prefix length used to compute the number of keys each token may attend to.
• ​total_seq_len (Int32): Number of token rows in output_indices.
• ​top_k (Int32): Row stride of output_indices and the requested number of selections per token.
• ​effective_k (Int32): Row stride of topk_indices and the actual number of computed selections, min(top_k, max_num_keys).