[ET-VK] RMS-norm fusion folds a non-leaf multiplier and then prepacks it, breaking adaptive norms

[giuliocorradi]

🐛 Describe the bug

Summary

The Vulkan AOT fuses x * rsqrt(mean(x²) + eps) into et_vk.rms_norm and folds
the multiply that follows the norm in as that norm's weight — without
checking that the multiplier is a constant it can prepack. When it is not, the
delegate aborts at the first inference:

prepack_standard at backends/vulkan/runtime/graph/ops/impl/Staging.cpp:229:
  (graph.val_is_tref(tensor_data)) is false!

This makes every adaptive normalisation unlowerable, and also breaks Gemma's
ordinary RMSNorm, whose scale is written 1.0 + weight.

Reproduction

No model needed. Three cases; the only difference is what the norm is multiplied
by.

import torch
from executorch.exir import to_edge
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.extension.pybindings.portable_lib import _load_for_executorch

D, EPS = 64, 1e-6

def norm(x):
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS)

class TimesParam(torch.nn.Module):            # OK
    def __init__(self):
        super().__init__()
        self.w = torch.nn.Parameter(torch.ones(D))
    def forward(self, x):
        return norm(x) * self.w

class TimesComputedConst(torch.nn.Module):    # ASSERTS  (Gemma's RMSNorm)
    def __init__(self):
        super().__init__()
        self.w = torch.nn.Parameter(torch.zeros(D))
    def forward(self, x):
        return norm(x) * (1.0 + self.w)

class TimesRuntime(torch.nn.Module):          # ASSERTS  (adaptive RMS norm)
    def forward(self, x, scale):
        return norm(x) * (1 + scale)

def run(mod, args, tag):
    low = to_edge(torch.export.export(mod, args)).to_backend(VulkanPartitioner())
    path = f"/tmp/{tag}.pte"
    with open(path, "wb") as f:
        low.to_executorch().write_to_file(f)
    out = _load_for_executorch(path).forward(args)[0]
    print(tag, "ok, max|d| =", (out - mod(*args)).abs().max().item())

x = torch.randn(1, 8, D)
run(TimesParam(), (x,), "times_param")
run(TimesComputedConst(), (x,), "times_computed_const")
run(TimesRuntime(), (x, torch.rand(1, 1, D) * 0.1), "times_runtime")

Result

case	multiplier	outcome
norm(x) * self.w	a leaf parameter	runs, max|d| 2.4e-07
norm(x) * (1.0 + w)	constant, but an intermediate	asserts in prepack_standard
norm(x) * scale	computed at inference	asserts in prepack_standard

So it is not "constant vs non-constant" that decides it — it is whether the
multiplier is a leaf the prepacker can see. 1.0 + w is constant-valued and
still fails.

Expected behaviour

The fusion should fold the trailing multiply into et_vk.rms_norm only when the
multiplier is prepackable, and otherwise leave it as a separate elementwise
multiply. Both graphs are legal; only one is fusable.

Suggested fix

In the rms-norm pattern (backends/vulkan/patterns/rms_norm.py), guard the fold
on the multiplier being a constant/leaf — the same val_is_tref-style
predicate the runtime later asserts on. Falling back to an unfused norm plus a
separate mul is correct and costs one extra dispatch.

Impact

Adaptive normalisation — a norm whose scale and shift are produced at inference
from a conditioning signal — appears in:

• π₀.₅ / openpi (PiGemmaRMSNorm, scale derived from the flow-matching
  timestep)
• DiT and most diffusion transformers (AdaLN, AdaLN-Zero)
• any Gemma-family model, via the ordinary 1.0 + weight form

None of these can be lowered to the Vulkan delegate without rewriting the model.

Note on the obvious workaround

Rewriting the norm as F.rms_norm(x, ones) avoids the assert but stops the
fusion firing, and the + eps then survives as an aten.add.Scalar that the
partitioner leaves on the CPU — one graph break per norm, 35 of them in the
model that prompted this report. A workaround has to keep the fusion working,
not merely stop it crashing, which is why the guard belongs in the pattern.

Versions

collect_env.py output

Collecting environment information...
PyTorch version: 2.12.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCm SDK used to build PyTorch: N/A
HIP used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.39

Python version: 3.12.3 (main, Jul 15 2026, 23:46:41) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-7.0.0-30-generic-x86_64-with-glibc2.39
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A
CPU: AMD RYZEN AI MAX+ 395 w/ Radeon 8060S, 16 cores / 32 threads (full lscpu elided)

Versions of relevant libraries:
[pip3] executorch==1.4.0a0+b20f16a
[pip3] numpy==2.4.6
[pip3] pytorch_tokenizers==1.4.1
[pip3] torch==2.12.1+cpu
[pip3] torchao==0.18.0.dev20260715+cpu
[conda] Could not collect

cc @SS-JIA @manuelcandales @digantdesai @cbilgin