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

def flash_attention[dtype: DType, q_layout: Layout, //, config: MHAConfig[dtype] = MHAConfig(SIMD[IntTuple](q_layout.shape[2]), SIMD[IntTuple](q_layout.shape[3]), Optional(None), Optional(None), Optional(None), Optional(None), Optional(None), Int(4), Int(1), FlashAttentionAlgorithm(Int(-1)), TensorMapSwizzle.SWIZZLE_128B), decoding_warp_split_k: Bool = False, naive_kernel: Bool = False, sink: 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: LayoutTensor[element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment], v: LayoutTensor[element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment], mask: LayoutTensor[element_layout=mask.element_layout, layout_int_type=mask.layout_int_type, linear_idx_type=mask.linear_idx_type, masked=mask.masked, alignment=mask.alignment], scale: Float32, context: DeviceContext, num_partitions: Optional[Int] = None, sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

Run flash attention with a dense mask tensor on the current device.

Wraps the mask tensor in a MaterializedMask and delegates to the mask-typed overload. Selects the flash-attention algorithm variant (FA2 / FA3 / naive) based on config.algorithm and the detected GPU.

Parameters:

dtype (DType): Element type shared by Q, K, V, and the output.
q_layout (Layout): Compile-time layout of the query tensor.
config (MHAConfig[dtype]): Tile/pipeline configuration; defaults are derived from the query layout's last two dimensions.
decoding_warp_split_k (Bool): Enable warp-level split-K for decode.
naive_kernel (Bool): Force the fallback naive attention kernel.
sink (Bool): Enable attention-sink mode (first tokens always attend).

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]): Destination tensor for attention output.
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 (LayoutTensor[element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment]): Key tensor.
v (LayoutTensor[element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment]): Value tensor.
mask (LayoutTensor[element_layout=mask.element_layout, layout_int_type=mask.layout_int_type, linear_idx_type=mask.linear_idx_type, masked=mask.masked, alignment=mask.alignment]): Dense attention mask tensor.
scale (Float32): Softmax temperature scale applied to Q·Kᵀ.
context (DeviceContext): GPU device context for kernel dispatch.
num_partitions (Optional[Int]): Override the number of split-K partitions; None selects automatically.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink-token weight tensor for attention sinks.

def flash_attention[cache_t: KVCacheT, mask_t: MHAMask, dtype: DType, q_layout: Layout, //, 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, decoding_warp_split_k: Bool = False, naive_kernel: 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: cache_t, v: cache_t, mask_functor: mask_t, valid_length: LayoutTensor[DType.uint32, element_layout=valid_length.element_layout, layout_int_type=valid_length.layout_int_type, linear_idx_type=valid_length.linear_idx_type, masked=valid_length.masked, alignment=valid_length.alignment], scale: Float32, ctx: DeviceContext, q_max_seq_len: Optional[Int] = 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)

Flash attention 2 algorithm. Compute: (1) Transpose (Q) BSHD -> BHSD; (2) Transpose (K) BSHD -> BHSD; (3) Transpose (V) BSHD -> BHSD; (4) P = Bmm(Q, K), P is also called "score"; (5) P = P * scale + mask; (6) P = softmax(P); (7) O = Bmm(P, V) (8) Output = Transpose(O).

B, S, H, D denote batch size, sequence length, head count and depth, respectively. (1), (2), (3) happens while loading the data into shared memory. (8) happens when writing output to global memory.

All inputs (query, key, and value) must have BSHD layout. The mask can be BSS or BHSS.

This kernel also handles grouped attention optimization. In this case the shape of K and V are BShD where h = H / num_groups.

This kernels handles batches with different valid lengths (i.e., before the padding). Such lengths are passed in valid_length argument.

Parameters:

cache_t (KVCacheT): KV-cache type backing the key and value tensors (inferred).
mask_t (MHAMask): Attention mask type implementing MHAMask (inferred).
dtype (DType): Element type shared by Q, K, V, and the output (inferred).
q_layout (Layout): Compile-time layout of the query tensor (inferred).
config (MHAConfig[dtype]): Tile/pipeline configuration; defaults are derived from the query layout's last two dimensions.
ragged (Bool): True for ragged-batch (variable-length) inputs (defaults to False).
sink (Bool): True to enable attention-sink mode where the first tokens always attend (defaults to False).
decoding_warp_split_k (Bool): True to enable warp-level split-K for decode (defaults to False).
naive_kernel (Bool): True to force the fallback naive attention kernel (defaults to False).

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 destination tensor for the attention output.
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 with BSHD layout.
k (cache_t): Key operand backed by a KV cache.
v (cache_t): Value operand backed by a KV cache.
mask_functor (mask_t): Mask instance used to apply the attention mask.
valid_length (LayoutTensor[DType.uint32, element_layout=valid_length.element_layout, layout_int_type=valid_length.layout_int_type, linear_idx_type=valid_length.linear_idx_type, masked=valid_length.masked, alignment=valid_length.alignment]): Per-sequence valid lengths for masking padded batches.
scale (Float32): Softmax temperature scale applied to Q·Kᵀ.
ctx (DeviceContext): GPU device context for kernel dispatch.
q_max_seq_len (Optional[Int]): Maximum query sequence length in the batch; None infers it from the KV cache.
kv_input_row_offsets (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Row offsets for ragged KV inputs; None for self-attention.
num_partitions (Optional[Int]): Override the number of split-K partitions; None selects automatically.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink-token weight tensor for attention sinks.
decode_dispatch_metadata (OptionalReg[MHADecodeDispatchMetadata]): Pre-computed decode dispatch metadata; None recomputes it.

def flash_attention[mask_t: MHAMask, dtype: DType, q_layout: Layout, //, config: MHAConfig[dtype] = MHAConfig(SIMD[IntTuple](q_layout.shape[2]), SIMD[IntTuple](q_layout.shape[3]), Optional(None), Optional(None), Optional(None), Optional(None), Optional(None), Int(4), Int(1), FlashAttentionAlgorithm(Int(-1)), TensorMapSwizzle.SWIZZLE_128B), decoding_warp_split_k: Bool = False, _use_valid_length: Bool = False, _padded_ndbuffer: Bool = False, naive_kernel: Bool = False, sink: 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: LayoutTensor[element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment], v: LayoutTensor[element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment], mask_functor: mask_t, scale: Float32, ctx: DeviceContext, num_partitions: Optional[Int] = None, valid_length: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None, sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

Run flash attention with dense LayoutTensor K/V operands.

Wraps K and V in LayoutTensorMHAOperand adapters and delegates to flash_attention_dispatch. Handles zero-sized attention (e.g. VAE mid-block on a placeholder image) by returning early.

Parameters:

mask_t (MHAMask): Attention mask type implementing MHAMask (inferred).
dtype (DType): Element type shared by Q, K, V, and the output (inferred).
q_layout (Layout): Compile-time layout of the query tensor (inferred).
config (MHAConfig[dtype]): Tile/pipeline configuration; defaults are derived from the query layout's last two dimensions.
decoding_warp_split_k (Bool): True to enable warp-level split-K for decode (defaults to False).
_use_valid_length (Bool): True to mask output with per-sequence lengths (defaults to False).
_padded_ndbuffer (Bool): True when the NBuffer holds padded dense inputs (defaults to False).
naive_kernel (Bool): True to force the fallback naive attention kernel (defaults to False).
sink (Bool): True to enable attention-sink mode where the first tokens always attend (defaults to False).

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 destination tensor for the attention output.
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 with BSHD layout.
k (LayoutTensor[element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment]): Key tensor with BSHD layout.
v (LayoutTensor[element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment]): Value tensor with BSHD layout.
mask_functor (mask_t): Mask instance used to apply the attention mask.
scale (Float32): Softmax temperature scale applied to Q·Kᵀ.
ctx (DeviceContext): GPU device context for kernel dispatch.
num_partitions (Optional[Int]): Override the number of split-K partitions; None selects automatically.
valid_length (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Per-sequence valid lengths for masking padded batches.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink-token weight tensor for attention sinks.

def flash_attention[mask_t: MHAMask, dtype: DType, output_type: DType, q_tt_layout: TensorLayout, k_tt_layout: TensorLayout, v_tt_layout: TensorLayout, output_tt_layout: TensorLayout, //, config: MHAConfig[dtype] = MHAConfig(q_tt_layout.static_shape[Int(2)], q_tt_layout.static_shape[Int(3)], Optional(None), Optional(None), Optional(None), Optional(None), Optional(None), Int(4), Int(1), FlashAttentionAlgorithm(Int(-1)), TensorMapSwizzle.SWIZZLE_128B), decoding_warp_split_k: Bool = False, _use_valid_length: Bool = False, _padded_ndbuffer: Bool = False, naive_kernel: Bool = False, sink: Bool = False](output: TileTensor[output_type, output_tt_layout, Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[dtype, q_tt_layout, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: TileTensor[dtype, k_tt_layout, Storage=k.Storage, linear_idx_type=k.linear_idx_type], v: TileTensor[dtype, v_tt_layout, Storage=v.Storage, linear_idx_type=v.linear_idx_type], mask_functor: mask_t, scale: Float32, ctx: DeviceContext, num_partitions: Optional[Int] = None, valid_length: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None, sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

TileTensor overload of flash attention.

Converts TileTensor operands to LayoutTensor and delegates to the dense LayoutTensor overload.

Parameters:

mask_t (MHAMask): Attention mask type implementing MHAMask (inferred).
dtype (DType): Element type of Q, K, and V (inferred).
output_type (DType): Element type of the output tensor, which may differ from dtype (inferred).
q_tt_layout (TensorLayout): Compile-time TensorLayout of the query tensor (inferred).
k_tt_layout (TensorLayout): Compile-time TensorLayout of the key tensor (inferred).
v_tt_layout (TensorLayout): Compile-time TensorLayout of the value tensor (inferred).
output_tt_layout (TensorLayout): Compile-time TensorLayout of the output tensor (inferred).
config (MHAConfig[dtype]): Tile/pipeline configuration; defaults are derived from the query layout's last two dimensions.
decoding_warp_split_k (Bool): True to enable warp-level split-K for decode (defaults to False).
_use_valid_length (Bool): True to mask output with per-sequence lengths (defaults to False).
_padded_ndbuffer (Bool): True when the NBuffer holds padded dense inputs (defaults to False).
naive_kernel (Bool): True to force the fallback naive attention kernel (defaults to False).
sink (Bool): True to enable attention-sink mode where the first tokens always attend (defaults to False).

Args:

output (TileTensor[output_type, output_tt_layout, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Mutable destination TileTensor for the attention output.
q (TileTensor[dtype, q_tt_layout, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Query TileTensor.
k (TileTensor[dtype, k_tt_layout, Storage=k.Storage, linear_idx_type=k.linear_idx_type]): Key TileTensor.
v (TileTensor[dtype, v_tt_layout, Storage=v.Storage, linear_idx_type=v.linear_idx_type]): Value TileTensor.
mask_functor (mask_t): Mask instance used to apply the attention mask.
scale (Float32): Softmax temperature scale applied to Q·Kᵀ.
ctx (DeviceContext): GPU device context for kernel dispatch.
num_partitions (Optional[Int]): Override the number of split-K partitions; None selects automatically.
valid_length (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Per-sequence valid lengths for masking padded batches.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink-token weight tensor for attention sinks.