Repository: pytorch
Issue: Issue: Illegal Memory Access in Backward Pass of `scaled_dot_product_attention` with Custom Attention Mask
Labels: triaged,module: sdpa

### 🐛 Describe the bug

**Bug Description:**
When using a custom attention mask in the `scaled_dot_product_attention` function, an illegal memory access error (`an illegal memory access was encountered`) occurs during the backward pass when the sequence length of `QK` (query-key) is greater than or equal to 65,536.

**reproducing code:**
```python
import torch
import torch.nn.functional as F
def torch_attention(q, k,v, n, ks, ts, upcast=False):
	if upcast:
		q = q.to(torch.float32)
		k = k.to(torch.float32)
		v = v.to(torch.float32)
	attn_mask = generate_mask(n, ks, max((n-ts)//ks,0), ts, "cuda", dtype=q.dtype)
	with torch.backends.cuda.sdp_kernel(enable_flash=True):
		attention_torch = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
	return attention_torch

def generate_mask(seq_len, ks, nk, ts, device='cpu', dtype=torch.bfloat16):
	row_ind = torch.arange(seq_len, device=device, dtype=torch.long)+1
	k_ind = torch.arange(nk, device=device, dtype=torch.long)+1
	mask_k = (row_ind.unsqueeze(1)>ts) * (torch.floor_divide(row_ind-ts, ks).unsqueeze(1) >= k_ind.unsqueeze(0))
	
	col_ind = torch.arange(seq_len, device=device, dtype=torch.long)+1
	ts = torch.tensor([ts]*seq_len, device=device, dtype=torch.long)
	nking_token = torch.maximum(torch.floor_divide(row_ind-ts, ks)*ks, torch.tensor([0]*seq_len, device=device, dtype=torch.long))
	remain_num = torch.maximum(row_ind-nking_token-ts, torch.tensor([0]*seq_len, device=device, dtype=torch.long))
	ts = ts+remain_num
	mask_t = (row_ind.unsqueeze(1)>=col_ind.unsqueeze(0)) * ((row_ind-ts).unsqueeze(1)<col_ind.unsqueeze(0))
	bool_mask = torch.concat([mask_k, mask_t], dim=1)
	final_mask = torch.zeros((seq_len, seq_len+nk), device=device, dtype=dtype)
	final_mask = torch.masked_fill(final_mask, ~bool_mask, -torch.inf)
	return final_mask


def test_torch_attn(bz, h, n, d, ks, ts):
	print(f"{bz=}, {h=}, {n=}, {d=}, {ks=}, {ts=}")
	nk = (n-ts)//ks
	torch.manual_seed(20)
	q = (torch.empty((bz, h, n, d), dtype=torch.bfloat16, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
	k = (torch.empty((bz, h, n+nk, d), dtype=torch.bfloat16, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
	v = (torch.empty((bz, h, n+nk, d), dtype=torch.bfloat16, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())
	do = torch.randn((bz, h, n, d), dtype=torch.bfloat16, device="cuda")
	
	attention_torch = torch_attention(q,k, v, n, ks, ts, upcast=True)

	gq_torch, gk_torch, gv_torch = torch.autograd.grad(attention_torch, (q, k, v), do)
	print(gq_torch-torch.zeros_like(gq_torch))

test_torch_attn(1,32,1024*64,128, ks=16, ts=1024)
```


**Erros:**
```
Traceback (most recent call last):
  File "/debug_report.py", line 47, in test_torch_attn
    print(gq_torch-torch.zeros_like(gq_torch))
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor.py", line 461, in __repr__
    return torch._tensor_str._str(self, tensor_contents=tensor_contents)
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 677, in _str
    return _str_intern(self, tensor_contents=tensor_contents)
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 597, in _str_intern
    tensor_str = _tensor_str(self, indent)
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 349, in _tensor_str
    formatter = _Formatter(get_summarized_data(self) if summarize else self)
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 387, in get_summarized_data
    return torch.stack([get_summarized_data(x) for x in self])
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 387, in <listcomp>
    return torch.stack([get_summarized_data(x) for x in self])
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 385, in get_summarized_data
    return torch.stack([get_summarized_data(x) for x in (start + end)])
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 385, in <listcomp>
    return torch.stack([get_summarized_data(x) for x in (start + end)])
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 385, in get_summarized_data
    return torch.stack([get_summarized_data(x) for x in (start + end)])
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 385, in <listcomp>
    return torch.stack([get_summarized_data(x) for x in (start + end)])
  File "/opt/conda/envs/torch220/lib/python3.9/site-packages/torch/_tensor_str.py", line 375, in get_summarized_data
    return torch.cat(
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

### Versions

This issue is observed on an **NVIDIA A800 GPU** with **PyTorch 2.2.0**.