Skip to content

GitLab

Unverified Commit a5b226ce authored by Cyril Vallez's avatar Cyril Vallez Committed by GitHub
Browse files

Fix flash attention speed issue (#32028) 

Add the lru_cache for speed
parent a1844a32
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
src/transformers/utils/import_utils.py
View file @ a5b226ce
...@@ -820,6 +820,7 @@ def is_flash_attn_greater_or_equal_2_10(): ...@@ -820,6 +820,7 @@ def is_flash_attn_greater_or_equal_2_10():
return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.1.0") return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.1.0")
@lru_cache()
def is_flash_attn_greater_or_equal(library_version: str): def is_flash_attn_greater_or_equal(library_version: str):
if not _is_package_available("flash_attn"): if not _is_package_available("flash_attn"):
return False return False
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment