quantizedSdpaUnfused with the key matmul over three GQA heads per batch:
the queries of each KV head enter the key quantized_matmul as [2, 9] rows
instead of 18, and the scores are restored to [nRep, L] before the mask,
softmax and value matmul. On the M4 Pro GPU (applegpu_g16s) MLX's native
key matvec keeps exact arithmetic in that grouping at 8192 or more keys.
A kernel for one composition: depth-2 speculative verify (L = 3) on 24
query heads, 4 KV heads, head dim 256, B <= 2, a 4-bit group-64 cache and
bf16 or f32 queries, built only by the cache composed for it. Other shapes
are outside its contract (the [2, 9] split needs nRep * L = 18).
quantizedSdpaUnfusedwith the key matmul over three GQA heads per batch: the queries of each KV head enter the key quantized_matmul as [2, 9] rows instead of 18, and the scores are restored to [nRep, L] before the mask, softmax and value matmul. On the M4 Pro GPU (applegpu_g16s) MLX's native key matvec keeps exact arithmetic in that grouping at 8192 or more keys. A kernel for one composition: depth-2 speculative verify (L = 3) on 24 query heads, 4 KV heads, head dim 256, B <= 2, a 4-bit group-64 cache and bf16 or f32 queries, built only by the cache composed for it. Other shapes are outside its contract (the [2, 9] split needs nRep * L = 18).