Update on "Add ONNX Runtime GQA-style SDPA benchmark"

Add run_onnx_gqa_sdpa() which faithfully ports the algorithm from
onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h:
  - Scale baked into GEMM alpha (no separate scaling pass)
  - Scores buffer padded to max_seq_len columns
  - Causal mask: zero out future positions, softmax on valid window only
  - Output always in [B, S, Hq, D] format

Extends validation to verify ONNX GQA output matches custom_sdpa_out
reference. Adds OnnxGQABenchFixture for benchmarking both layouts.

Differential Revision: [D96044317](https://our.internmc.facebook.com/intern/diff/D96044317/)

[ghstack-poisoned]

Author: kimishpatel

.github/workflows/android-release-artifacts.yml

@@ -53,6 +53,9 @@ jobs:
         run: |
           VERSION="${{ inputs.version }}"
           FLAVOR="${{ inputs.flavor }}"
+          if [ -z "$FLAVOR" ]; then
+            FLAVOR="xnnpack"
+          fi
           if [ -z "$VERSION" ]; then
             echo "No version name specified. Will create a snapshot AAR"
             echo "should-skip=false" >> $GITHUB_OUTPUT
@@ -115,9 +118,10 @@ jobs:
         fi
 
         FLAVOR="${{ inputs.flavor }}"
-        if [ ! -z "$FLAVOR" ]; then
-          GRADLE_ARGS+=" -Dflavor=${FLAVOR}"
+        if [ -z "$FLAVOR" ]; then
+          FLAVOR="xnnpack"
         fi
+        GRADLE_ARGS+=" -Dflavor=${FLAVOR}"
 
         if [[ "$FLAVOR" == "vulkan" || "$FLAVOR" == "all" || -z "$FLAVOR" ]]; then
           curl -O https://sdk.lunarg.com/sdk/download/1.4.321.1/linux/vulkansdk-linux-x86_64-1.4.321.1.tar.xz

.gitignore

@@ -18,6 +18,7 @@ cmake-out*
 cmake-out-android/
 build-android/
 build-x86/
+build-hexagon/
 dist/
 arm-scratch/
 executorch.egg-info

CMakeLists.txt

@@ -50,6 +50,17 @@ project(executorch)
 
 set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR})
 
+# Hexagon toolchain with release build complains about code in third party
+# libraries.
+if("${CMAKE_SYSTEM_PROCESSOR}" MATCHES "Hexagon" AND "${CMAKE_BUILD_TYPE}"
+                                                     STREQUAL "Release"
+)
+  add_compile_options(
+    -Wno-error=format -Wno-error=implicit-int-conversion
+    -Wno-error=unused-variable -Wno-error=unused-function
+  )
+endif()
+
 # --- ExecuTorch Version ---
 # Parse version from version.txt (single source of truth)
 file(READ "${EXECUTORCH_ROOT}/version.txt" ET_VERSION_STRING)

backends/arm/_passes/__init__.py

@@ -85,7 +85,6 @@
 from .decompose_sinh_pass import DecomposeSinhPass  # noqa
 from .decompose_slice_scatter_pass import DecomposeSliceScatterPass  # noqa
 from .decompose_softmax_pass import DecomposeSoftmaxPass  # noqa
-from .decompose_softmax_unstable_pass import DecomposeSoftmaxUnstablePass  # noqa
 from .decompose_sqrt_pass import DecomposeSqrtPass  # noqa
 from .decompose_strided_slice_copy_pass import DecomposeStridedSliceCopyPass  # noqa
 from .decompose_sum_pass import DecomposeSumPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -86,7 +86,6 @@
     DecomposeSinhPass,
     DecomposeSliceScatterPass,
     DecomposeSoftmaxPass,
-    DecomposeSoftmaxUnstablePass,
     DecomposeSqrtPass,
     DecomposeStridedSliceCopyPass,
     DecomposeSumPass,
@@ -196,12 +195,8 @@ def configure_skip_passes(
 
         match config.softmax:
             case SoftmaxDecompositionConfig.MASKED:
-                skip_set.add(DecomposeSoftmaxUnstablePass)
-            case SoftmaxDecompositionConfig.UNSTABLE:
-                skip_set.add(DecomposeSoftmaxPass)
-                skip_set.add(DecomposeMaskedFillPass)
+                pass
             case SoftmaxDecompositionConfig.STABLE:
-                skip_set.add(DecomposeSoftmaxUnstablePass)
                 skip_set.add(DecomposeMaskedFillPass)
 
         if config.fuse_duplicate_users is FuseDuplicateUsersConfig.DISABLED:
@@ -461,9 +456,7 @@ def _tosa_pipeline(
                 ConvertMmToBmmPass(),
                 DecomposeGluPass(),
                 DecomposeDivPass(),
-                # _safe_softmax results in a ReduceMax
-                # which is not currently supported by TOSA in U55
-                DecomposeSoftmaxPass(skip_safe_softmax=self.tosa_spec.is_U55_subset),
+                DecomposeSoftmaxPass(),
                 ConvertMinMaxPass(),
                 DecomposeAnyPass(),
                 DecomposeAdaptiveAvgPool2dPass(),
@@ -593,9 +586,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 DecomposeSqrtPass(tfa_pass=True),
                 DecomposeAdaptiveAvgPool2dPass(tfa_pass=True),
                 DecomposeAvgPool2dPass(tfa_pass=True),
-                DecomposeSoftmaxUnstablePass(tfa_pass=True),
                 DecomposeSoftmaxPass(
-                    skip_safe_softmax=self.tosa_spec.is_U55_subset,
                     tfa_pass=True,
                 ),
                 ConvertMinMaxPass(tfa_pass=True),

backends/arm/_passes/decompose_softmax_pass.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 from typing import Set, Type
 
 import torch
@@ -25,6 +26,8 @@
 )
 log_softmax = (torch.ops.aten.log_softmax.int, exir_ops.edge.aten._log_softmax.default)
 
+logger = logging.getLogger(__name__)
+
 
 def _get_logsoftmax_ops(op) -> tuple:
     """Returns the (log_op, sub_op, amax_op, expo_op, sum_op, reciprocal_op),
@@ -78,6 +81,7 @@ class DecomposeSoftmaxPass(ArmPass):
     def __init__(self, skip_safe_softmax: bool = False, **kwargs):
         super().__init__(**kwargs)
         self._skip_safe_softmax = skip_safe_softmax
+        self._warned_safe_softmax = False
 
     def call_operator(self, op, args, kwargs, meta):
         if op not in torch_softmax + edge_softmax or not self.allowed_to_transform(
@@ -88,6 +92,18 @@ def call_operator(self, op, args, kwargs, meta):
         if self._skip_safe_softmax and op == torch.ops.aten._safe_softmax.default:
             return super().call_operator(op, args, kwargs, meta)
 
+        if (
+            self.is_tfa_pass
+            and op == torch.ops.aten._safe_softmax.default
+            and not self._warned_safe_softmax
+        ):
+            logger.warning(
+                "aten._safe_softmax is being decomposed as regular softmax in "
+                "the annotation pipeline; this is only semantics-preserving "
+                "when no row is fully masked at runtime."
+            )
+            self._warned_safe_softmax = True
+
         log_op, sub_op, max_op, exp_op, sum_op, reciprocal_op, mul_op = (
             _get_logsoftmax_ops(op)
         )

backends/arm/_passes/decompose_softmax_unstable_pass.py

@@ -39,6 +39,7 @@ class TableOps:
         exir_ops.edge.aten.floor.default: torch.floor,
         exir_ops.edge.aten.log.default: torch.log,
         exir_ops.edge.aten.log1p.default: torch.log1p,
+        exir_ops.edge.aten.log10.default: torch.log10,
         exir_ops.edge.aten.reciprocal.default: torch.reciprocal,
         exir_ops.edge.aten.rsqrt.default: torch.rsqrt,
         exir_ops.edge.aten.sigmoid.default: torch.sigmoid,

backends/arm/_passes/insert_table_ops.py

@@ -16,7 +16,10 @@
 from dataclasses import dataclass, field
 from enum import Enum
 
-from executorch.backends.arm.common.pipeline_config import ArmPassPipelineConfig
+from executorch.backends.arm.common.pipeline_config import (
+    ArmPassPipelineConfig,
+    SoftmaxDecompositionConfig,
+)
 from executorch.backends.arm.tosa import TosaSpecification
 from executorch.exir._warnings import deprecated
 
@@ -250,7 +253,10 @@ def set_pass_pipeline_config(self, config: ArmPassPipelineConfig) -> None:
     def _create_default_pipeline_config(self) -> ArmPassPipelineConfig:
         config = ArmPassPipelineConfig()
         if self.tosa_spec.is_U55_subset:
-            config.disable_masked_softmax()
+            # Keep U55 on STABLE instead of the generic MASKED default:
+            # MASKED also enables masked_fill decomposition, which lowers to
+            # where/full_like and is not a good default fit for U55.
+            config.softmax = SoftmaxDecompositionConfig.STABLE
         return config
 
     def _get_intermediate_path(self) -> str | None:

backends/arm/common/arm_compile_spec.py

@@ -8,12 +8,9 @@
 from enum import auto, Enum
 from typing import Any
 
-from executorch.exir._warnings import deprecated
-
 
 class SoftmaxDecompositionConfig(Enum):
     MASKED = auto()  # Stable softmax + masked fill decomposition
-    UNSTABLE = auto()  # Unstable softmax, no masked fill decomposition
     STABLE = auto()  # Stable softmax, no masked fill decomposition
 
 
@@ -27,18 +24,6 @@ class ArmPassPipelineConfig:
     softmax: SoftmaxDecompositionConfig = SoftmaxDecompositionConfig.MASKED
     fuse_duplicate_users: FuseDuplicateUsersConfig = FuseDuplicateUsersConfig.ENABLED
 
-    @deprecated(
-        "The stable softmax decomposition is now supported by all arm targets and will be made default in a future release. Overwrite the default config using `compile_spec.set_pass_pipeline_config(ArmPassPipelineConfig())` to use the stable algorithm and avoid this error."
-    )
-    def disable_masked_softmax(self) -> None:
-        """
-            .. warning::
-
-        The stable softmax decomposition is now supported by all arm targets and will be made default in a future release. Overwrite the default config using `compile_spec.set_pass_pipeline_config(ArmPassPipelineConfig())` to use the stable algorithm and avoid this error."
-        """
-
-        self.softmax = SoftmaxDecompositionConfig.STABLE
-
     def disable_fuse_duplicate_users(self) -> None:
         self.fuse_duplicate_users = FuseDuplicateUsersConfig.DISABLED