Fix flash attention speed issue (#32028)

Add the lru_cache for speed
This commit is contained in:
Cyril Vallez 2024-07-23 12:21:23 +02:00 committed by GitHub
parent a1844a3209
commit a5b226ce98
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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")
@lru_cache()
def is_flash_attn_greater_or_equal(library_version: str):
if not _is_package_available("flash_attn"):
return False