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, rank: Int, mask_rank: Int, q_origin: ImmOrigin, output_origin: MutOrigin, //, input_k_fn: def[simd_width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, simd_width], input_v_fn: def[simd_width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, simd_width], input_mask_fn: def[simd_width: Int, mask_rank: Int](IndexList[mask_rank]) capturing thin -> SIMD[dtype, simd_width]](q: LayoutTensor[dtype, q_origin, 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_shape: IndexList[rank], v_shape: IndexList[rank], mask_shape: IndexList[mask_rank], output: LayoutTensor[dtype, output_origin, 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], scale: Float32, sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None, ctx: Optional[DeviceContext] = None)

Computes scaled dot-product flash attention on CPU for the given query, key, value, and mask accessors.

Parameters:

dtype (DType): The element type of the query, key, value, and output tensors (inferred).
rank (Int): The number of dimensions in the query, key, value, and output tensors, either 3 or 4 (inferred).
mask_rank (Int): The number of dimensions in the attention mask tensor (inferred).
q_origin (ImmOrigin): The memory origin of the read-only query tensor (inferred).
output_origin (MutOrigin): The memory origin of the writable output tensor (inferred).
input_k_fn (def[simd_width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, simd_width]): Compile-time function loading a SIMD vector of key elements at a given IndexList index.
input_v_fn (def[simd_width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, simd_width]): Compile-time function loading a SIMD vector of value elements at a given IndexList index.
input_mask_fn (def[simd_width: Int, mask_rank: Int](IndexList[mask_rank]) capturing thin -> SIMD[dtype, simd_width]): Compile-time function loading a SIMD vector of additive mask values at a given IndexList index.

Args:

q (LayoutTensor[dtype, q_origin, 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 in BSHD or BSD layout.
k_shape (IndexList[rank]): Shape of the key tensor.
v_shape (IndexList[rank]): Shape of the value tensor.
mask_shape (IndexList[mask_rank]): Shape of the attention mask tensor.
output (LayoutTensor[dtype, output_origin, 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]): Output tensor to write the attention results into.
scale (Float32): Scaling factor applied to the query-key dot products.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional per-head attention sink weights.
ctx (Optional[DeviceContext]): Optional device context for controlling parallelism.