Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-19 17:24:57 +02:00
Code Issues Packages Projects Releases Wiki Activity

ggml-metal: FA tensor path: gate to the (256, 256) shape

Browse Source 
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.
This commit is contained in:
ggerganov
2026-08-26 19:39:29 +03:00
parent 3f56636c8b
commit 60a1474ad9
1 changed files with 11 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-ops.cpp
+11 -6
View File
@@ -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;
}