Skip to content

vulkan : dequant q8_0 KV once in coopmat1 #25494

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
0cc4m merged 9 commits into from
Aug 19, 2026
Merged

vulkan : dequant q8_0 KV once in coopmat1 #25494

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
484ad9b
vulkan : dequant q8_0 KV once in coopmat1
Nathanw1014 Jul 9, 2026
b805834
vulkan : fall back instead of aborting when FA scratch exceeds maxSto…
Nathanw1014 Jul 10, 2026
d7f6bce
vulkan : require KV-cache layout in FA dequant path
Nathanw1014 Jul 21, 2026
643e95e
vulkan : skip FA dequant path on coopmat2
Nathanw1014 Jul 21, 2026
c3b3bb9
tests : add contiguously-allocated quant K/V FA tests
Nathanw1014 Jul 21, 2026
dd39add
vulkan : trim comments
Nathanw1014 Aug 9, 2026
f9f99bf
vulkan : tighten permutation checks for FA path
Nathanw1014 Aug 10, 2026
500c887
vulkan : set prealloc_x_need_sync after the FA dispatch
Nathanw1014 Aug 10, 2026
b890946
vulkan : exclude Intel Xe1 from FA dequant path
Nathanw1014 Aug 16, 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
68 changes: 63 additions & 5 deletions ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -790,6 +790,7 @@ struct vk_device_struct {
vk_pipeline pipeline_quantize_q8_1_x4;

vk_pipeline pipeline_dequant[GGML_TYPE_COUNT];
vk_pipeline pipeline_dequant_transpose[GGML_TYPE_COUNT]; // fused dequant+transpose for FA quant-KV
vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols];
vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols];
vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT];
Expand Down Expand Up @@ -5000,6 +5001,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q8_0], "dequant_q8_0_transpose", dequant_q8_0_transpose_len, dequant_q8_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1);
Expand Down Expand Up @@ -10258,9 +10260,25 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx

const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16;

// dequant K/V once into an f16 scratch, reordered KV layout so FA can read without a stride
const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32;
const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32;
const bool use_dequant_kv = k_quant && v_quant && neq1 >= 64 &&
k->nb[0] == ggml_type_size(k->type) && v->nb[0] == ggml_type_size(v->type) &&
k->nb[1] >= k->nb[2] && v->nb[1] >= v->nb[2] &&

@jeffbolznv jeffbolznv Aug 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the shader is assuming a specific permutation, but this check is looser than that.

ggml_is_contiguously_allocated(k) && ggml_is_contiguously_allocated(v) &&
(uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange &&
(uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange &&
ctx->device->pipeline_dequant_transpose[k->type] != nullptr &&
ctx->device->pipeline_dequant_transpose[v->type] != nullptr &&
// coopmat2 already decodes quant K/V during the matrix load, so the f16 scratch is pure overhead there

@0cc4m 0cc4m Aug 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think coopmat2 does anything different here. This comment is just speculation. It should just say that the coopmat2 path does not benefit from this feature.

!ctx->device->coopmat2;

@0cc4m 0cc4m Aug 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please limit the feature to non-Intel and Intel+coopmat support here.

const ggml_type k_type_eff = use_dequant_kv ? GGML_TYPE_F16 : k->type;
const ggml_type v_type_eff = use_dequant_kv ? GGML_TYPE_F16 : v->type;

// For scalar/coopmat1 FA, we can use the "large" size to accommodate qga.
// For coopmat2 FA, we always use the small size (which is still pretty large for gqa).
vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k->type, v->type, f32acc);
vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k_type_eff, v_type_eff, f32acc);
const uint32_t max_gqa = std::min(tuning_params.block_rows, 32u);

if (N <= 8 && qk_ratio > 1 && qk_ratio <= max_gqa &&
Expand All @@ -10273,7 +10291,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
workgroups_y /= gqa_ratio;
}

tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k->type, v->type, f32acc);
tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc);

const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type));
uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type));
Expand All @@ -10287,6 +10305,17 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
v_stride /= 4;
}

uint32_t nbk2_eff = (uint32_t)nbk2, nbk3_eff = (uint32_t)nbk3;
uint32_t nbv2_eff = (uint32_t)nbv2, nbv3_eff = (uint32_t)nbv3;
if (use_dequant_kv) {
k_stride = HSK;
v_stride = HSV;
nbk2_eff = (uint32_t)((uint64_t)HSK * KV * sizeof(ggml_fp16_t));
nbk3_eff = (uint32_t)((uint64_t)HSK * KV * nek2 * sizeof(ggml_fp16_t));
nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t));
nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t));
}

const uint32_t alignment = tuning_params.block_cols;
bool aligned = (KV % alignment) == 0 &&
// the "aligned" shader variant will forcibly align strides, for performance
Expand All @@ -10313,7 +10342,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16
&& (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256);
vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc,
mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type);
mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff);

vk_pipeline pipeline = nullptr;

Expand Down Expand Up @@ -10417,6 +10446,35 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf;
vk_subbuffer mask_opt_buf = use_mask_opt ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf;

if (use_dequant_kv) {
const uint64_t fp = sizeof(ggml_fp16_t);
const uint64_t k_f16_sz = (uint64_t)ggml_nelements(k) * fp;
const uint64_t v_f16_sz = (uint64_t)ggml_nelements(v) * fp;
if (ctx->prealloc_size_x < k_f16_sz + v_f16_sz) {
ctx->prealloc_size_x = k_f16_sz + v_f16_sz;
ggml_vk_preallocate_buffers(ctx, subctx);
}
vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type];
vk_pipeline tr_v = ctx->device->pipeline_dequant_transpose[v->type];
ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1);
ggml_pipeline_request_descriptor_sets(ctx, tr_v, 1);
if (ctx->prealloc_x_need_sync) {
ggml_vk_sync_buffers(ctx, subctx);
}
vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0, k_f16_sz };
vk_subbuffer v_dst = vk_subbuffer{ ctx->prealloc_x, k_f16_sz, v_f16_sz };
const uint32_t k_nel = (uint32_t)ggml_nelements(k);
const uint32_t v_nel = (uint32_t)ggml_nelements(v);
{ const std::vector<uint32_t> pc = { (uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, 0, k_nel };
ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, pc, { k_nel, 1, 1 }); }
{ const std::vector<uint32_t> pc = { (uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, 0, v_nel };
ggml_vk_dispatch_pipeline(ctx, subctx, tr_v, { v_buf, v_dst }, pc, { v_nel, 1, 1 }); }
ggml_vk_sync_buffers(ctx, subctx);
ctx->prealloc_x_need_sync = true;

@jeffbolznv jeffbolznv Aug 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs to be set at the end of the function, since the contents of prealloc_x are used in the FA dispatch. Otherwise it can be clobbered in the next ggml_vk_sync_buffers. (Codex found this one)

k_buf = k_dst;
v_buf = v_dst;
}

uint32_t mask_n_head_log2 = ((sinks != nullptr) << 24) | n_head_log2;

if (use_mask_opt)
Expand Down Expand Up @@ -10446,8 +10504,8 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
(uint32_t)nev2, (uint32_t)nev3,
nem1, nem2, nem3,
q_stride, (uint32_t)nbq2, (uint32_t)nbq3,
k_stride, (uint32_t)nbk2, (uint32_t)nbk3,
v_stride, (uint32_t)nbv2, (uint32_t)nbv3,
k_stride, nbk2_eff, nbk3_eff,
v_stride, nbv2_eff, nbv3_eff,
scale, max_bias, logit_softcap,
mask_n_head_log2, m0, m1,
gqa_ratio, split_kv, split_k };
Expand Down
11 changes: 11 additions & 0 deletions ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,18 @@ void main() {
return;
}

#ifdef DEQUANT_TRANSPOSE
// read [HS, NH, KV, NS], write [HS, KV, NH, NS]
const uint HS = p.M, NH = p.K, KVn = p.stride_a;
const uint e0 = ib * 32;
const uint b_idx = (e0 % HS)
+ ((e0 / (HS * NH)) % KVn) * HS
+ ((e0 / HS) % NH) * (HS * KVn)
+ (e0 / (HS * NH * KVn)) * (HS * KVn * NH)
+ 16 * il;
#else
const uint b_idx = 1024*i + 32*ir + 16*il;
#endif

const float d = float(data_a[ib].d);

Expand Down
4 changes: 4 additions & 0 deletions ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -753,6 +753,10 @@ void process_shaders() {
if (tname != "f16" && tname != "bf16") {
string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}}));
}
// Fused dequant+transpose variant for FA quant-KV (per-head-contiguous f16 scratch).
if (tname == "q8_0") {
string_to_spv("dequant_" + tname + "_transpose", "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}, {"DEQUANT_TRANSPOSE", "1"}}));
}

shader = (tname == "f32" || tname == "f16" || tname == "bf16") ? "get_rows.comp" : "get_rows_quant.comp";

Expand Down
18 changes: 13 additions & 5 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6627,9 +6627,10 @@ struct test_flash_attn_ext : public test_case {
const ggml_type type_K;
const ggml_type type_V;
std::array<int32_t, 4> permute;
const bool kv_view; // create K/V as views of a larger buffer (like a KV cache)

std::string vars() override {
return VARS_TO_STR14(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute);
return VARS_TO_STR15(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view);
}

double max_nmse_err() override {
Expand All @@ -6645,9 +6646,10 @@ struct test_flash_attn_ext : public test_case {

test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array<int64_t, 2> nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8,
bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32,
ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array<int32_t, 4> permute = {0, 1, 2, 3})
ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array<int32_t, 4> permute = {0, 1, 2, 3},
bool kv_view = true)
: hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec),
type_K(type_K), type_V(type_V), permute(permute) {}
type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view) {}

ggml_tensor * build_graph(ggml_context * ctx) override {
const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K));
Expand Down Expand Up @@ -6675,7 +6677,7 @@ struct test_flash_attn_ext : public test_case {
ggml_tensor * q = create_permuted(GGML_TYPE_F32, hsk_padded, nb, nh*nr23[0], nr23[1], false);
ggml_set_name(q, "q");

ggml_tensor * k = create_permuted(type_K, hsk_padded, kv, nh, nr23[1], true); // the K tensor is usually a view of the K cache
ggml_tensor * k = create_permuted(type_K, hsk_padded, kv, nh, nr23[1], kv_view); // the K tensor is usually a view of the K cache
ggml_set_name(k, "k");

ggml_tensor * v = nullptr;
Expand All @@ -6689,7 +6691,7 @@ struct test_flash_attn_ext : public test_case {
// - https://github.com/ggml-org/llama.cpp/pull/18986
v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], 0);
} else {
v = create_permuted(type_V, hsv_padded, kv, nh, nr23[1], true); // the V tensor is usually a view of the V cache
v = create_permuted(type_V, hsv_padded, kv, nh, nr23[1], kv_view); // the V tensor is usually a view of the V cache
}
ggml_set_name(v, "v");

Expand Down Expand Up @@ -9280,6 +9282,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q1_0));
test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, GGML_TYPE_F16));

// dense-allocated (non-view) quant K/V at batch >= 64, in cache and native layouts
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {4, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false));

test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3}));
test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1}));
test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3}));
Expand Down