Arm backend: Enable quantized softplus lowering through TOSA TABLE (pytorch#23522)

Keep quantized softplus intact through partitioning and lower it to a
TOSA TABLE generated from its beta and threshold parameters. This lets
Ethos-U85 delegate the Qwen3.5 DeltaNet gate with INT8 or INT16
activations.

Cover TOSA INT8, INT16, and FP paths, U55 and U85 FVP runs, and VGF
emulation. Map the exported op to its public PyTorch API in the
generated support pages.

AI-Assisted-by: Codex


cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Xingguo Li <xingguo.li@arm.com>

Author: xingguo01

backends/arm/_passes/arm_pass.py

@@ -89,6 +89,7 @@ def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False)
         ops_without_quantized_fake_kernel = {
             exir_ops.edge.aten.bmm.default,
             exir_ops.edge.aten.leaky_relu.default,
+            exir_ops.edge.aten.softplus.default,
         }
         if (
             op in ops_without_quantized_fake_kernel

backends/arm/_passes/insert_table_ops.py

@@ -66,6 +66,7 @@ class TableOps:
 
     # Targets that must be treated explicitly
     special_table_ops: Set[EdgeOpOverload] = {
+        exir_ops.edge.aten.softplus.default,
         exir_ops.edge.aten.pow.Tensor_Scalar,
         exir_ops.edge.aten.gelu.default,
         exir_ops.edge.aten.elu.default,
@@ -87,6 +88,26 @@ def __getitem__(self, node: Node):
             return self.unary_table_ops[target]
         elif target in self.special_table_ops:
             match target:
+                case exir_ops.edge.aten.softplus.default:
+                    beta = cast(
+                        float,
+                        (
+                            node.args[1]
+                            if len(node.args) > 1
+                            else node.kwargs.get("beta", 1)
+                        ),
+                    )
+                    threshold = cast(
+                        float,
+                        (
+                            node.args[2]
+                            if len(node.args) > 2
+                            else node.kwargs.get("threshold", 20)
+                        ),
+                    )
+                    return lambda x: torch.nn.functional.softplus(
+                        x, beta=beta, threshold=threshold
+                    ).flatten()
                 case exir_ops.edge.aten.pow.Tensor_Scalar:
                     # Exponent is a constant. Embed it into a lambda.
                     exp = cast(int, node.args[1])

backends/arm/operator_support/tosa_profile_supported_op_lists.py

@@ -56,6 +56,7 @@
     exir_ops.edge.aten.expm1.default,
     exir_ops.edge.aten.log.default,
     exir_ops.edge.aten.log1p.default,
+    exir_ops.edge.aten.softplus.default,
     exir_ops.edge.aten.log10.default,
     exir_ops.edge.aten.linear.default,
     exir_ops.edge.aten.split_with_sizes_copy.default,

backends/arm/quantizer/quantization_annotator.py

@@ -556,6 +556,7 @@ def _get_fixed_qparams_qspec(
     torch.ops.aten.pow.Tensor_Scalar,
     torch.ops.aten.gelu.default,
     torch.ops.aten.silu.default,
+    torch.ops.aten.softplus.default,
     torch.ops.aten.silu_.default,
     torch.ops.aten.sinh.default,
     torch.ops.aten.atan.default,

backends/arm/scripts/docgen/generate_op_support.py

@@ -598,6 +598,7 @@ def _active_explicit_backend_coverage() -> dict[tuple[str, str], dict[str, set[s
     "torch.ops.aten.logical_xor.default": ("torch.logical_xor",),
     "torch.ops.aten.logical_not.default": ("torch.logical_not",),
     # Linear algebra and neural-network layers.
+    "torch.ops.aten.softplus.default": ("torch.nn.functional.softplus",),
     "torch.ops.aten.linear.default": ("torch.nn.Linear", "torch.nn.functional.linear"),
     "torch.ops.aten.mm.default": ("torch.mm",),
     "torch.ops.aten.bmm.default": ("torch.bmm",),

backends/arm/test/misc/test_partition_decomposed_quantized_ops.py

@@ -4,9 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 # Test that tosa_supported_operators reject operators that are not
-# quantized properly. This is typically a consequence of a torch op
-# such a Softplus that is decompsed into many other ops without
-# surrounding q/dq nodes.
+# quantized properly after decomposition.
 
 import torch
 from executorch.backends.arm.test import common
@@ -48,11 +46,7 @@
 
 
 class SoftplusModule(torch.nn.Module):
-    """Module containing an addition followed by a Softplus.
-
-    Softplus is currently not supported by TosaBackend.
-
-    """
+    """Addition followed by Softplus for FP decomposition and INT lowering."""
 
     def __init__(self):
         super().__init__()
@@ -85,8 +79,6 @@ def forward(self, x: torch.Tensor):
         return x1 + x3
 
 
-# Softplus is decomposed which messes up the quantization. This test tests that CheckProperQuantization does not
-# partition nodes where quantization is not as expected.
 @common.parametrize("test_data", test_data)
 def test_softplus_tosa_FP(test_data: input_t1):
     pipeline = TosaPipelineFP[input_t1](
@@ -106,15 +98,7 @@ def test_softplus_tosa_INT(test_data: input_t1):
         SoftplusModule(),
         test_data=test_data,
         aten_op=softplus_aten_op,
-        exir_op=softplus_exir_op,
-    )
-    pipeline.pop_stage("check_not.exir")
-    # check that all ops in softplus_exir_op except add are rejected
-    pipeline.add_stage_after(
-        "to_edge_transform_and_lower",
-        pipeline.tester.check,
-        softplus_exir_op[1:],
-        suffix="exir_post_partition",
+        exir_op="executorch_exir_dialects_edge__ops_aten_softplus_default",
     )
     pipeline.run()
 

backends/arm/test/modules/test_qwen_delta_gate.py

@@ -0,0 +1,106 @@
+# 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 torch
+import torch.nn.functional as F
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
+    TosaPipelineINT,
+    VgfPipeline,
+)
+
+
+class QwenDeltaGate(torch.nn.Module):
+    a_log: torch.Tensor
+    dt_bias: torch.Tensor
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.register_buffer("a_log", torch.log(torch.tensor([0.5, 1.0, 1.5, 2.0])))
+        self.register_buffer("dt_bias", torch.tensor([-2.0, -0.5, 0.5, 2.0]))
+
+    def forward(self, a: torch.Tensor) -> torch.Tensor:
+        g = -self.a_log.exp() * F.softplus(a + self.dt_bias)
+        return g.exp()
+
+
+@common.XfailIfNoCorstone300
+def test_qwen_gate_u55_int():
+    values = torch.linspace(-4.0, 4.0, 32).reshape(1, 8, 4)
+    pipeline = EthosU55PipelineINT(
+        QwenDeltaGate(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        "executorch_exir_dialects_edge__ops_aten_softplus_default",
+        run_on_fvp=True,
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone320
+def test_qwen_gate_u85_int8():
+    values = torch.linspace(-4.0, 4.0, 32).reshape(1, 8, 4)
+    pipeline = EthosU85PipelineINT(
+        QwenDeltaGate(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        "executorch_exir_dialects_edge__ops_aten_softplus_default",
+        run_on_fvp=True,
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone320
+def test_qwen_gate_u85_int16():
+    values = torch.linspace(-4.0, 4.0, 32).reshape(1, 8, 4)
+    pipeline = EthosU85PipelineINT(
+        QwenDeltaGate(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        "executorch_exir_dialects_edge__ops_aten_softplus_default",
+        run_on_fvp=True,
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+def test_qwen_gate_vgf_quant():
+    values = torch.linspace(-4.0, 4.0, 32).reshape(1, 8, 4)
+    pipeline = VgfPipeline(
+        QwenDeltaGate(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        "executorch_exir_dialects_edge__ops_aten_softplus_default",
+        run_on_vulkan_runtime=True,
+    )
+    pipeline.run()
+
+
+def test_qwen_gate_tosa_int():
+    values = torch.linspace(-4.0, 4.0, 32).reshape(1, 8, 4)
+    pipeline = TosaPipelineINT(
+        QwenDeltaGate(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        "executorch_exir_dialects_edge__ops_aten_softplus_default",
+    )
+    pipeline.count_tosa_ops({"TABLE": 2})
+    pipeline.run()
+
+
+def test_qwen_gate_tosa_a16w8():
+    values = torch.linspace(-4.0, 4.0, 32).reshape(1, 8, 4)
+    pipeline = TosaPipelineINT(
+        QwenDeltaGate(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        "executorch_exir_dialects_edge__ops_aten_softplus_default",
+        tosa_extensions=["int16"],
+    )
+    pipeline.count_tosa_ops({"TABLE": 2})
+    pipeline.run()

backends/arm/test/ops/test_softplus.py

@@ -0,0 +1,52 @@
+# 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 torch
+import torch.nn.functional as F
+from executorch.backends.arm.test.tester.test_pipeline import (
+    TosaPipelineFP,
+    TosaPipelineINT,
+)
+
+
+class SoftplusWithParameters(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.softplus(x, beta=2.0, threshold=5.0)
+
+
+def test_softplus_parameters_tosa_int():
+    values = torch.tensor([[-10.0, -3.0, -1.0, 0.0, 1.0, 3.0, 10.0]])
+    pipeline = TosaPipelineINT(
+        SoftplusWithParameters(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        "executorch_exir_dialects_edge__ops_aten_softplus_default",
+    )
+    pipeline.count_tosa_ops({"TABLE": 1})
+    pipeline.run()
+
+
+def test_softplus_parameters_tosa_a16w8():
+    values = torch.tensor([[-10.0, -3.0, -1.0, 0.0, 1.0, 3.0, 10.0]])
+    pipeline = TosaPipelineINT(
+        SoftplusWithParameters(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        "executorch_exir_dialects_edge__ops_aten_softplus_default",
+        tosa_extensions=["int16"],
+    )
+    pipeline.count_tosa_ops({"TABLE": 1})
+    pipeline.run()
+
+
+def test_softplus_parameters_tosa_fp():
+    values = torch.tensor([[-10.0, -3.0, -1.0, 0.0, 1.0, 3.0, 10.0]])
+    pipeline = TosaPipelineFP(
+        SoftplusWithParameters(),
+        (values,),
+        "torch.ops.aten.softplus.default",
+        [],
+    )
+    pipeline.run()

backends/arm/test/targets.bzl

@@ -32,6 +32,7 @@ def define_arm_tests():
         "ops/test_slice.py",
         "ops/test_sigmoid.py",
         "ops/test_softmax.py",
+        "ops/test_softplus.py",
         "ops/test_sub.py",
         "ops/test_sum.py",
         "ops/test_tanh.py",
@@ -51,6 +52,11 @@ def define_arm_tests():
         "ops/test_split.py",
     ]
 
+    # Modules
+    test_files += [
+        "modules/test_qwen_delta_gate.py",
+    ]
+
     # Export recipes
     test_files += [
         "recipes/test_arm_recipes.py",

backends/arm/tosa/partitioner.py

@@ -879,6 +879,7 @@ def ops_to_not_decompose(  # noqa: C901
             torch.ops.aten.linear.default,
             torch.ops.aten.linspace.default,
             torch.ops.aten.silu.default,
+            torch.ops.aten.softplus.default,
             torch.ops.aten.pad.default,
         }
         ops_to_not_decompose_if_fp = {