SageAttention for attention implementation?
Author: Many0therFunctions
Pull request: no
Labels: name: Feature request

### Feature request

I've noticed it's been a while now, but transformers still only has flash attention as the fastest attention backend for calls like these: 

<img width="1307" height="780" alt="Image" src="https://github.com/user-attachments/assets/3f3d62f6-a166-4ca6-97a0-49263fd93299" />

Are there any plans to add sageattention as well? 

### Motivation

It's become increasingly involved to have to monkey patch sage attention support for every new model that comes out, and for older models that used older versions of transformers, I've had to do unholy things like this:

<img width="1296" height="705" alt="Image" src="https://github.com/user-attachments/assets/c5f4ff6a-094a-48f4-9339-17de1ece43d0" />


### Your contribution

I have an example of a patch I had to do so I will upload that here

[llama_nar.py.txt](https://github.com/user-attachments/files/21393926/llama_nar.py.txt)

Comments:
I don't think we have any plans to add SageAttention! It's more of a quantization method than a fully equivalent attention layer, right? cc @arthurzucker @MekkCyber 

It more that our goal is to be able to enable any `attention` method, without you having to monkeypatch anything!

Normally the only thing you should have to do is this: ( as per the documentation https://huggingface.co/docs/transformers/en/attention_interface) 
```python 
from transformers import AutoModelForCausalLM, AttentionInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
import torch
from sageattention import sageattn

model_id = "meta-llama/Llama-3.2-1B"

def sage_attention(module, query_states, key_states, value_states, _, **kwargs):
    return sageattn(query_states, key_states, value_states,  tensor_layout="HND")

AttentionInterface.register("sage_attention", sage_attention)

model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="sage_attention", device_map="auto")
model(torch.ones(1, 5, dtype=int, device=model.device))
```

We are really striving to make sure we get enough flexibility for powerful users like you, any feedback is welcome!
If the ` attention_mask` is not even needed, I think the automatic creation should be disabled

( I am getting an error but on sageattn's side:
```python
cuda:7 cuda:7 cuda:7
Traceback (most recent call last):
  File "/fsx/arthur/test_sage.py", line 19, in <module>
    model(torch.ones(1, 5, dtype=int, device=model.device))
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/transformers/src/transformers/utils/generic.py", line 961, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 460, in forward
    outputs: BaseModelOutputWithPast = self.model(
                                       ^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/transformers/src/transformers/utils/generic.py", line 1067, in wrapper
    outputs = func(self, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 390, in forward
    hidden_states = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/fsx/arthur/transformers/src/transformers/modeling_layers.py", line 94, in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 289, in forward
    hidden_states, _ = self.self_attn(
                       ^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/transformers/src/transformers/models/llama/modeling_llama.py", line 248, in forward
    attn_output, attn_weights = attention_interface(
                                ^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/test_sage.py", line 11, in sage_attention
    return sageattn(query_states, key_states, value_states, tensor_layout="HND")
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/sageattention/core.py", line 105, in sageattn
    q_int8, q_scale, k_int8, k_scale = per_block_int8(q, k, sm_scale=sm_scale, tensor_layout=tensor_layout)
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/sageattention/quant_per_block.py", line 63, in per_block_int8
    quant_per_block_int8_kernel[grid](
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/triton/runtime/jit.py", line 330, in <lambda>
    return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/triton/runtime/jit.py", line 653, in run
    kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata,
  File "/fsx/arthur/.venv/lib/python3.12/site-packages/triton/backends/nvidia/driver.py", line 444, in __call__
    self.launch(*args, **kwargs)
ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
```
)