ImportError `flash_attn_3` while setting `attn_implementation="flash_attention_3"` on hopper GPU
Author: allenphilipj
Pull request: no
Labels: name: bug

### System Info

```bash
Package                            Version
---------------------------------- ----------------------------
accelerate                         1.10.0
datasets                           4.0.0
flash-attn-3                       3.0.0b1+cu.12.8.torch.2.7
huggingface-hub                    0.34.4
transformers                       4.55.1
```

The check [here](https://github.com/huggingface/transformers/blob/50d2448a1a7b75354c3d0ca879afd124abd244ac/src/transformers/modeling_utils.py#L2468) is causing failures while setting `attn_implementation="flash_attention_3"` while loading the model.


### Who can help?

@Cyrilvallez 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I'm trying to use FA3 for my train task by loading the model as shown below:
```python
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_3",
)
```


I have FA3 installed:
```bash
$ uv pip list | grep flash
flash-attn-3                       3.0.0b1+cu.12.8.torch.2.7
```


I noticed that the actual FA3 implementation is mapped [here](https://github.com/huggingface/transformers/blob/53838edde77cb10f3a360150aa85a457637e9ac3/src/transformers/modeling_flash_attention_utils.py#L92C46-L92C61) and I'm able to confirm that `from flash_attn_interface import flash_attn_func, flash_attn_varlen_func` works on my venv. So I'm unable to understand the validation check for `flash_attn_3` module.

Interestingly, if I hack around the validation check for `flash_attn_3` module, the profiler shows it's correctly using the sm90 fa3 kernels.


### Expected behavior

Setting `attn_implementation="flash_attention_3"` should not fail due to `flash_attn_3` import error.

Comments:
This is weird, that means that it already fails here https://github.com/huggingface/transformers/blob/50d2448a1a7b75354c3d0ca879afd124abd244ac/src/transformers/utils/import_utils.py#L1136

Can you check if you can modify https://github.com/huggingface/transformers/blob/50d2448a1a7b75354c3d0ca879afd124abd244ac/src/transformers/utils/import_utils.py#L1140
to `if not _is_package_available("flash-attn-3"):` and if that works? Might be some name clashing based on the env?

## Flash Attention 3 Import Resolution - Complete Solution

Hi @allenphilipj,

This is a known issue with the Flash Attention 3 module detection logic in Transformers. The problem stems from a mismatch between the package name `flash-attn-3` (with hyphens) and the Python import module name `flash_attn_interface` (with underscores). Let me provide you with several solutions:

## Root Cause Analysis

The issue is in the `_is_package_available()` function which tries to import `flash_attn_3` as a module, but FA3 is actually structured differently:

```python
# What Transformers tries to check:
import flash_attn_3  # ❌ This module doesn't exist

# What actually works:
from flash_attn_interface import flash_attn_func  # ✅ This works
```

The FA3 package structure is:
- Package name: `flash-attn-3` 
- Import path: `flash_attn_interface`
- Kernels: `flash_attn_func`, `flash_attn_varlen_func`

## Solution 1: Quick Fix (Patch Import Utils)

As @vasqu suggested, you can modify the import check temporarily:

```python
# In transformers/src/transformers/utils/import_utils.py line ~1140
# Change from:
if not _is_package_available("flash_attn_3"):
    return False

# To:
if not _is_package_available("flash-attn-3"):
    return False

# Or better yet, check the actual interface:
try:
    from flash_attn_interface import flash_attn_func
    return True
except ImportError:
    return False
```

## Solution 2: Environment-Based Workaround

You can monkey-patch the function before loading your model:

```python
import transformers.utils.import_utils as import_utils

def custom_is_flash_attn_3_available():
    """Custom FA3 availability check"""
    try:
        from flash_attn_interface import flash_attn_func, flash_attn_varlen_func
        import flash_attn_3_cuda  # Check CUDA kernels too
        return True
    except ImportError:
        return False

# Override the original function
import_utils.is_flash_attn_3_available = custom_is_flash_attn_3_available

# Now load your model normally
from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_3",
)
```

## Solution 3: Dynamic Import Fix

For a more robust solution, replace the detection logic:

```python
def check_fa3_availability():
    """Comprehensive FA3 availability check"""
    checks = {
        'flash_attn_interface': False,
        'flash_attn_3_cuda': False,
        'sm90_support': False
    }
    
    try:
        # Check core interface
        from flash_attn_interface import flash_attn_func, flash_attn_varlen_func
        checks['flash_attn_interface'] = True
        
        # Check CUDA kernels
        import flash_attn_3_cuda
        checks['flash_attn_3_cuda'] = True
        
        # Check SM90 support (H100/H200)
        import torch
        if torch.cuda.is_available():
            device_capability = torch.cuda.get_device_capability()
            sm_version = device_capability[0] * 10 + device_capability[1]
            checks['sm90_support'] = sm_version >= 90
            
    except ImportError as e:
        print(f"FA3 import error: {e}")
        return False, checks
    
    all_available = all([checks['flash_attn_interface'], checks['flash_attn_3_cuda']])
    return all_available, checks

# Test before using
fa3_available, details = check_fa3_availability()
print(f"FA3 Available: {fa3_available}")
print(f"Details: {details}")
```

## Solution 4: Proper Fix for Transformers (Recommended)

This is what should be implemented in the Transformers codebase:

```python
# In utils/import_utils.py
def is_flash_attn_3_available() -> bool:
    """Check if Flash Attention 3 is available with proper interface detection"""
    try:
        # Check the actual interface that FA3 provides
        from flash_attn_interface import flash_attn_func, flash_attn_varlen_func
        
        # Verify CUDA kernels are available
        import flash_attn_3_cuda
        
        # Additional check for proper FA3 installation
        import importlib.util
        spec = importlib.util.find_spec("flash_attn_interface")
        if spec is None:
            return False
            
        return True
    except ImportError:
        return False
```

## Verification Script

To confirm everything works after applying the fix:

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Apply the fix first
import transformers.utils.import_utils as import_utils

def custom_is_flash_attn_3_available():
    try:
        from flash_attn_interface import flash_attn_func
        return True
    except ImportError:
        return False

import_utils.is_flash_attn_3_available = custom_is_flash_attn_3_available

# Test model loading
print(f"FA3 Available: {import_utils.is_flash_attn_3_available()}")

model = AutoModelForCausalLM.from_pretrained(
    "microsoft/DialoGPT-medium",  # Small test model
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_3",
    device_map="auto"
)

print("✅ Model loaded successfully with FA3!")
print(f"Model device: {next(model.parameters()).device}")
print(f"Model dtype: {next(model.parameters()).dtype}")

# Test inference to confirm FA3 kernels are used
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
inputs = tokenizer("Hello, how are you?", return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=10)
    
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"Generated: {response}")
```

## Performance Verification

To confirm FA3 kernels are actually being used:

```python
# Use the NVIDIA profiler to verify FA3 kernel usage
import torch.profiler as profiler

with profiler.profile(
    activities=[profiler.ProfilerActivity.CUDA],
    record_shapes=True,
    with_stack=True
) as prof:
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=10)

# Look for FA3-specific kernel names
fa3_kernels = [event for event in prof.events() 
              if 'flash_attn' in event.name and 'sm90' in event.name]
              
print(f"FA3 kernels found: {len(fa3_kernels)}")
for kernel in fa3_kernels[:3]:  # Show first 3
    print(f"  - {kernel.name}")
```

This solution addresses the core import detection issue and provides multiple approaches depending on your needs. The monkey-patch solution (Solution 2) is quickest for immediate use, while the proper fix (Solution 4) should be implemented upstream.

Let me know if you need clarification on any approach!

**Gabriel**

Thank you so much for the responses!

> This is weird, that means that it already fails here

Yes, I had previously used a version of Solution 2 suggested by @galafis and I was able to use FA3 kernels. If I try using `_is_package_available("flash-attn-3")`, it still shows up as `False` though I have FA3 installed. 

I think this is because the package name and the importable modules are different for FA3 pkg installation. As specified in the `top_level.txt`:
```txt
flash_attn_3_cuda
flash_attn_interface
```

I agree with @galafis that Solution 4 looks like the right approach.

Hmm, not really sure about the 4th solution as it looks more like a workaround tbh. Based on the `setup.py` in the fa repo https://github.com/Dao-AILab/flash-attention/blob/5059fd53e602bcc00336bb5cc8a85e50940485cb/hopper/setup.py#L36
there should be some associated name we should be able to find.

Can you check
1. With the latest transformers release? v4.55 had a few issues and refactors which should be stable in 4.56 now
2. What fa packages your system can find, e.g. using python3.10+ 
```python
from importlib.metadata import distributions

for dist in distributions():
    name = dist.metadata["Name"]
    if "flash" in name.lower():
        print(name)
```

1. I tried with `v4.56.2` and I'm still seeing the same error
2. The system is able to find `flash_attn_3` which matches with the METADATA in dist-info

Thanks for checking and bearing with me. This probably means that the package name `flash_attn_3` is not an importable module (since find_spec can't find it). Seems like we do need a workaround one way or another - I would still avoid using `flash_attn_3_cuda` --> I'm not sure how stable the API is there and whether it may change, the package name is more reliable imo.

Does something along the following work instead?
```python
@lru_cache
def is_flash_attn_3_available() -> bool:
    if not is_torch_available():
        return False

    # See https://github.com/huggingface/transformers/issues/41179
    # `_is_package_available` relies on the package to be an importable module, 
    # but fa3 only provides its package name as metadata, not importable. 
    # This is a workaround to check for the package's existence by its metadata.
    try:
        importlib.metadata.version("flash_attn_3")
    except importlib.metadata.PackageNotFoundError:
        return False

    import torch

    if not torch.cuda.is_available():
        return False

    # TODO: Check for a minimum version when FA3 is stable
    # return version.parse(importlib.metadata.version("flash_attn_3")) >= version.parse("3.0.0")

    return True
```

Edit: I installed the latest fa3 and I can find it in my env (without the workaround). Can you also share how you installed fa3, maybe even at which commit? This kinda looks like it uses a wheel in your scenario. As context, I installed via `git clone -> cd hopper -> python setup.py install`

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Any update fix plan in transformers  for this? I still encounter this in `4.57.1`

The problem is that I couldn't reproduce this, but I'm open for any fixes if someone where to submit a PR or similar

> The problem is that I couldn't reproduce this, but I'm open for any fixes if someone where to submit a PR or similar

Could you share the procedure you set up flash-attention3?

https://github.com/huggingface/transformers/issues/41179#issuecomment-3347845343 see the edit

Also on which version of transformers? I remember something being temporarily being broken on main, not sure if that got into the last v4 version 

I am using transformers==4.57.1. Also I tried torch 2.9 and 2.8, but both failed. Actually I use the same procedure as you provided. All under uv with python 3.12

It's very weird, I use python 3.11 with torch 2.6 iirc. 

What do you get for 
```python
import importlib

spec = importlib.util.find_spec("flash_attn_3")  # or "flash-attn-3" <-- not used by us but could be due to naming conventions?
```
The spec should be not `None`; otherwise it will explain why it cannot find it. Still open to any fixes, I just can't because I cannot recheck if what I did is correct to work on your systems