Update fused kernels and call _safe_softmax from SDPA

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

Author: drisspg

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -68,16 +68,23 @@ def test_vit_skip_conv(self):
             )
         )
 
+        conv_block = ["aten.convolution.default", "executorch_call_delegate"]
+        safe_softmax_block = [
+            "getitem",
+            "getitem",
+            "getitem",
+            "getitem",
+            "aten.any.dim",
+            "executorch_call_delegate",
+        ]
+        final_block = ["getitem"]
+        total = conv_block + 12 * safe_softmax_block + final_block
+
         assert [
             node.target.__name__
             for node in delegated_program_manager.exported_program().graph.nodes
             if node.op == "call_function"
-        ] == [
-            "aten.convolution.default",
-            "executorch_call_delegate",
-            "getitem",
-        ]
-
+        ] == total
 
 if __name__ == "__main__":
     test_runner = TestCoreMLPartitioner()