🐛 Describe the bug

NaN values in the input to torch.nn.functional.scaled_dot_product_attention result in NaN output, even when input NaN elements are masked out.

Expected behaviour: the output of SDPA doesn't depend on masked-out values.

In the example below I first change one element of the input to 1e10 and confirm that the output doesn't change. Then I change the same input element to NaN and see that the output did change - to NaN

import torch

# Create sample input
B = 1
L = 2
D = 4
q = torch.arange(8, dtype=torch.float).reshape(B, L, D)
k = torch.arange(8, dtype=torch.float).reshape(B, L, D)
v = torch.arange(8, dtype=torch.float).reshape(B, L, D)

# Use mask of shape (B, L, L)
# Mask out second source element for both elements in target
attn_mask = torch.tensor([[[True, False],
                           [True, False]]])
res1 = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)

# Change masked out element to a very large number - the result doesn't change
k[0, 1, 0] = 1e10
res2 = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
# Result didn't change
print(f"{res1 - res2 = }")

# Change masked out element to NaN - the result becomes NaN
k[0, 1, 0] = torch.nan
res3 = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
# Result became NaN
print(f"{res1 - res3 = }")

Output:

res1 - res2 = tensor([[[0., 0., 0., 0.],
         [0., 0., 0., 0.]]])
res1 - res3 = tensor([[[nan, nan, nan, nan],
         [nan, nan, nan, nan]]])

One scenario in which this comes up is when I use torch.empty to initialize a buffer for processing input of variable length. I can fill the beginning of the buffer with meaningful values and use the mask to tell SDPA to ignore the rest. However, this wouldn't work if masked-out elements of the buffer accidentally contain NaN, because of the described behaviour

The reason might be nan + x -> nan in

or similar code for other backends.

Versions

PyTorch version: 2.1.0.dev20230520
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.3.1 (arm64)
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.24.3
Libc version: N/A

Python version: 3.8.13 (default, Oct 19 2022, 17:52:09) [Clang 12.0.0 ] (64-bit runtime)
Python platform: macOS-13.3.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 M1 Pro

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.4
[pip3] torch==2.1.0.dev20230520
[conda] numpy 1.23.4 pypi_0 pypi
[conda] pytorch 2.1.0.dev20230520 py3.8_0 pytorch-nightly