For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

mha_sm90_dispatch

def mha_sm90_dispatch[q_type: DType, KVType: MHAOperand, MaskType: MHAMask, output_type: DType, MaxPromptLenType: OptionallyStaticInt, PartitionType: MHAPartitionScheme, //, config: MHAConfig[config.dtype], group: Int, ragged: Bool, sink: Bool, _is_cache_length_accurate: Bool](output: DeviceBuffer[output_type], q_arg: DeviceBuffer[q_type], k: KVType, v: KVType, num_rows_q: Int, mask_functor: MaskType, valid_length: DeviceBuffer[DType.uint32], max_prompt_len_arg: MaxPromptLenType, max_cache_valid_length_arg: Int, scale: Float32, kv_input_row_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]], batch_size_arg: Int, partition: PartitionType, ctx: DeviceContext, sink_weights: OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]])

Dispatches the SM90 FlashAttention-3 MHA kernel for a single batch.

Selects a transient, tiled, or queued tile scheduler based on the persistent-kernel configuration, builds the Q/K/V TMA tile descriptors, and forwards the request down the dispatch chain to the enqueued kernel.

Parameters:

q_type (DType): The dtype of the query tensor (inferred).
KVType (MHAOperand): The K/V operand type encoding dtype, page size, and layout (inferred).
MaskType (MHAMask): The mask functor type applied to attention scores (inferred).
output_type (DType): The dtype of the output tensor (inferred).
MaxPromptLenType (OptionallyStaticInt): The maximum prompt length, possibly known at compile time (inferred).
PartitionType (MHAPartitionScheme): The scheme for partitioning attention work across SMs (inferred).
config (MHAConfig[config.dtype]): The MHA configuration holding block sizes, head count, depth, and algorithm.
group (Int): The query grouping factor, the number of query heads per KV head.
ragged (Bool): Whether per-row valid lengths vary and require ragged masking.
sink (Bool): Whether attention sink weights are applied.
_is_cache_length_accurate (Bool): Whether the supplied cache length is exact and needs no clamping.

Args:

output (DeviceBuffer[output_type]): The device buffer that receives the attention output.
q_arg (DeviceBuffer[q_type]): The device buffer holding the query tensor.
k (KVType): The key operand lookup table.
v (KVType): The value operand lookup table.
num_rows_q (Int): The number of query rows in the batch.
mask_functor (MaskType): The mask functor instance applied to attention scores.
valid_length (DeviceBuffer[DType.uint32]): The per-batch valid sequence length buffer.
max_prompt_len_arg (MaxPromptLenType): The maximum prompt length after padding.
max_cache_valid_length_arg (Int): The maximum valid length of the KV cache.
scale (Float32): The scaling factor applied to the QK^T product.
kv_input_row_offsets (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional row offsets into the KV cache for ragged layouts.
batch_size_arg (Int): The number of sequences in the batch.
partition (PartitionType): The partition instance describing the work split.
ctx (DeviceContext): The device context used to enqueue the kernel.
sink_weights (OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink weights tensor applied when sink is enabled.