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_ring_close_kernel

def kpool_ring_close_kernel[dtype: DType, OutLayoutType: TensorLayout, out_origin: MutOrigin, ClosedLayoutType: TensorLayout, closed_origin: MutOrigin, TailLayoutType: TensorLayout, tail_origin: ImmOrigin, KLayoutType: TensorLayout, k_origin: ImmOrigin, GateLayoutType: TensorLayout, gate_origin: ImmOrigin, ApeLayoutType: TensorLayout, ape_origin: ImmOrigin, IROLayoutType: TensorLayout, iro_origin: ImmOrigin, CacheLenLayoutType: TensorLayout, SlotLayoutType: TensorLayout, slot_origin: ImmOrigin, OutEngine: TensorEngine, ClosedEngine: TensorEngine, TailEngine: TensorEngine, KEngine: TensorEngine, GateEngine: TensorEngine, ApeEngine: TensorEngine, IROEngine: TensorEngine, CacheLenEngine: TensorEngine, SlotEngine: TensorEngine, head_dim: Int, kpool: Int](pooled: TileTensor[dtype, OutLayoutType, out_origin, Engine=OutEngine], closed_pool: TileTensor[.int32, ClosedLayoutType, closed_origin, Engine=ClosedEngine], tail: TileTensor[dtype, TailLayoutType, tail_origin, Engine=TailEngine], 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], cache_lengths: TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin, Engine=CacheLenEngine], slot_idx: TileTensor[.uint32, SlotLayoutType, slot_origin, Engine=SlotEngine], num_requests: Int32)

Closes each request's pending tail-ring pool, ragged and unconditional.

A cached prefix ending mid-pool leaves that pool's earlier members in tail (see kpool_seed_tail_kernel), and neither kpool_compress_kernel nor kpool_seed_tail_kernel reads them back -- kpool_compress_kernel only ever builds pools entirely from this call's own tokens, skipping the leading remainder the ring already holds. This kernel is the one that reads that remainder, so a request whose new tokens complete it is not stuck forever.

Unlike kpool_tail_update_kernel, this takes a ragged, arbitrary-width chunk per request (input_row_offsets, not a compile-time next_n), and closes at most one pool per request per call -- the ring never holds more than kpool - 1 members between calls, so there is never a second one to close. Any further pools this call completes are entirely new tokens and belong to kpool_compress_kernel, which already derives its own pool count and starting row by skipping exactly the tokens this kernel consumes (align below, computed the same way in both).

Parameters:

• ​dtype (DType): Element type of tail, k, gate and pooled.
• ​OutLayoutType (TensorLayout): Layout of pooled.
• ​out_origin (MutOrigin): Origin of pooled.
• ​ClosedLayoutType (TensorLayout): Layout of closed_pool.
• ​closed_origin (MutOrigin): Origin of closed_pool.
• ​TailLayoutType (TensorLayout): Layout of tail.
• ​tail_origin (ImmOrigin): Origin of tail.
• ​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.
• ​CacheLenLayoutType (TensorLayout): Layout of cache_lengths.
• ​SlotLayoutType (TensorLayout): Layout of slot_idx.
• ​slot_origin (ImmOrigin): Origin of slot_idx.
• ​OutEngine (TensorEngine): Engine of pooled.
• ​ClosedEngine (TensorEngine): Engine of closed_pool.
• ​TailEngine (TensorEngine): Engine of tail.
• ​KEngine (TensorEngine): Engine of k.
• ​GateEngine (TensorEngine): Engine of gate.
• ​ApeEngine (TensorEngine): Engine of ape.
• ​IROEngine (TensorEngine): Engine of input_row_offsets.
• ​CacheLenEngine (TensorEngine): Engine of cache_lengths.
• ​SlotEngine (TensorEngine): Engine of slot_idx.
• ​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 [batch, head_dim], meaningful only where closed_pool is non-negative.
• ​closed_pool (TileTensor[.int32, ClosedLayoutType, closed_origin, Engine=ClosedEngine]): Output [batch]. The pool id this call closed for that request, or -1.
• ​tail (TileTensor[dtype, TailLayoutType, tail_origin, Engine=TailEngine]): Per-slot ring, [max_slots, 2, kpool, head_dim]. Index 0 holds keys, index 1 holds gate scores. Read only -- slot_idx[r]'s entry is only ever written by kpool_seed_tail_kernel.
• ​k (TileTensor[dtype, KLayoutType, k_origin, Engine=KEngine]): This call's layer-normed keys, [total_tokens, head_dim].
• ​gate (TileTensor[dtype, GateLayoutType, gate_origin, Engine=GateEngine]): This call's 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 + 1].
• ​cache_lengths (TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin, Engine=CacheLenEngine]): Cached-prefix length per request, [batch]. A pool covers absolute positions, so this is what tells each request whether it has a pending pool and how much of it the ring holds.
• ​slot_idx (TileTensor[.uint32, SlotLayoutType, slot_origin, Engine=SlotEngine]): Ring slot owned by each batch row, [batch], uint32.
• ​num_requests (Int32): Requests actually present.