Repository: transformers
Issue: Unexpected output of _flash_attention_forward() for cross attention
Labels: 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?