mirror of
https://github.com/ggml-org/whisper.cpp.git
synced 2026-10-21 09:41:29 +02:00
metal : fix excess threadgroup memory in quantized flash attention (llama/29340)
This commit is contained in:
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user