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_seed_tail_kernel

def kpool_seed_tail_kernel[dtype: DType, TailLayoutType: TensorLayout, tail_origin: MutOrigin, KLayoutType: TensorLayout, k_origin: ImmOrigin, GateLayoutType: TensorLayout, gate_origin: ImmOrigin, IROLayoutType: TensorLayout, iro_origin: ImmOrigin, CacheLenLayoutType: TensorLayout, SlotLayoutType: TensorLayout, slot_origin: ImmOrigin, head_dim: Int, kpool: Int](tail: TileTensor[dtype, TailLayoutType, tail_origin], k: TileTensor[dtype, KLayoutType, k_origin], gate: TileTensor[dtype, GateLayoutType, gate_origin], input_row_offsets: TileTensor[.uint32, IROLayoutType, iro_origin], cache_lengths: TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin], slot_idx: TileTensor[.uint32, SlotLayoutType, slot_origin], num_requests: Int32)

Stashes a prefill chunk's trailing tokens into the tail ring.

Compression writes only whole pools, so the tokens after the last complete pool have nowhere else to go. They are the members of the request's in-progress pool, and decode reads them back from the ring.

Only the trailing tokens this call owns are written, so a chunked prefill lands where a single one does.

Parameters:

• ​dtype (DType): Element type of tail, k and gate.
• ​TailLayoutType (TensorLayout): Layout of tail.
• ​tail_origin (MutOrigin): 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.
• ​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.
• ​head_dim (Int): Channels per key; also the block width.
• ​kpool (Int): Tokens per pool.

Args:

• ​tail (TileTensor[dtype, TailLayoutType, tail_origin]): Per-slot ring, [max_slots, 2, kpool, head_dim]. Index 0 holds keys, index 1 holds gate scores.
• ​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].
• ​input_row_offsets (TileTensor[.uint32, IROLayoutType, iro_origin]): Token row offsets per request, [batch_size + 1].
• ​cache_lengths (TileTensor[.uint32, CacheLenLayoutType, ImmutAnyOrigin]): Cached-prefix length per request, [batch_size].
• ​slot_idx (TileTensor[.uint32, SlotLayoutType, slot_origin]): Ring slot owned by each batch row, [num_requests], uint32.
• ​num_requests (Int32): Requests actually present.