Repository: pytorch
Issue: MultiheadAttention returns different results for binary mask and float mask
Labels: module: nn,triaged

### 🐛 Describe the bug
1. `MultiheadAttention` returns different results for binary mask and float mask in `eval` mode
2. For binary mask, `MultiheadAttention` returns different results in `train` mode and `eval` mode


```python
from torch import nn 
import torch 

x = torch.randn((1, 251, 96))
slen=251

pos1 = torch.arange(start=0, end=slen, dtype=torch.long, requires_grad=False).unsqueeze(1)
pos2 = torch.arange(start=0, end=slen, dtype=torch.long, requires_grad=False).unsqueeze(0)
relative_pos = pos1 - pos2
""" now, relative_pos=[
[0,-1,-2,...,-(T-1)],
[1, 0,-1,...,-(T-2)],
...
[T-1,T-2,...,  1, 0]
]
"""
attn_scope=251
maskb = ~((relative_pos >= 0) * (relative_pos < attn_scope)) # Lower triangular matrix 
maskf = torch.where((relative_pos >= 0) * (relative_pos < attn_scope), 0.0, -torch.inf) # Lower triangular matrix

mhsa=nn.MultiheadAttention(96,4,batch_first=True)

ybt,attnb=mhsa.forward(x,x,x, need_weights=False, average_attn_weights=False, attn_mask=maskb, is_causal=True)
yft,attnf=mhsa.forward(x,x,x, need_weights=False, average_attn_weights=False, attn_mask=maskf, is_causal=True)
print((ybt==yft).all(),(attnb==attnf)) # True, True

mhsa.eval()
with torch.no_grad():
    ybe,attnb=mhsa.forward(x,x,x, need_weights=False, average_attn_weights=False, attn_mask=maskb, is_causal=True)
    yfe,attnf=mhsa.forward(x,x,x, need_weights=False, average_attn_weights=False, attn_mask=maskf, is_causal=True)
    print((ybe==yfe).all(),(attnb==attnf)) # False, True
print((ybt==ybe).all()) # False: the output in training mode doesn't match the output in evaluation mode for binary mask
print((yft==yfe).all()) # True
```

### Versions

Collecting environment information...
PyTorch version: 2.1.2
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 10.5.0-1ubuntu1~22.04) 10.5.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-89-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.5.119
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 535.154.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.971
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.26.0
[pip3] pytorch-lightning==2.1.3
[pip3] torch==2.1.2
[pip3] torchaudio==2.1.2
[pip3] torchmetrics==1.2.0
[pip3] torchtnt==0.2.1
[pip3] torchvision==0.16.2
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-service               2.4.0           py310h5eee18b_1  
[conda] mkl_fft                   1.3.8           py310h5eee18b_0  
[conda] mkl_random                1.2.4           py310hdb19cb5_0  
[conda] numpy                     1.26.0          py310h5f9d8c6_0  
[conda] numpy-base                1.26.0          py310hb5e798b_0  
[conda] pytorch                   2.1.2           py3.10_cuda12.1_cudnn8.9.2_0    pytorch
[conda] pytorch-cuda              12.1                 ha16c6d3_5    pytorch
[conda] pytorch-lightning         2.1.3                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.1.2               py310_cu121    pytorch
[conda] torchmetrics              1.2.0                    pypi_0    pypi
[conda] torchtnt                  0.2.1                    pypi_0    pypi
[conda] torchtriton               2.1.0                     py310    pytorch
[conda] torchvision               0.16.2              py310_cu121    pytorch

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @drisspg @bhosmer @cpuhrsch @erichan1