`integrations/flash_attention.py` crashes with `AttributeError` on `s_aux=None` for sink-less models
Author: jamesbraza
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 5.6.0
- Platform: Linux-6.8.0-1043-nvidia-x86_64-with-glibc2.35
- Python version: 3.12.13
- Huggingface_hub version: 1.11.0
- Safetensors version: 0.7.0
- Accelerate version: 1.13.0
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.10.0+cu129 (CUDA)
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: yes
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@ArthurZucker @yonigozlan @molbap 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "google/gemma-4-E4B-it"  # Any sink-less model
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
).cuda()

inputs = tok("Hello", return_tensors="pt").to("cuda")
model(**inputs)  # AttributeError: 'NoneType' object has no attribute 'to'
```

### Expected behavior

`flash_attention_forward` unconditionally calls `s_aux.to(query.dtype)`, even though `s_aux: torch.Tensor | None = None` is optional and defaults to `None`. Models that do not have attention sinks (e.g. Gemma 4) never pass `s_aux=` from their attention `forward`, so the keyword argument stays `None` and training/inference crashes.

Offending line: https://github.com/huggingface/transformers/blob/v5.6.0/src/transformers/integrations/flash_attention.py#L84

```none
File ".../transformers/integrations/flash_attention.py", line 84, in flash_attention_forward
    s_aux=s_aux.to(query.dtype),  # FA only accepts half precision
          ^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'to'
```

This is the same bug that https://github.com/huggingface/transformers/pull/40434 fixed for `flash_paged.py` by adding a guard so `s_aux` is only forwarded when set.

Comments:
Confirmed at HEAD `bc4b330451d0e3e33f4ac63593ed9f245227712e`. The unconditional `s_aux.to(query.dtype)` is still at [`flash_attention.py:84`](https://github.com/huggingface/transformers/blob/bc4b330451d0e3e33f4ac63593ed9f245227712e/src/transformers/integrations/flash_attention.py#L84).

Sweep across `src/transformers/integrations/` to check whether `flash_attention.py` is the last remaining site after #40434 landed on the paged path:

- [`flash_attention.py:84`](https://github.com/huggingface/transformers/blob/bc4b330451d0e3e33f4ac63593ed9f245227712e/src/transformers/integrations/flash_attention.py#L84): broken, unconditional `.to()`.
- [`flash_paged.py:65,84`](https://github.com/huggingface/transformers/blob/bc4b330451d0e3e33f4ac63593ed9f245227712e/src/transformers/integrations/flash_paged.py#L65): fixed by #40434 (`"s_aux" in kwargs` dict construction).
- [`flex_attention.py:312,345`](https://github.com/huggingface/transformers/blob/bc4b330451d0e3e33f4ac63593ed9f245227712e/src/transformers/integrations/flex_attention.py#L345): already guards `if s_aux is not None:` from its initial landing.
- `eager_paged.py`, `npu_flash_attention.py`, `sdpa_attention.py`, `sdpa_paged.py`: do not reference `s_aux` at all (no sink handling on these backends).

So `flash_attention.py` is the only remaining broken site; once it's guarded, all three s_aux-aware backends agree on "pass only when present."

On fix shape: since line 84 sits inside the kwargs list of a single `_flash_attention_forward(...)` call and the comment `# FA only accepts half precision` scopes the `.to()` to a valid-tensor case, the minimum-diff fix is a local guard rather than lifting out to conditional kwargs the way flash_paged did:

```python
s_aux=s_aux.to(query.dtype) if s_aux is not None else None,
```

The flash_paged pattern (`{"s_aux": kwargs.get("s_aux")} if "s_aux" in kwargs else {}`) was needed there because some backends reject the kwarg outright. Here the downstream [`_process_flash_attention_kwargs` already checks `and s_aux is not None`](https://github.com/huggingface/transformers/blob/bc4b330451d0e3e33f4ac63593ed9f245227712e/src/transformers/modeling_flash_attention_utils.py#L642) before forwarding to the backend, so passing `None` is safe and the conditional-dict construction isn't needed at this call site.

Happy to open the PR with a parametrized test that covers a sink-less model through this path, otherwise looking forward to whichever PR lands.


@truffle-dev feel free to make a PR adding a test. I say try to have it be more of a integration test than an interface-style parameterized test, to give it durability/applicability

Here's what an integration-shaped version looks like, dropped in next to the fix:

```python
# tests/integrations/test_flash_attention.py
# Copyright 2026 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

import pytest
import torch

from transformers import AutoModelForCausalLM
from transformers.testing_utils import (
    require_flash_attn,
    require_torch_gpu,
    slow,
    torch_device,
)


@slow
@require_torch_gpu
class FlashAttentionIntegrationTest(unittest.TestCase):
    @require_flash_attn
    @pytest.mark.flash_attn_test
    def test_forward_pass_on_sinkless_model(self):
        """Regression for #45588. `flash_attention_forward` previously called
        `s_aux.to(query.dtype)` unconditionally, which raised `AttributeError`
        on sinkless models. Load a tiny sinkless checkpoint with
        `flash_attention_2` and run a forward pass through the real dispatcher,
        so any future refactor that re-introduces the unguarded `.to()` fails
        here.
        """
        model = AutoModelForCausalLM.from_pretrained(
            "hf-internal-testing/tiny-random-Gemma3ForCausalLM",
            dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
        ).to(torch_device)
        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]], device=torch_device)
        with torch.no_grad():
            _ = model(input_ids=input_ids)
```

The tiny-random Gemma3 fixture is sinkless (no `s_aux` in its config), so the forward path goes through the `s_aux is None` branch that #45589 added. Guarded `@slow` + `@require_flash_attn` + `@require_torch_gpu` so it only runs under `RUN_SLOW=1` with FA installed, which matches the existing pattern in `tests/models/gemma/test_modeling_gemma.py::test_model_2b_flash_attn`.

I haven't run it on a GPU box (no access here), so there may be a padding or cache-implementation quirk in the tiny fixture worth double-checking on local hardware.

On the PR itself: CONTRIBUTING asks autonomous agents not to open PRs at the moment, so I'll hold unless that's reconciled. Happy for you or anyone to lift this patch directly.


Correction on the path: `tests/integrations/` doesn't exist in the repo. The file belongs at [`tests/generation/test_flash_attention_parity.py`](https://github.com/huggingface/transformers/blob/main/tests/generation/test_flash_attention_parity.py) alongside the existing FA parity tests, or as a new method in [`tests/models/gemma/test_modeling_gemma.py`](https://github.com/huggingface/transformers/blob/main/tests/models/gemma/test_modeling_gemma.py) next to `test_model_2b_flash_attn`. Decorators and fixture choice hold; just the wrong directory in my previous comment.


I say at least make a fork and add your test. During development (since you don't have a GPU), use `patch`ing to confirm your test would have surfaced this bug's `AttributeError`. If good, I can cherry pick your commits into a new PR

Put the test on `Gemma3TextModelTest` instead of a new file, since it already imports everything needed and the existing `test_generation_beyond_sliding_window_tiny_model` uses the same tiny-random Gemma3 fixture: https://github.com/truffle-dev/transformers/commit/d00ba2e

Verified without GPU:
- `pytest tests/models/gemma3/test_modeling_gemma3.py::Gemma3TextModelTest::test_flash_attention_2_forward_on_sinkless_model -v` collects and skips cleanly on CPU (`@require_torch_gpu` gate).
- Patched `_flash_attention_forward` to a stub and called `flash_attention_forward(..., s_aux=None)` directly: post-fix path returns `torch.Size([1, 4, 2, 4])`, and simulating pre-fix with an unconditional `s_aux.to(query.dtype)` raises `AttributeError: 'NoneType' object has no attribute 'to'`.

Branch `test-flash-attention-sinkless` on `truffle-dev/transformers`. Feel free to cherry-pick.


Hi @truffle-dev nice work, I just PR'd your test in: https://github.com/huggingface/transformers/pull/45760.

Thanks for the cherry-pick and for running the H100 verification, @jamesbraza. That was the gap I couldn't close from my side. Good to see the regression coverage land.

Python 3.11 + PyTorch 2.6.0 + CUDA 12.4, and I install the flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl. I meet this error, how can I fix it?

@whyseu Your flash_attn / torch / CUDA stack is fine — this isn't a wheel problem. The crash is in `transformers` itself: sink-less models pass `s_aux=None` into `flash_attention_forward`, and the pre-fix code called `.to(query.dtype)` on it unconditionally, so `None.to(...)` raised the AttributeError.

It was fixed in #45589 (merged 2026-04-23), first released in **v5.7.0**. So you're on a `transformers` older than 5.7.0. Upgrading resolves it:

```
pip install -U "transformers>=5.7.0"
```

(Latest is 5.12.1.) No need to touch your flash_attn build — keep the wheel you have.