From 8da889e537d417cb72e0e97a97833ccef6180f65 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Sat, 14 Nov 2026 11:29:52 -0800 Subject: [PATCH] Update [ghstack-poisoned] --- .../runtime/vulkan/passes/LowerRMSNorm.cpp | 38 ++++++++---- .../runtime/vulkan/test_lower_rms_norm.cpp | 58 +++++++++++++++++-- .../runtime/vulkan/test_vulkan_engine.cpp | 50 +++++++++++++--- 3 files changed, 120 insertions(+), 26 deletions(-) diff --git a/backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp b/backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp index 59ad98974cd..a0455edc81e 100644 --- a/backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp +++ b/backends/native/runtime/vulkan/passes/LowerRMSNorm.cpp @@ -28,6 +28,24 @@ bool is_constant(const Value& value) { value.role == ValueRole::ConstantTensor; } +// A None weight scales by one; a present weight must be a constant that matches +// the normalized shape and the input dtype. +bool is_valid_weight( + const Graph& graph, + const Argument& weight, + const std::vector& shape, + ScalarType dtype) { + if (weight.kind() == ArgKind::None) { + return true; + } + if (weight.kind() != ArgKind::Tensor || !valid(weight.as_tensor().id)) { + return false; + } + const Value& value = graph.value(weight.as_tensor().id); + return value.is_tensor() && is_constant(value) && + value.tensor_meta().sizes == shape && value.tensor_meta().dtype == dtype; +} + } // namespace size_t lower_rms_norms(Graph& graph) { @@ -46,37 +64,33 @@ size_t lower_rms_norms(Graph& graph) { if (input == nullptr || input->kind() != ArgKind::Tensor || normalized_shape == nullptr || normalized_shape->kind() != ArgKind::IntList || weight == nullptr || - weight->kind() != ArgKind::Tensor || eps == nullptr || - eps->kind() != ArgKind::Float || valid(eps->as_float().id) || - !std::isfinite(eps->as_float().value) || eps->as_float().value < 0) { + eps == nullptr || eps->kind() != ArgKind::Float || + valid(eps->as_float().id) || !std::isfinite(eps->as_float().value) || + eps->as_float().value < 0) { continue; } const ValueId input_id = input->as_tensor().id; - const ValueId weight_id = weight->as_tensor().id; const ValueId output_id = node.outputs[0].value_id; - if (!valid(input_id) || !valid(weight_id)) { + if (!valid(input_id)) { continue; } const Value& input_value = graph.value(input_id); - const Value& weight_value = graph.value(weight_id); const Value& output_value = graph.value(output_id); - if (!input_value.is_tensor() || !weight_value.is_tensor() || - !output_value.is_tensor() || !is_constant(weight_value)) { + if (!input_value.is_tensor() || !output_value.is_tensor()) { continue; } const TensorMeta& input_meta = input_value.tensor_meta(); - const TensorMeta& weight_meta = weight_value.tensor_meta(); const TensorMeta& output_meta = output_value.tensor_meta(); const IntListArg& shape = normalized_shape->as_int_list(); - if (input_meta.sizes.empty() || !shape.ids.empty() || + if (input_meta.sizes.empty() || input_meta.sizes.back() < 1 || + !shape.ids.empty() || shape.values != std::vector{input_meta.sizes.back()} || - weight_meta.sizes != shape.values || (input_meta.dtype != ScalarType::Float && input_meta.dtype != ScalarType::Half && input_meta.dtype != ScalarType::BFloat16) || - weight_meta.dtype != input_meta.dtype || + !is_valid_weight(graph, *weight, shape.values, input_meta.dtype) || output_meta.dtype != input_meta.dtype || output_meta.sizes != input_meta.sizes) { continue; diff --git a/backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp b/backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp index 73279663311..db8e71ccc61 100644 --- a/backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp +++ b/backends/native/test/runtime/vulkan/test_lower_rms_norm.cpp @@ -22,14 +22,14 @@ namespace ptn::vulkan { namespace { -Method make_rms_norm(ScalarType dtype = ScalarType::Float) { +Method make_rms_norm(ScalarType dtype = ScalarType::Float, int64_t width = 4) { Method method; method.name = "forward"; Graph& graph = method.graph; graph.values = { - Value("input", dtype, {2, 4}), - Value("weight", dtype, {4}), - Value("output", dtype, {2, 4}), + Value("input", dtype, {2, width}), + Value("weight", dtype, {width}), + Value("output", dtype, {2, width}), }; graph.value(1).role = ValueRole::Parameter; graph.nodes = { @@ -49,7 +49,7 @@ Method make_rms_norm(ScalarType dtype = ScalarType::Float) { .inputs = { {.name = "input", .arg = TensorArg{0}}, - {.name = "normalized_shape", .arg = IntListArg{{4}}}, + {.name = "normalized_shape", .arg = IntListArg{{width}}}, {.name = "weight", .arg = TensorArg{1}}, {.name = "eps", .arg = FloatArg{1e-5}}, }, @@ -104,6 +104,20 @@ TEST(LowerRMSNormTest, RewritesBFloat16Norm) { EXPECT_NO_THROW(validate_graph(method.graph)); } +TEST(LowerRMSNormTest, RewritesNormWithoutWeight) { + Method method = make_rms_norm(); + method.graph.node(2).inputs[2].arg = NoneArg{}; + method.graph.rebuild_def_use(); + + EXPECT_EQ(lower_rms_norms(method.graph), 1); + + const Node& norm = method.graph.node(2); + EXPECT_EQ(norm.target, "torch.ops.et_vk.rms_norm.default"); + ASSERT_EQ(norm.inputs.size(), 3); + EXPECT_EQ(norm.inputs[1].arg.kind(), ArgKind::None); + EXPECT_NO_THROW(validate_graph(method.graph)); +} + TEST(LowerRMSNormTest, RejectsMismatchedNormalizedShape) { Method method = make_rms_norm(); method.graph.node(2).inputs[1].arg = IntListArg{{8}}; @@ -112,5 +126,39 @@ TEST(LowerRMSNormTest, RejectsMismatchedNormalizedShape) { EXPECT_EQ(method.graph.node(2).target, "torch.ops.aten.rms_norm.default"); } +TEST(LowerRMSNormTest, RejectsInvalidWeight) { + struct Case { + const char* name; + ValueRole role; + ScalarType dtype; + std::vector sizes; + }; + const std::vector cases = { + {"intermediate", ValueRole::Intermediate, ScalarType::Float, {4}}, + {"user input", ValueRole::UserInput, ScalarType::Float, {4}}, + {"wrong shape", ValueRole::Parameter, ScalarType::Float, {8}}, + {"wrong dtype", ValueRole::Parameter, ScalarType::Half, {4}}, + }; + for (const Case& c : cases) { + SCOPED_TRACE(c.name); + Method method = make_rms_norm(); + method.graph.value(1) = Value("weight", c.dtype, c.sizes); + method.graph.value(1).role = c.role; + method.graph.rebuild_def_use(); + + EXPECT_EQ(lower_rms_norms(method.graph), 0); + EXPECT_EQ(method.graph.node(2).target, "torch.ops.aten.rms_norm.default"); + } +} + +TEST(LowerRMSNormTest, RejectsZeroWidthNorm) { + Method method = make_rms_norm(ScalarType::Float, /*width=*/0); + method.graph.node(2).inputs[2].arg = NoneArg{}; + method.graph.rebuild_def_use(); + + EXPECT_EQ(lower_rms_norms(method.graph), 0); + EXPECT_EQ(method.graph.node(2).target, "torch.ops.aten.rms_norm.default"); +} + } // namespace } // namespace ptn::vulkan diff --git a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp index 5ef3fe32fb9..be4a961e0c4 100644 --- a/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp +++ b/backends/native/test/runtime/vulkan/test_vulkan_engine.cpp @@ -2833,12 +2833,15 @@ std::vector stored_floats(const std::vector& constants) { return tensors; } -// aten.rms_norm over the last dim of a [2, 4] input with a stored weight. -std::vector make_rms_norm_program(fbs::ScalarType dtype) { +// aten.rms_norm over the last dim of a [2, 4] input, with a stored weight if +// `affine` and a None weight otherwise. +std::vector make_rms_norm_program( + fbs::ScalarType dtype, + bool affine = true) { flatbuffers::FlatBufferBuilder builder; const auto meta = create_tensor_meta(builder, dtype, {2, 4}); const auto weight_meta = create_tensor_meta(builder, dtype, {4}); - const std::vector> tensor_values = { + std::vector> tensor_values = { fbs::CreateTensorValueDirect(builder, "input", meta), fbs::CreateTensorValueDirect(builder, "weight", weight_meta), fbs::CreateTensorValueDirect(builder, "normed", meta), @@ -2856,7 +2859,10 @@ std::vector make_rms_norm_program(fbs::ScalarType dtype) { fbs::CreateNamedArgumentDirect( builder, "normalized_shape", create_int_list_arg(builder, {4})), fbs::CreateNamedArgumentDirect( - builder, "weight", create_tensor_arg(builder, "weight")), + builder, + "weight", + affine ? create_tensor_arg(builder, "weight") + : create_none_arg(builder)), fbs::CreateNamedArgumentDirect( builder, "eps", @@ -2866,7 +2872,7 @@ std::vector make_rms_norm_program(fbs::ScalarType dtype) { const std::vector> output_inputs = { fbs::CreateNamedArgumentDirect( builder, "", create_tensor_arg(builder, "normed"))}; - const std::vector> nodes = { + std::vector> nodes = { fbs::CreateNodeDirect( builder, "input", @@ -2891,16 +2897,22 @@ std::vector make_rms_norm_program(fbs::ScalarType dtype) { fbs::CreateNodeDirect( builder, "output", fbs::OpKind::OUTPUT, "", &output_inputs), }; + std::vector> constants = { + fbs::CreateNamedTensorRefDirect( + builder, "weight", "weight", weight_meta, fbs::InputKind::PARAMETER), + }; + if (!affine) { + // Drop the weight value, its placeholder and its constant. + tensor_values.erase(tensor_values.begin() + 1); + nodes.erase(nodes.begin() + 1); + constants.clear(); + } const auto graph = fbs::CreateGraph( builder, builder.CreateVector(nodes), create_strings(builder, {"input"}), create_strings(builder, {"normed"}), builder.CreateVector(tensor_values)); - const std::vector> constants = { - fbs::CreateNamedTensorRefDirect( - builder, "weight", "weight", weight_meta, fbs::InputKind::PARAMETER), - }; return finish_program( builder, graph, @@ -3587,6 +3599,26 @@ TEST(VulkanEngineTest, ExecutesBFloat16RMSNormWithStoredWeight) { 0x3F00, 0xBFC0, 0xC000, 0xC040, 0x3F00, 0x3FC0, 0x4000, 0x4040})); } +TEST(VulkanEngineTest, ExecutesBFloat16RMSNormWithoutWeight) { + CompiledProgram compiled = compile_program( + make_rms_norm_program(fbs::ScalarType::BFLOAT16, /*affine=*/false)); + // bf16 bits: rows {2, -2, 2, -2} and {4, 4, -4, 4}, each with an RMS that + // normalizes it to +-1. + const std::vector input = { + 0x4000, 0xC000, 0x4000, 0xC000, 0x4080, 0x4080, 0xC080, 0x4080}; + compiled.executable->set_input( + 0, input.data(), input.size(), ScalarType::BFloat16); + compiled.executable->execute(); + std::vector output(8); + compiled.executable->get_output( + 0, output.data(), output.size(), ScalarType::BFloat16); + // {1, -1, 1, -1} and {1, 1, -1, 1}. + EXPECT_EQ( + output, + (std::vector{ + 0x3F80, 0xBF80, 0x3F80, 0xBF80, 0x3F80, 0x3F80, 0xBF80, 0x3F80})); +} + TEST(VulkanEngineTest, ExecutesMetadataOnlyAliasAfterAliasLowering) { CompiledProgram compiled = compile_program(make_alias_program()); const std::vector input = {1.0f, 2.0f, 3.0f, 4.0f};