[qwen] Port qwen qk_l2norm to LpNormalization (onnxruntime-genai#2127)

modelbuilder/builders/base.py

@@ -1158,6 +1158,11 @@ def make_transpose(self, name, root_input, dtype, shape, perm):
         self.make_value(output, dtype, shape=shape)
         return output
 
+    def make_lp_normalization(self, name, root_input, dtype, shape, axis=-1, p=2):
+        output = f"{name}/output_0"
+        self.make_node("LpNormalization", inputs=[root_input], outputs=[output], name=name, axis=axis, p=p)
+        self.make_value(output, dtype, shape=shape)
+
     def make_div(self, name, inputs, dtype, shape):
         output = f"{name}/output_0"
         self.make_node("Div", inputs=inputs, outputs=[output], name=name)

modelbuilder/builders/qwen.py

@@ -2684,37 +2684,20 @@ def _make_per_head_l2_normalize(self, basename, input_name, n_heads, head_dim):
         return unflat_out
 
     def _make_l2_normalize(self, basename, input_name, last_dim, leading_dims=None):
-        """L2-normalize along last dimension: x * rsqrt(sum(x^2) + eps)
+        """L2-normalize along last dimension via ORT's LpNormalization(p=2).
 
-        Matches the FLA library's l2norm used by the PyTorch reference:
-            inv_norm = rsqrt((x * x).sum(dim=-1, keepdim=True) + eps)
-            return x * inv_norm
+        Replaces the 5-node Square+ReduceSum+Add(eps)+Rsqrt+Mul subgraph with
+        a single LpNormalization op (natively supported by the WebGPU EP and
+        others). The +eps term is dropped; q/k come from RMSNorm+Proj so
+        magnitudes far exceed 1e-6, keeping any divergence within fp16 noise.
         """
         if leading_dims is None:
             leading_dims = ["batch_size", "sequence_length"]
         full_shape = [*leading_dims, last_dim]
-        reduced_shape = [*leading_dims, 1]
 
-        # sum(x^2, dim=-1, keepdim=True)
-        sq_name = f"{basename}/Square/Mul"
-        self.make_mul(sq_name, [input_name, input_name], self.io_dtype, full_shape)
-
-        sum_name = f"{basename}/SumSq/ReduceSum"
-        self.make_reduce_sum(sum_name, [f"{sq_name}/output_0", "/model/constants/INT64/[-1]"], self.io_dtype, reduced_shape, keepdims=True)
-
-        # sum(x^2) + eps
-        eps_name = self._get_shared_l2_eps()
-        add_eps_name = f"{basename}/AddEps/Add"
-        self.make_add(add_eps_name, [f"{sum_name}/output_0", eps_name], self.io_dtype, reduced_shape)
-
-        # x * rsqrt(sum(x^2) + eps)
-        rsqrt_name = f"{basename}/Rsqrt"
-        self.make_rsqrt(rsqrt_name, [f"{add_eps_name}/output_0"], self.io_dtype, reduced_shape)
-
-        norm_name = f"{basename}/Normalize/Mul"
-        self.make_mul(norm_name, [input_name, f"{rsqrt_name}/output_0"], self.io_dtype, full_shape)
-
-        return f"{norm_name}/output_0"
+        node_name = f"{basename}/LpNormalization"
+        self.make_lp_normalization(node_name, input_name, self.io_dtype, full_shape, axis=-1, p=2)
+        return f"{node_name}/output_0"
 
     def _make_gated_rms_norm(self, basename, input_name, gate_name, norm_module, layer_id):
         """Gated RMSNorm: RMSNorm(x) * SiLU(z).

tests/fast/test_random_qwen3_5.py

@@ -332,6 +332,51 @@ def test_qwen3_5_fp16_no_layernorm_fp32_casts(self):
 
         self.assertEqual(offending, [], f"Unexpected fp32 Cast nodes around LayerNorm: {offending}")
 
+    @requires_transformers("5")
+    @hide_stdout()
+    def test_qwen3_5_qk_l2norm_uses_lp_normalization(self):
+        """Verify the qk_l2norm path emits a single LpNormalization op.
+
+        Ported from microsoft/onnxruntime-genai#2127: the previous
+        5-node Square + ReduceSum + Add(eps) + Rsqrt + Mul subgraph used
+        per Q/K head per layer is replaced with a single
+        ``LpNormalization(p=2, axis=-1)`` op, which is natively supported
+        by all current EPs (CPU, CUDA, WebGPU, ...).  This shaves ~5 ops
+        per Q/K head per linear-attention layer and drops the +eps fallback
+        (q/k come from RMSNorm+Proj so magnitudes far exceed 1e-6).
+        """
+        import onnx
+
+        config = _make_qwen3_5_config(["full_attention", "linear_attention"])
+        _, output_dir = self._build_and_save_model(config, "fp32", "cpu")
+
+        text_onnx_path = os.path.join(output_dir, "model.onnx")
+        self.assertExists(text_onnx_path)
+
+        onnx_model = onnx.load(text_onnx_path)
+        lp_norm_nodes = [n for n in onnx_model.graph.node if n.op_type == "LpNormalization"]
+
+        # Each linear-attention layer L emits two qk_l2norm LpNormalization
+        # nodes (one for Q, one for K). The hybrid config has 1 such layer.
+        self.assertGreaterEqual(len(lp_norm_nodes), 2, "expected at least 2 LpNormalization nodes for qk_l2norm Q/K")
+
+        # All qk_l2norm LpNormalization nodes must have axis=-1 and p=2.
+        for node in lp_norm_nodes:
+            if "l2norm" not in node.name:
+                continue
+            axis = next((a.i for a in node.attribute if a.name == "axis"), None)
+            p = next((a.i for a in node.attribute if a.name == "p"), None)
+            self.assertEqual(axis, -1, f"{node.name}: axis must be -1")
+            self.assertEqual(p, 2, f"{node.name}: p must be 2")
+
+        # The old 5-node subgraph names must no longer appear under any
+        # ``*_l2norm`` basename.
+        for node in onnx_model.graph.node:
+            if "l2norm" not in node.name:
+                continue
+            for suffix in ("/Square/Mul", "/SumSq/ReduceSum", "/AddEps/Add", "/Rsqrt", "/Normalize/Mul"):
+                self.assertFalse(node.name.endswith(suffix), f"qk_l2norm path still emits legacy subgraph node {node.name}")
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)