Call _safe_softmax from sdpa math path

[drisspg]

Stack from ghstack (oldest at bottom):

• -> Call _safe_softmax from sdpa math path #131844
• Update fused mem-eff kernel to set masked out rows to 0 #131867
• Update fused kernels and call _safe_softmax from SDPA #131863
• add impls for required for nt ops #132710
• Add a private _safe_softmax #131060

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❌ 3 New Failures

As of commit 49c0f12 with merge base 4a2cf50 ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner-noclang / linux-job (gh)
  >>> Lint for torch/_decomp/decompositions.py:
• pull / linux-focal-py3_8-clang9-xla / test (xla, 1, 1, amz2023.linux.12xlarge) (gh)
  test_sdpa_pattern
• pull / linux-focal-py3.8-clang10-onnx / test (default, 2, 2, amz2023.linux.2xlarge) (gh)
  'test/onnx/test_fx_to_onnx.py::TestFxToOnnx::test_checkpoint_cast'

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

[jbschlosser]

it's happening :0

[jbschlosser]

wheee