Commit d2b3eb9b6 for llama.cpp
commit d2b3eb9b6030d1cb6a2c05fa07cc71b85da69f34
Author: Jeff Bolz <jbolz@nvidia.com>
Date: Wed Nov 11 04:08:06 2026 -0600
vulkan: allow dequantized KV optimization with CM2 (#31526)
* additional FA test cases
Assisted-by: Codex
* vulkan: allow dequantized KV optimization with CM2
Enable the KV dequantization optimization with coopmat2, guarded by a heuristic
based on query tokens*GQA ratio.
While doing this we saw very large speedups for certain head sizes, and it was
because those head sizes were incorrectly treated as unaligned. So also fix
those to be treated as aligned, which should help for when the dequant
optimization doesn't apply.
* gqa uses nbq2 for the stride, so we need to check that value for alignment
* tests : use VARS_TO_STR18 for flash attention
Assisted-by: Codex
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 749f57e72..502f73344 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -8290,6 +8290,10 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
const size_t kv_scratch_ts = kv_raw_strided ? ggml_type_size(k->type) : sizeof(ggml_fp16_t);
// K and V share one scratch buffer, so the sum must fit in one allocation
const uint64_t dequant_kv_sz = ((uint64_t)ggml_nelements(k) + (uint64_t)ggml_nelements(v)) * kv_scratch_ts;
+ const uint32_t cm2_min_queries = KV >= 8192 ? 512 : 1024;
+ const uint32_t cm2_min_reuse = KV >= 8192 ? 4096 : 8192;
+ // Query reuse must cover the conversion and synchronization cost.
+ const bool cm2_dequant_kv = KV >= 4096 && neq1 >= cm2_min_queries && (uint64_t) neq1 * qk_ratio >= cm2_min_reuse;
const bool use_dequant_kv = ((k_quant && v_quant) || kv_raw_strided) && neq1 >= 64 &&
is_dense_kv_cache(k) && is_dense_kv_cache(v) &&
(uint64_t)ggml_nelements(k) * kv_scratch_ts <= ctx->device->properties.limits.maxStorageBufferRange &&
@@ -8299,11 +8303,10 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
(kv_raw_strided ||
(ctx->device->pipeline_dequant_transpose[k->type] != nullptr &&
ctx->device->pipeline_dequant_transpose[v->type] != nullptr)) &&
- // coopmat2 path does not benefit from the f16 scratch
- !ctx->device->coopmat2 &&
// Intel Xe1 regresses, see PR 25494
(ctx->device->vendor_id != VK_VENDOR_ID_INTEL ||
- (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1));
+ (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)) &&
+ (!ctx->device->coopmat2 || cm2_dequant_kv);
const ggml_type k_type_eff = (use_dequant_kv && !kv_raw_strided) ? GGML_TYPE_F16 : k->type;
const ggml_type v_type_eff = (use_dequant_kv && !kv_raw_strided) ? GGML_TYPE_F16 : v->type;
@@ -8356,6 +8359,7 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
(gqa_ratio > 1 || (tuning_params.path == FA_SCALAR && N == 1));
const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type));
+ const uint32_t q_stride_eff = gqa_ratio > 1 ? (uint32_t)(nbq2 / ggml_type_size(q->type)) : q_stride;
uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type));
uint32_t v_stride = (uint32_t)(nbv1 / ggml_type_size(v->type));
@@ -8378,9 +8382,13 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
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
- (q_stride & 7) == 0 && (k_stride & 7) == 0 && (v_stride & 7) == 0;
+ // CM2 decoded K/V strides count whole blocks and do not need to be multiples of eight blocks.
+ const bool cm2_decode_k = tuning_params.path == FA_COOPMAT2 &&
+ k_type_eff != GGML_TYPE_F16 && k_type_eff != GGML_TYPE_BF16;
+ const bool cm2_decode_v = tuning_params.path == FA_COOPMAT2 &&
+ v_type_eff != GGML_TYPE_F16 && v_type_eff != GGML_TYPE_BF16;
+ bool aligned = (KV % alignment) == 0 && (q_stride_eff & 7) == 0 &&
+ (cm2_decode_k || (k_stride & 7) == 0) && (cm2_decode_v || (v_stride & 7) == 0);
// Need to use the coopmat2 variant that clamps loads when HSK/HSV aren't sufficiently aligned.
if (((HSK | HSV) % 16) != 0 && tuning_params.path == FA_COOPMAT2) {
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index c8b1d02e4..4ab1525d0 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -493,6 +493,7 @@ static std::string var_to_str(ggml_scale_mode mode) {
#define VARS_TO_STR15(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o) VAR_TO_STR(a) + "," + VARS_TO_STR14(b, c, d, e, f, g, h, i, j, k, l, m, n, o)
#define VARS_TO_STR16(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) VAR_TO_STR(a) + "," + VARS_TO_STR15(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p)
#define VARS_TO_STR17(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q) VAR_TO_STR(a) + "," + VARS_TO_STR16(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q)
+#define VARS_TO_STR18(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r) VAR_TO_STR(a) + "," + VARS_TO_STR17(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r)
// accept FLT_MAX as infinity
static bool isinf_or_max(float f) {
@@ -8994,9 +8995,10 @@ struct test_flash_attn_ext : public test_case {
const bool kv_view; // create K/V as views of a larger buffer (like a KV cache)
const bool v_is_view_of_k;
const int64_t n_kv_max;
+ const int64_t q_head_padding;
std::string vars() override {
- return VARS_TO_STR17(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k, n_kv_max);
+ return VARS_TO_STR18(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k, n_kv_max, q_head_padding);
}
double max_nmse_err() override {
@@ -9013,9 +9015,9 @@ 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},
- bool kv_view = true, bool v_is_view_of_k = false, int64_t n_kv_max = 0)
+ bool kv_view = true, bool v_is_view_of_k = false, int64_t n_kv_max = 0, int64_t q_head_padding = 0)
: 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), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k), n_kv_max(n_kv_max) {}
+ type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k), n_kv_max(n_kv_max), q_head_padding(q_head_padding) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K));
@@ -9040,7 +9042,11 @@ struct test_flash_attn_ext : public test_case {
return t;
};
- ggml_tensor * q = create_permuted(GGML_TYPE_F32, hsk_padded, nb, nh*nr23[0], nr23[1], false);
+ ggml_tensor * q = create_permuted(GGML_TYPE_F32, hsk_padded + q_head_padding, nb, nh*nr23[0], nr23[1], false);
+ if (q_head_padding > 0) {
+ ggml_set_name(q, "q_padded");
+ q = ggml_view_4d(ctx, q, hsk_padded, q->ne[1], q->ne[2], q->ne[3], q->nb[1], q->nb[2], q->nb[3], 0);
+ }
ggml_set_name(q, "q");
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
@@ -12961,6 +12967,18 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1025, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}));
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0));
+ // Dense Q8 KV pre-dequantization: both query thresholds and a partial final tile.
+ test_cases.emplace_back(new test_flash_attn_ext(64, 64, 1, {8, 1}, 4096, 1024, 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, 1, {8, 1}, 8192, 512, 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, 1, {8, 1}, 8193, 513, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
+
+ // Grouped decode uses Q's head stride, which is not eight-element aligned here.
+ for (int nb : {1, 3}) {
+ for (ggml_type type_KV : {GGML_TYPE_F16, GGML_TYPE_Q8_0}) {
+ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, {2, 1}, 512, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV, {0, 2, 1, 3}, false, false, 0, 4));
+ }
+ }
+
// MLA shape: the V cache is a sub-view of the K cache, with quantized KV
test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {8, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true));
test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {8, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true));
@@ -13718,6 +13736,21 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 8192, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));
test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 8192, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));
+ test_cases.emplace_back(new test_flash_attn_ext(96, 96, 2, {8, 1}, 8192, 128, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
+ for (int64_t kv : {8192, 32768}) {
+ for (int64_t nq : {128, 256, 512, 1024}) {
+ test_cases.emplace_back(new test_flash_attn_ext(192, 192, 2, {8, 1}, kv, nq, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
+ }
+ }
+ // MLA uses V as a view of the compressed K cache.
+ test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 32768, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false, true));
+ for (int64_t nq : {256, 512, 1024, 2048}) {
+ test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {8, 1}, 32768, nq, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
+ }
+ for (int64_t nq : {512, 1024}) {
+ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 2, {16, 1}, 4096, nq, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
+ }
+
// q8_0 KV cases with long context (decode and prompt)
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 128, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0));
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0));