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

JohannesGaessler · arthw · commit 22d75c6bf957 · 2024-11-18T13:51:05.000+08:00
diff --git a/src/llama.cpp b/src/llama.cpp
@@ -8881,7 +8881,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);
         }
 

Original file line number	Diff line number	Diff line change
`@@ -8881,7 +8881,7 @@ static struct ggml_tensor * llm_build_kqv(`
`8881`	`8881`	`cur = ggml_flash_attn_ext(ctx, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,`
`8882`	`8882`	`hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);`
`8883`	`8883`
`8884`		`- if (model.arch == LLM_ARCH_PHI2 \|\| model.arch == LLM_ARCH_PHI3 \|\| model.arch == LLM_ARCH_GPTNEOX) {`
	`8884`	`+ if (model.arch == LLM_ARCH_PHI2 \|\| model.arch == LLM_ARCH_PHI3 \|\| model.arch == LLM_ARCH_GPTNEOX \|\| model.arch == LLM_ARCH_GEMMA2) {`
`8885`	`8885`	`ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32);`
`8886`	`8886`	`}`
`8887`	`8887`