Repository: pytorch
Issue: [BUG] Given boolean tgt_mask, TransformerDecoder produces wrong results with MPS backend
Labels: module: nn,triaged,module: mps

### 🐛 Describe the bug

Two tokens are decoded in this example. Ideally, the output feature on the first token should be the same regardless of the sequence length as a square subsequent mask is applied. Here are two ways to generate the tgt mask. One is from the official example where a float mask is generated first but converted into a boolean mask later. Another way is to generate the boolean mask directly.

:heavy_check_mark: If running with `cpu` backend, the three printed values should be the same:
```
Token_1 (Ground Truth): [[[ 1.1367519  -0.7223548   0.81752294 -1.2319201 ]]]
Ground Truth: [ 1.136752  -0.7223549  0.8175233 -1.2319204]
Buggy Results: [ 1.136752  -0.7223549  0.8175233 -1.2319204]
``` 

❌  However, if we run it on `mps` backend, the directly generated boolean mask would result in the wrong decoded feature for the first token. The result looked like this:
```
Token_1 (Ground Truth): [[[ 1.4028326  -1.0564857  -0.8273474   0.48100036]]]
Ground Truth: [ 1.4028326  -1.0564857  -0.8273474   0.48100036]
Buggy Results: [ 1.3671896  -1.1002903  -0.79846895  0.5315697 ]
```

The bug can be reproduced by this minimal script:
```python
import torch
import torch.nn as nn

# Define the parameters
dropout = 0.  # Dropout rate
d_model = 4
device = "mps"
# device = "cpu"
torch.manual_seed(0)
decoder_layer = nn.TransformerDecoderLayer(batch_first=True, d_model=d_model, nhead=1, dim_feedforward=32,
                                           dropout=dropout)
transformer_decoder = nn.TransformerDecoder(decoder_layer=decoder_layer, num_layers=1)
transformer_decoder = transformer_decoder.to(device)
for p in transformer_decoder.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)

# Dummy inputs
tgt = torch.rand(1, 2, d_model, device=device)  # [target sequence length, batch size, embedding size]
memory = torch.rand(1, 20, d_model, device=device)  # [source sequence length, batch size, embedding size]

# generate mask
def _mask_to_float(mask):
    return mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))


def generate_square_subsequent_mask(sz):
    mask = (torch.triu(torch.ones((sz, sz))) == 1).transpose(0, 1)
    return _mask_to_float(mask).bool()


def _generate_square_subsequent_mask_bug(sz):
    mask = (torch.triu(torch.ones((sz, sz))) == 1).transpose(0, 1)
    return ~mask


token_1_mask = generate_square_subsequent_mask(1).to(device)
correct_mask = generate_square_subsequent_mask(2).to(device)
bug_mask = _generate_square_subsequent_mask_bug(2).to(device)

# The masks are all boolean type and are exactly same
assert correct_mask.dtype == bug_mask.dtype
assert torch.all(correct_mask == bug_mask)

# Forward pass through the decoder
token_1 = transformer_decoder(tgt=tgt[:, :1, :], memory=memory, tgt_mask=token_1_mask)
output_bug = transformer_decoder(tgt=tgt, memory=memory, tgt_mask=bug_mask)
output = transformer_decoder(tgt=tgt, memory=memory, tgt_mask=correct_mask)

print("Token_1 (Ground Truth):", token_1.detach().cpu().numpy())
print("Ground Truth:", output[0, 0].detach().cpu().numpy())
print("Buggy Results:", output_bug[0, 0].detach().cpu().numpy())

```

### Versions

PyTorch version: 2.2.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.1.0.2.5)
CMake version: version 3.28.3
Libc version: N/A

Python version: 3.11.7 (main, Dec 15 2023, 12:09:56) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M3 Max

Versions of relevant libraries:
[pip3] numpy==1.26.3
[pip3] pytorch-lightning==2.2.0.post0
[pip3] torch==2.2.1
[pip3] torchaudio==2.2.1
[pip3] torchdata==0.7.1
[pip3] torchmetrics==1.3.0.post0
[pip3] torchtext==0.17.1
[pip3] torchvision==0.17.1
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] numpy-base                1.26.3          py311hfbfe69c_0
[conda] pytorch-lightning         2.2.0.post0              pypi_0    pypi
[conda] torch                     2.2.1                    pypi_0    pypi
[conda] torchaudio                2.2.1                    pypi_0    pypi
[conda] torchdata                 0.7.1                    pypi_0    pypi
[conda] torchmetrics              1.3.0.post0              pypi_0    pypi
[conda] torchtext                 0.17.1                   pypi_0    pypi
[conda] torchvision               0.17.1                   pypi_0    pypi

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @kulinseth @malfet @DenisVieriu97 @jhavukainen @bhosmer @cpuhrsch @erichan1 @drisspg @razarmehr