Repository: pytorch
Issue: MultiheadAttention returns incorrect results for windowed long causal mask for long input
Labels: module: nn,triaged

### 🐛 Describe the bug
Note, though the error is small (about 2e-8), but it will be magnified by other network layers, and it will eventually become large enough.

```python
from torch import nn, Tensor
import torch


def get_causal_mask(sig_len: int, mask_len: int, device=None):
    pos1 = torch.arange(start=0, end=sig_len, dtype=torch.long, device=device, requires_grad=False).unsqueeze(1)
    pos2 = torch.arange(start=0, end=sig_len, dtype=torch.long, device=device, 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]
    ]
    """
    mask = torch.where((relative_pos >= 0) * (relative_pos < mask_len), 0.0, -torch.inf)
    return mask


class MM(nn.Module):

    def __init__(self) -> None:
        super().__init__()
        self.l = nn.ModuleList([nn.MultiheadAttention(96, 4, batch_first=True) for i in range(1)])

    def forward(self, xi: Tensor, mask: Tensor):
        for m in self.l:
            xi, _ = m.forward(xi, xi, xi, attn_mask=mask)
        return xi


x = torch.randn((1, 100 * 20, 96)).cuda()
mm = MM().to(x.device)
mm.eval()

with torch.no_grad():
    print('###################')
    for masklen in [10, 80, 100]:
        print('masklen=', masklen)
        ys = []
        for i in range(2, 20, 1):
            xi = x[:, -masklen * i:, :]
            mask = get_causal_mask(xi.shape[1], masklen, xi.device)
            yi = mm(xi, mask)[0, -10:, :]
            if i > 2:
                print(f"{i-1}=={i}", (yi == ys[-1]).all().item()) # should be all True, as the causal mask has limited length
            ys.append(yi)
        print('###################')
```

Output (should always be True):
```
###################
masklen= 10
2==3 True
3==4 True
4==5 True
5==6 True
6==7 True
7==8 True
8==9 True
9==10 True
10==11 True
11==12 True
12==13 True
13==14 True
14==15 True
15==16 True
16==17 True
17==18 True
18==19 True
###################
masklen= 80
2==3 True
3==4 True
4==5 True
5==6 True
6==7 True
7==8 True
8==9 True
9==10 True
10==11 True
11==12 True
12==13 False
13==14 False
14==15 False
15==16 False
16==17 False
17==18 False
18==19 False
###################
masklen= 100
2==3 True
3==4 True
4==5 True
5==6 True
6==7 True
7==8 True
8==9 True
9==10 True
10==11 False
11==12 False
12==13 False
13==14 False
14==15 False
15==16 False
16==17 False
17==18 False
18==19 False
###################
```




### 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: 12.2.140
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: 545.23.08
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.4
[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.4                    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 @bhosmer @cpuhrsch @erichan1 @drisspg