feat: support setting the attention impl in hf_transformers

Signed-off-by: Travis Johnson <[email protected]>

Author: tjohnson31415

server/text_generation_server/inference_engine/hf_transformers.py

@@ -24,6 +24,10 @@ def __init__(
             "trust_remote_code": TRUST_REMOTE_CODE,
         }
 
+        # TODO: consider if Flash Attention should be enabled based on FLASH_ATTENTION=True
+        if attn_impl := os.getenv("TRANSFORMERS_ATTN_IMPL"):
+            kwargs["attn_implementation"] = attn_impl
+
         if model_config.model_type == "mpt":
             model_config.init_device = str(self.device)
             kwargs["config"] = model_config