[ET-VK] Synthesize the layer_norm affine parameters that are absent

backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp

@@ -15,6 +15,7 @@
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+#include <executorch/runtime/core/portable_type/half.h>
 
 namespace vkcompute {
 
@@ -50,7 +51,7 @@ void resize_native_layer_norm_node(
   graph->virtual_resize(rstd, mean_size);
 }
 
-GlobalWorkGrid layer_norm_buffer_gwg(
+utils::uvec3 layer_norm_buffer_global_wg_size(
     ComputeGraph* graph,
     const vkapi::ShaderInfo& shader,
     const std::vector<ArgGroup>& args,
@@ -60,8 +61,65 @@ GlobalWorkGrid layer_norm_buffer_gwg(
   const ValueRef mean_tensor = args.at(0).refs.at(1);
   const uint32_t num_rows =
       utils::safe_downcast<uint32_t>(graph->numel_of(mean_tensor));
-  return GlobalWorkGrid(
-      {1u, num_rows, 1u}, kTiledWorkGrid, LocalWorkGroup(64u, 1u, 1u));
+  return {1u, num_rows, 1u};
+}
+
+utils::uvec3 layer_norm_buffer_local_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const utils::uvec3& global_workgroup_size,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)graph;
+  (void)shader;
+  (void)global_workgroup_size;
+  (void)args;
+  (void)resize_args;
+  return {64u, 1u, 1u};
+}
+
+// Builds a TensorRef of `sizes` filled with `value`, for synthesizing a
+// layer_norm affine parameter the caller did not supply. Ownership of the
+// buffer passes to the FreeableBuffer, then to the TensorRef, then to the
+// graph.
+namespace etensor = executorch::runtime::etensor;
+
+ValueRef constant_affine_tensorref(
+    ComputeGraph& graph,
+    const std::vector<int64_t>& sizes,
+    const vkapi::ScalarType dtype,
+    const float value) {
+  int64_t numel = 1;
+  for (int64_t d : sizes) {
+    numel *= d;
+  }
+  const size_t total_bytes =
+      static_cast<size_t>(numel) * vkapi::element_size(dtype);
+  auto* data = new uint8_t[total_bytes]();
+  if (value != 0.0f) {
+    switch (dtype) {
+      case vkapi::kFloat: {
+        auto* typed = reinterpret_cast<float*>(data);
+        std::fill(typed, typed + numel, value);
+        break;
+      }
+      case vkapi::kHalf: {
+        auto* typed = reinterpret_cast<etensor::Half*>(data);
+        std::fill(typed, typed + numel, etensor::Half(value));
+        break;
+      }
+      default:
+        delete[] data;
+        VK_THROW(
+            "native_layer_norm cannot synthesize an affine parameter of dtype ",
+            static_cast<int>(dtype));
+    }
+  }
+  executorch::runtime::FreeableBuffer buffer(
+      data, total_bytes, [](void* /*ctx*/, void* ptr, size_t /*size*/) {
+        delete[] static_cast<uint8_t*>(ptr);
+      });
+  return graph.add_tensorref(sizes, dtype, std::move(buffer));
 }
 
 void add_native_layer_norm_node(
@@ -78,16 +136,33 @@ void add_native_layer_norm_node(
     VK_THROW("native_layer_norm only supports normalized_shape with dim == 1");
   }
 
+  // The shader reads a weight and a bias binding unconditionally, but either
+  // affine parameter can legitimately be absent:
+  //   - nn.LayerNorm(bias=False) (e.g. HF Gemma4 audio_encoder) has no bias.
+  //   - F.layer_norm called with neither, where the caller applies its own
+  //     scale and shift afterwards (e.g. kokoro's AdaLayerNorm), has neither.
+  // Synthesize whatever is missing: a unit weight and a zero bias reproduce
+  // out = (x - mean) * rstd exactly.
+  const std::vector<int64_t> affine_sizes = graph.val_is_none(weight_data)
+      ? std::vector<int64_t>{*graph.get_int_list(normalized_shape)}
+      : graph.sizes_of(weight_data);
+  const vkapi::ScalarType affine_dtype = graph.val_is_none(weight_data)
+      ? graph.dtype_of(in)
+      : graph.dtype_of(weight_data);
+
+  ValueRef synthesized_weight = weight_data;
   if (graph.val_is_none(weight_data)) {
-    VK_THROW("native_layer_norm requires weight to be non-None");
+    synthesized_weight =
+        constant_affine_tensorref(graph, affine_sizes, affine_dtype, 1.0f);
   }
-
+  ValueRef synthesized_bias = bias_data;
   if (graph.val_is_none(bias_data)) {
-    VK_THROW("native_layer_norm requires bias to be non-None");
+    synthesized_bias =
+        constant_affine_tensorref(graph, affine_sizes, affine_dtype, 0.0f);
   }
 
-  ValueRef arg_weight = prepack_standard_like(graph, weight_data, in);
-  ValueRef arg_bias = prepack_standard_like(graph, bias_data, in);
+  ValueRef arg_weight = prepack_standard_like(graph, synthesized_weight, in);
+  ValueRef arg_bias = prepack_standard_like(graph, synthesized_bias, in);
 
   const auto out_val = graph.get_value_list(out);
   const ValueRef out_tensor = out_val->at(0);
@@ -118,8 +193,9 @@ void add_native_layer_norm_node(
   graph.execute_nodes().emplace_back(new DynamicDispatchNode(
       graph,
       VK_KERNEL_FROM_STR(kernel_name),
-      is_buffer ? layer_norm_buffer_gwg : default_pick_gwg,
-      is_buffer ? pick_required_lwg : default_pick_lwg,
+      is_buffer ? layer_norm_buffer_global_wg_size
+                : default_pick_global_wg_size,
+      is_buffer ? layer_norm_buffer_local_wg_size : default_pick_local_wg_size,
       // Inputs and Outputs
       {{{out_tensor, mean_tensor, rstd_tensor}, vkapi::kWrite},
        {{in, arg_weight, arg_bias}, vkapi::kRead}},

backends/vulkan/test/op_tests/cases.py

@@ -761,6 +761,13 @@ def get_native_layer_norm_inputs():
             ((S1, S2), [S2], (S2), (S2), 0.001),
             ((M, M1, M2), [M2], (M2), (M2), 0.001),
             ((S, XL, M1, M2), [M2], (M2), (M2), 0.001),
+            # Either affine parameter may be absent: nn.LayerNorm(bias=False)
+            # has no bias, and F.layer_norm called with neither (kokoro's
+            # AdaLayerNorm applies its own scale and shift afterwards) has
+            # neither.
+            ((M, M1, M2), [M2], (M2), None, 0.001),
+            ((M, M1, M2), [M2], None, (M2), 0.001),
+            ((M, M1, M2), [M2], None, None, 0.001),
         ]
     )
     test_suite.layouts = [