Feature request
Support flash-attn-4 (flash_attn.cute) in Transformers attention backend selection
System Info
transformers==5.3.0torch==2.10.0+cu128flash-attn-4==4.0.0b4accelerate==1.13.0trl==0.29.0peft==0.18.0deepspeed==0.18.7tokenizers==0.22.2huggingface_hub==1.6.0- Python 3.12
- CUDA 12.8
- GPU: NVIDIA Blackwell (
sm120)
Information
Reproduction
I am testing a Blackwell environment with:
- PyTorch 2.10
- CUDA 12.8
flash-attn-4transformers 5.3.0
Model loading fails before training starts when I pass:
from transformers import AutoModelForCausalLM
import torch
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)Actual behavior
Transformers appears to still expect the FlashAttention v2-style top-level import:
from flash_attn import flash_attn_func, flash_attn_varlen_func
But flash-attn-4 exposes the relevant API under:
from flash_attn.cute import flash_attn_func, flash_attn_varlen_func
As a result, model initialization fails with:
ImportError: cannot import name 'flash_attn_func' from 'flash_attn' (unknown location)
This is the traceback I get during training:
Traceback (most recent call last):
File "/home/joshua/llm-fine-tune-hf/main.py", line 68, in <module>
main()
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/main.py", line 94, in decorated_main
_run_hydra(
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
_run_app(
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/_internal/utils.py", line 457, in _run_app
run_and_report(
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
raise ex
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
return func()
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
lambda: hydra.run(
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/_internal/hydra.py", line 132, in run
_ = ret.return_value
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/core/utils.py", line 260, in return_value
raise self._return_value
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/hydra/core/utils.py", line 186, in run_job
ret.return_value = task_function(task_cfg)
File "/home/joshua/llm-fine-tune-hf/main.py", line 54, in main
return train(config)
File "/home/joshua/llm-fine-tune-hf/src/pipelines/pipeline.py", line 55, in train
model = setup.get_model()
File "/home/joshua/llm-fine-tune-hf/src/utils/setup.py", line 116, in get_model
model = AutoModelForCausalLM.from_pretrained(
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/transformers/models/auto/auto_factory.py", line 374, in from_pretrained
return model_class.from_pretrained(
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/transformers/modeling_utils.py", line 4094, in from_pretrained
model = cls(config, *model_args, **model_kwargs)
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 461, in __init__
super().__init__(config)
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/transformers/modeling_utils.py", line 1260, in __init__
self.config._attn_implementation_internal = self._check_and_adjust_attn_implementation(
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/transformers/modeling_utils.py", line 1893, in _check_and_adjust_attn_implementation
lazy_import_flash_attention(applicable_attn_implementation)
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py", line 171, in lazy_import_flash_attention
_flash_fn, _flash_varlen_fn, _pad_fn, _unpad_fn = _lazy_imports(
File "/home/joshua/anaconda3/envs/joshpp/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py", line 96, in _lazy_imports
from flash_attn import flash_attn_func, flash_attn_varlen_func
ImportError: cannot import name 'flash_attn_func' from 'flash_attn' (unknown location)
Expected behavior
One of the following would solve this cleanly:
- Detect
flash-attn-4 and import from flash_attn.cute when that package is installed.
- Introduce an explicit backend such as
attn_implementation="flash_attention_4".
- Document that
flash-attn-4 is not yet supported by the current attention backend selection logic.
Why this matters
Blackwell users moving to newer CUDA / PyTorch stacks are likely to try flash-attn-4, but the current import path fails before training begins. This makes the newer FA4 stack unusable from stock Transformers attention selection even though the FA4 functions are present and importable from flash_attn.cute.
Additional notes
In the same machine, an older stack works:
transformers==4.57.3flash_attn==2.8.3
That older stack exports flash_attn_func at the package top level, so the current Transformers import path works there.
By contrast, in the newer environment:
from flash_attn.cute import flash_attn_func, flash_attn_varlen_func
works, while:
from flash_attn import flash_attn_func
does not.
Motivation
I am trying to use a newer Blackwell training stack with torch==2.10.0+cu128, CUDA 12.8, and flash-attn-4==4.0.0b4.
At the moment, transformers==5.3.0 appears to assume the FlashAttention v2-style top-level API when attn_implementation="flash_attention_2" is selected. However, flash-attn-4 exposes its functions under flash_attn.cute instead of the older top-level import path. Because of this, model loading fails before training even starts with:
ImportError: cannot import name 'flash_attn_func' from 'flash_attn'
This makes the newer FA4 stack unusable from stock Transformers attention backend selection, even though the FA4 functions themselves are present and importable.
This seems related in spirit to earlier flash-attn compatibility/import issues, for example:
My main motivation is to make newer Blackwell-oriented FlashAttention stacks usable from Transformers without requiring users to patch library internals locally.
Your contribution
Yes, I can help with a PR.
I can test proposed changes on a local Blackwell environment using:
torch==2.10.0+cu128flash-attn-4==4.0.0b4transformers==5.3.0
If the maintainers agree on the intended direction, I can help with:
- validating a fix for
flash-attn-4 detection/import
- testing whether
flash_attn.cute can be supported safely
- verifying that the change does not break the existing
flash_attn v2 path
I have read the contribution guidance and can prepare a focused PR once the preferred approach is confirmed.
Feature request
Support
flash-attn-4(flash_attn.cute) in Transformers attention backend selectionSystem Info
transformers==5.3.0torch==2.10.0+cu128flash-attn-4==4.0.0b4accelerate==1.13.0trl==0.29.0peft==0.18.0deepspeed==0.18.7tokenizers==0.22.2huggingface_hub==1.6.0sm120)Information
Reproduction
I am testing a Blackwell environment with:
flash-attn-4transformers5.3.0Model loading fails before training starts when I pass:
Actual behavior
Transformers appears to still expect the FlashAttention v2-style top-level import:
But
flash-attn-4exposes the relevant API under:As a result, model initialization fails with:
This is the traceback I get during training:
Expected behavior
One of the following would solve this cleanly:
flash-attn-4and import fromflash_attn.cutewhen that package is installed.attn_implementation="flash_attention_4".flash-attn-4is not yet supported by the current attention backend selection logic.Why this matters
Blackwell users moving to newer CUDA / PyTorch stacks are likely to try
flash-attn-4, but the current import path fails before training begins. This makes the newer FA4 stack unusable from stock Transformers attention selection even though the FA4 functions are present and importable fromflash_attn.cute.Additional notes
In the same machine, an older stack works:
transformers==4.57.3flash_attn==2.8.3That older stack exports
flash_attn_funcat the package top level, so the current Transformers import path works there.By contrast, in the newer environment:
works, while:
does not.
Motivation
I am trying to use a newer Blackwell training stack with
torch==2.10.0+cu128, CUDA 12.8, andflash-attn-4==4.0.0b4.At the moment,
transformers==5.3.0appears to assume the FlashAttention v2-style top-level API whenattn_implementation="flash_attention_2"is selected. However,flash-attn-4exposes its functions underflash_attn.cuteinstead of the older top-level import path. Because of this, model loading fails before training even starts with:This makes the newer FA4 stack unusable from stock Transformers attention backend selection, even though the FA4 functions themselves are present and importable.
This seems related in spirit to earlier flash-attn compatibility/import issues, for example:
My main motivation is to make newer Blackwell-oriented FlashAttention stacks usable from Transformers without requiring users to patch library internals locally.
Your contribution
Yes, I can help with a PR.
I can test proposed changes on a local Blackwell environment using:
torch==2.10.0+cu128flash-attn-4==4.0.0b4transformers==5.3.0If the maintainers agree on the intended direction, I can help with:
flash-attn-4detection/importflash_attn.cutecan be supported safelyflash_attnv2 pathI have read the contribution guidance and can prepare a focused PR once the preferred approach is confirmed.