introduce triton sdpa kernel to cuda backend

[Gasoonjia]

Summary:
Introduce Triton SDPA Kernel to CUDA Backend

This diff introduces a kernel-generator (https://github.com/meta-pytorch/KernelAgent) driven, Triton-optimized implementation of scaled dot-product attention (SDPA) kernel to the CUDA backend. The new kernel is designed to replace the default Edge SDPA operator during graph transformation to accelerate the model inference and get rid of sdpa decomposition.

Changes

• Added a new file sdpa.py to fbcode/executorch/backends/cuda/triton/kernels and fbcode/executorch/backends/cuda/triton/kernels directories, which contains the Triton-optimized SDPA kernel implementation.
• Added a new fbcode/executorch/backends/cuda/triton/replacement_pass, which replaces the given existing edge ops with target triton kernels.
• Added tests for sdpa exporting with triton kernel. Without the triton kernel, sdpa model can not be exported.

Purpose

The purpose of this diff is to provide a high-performance SDPA kernel for the CUDA backend, which can be used to accelerate attention-based models on NVIDIA GPUs

Reviewed By: larryliu0820

Differential Revision: D87259044

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 2 Unrelated Failures

As of commit 35f22b8 with merge base f66b79f ():

NEW FAILURE - The following job has failed:

• Test CUDA Builds / test-model-cuda-e2e (openai, whisper-small, non-quantized) / linux-job (gh)
  RuntimeError: Command docker exec -t b9d54eeadc73cfde6d041139bc5cddd7c249da199fea94ba7bbe18aad047168b /exec failed with exit code 1

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

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

• pull / unittest / linux / linux-job (gh) (trunk failure)
  exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_emit_nested_lowered_backend_module
• pull / unittest-editable / linux / linux-job (gh) (trunk failure)
  exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_emit_nested_lowered_backend_module

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

[meta-codesync]

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

[larryliu0820]

Review automatically exported from Phabricator review in Meta.

[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.

[meta-codesync]

@Gasoonjia has imported this pull request. If you are a Meta employee, you can view this in D87259044.