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

latent_sparse_attention_ragged_paged

def latent_sparse_attention_ragged_paged[swa_t: KVCacheT, comp_t: KVCacheT, q_type: DType, out_type: DType, //, target: StringSpan[ImmStaticOrigin], window: Int](output: TileTensor[out_type, Engine=output.Engine, linear_idx_type=output.linear_idx_type], q: TileTensor[q_type, Engine=q.Engine, linear_idx_type=q.linear_idx_type], input_row_offsets: TileTensor[.uint32, Engine=input_row_offsets.Engine, linear_idx_type=input_row_offsets.linear_idx_type], comp_indices: TileTensor[.int32, Engine=comp_indices.Engine, linear_idx_type=comp_indices.linear_idx_type], attn_sink: TileTensor[.float32, Engine=attn_sink.Engine, linear_idx_type=attn_sink.linear_idx_type], swa_cache: swa_t, comp_cache: comp_t, scale: Float32, ctx: DeviceContext)

Attends every query row to its window keys and its listed compressed entries.

Parameters:

• ​swa_t (KVCacheT): The sliding-window leaf's cache type (inferred); paged by token position, one latent head.
• ​comp_t (KVCacheT): The compressed leaf's cache type (inferred); paged by entry, one latent head with the same head size.
• ​q_type (DType): Query element type (inferred).
• ​out_type (DType): Output element type (inferred).
• ​target (StringSpan[ImmStaticOrigin]): Compilation target string, selects the CPU or GPU path.
• ​window (Int): Sliding window length in tokens; a query at pos sees positions max(0, pos - window + 1) .. pos.

Args:

• ​output (TileTensor[out_type, Engine=output.Engine, linear_idx_type=output.linear_idx_type]): [num_rows, num_heads, head_dim].
• ​q (TileTensor[q_type, Engine=q.Engine, linear_idx_type=q.linear_idx_type]): [num_rows, num_heads, head_dim], the last axis contiguous.
• ​input_row_offsets (TileTensor[.uint32, Engine=input_row_offsets.Engine, linear_idx_type=input_row_offsets.linear_idx_type]): [batch + 1] ragged row offsets of q.
• ​comp_indices (TileTensor[.int32, Engine=comp_indices.Engine, linear_idx_type=comp_indices.linear_idx_type]): [num_rows, num_comp] entry indices into the compressed leaf; -1 marks an unused slot.
• ​attn_sink (TileTensor[.float32, Engine=attn_sink.Engine, linear_idx_type=attn_sink.linear_idx_type]): [num_heads] per-head sink logits.
• ​swa_cache (swa_t): This layer's sliding-window leaf.
• ​comp_cache (comp_t): This layer's compressed leaf.
• ​scale (Float32): Softmax scale applied to the scores.
• ​ctx (DeviceContext): Device context used to enqueue the GPU kernel.