Fix silent SDPA math-kernel fallback for GQA when key/value head_dim > 256 or differ
Author: Butterfingrz
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=46960)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46960)
<!-- ci-dashboard-badge:end -->

## What does this PR do?

For attention layers that use GQA with head_dim > 256 (notably Gemma4's head_dim=512 full-attention layers), when no attention_mask is passed (mask=None) during training/prefill, sdpa_attention_forward passes enable_gqa=True to SDPA. However, PyTorch's mem-efficient (EFFICIENT_ATTENTION) and flash SDPA kernels do not support the enable_gqa broadcast at head_dim > 256, so SDPA silently falls back to the MATH kernel, which materializes the S×S score matrix (O(S²) memory) and OOMs at long sequence lengths (e.g. 8192).

Because this fallback is completely silent (no warning), it is very hard to diagnose — it is effectively a "should run on the efficient path, but quietly degrades" correctness/resource bug. This PR makes use_gqa_in_sdpa return False for head_dim > 256, so we take the existing repeat_kv (MHA) path, which the mem-efficient kernel handles fine at head_dim=512, avoiding the silent MATH fallback. (This does not change behavior for head_dim ≤ 256 models; for head_dim > 256 no kernel supports the enable_gqa broadcast anyway, so repeat_kv is the strictly better and correct fallback.)

The same silent fallback also affects MLA-style models where the key and value head_dim differ (e.g. DeepSeek-V3 with qk=192 / v=128), as well as MLA where only `value` head_dim exceeds 256. Under `enable_gqa=True` the Flash kernel is disqualified (it requires q/k/v to share the last dim ≤ 256) and the mem-efficient kernel lacks the GQA broadcast, so only cuDNN can serve them — and they drop to MATH on any build/GPU without the cuDNN SDPA backend. This PR guards all of these: it checks both key and value head_dim, and requires them to be equal.

## Root cause

The root cause is that `use_gqa_in_sdpa` decides whether to use `enable_gqa=True` based only on:

* `torch >= 2.5`
* `attention_mask is None`

It does not take the attention head dimension into account.

```python
def use_gqa_in_sdpa(attention_mask, key):
    ...
    return _is_torch_greater_or_equal_than_2_5 and attention_mask is None
```

As a result, `sdpa_attention_forward` skips the existing `repeat_kv` path and passes `enable_gqa=True` whenever `mask is None`.

For layers with `head_dim > 256`, this requests an unsupported SDPA GQA-broadcast path. PyTorch's efficient/flash SDPA kernels reject that path, so SDPA silently falls back to the MATH kernel, which materializes the full `S × S` attention score matrix.

### Trigger chain for Gemma4

```text
Gemma4TextAttention full-attention layer
  ├─ config.global_head_dim = 512
  └─ no sliding_window / local_size
        ↓
create_causal_mask(...)
        ↓
_ignore_causal_mask_sdpa(...)
  ├─ kv_length == q_length
  └─ no padding
        ↓
returns None
        ↓
sdpa_attention_forward receives attention_mask=None
        ↓
use_gqa_in_sdpa(...) returns True
        ↓
SDPA is called with enable_gqa=True
        ↓
efficient/flash SDPA cannot handle enable_gqa at head_dim=512
        ↓
silent fallback to MATH
        ↓
O(S²) memory usage and OOM at long sequence lengths
```

Sliding-window layers are unaffected because they always materialize a 4D attention mask when `kv_length > sliding_window`. Therefore, they do not enter the `attention_mask is None` path and do not trigger `enable_gqa=True`.

This is why the issue is specifically observed on Gemma4's `head_dim=512` full-attention layers, while sliding layers do not OOM through this path.



## The fix

`use_gqa_in_sdpa` now also takes `value` and only opts into `enable_gqa=True` when the key and value head_dim are equal and ≤ 256:

```diff
-def use_gqa_in_sdpa(attention_mask, key):
+def use_gqa_in_sdpa(attention_mask, key, value):
     ...
     if _is_torch_xpu_available:
         return _is_torch_greater_or_equal_than_2_8
-    return _is_torch_greater_or_equal_than_2_5 and attention_mask is None
+    return _is_torch_greater_or_equal_than_2_5 and attention_mask is None and key.shape[-1] == value.shape[-1] <= 256
```

The call site in `sdpa_attention_forward` passes `value` too. A comment documents the boundary: GQA-in-SDPA is kept only when the Flash kernel can portably serve it (equal head_dim ≤ 256); at head_dim > 256, or when key/value head_dim differ, no kernel except cuDNN supports the `enable_gqa` broadcast, so `repeat_kv` (MHA) is the correct fallback.

The fix lives in the shared dispatch helper rather than a model file because the enable_gqa decision happens entirely there and the model side has no hook to influence it; this is also a general SDPA kernel limitation, not Gemma4-specific. This PR does not touch any modular_*.py/modeling_*.py (the change is confined to the integrations layer).

## Reproduction (Gemma4-shaped tensors, fp32 on CPU / bf16 on CUDA)

```python
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers.integrations.sdpa_attention import use_gqa_in_sdpa

torch.manual_seed(0)
B, Hq, Hkv, S, D = 1, 16, 8, 4096, 512  # GQA (groups=2), head_dim=512
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

q = torch.randn(B, Hq, S, D, device=device, dtype=dtype)
k = torch.randn(B, Hkv, S, D, device=device, dtype=dtype)
v = torch.randn(B, Hkv, S, D, device=device, dtype=dtype)

print(f"device={device}, head_dim={D}, mask=None")
print("use_gqa_in_sdpa(None, k, v) ->", use_gqa_in_sdpa(None, k, v))

# MLA-style differing head_dim (both <= 256) is also guarded now:
km = torch.randn(B, Hkv, S, 192, device=device, dtype=dtype)
vm = torch.randn(B, Hkv, S, 128, device=device, dtype=dtype)
print("use_gqa_in_sdpa(None, 192, 128) ->", use_gqa_in_sdpa(None, km, vm))  # False

if device == "cuda":
    # Buggy path: enable_gqa=True is rejected by the mem-efficient kernel -> SDPA would silently use MATH.
    try:
        with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
            torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)
        print("enable_gqa=True on EFFICIENT kernel: unexpectedly worked")
    except RuntimeError as e:
        print("enable_gqa=True on EFFICIENT kernel: REJECTED ->", str(e).splitlines()[0])

    # Fixed path: repeat_kv -> MHA, no enable_gqa -> the mem-efficient kernel accepts head_dim=512.
    kf = k.repeat_interleave(Hq // Hkv, dim=1)
    vf = v.repeat_interleave(Hq // Hkv, dim=1)
    with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
        out = torch.nn.functional.scaled_dot_product_attention(q, kf, vf, is_causal=True)
    print("repeat_kv (MHA) on EFFICIENT kernel: OK, out shape", tuple(out.shape))
```

## Before/after

* head_dim=512 (equal key/value): on `main` the dispatch decision is `True` (still takes the enable_gqa path); with this PR it is `False` (takes repeat_kv). Signature note: on `main` the helper is `use_gqa_in_sdpa(mask, key)`; this PR adds `value` → `use_gqa_in_sdpa(mask, key, value)`.
* MLA differing head_dim (both ≤ 256, e.g. 192/128): on `main` → `True` (only cuDNN can serve enable_gqa, so it silently drops to MATH where cuDNN is unavailable); with this PR → `False` (repeat_kv, served by the mem-efficient kernel). Verified on H20 / torch 2.11: `can_use_flash_attention` and `can_use_efficient_attention` are both False for 192/128.
* CUDA (needs a GPU): under the EFFICIENT_ATTENTION backend, enable_gqa=True at head_dim=512 is rejected by the kernel (the exact point where the buggy path silently falls back to MATH), while the repeat_kv (MHA) path succeeds — concretely showing that EA supports 512 but not enable_gqa.

## Added test

`tests/utils/test_modeling_utils.py::TestAttentionImplementation` adds three tests:

- `test_use_gqa_in_sdpa_skips_large_head_dim` — head_dim > 256 (checking both key and value) disables GQA-in-SDPA; equal head_dim ≤ 256 with no mask stays enabled; an explicit mask never enables it. (CPU)
- `test_use_gqa_in_sdpa_skips_mismatched_head_dim` — differing key/value head_dim ≤ 256 (e.g. 192/128) disables it too. (CPU)
- `test_gqa_in_sdpa_does_not_fall_back_to_math` (`@require_torch_gpu`) — forces the mem-efficient backend (excluding MATH) so a silent fallback becomes a hard error: the raw `enable_gqa=True` call raises, while `sdpa_attention_forward` (repeat_kv path) runs, for both head_dim=512 and 192/128.

The first two fail on main and pass with this fix. (The XPU branch is out of scope for this fix and is skipped.)

## Test commands and results (local, torch 2.12.1+cpu)

```bash
# Dispatch unit tests (CPU)
$ python -m pytest tests/utils/test_modeling_utils.py -k use_gqa_in_sdpa -q
2 passed

# Incl. the GPU functional test (run on a GPU, e.g. H20 / torch 2.11)
$ python -m pytest tests/utils/test_modeling_utils.py -k gqa_in_sdpa -q
3 passed

# Gemma4 attention (incl. eager_matches_sdpa no_attn_mask variants)
$ python -m pytest tests/models/gemma4/test_modeling_gemma4.py -k "sdpa or eager or attention" -q
108 passed, 63 skipped

# Representative GQA models regression (head_dim <= 256, covers the unchanged enable_gqa path)
$ python -m pytest tests/models/{llama,gemma2,qwen2}/test_modeling_*.py \
    -k "eager_matches_sdpa or sdpa_can_dispatch or attention_outputs" -q
81 passed, 9 skipped

# Attention-interface unit tests
$ python -m pytest tests/utils/test_modeling_utils.py -k "attn or attention or sdpa or interface" -q
14 passed, 7 skipped

# Repo consistency / typing / structure checks
$ make fix-repo     # All 15 checks passed
$ make typing       # All 2 checks passed
$ make check-repo   # All 24 checks passed
```

utils/tests_fetcher.py maps the change to ~477 model test files because it lives in the core shared file integrations/sdpa_attention.py; the runs above are a high-signal subset covering the changed path (head_dim=512), the unchanged path (head_dim ≤ 256 GQA), and the interface. ruff is clean.

## Scope of impact

* Models with equal key/value head_dim ≤ 256: behavior unchanged.
* GQA models with head_dim > 256 (e.g. Gemma4 head_dim=512): with mask=None they now take the repeat_kv (MHA) path, avoiding the silent MATH fallback and O(S²) OOM; numerical results are unchanged.
* MLA-style models with differing key/value head_dim (e.g. DeepSeek-V3 192/128), or with value head_dim > 256: also take repeat_kv instead of relying on cuDNN, which the mem-efficient kernel serves at differing head_dim; numerical results unchanged.
* Only the SDPA path is affected: eager / flash_attention / flex_attention / *_paged and the XPU dispatch are untouched.

## Code Agent Policy

* [✔ ] I confirm that this is not a pure code agent PR.

## Before submitting

* [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
* [✔ ] Did you read the contributor guideline (https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests) and the Pull Request checks (https://huggingface.co/docs/transformers/pr_checks)?
* [ ] Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
* [✔ ] Did you make sure to update the documentation with your changes according to the guidelines?
* [ ✔] Did you write any new necessary tests?

## Review request
 Would appreciate a review from @vasqu @Cyrilvallez @zucchini-nlp Thank you!

Comments:
> It reminds me of #45776 😢 agree this one does make sense regardless as base FA does not support this
> 
> Just a few smaller comments first

@vasqu  Thanks for your review , Gemma4's non-standard attention has been a real headache. I hit this on the AutoModel Gemma4 benchmark: when the attention path passes no mask (`attention_mask is None`), we opt into `enable_gqa=True`, and at `head_dim=512`, SDPA's memory-efficient (EA) kernel silently falls back to the `math` backend because it does not support the GQA broadcast there.

This showed up with a fully packed dataset: every sample was already at `max_length`, so there was no padding mask. That triggered an OOM I would not see in normal training, purely from the `O(S^2)` math fallback.

Hence the minimal change here to keep us off that path. 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46960). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@vasqu  Thanks for looking! The `tests_peft_integration [shard 0/1]` failure isn't from this PR — it's an incompatibility between current `main` and `PEFT 0.19.1`, the minimum pinned by `#46442`.

CI installs `peft==0.19.1`, whose `build_peft_weight_mapping` calls:

```python
WeightConverter(
    ...,
    distributed_operation=...,
    quantization_operation=...,
)
```

On current `main`, that constructor only accepts:

```python
(source_patterns, target_patterns, operations)
```

so it raises inside PEFT:

```text
peft/utils/transformers_weight_conversion.py:268
E   TypeError: WeightConverter.__init__() got an unexpected keyword argument 'distributed_operation'
```

This PR only touches:

- `integrations/sdpa_attention.py`
- `tests/utils/test_modeling_utils.py`

It does not touch any PEFT or weight-loading code. The failure only surfaces here because `#47021` wired the PEFT shard into the fetcher and this PR hits a core shared file.

Any PR whose CI runs that shard on current `main` fails the same way. It looks like this needs a fix on the PEFT side, or restoring the two kwargs.

Yep on it, not on you! Will merge as soon as I can

#47027

Thank you so much!

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46960)
**Latest run:** [28613003356:2](https://github.com/huggingface/transformers/actions/runs/28613003356)
**Result:** `failure` | **Jobs:** 14 | **Tests:** 46,801 | **Failures:** 0 | **Duration:** 14h 38m

<!-- ci-dashboard-recap:end -->

Hi @Butterfingrz and @vasqu,

Just wanted to mention that DeepseekV3Attention derives `num_key_value_groups` from `num_key_value_heads`, which MLA never honors, so any `kv_heads != attn_heads` config silently produces a broken model.

I don't know if you would consider to take this into account: maybe force `groups=1` or validate in `DeepseekV3Config.__post_init__`?

On the other hand, I think this PR is correct and shouldn't be touched.

Hey @albertvillanova, that would affect pretty much any model using MLA correct? `num_key_value_heads` is actually never used it seems but still derived the `num_key_value_groups`.

Imo, we should add something to `validate_architecture` that raises an error directly when the config has mismatching num heads vs num kv heads. Wdyt?