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

kpool_expand_topk_kernel

def kpool_expand_topk_kernel[OutLayoutType: TensorLayout, out_origin: MutOrigin, PoolLayoutType: TensorLayout, pool_origin: ImmOrigin, IROLayoutType: TensorLayout, iro_origin: ImmOrigin, CacheLenLayoutType: TensorLayout, OutEngine: TensorEngine, PoolEngine: TensorEngine, IROEngine: TensorEngine, CacheLenEngine: TensorEngine, kpool: Int, pool_topk: Int, always_select_tail: Bool](out_indices: TileTensor[.int32, OutLayoutType, out_origin, Engine=OutEngine], pool_ids: TileTensor[.int32, PoolLayoutType, pool_origin, Engine=PoolEngine], input_row_offsets: TileTensor[.uint32, IROLayoutType, iro_origin, Engine=IROEngine], cache_lengths: TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin, Engine=CacheLenEngine], total_seq_len: Int32)

Turns selected pool ids back into the token positions they cover.

The indexer selects pools; attention reads tokens. Each selected pool expands to the kpool consecutive positions it covers.

An unselected slot expands to -1 in every one of its positions, never to a clamped valid one, which would point attention at a token the indexer did not choose.

With always_select_tail the output carries kpool - 1 further columns holding the query's most recent positions, the ones no complete pool covers yet. Their location comes from the query's visible count, so it tracks the pool currently being filled.

Parameters:

• ​OutLayoutType (TensorLayout): Layout of out_indices.
• ​out_origin (MutOrigin): Origin of out_indices.
• ​PoolLayoutType (TensorLayout): Layout of pool_ids.
• ​pool_origin (ImmOrigin): Origin of pool_ids.
• ​IROLayoutType (TensorLayout): Layout of input_row_offsets.
• ​iro_origin (ImmOrigin): Origin of input_row_offsets.
• ​CacheLenLayoutType (TensorLayout): Layout of cache_lengths.
• ​OutEngine (TensorEngine): Engine of out_indices.
• ​PoolEngine (TensorEngine): Engine of pool_ids.
• ​IROEngine (TensorEngine): Engine of input_row_offsets.
• ​CacheLenEngine (TensorEngine): Engine of cache_lengths.
• ​kpool (Int): Tokens per pool.
• ​pool_topk (Int): Selected pools per token, index_topk // kpool.
• ​always_select_tail (Bool): Whether to append the incomplete trailing pool.

Args:

• ​out_indices (TileTensor[.int32, OutLayoutType, out_origin, Engine=OutEngine]): Output [total_seq_len, pool_topk * kpool + tail], where tail is kpool - 1 when always_select_tail and 0 otherwise.
• ​pool_ids (TileTensor[.int32, PoolLayoutType, pool_origin, Engine=PoolEngine]): Selected pool ids, [total_seq_len, pool_topk], -1 where fewer than pool_topk pools were available.
• ​input_row_offsets (TileTensor[.uint32, IROLayoutType, iro_origin, Engine=IROEngine]): Token row offsets per request, [batch_size + 1].
• ​cache_lengths (TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin, Engine=CacheLenEngine]): Cached-prefix length per request, [batch_size].
• ​total_seq_len (Int32): Number of token rows.