🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/140682

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit 9f8c51c with merge base 99c8d5a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

The segfault can be reproduced on ROCM platform reliably with the following env vars

PYTORCH_NO_HIP_MEMORY_CACHING=1 HSA_SVM_GUARD_PAGES=1 HSA_DISABLE_FRAGMENT_ALLOCATOR=1 AMD_SERIALIZE_KERNEL=3 

, which add guard pages to detect OOB memory access.

@jithunnair-amd @pruthvistony @jataylo @jeffdaily
I'm speculating our large SDPA error is partially due to the OOB access in the softmax operator, which reads extra data when the input is irregular.

If you have a reproducer, could you add a test case for this change please?

@eqy I have a testing case in mind, but where should I add it? (or according to @malfet we cannot add it)

I have a testing case in mind, but where should I add it? (or according to @malfet we cannot add it)

I think PR description should be sufficient, testing it with compute sanitizer now. And it fails:

PYTORCH_NO_CUDA_MEMORY_CACHING=1 /usr/local/cuda-12.3/bin/compute-sanitizer python -c "import torch;qk = torch.randn((1024,587), dtype=torch.float64, device='cuda'); print(torch.softmax(qk, dim=-1))"
========= COMPUTE-SANITIZER
========= Invalid __global__ read of size 8 bytes
=========     at void at::native::<unnamed>::cunn_SoftMaxForward<(int)2, double, double, double, at::native::<unnamed>::SoftMaxForwardEpilogue>(T4 *, const T2 *, int)+0x190
=========     by thread (588,0,0) in block (1023,0,0)
=========     Address 0x7fcb5f696000 is out of bounds
=========     and is 1 bytes after the nearest allocation at 0x7fcb5f200000 of size 4808704 bytes
...

@xinyazhang please add some sort of a test to the PR description, but otherwise looks good to me (we can not have unittest unfortunately, because at least on CUDA something like that would not be detectable without compute sanitizer)

Would we get some signal with torch.cuda.empty_cache() around the test?

@xinyazhang please add some sort of a test to the PR description, but otherwise looks good to me (we can not have unittest unfortunately, because at least on CUDA something like that would not be detectable without compute sanitizer)

Would we get some signal with torch.cuda.empty_cache() around the test?

Nope, as GPU maps memory into its address space in a pretty large chunks(A100 page size is probably 2Mb or something), see

$ PYTORCH_NO_CUDA_MEMORY_CACHING=1 python -c "import torch;qk = torch.randn((1024,587), dtype=torch.float64, device='cuda'); print(torch.softmax(qk, dim=-1))"; echo $?
0

vs launching the same with compute sanitizer, that will catch an error

@pytorchbot merge -f "Looks reasonable"

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

For documentation @malfet @eqy @jeffdaily @jithunnair-amd here is the independent unit test code.

#!/usr/bin/env python

import torch
import pytest

@pytest.mark.parametrize("dtype",
                         [torch.float16, torch.bfloat16, torch.float32, torch.float64],
                         ids=['fp16', 'bf16', 'fp32', 'fp64'])
def test_softmax_oob_access(dtype):
    qk_with_margin = torch.randn((1024+1, 587), dtype=dtype, device="cuda");
    qk_with_margin[-1].fill_(float('nan'))
    qk = qk_with_margin[:-1, :]
    smqk = torch.softmax(qk, dim=-1)
    assert not torch.isnan(smqk).any(), 'NaN indicates OOB memory access'

It only fails on float64 apparently, (but our SDPA's UT suffers from it and has to use CPU implementation on certain cases)

@xinyazhang that's smart and indeed reliable, please do not hesitate to propose a PR that adds this unites