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 struct
QSASparseAttentionRaggedPaged
struct QSASparseAttentionRaggedPaged
Gather-GQA over the token indices a QSA selection names.
Tensor shapes: - output : [total_q, num_q_heads, head_dim] (OUT) - q : [total_q, num_q_heads, head_dim] - input_row_offsets : [batch + 1] (uint32) - kv_blocks : [num_pages, 2, num_layers, page_size, num_kv_heads, head_dim] (mut in) - cache_lengths : [batch] (uint32) - kv_lookup_table : [batch, max_pages] (uint32) - max_prompt_length : [1] (uint32) - max_cache_length : [1] (uint32) - token_indices : [total_q, selection_width] (int32) - counts : [total_q] (int32)
token_indices holds positions within each query's own sequence, valid
entries first; entries at or past counts are never read.
Implemented traits
Methods
execute
static def execute[dtype: DType, out_dtype: DType, target: StringSpan[ImmStaticOrigin], group: Int, warps: Int, unroll: Int](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], q: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=q.static_spec], input_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec], kv_blocks: ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=kv_blocks.static_spec], page_stride: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=page_stride.static_spec], cache_lengths: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_lengths.static_spec], kv_lookup_table: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=kv_lookup_table.static_spec], max_prompt_length: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_prompt_length.static_spec], max_cache_length: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_cache_length.static_spec], layer_idx: UInt32, token_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=token_indices.static_spec], counts: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=counts.static_spec], scale: Float32, ctx: DeviceContext)
Runs qsa_sparse_attention for one layer.
Parameters:
- dtype (
DType): Element dtype ofqand the KV cache (inferred). - out_dtype (
DType): Element dtype ofoutput(inferred). - target (
StringSpan[ImmStaticOrigin]): Compilation target. - group (
Int): Query heads per kv head. - warps (
Int): Warps per CTA; must satisfy the kernel's divisibility asserts. - unroll (
Int): Gathered keys loaded before any is consumed.
Args:
- output (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec]): Attention output. - q (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=q.static_spec]): Queries. - input_row_offsets (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec]): Ragged query offsets. - kv_blocks (
ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=kv_blocks.static_spec]): Paged KV blocks. - page_stride (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=page_stride.static_spec]): Page-to-page distance in elements; -1 when packed. - cache_lengths (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_lengths.static_spec]): Per-sequence cached-key count. - kv_lookup_table (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=kv_lookup_table.static_spec]): Per-sequence page table. - max_prompt_length (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_prompt_length.static_spec]): Max new query tokens this step. - max_cache_length (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_cache_length.static_spec]): Max cached context this step. - layer_idx (
UInt32): Layer index into the KV cache. - token_indices (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=token_indices.static_spec]): Selected positions per query row. - counts (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=counts.static_spec]): Valid entries per query row. - scale (
Float32): QK scale. - ctx (
DeviceContext): Device context.
Raises:
Error: If the operand shapes disagree with the parameters.