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 module
fa_prefill
Apple (Metal) M5 MMA-based flash-attention (FA2) prefill kernel.
Two MmaOpApple 16x16 simdgroup GEMMs (Q.K^T and P.V) sandwiching an online
softmax, for Apple silicon prefill (compute_capability == 5), paged KV via
MHAOperand, BF16/FP16 storage with FP32 accumulation.
4 simdgroups share one threadgroup with no threadgroup memory and no
barrier(); they are co-resident only to share KV reads through the L2, which
beat both a single-simdgroup launch and an SMEM-staged variant. Each simdgroup
independently owns an Sq x depth query-row tile for one (batch, head) and
streams the KV range online in Sk-wide tiles, computing QK one 16-key
column-block at a time. Each column-block resolves its own page-aligned KV
sub-tile, so KV need only satisfy page_size == 0 or page_size % 16 == 0;
other page sizes fall back to mha_gpu_naive via the dispatcher.
The per-row max/sum over a score tile whose rows are scattered across the MMA
fragment uses an XOR-butterfly reduction; air.simd_sum (the decode reduction)
would reduce all 32 lanes and mix rows.
P.V keeps P register-resident: the QK score fragment, cast to q_type, is fed
directly as the P.V A-operand -- the C-output and A-input fragment layouts are
the identical _apple_frag_layout bit-scatter, so P never round-trips memory.
Q, K, and V load from memory each tile.
flash_attention_dispatch selects fa_prefill_apple for Apple prefill by
default; set MODULAR_ENABLE_APPLE_FA_PREFILL=0 to fall back to mha_gpu_naive.
comptime values
DEFAULT_NUM_M_MMAS
comptime DEFAULT_NUM_M_MMAS = 1
DEFAULT_NUM_N_MMAS
comptime DEFAULT_NUM_N_MMAS = 2
FA_PREFILL_APPLE_MAX_HEAD_DIM
comptime FA_PREFILL_APPLE_MAX_HEAD_DIM = 256
MMA_DIM
comptime MMA_DIM = 16
NEG_INF
comptime NEG_INF = SIMD(-3.0000000000000001E+38)
Functions
-
fa_prefill_apple: Host launcher for the Apple M5 flash-attention prefill kernel. -
fa_prefill_apple_core: Per-simdgroup core of the prefill kernel.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!