Skip to content

Metal: Optimize q8_0 KV #25556

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

Draft
tarruda wants to merge 3 commits into
base: master
Choose a base branch
from
Draft

Metal: Optimize q8_0 KV #25556

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a967948
metal: pack q8_0 loads for 256-wide KV
tarruda Jul 10, 2026
9cde680
metal: materialize q8_0 KV for high-GQA decode
tarruda Jul 10, 2026
793f141
metal: materialize q8_0 KV for deep prompts
tarruda Jul 11, 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
38 changes: 29 additions & 9 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 @@ -1339,6 +1339,18 @@ 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_q8_0_to_f16(
ggml_metal_library_t lib) {
const char * base = "kernel_flash_attn_ext_q8_0_to_f16";

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 @@ -1390,7 +1402,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];
Expand All @@ -1399,15 +1412,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);

Expand Down Expand Up @@ -1456,7 +1471,9 @@ 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,
int32_t nhptg,
bool use_f16_kv) {
assert(op->op == GGML_OP_FLASH_ATTN_EXT);

char base[256];
Expand All @@ -1465,12 +1482,15 @@ 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 * kernel = nhptg == 2 ? "flash_attn_ext_vec_gqa2" : "flash_attn_ext_vec";
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),
kernel,
type,
dk,
dv);

Expand Down
10 changes: 8 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 @@ -171,6 +171,9 @@ 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_q8_0_to_f16(
ggml_metal_library_t lib);

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 @@ -185,7 +188,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,
Expand All @@ -196,7 +200,9 @@ 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,
int32_t nhptg,
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,
Expand Down
16 changes: 16 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 @@ -330,6 +330,22 @@ typedef struct {
bool src2;
} ggml_metal_kargs_rope;

typedef struct {
int32_t ne10;
int32_t ne11;
int32_t ne12;
int32_t ne13;
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 nblocks;
} ggml_metal_kargs_flash_attn_ext_q8_0_to_f16;

typedef struct {
int32_t ne11;
int32_t ne_12_2; // assume K and V are same shape
Expand Down
Loading