Skip to content

Fix detection of flash-attention 2 and 4 #771

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
BrownianNotion wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix detection of flash-attention 2 and 4 #771

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
27bf233
Fix: fa2/fa4 detection
BrownianNotion Jul 15, 2026
23f6ad5
Fix
BrownianNotion Jul 15, 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
22 changes: 18 additions & 4 deletions src/olmo_core/nn/attention/flash_attn_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import importlib.metadata
import logging
from typing import Optional, Tuple

Expand All @@ -6,10 +7,16 @@

from .ring import RingAttentionLoadBalancerType

_flash_attn_pkgs = importlib.metadata.packages_distributions().get("flash_attn", [])

flash_attn_2 = None
try:
import flash_attn as flash_attn_2 # type: ignore
_is_flash_attn_2_available = "flash-attn" in [pkg.replace("_", "-") for pkg in _flash_attn_pkgs]

if _is_flash_attn_2_available:
import flash_attn as flash_attn_2 # type: ignore
Comment thread
BrownianNotion marked this conversation as resolved.
except ImportError:
flash_attn_2 = None # type: ignore
pass

try:
# NOTE: The flash-attn 3 API might be available under the name 'flash_attn_3.flash_attn_interface'
Expand All @@ -21,10 +28,17 @@
except ImportError:
flash_attn_3 = None # type: ignore

# match transformers fa4 detection - cute submodule available in both fa2/fa4
flash_attn_4 = None
try:
import flash_attn.cute as flash_attn_4 # type: ignore
_is_flash_attn_4_available = "flash-attn-4" in [
pkg.replace("_", "-") for pkg in _flash_attn_pkgs
]

if _is_flash_attn_4_available:
import flash_attn.cute as flash_attn_4 # type: ignore
except ImportError:
flash_attn_4 = None # type: ignore
pass

try:
import ring_flash_attn # type: ignore
Expand Down