From 5d6b00b7f3cfaeac6ce322f358e45cf112c22dbf Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 19 Aug 2026 14:10:25 +0300 Subject: [PATCH 1/9] metal: dequantize q8_0 KV to f16 before flash attention Add a preprocessing pass for GGML_OP_FLASH_ATTN_EXT on the Metal backend: when the KV cache is quantized (Q8_0 for now), dequantize K and V into a contiguous F16 scratch buffer and run the existing F16 flash attention kernels on it, instead of the in-kernel dequantization path. - new kernel kernel_flash_attn_ext_dequant_to_f16: one thread per quant block (K then V), stride-aware so permuted KV is supported; instantiated for Q8_0 (extending to Q4_0/Q4_1/Q5_0/Q5_1 is one instantiation + one gate case) - the gate is type-only: dequantize whenever the KV is quantized, regardless of head sizes, GQA ratio or n_kv; the attention kernels themselves are untouched - the F16 copies live in the op's own scratch allocation (ggml_metal_op_flash_attn_ext_extra_dequant_f16); the KV pad kernel reads the dequantized buffers when the path is active - the FA pipeline getters gain a use_f16_kv flag selecting the existing f16 kernels and contiguous strides - ref: https://github.com/ggml-org/llama.cpp/pull/25556 Verification (M2 Ultra): - test-backend-ops test -o FLASH_ATTN_EXT: 4798/4798 pass, including the new q8_0 eval cases (decode/prompt, permuted, sinks+ALiBi+softcap, kv=113 pad path, kv=16384) - llama-perplexity on Qwen2.5-0.5B with -ctk q8_0 -ctv q8_0 matches the f16 KV reference (PPL 1.0008 vs 1.0008) Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/src/ggml-metal/ggml-metal-device.cpp | 39 +++- ggml/src/ggml-metal/ggml-metal-device.h | 10 +- ggml/src/ggml-metal/ggml-metal-impl.h | 21 +++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 208 +++++++++++++++++----- ggml/src/ggml-metal/ggml-metal-ops.h | 1 + ggml/src/ggml-metal/ggml-metal.cpp | 1 + ggml/src/ggml-metal/ggml-metal.metal | 58 ++++++ tests/test-backend-ops.cpp | 24 +++ 8 files changed, 309 insertions(+), 53 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 953c757558a..905f0ec1bd6 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1409,6 +1409,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_p return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_dequant_to_f16( + ggml_metal_library_t lib, + const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + char base[256]; + + snprintf(base, 256, "kernel_flash_attn_ext_dequant_%s_to_f16", ggml_type_name(op->src[1]->type)); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, base, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_blk( ggml_metal_library_t lib, const struct ggml_tensor * op, @@ -1460,7 +1477,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( bool has_bias, bool has_scap, bool has_kvpad, - int32_t nsg) { + int32_t nsg, + bool use_f16_kv) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; @@ -1469,15 +1487,17 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( const int32_t dk = (int32_t) op->src[1]->ne[0]; const int32_t dv = (int32_t) op->src[2]->ne[0]; - const int32_t ns10 = op->src[1]->nb[1]/op->src[1]->nb[0]; - const int32_t ns20 = op->src[2]->nb[1]/op->src[2]->nb[0]; + const int32_t ns10 = use_f16_kv ? dk : op->src[1]->nb[1]/op->src[1]->nb[0]; + const int32_t ns20 = use_f16_kv ? dv : op->src[2]->nb[1]/op->src[2]->nb[0]; + + const char * type = use_f16_kv ? "f16" : ggml_type_name(op->src[1]->type); // do bounds checks for the mask? const bool bc_mask = op->src[3] && (op->src[3]->ne[1] % 8 != 0); snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d", "flash_attn_ext", - ggml_type_name(op->src[1]->type), + type, dk, dv); @@ -1526,7 +1546,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v bool has_scap, bool has_kvpad, int32_t nsg, - int32_t nwg) { + int32_t nwg, + bool use_f16_kv) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; @@ -1535,12 +1556,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v const int32_t dk = (int32_t) op->src[1]->ne[0]; const int32_t dv = (int32_t) op->src[2]->ne[0]; - const int32_t ns10 = op->src[1]->nb[1]/op->src[1]->nb[0]; - const int32_t ns20 = op->src[2]->nb[1]/op->src[2]->nb[0]; + const int32_t ns10 = use_f16_kv ? dk : op->src[1]->nb[1]/op->src[1]->nb[0]; + const int32_t ns20 = use_f16_kv ? dv : op->src[2]->nb[1]/op->src[2]->nb[0]; + + const char * type = use_f16_kv ? "f16" : ggml_type_name(op->src[1]->type); snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d", "flash_attn_ext_vec", - ggml_type_name(op->src[1]->type), + type, dk, dv); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 7e1deeaa210..26fd5760210 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -176,6 +176,10 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_mask, int32_t ncpsg); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_dequant_to_f16( + ggml_metal_library_t lib, + const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_blk( ggml_metal_library_t lib, const struct ggml_tensor * op, @@ -190,7 +194,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_bias, bool has_scap, bool has_kvpad, - int32_t nsg); + int32_t nsg, + bool use_f16_kv); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec( ggml_metal_library_t lib, @@ -201,7 +206,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_scap, bool has_kvpad, int32_t nsg, - int32_t nwg); + int32_t nwg, + bool use_f16_kv); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce( ggml_metal_library_t lib, diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 1f6e8c48bcb..ebb80d53612 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -343,6 +343,27 @@ typedef struct { bool src2; } ggml_metal_kargs_rope; +typedef struct { + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + int32_t ne20; + int32_t ne21; + int32_t ne22; + int32_t ne23; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + uint64_t nb20; + uint64_t nb21; + uint64_t nb22; + uint64_t nb23; + int32_t nblocks1; // number of blocks in K + int32_t nblocks2; // number of blocks in V +} ggml_metal_kargs_flash_attn_ext_dequant_to_f16; + typedef struct { int32_t ne11; int32_t ne_12_2; // assume K and V are same shape diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index b7f9b2d0d9c..5d2f13cb839 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2801,6 +2801,29 @@ bool ggml_metal_op_flash_attn_ext_use_vec(const ggml_tensor * op) { return (ne01 < 20) && (ne00 % 32 == 0); } +// ref: https://github.com/ggml-org/llama.cpp/pull/25556 +// dequantize the quantized KV cache to F16 before running the F16 flash attention kernels +static bool ggml_metal_op_flash_attn_ext_use_dequant_f16(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + switch (op->src[1]->type) { + case GGML_TYPE_Q8_0: + return true; + // extend to the other quantized KV types (q4_0, q4_1, q5_0, q5_1) as dequant kernels are added + default: + return false; + } +} + +// size of the F16 dequantized K tensor; the dequantized V tensor follows it in the same scratch buffer +static size_t ggml_metal_op_flash_attn_ext_dequant_f16_k_size(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + + return GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne10*ne11*ne12*ne13, 16); +} + size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); @@ -2816,6 +2839,11 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) { size_t res = 0; const bool has_mask = op->src[3] != nullptr; + const bool use_dequant_f16 = ggml_metal_op_flash_attn_ext_use_dequant_f16(op); + + // when the KV is dequantized to F16, the pad kernel copies the tail chunk from the F16 scratch buffer + const uint64_t nb11_pad = use_dequant_f16 ? sizeof(ggml_fp16_t)*ne10 : nb11; + const uint64_t nb21_pad = use_dequant_f16 ? sizeof(ggml_fp16_t)*ne20 : nb21; // note: the non-vec kernel requires more extra memory, so always reserve for it GGML_ASSERT(OP_FLASH_ATTN_EXT_NCPSG >= OP_FLASH_ATTN_EXT_VEC_NCPSG); @@ -2828,8 +2856,8 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) { if (has_kvpad) { res += OP_FLASH_ATTN_EXT_VEC_NCPSG*( - nb11*ne12*ne13 + - nb21*ne22*ne23 + + nb11_pad*ne12*ne13 + + nb21_pad*ne22*ne23 + (has_mask ? ggml_type_size(GGML_TYPE_F16)*ne31*ne32*ne33 : 0)); } } else { @@ -2838,8 +2866,8 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) { if (has_kvpad) { res += OP_FLASH_ATTN_EXT_NCPSG*( - nb11*ne12*ne13 + - nb21*ne22*ne23 + + nb11_pad*ne12*ne13 + + nb21_pad*ne22*ne23 + (has_mask ? ggml_type_size(GGML_TYPE_F16)*ne31*ne32*ne33 : 0)); } } @@ -2915,6 +2943,21 @@ size_t ggml_metal_op_flash_attn_ext_extra_tmp(const ggml_tensor * op) { return res; } +size_t ggml_metal_op_flash_attn_ext_extra_dequant_f16(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + if (!ggml_metal_op_flash_attn_ext_use_dequant_f16(op)) { + return 0; + } + + GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); + + const size_t k_size = ggml_metal_op_flash_attn_ext_dequant_f16_k_size(op); + const size_t v_size = GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne20*ne21*ne22*ne23, 16); + + return k_size + v_size; +} + int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -2989,6 +3032,85 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_buffer_id bid_tmp = bid_blk; bid_tmp.offs += ggml_metal_op_flash_attn_ext_extra_blk(op); + ggml_metal_buffer_id bid_dequant_f16 = bid_tmp; + bid_dequant_f16.offs += ggml_metal_op_flash_attn_ext_extra_tmp(op); + + const bool use_dequant_f16 = ggml_metal_op_flash_attn_ext_use_dequant_f16(op); + + ggml_metal_buffer_id bid_k = bid_src1; + ggml_metal_buffer_id bid_v = bid_src2; + + uint64_t nb10_attn = nb10; + uint64_t nb11_attn = nb11; + uint64_t nb12_attn = nb12; + uint64_t nb13_attn = nb13; + uint64_t nb20_attn = nb20; + uint64_t nb21_attn = nb21; + uint64_t nb22_attn = nb22; + uint64_t nb23_attn = nb23; + + if (use_dequant_f16) { + assert(ggml_metal_op_flash_attn_ext_extra_dequant_f16(op) != 0); + + const int64_t nblocks1_64 = (ne10/ggml_blck_size(op->src[1]->type))*(int64_t) ne11*ne12*ne13; + const int64_t nblocks2_64 = (ne20/ggml_blck_size(op->src[2]->type))*(int64_t) ne21*ne22*ne23; + GGML_ASSERT(nblocks1_64 + nblocks2_64 <= INT32_MAX/2); + const int32_t nblocks1 = nblocks1_64; + const int32_t nblocks2 = nblocks2_64; + + ggml_metal_buffer_id bid_v_f16 = bid_dequant_f16; + bid_v_f16.offs += ggml_metal_op_flash_attn_ext_dequant_f16_k_size(op); + + ggml_metal_kargs_flash_attn_ext_dequant_to_f16 args0 = { + /*.ne10 =*/ ne10, + /*.ne11 =*/ ne11, + /*.ne12 =*/ ne12, + /*.ne13 =*/ ne13, + /*.ne20 =*/ ne20, + /*.ne21 =*/ ne21, + /*.ne22 =*/ ne22, + /*.ne23 =*/ ne23, + /*.nb10 =*/ nb10, + /*.nb11 =*/ nb11, + /*.nb12 =*/ nb12, + /*.nb13 =*/ nb13, + /*.nb20 =*/ nb20, + /*.nb21 =*/ nb21, + /*.nb22 =*/ nb22, + /*.nb23 =*/ nb23, + /*.nblocks1 =*/ nblocks1, + /*.nblocks2 =*/ nblocks2, + }; + + auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_dequant_to_f16(lib, op); + const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline0), 256); + + ggml_metal_encoder_set_pipeline(enc, pipeline0); + ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); + ggml_metal_encoder_set_buffer (enc, bid_src1, 1); + ggml_metal_encoder_set_buffer (enc, bid_src2, 2); + ggml_metal_encoder_set_buffer (enc, bid_dequant_f16, 3); + ggml_metal_encoder_set_buffer (enc, bid_v_f16, 4); + + ggml_metal_encoder_dispatch_threadgroups(enc, (nblocks1 + nblocks2 + nth - 1)/nth, 1, 1, nth, 1, 1); + + // the pad and attention kernels read the dequantized KV + ggml_metal_op_concurrency_reset(ctx); + + bid_k = bid_dequant_f16; + bid_v = bid_v_f16; + + // contiguous F16 layout of the dequantized KV + nb10_attn = sizeof(ggml_fp16_t); + nb11_attn = nb10_attn*ne10; + nb12_attn = nb11_attn*ne11; + nb13_attn = nb12_attn*ne12; + nb20_attn = sizeof(ggml_fp16_t); + nb21_attn = nb20_attn*ne20; + nb22_attn = nb21_attn*ne21; + nb23_attn = nb22_attn*ne22; + } + if (!ggml_metal_op_flash_attn_ext_use_vec(op)) { // half8x8 kernel const int nqptg = OP_FLASH_ATTN_EXT_NQPSG; // queries per threadgroup @@ -3009,12 +3131,12 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.ne11 =*/ne11, /*.ne_12_2 =*/ne12, /*.ne_12_3 =*/ne13, - /*.nb11 =*/nb11, - /*.nb12 =*/nb12, - /*.nb13 =*/nb13, - /*.nb21 =*/nb21, - /*.nb22 =*/nb22, - /*.nb23 =*/nb23, + /*.nb11 =*/nb11_attn, + /*.nb12 =*/nb12_attn, + /*.nb13 =*/nb13_attn, + /*.nb21 =*/nb21_attn, + /*.nb22 =*/nb22_attn, + /*.nb23 =*/nb23_attn, /*.ne31 =*/ne31, /*.ne32 =*/ne32, /*.ne33 =*/ne33, @@ -3027,8 +3149,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_pipeline(enc, pipeline0); ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); - ggml_metal_encoder_set_buffer (enc, bid_src1, 1); - ggml_metal_encoder_set_buffer (enc, bid_src2, 2); + ggml_metal_encoder_set_buffer (enc, bid_k, 1); + ggml_metal_encoder_set_buffer (enc, bid_v, 2); ggml_metal_encoder_set_buffer (enc, bid_src3, 3); ggml_metal_encoder_set_buffer (enc, bid_pad, 4); @@ -3073,7 +3195,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_op_concurrency_reset(ctx); } - const int is_q = ggml_is_quantized(op->src[1]->type) ? 1 : 0; + const int is_q = !use_dequant_f16 && ggml_is_quantized(op->src[1]->type) ? 1 : 0; // 2*(2*ncpsg) // ncpsg soft_max values + ncpsg mask values @@ -3114,14 +3236,14 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.ne11 =*/ ne11, /*.ne_12_2 =*/ ne12, /*.ne_12_3 =*/ ne13, - /*.ns10 =*/ int32_t(nb11/nb10), - /*.nb11 =*/ nb11, - /*.nb12 =*/ nb12, - /*.nb13 =*/ nb13, - /*.ns20 =*/ int32_t(nb21/nb20), - /*.nb21 =*/ nb21, - /*.nb22 =*/ nb22, - /*.nb23 =*/ nb23, + /*.ns10 =*/ int32_t(nb11_attn/nb10_attn), + /*.nb11 =*/ nb11_attn, + /*.nb12 =*/ nb12_attn, + /*.nb13 =*/ nb13_attn, + /*.ns20 =*/ int32_t(nb21_attn/nb20_attn), + /*.nb21 =*/ nb21_attn, + /*.nb22 =*/ nb22_attn, + /*.nb23 =*/ nb23_attn, /*.ne31 =*/ ne31, /*.ne32 =*/ ne32, /*.ne33 =*/ ne33, @@ -3139,13 +3261,13 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, use_dequant_f16); ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ggml_metal_encoder_set_buffer (enc, bid_src0, 1); - ggml_metal_encoder_set_buffer (enc, bid_src1, 2); - ggml_metal_encoder_set_buffer (enc, bid_src2, 3); + ggml_metal_encoder_set_buffer (enc, bid_k, 2); + ggml_metal_encoder_set_buffer (enc, bid_v, 3); ggml_metal_encoder_set_buffer (enc, bid_src3, 4); ggml_metal_encoder_set_buffer (enc, bid_src4, 5); ggml_metal_encoder_set_buffer (enc, bid_pad, 6); @@ -3177,12 +3299,12 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.ne11 =*/ne11, /*.ne_12_2 =*/ne12, /*.ne_12_3 =*/ne13, - /*.nb11 =*/nb11, - /*.nb12 =*/nb12, - /*.nb13 =*/nb13, - /*.nb21 =*/nb21, - /*.nb22 =*/nb22, - /*.nb23 =*/nb23, + /*.nb11 =*/nb11_attn, + /*.nb12 =*/nb12_attn, + /*.nb13 =*/nb13_attn, + /*.nb21 =*/nb21_attn, + /*.nb22 =*/nb22_attn, + /*.nb23 =*/nb23_attn, /*.ne31 =*/ne31, /*.ne32 =*/ne32, /*.ne33 =*/ne33, @@ -3195,8 +3317,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_pipeline(enc, pipeline0); ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); - ggml_metal_encoder_set_buffer (enc, bid_src1, 1); - ggml_metal_encoder_set_buffer (enc, bid_src2, 2); + ggml_metal_encoder_set_buffer (enc, bid_k, 1); + ggml_metal_encoder_set_buffer (enc, bid_v, 2); ggml_metal_encoder_set_buffer (enc, bid_src3, 3); ggml_metal_encoder_set_buffer (enc, bid_pad, 4); @@ -3252,14 +3374,14 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.ne11 =*/ ne11, /*.ne_12_2 =*/ ne12, /*.ne_12_3 =*/ ne13, - /*.ns10 =*/ int32_t(nb11/nb10), - /*.nb11 =*/ nb11, - /*.nb12 =*/ nb12, - /*.nb13 =*/ nb13, - /*.ns20 =*/ int32_t(nb21/nb20), - /*.nb21 =*/ nb21, - /*.nb22 =*/ nb22, - /*.nb23 =*/ nb23, + /*.ns10 =*/ int32_t(nb11_attn/nb10_attn), + /*.nb11 =*/ nb11_attn, + /*.nb12 =*/ nb12_attn, + /*.nb13 =*/ nb13_attn, + /*.ns20 =*/ int32_t(nb21_attn/nb20_attn), + /*.nb21 =*/ nb21_attn, + /*.nb22 =*/ nb22_attn, + /*.nb23 =*/ nb23_attn, /*.ne31 =*/ ne31, /*.ne32 =*/ ne32, /*.ne33 =*/ ne33, @@ -3277,15 +3399,15 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg, use_dequant_f16); GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ggml_metal_encoder_set_buffer (enc, bid_src0, 1); - ggml_metal_encoder_set_buffer (enc, bid_src1, 2); - ggml_metal_encoder_set_buffer (enc, bid_src2, 3); + ggml_metal_encoder_set_buffer (enc, bid_k, 2); + ggml_metal_encoder_set_buffer (enc, bid_v, 3); ggml_metal_encoder_set_buffer (enc, bid_src3, 4); ggml_metal_encoder_set_buffer (enc, bid_src4, 5); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index b03b59e0bd9..2b3d94a21f9 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -42,6 +42,7 @@ bool ggml_metal_op_flash_attn_ext_use_vec(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_pad(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_blk(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op); +size_t ggml_metal_op_flash_attn_ext_extra_dequant_f16(const struct ggml_tensor * op); int ggml_metal_op_concat (ggml_metal_op_t ctx, int idx); int ggml_metal_op_repeat (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index ef3c92f2712..9b70d7f4e48 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -225,6 +225,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_ res += ggml_metal_op_flash_attn_ext_extra_pad(tensor); res += ggml_metal_op_flash_attn_ext_extra_blk(tensor); res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor); + res += ggml_metal_op_flash_attn_ext_extra_dequant_f16(tensor); } break; case GGML_OP_CUMSUM: case GGML_OP_ARGSORT: diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index c3a95dd6e52..e5d1f7a3908 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -6303,6 +6303,64 @@ template [[host_name("kernel_fwht_f32_128")]] kernel kernel_fwht_t kernel_fwht_f template [[host_name("kernel_fwht_f32_256")]] kernel kernel_fwht_t kernel_fwht_f32<256>; template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f32<512>; +// dequantize the quantized KV cache to contiguous F16 before running the F16 flash attention kernels +// - one thread per block; the first args.nblocks1 threads dequantize K, the next args.nblocks2 dequantize V +// - ref: https://github.com/ggml-org/llama.cpp/pull/25556 +template < + typename block_t, + short QK, + void (*deq_t4x4)(device const block_t *, short, thread float4x4 &)> +kernel void kernel_flash_attn_ext_dequant_to_f16( + constant ggml_metal_kargs_flash_attn_ext_dequant_to_f16 & args, + device const char * k, + device const char * v, + device half * k_dst, + device half * v_dst, + uint gid [[thread_position_in_grid]]) { + if (gid >= (uint) args.nblocks1 + (uint) args.nblocks2) { + return; + } + + const bool is_v = gid >= (uint) args.nblocks1; + const uint idst = is_v ? gid - (uint) args.nblocks1 : gid; + uint ib = idst; + + const int32_t ne1 = is_v ? args.ne20 : args.ne10; + const int32_t ne2 = is_v ? args.ne21 : args.ne11; + const int32_t ne3 = is_v ? args.ne22 : args.ne12; + + const uint64_t nb1 = is_v ? args.nb20 : args.nb10; + const uint64_t nb2 = is_v ? args.nb21 : args.nb11; + const uint64_t nb3 = is_v ? args.nb22 : args.nb12; + const uint64_t nb4 = is_v ? args.nb23 : args.nb13; + + const uint nb = ne1/QK; + const uint i0 = ib%nb; + ib /= nb; + const uint i1 = ib%ne2; + ib /= ne2; + const uint i2 = ib%ne3; + const uint i3 = ib/ne3; + + const uint64_t offs = i0*nb1 + i1*nb2 + i2*nb3 + i3*nb4; + + device const block_t * src = (device const block_t *) ((is_v ? v : k) + offs); + device half4 * dst = (device half4 *) (is_v ? v_dst : k_dst) + (QK/4)*idst; + + for (short i = 0; i < QK/16; ++i) { + float4x4 reg; + deq_t4x4(src, i, reg); + dst[4*i + 0] = (half4) reg[0]; + dst[4*i + 1] = (half4) reg[1]; + dst[4*i + 2] = (half4) reg[2]; + dst[4*i + 3] = (half4) reg[3]; + } +} + +typedef decltype(kernel_flash_attn_ext_dequant_to_f16) kernel_flash_attn_ext_dequant_to_f16_t; + +template [[host_name("kernel_flash_attn_ext_dequant_q8_0_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; + constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]]; constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]]; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 52cd747d561..9254ba09d9c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9910,6 +9910,15 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q2_0)); test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_F16)); + // q8_0 KV cases covering the Metal dequant->F16 FA path: decode and prompt batches, KV pad, + // permuted KV, feature flags, and long context + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 2}, 1025, 1, true, true, 8, 30, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1025, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). for (int64_t kv : { 4096, 16384 }) { @@ -10295,6 +10304,21 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // q8_0 KV dequant->F16 path (Metal) — long context, decode and prompt, with F16 references + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 128, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 2048, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + for (int kv : { 4096, 8192, 16384, }) { for (int hs : { 64, 128, }) { for (int nr : { 1, 4, }) { From ad615b62c1b1ba9cfb5f5de0101dccb68c08413d Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 19 Aug 2026 16:18:07 +0300 Subject: [PATCH 2/9] metal : launch the FA KV dequant kernel separately for K and V Simplify kernel_flash_attn_ext_dequant_to_f16: it now dequantizes a single tensor (its own ne/nb and dst) with no is_v branching, and the op dispatches it twice with the same pipeline - once for K and once for V. The kargs struct shrinks to a single ne/nb set plus nblocks. Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/src/ggml-metal/ggml-metal-impl.h | 27 ++++------- ggml/src/ggml-metal/ggml-metal-ops.cpp | 65 +++++++++++++++----------- ggml/src/ggml-metal/ggml-metal.metal | 45 ++++++------------ 3 files changed, 62 insertions(+), 75 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index ebb80d53612..caf6fc202d8 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -344,24 +344,15 @@ typedef struct { } ggml_metal_kargs_rope; typedef struct { - int32_t ne10; - int32_t ne11; - int32_t ne12; - int32_t ne13; - int32_t ne20; - int32_t ne21; - int32_t ne22; - int32_t ne23; - uint64_t nb10; - uint64_t nb11; - uint64_t nb12; - uint64_t nb13; - uint64_t nb20; - uint64_t nb21; - uint64_t nb22; - uint64_t nb23; - int32_t nblocks1; // number of blocks in K - int32_t nblocks2; // number of blocks in V + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t nblocks; // number of blocks in the tensor } ggml_metal_kargs_flash_attn_ext_dequant_to_f16; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 5d2f13cb839..d4af32ca7b4 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -3054,45 +3054,56 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { const int64_t nblocks1_64 = (ne10/ggml_blck_size(op->src[1]->type))*(int64_t) ne11*ne12*ne13; const int64_t nblocks2_64 = (ne20/ggml_blck_size(op->src[2]->type))*(int64_t) ne21*ne22*ne23; - GGML_ASSERT(nblocks1_64 + nblocks2_64 <= INT32_MAX/2); + GGML_ASSERT(nblocks1_64 <= INT32_MAX); + GGML_ASSERT(nblocks2_64 <= INT32_MAX); const int32_t nblocks1 = nblocks1_64; const int32_t nblocks2 = nblocks2_64; ggml_metal_buffer_id bid_v_f16 = bid_dequant_f16; bid_v_f16.offs += ggml_metal_op_flash_attn_ext_dequant_f16_k_size(op); - ggml_metal_kargs_flash_attn_ext_dequant_to_f16 args0 = { - /*.ne10 =*/ ne10, - /*.ne11 =*/ ne11, - /*.ne12 =*/ ne12, - /*.ne13 =*/ ne13, - /*.ne20 =*/ ne20, - /*.ne21 =*/ ne21, - /*.ne22 =*/ ne22, - /*.ne23 =*/ ne23, - /*.nb10 =*/ nb10, - /*.nb11 =*/ nb11, - /*.nb12 =*/ nb12, - /*.nb13 =*/ nb13, - /*.nb20 =*/ nb20, - /*.nb21 =*/ nb21, - /*.nb22 =*/ nb22, - /*.nb23 =*/ nb23, - /*.nblocks1 =*/ nblocks1, - /*.nblocks2 =*/ nblocks2, - }; - auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_dequant_to_f16(lib, op); const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline0), 256); + // K + ggml_metal_kargs_flash_attn_ext_dequant_to_f16 args_k = { + /*.ne0 =*/ ne10, + /*.ne1 =*/ ne11, + /*.ne2 =*/ ne12, + /*.ne3 =*/ ne13, + /*.nb0 =*/ nb10, + /*.nb1 =*/ nb11, + /*.nb2 =*/ nb12, + /*.nb3 =*/ nb13, + /*.nblocks =*/ nblocks1, + }; + ggml_metal_encoder_set_pipeline(enc, pipeline0); - ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); + ggml_metal_encoder_set_bytes (enc, &args_k, sizeof(args_k), 0); ggml_metal_encoder_set_buffer (enc, bid_src1, 1); - ggml_metal_encoder_set_buffer (enc, bid_src2, 2); - ggml_metal_encoder_set_buffer (enc, bid_dequant_f16, 3); - ggml_metal_encoder_set_buffer (enc, bid_v_f16, 4); + ggml_metal_encoder_set_buffer (enc, bid_dequant_f16, 2); + + ggml_metal_encoder_dispatch_threadgroups(enc, (nblocks1 + nth - 1)/nth, 1, 1, nth, 1, 1); + + // V + ggml_metal_kargs_flash_attn_ext_dequant_to_f16 args_v = { + /*.ne0 =*/ ne20, + /*.ne1 =*/ ne21, + /*.ne2 =*/ ne22, + /*.ne3 =*/ ne23, + /*.nb0 =*/ nb20, + /*.nb1 =*/ nb21, + /*.nb2 =*/ nb22, + /*.nb3 =*/ nb23, + /*.nblocks =*/ nblocks2, + }; + + ggml_metal_encoder_set_pipeline(enc, pipeline0); + ggml_metal_encoder_set_bytes (enc, &args_v, sizeof(args_v), 0); + ggml_metal_encoder_set_buffer (enc, bid_src2, 1); + ggml_metal_encoder_set_buffer (enc, bid_v_f16, 2); - ggml_metal_encoder_dispatch_threadgroups(enc, (nblocks1 + nblocks2 + nth - 1)/nth, 1, 1, nth, 1, 1); + ggml_metal_encoder_dispatch_threadgroups(enc, (nblocks2 + nth - 1)/nth, 1, 1, nth, 1, 1); // the pad and attention kernels read the dequantized KV ggml_metal_op_concurrency_reset(ctx); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index e5d1f7a3908..29c8dd19024 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -6303,8 +6303,8 @@ template [[host_name("kernel_fwht_f32_128")]] kernel kernel_fwht_t kernel_fwht_f template [[host_name("kernel_fwht_f32_256")]] kernel kernel_fwht_t kernel_fwht_f32<256>; template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f32<512>; -// dequantize the quantized KV cache to contiguous F16 before running the F16 flash attention kernels -// - one thread per block; the first args.nblocks1 threads dequantize K, the next args.nblocks2 dequantize V +// dequantize a quantized KV cache tensor to contiguous F16 before running the F16 flash attention kernels +// - one thread per block; dispatched separately for K and V // - ref: https://github.com/ggml-org/llama.cpp/pull/25556 template < typename block_t, @@ -6312,40 +6312,25 @@ template < void (*deq_t4x4)(device const block_t *, short, thread float4x4 &)> kernel void kernel_flash_attn_ext_dequant_to_f16( constant ggml_metal_kargs_flash_attn_ext_dequant_to_f16 & args, - device const char * k, - device const char * v, - device half * k_dst, - device half * v_dst, + device const char * x, + device half * x_dst, uint gid [[thread_position_in_grid]]) { - if (gid >= (uint) args.nblocks1 + (uint) args.nblocks2) { + if (gid >= (uint) args.nblocks) { return; } - const bool is_v = gid >= (uint) args.nblocks1; - const uint idst = is_v ? gid - (uint) args.nblocks1 : gid; - uint ib = idst; - - const int32_t ne1 = is_v ? args.ne20 : args.ne10; - const int32_t ne2 = is_v ? args.ne21 : args.ne11; - const int32_t ne3 = is_v ? args.ne22 : args.ne12; - - const uint64_t nb1 = is_v ? args.nb20 : args.nb10; - const uint64_t nb2 = is_v ? args.nb21 : args.nb11; - const uint64_t nb3 = is_v ? args.nb22 : args.nb12; - const uint64_t nb4 = is_v ? args.nb23 : args.nb13; - - const uint nb = ne1/QK; - const uint i0 = ib%nb; - ib /= nb; - const uint i1 = ib%ne2; - ib /= ne2; - const uint i2 = ib%ne3; - const uint i3 = ib/ne3; + const uint nb = args.ne0/QK; + const uint i0 = gid%nb; + uint ib = gid/nb; + const uint i1 = ib%args.ne1; + ib /= args.ne1; + const uint i2 = ib%args.ne2; + const uint i3 = ib/args.ne2; - const uint64_t offs = i0*nb1 + i1*nb2 + i2*nb3 + i3*nb4; + const uint64_t offs = i0*args.nb0 + i1*args.nb1 + i2*args.nb2 + i3*args.nb3; - device const block_t * src = (device const block_t *) ((is_v ? v : k) + offs); - device half4 * dst = (device half4 *) (is_v ? v_dst : k_dst) + (QK/4)*idst; + device const block_t * src = (device const block_t *) (x + offs); + device half4 * dst = (device half4 *) x_dst + (QK/4)*gid; for (short i = 0; i < QK/16; ++i) { float4x4 reg; From 74c13eea3d62703545ad7151880572d9cabf0cea Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 19 Aug 2026 16:23:54 +0300 Subject: [PATCH 3/9] metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash attention The dequant pass now covers all quantized KV types supported by the Metal flash attention kernels. The dequant kernel, kargs, scratch allocation and dispatch are type-generic, so each type is one kernel instantiation plus one gate case. Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 5 ++++- ggml/src/ggml-metal/ggml-metal.metal | 4 ++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index d4af32ca7b4..4de03ce1534 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2807,9 +2807,12 @@ static bool ggml_metal_op_flash_attn_ext_use_dequant_f16(const ggml_tensor * op) assert(op->op == GGML_OP_FLASH_ATTN_EXT); switch (op->src[1]->type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: return true; - // extend to the other quantized KV types (q4_0, q4_1, q5_0, q5_1) as dequant kernels are added default: return false; } diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 29c8dd19024..5ad09bd4e9d 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -6344,6 +6344,10 @@ kernel void kernel_flash_attn_ext_dequant_to_f16( typedef decltype(kernel_flash_attn_ext_dequant_to_f16) kernel_flash_attn_ext_dequant_to_f16_t; +template [[host_name("kernel_flash_attn_ext_dequant_q4_0_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; +template [[host_name("kernel_flash_attn_ext_dequant_q4_1_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; +template [[host_name("kernel_flash_attn_ext_dequant_q5_0_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; +template [[host_name("kernel_flash_attn_ext_dequant_q5_1_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; template [[host_name("kernel_flash_attn_ext_dequant_q8_0_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]]; From cd644cf1344e49bb1f7fb7ea8d580eac42f7067f Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 19 Aug 2026 16:46:57 +0300 Subject: [PATCH 4/9] metal : skip the redundant V dequant when V is a view of K In MLA-based models, the V of the FA op is a view of K (the first ne20 elements of each K row); the dequantized V is then a view of the dequantized K, so skip the second dequant dispatch, do not reserve the V scratch region, and let the pad and attention kernels read V from the K F16 buffer with K's strides. The detection follows the CUDA backend: V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)) Also fix the FA pipeline getters: ns10/ns20 are function constants baked into the kernels and must be the actual K/V row widths as seen by the kernel. The dispatch now passes them explicitly (nb11_attn/nb10_attn, nb21_attn/nb20_attn) instead of the getters assuming contiguous F16 KV (ns20 = dv), which was wrong when V is read from K with K's row pitch (e.g. 576 vs 512). New test cases: 576/512 q8_0 (MLA shape, V is a view of K) at kv=113 (KV pad), nb=1 (vec) and nb=64 (non-vec). Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/src/ggml-metal/ggml-metal-device.cpp | 14 ++- ggml/src/ggml-metal/ggml-metal-device.h | 8 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 101 +++++++++++++++------- tests/test-backend-ops.cpp | 6 ++ 4 files changed, 89 insertions(+), 40 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 905f0ec1bd6..b27b12e3194 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1478,7 +1478,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( bool has_scap, bool has_kvpad, int32_t nsg, - bool use_f16_kv) { + bool use_f16_kv, + int32_t ns10, // actual row width of K in elements, as seen by the kernel + int32_t ns20) { // actual row width of V in elements, as seen by the kernel assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; @@ -1487,9 +1489,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( const int32_t dk = (int32_t) op->src[1]->ne[0]; const int32_t dv = (int32_t) op->src[2]->ne[0]; - const int32_t ns10 = use_f16_kv ? dk : op->src[1]->nb[1]/op->src[1]->nb[0]; - const int32_t ns20 = use_f16_kv ? dv : op->src[2]->nb[1]/op->src[2]->nb[0]; - const char * type = use_f16_kv ? "f16" : ggml_type_name(op->src[1]->type); // do bounds checks for the mask? @@ -1547,7 +1546,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v bool has_kvpad, int32_t nsg, int32_t nwg, - bool use_f16_kv) { + bool use_f16_kv, + int32_t ns10, // actual row width of K in elements, as seen by the kernel + int32_t ns20) { // actual row width of V in elements, as seen by the kernel assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; @@ -1556,9 +1557,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v const int32_t dk = (int32_t) op->src[1]->ne[0]; const int32_t dv = (int32_t) op->src[2]->ne[0]; - const int32_t ns10 = use_f16_kv ? dk : op->src[1]->nb[1]/op->src[1]->nb[0]; - const int32_t ns20 = use_f16_kv ? dv : op->src[2]->nb[1]/op->src[2]->nb[0]; - const char * type = use_f16_kv ? "f16" : ggml_type_name(op->src[1]->type); snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d", diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 26fd5760210..c4456e66fde 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -195,7 +195,9 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_scap, bool has_kvpad, int32_t nsg, - bool use_f16_kv); + bool use_f16_kv, + int32_t ns10, // actual row width of K in elements, as seen by the kernel + int32_t ns20); // actual row width of V in elements, as seen by the kernel struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec( ggml_metal_library_t lib, @@ -207,7 +209,9 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_kvpad, int32_t nsg, int32_t nwg, - bool use_f16_kv); + bool use_f16_kv, + int32_t ns10, // actual row width of K in elements, as seen by the kernel + int32_t ns20); // actual row width of V in elements, as seen by the kernel struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce( ggml_metal_library_t lib, diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 4de03ce1534..c2dd332ea9d 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2818,6 +2818,18 @@ static bool ggml_metal_op_flash_attn_ext_use_dequant_f16(const ggml_tensor * op) } } +// in some models (e.g. MLA-based), V is a view of K (the first ne20 elements of each K row); +// the dequantized V is then a view of the dequantized K and does not need its own dequant or scratch +// - ref: https://github.com/ggml-org/llama.cpp/pull/13435 +static bool ggml_metal_op_flash_attn_ext_v_is_view_of_k(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + const ggml_tensor * K = op->src[1]; + const ggml_tensor * V = op->src[2]; + + return V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)); +} + // size of the F16 dequantized K tensor; the dequantized V tensor follows it in the same scratch buffer static size_t ggml_metal_op_flash_attn_ext_dequant_f16_k_size(const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); @@ -2845,8 +2857,15 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) { const bool use_dequant_f16 = ggml_metal_op_flash_attn_ext_use_dequant_f16(op); // when the KV is dequantized to F16, the pad kernel copies the tail chunk from the F16 scratch buffer - const uint64_t nb11_pad = use_dequant_f16 ? sizeof(ggml_fp16_t)*ne10 : nb11; - const uint64_t nb21_pad = use_dequant_f16 ? sizeof(ggml_fp16_t)*ne20 : nb21; + // note: when V is a view of K, the dequantized V is read from the dequantized K with K's row stride + const bool v_is_view_of_k = use_dequant_f16 && ggml_metal_op_flash_attn_ext_v_is_view_of_k(op); + uint64_t nb11_pad = nb11; + uint64_t nb21_pad = nb21; + + if (use_dequant_f16) { + nb11_pad = sizeof(ggml_fp16_t)*ne10; + nb21_pad = sizeof(ggml_fp16_t)*(v_is_view_of_k ? ne10 : ne20); + } // note: the non-vec kernel requires more extra memory, so always reserve for it GGML_ASSERT(OP_FLASH_ATTN_EXT_NCPSG >= OP_FLASH_ATTN_EXT_VEC_NCPSG); @@ -2956,6 +2975,13 @@ size_t ggml_metal_op_flash_attn_ext_extra_dequant_f16(const ggml_tensor * op) { GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); const size_t k_size = ggml_metal_op_flash_attn_ext_dequant_f16_k_size(op); + + // when V is a view of K, the dequantized V is a view of the dequantized K + const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op); + if (v_is_view_of_k) { + return k_size; + } + const size_t v_size = GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne20*ne21*ne22*ne23, 16); return k_size + v_size; @@ -3055,12 +3081,11 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { if (use_dequant_f16) { assert(ggml_metal_op_flash_attn_ext_extra_dequant_f16(op) != 0); + const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op); + const int64_t nblocks1_64 = (ne10/ggml_blck_size(op->src[1]->type))*(int64_t) ne11*ne12*ne13; - const int64_t nblocks2_64 = (ne20/ggml_blck_size(op->src[2]->type))*(int64_t) ne21*ne22*ne23; GGML_ASSERT(nblocks1_64 <= INT32_MAX); - GGML_ASSERT(nblocks2_64 <= INT32_MAX); const int32_t nblocks1 = nblocks1_64; - const int32_t nblocks2 = nblocks2_64; ggml_metal_buffer_id bid_v_f16 = bid_dequant_f16; bid_v_f16.offs += ggml_metal_op_flash_attn_ext_dequant_f16_k_size(op); @@ -3088,41 +3113,57 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_dispatch_threadgroups(enc, (nblocks1 + nth - 1)/nth, 1, 1, nth, 1, 1); - // V - ggml_metal_kargs_flash_attn_ext_dequant_to_f16 args_v = { - /*.ne0 =*/ ne20, - /*.ne1 =*/ ne21, - /*.ne2 =*/ ne22, - /*.ne3 =*/ ne23, - /*.nb0 =*/ nb20, - /*.nb1 =*/ nb21, - /*.nb2 =*/ nb22, - /*.nb3 =*/ nb23, - /*.nblocks =*/ nblocks2, - }; + // V (skip when V is a view of K: the dequantized V is a view of the dequantized K) + if (!v_is_view_of_k) { + const int64_t nblocks2_64 = (ne20/ggml_blck_size(op->src[2]->type))*(int64_t) ne21*ne22*ne23; + GGML_ASSERT(nblocks2_64 <= INT32_MAX); + const int32_t nblocks2 = nblocks2_64; + + ggml_metal_kargs_flash_attn_ext_dequant_to_f16 args_v = { + /*.ne0 =*/ ne20, + /*.ne1 =*/ ne21, + /*.ne2 =*/ ne22, + /*.ne3 =*/ ne23, + /*.nb0 =*/ nb20, + /*.nb1 =*/ nb21, + /*.nb2 =*/ nb22, + /*.nb3 =*/ nb23, + /*.nblocks =*/ nblocks2, + }; - ggml_metal_encoder_set_pipeline(enc, pipeline0); - ggml_metal_encoder_set_bytes (enc, &args_v, sizeof(args_v), 0); - ggml_metal_encoder_set_buffer (enc, bid_src2, 1); - ggml_metal_encoder_set_buffer (enc, bid_v_f16, 2); + ggml_metal_encoder_set_pipeline(enc, pipeline0); + ggml_metal_encoder_set_bytes (enc, &args_v, sizeof(args_v), 0); + ggml_metal_encoder_set_buffer (enc, bid_src2, 1); + ggml_metal_encoder_set_buffer (enc, bid_v_f16, 2); - ggml_metal_encoder_dispatch_threadgroups(enc, (nblocks2 + nth - 1)/nth, 1, 1, nth, 1, 1); + ggml_metal_encoder_dispatch_threadgroups(enc, (nblocks2 + nth - 1)/nth, 1, 1, nth, 1, 1); + } // the pad and attention kernels read the dequantized KV ggml_metal_op_concurrency_reset(ctx); bid_k = bid_dequant_f16; - bid_v = bid_v_f16; + bid_v = v_is_view_of_k ? bid_k : bid_v_f16; - // contiguous F16 layout of the dequantized KV + // contiguous F16 layout of the dequantized K nb10_attn = sizeof(ggml_fp16_t); nb11_attn = nb10_attn*ne10; nb12_attn = nb11_attn*ne11; nb13_attn = nb12_attn*ne12; - nb20_attn = sizeof(ggml_fp16_t); - nb21_attn = nb20_attn*ne20; - nb22_attn = nb21_attn*ne21; - nb23_attn = nb22_attn*ne22; + + // if V is a view of K, the dequantized V is read from the dequantized K with K's strides + if (v_is_view_of_k) { + nb20_attn = nb10_attn; + nb21_attn = nb11_attn; + nb22_attn = nb12_attn; + nb23_attn = nb13_attn; + } else { + // contiguous F16 layout of the dequantized V + nb20_attn = sizeof(ggml_fp16_t); + nb21_attn = nb20_attn*ne20; + nb22_attn = nb21_attn*ne21; + nb23_attn = nb22_attn*ne22; + } } if (!ggml_metal_op_flash_attn_ext_use_vec(op)) { @@ -3275,7 +3316,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, use_dequant_f16); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, use_dequant_f16, nb11_attn/nb10_attn, nb21_attn/nb20_attn); ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); @@ -3413,7 +3454,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg, use_dequant_f16); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg, use_dequant_f16, nb11_attn/nb10_attn, nb21_attn/nb20_attn); GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 9254ba09d9c..081b8fad256 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9919,6 +9919,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1025, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // MLA shape (V is a view of K) with quantized KV - exercises the skipped-V dequant path + // (the test harness builds V as a view of K for this shape; see build_graph) + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). for (int64_t kv : { 4096, 16384 }) { From 34ad8da7ed9c5836304b9dc95d38f663ba4aa7a7 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 19 Aug 2026 16:58:24 +0300 Subject: [PATCH 5/9] test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B --- tests/test-backend-ops.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 081b8fad256..9da7ac80bda 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9910,8 +9910,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q2_0)); test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_F16)); - // q8_0 KV cases covering the Metal dequant->F16 FA path: decode and prompt batches, KV pad, - // permuted KV, feature flags, and long context + // q8_0 KV cases: decode and prompt batches, KV pad, permuted KV, feature flags, and long context test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); @@ -9919,7 +9918,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1025, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - // MLA shape (V is a view of K) with quantized KV - exercises the skipped-V dequant path + // MLA shape (V is a view of K) with quantized KV // (the test harness builds V as a view of K for this shape; see build_graph) test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); @@ -10310,7 +10309,7 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - // q8_0 KV dequant->F16 path (Metal) — long context, decode and prompt, with F16 references + // q8_0 KV cases with long context (decode and prompt) test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 128, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); From 275768fb5ac887fda5002f5869a42c763bffcdaf Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 19 Aug 2026 17:00:02 +0300 Subject: [PATCH 6/9] pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B --- .pi/gg/SYSTEM.md | 1 + 1 file changed, 1 insertion(+) diff --git a/.pi/gg/SYSTEM.md b/.pi/gg/SYSTEM.md index d39afbe0338..6a757c86942 100644 --- a/.pi/gg/SYSTEM.md +++ b/.pi/gg/SYSTEM.md @@ -9,6 +9,7 @@ General: Coding: - When in doubt, always refer to the CONTRIBUTING.md file of the project +- In `test-backend-ops.cpp`, do not mention specific backends (e.g. Metal, CUDA) in comments - When referencing issues or PRs in comments, use the format: - C/C++ code: `// ref: ` - Other (CMake, etc.): `# ref: ` From 101428bb32fe5881a01cd2deae1eee31c0978fc9 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 19 Aug 2026 17:05:45 +0300 Subject: [PATCH 7/9] metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/src/ggml-metal/ggml-metal-device.cpp | 12 +++--- ggml/src/ggml-metal/ggml-metal-device.h | 6 +-- ggml/src/ggml-metal/ggml-metal-impl.h | 2 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 46 +++++++++++------------ ggml/src/ggml-metal/ggml-metal-ops.h | 2 +- ggml/src/ggml-metal/ggml-metal.cpp | 2 +- ggml/src/ggml-metal/ggml-metal.metal | 16 ++++---- 7 files changed, 43 insertions(+), 43 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index b27b12e3194..fed6abe1c2a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1409,14 +1409,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_p return res; } -ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_dequant_to_f16( +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_kv_f16( ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; - snprintf(base, 256, "kernel_flash_attn_ext_dequant_%s_to_f16", ggml_type_name(op->src[1]->type)); + snprintf(base, 256, "kernel_flash_attn_ext_kv_%s_f16", ggml_type_name(op->src[1]->type)); ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base); if (!res.pipeline) { @@ -1478,7 +1478,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( bool has_scap, bool has_kvpad, int32_t nsg, - bool use_f16_kv, + bool use_kv_f16, int32_t ns10, // actual row width of K in elements, as seen by the kernel int32_t ns20) { // actual row width of V in elements, as seen by the kernel assert(op->op == GGML_OP_FLASH_ATTN_EXT); @@ -1489,7 +1489,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( const int32_t dk = (int32_t) op->src[1]->ne[0]; const int32_t dv = (int32_t) op->src[2]->ne[0]; - const char * type = use_f16_kv ? "f16" : ggml_type_name(op->src[1]->type); + const char * type = use_kv_f16 ? "f16" : ggml_type_name(op->src[1]->type); // do bounds checks for the mask? const bool bc_mask = op->src[3] && (op->src[3]->ne[1] % 8 != 0); @@ -1546,7 +1546,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v bool has_kvpad, int32_t nsg, int32_t nwg, - bool use_f16_kv, + bool use_kv_f16, int32_t ns10, // actual row width of K in elements, as seen by the kernel int32_t ns20) { // actual row width of V in elements, as seen by the kernel assert(op->op == GGML_OP_FLASH_ATTN_EXT); @@ -1557,7 +1557,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v const int32_t dk = (int32_t) op->src[1]->ne[0]; const int32_t dv = (int32_t) op->src[2]->ne[0]; - const char * type = use_f16_kv ? "f16" : ggml_type_name(op->src[1]->type); + const char * type = use_kv_f16 ? "f16" : ggml_type_name(op->src[1]->type); snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d", "flash_attn_ext_vec", diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index c4456e66fde..a6370f34659 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -176,7 +176,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_mask, int32_t ncpsg); -struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_dequant_to_f16( +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_kv_f16( ggml_metal_library_t lib, const struct ggml_tensor * op); @@ -195,7 +195,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_scap, bool has_kvpad, int32_t nsg, - bool use_f16_kv, + bool use_kv_f16, int32_t ns10, // actual row width of K in elements, as seen by the kernel int32_t ns20); // actual row width of V in elements, as seen by the kernel @@ -209,7 +209,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_kvpad, int32_t nsg, int32_t nwg, - bool use_f16_kv, + bool use_kv_f16, int32_t ns10, // actual row width of K in elements, as seen by the kernel int32_t ns20); // actual row width of V in elements, as seen by the kernel diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index caf6fc202d8..2bbbf0c6f15 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -353,7 +353,7 @@ typedef struct { uint64_t nb2; uint64_t nb3; int32_t nblocks; // number of blocks in the tensor -} ggml_metal_kargs_flash_attn_ext_dequant_to_f16; +} ggml_metal_kargs_flash_attn_ext_kv_f16; typedef struct { int32_t ne11; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index c2dd332ea9d..10e4fe81cca 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2803,7 +2803,7 @@ bool ggml_metal_op_flash_attn_ext_use_vec(const ggml_tensor * op) { // ref: https://github.com/ggml-org/llama.cpp/pull/25556 // dequantize the quantized KV cache to F16 before running the F16 flash attention kernels -static bool ggml_metal_op_flash_attn_ext_use_dequant_f16(const ggml_tensor * op) { +static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); switch (op->src[1]->type) { @@ -2831,7 +2831,7 @@ static bool ggml_metal_op_flash_attn_ext_v_is_view_of_k(const ggml_tensor * op) } // size of the F16 dequantized K tensor; the dequantized V tensor follows it in the same scratch buffer -static size_t ggml_metal_op_flash_attn_ext_dequant_f16_k_size(const ggml_tensor * op) { +static size_t ggml_metal_op_flash_attn_ext_kv_f16_k_size(const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); @@ -2854,15 +2854,15 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) { size_t res = 0; const bool has_mask = op->src[3] != nullptr; - const bool use_dequant_f16 = ggml_metal_op_flash_attn_ext_use_dequant_f16(op); + const bool use_kv_f16 = ggml_metal_op_flash_attn_ext_use_kv_f16(op); // when the KV is dequantized to F16, the pad kernel copies the tail chunk from the F16 scratch buffer // note: when V is a view of K, the dequantized V is read from the dequantized K with K's row stride - const bool v_is_view_of_k = use_dequant_f16 && ggml_metal_op_flash_attn_ext_v_is_view_of_k(op); + const bool v_is_view_of_k = use_kv_f16 && ggml_metal_op_flash_attn_ext_v_is_view_of_k(op); uint64_t nb11_pad = nb11; uint64_t nb21_pad = nb21; - if (use_dequant_f16) { + if (use_kv_f16) { nb11_pad = sizeof(ggml_fp16_t)*ne10; nb21_pad = sizeof(ggml_fp16_t)*(v_is_view_of_k ? ne10 : ne20); } @@ -2965,16 +2965,16 @@ size_t ggml_metal_op_flash_attn_ext_extra_tmp(const ggml_tensor * op) { return res; } -size_t ggml_metal_op_flash_attn_ext_extra_dequant_f16(const ggml_tensor * op) { +size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); - if (!ggml_metal_op_flash_attn_ext_use_dequant_f16(op)) { + if (!ggml_metal_op_flash_attn_ext_use_kv_f16(op)) { return 0; } GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); - const size_t k_size = ggml_metal_op_flash_attn_ext_dequant_f16_k_size(op); + const size_t k_size = ggml_metal_op_flash_attn_ext_kv_f16_k_size(op); // when V is a view of K, the dequantized V is a view of the dequantized K const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op); @@ -3061,10 +3061,10 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_buffer_id bid_tmp = bid_blk; bid_tmp.offs += ggml_metal_op_flash_attn_ext_extra_blk(op); - ggml_metal_buffer_id bid_dequant_f16 = bid_tmp; - bid_dequant_f16.offs += ggml_metal_op_flash_attn_ext_extra_tmp(op); + ggml_metal_buffer_id bid_kv_f16 = bid_tmp; + bid_kv_f16.offs += ggml_metal_op_flash_attn_ext_extra_tmp(op); - const bool use_dequant_f16 = ggml_metal_op_flash_attn_ext_use_dequant_f16(op); + const bool use_kv_f16 = ggml_metal_op_flash_attn_ext_use_kv_f16(op); ggml_metal_buffer_id bid_k = bid_src1; ggml_metal_buffer_id bid_v = bid_src2; @@ -3078,8 +3078,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { uint64_t nb22_attn = nb22; uint64_t nb23_attn = nb23; - if (use_dequant_f16) { - assert(ggml_metal_op_flash_attn_ext_extra_dequant_f16(op) != 0); + if (use_kv_f16) { + assert(ggml_metal_op_flash_attn_ext_extra_kv_f16(op) != 0); const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op); @@ -3087,14 +3087,14 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { GGML_ASSERT(nblocks1_64 <= INT32_MAX); const int32_t nblocks1 = nblocks1_64; - ggml_metal_buffer_id bid_v_f16 = bid_dequant_f16; - bid_v_f16.offs += ggml_metal_op_flash_attn_ext_dequant_f16_k_size(op); + ggml_metal_buffer_id bid_v_f16 = bid_kv_f16; + bid_v_f16.offs += ggml_metal_op_flash_attn_ext_kv_f16_k_size(op); - auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_dequant_to_f16(lib, op); + auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_kv_f16(lib, op); const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline0), 256); // K - ggml_metal_kargs_flash_attn_ext_dequant_to_f16 args_k = { + ggml_metal_kargs_flash_attn_ext_kv_f16 args_k = { /*.ne0 =*/ ne10, /*.ne1 =*/ ne11, /*.ne2 =*/ ne12, @@ -3109,7 +3109,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_pipeline(enc, pipeline0); ggml_metal_encoder_set_bytes (enc, &args_k, sizeof(args_k), 0); ggml_metal_encoder_set_buffer (enc, bid_src1, 1); - ggml_metal_encoder_set_buffer (enc, bid_dequant_f16, 2); + ggml_metal_encoder_set_buffer (enc, bid_kv_f16, 2); ggml_metal_encoder_dispatch_threadgroups(enc, (nblocks1 + nth - 1)/nth, 1, 1, nth, 1, 1); @@ -3119,7 +3119,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { GGML_ASSERT(nblocks2_64 <= INT32_MAX); const int32_t nblocks2 = nblocks2_64; - ggml_metal_kargs_flash_attn_ext_dequant_to_f16 args_v = { + ggml_metal_kargs_flash_attn_ext_kv_f16 args_v = { /*.ne0 =*/ ne20, /*.ne1 =*/ ne21, /*.ne2 =*/ ne22, @@ -3142,7 +3142,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { // the pad and attention kernels read the dequantized KV ggml_metal_op_concurrency_reset(ctx); - bid_k = bid_dequant_f16; + bid_k = bid_kv_f16; bid_v = v_is_view_of_k ? bid_k : bid_v_f16; // contiguous F16 layout of the dequantized K @@ -3250,7 +3250,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_op_concurrency_reset(ctx); } - const int is_q = !use_dequant_f16 && ggml_is_quantized(op->src[1]->type) ? 1 : 0; + const int is_q = !use_kv_f16 && ggml_is_quantized(op->src[1]->type) ? 1 : 0; // 2*(2*ncpsg) // ncpsg soft_max values + ncpsg mask values @@ -3316,7 +3316,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, use_dequant_f16, nb11_attn/nb10_attn, nb21_attn/nb20_attn); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, use_kv_f16, nb11_attn/nb10_attn, nb21_attn/nb20_attn); ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); @@ -3454,7 +3454,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg, use_dequant_f16, nb11_attn/nb10_attn, nb21_attn/nb20_attn); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg, use_kv_f16, nb11_attn/nb10_attn, nb21_attn/nb20_attn); GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 2b3d94a21f9..159a628d04a 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -42,7 +42,7 @@ bool ggml_metal_op_flash_attn_ext_use_vec(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_pad(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_blk(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op); -size_t ggml_metal_op_flash_attn_ext_extra_dequant_f16(const struct ggml_tensor * op); +size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op); int ggml_metal_op_concat (ggml_metal_op_t ctx, int idx); int ggml_metal_op_repeat (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 9b70d7f4e48..0e8d409e0b8 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -225,7 +225,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_ res += ggml_metal_op_flash_attn_ext_extra_pad(tensor); res += ggml_metal_op_flash_attn_ext_extra_blk(tensor); res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor); - res += ggml_metal_op_flash_attn_ext_extra_dequant_f16(tensor); + res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor); } break; case GGML_OP_CUMSUM: case GGML_OP_ARGSORT: diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 5ad09bd4e9d..976c54f2185 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -6310,8 +6310,8 @@ template < typename block_t, short QK, void (*deq_t4x4)(device const block_t *, short, thread float4x4 &)> -kernel void kernel_flash_attn_ext_dequant_to_f16( - constant ggml_metal_kargs_flash_attn_ext_dequant_to_f16 & args, +kernel void kernel_flash_attn_ext_kv_f16( + constant ggml_metal_kargs_flash_attn_ext_kv_f16 & args, device const char * x, device half * x_dst, uint gid [[thread_position_in_grid]]) { @@ -6342,13 +6342,13 @@ kernel void kernel_flash_attn_ext_dequant_to_f16( } } -typedef decltype(kernel_flash_attn_ext_dequant_to_f16) kernel_flash_attn_ext_dequant_to_f16_t; +typedef decltype(kernel_flash_attn_ext_kv_f16) kernel_flash_attn_ext_kv_f16_t; -template [[host_name("kernel_flash_attn_ext_dequant_q4_0_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; -template [[host_name("kernel_flash_attn_ext_dequant_q4_1_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; -template [[host_name("kernel_flash_attn_ext_dequant_q5_0_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; -template [[host_name("kernel_flash_attn_ext_dequant_q5_1_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; -template [[host_name("kernel_flash_attn_ext_dequant_q8_0_to_f16")]] kernel kernel_flash_attn_ext_dequant_to_f16_t kernel_flash_attn_ext_dequant_to_f16; +template [[host_name("kernel_flash_attn_ext_kv_q4_0_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16; +template [[host_name("kernel_flash_attn_ext_kv_q4_1_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16; +template [[host_name("kernel_flash_attn_ext_kv_q5_0_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16; +template [[host_name("kernel_flash_attn_ext_kv_q5_1_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16; +template [[host_name("kernel_flash_attn_ext_kv_q8_0_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16; constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]]; From ca1ca2fc1ab6164038992d896c572541c686f5d4 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 20 Aug 2026 10:23:35 +0300 Subject: [PATCH 8/9] cont : clean-up --- ggml/src/ggml-metal/ggml-metal-device.cpp | 8 ++++---- ggml/src/ggml-metal/ggml-metal-device.h | 8 ++++---- ggml/src/ggml-metal/ggml-metal-impl.h | 2 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 20 +++++++++++++------- ggml/src/ggml-metal/ggml-metal.metal | 3 ++- 5 files changed, 24 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index fed6abe1c2a..52043696eb5 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1479,8 +1479,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( bool has_kvpad, int32_t nsg, bool use_kv_f16, - int32_t ns10, // actual row width of K in elements, as seen by the kernel - int32_t ns20) { // actual row width of V in elements, as seen by the kernel + int32_t ns10, + int32_t ns20) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; @@ -1547,8 +1547,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v int32_t nsg, int32_t nwg, bool use_kv_f16, - int32_t ns10, // actual row width of K in elements, as seen by the kernel - int32_t ns20) { // actual row width of V in elements, as seen by the kernel + int32_t ns10, + int32_t ns20) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index a6370f34659..b7d46605882 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -196,8 +196,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_kvpad, int32_t nsg, bool use_kv_f16, - int32_t ns10, // actual row width of K in elements, as seen by the kernel - int32_t ns20); // actual row width of V in elements, as seen by the kernel + int32_t ns10, + int32_t ns20); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec( ggml_metal_library_t lib, @@ -210,8 +210,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att int32_t nsg, int32_t nwg, bool use_kv_f16, - int32_t ns10, // actual row width of K in elements, as seen by the kernel - int32_t ns20); // actual row width of V in elements, as seen by the kernel + int32_t ns10, + int32_t ns20); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce( ggml_metal_library_t lib, diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 2bbbf0c6f15..73a8c09c542 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -352,7 +352,7 @@ typedef struct { uint64_t nb1; uint64_t nb2; uint64_t nb3; - int32_t nblocks; // number of blocks in the tensor + int32_t nblocks; } ggml_metal_kargs_flash_attn_ext_kv_f16; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 10e4fe81cca..ef5b49f1068 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2801,7 +2801,7 @@ bool ggml_metal_op_flash_attn_ext_use_vec(const ggml_tensor * op) { return (ne01 < 20) && (ne00 % 32 == 0); } -// ref: https://github.com/ggml-org/llama.cpp/pull/25556 +// ref: https://github.com/ggml-org/llama.cpp/pull/27390 // dequantize the quantized KV cache to F16 before running the F16 flash attention kernels static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); @@ -3281,6 +3281,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { const size_t smem = FATTN_SMEM(nsg); + const int32_t ns10 = nb11_attn/nb10_attn; + const int32_t ns20 = nb21_attn/nb20_attn; + ggml_metal_kargs_flash_attn_ext args = { /*.ne01 =*/ ne01, /*.ne02 =*/ ne02, @@ -3291,11 +3294,11 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.ne11 =*/ ne11, /*.ne_12_2 =*/ ne12, /*.ne_12_3 =*/ ne13, - /*.ns10 =*/ int32_t(nb11_attn/nb10_attn), + /*.ns10 =*/ ns10, /*.nb11 =*/ nb11_attn, /*.nb12 =*/ nb12_attn, /*.nb13 =*/ nb13_attn, - /*.ns20 =*/ int32_t(nb21_attn/nb20_attn), + /*.ns20 =*/ ns20, /*.nb21 =*/ nb21_attn, /*.nb22 =*/ nb22_attn, /*.nb23 =*/ nb23_attn, @@ -3316,7 +3319,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, use_kv_f16, nb11_attn/nb10_attn, nb21_attn/nb20_attn); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, use_kv_f16, ns10, ns20); ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); @@ -3419,6 +3422,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { } } + const int32_t ns10 = nb11_attn/nb10_attn; + const int32_t ns20 = nb21_attn/nb20_attn; + ggml_metal_kargs_flash_attn_ext_vec args = { /*.ne01 =*/ ne01, /*.ne02 =*/ ne02, @@ -3429,11 +3435,11 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.ne11 =*/ ne11, /*.ne_12_2 =*/ ne12, /*.ne_12_3 =*/ ne13, - /*.ns10 =*/ int32_t(nb11_attn/nb10_attn), + /*.ns10 =*/ ns10, /*.nb11 =*/ nb11_attn, /*.nb12 =*/ nb12_attn, /*.nb13 =*/ nb13_attn, - /*.ns20 =*/ int32_t(nb21_attn/nb20_attn), + /*.ns20 =*/ ns20, /*.nb21 =*/ nb21_attn, /*.nb22 =*/ nb22_attn, /*.nb23 =*/ nb23_attn, @@ -3454,7 +3460,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg, use_kv_f16, nb11_attn/nb10_attn, nb21_attn/nb20_attn); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg, use_kv_f16, ns10, ns20); GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 976c54f2185..565d1e7437e 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -6305,7 +6305,8 @@ template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f // dequantize a quantized KV cache tensor to contiguous F16 before running the F16 flash attention kernels // - one thread per block; dispatched separately for K and V -// - ref: https://github.com/ggml-org/llama.cpp/pull/25556 +// - ref: https://github.com/ggml-org/llama.cpp/pull/27390 +// TODO: try to dispatch more threadgroups - one for each 16 elements, to see if perf improves template < typename block_t, short QK, From 7daa467fd490a9bb603077c154007251413e01e3 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 20 Aug 2026 11:35:25 +0300 Subject: [PATCH 9/9] cont : remove TODO [no ci] --- ggml/src/ggml-metal/ggml-metal.metal | 1 - 1 file changed, 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 565d1e7437e..f8840dde205 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -6306,7 +6306,6 @@ template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f // dequantize a quantized KV cache tensor to contiguous F16 before running the F16 flash attention kernels // - one thread per block; dispatched separately for K and V // - ref: https://github.com/ggml-org/llama.cpp/pull/27390 -// TODO: try to dispatch more threadgroups - one for each 16 elements, to see if perf improves template < typename block_t, short QK,