From 3a2471811f7a91511b6907dd5050e2ba1de475e0 Mon Sep 17 00:00:00 2001 From: Andrei Date: Fri, 28 Jun 2024 16:07:47 -0400 Subject: [PATCH] Update src/llama.cpp Co-authored-by: slaren --- src/llama.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/llama.cpp b/src/llama.cpp index d2ee2d6ad..cb8ce3b8c 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -17402,11 +17402,12 @@ struct llama_context * llama_new_context_with_model( params.flash_attn = false; } - if (params.flash_attn && model->arch == LLM_ARCH_GEMMA2) { - LLAMA_LOG_WARN("%s: flash_attn is not compatible with Gemma2 - forcing off\n", __func__); + if (params.flash_attn && model->hparams.attn_soft_cap) { + LLAMA_LOG_WARN("%s: flash_attn is not compatible with attn_soft_cap - forcing off\n", __func__); params.flash_attn = false; } + if (params.flash_attn && model->hparams.n_embd_head_k != model->hparams.n_embd_head_v) { LLAMA_LOG_WARN("%s: flash_attn requires n_embd_head_k == n_embd_head_v - forcing off\n", __func__); params.flash_attn = false;