From ddad9667af63e78c4a533f8031330102c8be7340 Mon Sep 17 00:00:00 2001 From: Foad Abo Dahood <32059146+masterFoad@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:40:16 +0300 Subject: [PATCH] metal : fix excess threadgroup memory in quantized flash attention (llama/29340) --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index dfe46bac6..2a14cd0c2 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -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);