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

AnyType, Deinitable, Movable

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], 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 of q and the KV cache (inferred).
• ​out_dtype (DType): Element dtype of output (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.
• ​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.