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_stitch

def naive_fa_decode_apple_stitch[output_type: DType, p_type: DType, k_t: MHAOperand, v_t: MHAOperand, mask_t: MHAMask, output_layout: TensorLayout, p_layout: TensorLayout, valid_length_layout: TensorLayout, ragged: Bool = False, sink: Bool = False, _use_valid_length: Bool = False, _is_cache_length_accurate: Bool = False, *, SplitSize: Int](output: TileTensor[output_type, output_layout, MutAnyOrigin], o_partial: TileTensor[p_type, p_layout, ImmutAnyOrigin], m_partial: TileTensor[p_type, p_layout, ImmutAnyOrigin], l_partial: TileTensor[p_type, p_layout, ImmutAnyOrigin], k: k_t, valid_length: TileTensor[DType.uint32, valid_length_layout, ImmutAnyOrigin], max_prompt_len: Int, max_cache_size: Int, num_heads: Int, depth: Int, num_partitions: Int)

Combines per-partition partials into the final decode attention output.

Grid (num_heads, batch_size), block depth: one thread per head-dim element. Reads the contiguous per-split (o_partial, m_partial, l_partial) buffers written by naive_fa_decode_apple_core and merges them with a log-sum-exp (LSE) reduction in FP32, then writes the normalized acc / l row back into output. The active split count mirrors the producer's cur_cache_len so only the partials actually written are combined.

Constraints:

o_partial, m_partial, and output must be flat 1D TileTensors.

Parameters:

output_type (DType): Element type of the output tensor (inferred).
p_type (DType): Accumulation and partials element type (inferred).
k_t (MHAOperand): MHAOperand type of the key cache operand (inferred).
v_t (MHAOperand): Unused; mirrors mha_gpu_naive for dispatch uniformity (inferred).
mask_t (MHAMask): Unused; mirrors mha_gpu_naive for dispatch uniformity (inferred).
output_layout (TensorLayout): TensorLayout of the output tensor (inferred).
p_layout (TensorLayout): TensorLayout of the partials buffers (inferred).
valid_length_layout (TensorLayout): TensorLayout of the valid_length tensor (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; the producer already bakes the sink contribution into split 0 partials, so the stitch does not act on this flag (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).
SplitSize (Int): Per-partition KV span in keys; must match the producer's SplitSize.

Args:

output (TileTensor[output_type, output_layout, MutAnyOrigin]): Flat 1D output tensor; written with the normalized attention output per (batch, head, depth).
o_partial (TileTensor[p_type, p_layout, ImmutAnyOrigin]): Flat 1D partial output buffer from the producer; one accumulator per (batch, head, depth, split).
m_partial (TileTensor[p_type, p_layout, ImmutAnyOrigin]): Flat 1D partial row-max buffer from the producer; one running max per (batch, head, split).
l_partial (TileTensor[p_type, p_layout, ImmutAnyOrigin]): Flat 1D partial row-sum buffer from the producer; one running denominator per (batch, head, split).
k (k_t): Key cache operand implementing the MHAOperand contract; used for cache_length in the ragged and KVCache paths.
valid_length (TileTensor[DType.uint32, valid_length_layout, ImmutAnyOrigin]): Per-sequence row offsets or query lengths (uint32); meaning depends on ragged and _use_valid_length.
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): Runtime head dimension.
num_partitions (Int): Number of KV splits; must match the producer's partition count.