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_depth512_dispatch

def mha_sm100_depth512_dispatch[q_type: DType, KVType: MHAOperand, MaskType: MHAMask, output_type: DType, MaxPromptLenType: OptionallyStaticInt, PartitionType: MHAPartitionScheme, KVRowOffsetsEngine: TensorEngine, SinkEngine: TensorEngine, //, config: MHAConfig[config.dtype], group: Int, ragged: Bool, _is_cache_length_accurate: Bool, sink: Bool = False](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, partition: PartitionType, ctx: DeviceContext, sink_weights: OptionalReg[TileTensor[q_type, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=SinkEngine]] = None)

Dispatches the pair-CTA SM100 depth=256/512 MHA prefill kernel.

Builds the Depth512SM100Config, constructs TMA tile descriptors for the Q, K, V, and output operands, creates a TransientScheduler with pair_cta=True, and enqueues the SM100MHADepth512 kernel with cluster_dim=(2, 1, 1) so both CTAs in a cluster derive the same tile index from block_idx.x >> 1. Only prefill is supported; decoding is rejected at compile time.

Parameters:

• ​q_type (DType): The query tensor element type (inferred).
• ​KVType (MHAOperand): The paged KV cache operand type providing the tile factory and page size (inferred).
• ​MaskType (MHAMask): The attention mask type applied to the scores (inferred).
• ​output_type (DType): The output buffer element type (inferred).
• ​MaxPromptLenType (OptionallyStaticInt): The maximum prompt length as a static or runtime value (inferred).
• ​PartitionType (MHAPartitionScheme): The KV cache partition scheme (inferred).
• ​KVRowOffsetsEngine (TensorEngine): TensorEngine policy of kv_input_row_offsets (inferred).
• ​SinkEngine (TensorEngine): TensorEngine policy of sink_weights (inferred).
• ​config (MHAConfig[config.dtype]): The MHA configuration with head count, depth, and swizzle mode used to build the Depth512SM100Config.
• ​group (Int): Number of query heads per KV head for grouped-query attention.
• ​ragged (Bool): Whether the batch uses ragged sequence lengths with a non-null valid_length pointer.
• ​_is_cache_length_accurate (Bool): Whether the per-batch cache length values are exact.
• ​sink (Bool): Whether to apply per-head attention-sink weights.

Args:

• ​output (DeviceBuffer[output_type]): Device buffer that receives the attention output.
• ​q_arg (Pointer[Scalar[q_type]]): Pointer to the query tensor data.
• ​k (KVType): Key operand providing the paged KV cache tile factory.
• ​v (KVType): Value operand providing the paged KV cache tile factory.
• ​num_rows_q (Int): Number of query rows to process.
• ​mask (MaskType): Causal or padding mask applied to the attention scores.
• ​valid_length (Pointer[UInt32]): Per-batch pointer to valid cache lengths (used when ragged is true).
• ​max_prompt_len_arg (MaxPromptLenType): Maximum prompt length, static or runtime.
• ​max_cache_valid_length_arg (Int): Maximum valid cache length across the batch.
• ​scale (Float32): Scaling factor applied to the QK dot product.
• ​kv_input_row_offsets (OptionalReg[TileTensor[.uint32, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=KVRowOffsetsEngine]]): Optional ragged row-offset tensor for KV input rows.
• ​batch_size_arg (Int): Number of sequences in the batch.
• ​partition (PartitionType): Partition scheme for the KV cache.
• ​ctx (DeviceContext): Device context used to create TMA descriptors and enqueue the kernel.
• ​sink_weights (OptionalReg[TileTensor[q_type, Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ImmutAnyOrigin, Engine=SinkEngine]]): Optional per-head attention-sink weights, used when sink is set.