[Qualcomm] Support native_layer_norm and affine-free LayerNorm in QNN backend (pytorch#18990)

# [Qualcomm] Support native_layer_norm and affine-free LayerNorm in QNN
backend

## Summary

Adds QNN backend support for `aten.native_layer_norm.default` (which is
the decomposed form of `torch.nn.LayerNorm`) and handles models where
weight/bias are not provided (`elementwise_affine=False`).

## Problem

When exporting models with `torch.native_layer_norm` or
`torch.nn.LayerNorm(affine=False)` to the QNN backend, the following
issues occur:

1. **Missing `native_layer_norm` visitor**: The original
`LayerNormVisitor` only targets `aten.layer_norm.default`, but PyTorch
decomposes `torch.nn.LayerNorm` to `aten.native_layer_norm.default`
during export.

2. **None weight/bias**: When `elementwise_affine=False`, the weight and
bias arguments are `None`. QNN x86_64 runtime cannot handle `None`
tensor inputs, causing `AttributeError` when calling `get_parameter()`.

## Solution

### 1. Update visitor target (`op_layer_norm.py`)

Change the visitor target from `aten.layer_norm.default` to
`aten.native_layer_norm.default`:

```python
# Before
target = ["aten.layer_norm.default"]

# After
target = ["aten.native_layer_norm.default"]
```

This is correct because during ExecuTorch export,
`aten.layer_norm.default` is decomposed to
`aten.native_layer_norm.default` **before** the QNN lowering stage.

### 2. Handle None weight/bias (`op_layer_norm.py`)

When weight/bias are `None`, create synthetic tensors:

- Missing weight → `torch.ones(normalized_shapes)` (identity transform)
- Missing bias → `torch.zeros(normalized_shapes)` (no offset)

Create synthetic `fx.Node` objects to register these as QNN static
tensors:

```python
weight_tensor = torch.ones(normalized_shapes, dtype=torch.float32)
weight_node = torch.fx.Node(
    node.graph,
    node.name + "_runtime_weight",
    "call_function",
    exir_ops.edge.aten.tensor.default,
    (),
    {},
)
# Preserve quant_attrs with zero_point=0 for QNN compatibility
```

### 3. Use same annotator for both ops (`htp_rules.py`)

The quantizer annotator registers both `aten.layer_norm.default` and
`aten.native_layer_norm.default` to the same `LayerNorm` class, since
both ops have identical argument schemas:

```python
@register_annotator(
    [torch.ops.aten.layer_norm.default, torch.ops.aten.native_layer_norm.default],
    QnnConstants.OpLayerNorm.op_name,
)
```

### 4. Add None check to `get_parameter()` (`utils.py`)

Guard against `None` nodes to prevent `AttributeError`:

```python
if node is None:
    return None
```

## Files Changed

| File | Changes |
|------|---------|
| `builders/op_layer_norm.py` | Add `native_layer_norm` support + handle
None weight/bias |
| `builders/utils.py` | Add None guard in `get_parameter()` |
| `quantizer/annotators/htp_rules.py` | Register annotator for both ops
|
| `tests/models.py` | Add `NativeLayerNorm` test model |
| `tests/test_qnn_delegate.py` | Add floating-point and quantized tests
|

## Test Plan

Run QNN delegate tests for layer_norm:

```bash

python backends/qualcomm/tests/test_qnn_delegate.py \
    -k "test_qnn_backend_layer_norm or test_qnn_backend_native_layer_norm" \
    --soc_model SM8650 \
    --build_folder build-x86/ \
    --executorch_root . \
    --enable_x86_64
```

Expected: 4 tests pass (2 floating-point, 2 quantized).

## Release Notes

- `Release notes: qualcomm`

## Related Issues

This resolves the issue where FLUX2 transformer export fails with:
- `[QNN Delegate Op Builder]: LayerNorm weight is None, skipping`
- `AttributeError: 'NoneType' object has no attribute 'name'`

Fixes pytorch#18989 
  - Labels: bug, module:qnn

@abhinaykukkadapu

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: KevinUW114514

backends/qualcomm/_passes/annotate_quant_attrs.py

@@ -8,7 +8,7 @@
 
 import torch
 from executorch.backends.qualcomm.builders.node_visitor import dq_ops, q_ops
-from executorch.backends.qualcomm.builders.utils import get_parameter
+from executorch.backends.qualcomm.builders.utils import get_parameter, is_parameter
 from executorch.backends.qualcomm.utils.constants import (
     QCOM_DTYPE,
     QCOM_ENCODING,
@@ -130,7 +130,7 @@ def _annotate_quant_attrs(
             self._annotate_requant(n)
             # With fold_quant enabled, check if the input of dq op is quantized param.
             param = None
-            if n.target in dq_ops:
+            if n.target in dq_ops and is_parameter(n.args[0], self.edge_program):
                 param = get_parameter(n.args[0], self.edge_program)
             if n.target not in q_ops and param is None:
                 continue

backends/qualcomm/builders/op_layer_norm.py

@@ -11,7 +11,12 @@
 
 import numpy as np
 import torch
-from executorch.backends.qualcomm.utils.constants import QCOM_DATA
+from executorch.backends.qualcomm.utils.constants import (
+    QCOM_DATA,
+    QCOM_QUANT_ATTRS,
+    QCOM_ZERO_POINT,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
 
 from .node_visitor import NodeVisitor
 from .node_visitor_manager import register_node_visitor
@@ -31,6 +36,7 @@ def define_node(
         node: torch.fx.Node,
         nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
     ) -> PyQnnManager.PyQnnOpWrapper:
+        # args of node : ['input', 'normalized_shape', 'weight', 'bias', 'eps']
         input_node = self.get_node(node.args[0])
         input_tensor = self.get_tensor(input_node, node)
         input_tensor_wrapper = self.define_tensor(
@@ -53,22 +59,41 @@ def define_node(
             return
         axis = [len(input_tensor.shape) - 1]
         axis_shape = [len(axis)]
-        layer_norm_input_tensors = [input_tensor_wrapper]
 
-        weight_node = self.get_node(node.args[2])
-        if weight_node:
+        has_weight = len(node.args) > 2 and node.args[2] is not None
+        if has_weight:
+            weight_node = self.get_node(node.args[2])
+            assert weight_node is not None
             weight_tensor = get_parameter(weight_node, self.edge_program)
-            weight_tensor_wrapper = self.define_tensor(
-                weight_node,
-                node,
-                weight_tensor,
-                PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
-                nodes_to_wrappers,
+        else:
+            # elementwise_affine=False: use all-ones weight as identity
+            weight_tensor = torch.ones(normalized_shapes, dtype=torch.float32)
+            weight_node = torch.fx.Node(
+                node.graph,
+                node.name + "_runtime_weight",
+                "call_function",
+                exir_ops.edge.aten.tensor.default,
+                (),
+                {},
             )
-            layer_norm_input_tensors.append(weight_tensor_wrapper)
+            if quant_attrs := node.meta.get(QCOM_QUANT_ATTRS):
+                quant_attrs = quant_attrs.copy()
+                quant_attrs[QCOM_ZERO_POINT] = 0
+                weight_node.meta[QCOM_QUANT_ATTRS] = quant_attrs
+        weight_tensor_wrapper = self.define_tensor(
+            weight_node,
+            node,
+            weight_tensor,
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
+            nodes_to_wrappers,
+        )
+
+        layer_norm_input_tensors = [input_tensor_wrapper, weight_tensor_wrapper]
 
-        bias_node = self.get_node(node.args[3])
-        if bias_node:
+        has_bias = len(node.args) > 3 and node.args[3] is not None
+        if has_bias:
+            bias_node = self.get_node(node.args[3])
+            assert bias_node is not None
             bias_tensor = get_parameter(bias_node, self.edge_program)
             bias_tensor_wrapper = self.define_tensor(
                 bias_node,
@@ -79,7 +104,7 @@ def define_node(
             )
             layer_norm_input_tensors.append(bias_tensor_wrapper)
 
-        epsilon = node.args[4]
+        epsilon = node.args[4] if len(node.args) > 4 else 1e-05
 
         output_tensor = self.get_tensor(node, node, 0)
         output_tensor_wrapper = self.define_tensor(

backends/qualcomm/builders/utils.py

@@ -37,10 +37,12 @@ def get_parameter(
         param = get_buffer(edge_program, node)
     if is_lifted_tensor_constant(edge_program, node):
         param = get_lifted_tensor_constant(edge_program, node)
-    if param is not None:
-        # update node.meta["val"] to qualified QNN datatype (e.g. i64 to i32)
-        assert isinstance(param, torch.Tensor), "Expect parameter to be tensor"
-        param = param.type(node.meta["val"].dtype)
+    assert (
+        param is not None
+    ), f"Expect {node.name} to be parameter, buffer, or lifted tensor constant"
+    # update node.meta["val"] to qualified QNN datatype (e.g. i64 to i32)
+    assert isinstance(param, torch.Tensor), "Expect parameter to be tensor"
+    param = param.type(node.meta["val"].dtype)
     return param
 
 

backends/qualcomm/quantizer/annotators/htp_rules.py

@@ -851,7 +851,8 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
 
 
 @register_annotator(
-    [torch.ops.aten.layer_norm.default], QnnConstants.OpLayerNorm.op_name
+    [torch.ops.aten.layer_norm.default, torch.ops.aten.native_layer_norm.default],
+    QnnConstants.OpLayerNorm.op_name,
 )
 class LayerNorm(GeneralOpDef):
     @staticmethod

backends/qualcomm/quantizer/annotators/lpai_rules.py

@@ -478,10 +478,8 @@ class LayerNorm(GeneralOpDef):
     @staticmethod
     def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
         act_node = node.args[0]
-        weight_node = node.args[2]
-        bias_node = None
-        if len(node.args) > 2:
-            bias_node = node.args[3]
+        weight_node = node.args[2] if len(node.args) > 2 else None
+        bias_node = node.args[3] if len(node.args) > 3 else None
 
         if _is_annotated([node]):
             return
@@ -492,20 +490,22 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
             act_node,
             input_act_qspec,
         )
-        if input_act_qspec.dtype == torch.int32:
-            annotate_input_qspec_map(
-                node,
-                weight_node,
-                get_16a16w_qnn_ptq_config().weight,
-            )
-        else:
-            annotate_input_qspec_map(
-                node,
-                weight_node,
-                input_act_qspec,
-            )
-        nodes_to_mark_annotated = [node, weight_node]
-        if bias_node:
+        nodes_to_mark_annotated = [node]
+        if isinstance(weight_node, Node):
+            if input_act_qspec.dtype == torch.int32:
+                annotate_input_qspec_map(
+                    node,
+                    weight_node,
+                    get_16a16w_qnn_ptq_config().weight,
+                )
+            else:
+                annotate_input_qspec_map(
+                    node,
+                    weight_node,
+                    input_act_qspec,
+                )
+            nodes_to_mark_annotated.append(weight_node)
+        if isinstance(bias_node, Node):
             annotate_input_qspec_map(
                 node,
                 bias_node,

backends/qualcomm/tests/models.py

@@ -1495,9 +1495,14 @@ def forward(self, x):
 
 
 class LayerNorm(torch.nn.Module):
-    def __init__(self, bias=True):
+    def __init__(self, elementwise_affine=True, bias=True):
         super().__init__()
-        self.layer_norm = torch.nn.LayerNorm([768], eps=1e-6, bias=bias)
+        self.layer_norm = torch.nn.LayerNorm(
+            [768],
+            eps=1e-6,
+            elementwise_affine=elementwise_affine,
+            bias=bias,
+        )
         self.linear = torch.nn.Linear(768, 196)
 
     def forward(self, x):

backends/qualcomm/tests/test_qnn_delegate.py

@@ -1564,7 +1564,11 @@ def test_qnn_backend_up_sampling_nearest_2d_with_size(self):
         self.lower_module_and_test_output(module, sample_input)
 
     def test_qnn_backend_layer_norm(self):
-        modules = [LayerNorm(), LayerNorm(bias=False)]  # noqa: F405
+        modules = [
+            LayerNorm(),  # noqa: F405
+            LayerNorm(bias=False),  # noqa: F405
+            LayerNorm(elementwise_affine=False),  # noqa: F405
+        ]
         sample_input = (torch.randn(196, 768),)
         for i, module in enumerate(modules):
             with self.subTest(i=i):
@@ -4339,7 +4343,11 @@ def test_qnn_backend_up_sampling_nearest_2d_with_size(self):
         self.lower_module_and_test_output(module, sample_input)
 
     def test_qnn_backend_layer_norm(self):
-        modules = [LayerNorm(), LayerNorm(bias=False)]  # noqa: F405
+        modules = [
+            LayerNorm(),  # noqa: F405
+            LayerNorm(bias=False),  # noqa: F405
+            LayerNorm(elementwise_affine=False),  # noqa: F405
+        ]
         sample_input = (torch.randn(196, 768),)
         for i, module in enumerate(modules):
             with self.subTest(i=i):