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

Mojo function

fill_invalid_topk_kernel

def fill_invalid_topk_kernel[IROLayoutType: TensorLayout, iro_origin: ImmOrigin, cache_lengths_layout: TensorLayout, use_causal_mask: Bool](output_indices: Pointer[Int32, MutAnyOrigin, _safe=False], topk_indices: Pointer[Int32, MutAnyOrigin, _safe=False], input_row_offsets: TileTensor[DType.uint32, IROLayoutType, iro_origin], cache_lengths: TileTensor[DType.uint32, cache_lengths_layout, ImmutAnyOrigin], total_seq_len: Int, top_k: Int, effective_k: Int)

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.
use_causal_mask (Bool): Whether each token is restricted to keys up to its own position.

Args:

output_indices (Pointer[Int32, MutAnyOrigin, _safe=False]): Output buffer of shape [total_seq_len, top_k] with invalid positions set to -1.
topk_indices (Pointer[Int32, MutAnyOrigin, _safe=False]): Compact top-k index buffer of shape [total_seq_len, effective_k] produced by topk_gpu.
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 used to compute the number of keys each token may attend to.
total_seq_len (Int): Number of token rows in output_indices.
top_k (Int): Row stride of output_indices and the requested number of selections per token.
effective_k (Int): Row stride of topk_indices and the actual number of computed selections, min(top_k, max_num_keys).