public_git
/
whisper.cpp
mirror of https://github.com/ggml-org/whisper.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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<block_t, QK, deq_t4x4>:
  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
This commit is contained in:
Georgi Gerganov 2026-08-20 13:43:59 +03:00
parent 02be8f5da8
commit 283775e0a5
7 changed files with 328 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
ggml/src/ggml-metal/ggml-metal-device.cpp
View File

@ -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);

14
ggml/src/ggml-metal/ggml-metal-device.h
View File

@ -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,

12
ggml/src/ggml-metal/ggml-metal-impl.h
View File

@ -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

269
ggml/src/ggml-metal/ggml-metal-ops.cpp
View File

@ -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);

1
ggml/src/ggml-metal/ggml-metal-ops.h
View File

@ -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);

1
ggml/src/ggml-metal/ggml-metal.cpp
View File

@ -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:

47
ggml/src/ggml-metal/ggml-metal.metal
View File

@ -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<block_q8_0, 32, dequantize_q8_0>) 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<block_q4_0, 32, dequantize_q4_0>;
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<block_q4_1, 32, dequantize_q4_1>;
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<block_q5_0, 32, dequantize_q5_0>;
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<block_q5_1, 32, dequantize_q5_1>;
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<block_q8_0, 32, dequantize_q8_0>;
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)]];