[executorch][vulkan] Make the et_vk.rms_norm weight optional

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)
 

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);
   }
 }

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

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));
   }
 }

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

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<ValueRef> 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<float>(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

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<size_t>(hidden_size);
 
   std::vector<float> x_data = rand_floats(input_numel, 42);
-  std::vector<float> w_data = rand_floats(weight_numel, 123);
+  // The reference always multiplies by a weight, so no weight means ones.
+  std::vector<float> w_data = has_weight ? rand_floats(weight_numel, 123)
+                                         : std::vector<float>(weight_numel, 1);
 
   // For fp16: round-trip through Half so the reference uses the same precision
   // as the GPU input.
   std::vector<Half> 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<Half>(x_data[i]);
       x_data[i] = static_cast<float>(x_half[i]);
     }
-    for (size_t i = 0; i < weight_numel; ++i) {
-      w_half[i] = static_cast<Half>(w_data[i]);
-      w_data[i] = static_cast<float>(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<Half>(w_data[i]);
+        w_data[i] = static_cast<float>(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<double>(static_cast<double>(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
 //

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
 

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(),

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()