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_sm100_dispatch

def mha_sm100_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: MaskType, valid_length: DeviceBuffer[DType.uint32], max_prompt_len_arg: MaxPromptLenType, max_cache_valid_length_arg: Int, scale: Float32, kv_input_row_offsets: OptionalReg[TileTensor[DType.uint32, Layout[*?, *?], ImmutAnyOrigin]], batch_size_arg: Int, partition: PartitionType, ctx: DeviceContext, sink_weights: OptionalReg[TileTensor[q_type, Layout[*?, *?], ImmutAnyOrigin]])

Dispatches single-query (decode) multi-head attention to the SM100 kernel, selecting TMA tiles, scheduler, and partition configuration.

Parameters:

q_type (DType): The element type of the query tensor.
KVType (MHAOperand): The operand type describing the key and value memory layout (for example, paged or ragged).
MaskType (MHAMask): The mask type applied to the attention scores.
output_type (DType): The element type of the output tensor.
MaxPromptLenType (OptionallyStaticInt): The type representing the maximum prompt length, which may be statically or dynamically known.
PartitionType (MHAPartitionScheme): The scheme for partitioning work across CTAs.
config (MHAConfig[config.dtype]): The MHA configuration holding tile sizes, depth, and head count.
group (Int): The group-query attention group size, equal to the number of query heads sharing each KV head.
ragged (Bool): Whether the input sequences have variable lengths.
sink (Bool): Whether attention sink weights are applied.
_is_cache_length_accurate (Bool): Whether the reported cache length is exact, affecting position computation.

Args:

output (DeviceBuffer[output_type]): The output buffer for the attention results.
q_arg (DeviceBuffer[q_type]): The query tensor buffer.
k (KVType): The key operand, with layout described by KVType.
v (KVType): The value operand, with layout described by KVType.
num_rows_q (Int): Number of query rows in the batch.
mask (MaskType): The attention mask instance.
valid_length (DeviceBuffer[DType.uint32]): The per-sequence valid length buffer.
max_prompt_len_arg (MaxPromptLenType): The maximum prompt length across the batch.
max_cache_valid_length_arg (Int): The maximum valid cache length.
scale (Float32): The scaling factor applied to the query-key dot product.
kv_input_row_offsets (OptionalReg[TileTensor[DType.uint32, Layout[*?, *?], ImmutAnyOrigin]]): Optional row offsets into the KV input. for ragged layouts.
batch_size_arg (Int): The number of sequences in the batch.
partition (PartitionType): The partition descriptor for splitting work across CTAs.
ctx (DeviceContext): The device context used to launch the kernel.
sink_weights (OptionalReg[TileTensor[q_type, Layout[*?, *?], ImmutAnyOrigin]]): Optional sink weights for attention sinks.