CUDA: fix Gemma 2 numerical issues for FA (#9166)

diff --git a/src/llama.cpp b/src/llama.cpp
index aeea54cffe02032bee08d3ed3d52b63e6f28550e..fc8fb3e0ddef2a87ca274fb676279678f7b67ecc 100644 (file)
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -8877,7 +8877,7 @@ static struct ggml_tensor * llm_build_kqv(
         cur = ggml_flash_attn_ext(ctx, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,
                                   hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
 
-        if (model.arch == LLM_ARCH_PHI2 || model.arch == LLM_ARCH_PHI3 || model.arch == LLM_ARCH_GPTNEOX) {
+        if (model.arch == LLM_ARCH_PHI2 || model.arch == LLM_ARCH_PHI3 || model.arch == LLM_ARCH_GPTNEOX || model.arch == LLM_ARCH_GEMMA2) {
             ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32);
         }