Fix attention sink implementation in flex attention
Author: SamuelBarryCS
Pull request: yes
Labels: name: flex attention

## What
- Fixes #41026 where @jonny-so correctly identified that attention sinks were being incorrectly applied within the [score_mod] function of flex attention
- The previous implementation attempted to apply attention sinks by manipulating pre-softmax scores within score_mod, which is incorrect because attention sinks require access to the full attention matrix after softmax normalization ==> this fix moves attention sink application after flex attention computation to renormalize the output correctly
- As a result, we can now run flex attention with attention sink (which was impossible before because it would crash when s_aux was provided)


## Test performed
- No major logic change, existing tests are thus still passing
- Created a small script `test_attention_sink.py` to test the old vs. new version of the attention sink implementation in flex attention.` flex_attention_old.py` is simply a copy/paste from the file `flex_attention.py` on main, and this script + `test_attention_sink.py`  are of course both meant to be deleted before merging. The fix is confirmed by the script output:

```
Testing attention sinks implementation...
Input shapes: query torch.Size([1, 2, 4, 16]), key torch.Size([1, 2, 4, 16]), value torch.Size([1, 2, 4, 16])
s_aux shape: torch.Size([2])

--- Old Implementation ---
Status: FAILED - ['LoweringException: AssertionError: wrong ndim () [2]', '  target: flex_attention']

--- New Implementation ---
Before calling flex_attention_new...
torch.Size([1, 2, 4, 16]) torch.Size([1, 2, 4]) torch.Size([2])
After calling flex_attention_new...
Status: SUCCESS
Output: (tensor([[[[-0.4479, -0.0579,  0.0891, -0.1978,  0.4921, -0.0836, -0.4697,
            0.2943,  0.4493, -0.2208, -0.1416, -0.0743, -0.2410,  0.4453,
           -0.3093,  0.0605],
[...]
```


## How to review
- Check diff
- Run `test_attention_sink.py`



Comments:
cc @Cyrilvallez 

Nice work! Maybe it also makes sense to raise an error if `return_lse==False` and `s_aux is not None`, otherwise `s_aux` gets silently ignored when not running on CUDA.

cc @ArthurZucker this looks good but I believe you're more familiar with sinks!

> Nice work! Maybe it also makes sense to raise an error if `return_lse==False` and `s_aux is not None`, otherwise `s_aux` gets silently ignored when not running on CUDA.

Very good point @jonny-so :) I also agree that this silent ignore was not ideal. Just addressed it here https://github.com/huggingface/transformers/pull/41083/commits/f37738fb15bc19d187443502e5000b11d0374d76 

FYI I removed the test script and duplicated old version of flex_attention in https://github.com/huggingface/transformers/pull/41083/commits/cdf3325898a9178ac589bc21629907626c8db488 to make the PR merge-ready. Feel free to switch to an earlier commit if you want to run the test on your side guys, otherwise good to merge on my side!

> Great PR, thanks a lot 🫡
> If you want we can add a small educational comment in the gpt_oss.md file?

Thanks Arthur!
Sure thing, done here: https://github.com/huggingface/transformers/pull/41083/commits/58799e2344b8ac2705fb4bd6f6ded2707e2ff565. Ready to be merged on my side 🫡

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