Fix flash_attention_3 detection and import for hopper wheel installs

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"),
@@ -93,16 +95,17 @@ 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,
         "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
     },
@@ -163,7 +166,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
 

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