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

fa_prefill_apple

def fa_prefill_apple[output_type: DType, k_t: MHAOperand, v_t: MHAOperand, mask_t: MHAMask, //, ragged: Bool = False, sink: Bool = False, _use_valid_length: Bool = False, _is_cache_length_accurate: Bool = False, num_simdgroups: Int = Int(4)](q: LayoutTensor[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, output: LayoutTensor[output_type, 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], 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, batch_size: Int, max_prompt_len: Int, max_cache_size: Int, num_heads: Int, depth: Int, group: Int, ctx: DeviceContext, sink_weights: OptionalReg[LayoutTensor[q.dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

Host launcher for the Apple M5 flash-attention prefill kernel.

Mirrors mha_gpu_naive's MHAOperand overload so flash_attention_dispatch routes to it like the fallback, and specializes one kernel per supported depth (a multiple of 16 up to FA_PREFILL_APPLE_MAX_HEAD_DIM). The external LayoutTensor ABI is converted to TileTensor at the enqueue boundary so the kernel is TileTensor-only.

Parameters:

output_type (DType): The dtype of the output tensor (inferred).
k_t (MHAOperand): The MHAOperand type of the paged key cache (inferred).
v_t (MHAOperand): The MHAOperand type of the paged value cache (inferred).
mask_t (MHAMask): The MHAMask type of the attention mask functor (inferred).
ragged (Bool): If True, valid_length is a cumulative offset buffer over variable-length sequences in the batch (defaults to False).
sink (Bool): If True, pre-seed the softmax state from per-head sink weights (defaults to False).
_use_valid_length (Bool): If True, read each batch's query length from valid_length[batch_id] (defaults to False).
_is_cache_length_accurate (Bool): If True, the total attention length equals the query length (no cached prefix); if False, add the key operand's per-batch cache length (defaults to False).
num_simdgroups (Int): Number of simdgroups co-resident per threadgroup, controlling per-launch occupancy (defaults to 4).

Args:

q (LayoutTensor[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]): The query tensor in flattened BSHD layout (batch, seq, head, depth).
k (k_t): The paged key cache operand.
v (v_t): The paged value cache operand.
mask_functor (mask_t): The attention mask functor applied to each score tile.
output (LayoutTensor[output_type, 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]): The output tensor, written in the same BSHD layout as q.
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-batch uint32 sequence lengths or cumulative offsets, depending on ragged and _use_valid_length.
scale (Float32): The softmax scale factor applied to QK score products.
batch_size (Int): Number of sequences in the batch.
max_prompt_len (Int): Maximum query length across the batch, in tokens.
max_cache_size (Int): Maximum KV cache length across the batch, in tokens.
num_heads (Int): Number of query attention heads.
depth (Int): Head dimension; must be a multiple of 16 and at most FA_PREFILL_APPLE_MAX_HEAD_DIM.
group (Int): Number of query heads per KV head (GQA group size).
ctx (DeviceContext): The device context used to enqueue the kernel.
sink_weights (OptionalReg[LayoutTensor[q.dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional per-head sink weights indexed by head id (defaults to None).