Repository navigation
Open to contribution: adding torch.nn.functional.scaled_dot_product_attention support for more architectures #28005
Description
Feature request
In Transformers 4.36, we started adding native support of torch.nn.functional.scaled_dot_product_attention (SDPA), enabled by default in Transformers: https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-and-memory-efficient-attention-through-pytorchs-scaleddotproductattention
SDPA allows to dispatch to memory-efficient attention, flash attention on supported GPUs (currently NVIDIA-only), and even on Intel CPUs.
For the record, here's a benchmark on some currently supported models:
Training benchmark, run on A100-SXM4-80GB.
| Model | Batch size | Sequence length | Time per batch ("eager", s) |
Time per batch ("sdpa", s) |
Speedup | Peak memory ("eager", MB) |
Peak memory ("sdpa", MB) |
Memory savings |
|---|---|---|---|---|---|---|---|---|
| llama2 7b | 4 | 1024 | 1.065 | 0.90 | 19.4% | 73878.28 | 45977.81 | 60.7% |
| llama2 7b | 4 | 2048 | OOM | 1.87 | / | OOM | 78394.58 | SDPA does not OOM |
| llama2 7b | 1 | 2048 | 0.64 | 0.48 | 32.0% | 55557.01 | 29795.63 | 86.4% |
| llama2 7b | 1 | 3072 | OOM | 0.75 | / | OOM | 37916.08 | SDPA does not OOM |
| llama2 7b | 1 | 4096 | OOM | 1.03 | / | OOM | 46028.14 | SDPA does not OOM |
| llama2 7b | 2 | 4096 | OOM | 2.05 | / | OOM | 78428.14 | SDPA does not OOM |
Inference benchmark, run on A100-SXM4-80GB.
| Model | Batch size | Prompt length | Num new tokens | Per token latency "eager" (ms) |
Per token latency "sdpa" (ms) |
Speedup |
|---|---|---|---|---|---|---|
| llama2 13b | 1 | 1024 | 1 (prefill) | 178.66 | 159.36 | 12.11% |
| llama2 13b | 1 | 100 | 100 | 40.35 | 37.62 | 7.28% |
| llama2 13b | 8 | 100 | 100 | 40.55 | 38.06 | 6.53% |
| Whisper v3 large | 1 | / | 62 | 20.05 | 18.90 | 6.10% |
| Whisper v3 large | 8 | / | 77 | 25.42 | 24.77 | 2.59% |
| Whisper v3 large | 16 | / | 77 | 28.51 | 26.32 | 8.34% |
Previously, we had a partial support of SDPA in Optimum BetterTransformer but we are now looking to slowly deprecate it in favor of upstream support of SDPA directly in Transformers.
Here are the architectures for which support has been requested:
- Codegen (BetterTransformer not supporting CodeGen2 optimum#1050)
- LLAVA (Can optimum.bettertransformer supports LLAVA model? optimum#1592)
- Marian ( BetterTransforer not Support Marian optimum#1142)
- Mistral (Add support for mistral type Model to use Mistral and Zephyr optimum#1553)
- LongT5 (longT5 BetterTransformer implementation optimum#1506)
- ViT (Add support for mistral type Model to use Mistral and Zephyr optimum#1553)
The integration could take inspiration from https://github.com/huggingface/optimum/blob/main/optimum/bettertransformer/models/decoder_models.py & https://github.com/huggingface/optimum/blob/main/optimum/bettertransformer/models/attention.py
Motivation
Faster training & inference, lower memory requirement
Your contribution
I may work on some at some point, but contributions are most welcome.
You should refer to #26572 to add the support of SDPA for a model, roughly following these steps:
- Create a
XxxSdpaAttentionclass inheriting fromXxxAttentionand implement the attention logic using SDPA - Use
_prepare_4d_causal_attention_mask_for_sdpainstead of_prepare_4d_causal_attention_maskfor SDPA - Use
_prepare_4d_attention_mask_for_sdpainstead of_prepare_4d_attention_maskfor SDPA - Add
_supports_sdpa = TruetoXxxPreTrainedModel - Add
"sdpa"key toXXX_ATTENTION_CLASSESin the model modeling file
Activity
ENate commentedon Dec 14, 2023
Hi @fxmarty I can take a look at this issue. Of I can ask questions if necessary. Or has anyone taken it already?
davidan5 commentedon Dec 18, 2023
does someone know if longT5 and all T5 models are blocked by bias support in flash attention ?
Dao-AILab/flash-attention#617
ENate commentedon Dec 19, 2023
Hi @davidan5 are you working on the implementation?
davidan5 commentedon Dec 19, 2023
@ENate I was trying to understand the status and have an estimation of the code change to see if I can contribute.
ENate commentedon Dec 19, 2023
I see.
hackyon commentedon Jan 29, 2024
I'm interested in taking a look at this for the Mistral model if that's still needed. Otherwise, please let me know if there are any other models that still need some work. Thanks
ENate commentedon Jan 29, 2024
Is LongT5 still open?
ArthurZucker commentedon Jan 30, 2024
Mistral is already covered! LongT5 if it is like T5 and has attention bias that might not be supported
hackyon commentedon Jan 30, 2024
Oh yea, looks like you added support for Mistral/Mixtral last month.
It doesn't seem to be supported for BERT yet (I think someone else is working on FA2 but not SDPA), so I'll take a crack at it. It looks like there is a config for relative position embeddings for BERT, so I'll just have it fallback to the original attention for configs using relative position embeddings.
@ArthurZucker - Please let me know if you know if someone else is already working on SDPA for BERT and I can look for something else to do. Thanks!
ArthurZucker commentedon Jan 31, 2024
Not sure anyone is working on that but bert is already so small that I doubt it will have a lot of impact on perf!
abdulfatir commentedon Mar 31, 2024
@ArthurZucker for the T5 family of models, attention bias is required, so flash-attention won't work for now but torch SDPA can still use the memory efficient kernel from xformers, right? I did some benchmarking with Chronos models (based on T5 architecture) here (amazon-science/chronos-forecasting#33) and there's a clear speedup when using torch SDPA.
27 remaining items
Disable `sdpa` for `blip2` (#579)
github-actions commentedon Oct 25, 2024
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 are likely to be ignored.
ArthurZucker commentedon Oct 29, 2024
It should not be super complicated, and a good refactoring is quite needed. Once #22105 is merged, we can easily add it!
toneyfox commentedon Nov 5, 2024
ailab_OmniGen
Phi3Transformer does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention yet. Please request the support for this architecture: #28005. If you believe this error is a bug, please open an issue in Transformers GitHub repository and load your model with the argument
attn_implementation="eager"meanwhile. Example:model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="eager")HELP ME。
The boss asked how to solve this problem？
ENate commentedon Nov 5, 2024
Hi there. Are there any open models?
michaelshekasta commentedon Nov 5, 2024
@ArthurZucker
Okay, do we have something like a roadmap? I understand that we should finish the DeBERTaV2 refactoring before implementing FA2 for DeBERTaV3. Do I understand it correctly
ArthurZucker commentedon Nov 25, 2024
Yep! The refactor is now merged, I'll have a look at #34826!
github-actions commentedon Dec 20, 2024
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 are likely to be ignored.
sneha09-hub commentedon Apr 3, 2025
I am using "unsloth/Pixtral-12B-2409-bnb-4bit"
Today I am facing error while loading the model sayin :
'PixtralVisionModel does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention yet'
Gorodecki commentedon Jun 4, 2025
I am using audio class model Wav2Vec2ForCTC
Error:
ValueError: Wav2Vec2ForCTC does not support Flash Attention 2.0 yet.
or
Wav2Vec2ForCTC does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention yet.
halfdayoff commentedon Aug 15, 2025
Will GptOssForCausalLM support SDPA？
Error:
raise ValueError(
ValueError: GptOssForCausalLM does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention yet. Please request the support for this architecture: #28005. If you believe this error is a bug, please open an issue in Transformers GitHub repository and load your model with the argument
attn_implementation="eager"meanwhile. Example:model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="eager")ArthurZucker commentedon Aug 21, 2025
No!
sdpadoes not allow the usage of attention sink 😉Disable `sdpa` for `blip2` (#579)
training: gpt-oss-20b — force attn_implementation=eager