Inconsistent Outputs When Using Flash Attention 2 and SDPA Attention with Attention Mask
Author: tartarleft
Pull request: no
Labels: name: bug

### System Info


- `transformers` version: 4.46.3
- Platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.10
- Python version: 3.8.18
- Huggingface_hub version: 0.25.2
- Safetensors version: 0.4.5
- Accelerate version: 1.0.1
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - use_cpu: False
        - debug: False
        - num_processes: 8
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.4.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: No
- Using GPU in script?: True
- GPU type: NVIDIA A800-SXM4-40GB

### Who can help?

 @ylacombe, @eustlb

### Information

- [ ] The official example scripts
- [x] 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

Hi,
I am encountering an issue where the outputs of flash_attention_2 and sdpa attention implementations are inconsistent when an attention mask is used. Below is the code I used to test this:

```
import torch
import torch.nn as nn
from transformers import WhisperModel

def make_pad_mask(lengths: torch.Tensor) -> torch.Tensor:
    """
    Args:
      lengths:
        A 1-D tensor containing sentence lengths.
    Returns:
      Return a 2-D bool tensor, where masked positions
      are filled with `True` and non-masked positions are
      filled with `False`.

    >>> lengths = torch.tensor([1, 3, 2, 5])
    >>> make_pad_mask(lengths)
    tensor([[False,  True,  True,  True,  True],
            [False, False, False,  True,  True],
            [False, False,  True,  True,  True],
            [False, False, False, False, False]])
    """
    assert lengths.ndim == 1, lengths.ndim

    max_len = lengths.max()
    n = lengths.size(0)

    expaned_lengths = torch.arange(max_len).expand(n, max_len).to(lengths)

    return expaned_lengths >= lengths.unsqueeze(1)


def whisper_encoder_dynamic_length_forward(
    self,
    input_features,
    attention_mask=None,
    head_mask=None,
    output_attentions=None,
    output_hidden_states=None,
    return_dict=None,
):
    expected_seq_length = (
        self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0]
    )
    if input_features.shape[-1] > expected_seq_length:
        raise ValueError(
            f"Whisper expects the mel input features to be of length less equal {expected_seq_length}, but found {input_features.shape[-1]}."
        )

    output_attentions = (
        output_attentions
        if output_attentions is not None
        else self.config.output_attentions
    )
    output_hidden_states = (
        output_hidden_states
        if output_hidden_states is not None
        else self.config.output_hidden_states
    )
    return_dict = (
        return_dict if return_dict is not None else self.config.use_return_dict
    )
    inputs_embeds = nn.functional.gelu(self.conv1(input_features))
    inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))
    
    inputs_embeds = inputs_embeds.permute(0, 2, 1)
    embed_pos = self.embed_positions.weight

    hidden_states = (
        inputs_embeds + embed_pos[: inputs_embeds.shape[1], :]
    )  # NOTE: modified here
    hidden_states = nn.functional.dropout(
        hidden_states, p=self.dropout, training=self.training
    )

    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None

    # check if head_mask has a correct number of layers specified if desired
    if head_mask is not None:
        assert head_mask.size()[0] == (
            len(self.layers)
        ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."

    for idx, encoder_layer in enumerate(self.layers):
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)
        # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
        to_drop = False
        if self.training:
            dropout_probability = torch.rand([])
            if dropout_probability < self.layerdrop:  # skip the layer
                to_drop = True

        if to_drop:
            layer_outputs = (None, None)
        else:
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    encoder_layer.__call__,
                    hidden_states,
                    attention_mask,
                    (head_mask[idx] if head_mask is not None else None),
                    output_attentions,
                )
            else:
                layer_outputs = encoder_layer(
                    hidden_states,
                    attention_mask,
                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                    output_attentions=output_attentions,
                )

            hidden_states = layer_outputs[0]

        if output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)

    hidden_states = self.layer_norm(hidden_states)
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)

    if not return_dict:
        return tuple(
            v for v in [hidden_states, encoder_states, all_attentions] if v is not None
        )
    return BaseModelOutput(
        last_hidden_state=hidden_states,
        hidden_states=encoder_states,
        attentions=all_attentions,
    )

def forward_whisper_encoder_sdpa(
    encoder,
    x: torch.Tensor,
    x_lens: torch.Tensor,
):
    x = x.permute(0, 2, 1)

    audio_feat_lengths = (x_lens - 1) // 2 + 1
    encoder_out_lens = audio_feat_lengths

    batch_size, _, max_mel_seq_len = x.shape
    max_seq_len = (max_mel_seq_len - 1) // 2 + 1
    # Create a sequence tensor of shape (batch_size, max_seq_len)
    seq_range = (
        torch.arange(
            0,
            max_seq_len,
            dtype=audio_feat_lengths.dtype,
            device=audio_feat_lengths.device,
        )
        .unsqueeze(0)
        .expand(batch_size, max_seq_len)
    )
    lengths_expand = audio_feat_lengths.unsqueeze(1).expand(batch_size, max_seq_len)
    #Create mask
    padding_mask = seq_range >= lengths_expand

    audio_attention_mask_ = padding_mask.view(batch_size, 1, 1, max_seq_len).expand(
        batch_size, 1, max_seq_len, max_seq_len
    )
    audio_attention_mask = audio_attention_mask_.to(
        dtype=encoder.conv1.weight.dtype,
        device=encoder.conv1.weight.device,
    )
    audio_attention_mask[audio_attention_mask_] = float("-inf")
    encoder_out = encoder(
        input_features=x, attention_mask=audio_attention_mask
    ).last_hidden_state

    assert encoder_out_lens.max() == encoder_out.size(1)

    return encoder_out

def forward_whisper_encoder_flashatten2(
    encoder,
    x: torch.Tensor,
    x_lens: torch.Tensor,
):
    x = x.permute(0, 2, 1)

    audio_feat_lengths = (x_lens - 1) // 2 + 1
    attention_mask = ~make_pad_mask(audio_feat_lengths)
    print(attention_mask)
    encoder_out = encoder(
        input_features=x, attention_mask=attention_mask
    ).last_hidden_state

    return encoder_out

device = torch.device("cuda:0")

speech_encoder_sdpa = WhisperModel.from_pretrained("speech_ssl_models/whisper-large-v3", attn_implementation='sdpa', torch_dtype=torch.float16).encoder
speech_encoder_sdpa.__class__.forward = whisper_encoder_dynamic_length_forward
speech_encoder_sdpa.to(device)

speech_encoder_flash = WhisperModel.from_pretrained("speech_ssl_models/whisper-large-v3", attn_implementation='flash_attention_2', torch_dtype=torch.float16).encoder
speech_encoder_flash.__class__.forward = whisper_encoder_dynamic_length_forward
speech_encoder_flash.to(device)

x = torch.randn(2, 100, 128).half().to(device)
x_lens = torch.tensor([100, 60]).long().to(device)

out1 = forward_whisper_encoder_sdpa(speech_encoder_sdpa, x, x_lens)
out2 = forward_whisper_encoder_flashatten2(speech_encoder_flash, x, x_lens)

print("*******************whisper-sdpa**********************************\n", out1.shape, out1)
print("*******************whisper-flash-attentions2*********************\n", out2.shape, out2)
print("********************outputs allclose*****************************\n",torch.allclose(out1,out2,atol=1e-5))

```

### Expected behavior

The outputs of the two attention implementations are not consistent when an attention mask is used. Could you please help me understand why this is happening and how to resolve it?

Thank you!

You can replace the model path "speech_ssl_models/whisper-large-v3" with the appropriate model path you are using.

Comments:
Possibly related to #36477

On a second look the bounds are pretty small (atol 1e-5). The numerical differences might be inherently larger especially the bigger the model the more those inaccuracies will accumulate. 

> On a second look the bounds are pretty small (atol 1e-5). The numerical differences might be inherently larger especially the bigger the model the more those inaccuracies will accumulate.

Thanks for the reply. However, this is not the case. Here are the reproduction results after increasing atol to 1e-1:
```
*******************whisper-sdpa**********************************
torch.Size([2, 50, 1280]) tensor([[[ 0.3782, 0.2109, -0.0705, ..., 0.2100, -0.1501, -0.3308],
     [ 0.3062, -0.0104, -0.1207, ..., 0.0685, -0.2141, -0.0621],
     [ 0.4434, -0.1312, -0.1172, ..., -0.2229, -0.2133, -0.0853],
     ...,
     [ 0.3545, 0.0716, -0.5259, ..., 0.0616, -0.1079, -0.2010],
     [ 0.3269, 0.0684, -0.4714, ..., 0.2563, -0.0019, -0.3396],
     [ 0.4163, 0.0390, -0.1763, ..., -0.0389, 0.0453, -0.0035]],

    [[ 0.4048, 0.0847, 0.0086, ..., -0.1125, 0.0444, -0.2954],
     [ 0.3779, 0.0324, -0.1550, ..., 0.0249, -0.2200, -0.1788],
     [ 0.3748, -0.1028, -0.2502, ..., -0.0160, -0.0489, 0.0163],
     ...,
     [ 0.2771, 0.1885, -0.3101, ..., 0.1229, 0.0491, -0.1340],
     [ 0.2815, 0.0947, -0.2576, ..., 0.1504, -0.0321, -0.3428],
     [ 0.1241, 0.0125, -0.0092, ..., -0.0206, 0.0148, -0.0406]]],
    device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)

*******************whisper-flash-attentions2*********************
torch.Size([2, 50, 1280]) tensor([[[ 0.3792, 0.2114, -0.0708, ..., 0.2108, -0.1500, -0.3315],
     [ 0.3069, -0.0102, -0.1208, ..., 0.0687, -0.2136, -0.0629],
     [ 0.4436, -0.1316, -0.1169, ..., -0.2229, -0.2131, -0.0850],
     ...,
     [ 0.3542, 0.0715, -0.5273, ..., 0.0621, -0.1078, -0.2028],
     [ 0.3279, 0.0682, -0.4709, ..., 0.2559, -0.0014, -0.3406],
     [ 0.4158, 0.0393, -0.1766, ..., -0.0390, 0.0457, -0.0042]],

    [[ 0.4048, 0.0830, 0.0085, ..., -0.1135, 0.0453, -0.2949],
     [ 0.3789, 0.0319, -0.1552, ..., 0.0239, -0.2197, -0.1782],
     [ 0.3757, -0.1034, -0.2499, ..., -0.0165, -0.0490, 0.0168],
     ...,
     [-0.0271, 0.0034, 0.0042, ..., 0.0022, 0.0200, 0.0111],
     [-0.0403, 0.0034, -0.0043, ..., 0.0024, 0.0016, 0.0084],
     [-0.0390, 0.0064, -0.0042, ..., 0.0042, 0.0078, 0.0070]]],
    device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)

********************outputs allclose*****************************
 False
```

The developer in #36477 submitted a code fix for a similar bug:
```
causal_mask = attention_mask[:, : key_states.shape[-2]] -> causal_mask = attention_mask[:, : key_states.shape[-3]]
```

 but this is still incorrect in this case:

```
*******************whisper-sdpa**********************************
torch.Size([2, 50, 1280]) tensor([[[ 4.3506e-01, 2.7313e-02, -7.0679e-02, ..., 2.0935e-01,
     -1.6223e-01, -2.1130e-01],
     [ 4.4922e-01, -1.2524e-01, -5.6000e-02, ..., 5.3711e-02,
     -1.8323e-01, 1.5955e-01],
     [ 3.6133e-01, -8.2886e-02, -9.3933e-02, ..., 8.0872e-02,
     -2.5708e-01, 2.4277e-02],
     ...,
     [ 3.5522e-01, -1.1133e-01, -1.3159e-01, ..., 3.9886e-02,
     -8.1360e-02, -1.8408e-01],
     [ 3.6377e-01, 2.2751e-02, -1.2201e-01, ..., 4.0430e-01,
     -9.5291e-03, -2.7905e-01],
     [ 4.8682e-01, 2.9724e-02, -3.2043e-02, ..., 1.8204e-02,
     -1.0150e-01, -3.2715e-01]],

    [[ 4.1089e-01, 3.7594e-03, 3.0426e-02, ..., -1.7725e-01,
      8.2825e-02, -3.6353e-01],
     [ 4.3872e-01, -1.5149e-01, 1.3489e-02, ..., -1.6406e-01,
     -1.2421e-01, -1.1749e-01],
     [ 4.2456e-01, -3.0249e-01, -7.4646e-02, ..., 9.7107e-02,
      3.1104e-01, -1.3159e-01],
     ...,
     [ 2.6123e-01, -1.6650e-01, 1.0895e-01, ..., -3.1055e-01,
      2.4670e-01, 4.9400e-03],
     [ 3.0859e-01, 2.7561e-04, -2.3315e-02, ..., -3.9935e-04,
      5.2917e-02, -5.5817e-02],
     [ 3.2715e-01, 5.0690e-02, -1.6040e-01, ..., -1.2915e-01,
     -1.5210e-01, -8.7769e-02]]], device='cuda:0', dtype=torch.float16,
    grad_fn=<NativeLayerNormBackward0>)

*******************whisper-flash-attentions2*********************
torch.Size([2, 50, 1280]) tensor([[[ 4.3506e-01, 2.6825e-02, -7.1350e-02, ..., 2.0996e-01,
     -1.6235e-01, -2.1082e-01],
     [ 4.4971e-01, -1.2585e-01, -5.5573e-02, ..., 5.4138e-02,
     -1.8274e-01, 1.6150e-01],
     [ 3.6133e-01, -8.2825e-02, -9.3811e-02, ..., 8.0261e-02,
     -2.5684e-01, 2.4200e-02],
     ...,
     [ 3.5596e-01, -1.1072e-01, -1.3074e-01, ..., 4.0039e-02,
     -8.1360e-02, -1.8384e-01],
     [ 3.6304e-01, 2.2934e-02, -1.2183e-01, ..., 4.0503e-01,
     -9.9640e-03, -2.7930e-01],
     [ 4.8730e-01, 2.9480e-02, -3.1281e-02, ..., 1.8280e-02,
     -1.0150e-01, -3.2739e-01]],

    [[ 4.1138e-01, 5.3368e-03, 3.0624e-02, ..., -1.7822e-01,
      8.2520e-02, -3.6279e-01],
     [ 4.3799e-01, -1.5027e-01, 1.1955e-02, ..., -1.6492e-01,
     -1.2457e-01, -1.1682e-01],
     [ 4.2432e-01, -3.0151e-01, -7.3914e-02, ..., 9.6375e-02,
      3.1152e-01, -1.3037e-01],
     ...,
     [-3.3539e-02, 2.4052e-03, -4.8828e-03, ..., 2.9602e-03,
      9.1476e-03, 9.4528e-03],
     [-2.8152e-02, -4.5395e-03, 3.2406e-03, ..., -1.3189e-03,
      2.3315e-02, 1.1787e-02],
     [-3.5400e-02, -1.5676e-04, -4.1428e-03, ..., 2.5291e-03,
      7.4196e-03, 8.2474e-03]]], device='cuda:0', dtype=torch.float16,
    grad_fn=<NativeLayerNormBackward0>)

********************outputs allclose*****************************
 False
```

As you can see, in the first tensor of the batch, there is no mask on the time dimension, so the results of the two attention calculations should only differ numerically. However, in the second tensor of the batch, the last part of the time dimension is masked as False, leading to significant differences in the results after the two attention calculations.




Yea it seems that the other PR tackles a different issue. 

Just to be sure, since you add some custom code around the mask creation etc - do the implementations not differ when using no attn mask? And could you use the same base 2d attn masks for both flash attn and sdpa but add something like https://github.com/huggingface/transformers/blob/51ed61e2f05176f81fa7c9decba10cc28e138f61/src/transformers/models/bert/modeling_bert.py#L1108-L1110 in the sdpa path. 

1. When using no attn mask：

- the implementations do not differ when atol is set to 1e-1:

```
*******************whisper-sdpa**********************************
 torch.Size([2, 50, 1280]) tensor([[[ 0.3914, -0.1010, -0.1462,  ..., -0.1470, -0.1230, -0.3005],
         [ 0.3359, -0.1205, -0.0433,  ...,  0.2209, -0.1561, -0.1725],
         [ 0.3872, -0.1057, -0.1221,  ..., -0.0051, -0.1808, -0.4028],
         ...,
         [ 0.3757,  0.0111, -0.1469,  ...,  0.0546, -0.2401,  0.3813],
         [ 0.2913, -0.0045, -0.1113,  ...,  0.0077, -0.1987,  0.2749],
         [ 0.3704,  0.0057,  0.1202,  ...,  0.2708, -0.3875, -0.1508]],

        [[ 0.4426,  0.0720, -0.0415,  ...,  0.0444,  0.0572, -0.1478],
         [ 0.4270,  0.0862, -0.0427,  ...,  0.0035, -0.0510, -0.0044],
         [ 0.2325,  0.4360, -0.2979,  ..., -0.1641,  0.0156, -0.3977],
         ...,
         [ 0.4546,  0.0878, -0.0496,  ...,  0.1422, -0.2434, -0.4148],
         [ 0.4119,  0.0057, -0.3618,  ..., -0.1039, -0.3323, -0.0682],
         [ 0.3638, -0.0507, -0.1096,  ...,  0.1890, -0.3342, -0.2095]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)
*******************whisper-flash-attentions2*********************
 torch.Size([2, 50, 1280]) tensor([[[ 0.3914, -0.1011, -0.1473,  ..., -0.1472, -0.1248, -0.3015],
         [ 0.3362, -0.1194, -0.0443,  ...,  0.2211, -0.1575, -0.1730],
         [ 0.3879, -0.1055, -0.1223,  ..., -0.0047, -0.1813, -0.4043],
         ...,
         [ 0.3760,  0.0114, -0.1476,  ...,  0.0546, -0.2405,  0.3801],
         [ 0.2910, -0.0040, -0.1114,  ...,  0.0080, -0.1982,  0.2742],
         [ 0.3711,  0.0060,  0.1202,  ...,  0.2720, -0.3875, -0.1516]],

        [[ 0.4434,  0.0720, -0.0412,  ...,  0.0452,  0.0571, -0.1487],
         [ 0.4280,  0.0856, -0.0429,  ...,  0.0042, -0.0511, -0.0047],
         [ 0.2324,  0.4351, -0.2981,  ..., -0.1644,  0.0147, -0.3972],
         ...,
         [ 0.4546,  0.0874, -0.0493,  ...,  0.1420, -0.2427, -0.4143],
         [ 0.4119,  0.0053, -0.3625,  ..., -0.1036, -0.3315, -0.0685],
         [ 0.3640, -0.0507, -0.1107,  ...,  0.1888, -0.3347, -0.2106]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)
********************outputs allclose*****************************
 True
```

- the implementations differ when atol is set to 1e-2：
```
*******************whisper-sdpa**********************************
 torch.Size([2, 50, 1280]) tensor([[[ 0.3914, -0.1010, -0.1462,  ..., -0.1470, -0.1230, -0.3005],
         [ 0.3359, -0.1205, -0.0433,  ...,  0.2209, -0.1561, -0.1725],
         [ 0.3872, -0.1057, -0.1221,  ..., -0.0051, -0.1808, -0.4028],
         ...,
         [ 0.3757,  0.0111, -0.1469,  ...,  0.0546, -0.2401,  0.3813],
         [ 0.2913, -0.0045, -0.1113,  ...,  0.0077, -0.1987,  0.2749],
         [ 0.3704,  0.0057,  0.1202,  ...,  0.2708, -0.3875, -0.1508]],

        [[ 0.4426,  0.0720, -0.0415,  ...,  0.0444,  0.0572, -0.1478],
         [ 0.4270,  0.0862, -0.0427,  ...,  0.0035, -0.0510, -0.0044],
         [ 0.2325,  0.4360, -0.2979,  ..., -0.1641,  0.0156, -0.3977],
         ...,
         [ 0.4546,  0.0878, -0.0496,  ...,  0.1422, -0.2434, -0.4148],
         [ 0.4119,  0.0057, -0.3618,  ..., -0.1039, -0.3323, -0.0682],
         [ 0.3638, -0.0507, -0.1096,  ...,  0.1890, -0.3342, -0.2095]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)
*******************whisper-flash-attentions2*********************
 torch.Size([2, 50, 1280]) tensor([[[ 0.3914, -0.1011, -0.1473,  ..., -0.1472, -0.1248, -0.3015],
         [ 0.3362, -0.1194, -0.0443,  ...,  0.2211, -0.1575, -0.1730],
         [ 0.3879, -0.1055, -0.1223,  ..., -0.0047, -0.1813, -0.4043],
         ...,
         [ 0.3760,  0.0114, -0.1476,  ...,  0.0546, -0.2405,  0.3801],
         [ 0.2910, -0.0040, -0.1114,  ...,  0.0080, -0.1982,  0.2742],
         [ 0.3711,  0.0060,  0.1202,  ...,  0.2720, -0.3875, -0.1516]],

        [[ 0.4434,  0.0720, -0.0412,  ...,  0.0452,  0.0571, -0.1487],
         [ 0.4280,  0.0856, -0.0429,  ...,  0.0042, -0.0511, -0.0047],
         [ 0.2324,  0.4351, -0.2981,  ..., -0.1644,  0.0147, -0.3972],
         ...,
         [ 0.4546,  0.0874, -0.0493,  ...,  0.1420, -0.2427, -0.4143],
         [ 0.4119,  0.0053, -0.3625,  ..., -0.1036, -0.3315, -0.0685],
         [ 0.3640, -0.0507, -0.1107,  ...,  0.1888, -0.3347, -0.2106]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)
********************outputs allclose*****************************
 False
```

2. use the same base 2d attn masks for both flash attn and sdpa,  add use _prepare_4d_attention_mask_for_sdpa for sdpa:

```
# For flash attn:
    audio_feat_lengths = (x_lens - 1) // 2 + 1
    attention_mask = ~make_pad_mask(audio_feat_lengths)
    encoder_out = encoder(
        input_features=x, attention_mask=attention_mask
    ).last_hidden_state

# For sdpa:
    audio_feat_lengths = (x_lens - 1) // 2 + 1
    batch_size, _, max_mel_seq_len = x.shape
    max_seq_len = (max_mel_seq_len - 1) // 2 + 1
    attention_mask = ~make_pad_mask(audio_feat_lengths)
    audio_attention_mask = _prepare_4d_attention_mask_for_sdpa(
                attention_mask, x.dtype, tgt_len=max_seq_len)
    encoder_out = encoder(
        input_features=x, attention_mask=audio_attention_mask
    ).last_hidden_state
```
the results are:
```
 torch.Size([2, 50, 1280]) tensor([[[ 0.3916, -0.1014, -0.1471,  ..., -0.1466, -0.1238, -0.3008],
         [ 0.3357, -0.1197, -0.0441,  ...,  0.2205, -0.1571, -0.1730],
         [ 0.3877, -0.1057, -0.1222,  ..., -0.0049, -0.1810, -0.4041],
         ...,
         [ 0.3750,  0.0126, -0.1473,  ...,  0.0546, -0.2402,  0.3818],
         [ 0.2910, -0.0038, -0.1114,  ...,  0.0078, -0.1991,  0.2744],
         [ 0.3708,  0.0057,  0.1197,  ...,  0.2720, -0.3875, -0.1512]],

        [[ 0.4521,  0.1001,  0.0033,  ..., -0.0084,  0.0373, -0.1426],
         [ 0.3301,  0.1168, -0.5063,  ..., -0.2871, -0.2284, -0.0176],
         [ 0.1449,  0.4624, -0.3330,  ..., -0.0737,  0.0264, -0.3853],
         ...,
         [ 0.1848, -0.0013, -0.2993,  ...,  0.1217, -0.3250, -0.7983],
         [ 0.1735,  0.1295, -0.3853,  ..., -0.2010, -0.2888, -0.3525],
         [ 0.2339, -0.2922, -0.3665,  ...,  0.2241,  0.0216, -0.5244]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)
*******************whisper-flash-attentions2*********************
 torch.Size([2, 50, 1280]) tensor([[[ 0.3914, -0.1011, -0.1473,  ..., -0.1472, -0.1248, -0.3015],
         [ 0.3362, -0.1194, -0.0443,  ...,  0.2211, -0.1575, -0.1730],
         [ 0.3879, -0.1055, -0.1223,  ..., -0.0047, -0.1813, -0.4043],
         ...,
         [ 0.3760,  0.0114, -0.1476,  ...,  0.0546, -0.2405,  0.3801],
         [ 0.2910, -0.0040, -0.1114,  ...,  0.0080, -0.1982,  0.2742],
         [ 0.3711,  0.0060,  0.1202,  ...,  0.2720, -0.3875, -0.1516]],

        [[ 0.4526,  0.1000,  0.0032,  ..., -0.0083,  0.0372, -0.1416],
         [ 0.3306,  0.1155, -0.5054,  ..., -0.2861, -0.2278, -0.0183],
         [ 0.1453,  0.4612, -0.3328,  ..., -0.0748,  0.0270, -0.3833],
         ...,
         [-0.0340,  0.0058,  0.0025,  ..., -0.0008,  0.0114,  0.0088],
         [-0.0376,  0.0034, -0.0069,  ...,  0.0009,  0.0066,  0.0086],
         [-0.0283,  0.0038, -0.0040,  ..., -0.0006,  0.0123,  0.0087]]],
       device='cuda:0', dtype=torch.float16, grad_fn=<NativeLayerNormBackward0>)
********************outputs allclose*****************************
 False
```


For 1. it's also advisable to adjust the rtol as well - the values seem ok to me overall.

For 2. there definitely is something wrong going on there (seems like my quick idea was just doing worse). Would have to dig deeper to see what's going wrong there. Just as a quick suggestion to check which implementation might be wrong is to compare it to eager attention. Will try to look into it later.

@tartarleft So it seems that #36477 is indeed relevant, but otherwise the differences stem from how sdpa and flash attn handle padding:
- SDPA does not change any embeddings at all and carries it through all forwards.
- Flash attn 2 uses the so-called varlen function which removes unnecessary padding tokens -> calculates the attn -> adds 0 embeddings back to where the paddings where supposed to be.

This means, in essence, that all values that are non-padding tokens should be close to each other while padding tokens will have differences depending on your implementation (could be contributed due to the padding embedding having a different initial value, some operations accumulating, etc). You can check this by extracting the seq_lens and then checking for equality, e.g.:
```python
seq_len_1 = attention_mask.shape[-1] - (attention_mask[0, :] == False).sum()
seq_len_2 = attention_mask.shape[-1] - (attention_mask[1, :] == False).sum()
...
print("*******************whisper-sdpa**********************************\n", out1.shape, out1)
print("*******************whisper-flash-attentions2*********************\n", out2.shape, out2)
print("********************outputs allclose*****************************\n", torch.allclose(out1[0, :seq_len_1], out2[0, :seq_len_1], atol=1e-1))
print("********************outputs allclose*****************************\n", torch.allclose(out1[1, :seq_len_2], out2[1, :seq_len_2], atol=1e-1))
```

Depending on your hardware `atol=4e-2, rtol=4e-2` might be good bounds as well.

> [@tartarleft](https://github.com/tartarleft) So it seems that [#36477](https://github.com/huggingface/transformers/pull/36477) is indeed relevant, but otherwise the differences stem from how sdpa and flash attn handle padding:
> 
> * SDPA does not change any embeddings at all and carries it through all forwards.
> * Flash attn 2 uses the so-called varlen function which removes unnecessary padding tokens -> calculates the attn -> adds 0 embeddings back to where the paddings where supposed to be.
> 
> This means, in essence, that all values that are non-padding tokens should be close to each other while padding tokens will have differences depending on your implementation (could be contributed due to the padding embedding having a different initial value, some operations accumulating, etc). You can check this by extracting the seq_lens and then checking for equality, e.g.:
> 
> seq_len_1 = attention_mask.shape[-1] - (attention_mask[0, :] == False).sum()
> seq_len_2 = attention_mask.shape[-1] - (attention_mask[1, :] == False).sum()
> ...
> print("*******************whisper-sdpa**********************************\n", out1.shape, out1)
> print("*******************whisper-flash-attentions2*********************\n", out2.shape, out2)
> print("********************outputs allclose*****************************\n", torch.allclose(out1[0, :seq_len_1], out2[0, :seq_len_1], atol=1e-1))
> print("********************outputs allclose*****************************\n", torch.allclose(out1[1, :seq_len_2], out2[1, :seq_len_2], atol=1e-1))

Thank you for your help. I now fully understand the computation process. Both methods can maintain consistency at non-padding positions, and knowing this is sufficient for further development. However, I noticed that after inputting the attention mask, the computation time for flash_attn is significantly slower than sdpa, which seems to lose the advantage of the flash_attn algorithm design. Some results are as follows:
```
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats(device)
start_time = time.time()
out1 = forward_whisper_encoder_sdpa(speech_encoder_sdpa, x, x_lens)
end_time = time.time()
time_out1 = end_time - start_time
memory_used_sdpa = torch.cuda.max_memory_allocated(device)

torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats(device)
start_time = time.time()
out2 = forward_whisper_encoder_flashatten2(speech_encoder_flash, x, x_lens)
end_time = time.time()
time_out2 = end_time - start_time
memory_used_flash = torch.cuda.max_memory_allocated(device)

...

Time taken for forward_whisper_encoder_sdpa: 0.169308 seconds
Time taken for forward_whisper_encoder_flashatten2: 0.389256 seconds
Memory used by forward_whisper_encoder_sdpa: 2567.88 MB
Memory used by forward_whisper_encoder_flashatten2: 2696.47 MB
```
As you have mentioned in  [#36584](https://github.com/huggingface/transformers/issues/36584#issuecomment-2704300111), spda uses  xformers attn in this case. However, I found in the FlashAttention2 paper that it is faster than xFormers (could this be because the cases themselves are different?). Do you know why this might be happening? 

Indeed, but keep in mind that the other PR's patches are really needed. Atm the mask is incorrectly handled.

That's a nice observation! I'm not 100% sure (you'd need to benchmark each method's contributing time) but my hunch would be: 
- Flash attention when given an attn mask like I said has to remove and readd padding embeddings. This is pretty expensive and takes time. The attn calculation itself is very fast.
- Sdpa directly calculates the attention (likely with xformers). 

Combining this when given a small input (<2048 seq lens), I assume that the other operations in flash attn take too much time in comparison. If you were to go to longer sequences, I'd expect the times to come closer and eventually catch up / be better. 

However, like any program, there might be bugs/regressions in flash attention functions. Don't think this is the case here tho.

Thank you for sharing this insightful analysis that aligns perfectly with the performance characteristics I've observed:

- Sequence length dependency : The performance gap between Flash Attention and SDPA with attention masks indeed narrows as sequence lengths increase, with SDPA showing clear advantages in both speed and memory efficiency  at shorter sequence lengths (<3000 tokens)

- Practical implementation : For Whisper encoder's typical audio feature processing (where sequences are limited to 3000 tokens), SDPA currently provides the optimal balance of performance and memory usage in my observation

This discussion has been exceptionally valuable for refining my understanding of attention mechanism optimizations. Your technical rigor in analyzing these tradeoffs is truly appreciated!