Update fused kernels and call _safe_softmax from SDPA

[drisspg]

Summary:
X-link: pytorch/pytorch#131863

cc ezyang gchanan jgong5 mingfeima XiaobingSuper sanchitintel ashokei jingxu10 Valentine233

imported-using-ghimport

Reviewed By: Chillee

Differential Revision: D61418679

Pulled By: drisspg

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (30 Unrelated Failures)

As of commit 388a062 with merge base 6cb5726 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Apple / build-frameworks-ios / macos-job (gh) (similar failure)
• Apple / test-demo-ios / macos-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-binary-size-linux / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-binary-size-linux-gcc / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-custom-ops-linux (buck2) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-custom-ops-linux (cmake) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-export-llava-linux / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-llama-runner-linux (fp32, buck2, portable) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-llama-runner-linux (fp32, buck2, xnnpack+custom) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-llama-runner-linux (fp32, buck2, xnnpack+custom+qe) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-llama-runner-linux (fp32, cmake, portable) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-llama-runner-linux (fp32, cmake, xnnpack+custom) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-llama-runner-linux (fp32, cmake, xnnpack+custom+qe) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-llama-runner-linux-android (cmake) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-llama-runner-qnn-linux (fp32, cmake, qnn) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-models-linux (buck2, mv3, portable, linux.2xlarge, 90) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-models-linux (buck2, mv3, xnnpack-quantization-delegation, linux.2xlarge, 90) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-models-linux (cmake, mv3, portable, linux.2xlarge, 90) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-models-linux (cmake, mv3, xnnpack-quantization-delegation, linux.2xlarge, 90) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-models-linux (cmake, vit, portable, linux.2xlarge, 90) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-models-linux (cmake, vit, xnnpack-delegation, linux.2xlarge, 90) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-pybind-build-linux (cmake) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-quantized-aot-lib-linux (cmake) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-selective-build-linux (buck2) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-selective-build-linux (cmake) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-setup-linux-gcc (cmake) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / unittest / linux / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / unittest / macos / macos-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / unittest-arm (buck2) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Build documentation / build (buck2) / Build doc (gh) (trunk failure)

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

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D61418679

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D61418679

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D61418679

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D61418679

[shoumikhin]

@YifanShenSZ ptal when you have a chance!

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D61418679

[YifanShenSZ]

Awesome safe softmax 🎊 & thanks for the subsequent fix in executorch 🙏

[drisspg]

@shoumikhin not totally sure what to do about how red executorch trunk is on this one