Skip to content

[Auto] Handle flash_attn distribution mapping KeyError (cluster-45520-3): merged 1 of 2 PRs #37

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
evalstate wants to merge 2 commits into
base: main
Choose a base branch
from
+4 −4
Open

[Auto] Handle flash_attn distribution mapping KeyError (cluster-45520-3): merged 1 of 2 PRs #37

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
22be6ec
utils: stop crashing with KeyError when flash_attn is importable but …
SAY-5 Apr 20, 2026
4bbd240
Merge commit 'refs/pull/45524/head' of https://github.com/huggingface…
evalstate Apr 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions src/transformers/utils/import_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -948,7 +948,7 @@ def is_flash_attn_2_available() -> bool:
is_available, flash_attn_version = _is_package_available("flash_attn", return_version=True)
# FA4 is also distributed under "flash_attn", hence we need to check the naming here
is_available = is_available and "flash-attn" in [
pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]
pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn", [])
]

if not is_available or not (is_torch_cuda_available() or is_torch_mlu_available()):
Expand All @@ -967,7 +967,7 @@ def is_flash_attn_3_available() -> bool:
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"]
pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn_interface", [])
]
return is_available and is_torch_cuda_available()

Expand All @@ -979,7 +979,7 @@ def is_flash_attn_4_available() -> bool:
# NOTE: FA2 seems to distribute the `cute` subdirectory even if only FA2 has been installed
# -> check for the proper (normalized) distribution name
is_available = is_available and "flash-attn-4" in [
pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]
pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn", [])
]

return is_available and is_torch_cuda_available()
Expand All @@ -990,7 +990,7 @@ def is_flash_attn_greater_or_equal(library_version: str) -> bool:
is_available, flash_attn_version = _is_package_available("flash_attn", return_version=True)
# FA4 is also distributed under "flash_attn", hence we need to check the naming here
is_available = is_available and "flash-attn" in [
pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]
pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn", [])
]

if not is_available:
Expand Down