From 60a1474ad95d93891f81d69bd0379b29371e9bf7 Mon Sep 17 00:00:00 2001 From: ggerganov Date: Wed, 26 Aug 2026 19:39:29 +0300 Subject: [PATCH] ggml-metal: FA tensor path: gate to the (256, 256) shape the tensor kernel is only faster than the vec kernel for head size 256/256 (measured +3%..+13% across GQA ratios 8:8 .. 64:8 on M5 Max, nq = kv = 4096). all the other (dk, dv) pairs lose and now fall back to the vec kernel: 64/64: -18% 128/128: -12% MHA, -64% GQA 32:8 192/128: -3% 192/192: -53% 320/256: -20% GQA 512/512: -35%..-65% 576/512: -60%..-70% the large-dk losses are register spills (the f32 Q tile (dk, 8) costs dk/4 registers per thread). the f16 Q attempt to fix them is blocked by an MPP driver/compiler bug: with f16 QK operands the coop destination (P) tile element map changes from (2q x 8kv) to (4q x 4kv) per thread and the hardware P -> f16 right input conversion of get_right_input_cooperative_tensor does not match the new layout (reproduced standalone: the PV output contains the correct values at the wrong positions; the same class of bug breaks the coop element map when the K operand is strided). the smem P workaround (write P to threadgroup memory, PV from a tensor_inline right input, transR = true with the (k, n) k-contiguous tile) is correct (4809/4809 test-backend-ops) but 2-3x slower than the coop register path. correctness: 4809/4809 test-backend-ops FLASH_ATTN_EXT. --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index d6b1a8888d..a8b757d2c8 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2859,13 +2859,18 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const struct ggml_metal_devi return false; } - // only the (dk, dv) pairs with an instantiated kernel + // only the (dk, dv) pairs where the tensor path is at least on par with + // the vec kernel. Measured on M5 Max (nq = kv = 4096, several GQA + // ratios, 8:8 .. 64:8): + // (256, 256): +3%..+13% + // all the other pairs lose: 64/64 -18%, 128/128 -12% MHA / -64% GQA + // 32:8, 192/128 -3%, 192/192 -53%, 320/256 -20% GQA, 512/512 -35%..-65%, + // 576/512 -60%..-70% (per-d-block form: NBLKx the QK^T work; the f16 Q + // hoisted form is blocked by an MPP bug - the coop P -> f16 right input + // conversion is wrong for the P tile layout induced by f16 QK operands - + // and the smem P workaround is 2-3x slower). They fall back to vec. switch (dv) { - case 64: if (dk != 64) return false; break; - case 128: if (dk != 128 && dk != 192) return false; break; - case 192: if (dk != 192) return false; break; - case 256: if (dk != 256 && dk != 320) return false; break; - case 512: if (dk != 512 && dk != 576) return false; break; + case 256: if (dk != 256) return false; break; default: return false; }