From ac0d5f3c4d843182e7fd5c7b39b8beac8e87b63c Mon Sep 17 00:00:00 2001 From: Giulio Corradi Date: Sun, 13 Sep 2026 05:30:30 +0200 Subject: [PATCH] Only fuse an RMS norm whose weight can be prepacked et_vk.rms_norm prepacks its weight, so the multiply that the fusion folds in has to be a constant the prepacker can see. The pattern folded in any multiply that followed the norm, so a multiplier computed in the graph produced an op the runtime aborted on at the first inference: prepack_standard ... (graph.val_is_tref(tensor_data)) is false! This affects any adaptive normalization - a norm whose scale is produced at inference from a conditioning signal, as in DiT-style AdaLN - and also Gemma's ordinary RMSNorm, whose scale is written `1.0 + weight` and is therefore constant valued but still an intermediate node. Reject a non-prepackable weight in two places, because the two callers see different information. The detector is given only the graph, so it can require the weight to be a placeholder, which covers every computed multiplier. Only the replacement is given the exported program, so the distinction between a constant placeholder and a user input is made there. Not fusing is correct: the norm and the multiply are both supported ops, so the pattern simply stays unfused and costs one extra dispatch. Test Plan: python -m unittest backends.vulkan.test.test_vulkan_passes -v 12 tests pass. The new test, test_rms_norm_fuses_only_prepackable_weight, fails without this change with "1 != 0 : expected 0 fused rms_norm" for both the computed-constant and the graph-input case. Verified end to end on an AMD Radeon 8060S (RADV GFX1151), lowering each module and comparing the delegate's output against eager: multiplier before after leaf parameter fused, ok fused, ok leaf buffer fused, ok fused, ok parameter through a dtype cast prepack abort unfused, ok 1.0 + parameter (Gemma) prepack abort unfused, ok graph input (adaptive norm) prepack abort unfused, ok computed by a layer (AdaLN) prepack abort unfused, ok backends.vulkan.test.test_vulkan_delegate shows an identical set of results before and after the change. --- backends/vulkan/patterns/rms_norm.py | 22 +++++++ backends/vulkan/test/test_vulkan_passes.py | 71 ++++++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/backends/vulkan/patterns/rms_norm.py b/backends/vulkan/patterns/rms_norm.py index beb5e677ead..f06f2d8a6cb 100644 --- a/backends/vulkan/patterns/rms_norm.py +++ b/backends/vulkan/patterns/rms_norm.py @@ -6,6 +6,8 @@ from typing import Optional +import executorch.backends.vulkan.utils as utils + import torch from executorch.backends.vulkan.patterns.pattern_registry import ( @@ -70,6 +72,19 @@ def __init__(self, final_mul_node: torch.fx.Node) -> None: # noqa: C901 if norm_mul_node is None: return + # et_vk.rms_norm prepacks its weight, so the multiplier has to be a + # constant that the prepacker can see. A multiplier that is computed in + # the graph - an adaptive norm whose scale comes from a conditioning + # signal, or Gemma's `1.0 + weight` - is not prepackable, and folding it + # in anyway makes the delegate abort at the first inference with + # "prepack_standard ... (graph.val_is_tref(tensor_data)) is false". + # Leaving the multiply unfused is correct and costs one dispatch. + if ( + not isinstance(self.weight_node, torch.fx.Node) + or self.weight_node.op != "placeholder" + ): + return + self.all_nodes.append(norm_mul_node) # norm_mul: mul(x_f32, rstd_f32) @@ -263,6 +278,13 @@ def replace_rms_norm_with_fused_op( graph_module: torch.fx.GraphModule, match: RmsNormMatch, ): + # The detector only sees the graph, which cannot distinguish a constant + # placeholder from a user input; both look the same there. Only a constant + # is actually prepackable, so make the final check here, where the exported + # program is available. + if not utils.is_param_node(ep, match.weight_node): + return + eps_val = _extract_eps_value(match.eps_node) with graph_module.graph.inserting_before(match.anchor_node): diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index f030b9268a1..4ac9d3ad69b 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -779,3 +779,74 @@ def forward(self, x): gm = ep.graph_module self.assertEqual(op_node_count(gm, "q8ta_pixel_shuffle.default"), 0) + + def test_rms_norm_fuses_only_prepackable_weight(self): + """et_vk.rms_norm prepacks its weight, so the fusion must only fold in a + multiplier that is an actual constant. + + Folding in a computed multiplier - an adaptive norm whose scale comes + from a conditioning signal, or Gemma's `1.0 + weight` - produced an op + the runtime aborted on: + + prepack_standard ... (graph.val_is_tref(tensor_data)) is false! + + Leaving those unfused is correct; the norm and the multiply are both + supported on their own. + """ + eps = 1e-6 + dim = 64 + + def norm(x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) + + class TimesParameter(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = torch.nn.Parameter(torch.rand(dim) + 0.5) + + def forward(self, x): + return norm(x) * self.w + + class TimesComputedConstant(torch.nn.Module): + """Gemma's RMSNorm: constant valued, but an intermediate node.""" + + def __init__(self): + super().__init__() + self.w = torch.nn.Parameter(torch.rand(dim) * 0.1) + + def forward(self, x): + return norm(x) * (1.0 + self.w) + + class TimesGraphInput(torch.nn.Module): + """Adaptive norm: the scale is a runtime value.""" + + def forward(self, x, scale): + return norm(x) * scale + + x = torch.randn(1, 8, dim) + + for model, inputs, expected, why in [ + (TimesParameter(), (x,), 1, "a leaf parameter is prepackable"), + (TimesComputedConstant(), (x,), 0, "1.0 + w is an intermediate node"), + ( + TimesGraphInput(), + (x, torch.rand(1, 1, dim)), + 0, + "a graph input is not a constant", + ), + ]: + with self.subTest(model=type(model).__name__): + edge_program = to_edge( + torch.export.export(model.eval(), inputs, strict=True), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + ep = edge_program._edge_programs["forward"] + fuse_pass = FusePatternsPass() + fuse_pass._exported_program = ep + fuse_pass.call(ep.graph_module) + + self.assertEqual( + op_node_count(ep.graph_module, "rms_norm.default"), + expected, + f"expected {expected} fused rms_norm: {why}", + )