T5: add SDPA and Flash Attention 2 support - #46946
Conversation
9b7aa2e to
0ac246e
Compare
Refactor T5Attention to use ALL_ATTENTION_FUNCTIONS dispatch system, bringing T5 in line with BERT, BART, Llama and other models that already support multiple attention backends. Key design decisions: - T5 does not scale Q*K by 1/sqrt(d_kv); pass scaling=1.0 explicitly - position_bias is computed before dispatch and passed as additive attn_mask, which SDPA handles natively - Flash Attention 2 cannot accept additive bias, so when FA2 is selected we transparently fall back to SDPA for correctness - output_attentions=True forces eager path since SDPA/FA don't return attention probabilities This resolves the long-standing gap where T5 was the only major architecture without _supports_sdpa / _supports_flash_attn, causing ValueError for users requesting these backends (e.g. issue huggingface#46640). Partially addresses huggingface#26350.
The previous commit added SDPA/FA2 dispatch code to longt5, mt5, pop2piano, and udop but forgot to import ALL_ATTENTION_FUNCTIONS, sdpa_attention_forward, and define eager_attention_forward. This fixes all 16 F821 'Undefined name' errors from the CI code quality check.
…FUNCTIONS T5 was the canonical 'legacy' example but now supports SDPA and FA2 through the modern AttentionInterface. Replaced it with FSMT as the legacy model reference.
3b92c62 to
168d4d9
Compare
|
[For maintainers] Suggested jobs to run (before merge) run-slow: longt5, mt5, pop2piano, switch_transformers, t5, udop |
CI recapDashboard: View test results in Grafana |
|
Hi @ArthurZucker @Cyrilvallez, just a friendly ping on this PR. All CI tests are passing (3724 tests, 0 failures). Is there anything else you'd like me to address or clarify? Happy to make any adjustments. |
|
This competes with #47014 |
|
Looks very agentic to me and the fact that FA support was added cries to me that it is either an agent or someone who hasn't even tested it. I know #47014 at least the contributors and it has a history from a previous PR, so much more serious and likely to be correct |
|
Ha indeed, did not notice that the other one was from @jiqing-feng! Closing this one then in favor of #47014 for sure! (and indeed it has all the markers of just random agent code that we would waste time a lot of time to explain stuff to machine...) |
4 participants
What does this PR do?
T5 (and the whole T5 family — FLAN-T5, mT5, ByT5) has been the only major architecture in transformers without SDPA or Flash Attention support. Users hitting
attn_implementation="sdpa"or"flash_attention_2"get aValueError(see #46640). This PR fixes that by refactoringT5Attentionto dispatch throughALL_ATTENTION_FUNCTIONS, matching how BERT, BART, Llama and others already work.Partially addresses #26350 (community effort to add FA2 to more architectures).
There is an older attempt (#31167) but it uses the deprecated
T5SdpaAttentionsubclass approach. Arthur asked there to switch to theALL_ATTENTION_FUNCTIONSdispatch pattern — this PR does exactly that.How it works
T5 attention is trickier than most models because of two things:
1. No scaling. T5 does not divide Q·K by √d (the original paper found it unnecessary with their initialization). The
scaling=1.0is passed explicitly so SDPA/flash kernels do not apply the default 1/√d.2. Learned relative position bias. T5 adds a learned
position_biastensor to the attention scores before softmax. This bias has shape(batch, heads, q_len, k_len)and is additive.scaled_dot_product_attentionaccepts an additiveattn_mask, soposition_bias(with the causal/padding mask folded in) goes straight through.alibi_slopesfor linear bias). When FA2 is selected andposition_biasis present, we transparently fall back to SDPA, which still gets the memory-efficient fused kernel from PyTorch.Changes
eager_attention_forward(mirrors the pattern inmodeling_bart.py/modeling_bert.py)T5Attention.forwarddispatches throughALL_ATTENTION_FUNCTIONS.get_interface()instead of doing manualmatmul → softmax → matmuloutput_attentions=Trueforces eager (SDPA/FA do not return attention weights)_supports_sdpa = Trueand_supports_flash_attn = TrueonT5PreTrainedModelVerification
SDPA vs eager logits are bit-identical (max diff 0.0 on CPU with fp32). Also tested generation, KV cache, gradient checkpointing, encoder-only model, and cross-attention — all green.
Quick benchmark (T5-base config, seq_len=1024, CPU):
Before
After