From 2861e1031185133e8391652d046cdedccec9fbd2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= <mans.nilsson@arm.com>
Date: Thu, 28 May 2026 12:11:02 +0200
Subject: [PATCH 1/5] XNNPACK: Lift constant mul scalars for partitioning
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

XNNPACK supports the tensor overload for multiply, but plain
aten.mul.Scalar is not selected by the XNNPACK partitioner.

Adds a narrow scalar-lifting pass that rewrites aten.mul.Scalar to
aten.mul.Tensor by registering the scalar as a small constant buffer.
This avoids introducing an aten.full op while allowing the existing
multiply lowering path to partition the op.

Keep SDPA scale multipliers as aten.mul.Scalar so ConvertToSDPAPass can
still recover the attention scale before replacing the pattern. Add test
coverage for that guard.

Allow the XNNPACK tester to pass transform passes through to
to_edge_transform_and_lower. This keeps op tests on the same path as
existing XNNPACK model tests that already use explicit transform passes.

For DeIT Tiny, this removes 24 portable aten.mul.Scalar nodes and
reduces delegate count from 62 to 50. In current local timing checks
the latency impact is modest: about 1% faster on both Android SME2 and
the aarch64 XNNPACK/KleidiAI NEON-class host runner.
These are modest uplifts but may introduce more opportunities for
improvements.

Signed-off-by: Måns Nilsson <mans.nilsson@arm.com>
Change-Id: I83b6ad53925edb72afdf0077b5dbb99b5d9c4648
---
 .../lift_constant_scalar_operands_pass.py     | 151 ++++++++++++++++++
 backends/xnnpack/test/ops/test_multiply.py    |  31 ++++
 ...test_lift_constant_scalar_operands_pass.py |  74 +++++++++
 backends/xnnpack/test/tester/tester.py        |  46 +++++-
 backends/xnnpack/utils/configs.py             |   5 +-
 5 files changed, 304 insertions(+), 3 deletions(-)
 create mode 100644 backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
 create mode 100644 backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py

diff --git a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
new file mode 100644
index 00000000000..8c6be2b62b5
--- /dev/null
+++ b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
@@ -0,0 +1,151 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from numbers import Number
+from typing import Dict, Optional, Union
+
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
+from executorch.exir.pass_base import ExportPass, PassResult
+from torch._ops import OpOverload
+
+
+ScalarOp = Union[EdgeOpOverload, OpOverload]
+
+
+class LiftConstantScalarOperandsPass(ExportPass):
+    """
+    Lift scalar operands into tensor constants for selected binary ops.
+
+    XNNPACK already supports the tensor overloads for these binary operations.
+    This pass converts explicitly listed scalar overloads to their tensor
+    overloads by replacing constant scalar operands with small tensor constants.
+    The constants are registered as buffers so they do not become portable
+    ``full`` kernels. Keep the op map narrow until each new scalar overload is
+    covered by tests.
+    """
+
+    default_scalar_to_tensor_ops: Dict[ScalarOp, ScalarOp] = {
+        exir_ops.edge.aten.mul.Scalar: exir_ops.edge.aten.mul.Tensor,
+    }
+    sdpa_passthrough_ops = {
+        exir_ops.edge.aten.expand_copy.default,
+        exir_ops.edge.aten.view_copy.default,
+    }
+
+    def __init__(
+        self,
+        scalar_to_tensor_ops: Optional[Dict[ScalarOp, ScalarOp]] = None,
+    ) -> None:
+        super().__init__()
+        self.scalar_to_tensor_ops = (
+            scalar_to_tensor_ops
+            if scalar_to_tensor_ops is not None
+            else self.default_scalar_to_tensor_ops
+        )
+        self._modified = False
+
+    def _create_constant_node(
+        self,
+        graph_module: torch.fx.GraphModule,
+        node: torch.fx.Node,
+        value: Number,
+    ) -> torch.fx.Node:
+        input_node = node.args[0]
+        if not isinstance(input_node, torch.fx.Node):
+            raise RuntimeError("Expected scalar op input to be an FX node.")
+
+        input_value = input_node.meta["val"]
+        tensor = torch.tensor(value, dtype=input_value.dtype, device=input_value.device)
+        name = self._get_new_attr_name(graph_module)
+        graph_module.register_buffer(name, tensor)
+
+        fake_mode = node.meta["val"].fake_mode
+        with graph_module.graph.inserting_before(node):
+            constant_node = graph_module.graph.get_attr(name)
+            constant_node.meta["val"] = fake_mode.from_tensor(
+                tensor, static_shapes=True
+            )
+        return constant_node
+
+    def _get_new_attr_name(self, graph_module: torch.fx.GraphModule) -> str:
+        prefix = "_tensor_constant_"
+        index = 0
+        while hasattr(graph_module, f"{prefix}{index}"):
+            index += 1
+        return f"{prefix}{index}"
+
+    def _feeds_sdpa_qk_bmm(self, node: torch.fx.Node) -> bool:
+        """
+        Return true for the scale muls consumed by XNNPACK's SDPA pattern.
+
+        ConvertToSDPAPass recovers the user-specified attention scale from the
+        pre-QK^T ``aten.mul.Scalar`` nodes. Keep those scalar muls intact so
+        SDPA conversion can still find the scale before replacing the pattern.
+        """
+        users_to_visit = list(node.users)
+        visited = set()
+        while users_to_visit:
+            user = users_to_visit.pop()
+            if user in visited:
+                continue
+            visited.add(user)
+
+            if (
+                user.op == "call_function"
+                and user.target == exir_ops.edge.aten.bmm.default
+            ):
+                return True
+
+            if user.op == "call_function" and user.target in self.sdpa_passthrough_ops:
+                users_to_visit.extend(user.users)
+
+        return False
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        self._modified = False
+
+        for node in list(graph_module.graph.nodes):
+            if (
+                node.op != "call_function"
+                or node.target not in self.scalar_to_tensor_ops
+                or len(node.args) != 2
+                or not isinstance(node.args[0], torch.fx.Node)
+                or not isinstance(node.args[1], Number)
+            ):
+                continue
+
+            if (
+                node.target == exir_ops.edge.aten.mul.Scalar
+                and self._feeds_sdpa_qk_bmm(node)
+            ):
+                continue
+
+            input_value = node.args[0].meta.get("val")
+            output_value = node.meta.get("val")
+            if (
+                input_value is None
+                or output_value is None
+                or input_value.dtype != output_value.dtype
+            ):
+                continue
+
+            tensor_arg = self._create_constant_node(graph_module, node, node.args[1])
+            node.args = (node.args[0], tensor_arg)
+            node.target = self.scalar_to_tensor_ops[node.target]
+            self._modified = True
+
+        graph_module.graph.eliminate_dead_code()
+        graph_module.graph.lint()
+        graph_module.recompile()
+
+        modified = self._modified
+        self._modified = False
+        return PassResult(graph_module, modified)
diff --git a/backends/xnnpack/test/ops/test_multiply.py b/backends/xnnpack/test/ops/test_multiply.py
index 3315200005d..118136fcd08 100644
--- a/backends/xnnpack/test/ops/test_multiply.py
+++ b/backends/xnnpack/test/ops/test_multiply.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -8,6 +9,10 @@
 
 import torch
 from executorch.backends.xnnpack.test.tester import Tester
+from executorch.backends.xnnpack.utils.configs import (
+    get_transform_passes,
+    get_xnnpack_edge_compile_config,
+)
 
 
 class TestMul(unittest.TestCase):
@@ -29,6 +34,10 @@ def forward(self, x, y):
             z = torch.mul(x, y) * torch.functional.torch.mul(x, y)
             return z
 
+    class MulScalar(torch.nn.Module):
+        def forward(self, x):
+            return torch.ops.aten.mul.Scalar(x, 0.5)
+
     class MulRelu(torch.nn.Module):
         def forward(self, x, y):
             z = x * y
@@ -58,6 +67,28 @@ def test_fp32_mul(self):
         inputs = (torch.randn((1, 3)), torch.randn((4, 3)))
         self._test_mul(inputs)
 
+    def test_fp32_mul_scalar(self):
+        (
+            Tester(self.MulScalar(), (torch.randn(2, 3),))
+            .export()
+            .to_edge_transform_and_lower(
+                transform_passes=get_transform_passes(),
+                edge_compile_config=get_xnnpack_edge_compile_config(
+                    skip_dim_order=True
+                ),
+            )
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .check_not(
+                [
+                    "executorch_exir_dialects_edge__ops_aten_mul_Tensor",
+                    "executorch_exir_dialects_edge__ops_aten_mul_Scalar",
+                ]
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
     def test_qs8_mul(self):
         inputs = (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 1))
         (
diff --git a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
new file mode 100644
index 00000000000..5ad44f78af0
--- /dev/null
+++ b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
@@ -0,0 +1,74 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from copy import deepcopy
+
+import torch
+from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
+    LiftConstantScalarOperandsPass,
+)
+from executorch.backends.xnnpack.partition.graphs import sdpa
+from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config
+from executorch.exir import to_edge
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+class TestLiftConstantScalarOperandsPass(unittest.TestCase):
+    def setUp(self):
+        torch._dynamo.reset()
+
+    class MulScalar(torch.nn.Module):
+        def forward(self, x):
+            return torch.ops.aten.mul.Scalar(x, 0.5)
+
+    class AddScalar(torch.nn.Module):
+        def forward(self, x):
+            return torch.ops.aten.add.Scalar(x, 0.5)
+
+    def _to_edge_graph(self, module):
+        edge = to_edge(
+            torch.export.export(module, (torch.randn(2, 3),), strict=True),
+            compile_config=get_xnnpack_edge_compile_config(skip_dim_order=True),
+        )
+        return edge.transform([LiftConstantScalarOperandsPass()]).exported_program()
+
+    def test_lifts_mul_scalar_operand(self):
+        graph = self._to_edge_graph(self.MulScalar()).graph_module.graph
+
+        self.assertFalse(
+            any(node.target == exir_ops.edge.aten.mul.Scalar for node in graph.nodes)
+        )
+        self.assertTrue(
+            any(node.target == exir_ops.edge.aten.mul.Tensor for node in graph.nodes)
+        )
+        self.assertTrue(any(node.op == "get_attr" for node in graph.nodes))
+
+    def test_keeps_unmapped_scalar_op(self):
+        graph = self._to_edge_graph(self.AddScalar()).graph_module.graph
+
+        self.assertTrue(
+            any(node.target == exir_ops.edge.aten.add.Scalar for node in graph.nodes)
+        )
+
+    def test_keeps_sdpa_scale_mul_scalar(self):
+        graph_module = deepcopy(sdpa.get_graphs()[0])
+
+        LiftConstantScalarOperandsPass()(graph_module)
+
+        scale_mul_count = 0
+        lifted_mul_count = 0
+        for node in graph_module.graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target == exir_ops.edge.aten.mul.Scalar:
+                scale_mul_count += 1
+            if node.target == exir_ops.edge.aten.mul.Tensor:
+                lifted_mul_count += 1
+
+        self.assertEqual(scale_mul_count, 2)
+        self.assertEqual(lifted_mul_count, 0)
diff --git a/backends/xnnpack/test/tester/tester.py b/backends/xnnpack/test/tester/tester.py
index fc12da231c0..481864e265f 100644
--- a/backends/xnnpack/test/tester/tester.py
+++ b/backends/xnnpack/test/tester/tester.py
@@ -1,6 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
-# Copyright 2024-2025 Arm Limited and/or its affiliates.
+# Copyright 2024-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -24,9 +24,10 @@
     QuantizationConfig,
 )
 from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config
-from executorch.exir import EdgeCompileConfig
+from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
 from executorch.exir.backend.partitioner import Partitioner
 from torch._export.pass_base import PassType
+from torch.export import ExportedProgram
 from torchao.quantization.pt2e.quantizer import Quantizer
 
 
@@ -77,6 +78,7 @@ def __init__(
         self,
         partitioners: Optional[List[Partitioner]] = None,
         edge_compile_config: Optional[EdgeCompileConfig] = None,
+        transform_passes: Optional[List[PassType]] = None,
     ):
         super().__init__(
             default_partitioner_cls=XnnpackPartitioner,
@@ -84,6 +86,21 @@ def __init__(
             edge_compile_config=edge_compile_config
             or get_xnnpack_edge_compile_config(),
         )
+        self.transform_passes = transform_passes
+
+    def run(
+        self,
+        artifact: ExportedProgram,
+        inputs=None,
+        generate_etrecord: bool = False,
+    ) -> None:
+        self.edge_dialect_program = to_edge_transform_and_lower(
+            artifact,
+            transform_passes=self.transform_passes,
+            compile_config=self.edge_compile_conf,
+            partitioner=self.partitioners,
+            generate_etrecord=generate_etrecord,
+        )
 
 
 class Partition(BaseStages.Partition):
@@ -132,3 +149,28 @@ def __init__(
             dynamic_shapes=dynamic_shapes,
             **kwargs,
         )
+
+    def to_edge_transform_and_lower(
+        self,
+        to_edge_and_lower_stage: Optional[BaseStages.ToEdgeTransformAndLower] = None,
+        generate_etrecord: bool = False,
+        *,
+        partitioners: Optional[List[Partitioner]] = None,
+        edge_compile_config: Optional[EdgeCompileConfig] = None,
+        transform_passes: Optional[List[PassType]] = None,
+    ):
+        if to_edge_and_lower_stage is None:
+            to_edge_and_lower_stage = ToEdgeTransformAndLower(
+                partitioners=partitioners,
+                edge_compile_config=edge_compile_config,
+                transform_passes=transform_passes,
+            )
+        else:
+            if partitioners is not None:
+                to_edge_and_lower_stage.partitioners = partitioners
+            if edge_compile_config is not None:
+                to_edge_and_lower_stage.edge_compile_conf = edge_compile_config
+        return super().to_edge_transform_and_lower(
+            to_edge_and_lower_stage,
+            generate_etrecord=generate_etrecord,
+        )
diff --git a/backends/xnnpack/utils/configs.py b/backends/xnnpack/utils/configs.py
index 3016e94146b..ec47b81e835 100644
--- a/backends/xnnpack/utils/configs.py
+++ b/backends/xnnpack/utils/configs.py
@@ -9,6 +9,9 @@
 
 import executorch.exir as exir
 
+from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
+    LiftConstantScalarOperandsPass,
+)
 from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import (
     RemoveNoopExpandCopyPass,
 )
@@ -25,7 +28,7 @@ def get_xnnpack_edge_compile_config(
 
 
 def get_transform_passes(additional_passes=None) -> List[PassType]:
-    passes = [RemoveNoopExpandCopyPass()]
+    passes = [RemoveNoopExpandCopyPass(), LiftConstantScalarOperandsPass()]
     if additional_passes:
         passes.extend(additional_passes)
     return passes

From 3952854c6eacd2a4c517512f28d83f44a4efb9d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= <mans.nilsson@arm.com>
Date: Fri, 26 Jun 2026 10:12:14 +0200
Subject: [PATCH 2/5] Address review comments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Måns Nilsson <mans.nilsson@arm.com>
Change-Id: I4ab9a85590762c24f474b935a26f61d90b9a0565
---
 .../lift_constant_scalar_operands_pass.py     |  9 +++----
 ...test_lift_constant_scalar_operands_pass.py | 19 +++++++++++---
 backends/xnnpack/test/tester/tester.py        | 26 +++++++++++++------
 3 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
index 8c6be2b62b5..e256a31750e 100644
--- a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
+++ b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
@@ -50,7 +50,6 @@ def __init__(
             if scalar_to_tensor_ops is not None
             else self.default_scalar_to_tensor_ops
         )
-        self._modified = False
 
     def _create_constant_node(
         self,
@@ -65,6 +64,8 @@ def _create_constant_node(
         input_value = input_node.meta["val"]
         tensor = torch.tensor(value, dtype=input_value.dtype, device=input_value.device)
         name = self._get_new_attr_name(graph_module)
+        # Keep constants as module attributes so the portable path can emit them
+        # without introducing aten.full, while XNNPACK can still read them as params.
         graph_module.register_buffer(name, tensor)
 
         fake_mode = node.meta["val"].fake_mode
@@ -110,7 +111,7 @@ def _feeds_sdpa_qk_bmm(self, node: torch.fx.Node) -> bool:
         return False
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
-        self._modified = False
+        modified = False
 
         for node in list(graph_module.graph.nodes):
             if (
@@ -140,12 +141,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             tensor_arg = self._create_constant_node(graph_module, node, node.args[1])
             node.args = (node.args[0], tensor_arg)
             node.target = self.scalar_to_tensor_ops[node.target]
-            self._modified = True
+            modified = True
 
         graph_module.graph.eliminate_dead_code()
         graph_module.graph.lint()
         graph_module.recompile()
 
-        modified = self._modified
-        self._modified = False
         return PassResult(graph_module, modified)
diff --git a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
index 5ad44f78af0..5c61731a786 100644
--- a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
+++ b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
@@ -16,6 +16,7 @@
 from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config
 from executorch.exir import to_edge
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_manager import ExportedProgramPassManager
 
 
 class TestLiftConstantScalarOperandsPass(unittest.TestCase):
@@ -30,12 +31,17 @@ class AddScalar(torch.nn.Module):
         def forward(self, x):
             return torch.ops.aten.add.Scalar(x, 0.5)
 
-    def _to_edge_graph(self, module):
-        edge = to_edge(
+    def _to_edge_program_manager(self, module):
+        return to_edge(
             torch.export.export(module, (torch.randn(2, 3),), strict=True),
             compile_config=get_xnnpack_edge_compile_config(skip_dim_order=True),
         )
-        return edge.transform([LiftConstantScalarOperandsPass()]).exported_program()
+
+    def _to_edge_graph(self, module):
+        edge = self._to_edge_program_manager(module)
+        return ExportedProgramPassManager([LiftConstantScalarOperandsPass()])(
+            edge.exported_program()
+        ).exported_program
 
     def test_lifts_mul_scalar_operand(self):
         graph = self._to_edge_graph(self.MulScalar()).graph_module.graph
@@ -48,6 +54,13 @@ def test_lifts_mul_scalar_operand(self):
         )
         self.assertTrue(any(node.op == "get_attr" for node in graph.nodes))
 
+    def test_lifted_mul_scalar_can_emit_without_delegation(self):
+        edge = self._to_edge_program_manager(self.MulScalar()).transform(
+            (LiftConstantScalarOperandsPass(),)
+        )
+
+        self.assertIsNotNone(edge.to_executorch())
+
     def test_keeps_unmapped_scalar_op(self):
         graph = self._to_edge_graph(self.AddScalar()).graph_module.graph
 
diff --git a/backends/xnnpack/test/tester/tester.py b/backends/xnnpack/test/tester/tester.py
index 481864e265f..396e149565f 100644
--- a/backends/xnnpack/test/tester/tester.py
+++ b/backends/xnnpack/test/tester/tester.py
@@ -26,6 +26,7 @@
 from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config
 from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
 from executorch.exir.backend.partitioner import Partitioner
+from executorch.exir.pass_manager import PassType as ExirPassType
 from torch._export.pass_base import PassType
 from torch.export import ExportedProgram
 from torchao.quantization.pt2e.quantizer import Quantizer
@@ -78,7 +79,7 @@ def __init__(
         self,
         partitioners: Optional[List[Partitioner]] = None,
         edge_compile_config: Optional[EdgeCompileConfig] = None,
-        transform_passes: Optional[List[PassType]] = None,
+        transform_passes: Optional[List[ExirPassType]] = None,
     ):
         super().__init__(
             default_partitioner_cls=XnnpackPartitioner,
@@ -152,25 +153,34 @@ def __init__(
 
     def to_edge_transform_and_lower(
         self,
-        to_edge_and_lower_stage: Optional[BaseStages.ToEdgeTransformAndLower] = None,
+        to_edge_and_transform_stage: Optional[
+            BaseStages.ToEdgeTransformAndLower
+        ] = None,
         generate_etrecord: bool = False,
         *,
         partitioners: Optional[List[Partitioner]] = None,
         edge_compile_config: Optional[EdgeCompileConfig] = None,
-        transform_passes: Optional[List[PassType]] = None,
+        transform_passes: Optional[List[ExirPassType]] = None,
     ):
-        if to_edge_and_lower_stage is None:
-            to_edge_and_lower_stage = ToEdgeTransformAndLower(
+        if to_edge_and_transform_stage is None:
+            to_edge_and_transform_stage = ToEdgeTransformAndLower(
                 partitioners=partitioners,
                 edge_compile_config=edge_compile_config,
                 transform_passes=transform_passes,
             )
         else:
             if partitioners is not None:
-                to_edge_and_lower_stage.partitioners = partitioners
+                to_edge_and_transform_stage.partitioners = partitioners
             if edge_compile_config is not None:
-                to_edge_and_lower_stage.edge_compile_conf = edge_compile_config
+                to_edge_and_transform_stage.edge_compile_conf = edge_compile_config
+            if transform_passes is not None:
+                if not isinstance(to_edge_and_transform_stage, ToEdgeTransformAndLower):
+                    raise ValueError(
+                        "transform_passes requires the XNNPACK "
+                        "ToEdgeTransformAndLower stage."
+                    )
+                to_edge_and_transform_stage.transform_passes = transform_passes
         return super().to_edge_transform_and_lower(
-            to_edge_and_lower_stage,
+            to_edge_and_transform_stage,
             generate_etrecord=generate_etrecord,
         )

From 9eac80e90c8d6b67cc2970b8e99abeb1b8cae6e0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= <mans.nilsson@arm.com>
Date: Tue, 4 Aug 2026 17:22:46 +0200
Subject: [PATCH 3/5] Adress Claude review comments

---
 .../lift_constant_scalar_operands_pass.py     |  2 +-
 ...test_lift_constant_scalar_operands_pass.py | 33 ++++++++++---------
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
index e256a31750e..7d7eba44047 100644
--- a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
+++ b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
@@ -16,7 +16,6 @@
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch._ops import OpOverload
 
-
 ScalarOp = Union[EdgeOpOverload, OpOverload]
 
 
@@ -64,6 +63,7 @@ def _create_constant_node(
         input_value = input_node.meta["val"]
         tensor = torch.tensor(value, dtype=input_value.dtype, device=input_value.device)
         name = self._get_new_attr_name(graph_module)
+        # ExportPass has no ExportedProgram access to create a constant placeholder.
         # Keep constants as module attributes so the portable path can emit them
         # without introducing aten.full, while XNNPACK can still read them as params.
         graph_module.register_buffer(name, tensor)
diff --git a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
index 5c61731a786..749c89f0038 100644
--- a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
+++ b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
@@ -69,19 +69,20 @@ def test_keeps_unmapped_scalar_op(self):
         )
 
     def test_keeps_sdpa_scale_mul_scalar(self):
-        graph_module = deepcopy(sdpa.get_graphs()[0])
-
-        LiftConstantScalarOperandsPass()(graph_module)
-
-        scale_mul_count = 0
-        lifted_mul_count = 0
-        for node in graph_module.graph.nodes:
-            if node.op != "call_function":
-                continue
-            if node.target == exir_ops.edge.aten.mul.Scalar:
-                scale_mul_count += 1
-            if node.target == exir_ops.edge.aten.mul.Tensor:
-                lifted_mul_count += 1
-
-        self.assertEqual(scale_mul_count, 2)
-        self.assertEqual(lifted_mul_count, 0)
+        for graph in sdpa.get_graphs():
+            graph_module = deepcopy(graph)
+
+            LiftConstantScalarOperandsPass()(graph_module)
+
+            scale_mul_count = 0
+            lifted_mul_count = 0
+            for node in graph_module.graph.nodes:
+                if node.op != "call_function":
+                    continue
+                if node.target == exir_ops.edge.aten.mul.Scalar:
+                    scale_mul_count += 1
+                if node.target == exir_ops.edge.aten.mul.Tensor:
+                    lifted_mul_count += 1
+
+            self.assertEqual(scale_mul_count, 2)
+            self.assertEqual(lifted_mul_count, 0)

From 257ef3a853b2ccc4684c0c20910846f7601999f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= <mans.nilsson@arm.com>
Date: Thu, 27 Aug 2026 12:10:10 +0200
Subject: [PATCH 4/5] XNNPACK: Handle lifted scalar scales in SDPA pass
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Måns Nilsson <mans.nilsson@arm.com>
---
 backends/xnnpack/_passes/convert_to_sdpa.py   | 82 ++++++++++++-------
 .../lift_constant_scalar_operands_pass.py     | 37 ---------
 ...test_lift_constant_scalar_operands_pass.py | 40 ++++++++-
 3 files changed, 89 insertions(+), 70 deletions(-)

diff --git a/backends/xnnpack/_passes/convert_to_sdpa.py b/backends/xnnpack/_passes/convert_to_sdpa.py
index c7982db750f..a4c421bfede 100644
--- a/backends/xnnpack/_passes/convert_to_sdpa.py
+++ b/backends/xnnpack/_passes/convert_to_sdpa.py
@@ -1,19 +1,23 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # 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 copy import deepcopy
 from typing import Optional
 
 import torch
 from executorch.backends.transforms import get_shape
-
+from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
+    LiftConstantScalarOperandsPass,
+)
 from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
 from executorch.backends.xnnpack.partition.graphs import sdpa
+from executorch.backends.xnnpack.utils.utils import get_param_tensor
 from executorch.exir.dialects._ops import ops as exir_ops
-
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.fx.passes.utils.matcher_utils import InternalMatch, SubgraphMatcher
 
@@ -24,31 +28,43 @@
 class ConvertToSDPAPass(XNNPACKPass):
     def get_scale(self, match: InternalMatch) -> Optional[float]:
         """
-        Returns the scale of the SDPA op.
+        Return the SDPA scale recovered from the matched pre-QK^T multiplications.
 
-        Scale: Optional[float] doesn't change the graph pattern.
-        The default value can be calulated however we need to extract
-        it for lowering when it is the user supplied value anyway.
+        The multiplier may be a scalar literal or a constant tensor introduced by
+        scalar lifting. The decomposition applies the square root of the attention
+        scale before QK^T, so the extracted multiplier is squared to recover the
+        original value.
         """
         for node in match.nodes_map.values():
-            if (
-                node.op == "call_function"
-                and node.target == exir_ops.edge.aten.mul.Scalar
-            ):
-                scale = node.args[1]
-
-                dtype = torch.float
-                mul_val = node.meta.get("val", None)
-                if mul_val is not None:
-                    dtype = mul_val.dtype
-
-                if isinstance(scale, float):
-                    # Convert scale value to fp16 (reducing precision)
-                    scale = torch.tensor(scale, dtype=dtype).item()
-
-                    # since scale we extracted this before the QK^T.
-                    return scale**2
-                break
+            if node.op != "call_function" or node.target not in {
+                exir_ops.edge.aten.mul.Scalar,
+                exir_ops.edge.aten.mul.Tensor,
+            }:
+                continue
+
+            scale = node.args[1]
+
+            # Extract the scale from the constant tensor introduced by scalar
+            # lifting.
+            if node.target == exir_ops.edge.aten.mul.Tensor:
+                if not isinstance(scale, torch.fx.Node):
+                    continue
+                scale_tensor = get_param_tensor(self.exported_program, scale)
+                if scale_tensor is None or scale_tensor.numel() != 1:
+                    continue
+                scale = scale_tensor.item()
+
+            dtype = torch.float
+            mul_val = node.meta.get("val", None)
+            if mul_val is not None:
+                dtype = mul_val.dtype
+
+            if isinstance(scale, float):
+                # Convert scale value to fp16 (reducing precision)
+                scale = torch.tensor(scale, dtype=dtype).item()
+
+                # since scale we extracted this before the QK^T.
+                return scale**2
         return None
 
     def assert_2d_mask(self, match: InternalMatch) -> None:
@@ -99,11 +115,19 @@ def call(self, graph_module: torch.fx.GraphModule):
         logger.debug("ConvertToSDPA Begin: ")
         logger.debug(graph_module.print_readable(print_output=False))
 
-        for pattern in sdpa.get_graphs():
-            sm = SubgraphMatcher(pattern.graph, ignore_literals=True)
-            matches = list(sm.match(graph_module.graph))
-            for partition_to_replace in matches:
-                self.create_sdpa(graph_module, partition_to_replace)
+        for scalar_pattern in sdpa.get_graphs():
+            # Deep-copy the cached scalar pattern so lifting it does not modify the
+            # pattern used by non-lifted flows.
+            tensor_pattern = deepcopy(scalar_pattern)
+            tensor_pattern = LiftConstantScalarOperandsPass()(
+                tensor_pattern
+            ).graph_module
+
+            for pattern in (scalar_pattern, tensor_pattern):
+                sm = SubgraphMatcher(pattern.graph, ignore_literals=True)
+                matches = list(sm.match(graph_module.graph))
+                for partition_to_replace in matches:
+                    self.create_sdpa(graph_module, partition_to_replace)
 
         graph_module.recompile()
         graph_module = super().call(graph_module).graph_module
diff --git a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
index 7d7eba44047..b0a21b273d5 100644
--- a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
+++ b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
@@ -34,10 +34,6 @@ class LiftConstantScalarOperandsPass(ExportPass):
     default_scalar_to_tensor_ops: Dict[ScalarOp, ScalarOp] = {
         exir_ops.edge.aten.mul.Scalar: exir_ops.edge.aten.mul.Tensor,
     }
-    sdpa_passthrough_ops = {
-        exir_ops.edge.aten.expand_copy.default,
-        exir_ops.edge.aten.view_copy.default,
-    }
 
     def __init__(
         self,
@@ -83,33 +79,6 @@ def _get_new_attr_name(self, graph_module: torch.fx.GraphModule) -> str:
             index += 1
         return f"{prefix}{index}"
 
-    def _feeds_sdpa_qk_bmm(self, node: torch.fx.Node) -> bool:
-        """
-        Return true for the scale muls consumed by XNNPACK's SDPA pattern.
-
-        ConvertToSDPAPass recovers the user-specified attention scale from the
-        pre-QK^T ``aten.mul.Scalar`` nodes. Keep those scalar muls intact so
-        SDPA conversion can still find the scale before replacing the pattern.
-        """
-        users_to_visit = list(node.users)
-        visited = set()
-        while users_to_visit:
-            user = users_to_visit.pop()
-            if user in visited:
-                continue
-            visited.add(user)
-
-            if (
-                user.op == "call_function"
-                and user.target == exir_ops.edge.aten.bmm.default
-            ):
-                return True
-
-            if user.op == "call_function" and user.target in self.sdpa_passthrough_ops:
-                users_to_visit.extend(user.users)
-
-        return False
-
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         modified = False
 
@@ -123,12 +92,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             ):
                 continue
 
-            if (
-                node.target == exir_ops.edge.aten.mul.Scalar
-                and self._feeds_sdpa_qk_bmm(node)
-            ):
-                continue
-
             input_value = node.args[0].meta.get("val")
             output_value = node.meta.get("val")
             if (
diff --git a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
index 749c89f0038..536e558a53c 100644
--- a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
+++ b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
@@ -9,6 +9,7 @@
 from copy import deepcopy
 
 import torch
+from executorch.backends.xnnpack._passes.convert_to_sdpa import ConvertToSDPAPass
 from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
     LiftConstantScalarOperandsPass,
 )
@@ -31,6 +32,10 @@ class AddScalar(torch.nn.Module):
         def forward(self, x):
             return torch.ops.aten.add.Scalar(x, 0.5)
 
+    class SDPA(torch.nn.Module):
+        def forward(self, q, k, v, mask):
+            return torch.nn.functional.scaled_dot_product_attention(q, k, v, mask)
+
     def _to_edge_program_manager(self, module):
         return to_edge(
             torch.export.export(module, (torch.randn(2, 3),), strict=True),
@@ -68,11 +73,11 @@ def test_keeps_unmapped_scalar_op(self):
             any(node.target == exir_ops.edge.aten.add.Scalar for node in graph.nodes)
         )
 
-    def test_keeps_sdpa_scale_mul_scalar(self):
+    def test_lifts_sdpa_scale_mul_scalar(self):
         for graph in sdpa.get_graphs():
             graph_module = deepcopy(graph)
 
-            LiftConstantScalarOperandsPass()(graph_module)
+            graph_module = LiftConstantScalarOperandsPass()(graph_module).graph_module
 
             scale_mul_count = 0
             lifted_mul_count = 0
@@ -84,5 +89,32 @@ def test_keeps_sdpa_scale_mul_scalar(self):
                 if node.target == exir_ops.edge.aten.mul.Tensor:
                     lifted_mul_count += 1
 
-            self.assertEqual(scale_mul_count, 2)
-            self.assertEqual(lifted_mul_count, 0)
+            self.assertEqual(scale_mul_count, 0)
+            self.assertEqual(lifted_mul_count, 2)
+
+    def test_converts_sdpa_after_lifting_scale_mul_scalar(self):
+        q = torch.randn(2, 4, 8, 16)
+        k = torch.randn(2, 4, 8, 16)
+        v = torch.randn(2, 4, 8, 16)
+        mask = torch.randn(8, 8)
+        edge = to_edge(
+            torch.export.export(self.SDPA(), (q, k, v, mask), strict=True),
+            compile_config=get_xnnpack_edge_compile_config(),
+        )
+        exported_program = ExportedProgramPassManager(
+            [LiftConstantScalarOperandsPass()]
+        )(edge.exported_program()).exported_program
+        exported_program = ExportedProgramPassManager(
+            [ConvertToSDPAPass(exported_program)]
+        )(exported_program).exported_program
+
+        graph = exported_program.graph_module.graph
+        self.assertTrue(
+            any(
+                node.target == exir_ops.edge.aten.scaled_dot_product_attention.default
+                for node in graph.nodes
+            )
+        )
+        self.assertFalse(
+            any(node.target == exir_ops.edge.aten.bmm.default for node in graph.nodes)
+        )

From 9e8ef8c8a3894c90c87c2b458580d5db8db29856 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= <mans.nilsson@arm.com>
Date: Tue, 1 Sep 2026 17:27:49 +0200
Subject: [PATCH 5/5] XNNPACK: Preserve scalar multiply through partitioning

Preserve eligible mul.Scalar operations during partitioning and lift their
scalar operands to constant tensors during XNNPACK preprocessing. Ensure SDPA
conversion still works after the default transform passes.
---
 backends/xnnpack/_passes/__init__.py          |   5 +
 backends/xnnpack/_passes/convert_to_sdpa.py   |  38 ++---
 .../lift_constant_scalar_operands_pass.py     |  54 ++++---
 backends/xnnpack/partition/config/__init__.py |   2 +
 .../partition/config/generic_node_configs.py  |  30 ++++
 backends/xnnpack/test/ops/test_multiply.py    |  11 +-
 ...test_lift_constant_scalar_operands_pass.py | 140 ++++++++++--------
 .../xnnpack/test/test_xnnpack_partitioner.py  |  34 +++++
 backends/xnnpack/test/tester/tester.py        |  56 +------
 backends/xnnpack/utils/configs.py             |   5 +-
 10 files changed, 197 insertions(+), 178 deletions(-)

diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py
index 22147fa4215..f4aea344d26 100644
--- a/backends/xnnpack/_passes/__init__.py
+++ b/backends/xnnpack/_passes/__init__.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -28,6 +29,9 @@
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
 from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
 from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass
+from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
+    LiftConstantScalarOperandsPass,
+)
 from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass
 from executorch.backends.xnnpack._passes.propagate_custom_meta_pass import (
     PropagateCustomMetaPass,
@@ -76,6 +80,7 @@ def __init__(
                 ConvertToLinearPass,
                 PropagateCustomMetaPass,
                 ConvertToSDPAPass,
+                LiftConstantScalarOperandsPass,
                 ConstPropPass,
                 FuseBatchNormPass,
                 DecomposeBatchNorm,
diff --git a/backends/xnnpack/_passes/convert_to_sdpa.py b/backends/xnnpack/_passes/convert_to_sdpa.py
index a4c421bfede..b926d62df8f 100644
--- a/backends/xnnpack/_passes/convert_to_sdpa.py
+++ b/backends/xnnpack/_passes/convert_to_sdpa.py
@@ -11,12 +11,11 @@
 
 import torch
 from executorch.backends.transforms import get_shape
-from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
-    LiftConstantScalarOperandsPass,
+from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import (
+    RemoveNoopExpandCopyPass,
 )
 from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
 from executorch.backends.xnnpack.partition.graphs import sdpa
-from executorch.backends.xnnpack.utils.utils import get_param_tensor
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.fx.passes.utils.matcher_utils import InternalMatch, SubgraphMatcher
@@ -30,30 +29,18 @@ def get_scale(self, match: InternalMatch) -> Optional[float]:
         """
         Return the SDPA scale recovered from the matched pre-QK^T multiplications.
 
-        The multiplier may be a scalar literal or a constant tensor introduced by
-        scalar lifting. The decomposition applies the square root of the attention
-        scale before QK^T, so the extracted multiplier is squared to recover the
-        original value.
+        The decomposition applies the square root of the attention scale before
+        QK^T, so the extracted multiplier is squared to recover the original value.
         """
         for node in match.nodes_map.values():
-            if node.op != "call_function" or node.target not in {
-                exir_ops.edge.aten.mul.Scalar,
-                exir_ops.edge.aten.mul.Tensor,
-            }:
+            if (
+                node.op != "call_function"
+                or node.target != exir_ops.edge.aten.mul.Scalar
+            ):
                 continue
 
             scale = node.args[1]
 
-            # Extract the scale from the constant tensor introduced by scalar
-            # lifting.
-            if node.target == exir_ops.edge.aten.mul.Tensor:
-                if not isinstance(scale, torch.fx.Node):
-                    continue
-                scale_tensor = get_param_tensor(self.exported_program, scale)
-                if scale_tensor is None or scale_tensor.numel() != 1:
-                    continue
-                scale = scale_tensor.item()
-
             dtype = torch.float
             mul_val = node.meta.get("val", None)
             if mul_val is not None:
@@ -116,14 +103,11 @@ def call(self, graph_module: torch.fx.GraphModule):
         logger.debug(graph_module.print_readable(print_output=False))
 
         for scalar_pattern in sdpa.get_graphs():
-            # Deep-copy the cached scalar pattern so lifting it does not modify the
-            # pattern used by non-lifted flows.
-            tensor_pattern = deepcopy(scalar_pattern)
-            tensor_pattern = LiftConstantScalarOperandsPass()(
-                tensor_pattern
+            normalized_pattern = RemoveNoopExpandCopyPass()(
+                deepcopy(scalar_pattern)
             ).graph_module
 
-            for pattern in (scalar_pattern, tensor_pattern):
+            for pattern in (scalar_pattern, normalized_pattern):
                 sm = SubgraphMatcher(pattern.graph, ignore_literals=True)
                 matches = list(sm.match(graph_module.graph))
                 for partition_to_replace in matches:
diff --git a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
index b0a21b273d5..01524a3268d 100644
--- a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
+++ b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py
@@ -1,5 +1,3 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
 # Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -11,24 +9,27 @@
 from typing import Dict, Optional, Union
 
 import torch
+from executorch.backends.transforms.utils import create_constant_placeholder
+from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
-from executorch.exir.pass_base import ExportPass, PassResult
+from executorch.exir.pass_base import PassResult
 from torch._ops import OpOverload
+from torch.export import ExportedProgram
+from torch.export.graph_signature import InputKind
 
 ScalarOp = Union[EdgeOpOverload, OpOverload]
 
 
-class LiftConstantScalarOperandsPass(ExportPass):
+class LiftConstantScalarOperandsPass(XNNPACKPass):
     """
     Lift scalar operands into tensor constants for selected binary ops.
 
     XNNPACK already supports the tensor overloads for these binary operations.
     This pass converts explicitly listed scalar overloads to their tensor
     overloads by replacing constant scalar operands with small tensor constants.
-    The constants are registered as buffers so they do not become portable
-    ``full`` kernels. Keep the op map narrow until each new scalar overload is
-    covered by tests.
+    The constants are registered as exported-program constant tensor inputs.
+    Keep the op map narrow until each new scalar overload is covered by tests.
     """
 
     default_scalar_to_tensor_ops: Dict[ScalarOp, ScalarOp] = {
@@ -37,9 +38,10 @@ class LiftConstantScalarOperandsPass(ExportPass):
 
     def __init__(
         self,
+        exported_program: ExportedProgram,
         scalar_to_tensor_ops: Optional[Dict[ScalarOp, ScalarOp]] = None,
     ) -> None:
-        super().__init__()
+        super().__init__(exported_program)
         self.scalar_to_tensor_ops = (
             scalar_to_tensor_ops
             if scalar_to_tensor_ops is not None
@@ -58,24 +60,28 @@ def _create_constant_node(
 
         input_value = input_node.meta["val"]
         tensor = torch.tensor(value, dtype=input_value.dtype, device=input_value.device)
-        name = self._get_new_attr_name(graph_module)
-        # ExportPass has no ExportedProgram access to create a constant placeholder.
-        # Keep constants as module attributes so the portable path can emit them
-        # without introducing aten.full, while XNNPACK can still read them as params.
-        graph_module.register_buffer(name, tensor)
-
-        fake_mode = node.meta["val"].fake_mode
-        with graph_module.graph.inserting_before(node):
-            constant_node = graph_module.graph.get_attr(name)
-            constant_node.meta["val"] = fake_mode.from_tensor(
-                tensor, static_shapes=True
+        name = self._get_new_constant_name(graph_module)
+        first_placeholder = next(
+            graph_node
+            for graph_node in graph_module.graph.nodes
+            if graph_node.op == "placeholder"
+        )
+        with graph_module.graph.inserting_before(first_placeholder):
+            return create_constant_placeholder(
+                self.exported_program,
+                graph_module.graph,
+                name,
+                InputKind.CONSTANT_TENSOR,
+                tensor,
             )
-        return constant_node
 
-    def _get_new_attr_name(self, graph_module: torch.fx.GraphModule) -> str:
+    def _get_new_constant_name(self, graph_module: torch.fx.GraphModule) -> str:
         prefix = "_tensor_constant_"
+        existing_names = {node.name for node in graph_module.graph.nodes}
+        existing_names.update(self.exported_program.constants)
+        existing_names.update(self.exported_program.state_dict)
         index = 0
-        while hasattr(graph_module, f"{prefix}{index}"):
+        while f"{prefix}{index}" in existing_names:
             index += 1
         return f"{prefix}{index}"
 
@@ -95,8 +101,8 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             input_value = node.args[0].meta.get("val")
             output_value = node.meta.get("val")
             if (
-                input_value is None
-                or output_value is None
+                not isinstance(input_value, torch.Tensor)
+                or not isinstance(output_value, torch.Tensor)
                 or input_value.dtype != output_value.dtype
             ):
                 continue
diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py
index c6c54f083d6..651eba6234a 100644
--- a/backends/xnnpack/partition/config/__init__.py
+++ b/backends/xnnpack/partition/config/__init__.py
@@ -43,6 +43,7 @@
     MeanDimConfig,
     MinimumConfig,
     MulConfig,
+    MulScalarConfig,
     NegConfig,
     PermuteConfig,
     PowConfig,
@@ -106,6 +107,7 @@
     MinimumConfig,
     MMConfig,
     MulConfig,
+    MulScalarConfig,
     NegConfig,
     PermuteConfig,
     PowConfig,
diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py
index 8a769a8383e..8c1cab82eb3 100644
--- a/backends/xnnpack/partition/config/generic_node_configs.py
+++ b/backends/xnnpack/partition/config/generic_node_configs.py
@@ -8,6 +8,7 @@
 # pyre-unsafe
 
 import logging
+from numbers import Number
 from typing import cast, List, Optional
 
 import numpy as np
@@ -336,6 +337,35 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
         return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]
 
 
+class MulScalarConfig(GenericNodePartitionerConfig):
+    target_name = "mul.Scalar"
+
+    def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
+        if not self.check_common_constraints(node, ep):
+            return False
+
+        if (
+            len(node.args) != 2
+            or not isinstance(node.args[0], torch.fx.Node)
+            or not isinstance(node.args[1], Number)
+        ):
+            return False
+
+        input_value = node.args[0].meta.get("val")
+        output_value = node.meta.get("val")
+        return (
+            isinstance(input_value, torch.Tensor)
+            and isinstance(output_value, torch.Tensor)
+            and input_value.dtype == output_value.dtype
+        )
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+    def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
+        return torch.ops.aten.mul.Scalar
+
+
 class MaximumConfig(GenericNodePartitionerConfig):
     target_name = "maximum.default"
 
diff --git a/backends/xnnpack/test/ops/test_multiply.py b/backends/xnnpack/test/ops/test_multiply.py
index 118136fcd08..a44b5e6b406 100644
--- a/backends/xnnpack/test/ops/test_multiply.py
+++ b/backends/xnnpack/test/ops/test_multiply.py
@@ -9,10 +9,6 @@
 
 import torch
 from executorch.backends.xnnpack.test.tester import Tester
-from executorch.backends.xnnpack.utils.configs import (
-    get_transform_passes,
-    get_xnnpack_edge_compile_config,
-)
 
 
 class TestMul(unittest.TestCase):
@@ -71,12 +67,7 @@ def test_fp32_mul_scalar(self):
         (
             Tester(self.MulScalar(), (torch.randn(2, 3),))
             .export()
-            .to_edge_transform_and_lower(
-                transform_passes=get_transform_passes(),
-                edge_compile_config=get_xnnpack_edge_compile_config(
-                    skip_dim_order=True
-                ),
-            )
+            .to_edge_transform_and_lower()
             .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
             .check_not(
                 [
diff --git a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
index 536e558a53c..dd59aab5204 100644
--- a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
+++ b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py
@@ -1,23 +1,23 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
 # Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
 import unittest
-from copy import deepcopy
 
 import torch
+from executorch.backends.xnnpack._passes import XNNPACKPassManager
 from executorch.backends.xnnpack._passes.convert_to_sdpa import ConvertToSDPAPass
 from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
     LiftConstantScalarOperandsPass,
 )
-from executorch.backends.xnnpack.partition.graphs import sdpa
-from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config
+from executorch.backends.xnnpack.utils.configs import (
+    get_transform_passes,
+    get_xnnpack_edge_compile_config,
+)
 from executorch.exir import to_edge
 from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.pass_manager import ExportedProgramPassManager
+from torch.export.graph_signature import InputKind
 
 
 class TestLiftConstantScalarOperandsPass(unittest.TestCase):
@@ -42,14 +42,16 @@ def _to_edge_program_manager(self, module):
             compile_config=get_xnnpack_edge_compile_config(skip_dim_order=True),
         )
 
-    def _to_edge_graph(self, module):
-        edge = self._to_edge_program_manager(module)
-        return ExportedProgramPassManager([LiftConstantScalarOperandsPass()])(
-            edge.exported_program()
-        ).exported_program
+    def _lift(self, exported_program):
+        return XNNPACKPassManager(
+            exported_program, passes=[LiftConstantScalarOperandsPass]
+        ).transform()
 
     def test_lifts_mul_scalar_operand(self):
-        graph = self._to_edge_graph(self.MulScalar()).graph_module.graph
+        exported_program = self._lift(
+            self._to_edge_program_manager(self.MulScalar()).exported_program()
+        )
+        graph = exported_program.graph_module.graph
 
         self.assertFalse(
             any(node.target == exir_ops.edge.aten.mul.Scalar for node in graph.nodes)
@@ -57,64 +59,84 @@ def test_lifts_mul_scalar_operand(self):
         self.assertTrue(
             any(node.target == exir_ops.edge.aten.mul.Tensor for node in graph.nodes)
         )
-        self.assertTrue(any(node.op == "get_attr" for node in graph.nodes))
+        self.assertFalse(any(node.op == "get_attr" for node in graph.nodes))
+
+        constant_specs = [
+            spec
+            for spec in exported_program.graph_signature.input_specs
+            if spec.kind == InputKind.CONSTANT_TENSOR
+        ]
+        self.assertEqual(len(constant_specs), 1)
+        constant_spec = constant_specs[0]
+        self.assertIn(constant_spec.target, exported_program.constants)
+
+        placeholders = [node for node in graph.nodes if node.op == "placeholder"]
+        self.assertEqual(placeholders[0].name, constant_spec.arg.name)
+        mul_node = next(
+            node for node in graph.nodes if node.target == exir_ops.edge.aten.mul.Tensor
+        )
+        self.assertIs(mul_node.args[1], placeholders[0])
 
-    def test_lifted_mul_scalar_can_emit_without_delegation(self):
-        edge = self._to_edge_program_manager(self.MulScalar()).transform(
-            (LiftConstantScalarOperandsPass(),)
+    def test_is_idempotent(self):
+        exported_program = self._lift(
+            self._to_edge_program_manager(self.MulScalar()).exported_program()
         )
+        exported_program = self._lift(exported_program)
 
-        self.assertIsNotNone(edge.to_executorch())
+        constant_specs = [
+            spec
+            for spec in exported_program.graph_signature.input_specs
+            if spec.kind == InputKind.CONSTANT_TENSOR
+        ]
+        self.assertEqual(len(constant_specs), 1)
+        self.assertEqual(len(exported_program.constants), 1)
 
     def test_keeps_unmapped_scalar_op(self):
-        graph = self._to_edge_graph(self.AddScalar()).graph_module.graph
+        exported_program = self._lift(
+            self._to_edge_program_manager(self.AddScalar()).exported_program()
+        )
+        graph = exported_program.graph_module.graph
 
         self.assertTrue(
             any(node.target == exir_ops.edge.aten.add.Scalar for node in graph.nodes)
         )
+        self.assertFalse(exported_program.constants)
 
-    def test_lifts_sdpa_scale_mul_scalar(self):
-        for graph in sdpa.get_graphs():
-            graph_module = deepcopy(graph)
-
-            graph_module = LiftConstantScalarOperandsPass()(graph_module).graph_module
-
-            scale_mul_count = 0
-            lifted_mul_count = 0
-            for node in graph_module.graph.nodes:
-                if node.op != "call_function":
-                    continue
-                if node.target == exir_ops.edge.aten.mul.Scalar:
-                    scale_mul_count += 1
-                if node.target == exir_ops.edge.aten.mul.Tensor:
-                    lifted_mul_count += 1
-
-            self.assertEqual(scale_mul_count, 0)
-            self.assertEqual(lifted_mul_count, 2)
-
-    def test_converts_sdpa_after_lifting_scale_mul_scalar(self):
+    def test_converts_sdpa_after_default_transform_passes(self):
         q = torch.randn(2, 4, 8, 16)
         k = torch.randn(2, 4, 8, 16)
         v = torch.randn(2, 4, 8, 16)
         mask = torch.randn(8, 8)
-        edge = to_edge(
-            torch.export.export(self.SDPA(), (q, k, v, mask), strict=True),
-            compile_config=get_xnnpack_edge_compile_config(),
-        )
-        exported_program = ExportedProgramPassManager(
-            [LiftConstantScalarOperandsPass()]
-        )(edge.exported_program()).exported_program
-        exported_program = ExportedProgramPassManager(
-            [ConvertToSDPAPass(exported_program)]
-        )(exported_program).exported_program
-
-        graph = exported_program.graph_module.graph
-        self.assertTrue(
-            any(
-                node.target == exir_ops.edge.aten.scaled_dot_product_attention.default
-                for node in graph.nodes
-            )
-        )
-        self.assertFalse(
-            any(node.target == exir_ops.edge.aten.bmm.default for node in graph.nodes)
-        )
+        for use_default_transforms in (False, True):
+            with self.subTest(use_default_transforms=use_default_transforms):
+                edge = to_edge(
+                    torch.export.export(self.SDPA(), (q, k, v, mask), strict=True),
+                    compile_config=get_xnnpack_edge_compile_config(),
+                )
+                if use_default_transforms:
+                    edge = edge.transform(get_transform_passes())
+                exported_program = XNNPACKPassManager(
+                    edge.exported_program(),
+                    passes=[ConvertToSDPAPass, LiftConstantScalarOperandsPass],
+                ).transform()
+
+                graph = exported_program.graph_module.graph
+                self.assertTrue(
+                    any(
+                        node.target
+                        == exir_ops.edge.aten.scaled_dot_product_attention.default
+                        for node in graph.nodes
+                    )
+                )
+                self.assertFalse(
+                    any(
+                        node.target == exir_ops.edge.aten.bmm.default
+                        for node in graph.nodes
+                    )
+                )
+                self.assertFalse(
+                    any(
+                        node.target == exir_ops.edge.aten.mul.Scalar
+                        for node in graph.nodes
+                    )
+                )
diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py
index 894fab4098f..278646b9cf1 100644
--- a/backends/xnnpack/test/test_xnnpack_partitioner.py
+++ b/backends/xnnpack/test/test_xnnpack_partitioner.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -30,6 +31,39 @@ def __init__(self):
         def forward(self, x):
             return self.linear(x)
 
+    class MulScalar(torch.nn.Module):
+        def forward(self, x):
+            return torch.ops.aten.mul.Scalar(x, 0.5)
+
+    def test_mul_scalar_ops_to_not_decompose(self):
+        partitioner = XnnpackPartitioner()
+        exported_program = export(self.MulScalar(), (torch.randn(2, 3),))
+        ops, filter_fn = partitioner.ops_to_not_decompose(exported_program)
+
+        self.assertIn(torch.ops.aten.mul.Scalar, ops)
+        self.assertIsNotNone(filter_fn)
+        mul_node = next(
+            node
+            for node in exported_program.graph.nodes
+            if node.target == torch.ops.aten.mul.Scalar
+        )
+        self.assertTrue(filter_fn(mul_node))
+
+    def test_mul_scalar_ops_to_not_decompose_rejects_unsupported_dtype(self):
+        partitioner = XnnpackPartitioner()
+        exported_program = export(
+            self.MulScalar(), (torch.ones(2, 3, dtype=torch.int32),)
+        )
+        _, filter_fn = partitioner.ops_to_not_decompose(exported_program)
+
+        self.assertIsNotNone(filter_fn)
+        mul_node = next(
+            node
+            for node in exported_program.graph.nodes
+            if node.target == torch.ops.aten.mul.Scalar
+        )
+        self.assertFalse(filter_fn(mul_node))
+
     def test_deprecation_warning_for_to_backend_workflow(self):
         """
         Test that the deprecated to_edge + to_backend workflow shows a deprecation warning.
diff --git a/backends/xnnpack/test/tester/tester.py b/backends/xnnpack/test/tester/tester.py
index 396e149565f..fc12da231c0 100644
--- a/backends/xnnpack/test/tester/tester.py
+++ b/backends/xnnpack/test/tester/tester.py
@@ -1,6 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
-# Copyright 2024-2026 Arm Limited and/or its affiliates.
+# Copyright 2024-2025 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -24,11 +24,9 @@
     QuantizationConfig,
 )
 from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config
-from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
+from executorch.exir import EdgeCompileConfig
 from executorch.exir.backend.partitioner import Partitioner
-from executorch.exir.pass_manager import PassType as ExirPassType
 from torch._export.pass_base import PassType
-from torch.export import ExportedProgram
 from torchao.quantization.pt2e.quantizer import Quantizer
 
 
@@ -79,7 +77,6 @@ def __init__(
         self,
         partitioners: Optional[List[Partitioner]] = None,
         edge_compile_config: Optional[EdgeCompileConfig] = None,
-        transform_passes: Optional[List[ExirPassType]] = None,
     ):
         super().__init__(
             default_partitioner_cls=XnnpackPartitioner,
@@ -87,21 +84,6 @@ def __init__(
             edge_compile_config=edge_compile_config
             or get_xnnpack_edge_compile_config(),
         )
-        self.transform_passes = transform_passes
-
-    def run(
-        self,
-        artifact: ExportedProgram,
-        inputs=None,
-        generate_etrecord: bool = False,
-    ) -> None:
-        self.edge_dialect_program = to_edge_transform_and_lower(
-            artifact,
-            transform_passes=self.transform_passes,
-            compile_config=self.edge_compile_conf,
-            partitioner=self.partitioners,
-            generate_etrecord=generate_etrecord,
-        )
 
 
 class Partition(BaseStages.Partition):
@@ -150,37 +132,3 @@ def __init__(
             dynamic_shapes=dynamic_shapes,
             **kwargs,
         )
-
-    def to_edge_transform_and_lower(
-        self,
-        to_edge_and_transform_stage: Optional[
-            BaseStages.ToEdgeTransformAndLower
-        ] = None,
-        generate_etrecord: bool = False,
-        *,
-        partitioners: Optional[List[Partitioner]] = None,
-        edge_compile_config: Optional[EdgeCompileConfig] = None,
-        transform_passes: Optional[List[ExirPassType]] = None,
-    ):
-        if to_edge_and_transform_stage is None:
-            to_edge_and_transform_stage = ToEdgeTransformAndLower(
-                partitioners=partitioners,
-                edge_compile_config=edge_compile_config,
-                transform_passes=transform_passes,
-            )
-        else:
-            if partitioners is not None:
-                to_edge_and_transform_stage.partitioners = partitioners
-            if edge_compile_config is not None:
-                to_edge_and_transform_stage.edge_compile_conf = edge_compile_config
-            if transform_passes is not None:
-                if not isinstance(to_edge_and_transform_stage, ToEdgeTransformAndLower):
-                    raise ValueError(
-                        "transform_passes requires the XNNPACK "
-                        "ToEdgeTransformAndLower stage."
-                    )
-                to_edge_and_transform_stage.transform_passes = transform_passes
-        return super().to_edge_transform_and_lower(
-            to_edge_and_transform_stage,
-            generate_etrecord=generate_etrecord,
-        )
diff --git a/backends/xnnpack/utils/configs.py b/backends/xnnpack/utils/configs.py
index ec47b81e835..3016e94146b 100644
--- a/backends/xnnpack/utils/configs.py
+++ b/backends/xnnpack/utils/configs.py
@@ -9,9 +9,6 @@
 
 import executorch.exir as exir
 
-from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
-    LiftConstantScalarOperandsPass,
-)
 from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import (
     RemoveNoopExpandCopyPass,
 )
@@ -28,7 +25,7 @@ def get_xnnpack_edge_compile_config(
 
 
 def get_transform_passes(additional_passes=None) -> List[PassType]:
-    passes = [RemoveNoopExpandCopyPass(), LiftConstantScalarOperandsPass()]
+    passes = [RemoveNoopExpandCopyPass()]
     if additional_passes:
         passes.extend(additional_passes)
     return passes