apply simple sdpa to coreml/mps backend (#3660)

Summary:
Pull Request resolved: #3660

coreml and mps doesn't support sdpa at the moment, use simple sdpa to have a simpler decomposition. Observer 1.5x faster on emulator.

Reviewed By: shoumikhin, kirklandsign

Differential Revision: D57476985

fbshipit-source-id: 2dbcad1a6e8b744e0a95d60fcd740369d665eab0

Author: cccclai

examples/models/llama2/export_llama_lib.py

@@ -346,9 +346,12 @@ def _prepare_for_llama_export(modelname: str, args) -> LlamaEdgeManager:
     if args.use_sdpa_with_kv_cache:
         transforms.append(replace_sdpa_with_custom_op)
 
-    if args.qnn and args.use_kv_cache:
-        transforms.append(replace_sdpa_with_simple_sdpa)
-        transforms.append(replace_causal_mask)
+    if args.use_kv_cache:
+        if args.qnn or args.coreml or args.mps:
+            # Currently qnn/coreml/mps doesn't support sdpa op, use the simpler decomposition
+            # to get free perf gain.
+            transforms.append(replace_sdpa_with_simple_sdpa)
+            transforms.append(replace_causal_mask)
     return (
         load_llama_model(
             modelname=modelname,