From 5721f54c0513924683aae7d4cd54338440b8f7c2 Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Mon, 21 Sep 2026 17:21:17 -0400 Subject: [PATCH] [Vulkan] Fuse SwiGLU elementwise operations --- backends/vulkan/custom_ops_lib.py | 12 ++ backends/vulkan/op_registry.py | 10 ++ backends/vulkan/patterns/BUCK | 1 + backends/vulkan/patterns/__init__.py | 2 + backends/vulkan/patterns/swiglu.py | 65 +++++++++ .../graph/ops/glsl/binary_op_buffer.glsl | 5 + .../graph/ops/glsl/binary_op_buffer.yaml | 6 + .../graph/ops/glsl/binary_op_texture.glsl | 5 + .../graph/ops/glsl/binary_op_texture.yaml | 6 + .../runtime/graph/ops/impl/BinaryOp.cpp | 2 + backends/vulkan/test/test_vulkan_passes.py | 115 ++++++++++++++++ .../vulkan/test/vulkan_compute_api_test.cpp | 129 ++++++++++++++++++ 12 files changed, 358 insertions(+) create mode 100644 backends/vulkan/patterns/swiglu.py diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index a074597466a..3bef9899981 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -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 ## ################ diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index add4d01a78e..6e0cbff6bcb 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -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, diff --git a/backends/vulkan/patterns/BUCK b/backends/vulkan/patterns/BUCK index bc9f97a804c..b83b2401428 100644 --- a/backends/vulkan/patterns/BUCK +++ b/backends/vulkan/patterns/BUCK @@ -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/...", diff --git a/backends/vulkan/patterns/__init__.py b/backends/vulkan/patterns/__init__.py index 68df9905671..4b3e8d87f62 100644 --- a/backends/vulkan/patterns/__init__.py +++ b/backends/vulkan/patterns/__init__.py @@ -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 ( diff --git a/backends/vulkan/patterns/swiglu.py b/backends/vulkan/patterns/swiglu.py new file mode 100644 index 00000000000..11fc7b47c15 --- /dev/null +++ b/backends/vulkan/patterns/swiglu.py @@ -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) diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.glsl index ac4751a47e5..c4a4f8edeba 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.glsl @@ -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)) { diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml index d1f3600cfc9..d246dcb9f20 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml @@ -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 diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.glsl index cb224eb7a0a..13d5c6b26bc 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.glsl @@ -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); diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml index 07e3905af02..d6a72b543b8 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml @@ -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 diff --git a/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp b/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp index d78f4d6a3ff..55293087398 100644 --- a/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp @@ -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); diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index b89a2ab6dde..dda8fa1023e 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -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], diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index a0ca49cff6b..d77b8e8c95c 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -9,7 +9,9 @@ #include #include +#include #include +#include #include #include @@ -4239,3 +4241,130 @@ TEST(VulkanWorkGroupSizeTest, compute_graph_preserves_dispatch_intent) { EXPECT_EQ(graph.create_lwg(buffer_gwg), buffer_gwg.required_lwg_size()); EXPECT_FALSE(texture_gwg.is_linear()); } + +class VulkanSwiGLUTest : public ::testing::TestWithParam> {}; + +TEST_P(VulkanSwiGLUTest, MatchesUnfusedAcrossResizes) { + if (!api::available()) { + GTEST_SKIP(); + } + const auto storage = std::get<0>(GetParam()); + const auto dtype = std::get<1>(GetParam()); + const auto layout = std::get<2>(GetParam()); + const auto width = std::get<3>(GetParam()); + const auto broadcast = std::get<4>(GetParam()); + GraphConfig config; + config.set_storage_type_override(storage); + config.set_memory_layout_override(layout); + ComputeGraph graph(config); + const std::vector max_shape = {2, 63, width}; + const std::vector up_shape = + broadcast ? std::vector{1, 1, width} : max_shape; + const IOValueRef gate = graph.add_input_tensor(max_shape, dtype); + const IOValueRef up = graph.add_input_tensor(up_shape, dtype); + const ValueRef sigmoid = graph.add_tensor(max_shape, dtype); + const ValueRef silu = graph.add_tensor(max_shape, dtype); + const ValueRef unfused = graph.add_tensor(max_shape, dtype); + const ValueRef fused = graph.add_tensor(max_shape, dtype); + VK_GET_OP_FN("aten.sigmoid.default")(graph, {gate.value, sigmoid}); + VK_GET_OP_FN("aten.mul.Tensor")(graph, {gate.value, sigmoid, silu}); + VK_GET_OP_FN("aten.mul.Tensor")(graph, {silu, up.value, unfused}); + VK_GET_OP_FN("et_vk.swiglu.default")(graph, {gate.value, up.value, fused}); + const ValueRef unfused_staging = graph.set_output_tensor(unfused); + const ValueRef fused_staging = graph.set_output_tensor(fused); + graph.prepare(); + graph.prepack(); + + auto check_resizes = [&](auto scalar) { + using T = decltype(scalar); + for (const int64_t rows : {1, 7, 63, 1}) { + SCOPED_TRACE(rows); + const std::vector shape = {2, rows, width}; + graph.resize_input(0, shape); + if (!broadcast) { + graph.resize_input(1, shape); + } + graph.propagate_resize(); + ASSERT_EQ(graph.sizes_of(fused), shape); + const size_t count = 2 * rows * width; + std::vector gate_data(count); + std::vector up_data(broadcast ? width : count); + for (size_t i = 0; i < gate_data.size(); ++i) { + gate_data[i] = 20.0f * std::sin(static_cast(i)); + } + for (size_t i = 0; i < up_data.size(); ++i) { + up_data[i] = 2.0f * std::cos(static_cast(i)); + } + graph.maybe_cast_and_copy_into_staging( + gate.staging, gate_data.data(), gate_data.size(), dtype); + graph.maybe_cast_and_copy_into_staging( + up.staging, up_data.data(), up_data.size(), dtype); + graph.execute(); + std::vector expected(count); + std::vector actual(count); + graph.maybe_cast_and_copy_from_staging( + unfused_staging, expected.data(), count, dtype); + graph.maybe_cast_and_copy_from_staging( + fused_staging, actual.data(), count, dtype); + for (size_t i = 0; i < count; ++i) { + const float x = float(gate_data[i]); + const float y = float(up_data[broadcast ? i % width : i]); + const float reference = (x / (1.0f + std::exp(-x))) * y; + if (dtype == vkapi::kHalf) { + ASSERT_NEAR( + float(actual[i]), reference, 1e-3f * std::abs(reference) + 1e-7f) + << "element " << i; + // The unfused path rounds twice more, including subnormal sigmoid. + const float atol = 6e-8f * (std::abs(x * y) + std::abs(y) + 1.0f); + ASSERT_NEAR( + float(actual[i]), + float(expected[i]), + 3e-3f * std::abs(float(expected[i])) + atol) + << "element " << i; + } else { + ASSERT_NEAR( + float(actual[i]), reference, 2e-6f * std::abs(reference) + 1e-7f) + << "element " << i; + ASSERT_NEAR( + float(actual[i]), + float(expected[i]), + 2e-6f * std::abs(float(expected[i])) + 1e-7f) + << "element " << i; + } + } + } + }; + if (dtype == vkapi::kHalf) { + check_resizes(executorch::aten::Half{}); + } else { + check_resizes(float{}); + } +} + +INSTANTIATE_TEST_SUITE_P( + StorageDtypeLayoutWidthBroadcast, + VulkanSwiGLUTest, + ::testing::Combine( + ::testing::Values(utils::kBuffer, utils::kTexture3D), + ::testing::Values(vkapi::kFloat, vkapi::kHalf), + ::testing::Values( + utils::kWidthPacked, + utils::kHeightPacked, + utils::kChannelsPacked), + ::testing::Values(int64_t(13), int64_t(1023)), + ::testing::Bool())); + +INSTANTIATE_TEST_SUITE_P( + LlmHiddenSize, + VulkanSwiGLUTest, + ::testing::Combine( + ::testing::Values(utils::kBuffer, utils::kTexture3D), + ::testing::Values(vkapi::kFloat, vkapi::kHalf), + ::testing::Values(utils::kWidthPacked), + ::testing::Values(int64_t(3072)), + ::testing::Bool()));