From 3568c3cbf2d06944a29cfe565121a1764d9145d7 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Sat, 14 Nov 2026 11:29:46 -0800 Subject: [PATCH] Update [ghstack-poisoned] --- backends/vulkan/custom_ops_lib.py | 8 ++- .../graph/ops/glsl/rms_norm_buffer.glsl | 14 ++-- .../graph/ops/glsl/rms_norm_buffer.yaml | 3 + .../graph/ops/glsl/rms_norm_texture.glsl | 15 +++-- .../graph/ops/glsl/rms_norm_texture.yaml | 3 + .../vulkan/runtime/graph/ops/impl/RmsNorm.cpp | 21 +++--- .../vulkan/test/op_tests/rms_norm_test.cpp | 64 ++++++++++++++++--- backends/vulkan/test/test_vulkan_passes.py | 11 ++++ backends/webgpu/partitioner/partitioner.py | 10 +++ .../webgpu/test/test_webgpu_partitioner.py | 36 ++++++++++- 10 files changed, 154 insertions(+), 31 deletions(-) diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index bc54ae6f66c..fda5f07eec4 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -1232,17 +1232,19 @@ def swiglu_impl(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: def rms_norm_impl( x: torch.Tensor, - weight: torch.Tensor, + weight: Optional[torch.Tensor], eps: float, ) -> torch.Tensor: input_dtype = x.dtype variance = x.float().pow(2).mean(-1, keepdim=True) x_normed = x.float() * torch.rsqrt(variance + eps) - return (x_normed * weight.float()).to(input_dtype) + if weight is not None: + x_normed = x_normed * weight.float() + return x_normed.to(input_dtype) name = "rms_norm" -lib.define(f"{name}(Tensor x, Tensor weight, float eps) -> Tensor") +lib.define(f"{name}(Tensor x, Tensor? weight, float eps) -> Tensor") lib.impl(name, rms_norm_impl, "CompositeExplicitAutograd") rms_norm_op = getattr(getattr(torch.ops, namespace), name) diff --git a/backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.glsl index a869cc1188f..d911e9ce209 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.glsl @@ -16,6 +16,9 @@ ${define_required_extensions("buffer", DTYPE)} ${define_active_storage_type("buffer")} +$if HAS_WEIGHT: + #define HAS_WEIGHT + #extension GL_EXT_control_flow_attributes : require layout(std430) buffer; @@ -24,7 +27,8 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} ${layout_declare_tensor(B, "r", "t_in", DTYPE, "buffer")} -${layout_declare_tensor(B, "r", "t_weight", DTYPE, "buffer")} +$if HAS_WEIGHT: + ${layout_declare_tensor(B, "r", "t_weight", DTYPE, "buffer")} ${layout_declare_ubo(B, "BufferMetadata", "outp")} ${layout_declare_ubo(B, "BufferMetadata", "inp")} @@ -99,8 +103,10 @@ void main() { // Phase 2: Normalize and write output for (uint x = worker_id; x < row_width; x += NUM_WORKERS) { const uint in_bufi = base_bufi + x * width_stride; - const float val = float(t_in[in_bufi]); - const float w = float(t_weight[x]); - t_out[in_bufi] = T(val * rstd * w); + float out_val = float(t_in[in_bufi]) * rstd; +#ifdef HAS_WEIGHT + out_val *= float(t_weight[x]); +#endif + t_out[in_bufi] = T(out_val); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.yaml index 56358f406fd..7db12cc29be 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.yaml @@ -8,9 +8,12 @@ rms_norm_buffer: parameter_names_with_default_values: DTYPE: float STORAGE: buffer + HAS_WEIGHT: false generate_variant_forall: DTYPE: - VALUE: half - VALUE: float shader_variants: - NAME: rms_norm_buffer + HAS_WEIGHT: true + - NAME: rms_norm_no_weight_buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/rms_norm_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/rms_norm_texture.glsl index 838736919e2..115fdb9c83f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/rms_norm_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/rms_norm_texture.glsl @@ -17,6 +17,9 @@ ${define_required_extensions("texture3d", DTYPE)} ${define_active_storage_type("texture3d")} +$if HAS_WEIGHT: + #define HAS_WEIGHT + #extension GL_EXT_control_flow_attributes : require layout(std430) buffer; @@ -26,7 +29,8 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")} ${layout_declare_tensor(B, "r", "t_in", DTYPE, "texture3d")} -${layout_declare_tensor(B, "r", "t_weight", DTYPE, "texture3d")} +$if HAS_WEIGHT: + ${layout_declare_tensor(B, "r", "t_weight", DTYPE, "texture3d")} ${layout_declare_ubo(B, "TextureMetadata", "out_meta")} ${layout_declare_ubo(B, "TextureMetadata", "in_meta")} @@ -105,9 +109,10 @@ void main() { ivec3 pos = ivec3(tx, y, z); const VEC4_T in_val = texelFetch(t_in, pos, 0); - const VEC4_T weight = texelFetch(t_weight, ivec3(tx, 0, 0), 0); - const VEC4_T outtex = - VEC4_T(vec4(in_val) * rstd * vec4(weight)); - imageStore(t_out, pos, outtex); + vec4 outtex = vec4(in_val) * rstd; +#ifdef HAS_WEIGHT + outtex *= vec4(texelFetch(t_weight, ivec3(tx, 0, 0), 0)); +#endif + imageStore(t_out, pos, VEC4_T(outtex)); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/rms_norm_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/rms_norm_texture.yaml index 3cd2c2d9436..78f4e866a33 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/rms_norm_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/rms_norm_texture.yaml @@ -7,9 +7,12 @@ rms_norm_texture: parameter_names_with_default_values: DTYPE: float + HAS_WEIGHT: false generate_variant_forall: DTYPE: - VALUE: half - VALUE: float shader_variants: - NAME: rms_norm_texture3d + HAS_WEIGHT: true + - NAME: rms_norm_no_weight_texture3d diff --git a/backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp index c152b3d408c..ecb843bd35b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp @@ -51,17 +51,20 @@ void add_rms_norm_node( const ValueRef eps, const ValueRef out) { const bool is_buffer = graph.is_buffer_storage(in); - - // The buffer shader indexes the weight as a dense array. - ValueRef arg_weight = prepack_standard( - graph, - weight_data, - graph.storage_type_of(in), - is_buffer ? utils::kWidthPacked : graph.estimate_memory_layout_of(in)); + const bool has_weight = graph.val_is_not_none(weight_data); + std::vector read_inputs = {in}; + if (has_weight) { + // The buffer shader indexes the weight as a dense array. + read_inputs.push_back(prepack_standard( + graph, + weight_data, + graph.storage_type_of(in), + is_buffer ? utils::kWidthPacked : graph.estimate_memory_layout_of(in))); + } float epsilon = graph.extract_scalar(eps); - std::string kernel_name("rms_norm"); + std::string kernel_name(has_weight ? "rms_norm" : "rms_norm_no_weight"); kernel_name.reserve(kShaderNameReserve); add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(out)); @@ -82,7 +85,7 @@ void add_rms_norm_node( rms_norm_gwg, pick_required_lwg, // Inputs and Outputs - {{out, vkapi::kWrite}, {{in, arg_weight}, vkapi::kRead}}, + {{out, vkapi::kWrite}, {read_inputs, vkapi::kRead}}, // Shader params buffers param_ubos, // Push Constants diff --git a/backends/vulkan/test/op_tests/rms_norm_test.cpp b/backends/vulkan/test/op_tests/rms_norm_test.cpp index 057eda4445f..e0a4ca2b062 100644 --- a/backends/vulkan/test/op_tests/rms_norm_test.cpp +++ b/backends/vulkan/test/op_tests/rms_norm_test.cpp @@ -83,27 +83,32 @@ void test_rms_norm( const float eps = 1e-5f, const vkcompute::vkapi::ScalarType dtype = vkcompute::vkapi::kFloat, const vkcompute::utils::StorageType storage_type = - vkcompute::utils::kTexture3D) { + vkcompute::utils::kTexture3D, + const bool has_weight = true) { const int64_t hidden_size = input_shape.back(); const size_t input_numel = numel(input_shape); const size_t weight_numel = static_cast(hidden_size); std::vector x_data = rand_floats(input_numel, 42); - std::vector w_data = rand_floats(weight_numel, 123); + // The reference always multiplies by a weight, so no weight means ones. + std::vector w_data = has_weight ? rand_floats(weight_numel, 123) + : std::vector(weight_numel, 1); // For fp16: round-trip through Half so the reference uses the same precision // as the GPU input. std::vector x_half, w_half; if (dtype == vkcompute::vkapi::kHalf) { x_half.resize(input_numel); - w_half.resize(weight_numel); for (size_t i = 0; i < input_numel; ++i) { x_half[i] = static_cast(x_data[i]); x_data[i] = static_cast(x_half[i]); } - for (size_t i = 0; i < weight_numel; ++i) { - w_half[i] = static_cast(w_data[i]); - w_data[i] = static_cast(w_half[i]); + if (has_weight) { + w_half.resize(weight_numel); + for (size_t i = 0; i < weight_numel; ++i) { + w_half[i] = static_cast(w_data[i]); + w_data[i] = static_cast(w_half[i]); + } } } @@ -118,9 +123,14 @@ void test_rms_norm( IOValueRef r_x = graph.add_input_tensor(input_shape, dtype); - ValueRef r_weight = (dtype == vkapi::kHalf) - ? graph.add_tensorref({hidden_size}, vkapi::kHalf, w_half.data()) - : graph.add_tensorref({hidden_size}, vkapi::kFloat, w_data.data()); + ValueRef r_weight; + if (!has_weight) { + r_weight = graph.add_none(); + } else if (dtype == vkapi::kHalf) { + r_weight = graph.add_tensorref({hidden_size}, vkapi::kHalf, w_half.data()); + } else { + r_weight = graph.add_tensorref({hidden_size}, vkapi::kFloat, w_data.data()); + } const ValueRef r_eps = graph.add_scalar(static_cast(eps)); const ValueRef r_out = graph.add_tensor(input_shape, dtype); @@ -197,6 +207,24 @@ TEST(VulkanRmsNormTest, 4d_multi_z_slice_texture) { test_rms_norm({1, 5, 4, 128}); } +TEST(VulkanRmsNormTest, no_weight_texture) { + test_rms_norm( + {1, 5, 4, 128}, + 1e-5f, + vkcompute::vkapi::kFloat, + vkcompute::utils::kTexture3D, + /*has_weight=*/false); +} + +TEST(VulkanRmsNormTest, no_weight_fp16_texture) { + test_rms_norm( + {1, 1, 1, 896}, + 1e-6f, + vkcompute::vkapi::kHalf, + vkcompute::utils::kTexture3D, + /*has_weight=*/false); +} + // // Buffer storage tests // @@ -217,6 +245,24 @@ TEST(VulkanRmsNormTest, fp16_buffer) { vkcompute::utils::kBuffer); } +TEST(VulkanRmsNormTest, no_weight_buffer) { + test_rms_norm( + {1, 5, 4, 128}, + 1e-5f, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + /*has_weight=*/false); +} + +TEST(VulkanRmsNormTest, no_weight_fp16_buffer) { + test_rms_norm( + {1, 1, 7, 896}, + 1e-6f, + vkcompute::vkapi::kHalf, + vkcompute::utils::kBuffer, + /*has_weight=*/false); +} + // // Dynamic resize test // diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index 11e8872a8d2..39231dcc51e 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -794,6 +794,17 @@ def test_embedding_q4gsw_reference_rejects_invalid_scale_shape(self): True, ) + def test_rms_norm_reference_matches_aten(self): + import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + + x = torch.arange(48, dtype=torch.float32).reshape(2, 3, 8).sub(24) + for weight in (torch.linspace(-1, 1, 8), None): + with self.subTest(weight=weight is not None): + torch.testing.assert_close( + torch.ops.et_vk.rms_norm.default(x, weight, 1e-5), + torch.nn.functional.rms_norm(x, (8,), weight, 1e-5), + ) + def test_fuse_quantized_embedding_with_per_row_scales(self): import executorch.backends.vulkan.custom_ops_lib # noqa: F401 diff --git a/backends/webgpu/partitioner/partitioner.py b/backends/webgpu/partitioner/partitioner.py index 30aa2d1d630..69314fde64f 100644 --- a/backends/webgpu/partitioner/partitioner.py +++ b/backends/webgpu/partitioner/partitioner.py @@ -89,6 +89,13 @@ def _convolution_check( return vulkan_check(node) +def _rms_norm_check( + node: torch.fx.Node, vulkan_check: Callable[[torch.fx.Node], bool] +) -> bool: + # The WebGPU kernel always binds a weight buffer. + return node.args[1] is not None and vulkan_check(node) + + def _fp32_unary_features() -> OpFeatures: return OpFeatures( inputs_storage=vk_utils.ANY_STORAGE, @@ -115,6 +122,9 @@ def _fp32_unary_features() -> OpFeatures: _convolution_check, dtypes={torch.float32}, ), + exir_ops.edge.et_vk.rms_norm.default: _webgpu_features( + exir_ops.edge.et_vk.rms_norm.default, _rms_norm_check + ), exir_ops.edge.aten.elu.default: _fp32_unary_features(), exir_ops.edge.aten.erf.default: _fp32_unary_features(), exir_ops.edge.aten.silu.default: _fp32_unary_features(), diff --git a/backends/webgpu/test/test_webgpu_partitioner.py b/backends/webgpu/test/test_webgpu_partitioner.py index 50d453fb98b..d29c9bf4201 100644 --- a/backends/webgpu/test/test_webgpu_partitioner.py +++ b/backends/webgpu/test/test_webgpu_partitioner.py @@ -5,7 +5,7 @@ # LICENSE file in the root directory of this source tree. import unittest -from typing import List, Tuple +from typing import List, Optional, Tuple import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner @@ -70,6 +70,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x + torch.arange(0, x.shape[0], 1, dtype=torch.float32) +class RmsNormModule(torch.nn.Module): + def __init__(self, affine: bool) -> None: + super().__init__() + self.weight: Optional[torch.nn.Parameter] = ( + torch.nn.Parameter(torch.ones(8)) if affine else None + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.et_vk.rms_norm(x, self.weight, 1e-5) + + class TestWebGPUPartitioner(unittest.TestCase): def _export_add(self) -> ExportedProgram: return torch.export.export( @@ -488,6 +499,29 @@ def test_transposed_conv_stays_on_cpu(self) -> None: self._targets(non_lowered), ) + def _rms_norm_is_delegated(self, partitioner: Partitioner, affine: bool) -> bool: + exported_program = torch.export.export( + RmsNormModule(affine), (torch.rand(2, 8),) + ) + result = partitioner.partition(to_edge(exported_program).exported_program()) + (node,) = [ + node + for node in result.tagged_exported_program.graph.nodes + if node.target == exir_ops.edge.et_vk.rms_norm.default + ] + return "delegation_tag" in node.meta + + def test_rms_norm_without_weight_stays_on_cpu(self) -> None: + # Vulkan implements a None weight; the WebGPU kernel does not. + for affine in (True, False): + with self.subTest(affine=affine): + self.assertEqual( + self._rms_norm_is_delegated(WebGPUPartitioner(), affine), affine + ) + self.assertTrue( + self._rms_norm_is_delegated(VulkanPartitioner(), affine) + ) + if __name__ == "__main__": unittest.main()