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_tail_update_kernel

def kpool_tail_update_kernel[dtype: DType, TailLayoutType: TensorLayout, tail_origin: MutOrigin, OutLayoutType: TensorLayout, out_origin: MutOrigin, ClosedLayoutType: TensorLayout, closed_origin: MutOrigin, KLayoutType: TensorLayout, k_origin: ImmOrigin, GateLayoutType: TensorLayout, gate_origin: ImmOrigin, ApeLayoutType: TensorLayout, ape_origin: ImmOrigin, PosLayoutType: TensorLayout, pos_origin: ImmOrigin, SlotLayoutType: TensorLayout, slot_origin: ImmOrigin, TailEngine: TensorEngine, OutEngine: TensorEngine, ClosedEngine: TensorEngine, KEngine: TensorEngine, GateEngine: TensorEngine, ApeEngine: TensorEngine, PosEngine: TensorEngine, SlotEngine: TensorEngine, head_dim: Int, kpool: Int, next_n: Int = Int(1)](tail: TileTensor[dtype, TailLayoutType, tail_origin, Engine=TailEngine], pooled: TileTensor[dtype, OutLayoutType, out_origin, Engine=OutEngine], closed_pool: TileTensor[.int32, ClosedLayoutType, closed_origin, Engine=ClosedEngine], k: TileTensor[dtype, KLayoutType, k_origin, Engine=KEngine], gate: TileTensor[dtype, GateLayoutType, gate_origin, Engine=GateEngine], ape: TileTensor[.float32, ApeLayoutType, ape_origin, Engine=ApeEngine], positions: TileTensor[.int32, PosLayoutType, pos_origin, Engine=PosEngine], slot_idx: TileTensor[.uint32, SlotLayoutType, slot_origin, Engine=SlotEngine], num_requests: Int32)

Stashes a request's new tokens, and closes each pool as it fills.

A decoded token cannot be pooled on arrival, because its pool-mates arrived on earlier steps and have left the batch. Each request keeps its in-progress pool in tail, a ring of kpool slots addressed by position % kpool.

A speculative step appends next_n tokens at once, so several pools can close in one call.

The ring is indexed by slot_idx[r], not by r. A batch reorders between steps, so row r is not always the same request.

Every real token stashes, whether or not it closes a pool.

Rejected speculative tokens are the caller's problem. The ring holds no pointer to rewind, so a rejected token that has already stashed stays.

Parameters:

• ​dtype (DType): Element type of tail, k, gate and pooled.
• ​TailLayoutType (TensorLayout): Layout of tail.
• ​tail_origin (MutOrigin): Origin of tail.
• ​OutLayoutType (TensorLayout): Layout of pooled.
• ​out_origin (MutOrigin): Origin of pooled.
• ​ClosedLayoutType (TensorLayout): Layout of closed_pool.
• ​closed_origin (MutOrigin): Origin of closed_pool.
• ​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.
• ​PosLayoutType (TensorLayout): Layout of positions.
• ​pos_origin (ImmOrigin): Origin of positions.
• ​SlotLayoutType (TensorLayout): Layout of slot_idx.
• ​slot_origin (ImmOrigin): Origin of slot_idx.
• ​TailEngine (TensorEngine): Engine of tail.
• ​OutEngine (TensorEngine): Engine of pooled.
• ​ClosedEngine (TensorEngine): Engine of closed_pool.
• ​KEngine (TensorEngine): Engine of k.
• ​GateEngine (TensorEngine): Engine of gate.
• ​ApeEngine (TensorEngine): Engine of ape.
• ​PosEngine (TensorEngine): Engine of positions.
• ​SlotEngine (TensorEngine): Engine of slot_idx.
• ​head_dim (Int): Channels per key; also the block width.
• ​kpool (Int): Tokens per pool.
• ​next_n (Int): Tokens appended per request per call.

Args:

• ​tail (TileTensor[dtype, TailLayoutType, tail_origin, Engine=TailEngine]): Per-slot ring, [max_slots, 2, kpool, head_dim], sized by the engine's concurrent-request capacity. Index 0 holds keys, index 1 holds gate scores. Only slot slot_idx[r] is touched for row r. Persists across steps.
• ​pooled (TileTensor[dtype, OutLayoutType, out_origin, Engine=OutEngine]): Output [num_requests, ceil(next_n / kpool), head_dim], meaningful only where closed_pool is non-negative.
• ​closed_pool (TileTensor[.int32, ClosedLayoutType, closed_origin, Engine=ClosedEngine]): Output [num_requests, ceil(next_n / kpool)]. The pool ids this call completed, in order, padded with -1.
• ​k (TileTensor[dtype, KLayoutType, k_origin, Engine=KEngine]): This step's layer-normed keys, [num_requests, next_n, head_dim].
• ​gate (TileTensor[dtype, GateLayoutType, gate_origin, Engine=GateEngine]): This step's gate scores, [num_requests, next_n, head_dim].
• ​ape (TileTensor[.float32, ApeLayoutType, ape_origin, Engine=ApeEngine]): Within-pool position embedding, [kpool, head_dim], f32.
• ​positions (TileTensor[.int32, PosLayoutType, pos_origin, Engine=PosEngine]): Absolute position of each new token, [num_requests, next_n]. Negative marks a padded entry.
• ​slot_idx (TileTensor[.uint32, SlotLayoutType, slot_origin, Engine=SlotEngine]): Ring slot owned by each batch row, [num_requests], uint32.
• ​num_requests (Int32): Requests actually present.