Proposal: add sdpa_memeff attn_implementation for shape combinations no fast backend covers

[dvdimitrov13]

Summary

Proposal to add a new attn_implementation="sdpa_memeff" that pins torch's SDPA dispatcher to SDPBackend.EFFICIENT_ATTENTION (via sdpa_kernel([EFFICIENT_ATTENTION]) wrapping the existing sdpa_attention_forward). Filing as an issue to validate design direction before opening a PR.

Motivation — two independent failure modes that attn_implementation="sdpa" doesn't handle well

1. head_dim > 256

On every CUDA arch I've tested (RTX 5090 sm_120; same library caps apply on H100 sm_90):

Backend	head_dim=320 (Gemma 4)
FLASH	REJECT (library cap: head_dim ≤ 256)
CUDNN	REJECT (library cap: head_dim ≤ 256)
EFFICIENT (CUTLASS mem-eff)	ACCEPT
MATH	ACCEPT (but O(seq²) fp32 softmax — 74 GB at seq=32k, OOMs on 96 GB card)

So for Gemma 4 (google/gemma-4-31B, gemma-4-26B-A4B, gemma-4-E4B, gemma-4-E2B — all head_dim=320) there's only one fast-path backend, and stock sdpa has no way to pin it reliably.

2. Input layouts that disqualify FLASH at any head_dim

pytorch/pytorch#44928 — Qwen3.5 RLHF training sees dense 4D mask materialization (from 3D position_ids), which disqualifies FLASH's is_causal=True short-circuit. Dispatcher falls through to MATH and triggers NaN gradients in bf16.

Same root shape of problem as (1), driven by layout rather than head_dim. A manual with sdpa_kernel([EFFICIENT_ATTENTION]): ... workaround exists but requires model-code edits; process-global torch.backends.cuda.enable_math_sdp(False) has side effects.

Proposed fix

New file src/transformers/integrations/sdpa_memeff.py. Registers "sdpa_memeff" in ALL_ATTENTION_FUNCTIONS. Copies the existing sdpa_attention_forward with two additions:

1. Wraps the actual SDPA call in with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION]): ...
2. Unconditional GQA repeat_kv (EFFICIENT_ATTENTION rejects dense GQA where num_heads_q != num_heads_kv).

Verified locally on full 8B Gemma-4-E4B-it (bf16, real checkpoint): 100% top-1 agreement vs stock sdpa where stock works, and correct execution where stock falls back to MATH and OOMs.

Questions for maintainers

1. Shape of the contribution. Dedicated "sdpa_memeff" name vs. a config knob on sdpa (e.g. sdpa_preferred_backends=["EFFICIENT"]) vs. the pipe syntax introduced by PR [attn_implementation] remove recursive, allows custom kernels with wrappers #39823 ("sdpa|efficient_attention")?

2. Overlap with pytorch-side fixes. pytorch's SDPA dispatcher is missing sm_120 support for CUDNN head_dim=256 (filed separately). Once that lands, CUDNN handles head_dim ≤ 256 on sm_120 — but head_dim=320 (Gemma 4) is still library-capped on both FLASH and CUDNN, and [Bug] Catastrophic gradient explosion (NaN) in RLHF with Qwen3.5 due to 3D position_ids forcing SDPA Math fallback and BF16 collapse #44928-class layout-induced FLASH disqualification is unaffected. So sdpa_memeff isn't made redundant.

3. Test coverage you'd want. Convergence (loss-curve equivalence with stock sdpa on a small model)? Numerical (per-batch logit diff)? Both?

4. Anything first-time-contributor I should know — CLA, preferred test location, maintainer norms for adding to ALL_ATTENTION_FUNCTIONS, etc.?

Happy to open a PR if this direction is right.

---

🤖 Drafted with Claude Code (Claude Opus 4.7), reviewed and posted by me.

[Rocketknight1]

cc @ArthurZucker @Cyrilvallez for flashattention/SDPA

[vasqu]

This doesn't really make sense as SDPA selects the most efficient available backend (if not explicitly disabled with decorators or such)

So if the cudnn backend (which afaik is not enabled by default) and FA backend fail, it still uses the memory-efficient path if possible. It goes through a priority order of

• Cudnn
• FA
• Memory efficient
• Math

Also see https://github.com/pytorch/pytorch/blob/cf7c9b3a733314f2501678577a2a60041a54a214/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp#L966-L1002 for the actual priority order (sdp_utils.cpp has them under the different HW folders, just used cuda as the most common one)

[vasqu]

If you want to pin a specific backend, I don't see why you just don't call the backend decorator to force it, e.g.

with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION]):
    pass  # do your custom logic here

So there is no reason to open a separate implementation for this

The only thing that might make sense is the GQA part but then I need clarification and maybe we should tighten the condition there instead to allow the mem efficient path to be invoked. GQA flag only works in tandem with FA so we mightve been to lax there. But still no need for a whole new attn implementation.

[Cyrilvallez]

Agreed with @vasqu here. You can simply do:

with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION]):
    model.generate(...)