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

naive_fa_decode_apple

def naive_fa_decode_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](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 split-K decode attention pair (decode-only).

Parameters:

output_type (DType): The element type of the output tensor (inferred). Unused by the kernels; mirrors mha_gpu_naive for dispatch uniformity.
k_t (MHAOperand): The MHAOperand type of the key cache operand (inferred).
v_t (MHAOperand): The MHAOperand type of the value cache operand (inferred).
mask_t (MHAMask): The MHAMask functor type applied to attention scores (inferred).
ragged (Bool): Whether sequences are ragged with variable lengths and row offsets in valid_length (defaults to False).
sink (Bool): Whether attention sink is enabled, pre-seeding split 0 with per-head sink weights (defaults to False).
_use_valid_length (Bool): Whether to use valid_length for KVCache decode as per-sequence query lengths (defaults to False).
_is_cache_length_accurate (Bool): Whether the cache length equals the query length, so no new-token KV is added (defaults to False).

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; one token per sequence (decode).
k (k_t): The key cache operand implementing the MHAOperand contract.
v (v_t): The value cache operand implementing the MHAOperand contract.
mask_functor (mask_t): The mask functor applied to each attention score.
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 by the stitch kernel with the normalized attention output.
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 row offsets or query lengths (uint32); meaning depends on ragged and _use_valid_length.
scale (Float32): The softmax scale factor applied to Q.K^T scores.
batch_size (Int): Number of sequences in the batch.
max_prompt_len (Int): Maximum prompt length; the dense decode path's query length.
max_cache_size (Int): Full key count for the dense decode path (the K tensor's seq dim).
num_heads (Int): Number of query attention heads.
depth (Int): Head dimension; must be a multiple of WARP_SIZE and at most NAIVE_FA_DECODE_APPLE_MAX_HEAD_DIM.
group (Int): Number of query heads per KV head (GQA group size).
ctx (DeviceContext): The device context used to enqueue kernels and allocate partial buffers.
sink_weights (OptionalReg[LayoutTensor[q.dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Per-head sink weights (shape [num_heads]); read only when sink is True (defaults to None).