From 67fc9c5eb9da9d30a5a9869e869254d4fd35619d Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Mon, 6 Apr 2026 19:36:55 +0200 Subject: [PATCH] Fix Gemma4 quantized KV cache CPU FA performance (#1590) --- ggml/src/iqk/fa/iqk_fa_512_512.cpp | 137 ++++++----------------------- ggml/src/iqk/iqk_flash_attn.cpp | 3 + 2 files changed, 32 insertions(+), 108 deletions(-) diff --git a/ggml/src/iqk/fa/iqk_fa_512_512.cpp b/ggml/src/iqk/fa/iqk_fa_512_512.cpp index 6e2e7935..862a3a93 100644 --- a/ggml/src/iqk/fa/iqk_fa_512_512.cpp +++ b/ggml/src/iqk/fa/iqk_fa_512_512.cpp @@ -4,120 +4,41 @@ #include "iqk/fa/iqk_fa_templates.h" -namespace { - -template -inline void iqk_deepseek_helper(KHelper& kh, VHelper& vh, - int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, - const float * q, const char * mask, float scale, float softcap, float * qkv, - const float * sinkf, float * M, float * S) { - auto update = [&nq1, &mask, &q, &qkv, &M, &S, stride_q, stride_m, stride_qkv] (int n) { - nq1 -= n; - if (nq1 == 0) return true; - q += n*stride_q; - mask += n*stride_m; - qkv += n*stride_qkv; - if (M && S) { M += n; S += n; } - return false; - }; - if (nq1 >= 16) { - int n_step = nq1/16; - FlashAttn<512, 512, 16, step_k> fa(scale, softcap, sinkf); - fa.compute(kh, vh, 16*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); - if (update(16*n_step)) return; - } - if (nq1 >= 8) { - int n_step = nq1/8; - FlashAttn<512, 512, 8, step_k> fa(scale, softcap, sinkf); - fa.compute(kh, vh, 8*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); - if (update(8*n_step)) return; - } - if (nq1 >= 4) { - int n_step = nq1/4; - FlashAttn<512, 512, 4, step_k> fa(scale, softcap, sinkf); - fa.compute(kh, vh, 4*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); - if (update(4*n_step)) return; - } - if (nq1 == 3) { - FlashAttn<512, 512, 3, step_k> fa(scale, softcap, sinkf); - fa.compute(kh, vh, 3, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); - } - else if (nq1 == 2) { - FlashAttn<512, 512, 2, step_k> fa(scale, softcap, sinkf); - fa.compute(kh, vh, 2, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); - } else { - FlashAttn<512, 512, 1, step_k> fa(scale, softcap, sinkf); - fa.compute(kh, vh, 1, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); - } -} - -template -inline bool iqk_deepseek_helper(ggml_type type_k, - int nq1, int nk1, int stride_q, int stride_k, int stride_v, int stride_m, int stride_qkv, - const float * q, const char * k, const char * v, const char * mask, - float scale, float softcap, float * qkv, const float * sinkf, float * M, float * S) { - if (type_k == GGML_TYPE_Q8_0) { - HelperQ80 kh((const char *)k, stride_k); - HelperQ80 vh((const char *)v, stride_v); - iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); - return true; - } - if (type_k == GGML_TYPE_Q8_0_R8) { - HelperQ80R8<512> kh((const char *)k, stride_k); - HelperQ80 vh((const char *)v, stride_v); - iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); - return true; - } - if (type_k == GGML_TYPE_Q6_0) { - HelperQ60 kh((const char *)k, stride_k); - HelperQ60 vh((const char *)v, stride_v); - iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); - return true; - } -#if GGML_IQK_FA_ALL_QUANTS - if (type_k == GGML_TYPE_Q8_KV) { - HelperQ8KV<512> kh((const char *)k, stride_k); - HelperQ8KV<512> vh((const char *)v, stride_v); - iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); - return true; - } -#endif - if (type_k == GGML_TYPE_F16) { - HelperF16 kh((const char *)k, stride_k); - HelperF16 vh((const char *)v, stride_v); - iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); - return true; - } -#ifdef __AVX512BF16__ - if (type_k == GGML_TYPE_BF16) { - HelperBF16<512, step_k> kh((const char *)k, stride_k); - HelperBF16<512, step_k> vh((const char *)v, stride_v); - if (nq1 % 8 == 0) { - FlashAttnBF16<512, 512, 8, step_k> fa(scale, softcap, sinkf); - fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); - } else { - FlashAttnBF16<512, 512, 1, step_k> fa(scale, softcap, sinkf); - fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); - } - return true; - } -#endif - return false; -} - -} - IQK_FA_CASE(iqk_fa_512_512) { auto type_k = ggml_type(int_type_k); auto type_v = ggml_type(int_type_v); - if (!(type_k == type_v || (type_k == GGML_TYPE_Q8_0_R8 && type_v == GGML_TYPE_Q8_0))) { - return false; - } stride_q /= sizeof(float); // q stride as float - return iqk_deepseek_helper<32>(type_k, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, - q, (const char *)k, (const char *)v, (const char *)mask, scale, softcap, qkv, sinkf, M, S); + auto ck = (const char *)k; + auto cv = (const char *)v; + auto cm = (const char *)mask; + +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + if (type_v != GGML_TYPE_BF16) return false; // we do not support mixing bf16 k-cache with other types + if (nk%64 == 0) { + iqk_flash_helper_T<512, 512, 64>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } + iqk_flash_helper_T<512, 512, 32>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + + if (nk%128 == 0) { + return iqk_flash_helper_T<512, 512, 128>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + if (nk%64 == 0) { + return iqk_flash_helper_T<512, 512, 64>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + + return iqk_flash_helper_T<512, 512, 32>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); } diff --git a/ggml/src/iqk/iqk_flash_attn.cpp b/ggml/src/iqk/iqk_flash_attn.cpp index f30c2b80..02509095 100644 --- a/ggml/src/iqk/iqk_flash_attn.cpp +++ b/ggml/src/iqk/iqk_flash_attn.cpp @@ -215,6 +215,9 @@ extern "C" IQK_API bool iqk_flash_attn_noalibi(int type_q, int type_mask, float for (; last_k > first_k; --last_k) { if (umask[last_k-1] == 0) break; } + int non = 32*((last_k - first_k + 31)/32); + first_k = std::max(0, last_k - non); + last_k = std::min(first_k + non, nek1); //printf("nek1 = %d, first = %d, last = %d\n", nek1, first, last); if (last_k - first_k <= 3*nek1/4 && (last_k - first_k)%32 == 0) { //printf("Reducing from %d to %d\n", nek1, last_k - first_k);