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

flash_attention_dispatch

def flash_attention_dispatch[k_t: MHAOperand, v_t: MHAOperand, mask_t: MHAMask, dtype: DType, q_layout: Layout, //, kv_num_heads: Int, config: MHAConfig[dtype] = MHAConfig(SIMD[IntTuple](q_layout.shape[Int((add q_layout.rank(), -2))]), SIMD[IntTuple](q_layout.shape[Int((add q_layout.rank(), -1))]), Optional(None), Optional(None), Optional(None), Optional(None), Optional(None), Int(4), Int(1), FlashAttentionAlgorithm(Int(-1)), TensorMapSwizzle.SWIZZLE_128B), ragged: Bool = False, sink: Bool = False, _is_flash_attention_applicable: Bool = True, _is_cache_length_accurate: Bool = False, _use_valid_length: Bool = True, _padded_ndbuffer: Bool = False, decoding_warp_split_k: Bool = False](output: LayoutTensor[element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], q: LayoutTensor[dtype, q_layout, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment], k: k_t, v: v_t, mask_functor: mask_t, max_prompt_len: Int, max_cache_valid_length: Int, scale: Float32, is_token_generation: Bool, ctx: DeviceContext, valid_length: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None, kv_input_row_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None, num_partitions: Optional[Int] = None, sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None, decode_dispatch_metadata: OptionalReg[MHADecodeDispatchMetadata] = None)

Dispatch a flash-attention kernel for prefill or decode over dense or KV-cache operands.

Central dispatch point that inspects is_token_generation, dtype, depth, and the target GPU to select among FA2, FA3, and naive implementations. Handles both prefill (context encoding) and incremental decode, routing ragged-batch and paged-KV-cache inputs through appropriate kernel paths.

Parameters:

k_t (MHAOperand): KV-cache or dense operand type for the key tensor.
v_t (MHAOperand): KV-cache or dense operand type for the value tensor.
mask_t (MHAMask): Attention mask type implementing MHAMask.
dtype (DType): Element type of Q (K/V type is inferred from k_t).
q_layout (Layout): Compile-time layout of the query tensor.
kv_num_heads (Int): Number of key/value heads (for GQA).
config (MHAConfig[dtype]): Tile/pipeline configuration; defaults from query shape.
ragged (Bool): True for ragged-batch (variable-length) inputs.
sink (Bool): True to enable attention-sink mode.
_is_flash_attention_applicable (Bool): Internal flag to suppress FA path.
_is_cache_length_accurate (Bool): True when KV cache length already includes the newest tokens.
_use_valid_length (Bool): True to mask output with per-sequence lengths.
_padded_ndbuffer (Bool): True when the NBuffer has padded dense inputs.
decoding_warp_split_k (Bool): Enable warp-level split-K for decode.

Args:

output (LayoutTensor[element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment]): Mutable output tensor.
q (LayoutTensor[dtype, q_layout, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment]): Query tensor.
k (k_t): Key operand (KV cache or dense tensor).
v (v_t): Value operand (KV cache or dense tensor).
mask_functor (mask_t): Mask instance.
max_prompt_len (Int): Maximum query sequence length in the batch.
max_cache_valid_length (Int): Maximum key/value sequence length.
scale (Float32): Softmax temperature scale.
is_token_generation (Bool): True for decode mode; False for prefill.
ctx (DeviceContext): GPU device context.
valid_length (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Per-sequence valid lengths for masked output.
kv_input_row_offsets (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Row offsets for ragged KV inputs.
num_partitions (Optional[Int]): Override split-K partition count; None for auto.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink-token weight tensor.
decode_dispatch_metadata (OptionalReg[MHADecodeDispatchMetadata]): Pre-computed decode dispatch metadata.