Unexpected output of _flash_attention_forward() for cross attention
Author: Cypress-Yang
Pull request: no
Labels: name: bug

### System Info

My environment:

> transformers              4.44.1
> flash-attn                2.6.3



### Who can help?

_No response_

### Information

- [ ] My own modified scripts

### Tasks

- [ ] My own task or dataset (give details below)

### Reproduction

```py
import torch
from transformers.modeling_flash_attention_utils import _flash_attention_forward

query = torch.randn(4, 20,1, 32).cuda().half()
key = torch.randn(4, 10, 1, 32).cuda().half()
value = key

unmasked = _flash_attention_forward(query,key,value,
                                    attention_mask=None, 
                                    query_length=20, 
                                    is_causal=False)
masked = _flash_attention_forward(query,key,value, 
                                    attention_mask=torch.ones((4, 10)).cuda().bool(), 
                                    query_length=20, 
                                    is_causal=False)
breakpoint()
```


### Expected behavior

In my understanding, the attention_mask has a size of `(batch_size, seq_len)` where 1 stands for the position of non-padding tokens, so an all-one mask should lead to the same result as no mask provided. However, the outputs are significantly different.
```
(Pdb) print(unmasked, masked)
tensor([[[[ 0.4114, -0.3369,  0.6221,  ..., -0.4475, -0.2361,  0.3022]],

         [[ 0.2480, -0.1396,  0.1614,  ..., -0.0728, -0.2788,  0.3950]],

         [[ 0.3828, -0.1323,  0.2101,  ..., -0.4751, -0.0179,  0.3181]],

         ...,

         [[ 0.2654, -0.3137,  0.1637,  ...,  0.3464, -0.6318,  0.4377]],

         [[ 0.5464, -0.2251,  0.4897,  ..., -0.3184, -0.1769,  0.3203]],

         [[-0.1514,  0.3037, -0.1609,  ..., -0.4651, -0.1842,  0.3386]]],


        [[[ 0.1772,  0.3240, -1.1143,  ...,  0.1444,  0.5684,  0.3770]],

         [[ 0.4187,  0.2264,  0.2446,  ...,  0.7036,  0.3003,  0.2981]],

         [[ 0.1241,  0.1919, -0.5239,  ..., -0.1606,  0.5210, -0.1896]],

         ...,

         [[ 0.5225, -0.2333,  0.1004,  ...,  0.0297,  1.0059, -0.1329]],

         [[ 0.4304,  0.4819,  0.1232,  ...,  0.5234,  0.5210,  0.2379]],

         [[ 0.5361, -0.0976, -0.3975,  ...,  0.2217,  0.8481,  0.0780]]],


        [[[-0.6426, -0.1761,  0.3420,  ...,  0.4404,  0.5273,  0.0485]],

         [[-0.2313,  0.5249,  0.8975,  ...,  0.2517,  0.2163,  0.3628]],

         [[-0.9180, -0.7173, -0.3291,  ...,  0.0781,  1.0693, -0.5142]],

         ...,

         [[-0.8945, -0.1444, -0.0460,  ...,  0.2571,  0.8721, -0.0226]],

         [[-0.6978, -0.7417,  0.2061,  ...,  0.2173,  0.2798, -0.2246]],

         [[-0.3818, -0.7246,  0.7720,  ..., -0.3567,  0.0623, -0.0179]]],


        [[[-0.5347, -0.6885, -1.3604,  ..., -1.3672, -1.1768, -1.2275]],

         [[-0.2400, -0.5176, -0.4875,  ..., -0.2822,  0.1527, -0.0917]],

         [[-0.1940, -0.1766, -0.8022,  ..., -0.3743, -0.2607,  0.1602]],

         ...,

         [[-0.2346, -0.4260, -0.2166,  ...,  0.1776, -0.2793, -0.8052]],

         [[-0.3430, -0.9839,  0.3735,  ...,  0.3267,  0.4268, -0.5464]],

         [[-0.3293, -0.0431, -1.1631,  ..., -0.5742, -1.3242, -1.2441]]]],
       device='cuda:0', dtype=torch.float16) tensor([[[[ 0.4114, -0.3369,  0.6221,  ..., -0.4475, -0.2361,  0.3022]],

         [[ 0.2480, -0.1396,  0.1614,  ..., -0.0728, -0.2788,  0.3950]],

         [[ 0.3828, -0.1323,  0.2101,  ..., -0.4751, -0.0179,  0.3181]],

         ...,

         [[-0.0500, -0.0776,  0.3552,  ...,  0.6475,  0.1764, -0.2125]],

         [[ 0.7197,  0.4253,  0.0373,  ...,  0.7168,  0.5254,  0.2496]],

         [[ 0.0498,  0.1896, -0.1191,  ...,  0.1239,  0.5039, -0.0274]]],


        [[[-0.6572, -0.5571, -0.0978,  ...,  0.2294,  0.8623, -0.4048]],

         [[-0.5635, -0.4026,  0.1295,  ...,  0.1743,  0.6333, -0.0356]],

         [[-0.8359, -0.7275,  0.3054,  ..., -0.2150,  0.5693, -0.5825]],

         ...,

         [[-0.9004, -0.7935,  0.1372,  ..., -0.1024,  0.6543,  0.3892]],

         [[-0.1636,  0.1940, -0.9355,  ..., -0.2068, -0.7847, -0.1024]],

         [[-0.1720, -0.6123, -0.6470,  ..., -0.3550,  0.4495, -0.0429]]],


        [[[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         ...,

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]],


        [[[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         ...,

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]]],
       device='cuda:0', dtype=torch.float16)
```

Is this due to a bug? Or I just have some misunderstandings about this function?

Comments:
cc @ArthurZucker because I see you in the git blame for that function

I also find the same problem. When attention_mask is not `None`, the function `_flash_attention_forward` will firstly upad the inputs by function `_upad_input`, causing the dimension of query, key and value to be the same. In your case, the dimension of query, key and value should be `(40, 1, 32)` after upadding. And then attention is calculated like a self attention manner by function `flash_attn_varlen_func`.  Finally, it will pad back missing values by function `pad_input`. So you could find the masked result has so many 0.0 values.

So there should be some problem for the flash attention calculation when qkv is cross attention.

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.

So anyone could help with this issue? It is very important for me, please. @Rocketknight1 

Hey! Sorry for not answering.  These functions are not supposed to be used like that. 
The best model to have a look at is probably `BART`, which does use cross attention. 

Transformers, up until now, has always had to unpadd the input to go through flash attention, which is why you are seeing this behaviour!

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.

Same issue here. @ArthurZucker does FA2 support cross attention with both encoder padding mask and decoder padding mask? Just don't want the decoder to attend the padding tokens.

Hey! what do you mean? Flash attention can either compute cross or self attention. The mask defines whether you are causal or not, and the type of input (k and v being projected HS or being encoder states) defines if you are cross or self. 

@ArthurZucker I mean does it support 4D attention mask like SDPA with shape like [B, H, L, S], and I want to define in which positions in encoder states, the decoder query should not attend, namely using the [L, S] part to define the padding mask in SDPA.

The attention masks just need to be converted to `cumulative_q` and ` cumulative_k`.
It should be supported! 

So it looks like I should not depends on this logic https://github.com/huggingface/transformers/blob/cf97f6cfd1c9cff6f3f188e91b32dd3b2a1c5fae/src/transformers/modeling_flash_attention_utils.py#L400, instead, I should prepare cu_seq_lens_q and cu_seq_lens_k myself and pass through to https://github.com/huggingface/transformers/blob/cf97f6cfd1c9cff6f3f188e91b32dd3b2a1c5fae/src/transformers/modeling_flash_attention_utils.py#L355?

Yes if your usecase is specific. Otherwise decoding or just normal training with a 2d mask is supported! 
We have a datacollator with flattening as well