From f67021269cc5a09b4e0fa392c83226b571ec8025 Mon Sep 17 00:00:00 2001 From: Alberto Date: Sun, 12 Apr 2026 18:47:53 +0200 Subject: [PATCH 1/2] Fix flash_attention_3 detection and import for hopper wheel installs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit is_flash_attn_3_available() checked package distribution metadata for "flash-attn-3", but the hopper wheel (built from flash-attention/hopper/) doesn't register under that distribution name. Replace metadata check with actual import probing of both known module paths: 1. flash_attn_interface (standalone flash-attn-3 wheel) 2. hopper.flash_attn_interface (built from flash-attention/hopper/) Also fix _lazy_imports() to try both paths, and correct the compatibility matrix minimum CUDA version from 8 (Ampere) to 9 (Hopper) — FA3 kernels require sm90+. --- .../modeling_flash_attention_utils.py | 14 +++++++---- src/transformers/utils/import_utils.py | 24 +++++++++++++------ 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py index 9211ccb19a9e..874e5359d77a 100644 --- a/src/transformers/modeling_flash_attention_utils.py +++ b/src/transformers/modeling_flash_attention_utils.py @@ -93,10 +93,9 @@ def is_flash_attn_available(): 3: { "flash_attn_version": 3, "general_availability_check": is_flash_attn_3_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn_interface") is not None - and "flash-attn-3" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]], + "pkg_availability_check": is_flash_attn_3_available, # defers to actual import probing "supported_devices": ((is_torch_cuda_available, "cuda"),), - "cuda_min_major_version": 8, # Ampere + "cuda_min_major_version": 9, # Hopper (sm90+) — FA3 does NOT run on Ampere }, 4: { "flash_attn_version": 4, @@ -163,7 +162,14 @@ def _lazy_imports( from .integrations.npu_flash_attention import npu_flash_attn_with_kvcache as flash_attn_with_kvcache else: if implementation == "flash_attention_3" or (implementation is None and is_fa3 and not is_fa4): - from flash_attn_interface import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache + # FA3 Hopper kernels live under different module paths depending on + # install method. Try standalone wheel first, then hopper subpackage. + try: + from flash_attn_interface import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache + except ImportError: + from hopper.flash_attn_interface import flash_attn_func, flash_attn_varlen_func # type: ignore[no-redef] + + flash_attn_with_kvcache = None # hopper wheel may not expose this elif implementation == "flash_attention_4" or (implementation is None and is_fa4): from flash_attn.cute import flash_attn_func, flash_attn_varlen_func diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index 1e1ac2545f05..cf24b1e37bca 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -963,13 +963,23 @@ def is_flash_attn_2_available() -> bool: @lru_cache def is_flash_attn_3_available() -> bool: - # Universally available under `flash_attn_interface` - is_available = _is_package_available("flash_attn_interface")[0] - # Resolving and ensuring the proper name of FA3 being associated - is_available = is_available and "flash-attn-3" in [ - pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"] - ] - return is_available and is_torch_cuda_available() + if not is_torch_cuda_available(): + return False + + # FA3 Hopper kernels can live under different module paths depending on + # how they were installed: + # 1. "flash_attn_interface" — standalone flash-attn-3 wheel + # 2. "hopper.flash_attn_interface" — built from flash-attention/hopper/ + # The previous check only looked at package distribution metadata which + # doesn't match most real-world installs. Try the actual imports instead. + for mod_name in ("flash_attn_interface", "hopper.flash_attn_interface"): + try: + mod = __import__(mod_name, fromlist=["flash_attn_func"]) + if hasattr(mod, "flash_attn_func"): + return True + except Exception: + continue + return False @lru_cache From d4c3f29b38b419bbc55d2dea7296a520423d71db Mon Sep 17 00:00:00 2001 From: Alberto Date: Sun, 12 Apr 2026 19:19:31 +0200 Subject: [PATCH 2/2] Apply ruff formatting to compatibility matrix lambdas --- src/transformers/modeling_flash_attention_utils.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py index 874e5359d77a..7667efdb19a6 100644 --- a/src/transformers/modeling_flash_attention_utils.py +++ b/src/transformers/modeling_flash_attention_utils.py @@ -74,8 +74,10 @@ def is_flash_attn_available(): 2: { "flash_attn_version": 2, "general_availability_check": is_flash_attn_2_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None - and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], + "pkg_availability_check": lambda *args, **kwargs: ( + importlib.util.find_spec("flash_attn") is not None + and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]] + ), "supported_devices": ( (is_torch_cuda_available, "cuda"), (is_torch_mlu_available, "mlu"), @@ -100,8 +102,10 @@ def is_flash_attn_available(): 4: { "flash_attn_version": 4, "general_availability_check": is_flash_attn_4_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None - and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], + "pkg_availability_check": lambda *args, **kwargs: ( + importlib.util.find_spec("flash_attn") is not None + and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]] + ), "supported_devices": ((is_torch_cuda_available, "cuda"),), "cuda_min_major_version": 9, # Hopper },