From 283775e0a5deee4a212287b2858039385ab1dd2e Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 20 Aug 2026 13:43:59 +0300 Subject: [PATCH] metal : dequantize quantized KV to F16 before flash attention (llama/27390) * 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 * 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 * 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 * 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 * test : remove backend-specific wording from test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * pi : avoid backend mentions in test-backend-ops comments Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : rename the FA dequant_f16 identifiers to kv_f16 Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : clean-up * cont : remove TODO --- ggml/src/ggml-metal/ggml-metal-device.cpp | 37 ++- ggml/src/ggml-metal/ggml-metal-device.h | 14 +- ggml/src/ggml-metal/ggml-metal-impl.h | 12 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 269 ++++++++++++++++++---- ggml/src/ggml-metal/ggml-metal-ops.h | 1 + ggml/src/ggml-metal/ggml-metal.cpp | 1 + ggml/src/ggml-metal/ggml-metal.metal | 47 ++++ 7 files changed, 328 insertions(+), 53 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 953c75755..52043696e 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_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_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) { + 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,10 @@ 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_kv_f16, + int32_t ns10, + int32_t ns20) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; @@ -1469,15 +1489,14 @@ 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 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); 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 +1545,10 @@ 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_kv_f16, + int32_t ns10, + int32_t ns20) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); char base[256]; @@ -1535,12 +1557,11 @@ 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 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", - 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 7e1deeaa2..b7d466058 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_kv_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,10 @@ 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_kv_f16, + 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, @@ -201,7 +208,10 @@ 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_kv_f16, + 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 05ea7470e..f0b779979 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -345,6 +345,18 @@ typedef struct { bool inplace; } ggml_metal_kargs_rope; +typedef struct { + 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; +} ggml_metal_kargs_flash_attn_ext_kv_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 d8435e957..2dde14d8d 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2801,6 +2801,44 @@ 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/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); + + 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; + default: + return false; + } +} + +// 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_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); + + 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 +2854,18 @@ 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_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_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_kv_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); @@ -2828,8 +2878,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 +2888,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 +2965,28 @@ 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_kv_f16(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + 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_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); + 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; +} + int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -2989,6 +3061,111 @@ 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_kv_f16 = bid_tmp; + bid_kv_f16.offs += ggml_metal_op_flash_attn_ext_extra_tmp(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; + + 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_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); + + const int64_t nblocks1_64 = (ne10/ggml_blck_size(op->src[1]->type))*(int64_t) ne11*ne12*ne13; + GGML_ASSERT(nblocks1_64 <= INT32_MAX); + const int32_t nblocks1 = nblocks1_64; + + 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_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_kv_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, &args_k, sizeof(args_k), 0); + ggml_metal_encoder_set_buffer (enc, bid_src1, 1); + 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); + + // 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_kv_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, (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_kv_f16; + bid_v = v_is_view_of_k ? bid_k : bid_v_f16; + + // 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; + + // 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)) { // half8x8 kernel const int nqptg = OP_FLASH_ATTN_EXT_NQPSG; // queries per threadgroup @@ -3009,12 +3186,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 +3204,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 +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 = 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 @@ -3104,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, @@ -3114,14 +3294,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 =*/ ns10, + /*.nb11 =*/ nb11_attn, + /*.nb12 =*/ nb12_attn, + /*.nb13 =*/ nb13_attn, + /*.ns20 =*/ ns20, + /*.nb21 =*/ nb21_attn, + /*.nb22 =*/ nb22_attn, + /*.nb23 =*/ nb23_attn, /*.ne31 =*/ ne31, /*.ne32 =*/ ne32, /*.ne33 =*/ ne33, @@ -3139,13 +3319,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_kv_f16, ns10, ns20); 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 +3357,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 +3375,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); @@ -3242,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, @@ -3252,14 +3435,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 =*/ ns10, + /*.nb11 =*/ nb11_attn, + /*.nb12 =*/ nb12_attn, + /*.nb13 =*/ nb13_attn, + /*.ns20 =*/ ns20, + /*.nb21 =*/ nb21_attn, + /*.nb22 =*/ nb22_attn, + /*.nb23 =*/ nb23_attn, /*.ne31 =*/ ne31, /*.ne32 =*/ ne32, /*.ne33 =*/ ne33, @@ -3277,15 +3460,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_kv_f16, ns10, ns20); 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 b03b59e0b..159a628d0 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_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 ef3c92f27..0e8d409e0 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_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 0537fa4cf..949931c8d 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -6318,6 +6318,53 @@ 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 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 +template < + typename block_t, + short QK, + void (*deq_t4x4)(device const block_t *, short, thread float4x4 &)> +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]]) { + if (gid >= (uint) args.nblocks) { + return; + } + + 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*args.nb0 + i1*args.nb1 + i2*args.nb2 + i3*args.nb3; + + 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; + 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_kv_f16) kernel_flash_attn_ext_kv_f16_t; + +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)]]; constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]];