Fix scaled_dot_product_attention reference implementation to match MATH backend #163508
Conversation
MATH backend scales Q K pre-softmax, leading to numerical differences when comparing the ref impl with MATH. Now matches SDPBackend.MATH with `rotl=0., atol=0.`
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/163508
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New FailuresAs of commit 7cb5f4f with merge base 96a3afb ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "module: sdpa" |
|
@pytorchbot label "release notes: nn" |
There was a problem hiding this comment.
Thanks for the contribution! I agree with the fixes and it's nice to have bitwise-accurate validation against the math backend via tests.
One concern I have is that the test won't catch updates to the reference implementation in the docs. Running this via the doctest mechanism would address this. cc @svekars for insight on how to ensure this validation happens during doctest time
Don't think they are seeded the same way on some archs, leading to different dropout
|
Thanks! Agree, originally I had a doctest but I thought it might pollute the page since it would have to be immediately below the code block. This way it should still catch regressions to the MATH kernel, which seems more likely to breaking change. Happy to go either way though. |
If it's just a few lines to compare reference vs. math, my opinion is that this doesn't pollute the page too much. In fact, it makes it very clear that the reference in the docs is what users should expect from the math backend. |
|
Looking into this: I'm not sure there's a good way to integrate with doctest. SDPA is not defined in a The current sdpa examples aren't actually run with xdoctest, you can test this with: export XDOCTEST_GLOBAL_EXEC="from torch import nn\nimport torch.nn.functional as F\nimport torch"
xdoctest -m torch.nn.functional --analysis dynamic
<examples are run, but not sdpa>I could add a "fake" doctest in, but that might be more confusing, or I think we'd need to refactor. |
|
@drisspg Ok to be merged? |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 1 jobs have failed, first few of them are: trunk / linux-jammy-cuda12.8-py3.10-gcc11 / test (default, 2, 5, linux.g6.4xlarge.experimental.nvidia.gpu) Details for Dev Infra teamRaised by workflow job |
|
@drisspg my bad, needed to put |
torch.softmax(inp, dtype=torch.float32).to(torch.float16)is not equivalent totorch.softmax(inp)for fp16 input #123911 (document float32 upcasting behaviour)Summary
The documented reference implementation of SDPA doesn't numerically match the MATH. This causes confusion when testing numerical accuracy of kernels / code vs MATH.
Changes
Updated the reference implementation to match MATH's actual behavior. The key corrections are:
Added regression test
test_reference_implementation_bitwise_match_math_backend. Test verifies exact bitwise match (rtol=0, atol=0) between the reference implementation and MATH.