Skip to content

metal : dequantize quantized KV to F16 before flash attention #27390

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ggerganov merged 9 commits into from
Aug 20, 2026
Merged

metal : dequantize quantized KV to F16 before flash attention #27390

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
5d6b00b
metal: dequantize q8_0 KV to f16 before flash attention
ggerganov Aug 19, 2026
ad615b6
metal : launch the FA KV dequant kernel separately for K and V
ggerganov Aug 19, 2026
74c13ee
metal : dequantize q4_0, q4_1, q5_0 and q5_1 KV to f16 before flash a…
ggerganov Aug 19, 2026
cd644cf
metal : skip the redundant V dequant when V is a view of K
ggerganov Aug 19, 2026
34ad8da
test : remove backend-specific wording from test-backend-ops comments
ggerganov Aug 19, 2026
275768f
pi : avoid backend mentions in test-backend-ops comments
ggerganov Aug 19, 2026
101428b
metal : rename the FA dequant_f16 identifiers to kv_f16
ggerganov Aug 19, 2026
ca1ca2f
cont : clean-up
ggerganov Aug 20, 2026
7daa467
cont : remove TODO
ggerganov Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .pi/gg/SYSTEM.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: <url>`
- Other (CMake, etc.): `# ref: <url>`
Expand Down
37 changes: 29 additions & 8 deletions ggml/src/ggml-metal/ggml-metal-device.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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];
Expand All @@ -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);

Expand Down Expand Up @@ -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];
Expand All @@ -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);

Expand Down
14 changes: 12 additions & 2 deletions ggml/src/ggml-metal/ggml-metal-device.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
12 changes: 12 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -343,6 +343,18 @@ typedef struct {
bool src2;
} 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
Expand Down
Loading