CS348Project/llama.cpp
5
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2025-10-27 08:21:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

metal : remove mask padding requirement

Browse Source
This commit is contained in:
Georgi Gerganov
2025-09-28 18:08:56 +03:00
parent 50d2b21d7c
commit 46c338fc6c
3 changed files with 18 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -994,19 +994,23 @@ ggml_metal_pipeline_t ggml_metal_library_get_pipeline_flash_attn_ext(
const int32_t ns10 = op->src[1]->nb[1]/op->src[1]->nb[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 ns20 = op->src[2]->nb[1]/op->src[2]->nb[0];
// 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", snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d",
"flash_attn_ext", "flash_attn_ext",
ggml_type_name(op->src[1]->type), ggml_type_name(op->src[1]->type),
dk, dk,
dv); dv);
snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_ns10=%d_ns20=%d_nsg=%d", snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d",
base, base,
has_mask, has_mask,
has_sinks, has_sinks,
has_bias, has_bias,
has_scap, has_scap,
has_kvpad, has_kvpad,
bc_mask,
ns10, ns10,
ns20, ns20,
nsg); nsg);
@@ -1024,6 +1028,8 @@ ggml_metal_pipeline_t ggml_metal_library_get_pipeline_flash_attn_ext(
ggml_metal_cv_set_bool(cv, has_scap, FC_FLASH_ATTN_EXT + 3); ggml_metal_cv_set_bool(cv, has_scap, FC_FLASH_ATTN_EXT + 3);
ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT + 4); ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT + 4);
ggml_metal_cv_set_bool(cv, bc_mask, FC_FLASH_ATTN_EXT + 10);
ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT + 20); ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT + 20);
ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT + 21); ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT + 21);
ggml_metal_cv_set_int32(cv, nsg, FC_FLASH_ATTN_EXT + 22); ggml_metal_cv_set_int32(cv, nsg, FC_FLASH_ATTN_EXT + 22);

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

@@ -1979,8 +1979,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
GGML_ASSERT(ne12 == ne22); GGML_ASSERT(ne12 == ne22);
GGML_ASSERT(!op->src[3] || op->src[3]->type == GGML_TYPE_F16); GGML_ASSERT(!op->src[3] || op->src[3]->type == GGML_TYPE_F16);
GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] >= GGML_PAD(op->src[0]->ne[1], 8) && GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] >= op->src[0]->ne[1] &&
"the Flash-Attention Metal kernel requires the mask to be padded to 8 and at least n_queries big"); "the Flash-Attention Metal kernel requires the mask to be at least n_queries big");
float scale; float scale;
float max_bias; float max_bias;

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

@@ -4492,6 +4492,8 @@ constant bool FC_flash_attn_ext_has_bias [[function_constant(FC_FLASH_ATTN_EXT
constant bool FC_flash_attn_ext_has_scap [[function_constant(FC_FLASH_ATTN_EXT + 3)]]; constant bool FC_flash_attn_ext_has_scap [[function_constant(FC_FLASH_ATTN_EXT + 3)]];
constant bool FC_flash_attn_ext_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT + 4)]]; constant bool FC_flash_attn_ext_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT + 4)]];
constant bool FC_flash_attn_ext_bc_mask [[function_constant(FC_FLASH_ATTN_EXT + 10)]];
//constant float FC_flash_attn_ext_scale [[function_constant(FC_FLASH_ATTN_EXT + 10)]]; //constant float FC_flash_attn_ext_scale [[function_constant(FC_FLASH_ATTN_EXT + 10)]];
//constant float FC_flash_attn_ext_max_bias [[function_constant(FC_FLASH_ATTN_EXT + 11)]]; //constant float FC_flash_attn_ext_max_bias [[function_constant(FC_FLASH_ATTN_EXT + 11)]];
//constant float FC_flash_attn_ext_logit_softcap [[function_constant(FC_FLASH_ATTN_EXT + 12)]]; //constant float FC_flash_attn_ext_logit_softcap [[function_constant(FC_FLASH_ATTN_EXT + 12)]];
@@ -4678,7 +4680,7 @@ void kernel_flash_attn_ext_impl(
v += (ikv2 + ikv3*args.ne_12_2)*args.nb21*C; v += (ikv2 + ikv3*args.ne_12_2)*args.nb21*C;
if (!FC_flash_attn_ext_has_mask) { if (!FC_flash_attn_ext_has_mask) {
threadgroup half * sm = (threadgroup half *) (sm2); threadgroup half * sm = (threadgroup half *) (sm2);
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
const short j = jj*NSG + sgitg; const short j = jj*NSG + sgitg;
@@ -4708,7 +4710,12 @@ void kernel_flash_attn_ext_impl(
FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
const short j = jj*NSG + sgitg; const short j = jj*NSG + sgitg;
sm2[j*SH + tiisg] = pm2[jj][tiisg]; if (FC_flash_attn_ext_bc_mask) {
sm2[j*SH + tiisg] = (iq1 + j) < args.ne31 ? pm2[jj][tiisg] : half2(-MAXHALF, -MAXHALF);
} else {
sm2[j*SH + tiisg] = pm2[jj][tiisg];
}
pm2[jj] += NW; pm2[jj] += NW;
} }