From f3382e367b2aa4961dd13aba63256d51edd69f41 Mon Sep 17 00:00:00 2001 From: aanelluri Date: Wed, 4 Nov 2026 17:47:56 -0500 Subject: [PATCH] fix: guard skip-rmsnorm fusion against mismatched dtypes (#3155) --- .../ort_fusions/skip_normalization.py | 8 ++++ .../ort_fusions/skip_normalization_test.py | 43 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/onnxscript/rewriter/ort_fusions/skip_normalization.py b/onnxscript/rewriter/ort_fusions/skip_normalization.py index c76a7454cb..036fba4f51 100644 --- a/onnxscript/rewriter/ort_fusions/skip_normalization.py +++ b/onnxscript/rewriter/ort_fusions/skip_normalization.py @@ -88,6 +88,14 @@ def no_match(val: ir.Value, dims: Sequence[str]) -> bool: if stash_type != 1: return check_result.fail("Stash type is not supported.") + if ( + input.dtype is not None + and skip.dtype is not None + and gamma.dtype is not None + and (input.dtype != skip.dtype or input.dtype != gamma.dtype) + ): + return check_result.fail("Dtype mismatch", [input, skip, gamma]) + return check_result def rewrite( diff --git a/onnxscript/rewriter/ort_fusions/skip_normalization_test.py b/onnxscript/rewriter/ort_fusions/skip_normalization_test.py index 6ee80ce5dc..c17e3fa1bb 100644 --- a/onnxscript/rewriter/ort_fusions/skip_normalization_test.py +++ b/onnxscript/rewriter/ort_fusions/skip_normalization_test.py @@ -4,6 +4,8 @@ import unittest +import onnx +import onnx_ir import onnxscript.optimizer from onnxscript.rewriter.models._bart_encoder import bart_encoder_test from onnxscript.rewriter.models._smollm_1 import smollm_test_1 @@ -18,6 +20,47 @@ class TestSkipNormalization(unittest.TestCase): + def test_skip_rms_norm_dtype_mismatch(self): + def make_model(gamma_dtype): + graph = onnx.helper.make_graph( + [ + onnx.helper.make_node("Add", ["input", "skip"], ["skip_sum"]), + onnx.helper.make_node( + "SimplifiedLayerNormalization", + ["skip_sum", "gamma"], + ["output"], + axis=-1, + stash_type=1, + ), + ], + "skip_rms_norm", + [ + onnx.helper.make_tensor_value_info( + "input", onnx.TensorProto.FLOAT16, [2, 8, 64] + ), + onnx.helper.make_tensor_value_info( + "skip", onnx.TensorProto.FLOAT16, [2, 8, 64] + ), + onnx.helper.make_tensor_value_info("gamma", gamma_dtype, [64]), + ], + [ + onnx.helper.make_tensor_value_info( + "output", onnx.TensorProto.FLOAT16, [2, 8, 64] + ) + ], + ) + return onnx_ir.from_proto(onnx.helper.make_model(graph)) + + mixed_dtype_model = make_model(onnx.TensorProto.FLOAT) + fuse_skip_rms_normalization(mixed_dtype_model) + mixed_dtype_op_types = [node.op_type for node in mixed_dtype_model.graph] + self.assertNotIn("SkipSimplifiedLayerNormalization", mixed_dtype_op_types) + + matching_dtype_model = make_model(onnx.TensorProto.FLOAT16) + fuse_skip_rms_normalization(matching_dtype_model) + matching_dtype_op_types = [node.op_type for node in matching_dtype_model.graph] + self.assertIn("SkipSimplifiedLayerNormalization", matching_dtype_op_types) + def test_smollm(self): smollm_test = smollm_test_1() model = smollm_test.get_onnx_model()