pytorch · mergennachin · Sep 21, 2026
@@ -1131,6 +1131,18 @@ def sdpa_impl(
 lib.impl(name, sdpa_impl, "CompositeExplicitAutograd")
 sdpa_op = getattr(getattr(torch.ops, namespace), name)
 
+################
+## swiglu ##
+################
+
+
+def swiglu_impl(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    return (gate * torch.sigmoid(gate)) * up
+
+
+lib.define("swiglu(Tensor gate, Tensor up) -> Tensor")
+lib.impl("swiglu", swiglu_impl, "CompositeExplicitAutograd")
+
 ################
 ## rms_norm ##
 ################

@@ -271,6 +271,16 @@ def register_binaryop_cpp_ops():
     )
 
 
+@update_features(exir_ops.edge.et_vk.swiglu.default)
+def register_swiglu():
+    return OpFeatures(
+        inputs_storage=utils.ANY_STORAGE,
+        inputs_dtypes=utils.FP_T,
+        supports_resize=True,
+        supports_highdim=True,
+    )
+
+
 @update_features(
     [
         exir_ops.edge.aten.eq.Tensor,

@@ -21,6 +21,7 @@ fbcode_target(_kind = runtime.python_library,
         "weight_packing_utils.py",
         "sdpa.py",
         "select_as_symint.py",
+        "swiglu.py",
     ],
     visibility = [
         "//executorch/backends/...",

@@ -28,6 +28,8 @@
 
 import executorch.backends.vulkan.patterns.select_as_symint  # noqa
 
+import executorch.backends.vulkan.patterns.swiglu  # noqa
+
 import torch
 
 from executorch.backends.vulkan.patterns.pattern_registry import (

@@ -0,0 +1,65 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+import torch
+from executorch.exir import ExportedProgram
+from executorch.exir.dialects._ops import ops as exir_ops
+
+from .pattern_registry import (
+    PatternMatch,
+    register_pattern_detector,
+    register_pattern_replacement,
+)
+
+
+@register_pattern_detector("swiglu")
+def find_swiglu_pattern(node: torch.fx.Node) -> Optional[PatternMatch]:
+    if node.target != exir_ops.edge.aten.mul.Tensor:
+        return None
+
+    for silu, up in (node.args, node.args[::-1]):
+        if (
+            not isinstance(silu, torch.fx.Node)
+            or not isinstance(up, torch.fx.Node)
+            or silu.target != exir_ops.edge.aten.mul.Tensor
+            or len(silu.users) != 1
+        ):
+            continue
+        for gate, sigmoid in (silu.args, silu.args[::-1]):
+            if (
+                not isinstance(gate, torch.fx.Node)
+                or not isinstance(sigmoid, torch.fx.Node)
+                or sigmoid.target != exir_ops.edge.aten.sigmoid.default
+                or sigmoid.args != (gate,)
+                or len(sigmoid.users) != 1
+            ):
+                continue
+            values = [n.meta.get("val") for n in (gate, up, sigmoid, silu, node)]
+            if not all(isinstance(v, torch.Tensor) for v in values):
+                continue
+            if not all(v.dtype == values[0].dtype for v in values):
+                continue
+            if values[0].dtype not in (torch.float16, torch.float32):
+                continue
+            return PatternMatch(
+                [gate, up], [node], [sigmoid, silu, node], anchor_node=node
+            )
+    return None
+
+
+@register_pattern_replacement("swiglu")
+def replace_swiglu_pattern(
+    ep: ExportedProgram, graph_module: torch.fx.GraphModule, match: PatternMatch
+):
+    # The up projection can occur after sigmoid/mul, so insert at the final mul.
+    with graph_module.graph.inserting_before(match.anchor_node):
+        fused = graph_module.graph.call_function(
+            exir_ops.edge.et_vk.swiglu.default, tuple(match.input_nodes)
+        )
+    fused.meta = match.anchor_node.meta.copy()
+    match.anchor_node.replace_all_uses_with(fused)
@@ -62,6 +62,11 @@ ${layout_declare_spec_const(C, "int", "other_layout", "CONTIG_LAYOUT_INT")}
 ${layout_declare_spec_const(C, "int", "in_broadcast_packed_dim", "0")}
 ${layout_declare_spec_const(C, "int", "other_broadcast_packed_dim", "0")}
 
+$if OPERATOR == "swiglu(X, Y)":
+  float swiglu(float gate, float up) {
+    return (gate * (1.0 / (1.0 + exp(-gate)))) * up;
+  }
+
 void main() {
   const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {

@@ -22,6 +22,12 @@ binary_op_buffer:
       OPERATOR: X - A * Y
     - NAME: binary_mul_buffer
       OPERATOR: X * Y
+    - NAME: binary_swiglu_buffer
+      OPERATOR: swiglu(X, Y)
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
     - NAME: binary_div_buffer
       OPERATOR: X / Y
     - NAME: binary_pow_buffer

@@ -62,6 +62,11 @@ ${layout_declare_spec_const(C, "int", "other_broadcast_packed_dim", "0")}
 $if MASK_PADDING:
   #define MASK_PADDING
 
+$if OPERATOR == "swiglu(X, Y)":
+  vec4 swiglu(vec4 gate, vec4 up) {
+    return (gate * (1.0 / (1.0 + exp(-gate)))) * up;
+  }
+
 void main() {
   const ivec3 out_pos = ivec3(gl_GlobalInvocationID);
 

@@ -22,6 +22,12 @@ binary_op_texture:
       OPERATOR: X - A * Y
     - NAME: binary_mul_texture3d
       OPERATOR: X * Y
+    - NAME: binary_swiglu_texture3d
+      OPERATOR: swiglu(X, Y)
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
     - NAME: binary_div_texture3d
       OPERATOR: X / Y
       MASK_PADDING: 1

@@ -134,6 +134,7 @@ DEFINE_BINARY_OP_WITH_ALPHA_FN(sub);
 DEFINE_BINARY_OP_WITH_ALPHA_FN(floor_divide);
 
 DEFINE_BINARY_OP_FN(mul);
+DEFINE_BINARY_OP_FN(swiglu);
 DEFINE_BINARY_OP_FN(div);
 DEFINE_BINARY_OP_FN(pow);
 DEFINE_BINARY_OP_FN(minimum);
@@ -149,6 +150,7 @@ REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.add.Tensor, add);
   VK_REGISTER_OP(aten.sub.Tensor, sub);
   VK_REGISTER_OP(aten.mul.Tensor, mul);
+  VK_REGISTER_OP(et_vk.swiglu.default, swiglu);
   VK_REGISTER_OP(aten.div.Tensor, div);
   VK_REGISTER_OP(aten.div.Tensor_mode, floor_divide);
   VK_REGISTER_OP(aten.pow.Tensor_Tensor, pow);

@@ -1,4 +1,5 @@
 import unittest
+from itertools import product
 from typing import List, Optional, Tuple
 
 import torch
@@ -8,6 +9,7 @@
 from executorch.backends.vulkan._passes.remove_redundant_ops import (
     RemoveRedundantOpsTransform,
 )
+from executorch.backends.vulkan.patterns.swiglu import find_swiglu_pattern
 
 from executorch.exir import EdgeCompileConfig, EdgeProgramManager, to_edge
 
@@ -90,6 +92,119 @@ def op_node_count(graph_module: torch.fx.GraphModule, canonical_op_name: str) ->
     return count
 
 
+class TestSwiGLUFusion(unittest.TestCase):
+    def _fuse(self, model, inputs, dynamic_shapes=None):
+        program = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
+        edge = to_edge(
+            program, compile_config=EdgeCompileConfig(_check_ir_validity=False)
+        )
+        ep = edge.exported_program()
+        matches = [
+            match
+            for node in ep.graph_module.graph.nodes
+            if (match := find_swiglu_pattern(node)) is not None
+        ]
+        for match in matches:
+            self.assertTrue(set(match.input_nodes).isdisjoint(match.all_nodes))
+        fuse_pass = FusePatternsPass()
+        fuse_pass._exported_program = ep
+        result = fuse_pass(ep.graph_module)
+        result.graph_module.graph.lint()
+        return result.graph_module
+
+    def test_operand_order_and_broadcast(self):
+        class Model(torch.nn.Module):
+            def __init__(self, reverse_inner, reverse_outer):
+                super().__init__()
+                self.reverse_inner = reverse_inner
+                self.reverse_outer = reverse_outer
+
+            def forward(self, gate, up):
+                sigmoid = torch.sigmoid(gate)
+                silu = sigmoid * gate if self.reverse_inner else gate * sigmoid
+                return up * silu if self.reverse_outer else silu * up
+
+        for reverse_inner, reverse_outer, dtype, broadcast in product(
+            (False, True), (False, True), (torch.float16, torch.float32), (False, True)
+        ):
+            with self.subTest(
+                inner=reverse_inner,
+                outer=reverse_outer,
+                dtype=dtype,
+                broadcast=broadcast,
+            ):
+
+                inputs = (
+                    torch.randn(2, 7, 13, dtype=dtype),
+                    torch.randn(13 if broadcast else (2, 7, 13), dtype=dtype),
+                )
+                model = Model(reverse_inner, reverse_outer)
+                gm = self._fuse(model, inputs)
+                self.assertEqual(op_node_count(gm, "swiglu.default"), 1)
+                self.assertEqual(op_node_count(gm, "sigmoid.default"), 0)
+                self.assertEqual(op_node_count(gm, "mul.Tensor"), 0)
+                torch.testing.assert_close(
+                    gm(*inputs)[0], model(*inputs), rtol=0, atol=0
+                )
+
+    def test_silu_and_dynamic_projection_order(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, gate_weight, up_weight):
+                gate = torch.nn.functional.linear(x, gate_weight)
+                silu = torch.nn.functional.silu(gate)
+                up = torch.nn.functional.linear(x, up_weight)
+                return silu * up
+
+        model = Model()
+        inputs = (torch.randn(3, 16), torch.randn(13, 16), torch.randn(13, 16))
+        gm = self._fuse(
+            model, inputs, ({0: torch.export.Dim("tokens", min=1, max=255)}, None, None)
+        )
+        self.assertEqual(op_node_count(gm, "swiglu.default"), 1)
+        for tokens in (1, 7, 255, 1):
+            args = (torch.randn(tokens, 16), *inputs[1:])
+            torch.testing.assert_close(gm(*args)[0], model(*args))
+
+    def test_shared_intermediates_are_not_fused(self):
+        class Model(torch.nn.Module):
+            def __init__(self, shared_sigmoid):
+                super().__init__()
+                self.shared_sigmoid = shared_sigmoid
+
+            def forward(self, gate, up):
+                sigmoid = torch.sigmoid(gate)
+                silu = gate * sigmoid
+                return silu * up, sigmoid if self.shared_sigmoid else silu
+
+        for shared_sigmoid in (False, True):
+
+            inputs = (torch.randn(2, 13), torch.randn(2, 13))
+            model = Model(shared_sigmoid)
+            gm = self._fuse(model, inputs)
+            self.assertEqual(op_node_count(gm, "swiglu.default"), 0)
+            torch.testing.assert_close(gm(*inputs), model(*inputs))
+
+    def test_different_sigmoid_input_is_not_fused(self):
+        class Model(torch.nn.Module):
+            def forward(self, gate, up):
+                return (gate * torch.sigmoid(up)) * up
+
+        inputs = (torch.randn(2, 13), torch.randn(2, 13))
+        self.assertEqual(
+            op_node_count(self._fuse(Model(), inputs), "swiglu.default"), 0
+        )
+
+    def test_mixed_precision_is_not_fused(self):
+        class Model(torch.nn.Module):
+            def forward(self, gate, up):
+                return (gate * torch.sigmoid(gate)) * up
+
+        inputs = (torch.randn(2, 13, dtype=torch.float16), torch.randn(2, 13))
+        self.assertEqual(
+            op_node_count(self._fuse(Model(), inputs), "swiglu.default"), 0
+        )
+
+
 def run_conv1d_as_conv2d(
     model: torch.nn.Module,
     sample_inputs: Tuple[torch.Tensor],