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

mha_single_batch_pipelined

def mha_single_batch_pipelined[q_type: DType, k_t: MHAOperand, v_t: MHAOperand, output_type: DType, mask_t: MHAMask, *, config: MHAConfig[config.dtype], group: Int = Int(1), sink: Bool = False](q_ptr: Pointer[Scalar[q_type], ImmutAnyOrigin], k: k_t, v: v_t, output_ptr: Pointer[Scalar[output_type], MutAnyOrigin], scale: Float32, seq_len: Int, max_seq_len: Int, start_pos: UInt32, num_keys: Int, mask_tensor_col: Int, mask: mask_t, batch_idx: Int, sink_weights: OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]])

MHA for token gen where seqlen = 1 and num_keys >= 1.

The general data layout and steps conform to flash attention. Two exceptions:

1 Partition across B, H, and num_keys (TODO). The last one is split-K and will need a separate reduction kernel at the end.

2 First bmm becomes gemv and second bmm becomes gevm. TODO: use more optimized kernels for them

Parameters:

• ​q_type (DType): Element type of the query tensor.
• ​k_t (MHAOperand): Key operand type backing the key tensor (KV cache or dense).
• ​v_t (MHAOperand): Value operand type backing the value tensor (KV cache or dense).
• ​output_type (DType): Element type of the output tensor.
• ​mask_t (MHAMask): Attention mask type implementing MHAMask.
• ​config (MHAConfig[config.dtype]): Tile and pipeline configuration for the kernel.
• ​group (Int): GQA group size, the ratio of query heads to key/value heads (defaults to 1).
• ​sink (Bool): True to enable attention-sink mode where the first tokens always attend (defaults to False).

Args:

• ​q_ptr (Pointer[Scalar[q_type], ImmutAnyOrigin]): Pointer to the query tensor in global memory.
• ​k (k_t): Key operand backed by a KV cache or dense tensor.
• ​v (v_t): Value operand backed by a KV cache or dense tensor.
• ​output_ptr (Pointer[Scalar[output_type], MutAnyOrigin]): Pointer to the output tensor in global memory.
• ​scale (Float32): Softmax temperature scale applied to Q·Kᵀ.
• ​seq_len (Int): Valid query sequence length excluding padding.
• ​max_seq_len (Int): Padded query sequence length.
• ​start_pos (UInt32): Starting position of the current query in the KV cache, used for mask row indexing.
• ​num_keys (Int): Number of key entries in the KV cache.
• ​mask_tensor_col (Int): Second dimension of the mask tensor.
• ​mask (mask_t): Mask instance used to apply the attention mask.
• ​batch_idx (Int): Index of the sequence within the batch.
• ​sink_weights (OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink-token weight tensor for attention sinks; required when sink is True.