diff --git a/modelbuilder/builders/base.py b/modelbuilder/builders/base.py index 63b0708..e630b64 100644 --- a/modelbuilder/builders/base.py +++ b/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) diff --git a/modelbuilder/builders/qwen.py b/modelbuilder/builders/qwen.py index 8c15d42..c1910f2 100644 --- a/modelbuilder/builders/qwen.py +++ b/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). diff --git a/tests/fast/test_random_qwen3_5.py b/tests/fast/test_random_qwen3_5.py index 0235d40..d187332 100644 --- a/tests/fast/test_random_qwen3_5.py +++ b/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)