flex_attention does not output the full attention_weights with output_attention option
Author: gblackout
Pull request: no
Labels: name: bug

### System Info

transformers==4.45.1
torch==2.6.0


### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

```attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
                logger.warning_once(
                    "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
                    'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
                )
            else:
                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )

### Expected behavior

The output_attention option (e.g. in [llama](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L232)) is supposed to return the full attention weights of shape (bsize, nheads, len_q, len_k). Currently, transformers support this option with eager attention and flex_attention, but it looks like what [flex_attention](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/flex_attention.py) returns are just the logsumexp scores of shape (bsize, nheads, len_q). It seems they do not support outputting the full attention weight either.

Comments:
The output_attention option (e.g. in [llama](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L232)) is supposed to return the full attention weights of shape (bsize, nheads, len_q, len_k). Currently, transformers support this option with eager attention and flex_attention, but it looks like what [flex_attention](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/flex_attention.py) returns are just the logsumexp scores of shape (bsize, nheads, len_q). It seems they do not support outputting the full attention weight either.

I believe that flashattention and possibly flexattention don't support outputting weight scores because they aren't fully materialized during the computation, but I'm not certain about that! cc @ArthurZucker 

Yes, what I mean is that currently, it allows you to set output_attention=True and use flex_attention, but it does not function as expected because, just like you said, it does not fully materialize the attention. So in that case, maybe it should warn and switch  back to eager just like that with the sdpa attention.

I think this is related to the version of transformers no? We now fully pass the flag to `flex_attention` call that return the attention weights 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

The problem is `flex-attention` returns `LSE` rather than `attention weights`, and you can't recover `attention weights` from `LSE`