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}", + )