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_kv_cache

def flash_attention_kv_cache[dtype: DType, cache_t: KVCacheT, q_origin: ImmOrigin, output_origin: MutOrigin, //](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: cache_t, v: cache_t, mask: LayoutTensor[dtype, 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, 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], sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

Computes flash attention on CPU using a KV cache with an additive LayoutTensor mask.

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 layout.
k (cache_t): Key cache.
v (cache_t): Value cache.
mask (LayoutTensor[dtype, 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]): Additive attention mask tensor.
scale (Float32): Scaling factor applied to the query-key dot products.
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.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional per-head attention sink weights.

Computes flash attention on CPU using a KV cache with an MHAMask-based mask.

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 layout.
k (cache_t): Key cache.
v (cache_t): Value cache.
mask (mask_t): MHAMask applied to the attention scores.
scale (Float32): Scaling factor applied to the query-key dot products.
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.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional per-head attention sink weights.

def flash_attention_kv_cache[dtype: DType, cache_t: KVCacheT, mask_t: MHAMask, q_origin: ImmOrigin, output_origin: MutOrigin, //](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], q_input_row_offsets: LayoutTensor[DType.uint32, element_layout=q_input_row_offsets.element_layout, layout_int_type=q_input_row_offsets.layout_int_type, linear_idx_type=q_input_row_offsets.linear_idx_type, masked=q_input_row_offsets.masked, alignment=q_input_row_offsets.alignment], kv_input_row_offsets: LayoutTensor[DType.uint32, element_layout=kv_input_row_offsets.element_layout, layout_int_type=kv_input_row_offsets.layout_int_type, linear_idx_type=kv_input_row_offsets.linear_idx_type, masked=kv_input_row_offsets.masked, alignment=kv_input_row_offsets.alignment], k: cache_t, v: cache_t, mask: mask_t, scale: Float32, 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], sink_weights: OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

Computes flash attention on CPU for ragged tensors using a KV cache with an MHAMask-based mask.

Parameters:

dtype (DType): The element type of the query, key, value, and output tensors (inferred).
cache_t (KVCacheT): The KV cache type storing key and value states (inferred).
mask_t (MHAMask): The MHAMask type applied to the attention scores (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).

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]): Flattened query tensor indexed by (row_offset, head, depth).
q_input_row_offsets (LayoutTensor[DType.uint32, element_layout=q_input_row_offsets.element_layout, layout_int_type=q_input_row_offsets.layout_int_type, linear_idx_type=q_input_row_offsets.linear_idx_type, masked=q_input_row_offsets.masked, alignment=q_input_row_offsets.alignment]): Per-batch start offsets into the flattened query tensor; batch b spans rows [offsets[b], offsets[b + 1]).
kv_input_row_offsets (LayoutTensor[DType.uint32, element_layout=kv_input_row_offsets.element_layout, layout_int_type=kv_input_row_offsets.layout_int_type, linear_idx_type=kv_input_row_offsets.linear_idx_type, masked=kv_input_row_offsets.masked, alignment=kv_input_row_offsets.alignment]): Per-batch start offsets into the flattened KV tensors; batch b spans rows [offsets[b], offsets[b + 1]).
k (cache_t): Key cache storing per-head key states.
v (cache_t): Value cache storing per-head value states.
mask (mask_t): MHAMask applied additively to the query-key attention scores.
scale (Float32): Scaling factor applied to the query-key dot products.
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.
sink_weights (OptionalReg[LayoutTensor[dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional per-head attention sink weights.