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_compress_kernel

def kpool_compress_kernel[dtype: DType, KLayoutType: TensorLayout, k_origin: ImmOrigin, GateLayoutType: TensorLayout, gate_origin: ImmOrigin, ApeLayoutType: TensorLayout, ape_origin: ImmOrigin, IROLayoutType: TensorLayout, iro_origin: ImmOrigin, PROLayoutType: TensorLayout, pro_origin: ImmOrigin, CacheLenLayoutType: TensorLayout, OutLayoutType: TensorLayout, out_origin: MutOrigin, KEngine: TensorEngine, GateEngine: TensorEngine, ApeEngine: TensorEngine, IROEngine: TensorEngine, PROEngine: TensorEngine, CacheLenEngine: TensorEngine, OutEngine: TensorEngine, head_dim: Int, kpool: Int](pooled: TileTensor[dtype, OutLayoutType, out_origin, Engine=OutEngine], k: TileTensor[dtype, KLayoutType, k_origin, Engine=KEngine], gate: TileTensor[dtype, GateLayoutType, gate_origin, Engine=GateEngine], ape: TileTensor[.float32, ApeLayoutType, ape_origin, Engine=ApeEngine], input_row_offsets: TileTensor[.uint32, IROLayoutType, iro_origin, Engine=IROEngine], pool_row_offsets: TileTensor[.uint32, PROLayoutType, pro_origin, Engine=PROEngine], cache_lengths: TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin, Engine=CacheLenEngine])

Builds one pooled key per block; one thread per channel.

Parameters:

• ​dtype (DType): Element type of k, gate and pooled.
• ​KLayoutType (TensorLayout): Layout of k.
• ​k_origin (ImmOrigin): Origin of k.
• ​GateLayoutType (TensorLayout): Layout of gate.
• ​gate_origin (ImmOrigin): Origin of gate.
• ​ApeLayoutType (TensorLayout): Layout of ape.
• ​ape_origin (ImmOrigin): Origin of ape.
• ​IROLayoutType (TensorLayout): Layout of input_row_offsets.
• ​iro_origin (ImmOrigin): Origin of input_row_offsets.
• ​PROLayoutType (TensorLayout): Layout of pool_row_offsets.
• ​pro_origin (ImmOrigin): Origin of pool_row_offsets.
• ​CacheLenLayoutType (TensorLayout): Layout of cache_lengths.
• ​OutLayoutType (TensorLayout): Layout of pooled.
• ​out_origin (MutOrigin): Origin of pooled.
• ​KEngine (TensorEngine): Engine of k.
• ​GateEngine (TensorEngine): Engine of gate.
• ​ApeEngine (TensorEngine): Engine of ape.
• ​IROEngine (TensorEngine): Engine of input_row_offsets.
• ​PROEngine (TensorEngine): Engine of pool_row_offsets.
• ​CacheLenEngine (TensorEngine): Engine of cache_lengths.
• ​OutEngine (TensorEngine): Engine of pooled.
• ​head_dim (Int): Channels per key; also the block width.
• ​kpool (Int): Tokens per pool.

Args:

• ​pooled (TileTensor[dtype, OutLayoutType, out_origin, Engine=OutEngine]): Output [total_pools, head_dim], where total_pools is the last entry of pool_row_offsets.
• ​k (TileTensor[dtype, KLayoutType, k_origin, Engine=KEngine]): Layer-normed indexer keys, [total_tokens, head_dim].
• ​gate (TileTensor[dtype, GateLayoutType, gate_origin, Engine=GateEngine]): Per-token gate scores, [total_tokens, head_dim].
• ​ape (TileTensor[.float32, ApeLayoutType, ape_origin, Engine=ApeEngine]): Within-pool position embedding, [kpool, head_dim], f32.
• ​input_row_offsets (TileTensor[.uint32, IROLayoutType, iro_origin, Engine=IROEngine]): Token row offsets per request, [batch_size + 1].
• ​pool_row_offsets (TileTensor[.uint32, PROLayoutType, pro_origin, Engine=PROEngine]): Output row offsets per request, [batch_size + 1]. Request b owns output rows [pool_row_offsets[b], pool_row_offsets[b + 1]), one per pool built here.
• ​cache_lengths (TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin, Engine=CacheLenEngine]): Cached-prefix length per request, [batch_size]. A pool covers absolute positions, so this is what places the call's tokens on the pool grid.