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, head_dim: Int, kpool: Int](pooled: TileTensor[dtype, OutLayoutType, out_origin], k: TileTensor[dtype, KLayoutType, k_origin], gate: TileTensor[dtype, GateLayoutType, gate_origin], ape: TileTensor[.float32, ApeLayoutType, ape_origin], input_row_offsets: TileTensor[.uint32, IROLayoutType, iro_origin], pool_row_offsets: TileTensor[.uint32, PROLayoutType, pro_origin], cache_lengths: TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin])

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.
• ​head_dim (Int): Channels per key; also the block width.
• ​kpool (Int): Tokens per pool.

Args:

• ​pooled (TileTensor[dtype, OutLayoutType, out_origin]): Output [total_pools, head_dim], where total_pools is the last entry of pool_row_offsets.
• ​k (TileTensor[dtype, KLayoutType, k_origin]): Layer-normed indexer keys, [total_tokens, head_dim].
• ​gate (TileTensor[dtype, GateLayoutType, gate_origin]): Per-token gate scores, [total_tokens, head_dim].
• ​ape (TileTensor[.float32, ApeLayoutType, ape_origin]): Within-pool position embedding, [kpool, head_dim], f32.
• ​input_row_offsets (TileTensor[.uint32, IROLayoutType, iro_origin]): Token row offsets per request, [batch_size + 1].
• ​pool_row_offsets (TileTensor[.uint32, PROLayoutType, pro_origin]): 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]): 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.