SDPA: use exp_u20 for the softmax exponential (#22082)

[pssrawat]

Summary:

Switch the softmax exponential in the flash-attention custom SDPA
kernel (op_sdpa_impl.h) from Vectorized::exp() to
Vectorized::exp_u20(). The file already carried this as a TODO.

This is not bit-exact. exp_u20 is a ULP-20 approximation, so it differs
from exp outright, and fourteen autoregressive layers on that is enough
to flip a near-tie argmax.

Differential Revision: D117198987

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Cancelled Job, 1 Unrelated Failure, 13 Unclassified Failures

As of commit a8cac3c with merge base 2287a8b ():

NEW FAILURE - The following job has failed:

• Cadence Build & Test / hifi-build / hifi4 (gh)
  ##[error]Refusing to check out fork pull request code from a 'pull_request_target' workflow. This workflow runs with the base repository's GITHUB_TOKEN, secrets, default-branch cache scope, and runner access. Fetching and executing a fork's code in that trusted context commonly leads to "pwn request" vulnerabilities. To opt in, review the risks at https://gh.io/securely-using-pull_request_target and set 'allow-unsafe-pr-checkout: true' on the actions/checkout step.

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

• Build Aarch64 Linux CUDA Wheels / pytorch/executorch / build-wheel-py3_10-cuda-aarch6413_0-aarch64 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Build Aarch64 Linux CUDA Wheels / pytorch/executorch / upload / upload-wheel-py3_10-cuda-aarch6413_0-aarch64 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  Unable to download artifact(s): Artifact not found
• Build Aarch64 Linux Wheels / pytorch/executorch / build-wheel-py3_10-cpu-aarch64 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Build Aarch64 Linux Wheels / pytorch/executorch / upload / upload-wheel-py3_10-cpu-aarch64 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  Unable to download artifact(s): Artifact not found
• Build Linux CUDA Wheels / pytorch/executorch / build-manywheel-py3_10-cuda13_0 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Build Linux CUDA Wheels / pytorch/executorch / upload / upload-manywheel-py3_10-cuda13_0 (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  Unable to download artifact(s): Artifact not found
• Build Linux Wheels / pytorch/executorch / build-manywheel-py3_10-cpu (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Build Linux Wheels / pytorch/executorch / upload / upload-manywheel-py3_10-cpu (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  Unable to download artifact(s): Artifact not found
• Build macOS Wheels / pytorch/executorch / build-wheel-py3_10-cpu (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]The operation was canceled.
• Build macOS Wheels / pytorch/executorch / upload / upload-wheel-py3_10-cpu (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  Unable to download artifact(s): Artifact not found
• Build Windows Wheels / pytorch/executorch / build-wheel-py3_10-cpu (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Build Windows Wheels / pytorch/executorch / upload / upload-wheel-py3_10-cpu (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  Unable to download artifact(s): Artifact not found
• Check Labels (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)

CANCELLED JOB - The following job was cancelled. Please retry:

• Build documentation / build (buck2) / Build doc (gh)
  ##[error]The operation was canceled.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Cadence Build & Test / vision-build / vision (gh) (detected as infra flaky with no runner)

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

[meta-codesync]

@pssrawat has exported this pull request. If you are a Meta employee, you can view the originating Diff in D117198987.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.