Causal attention over bf16 keys and values with GQA folded into rows: the
rep·L query vectors of each KV head multiply its keys in one matmul and its
values in another, so the cache is read once per KV head instead of once per
query head (MLX's fallback broadcasts each KV head over its query heads).
Otherwise MLX's fallback arithmetic, op for op: the scale cast to the query
dtype, the bottom-right causal mask at -inf, the precise softmax. Bit for bit
ops.sdpa on the windows foldedSdpaEligible takes on applegpu_g13s, where
it was measured, and to bf16 precision elsewhere (rows (r, i): head
KVH·rep + r, query i, which sees keys j <= N - L + i). q [B, H, L, D], keys
and values [B, KVH, N, D] (any strides) → [B, H, L, D], owned by the caller.
Causal attention over bf16 keys and values with GQA folded into rows: the rep·L query vectors of each KV head multiply its keys in one matmul and its values in another, so the cache is read once per KV head instead of once per query head (MLX's fallback broadcasts each KV head over its query heads). Otherwise MLX's fallback arithmetic, op for op: the scale cast to the query dtype, the bottom-right causal mask at -inf, the precise softmax. Bit for bit
ops.sdpaon the windowsfoldedSdpaEligibletakes onapplegpu_g13s, where it was measured, and to bf16 precision elsewhere (rows (r, i): head KVH·rep + r, query i, which sees keys j <= N - L + i). q [B, H, L, D], keys and values [B, KVH, N, D] (any strides) → [B, H, L, D], owned by the caller.