Enables Flash Attention in TGI for gemma models (#235)

huggingface · Oct 18, 2024 · c5e3881 · c5e3881
1 parent 9ae5ad5
commit c5e3881
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/server/text_generation_server/models/causal_lm.py b/server/text_generation_server/models/causal_lm.py
@@ -694,11 +694,13 @@ def __init__(
             "return_dict": True,
         }
 
-        if model.config.model_type in ["llama", "mistral", "starcoder2", "qwen2", "falcon"]:
+        if model.config.model_type in ["llama", "mistral", "starcoder2", "qwen2", "falcon", "gemma"]:
 
             if model.config.model_type not in ["falcon"]:
                 kwargs["attn_softmax_bf16"] = True
-            kwargs["trim_logits"] = True
+
+            if model.config.model_type not in ["gemma"]:
+                kwargs["trim_logits"] = True
 
             if os.getenv("USE_FLASH_ATTENTION", "false").lower() == "true":
                 kwargs["use_flash_attention"] = True