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_sm100_dispatch

def mha_sm100_dispatch[q_type: DType, KVType: MHAOperand, MaskType: MHAMask, output_type: DType, MaxPromptLenType: OptionallyStaticInt, KVRowOffsetsEngine: TensorEngine, SinkEngine: TensorEngine, //, config: MHAConfig[config.dtype], group: Int, ragged: Bool, sink: Bool, _is_cache_length_accurate: Bool](output: DeviceBuffer[output_type], q_arg: Pointer[Scalar[q_type]], k: KVType, v: KVType, num_rows_q: Int, mask: MaskType, valid_length: Pointer[UInt32], max_prompt_len_arg: MaxPromptLenType, max_cache_valid_length_arg: Int, scale: Float32, kv_input_row_offsets: OptionalReg[TileTensor[.uint32, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=KVRowOffsetsEngine]], batch_size_arg: Int, ctx: DeviceContext, sink_weights: OptionalReg[TileTensor[q_type, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=SinkEngine]], num_partitions_override: Int = Int(0))

Dispatches the SM100 FA4 flash-attention kernel for a prefill or decode workload.

Selects between the 1Q split-K and 2Q FA4 configurations based on a occupancy and prompt-length heuristic, constructs the Q/K/V/O TMA tile descriptors and transient tile scheduler, threads optional ragged valid-length, KV-row-offset, and sink-attention arguments through to the compiled kernel, and enqueues the launch onto the supplied device context.

Decode is not a separate route: a single-token prompt is just the shortest prefill, so it takes the same ladder and lands on the smallest tile that single-tiles it. max_prompt_len_arg must be dynamic even when it is 1 -- see the note on the missing _is_decoding guard in the body.

Parameters:

• ​q_type (DType): Element type of the query tensor (inferred).
• ​KVType (MHAOperand): Key/value operand descriptor with dtype and page size (inferred).
• ​MaskType (MHAMask): Attention mask scheme applied to the Q@K' scores (inferred).
• ​output_type (DType): Element type of the attention output buffer (inferred).
• ​MaxPromptLenType (OptionallyStaticInt): Optionally-static type encoding the maximum prompt length (inferred).
• ​KVRowOffsetsEngine (TensorEngine): TensorEngine policy of kv_input_row_offsets (inferred).
• ​SinkEngine (TensorEngine): TensorEngine policy of sink_weights (inferred).
• ​config (MHAConfig[config.dtype]): MHA configuration supplying dtype, head count, depth, and swizzle mode.
• ​group (Int): Number of query heads per KV head (GQA group size).
• ​ragged (Bool): Whether to dispatch the variable-length valid-length path.
• ​sink (Bool): Whether to thread sink-attention weights into the kernel.
• ​_is_cache_length_accurate (Bool): Whether the supplied cache length is accurate, threaded to the compiled kernel.

Args:

• ​output (DeviceBuffer[output_type]): Device buffer that receives the attention output rows.
• ​q_arg (Pointer[Scalar[q_type]]): Pointer to the query tensor data.
• ​k (KVType): Key operand descriptor.
• ​v (KVType): Value operand descriptor.
• ​num_rows_q (Int): Number of query rows to attend over.
• ​mask (MaskType): Attention mask applied to the Q@K' scores.
• ​valid_length (Pointer[UInt32]): Per-row valid KV length pointer, used when ragged is set.
• ​max_prompt_len_arg (MaxPromptLenType): Maximum prompt length, optionally static.
• ​max_cache_valid_length_arg (Int): Maximum valid KV cache length across the batch.
• ​scale (Float32): Scalar applied to the Q@K' product before softmax.
• ​kv_input_row_offsets (OptionalReg[TileTensor[.uint32, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=KVRowOffsetsEngine]]): Optional per-row KV input offsets for ragged layouts.
• ​batch_size_arg (Int): Number of sequences in the batch.
• ​ctx (DeviceContext): Device context used to build TMA descriptors and enqueue the kernel.
• ​sink_weights (OptionalReg[TileTensor[q_type, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=SinkEngine]]): Optional sink-attention weights used when sink is set.
• ​num_partitions_override (Int): Exact split-K partition count to use, or 0 for the automatic ladder. A non-zero value is honored verbatim (cluster/DSMEM when it is a CLUSTER_SPLITK_CANDIDATES member, else the workspace route, which admits any P) and additionally forces a BM < 256 route, since 2Q cannot split at all. Overriding bypasses the capture-invariant ws_p_ceiling ceiling and the measured Layout-E workspace guards.