metal : fix excess threadgroup memory in quantized flash attention (llama/29340)

This commit is contained in:
Foad Abo Dahood
2026-10-19 09:47:10 +02:00
committed by Daniel Bevenius
parent 5ae08fcb13
commit ddad9667af
+5 -1
View File
@@ -3088,7 +3088,9 @@ static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) {
// depending on compute/bandwidth ratio, dequant to f16 kv is not always beneficial
// ref: https://github.com/ggml-org/llama.cpp/pull/27390#issuecomment-5355152767
// TODO: tune per device
if (op->src[0]->ne[1] < 32) {
// large heads need the upfront dequant to fit the non-vec threadgroup memory
if (op->src[0]->ne[1] < 32 &&
(op->src[0]->ne[0] < 512 || ggml_metal_op_flash_attn_ext_use_vec(op))) {
return false;
}
@@ -3827,6 +3829,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_set_buffer (enc, bid_blk, 7);
ggml_metal_encoder_set_buffer (enc, bid_dst, 8);
GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size);
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nqptg - 1)/nqptg, ne02, ne03, 32, nsg, 1);