From 11519e85a80fe7e5fd4b4c4d94539cfff7a81c2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Thu, 3 Sep 2026 08:49:55 +0200 Subject: [PATCH] [ET-VK] Synthesize the layer_norm affine parameters that are absent The layer_norm shader reads a weight and a bias binding unconditionally, and add_native_layer_norm_node already synthesized a zero bias for nn.LayerNorm(bias=False). A missing weight was still a hard error, so any model that calls F.layer_norm with no affine parameters at all aborts at prepack with "native_layer_norm requires weight to be non-None". That call is not unusual: kokoro's AdaLayerNorm normalizes with no affine parameters and applies its own style-conditioned scale and shift afterwards. Synthesize a unit weight the same way, sized from normalized_shape when there is no weight tensor to take a shape from. A unit weight and a zero bias reproduce out = (x - mean) * rstd exactly. The two paths now share one helper. --- .../graph/ops/impl/NativeLayerNorm.cpp | 96 +++++++++++++++++-- backends/vulkan/test/op_tests/cases.py | 7 ++ 2 files changed, 93 insertions(+), 10 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp index b1d8965a4d1..fd568f048b6 100644 --- a/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp @@ -15,6 +15,7 @@ #include #include +#include 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& 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(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& args, + const std::vector& 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& 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(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(data); + std::fill(typed, typed + numel, value); + break; + } + case vkapi::kHalf: { + auto* typed = reinterpret_cast(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(dtype)); + } + } + executorch::runtime::FreeableBuffer buffer( + data, total_bytes, [](void* /*ctx*/, void* ptr, size_t /*size*/) { + delete[] static_cast(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 affine_sizes = graph.val_is_none(weight_data) + ? std::vector{*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}}, diff --git a/backends/vulkan/test/op_tests/cases.py b/backends/vulkan/test/op_tests/cases.py index 4225619b330..dc551e8ff5e 100644 --- a/backends/vulkan/test/op_tests/cases.py +++ b/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 = [