introduce cuda sdpa (#15996)

Summary: Pull Request resolved: #15996

Test Plan: buck2 test mode/dev-nosan fbcode//executorch/backends/cuda/runtime/shims/tests:test_aoti_torch_cuda_scaled_dot_product_attention

Differential Revision: D87950475

Pulled By: Gasoonjia

Author: Gasoonjia

.ci/docker/ci_commit_pins/optimum-executorch.txt

@@ -1 +1 @@
-d03e90c2cd9048e6d9a75285c0355f033cd016fc
+8967fe914c252bf242b7d0ad4f5e098a007a6993

backends/cuda/cuda_backend.py

@@ -39,6 +39,8 @@ def get_device_name(cls) -> str:
     def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
         return {
             "at::_ops::_weight_int4pack_mm::call": None,
+            "at::_ops::_scaled_dot_product_flash_attention::call": None,
+            "at::_ops::_scaled_dot_product_efficient_attention::call": None,
         }
 
     @classmethod
@@ -68,7 +70,8 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                     )
                 triton_kernel_mode = mode
 
-        return [ReplaceEdgeOpWithTritonOpPass()] if triton_kernel_mode == "ON" else []
+        return []
+        # return [ReplaceEdgeOpWithTritonOpPass()] if triton_kernel_mode == "ON" else []
 
     @classmethod
     def get_aoti_compile_options(

backends/cuda/runtime/TARGETS

@@ -53,6 +53,7 @@ runtime.cxx_library(
         "shims/cuda_guard.cpp",
         "shims/int4mm.cu",
         "shims/memory.cpp",
+        "shims/sdpa.cu",
         "shims/tensor_attribute.cpp",
     ],
     headers = [
@@ -61,6 +62,8 @@ runtime.cxx_library(
         "shims/int4mm.cuh",
         "shims/int4mm.h",
         "shims/memory.h",
+        "shims/sdpa.cuh",
+        "shims/sdpa.h",
         "shims/tensor_attribute.h",
         "utils.h",
     ],
@@ -84,6 +87,7 @@ runtime.cxx_library(
     ],
     external_deps = [
         ("cuda", None, "cuda-lazy"),
+        ("cuda", None, "cublas-lazy"),
     ],
 )