From e1663a2885d564ba2d8641af451242ee91f02a8d Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 11:57:06 -0700 Subject: [PATCH 01/14] Refactor ONNX graph surgery base classes Move the surgeon registry and shared proto/rewrite-rule base classes into a dedicated package so new transformations can be implemented in focused modules without growing graph_surgeries.py. Signed-off-by: Xiaoyu Zhang --- olive/passes/onnx/graph_surgeries.py | 88 +---------------- olive/passes/onnx/graph_surgery/__init__.py | 7 ++ olive/passes/onnx/graph_surgery/base.py | 100 ++++++++++++++++++++ 3 files changed, 108 insertions(+), 87 deletions(-) create mode 100644 olive/passes/onnx/graph_surgery/__init__.py create mode 100644 olive/passes/onnx/graph_surgery/base.py diff --git a/olive/passes/onnx/graph_surgeries.py b/olive/passes/onnx/graph_surgeries.py index 44e1b7a939..127e09799b 100644 --- a/olive/passes/onnx/graph_surgeries.py +++ b/olive/passes/onnx/graph_surgeries.py @@ -34,6 +34,7 @@ from olive.model.utils import resolve_onnx_path from olive.passes import Pass from olive.passes.onnx.common import get_external_data_config, model_proto_to_olive_model +from olive.passes.onnx.graph_surgery import ProtoSurgeon, RewriteRuleSurgeon, Surgeon from olive.passes.pass_config import BasePassConfig, PassConfigParam if TYPE_CHECKING: @@ -49,93 +50,6 @@ # pylint: disable=W0621 -class Surgeon: - """Base class for surgeons that operate on the ONNX IR model.""" - - # Refer to https://microsoft.github.io/onnxscript/intermediate_representation/ir_api.html#onnxscript.ir.Model - # for the IR model API. - - registry: ClassVar[dict[str, type[Surgeon]]] = {} - - @classmethod - def __init_subclass__(cls, **kwargs): - super().__init_subclass__(**kwargs) - Surgeon.registry[cls.__name__.lower()] = cls - - def __init__(self): - pass - - def __call__(self, model: ModelProto) -> ModelProto: - return ir.to_proto(self.call_ir(ir.from_proto(model))) - - def call_ir(self, model: ir.Model) -> ir.Model: - # Implement this method in subclasses to operate on the IR model. - raise NotImplementedError - - -class ProtoSurgeon(Surgeon): - """Base class for surgeons that operate on the ONNX model proto directly.""" - - def __call__(self, model: ModelProto) -> ModelProto: - raise NotImplementedError - - def call_ir(self, model: ir.Model) -> ir.Model: - raise RuntimeError("Implement __call__ method instead of operator on onnx.ModelProto directly.") - - @staticmethod - def get_node_by_name(model, name: str, match_output: bool = False): - for node in model.graph.node: - if (match_output and node.output[0] == name) or (not match_output and node.name == name): - return node - return None - - @staticmethod - def get_tensor_shapes(model) -> dict[str, list[int]]: - return {info.name: [x.dim_value for x in info.type.tensor_type.shape.dim] for info in model.graph.value_info} - - @staticmethod - def get_tensor_types(model): - return {info.name: info.type.tensor_type.elem_type for info in model.graph.value_info} - - @staticmethod - def get_initializer_types(model): - return {initializer.name: initializer.data_type for initializer in model.graph.initializer} - - @staticmethod - def get_initializer_shapes(model) -> dict[str, list[int]]: - return {initializer.name: initializer.dims for initializer in model.graph.initializer} - - @staticmethod - def get_initializer_by_name(model, name: str): - for initializer in model.graph.initializer: - if initializer.name == name: - return initializer - return None - - @staticmethod - def create_new_name(name: str, old_op: str, new_op: str) -> str: - return name.replace(old_op, new_op) if old_op in name else f"{name}_{new_op}" - - -class RewriteRuleSurgeon(Surgeon): - """Base class for surgeons implemented as onnxscript rewrite rules. - - Subclasses implement :meth:`rules` to return an - :class:`onnxscript.rewriter.pattern.RewriteRuleSet`, expressing the match - pattern and its replacement with the ONNX IR op builder. ``call_ir`` applies - the rules to the IR model in place. Prefer this over manual proto/DAG - manipulation for local subgraph pattern replacements: the rewriter handles - operand commutativity, use-count bookkeeping, and dead-node cleanup. - """ - - def rules(self) -> pattern.RewriteRuleSet: - raise NotImplementedError - - def call_ir(self, model: ir.Model) -> ir.Model: - self.rules().apply_to_model(model) - return model - - # TODO(anyone): This is incorrect, remove or fix class RenameInputs(Surgeon): def __init__(self, old_names: list[str], new_names: list[str]): diff --git a/olive/passes/onnx/graph_surgery/__init__.py b/olive/passes/onnx/graph_surgery/__init__.py new file mode 100644 index 0000000000..ccd8d9b105 --- /dev/null +++ b/olive/passes/onnx/graph_surgery/__init__.py @@ -0,0 +1,7 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from olive.passes.onnx.graph_surgery.base import ProtoSurgeon, RewriteRuleSurgeon, Surgeon + +__all__ = ["ProtoSurgeon", "RewriteRuleSurgeon", "Surgeon"] diff --git a/olive/passes/onnx/graph_surgery/base.py b/olive/passes/onnx/graph_surgery/base.py new file mode 100644 index 0000000000..6acc780bd8 --- /dev/null +++ b/olive/passes/onnx/graph_surgery/base.py @@ -0,0 +1,100 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from onnxscript import ir + +if TYPE_CHECKING: + from onnx import ModelProto + from onnxscript.rewriter import pattern + + +class Surgeon: + """Base class for surgeons that operate on the ONNX IR model.""" + + # Refer to https://microsoft.github.io/onnxscript/intermediate_representation/ir_api.html#onnxscript.ir.Model + # for the IR model API. + + registry: ClassVar[dict[str, type[Surgeon]]] = {} + + @classmethod + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + Surgeon.registry[cls.__name__.lower()] = cls + + def __init__(self): + pass + + def __call__(self, model: ModelProto) -> ModelProto: + return ir.to_proto(self.call_ir(ir.from_proto(model))) + + def call_ir(self, model: ir.Model) -> ir.Model: + # Implement this method in subclasses to operate on the IR model. + raise NotImplementedError + + +class ProtoSurgeon(Surgeon): + """Base class for surgeons that operate on the ONNX model proto directly.""" + + def __call__(self, model: ModelProto) -> ModelProto: + raise NotImplementedError + + def call_ir(self, model: ir.Model) -> ir.Model: + raise RuntimeError("Implement __call__ method instead of operator on onnx.ModelProto directly.") + + @staticmethod + def get_node_by_name(model, name: str, match_output: bool = False): + for node in model.graph.node: + if (match_output and node.output[0] == name) or (not match_output and node.name == name): + return node + return None + + @staticmethod + def get_tensor_shapes(model) -> dict[str, list[int]]: + return {info.name: [x.dim_value for x in info.type.tensor_type.shape.dim] for info in model.graph.value_info} + + @staticmethod + def get_tensor_types(model): + return {info.name: info.type.tensor_type.elem_type for info in model.graph.value_info} + + @staticmethod + def get_initializer_types(model): + return {initializer.name: initializer.data_type for initializer in model.graph.initializer} + + @staticmethod + def get_initializer_shapes(model) -> dict[str, list[int]]: + return {initializer.name: initializer.dims for initializer in model.graph.initializer} + + @staticmethod + def get_initializer_by_name(model, name: str): + for initializer in model.graph.initializer: + if initializer.name == name: + return initializer + return None + + @staticmethod + def create_new_name(name: str, old_op: str, new_op: str) -> str: + return name.replace(old_op, new_op) if old_op in name else f"{name}_{new_op}" + + +class RewriteRuleSurgeon(Surgeon): + """Base class for surgeons implemented as onnxscript rewrite rules. + + Subclasses implement :meth:`rules` to return an + :class:`onnxscript.rewriter.pattern.RewriteRuleSet`, expressing the match + pattern and its replacement with the ONNX IR op builder. ``call_ir`` applies + the rules to the IR model in place. Prefer this over manual proto/DAG + manipulation for local subgraph pattern replacements: the rewriter handles + operand commutativity, use-count bookkeeping, and dead-node cleanup. + """ + + def rules(self) -> pattern.RewriteRuleSet: + raise NotImplementedError + + def call_ir(self, model: ir.Model) -> ir.Model: + self.rules().apply_to_model(model) + return model From b09abfc2abeb4deef16f897c176ee600d03fe1eb Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:06:25 -0700 Subject: [PATCH 02/14] Port Mobius activation and normalization surgeries Add exporter-independent rewrite-rule surgeons for Gelu, BiasGelu, LayerNormalization, SkipLayerNormalization, and SkipSimplifiedLayerNormalization. Cover matching variants, rejection cases, single-consumer residuals, and shared residual rewiring through the public GraphSurgeries pass. Signed-off-by: Xiaoyu Zhang --- .../passes/onnx/graph_surgery/activations.py | 164 +++++++++++ .../onnx/graph_surgery/normalization.py | 257 ++++++++++++++++++ .../onnx/test_graph_surgeries_activations.py | 183 +++++++++++++ .../test_graph_surgeries_normalization.py | 250 +++++++++++++++++ 4 files changed, 854 insertions(+) create mode 100644 olive/passes/onnx/graph_surgery/activations.py create mode 100644 olive/passes/onnx/graph_surgery/normalization.py create mode 100644 test/passes/onnx/test_graph_surgeries_activations.py create mode 100644 test/passes/onnx/test_graph_surgeries_normalization.py diff --git a/olive/passes/onnx/graph_surgery/activations.py b/olive/passes/onnx/graph_surgery/activations.py new file mode 100644 index 0000000000..23550afd0d --- /dev/null +++ b/olive/passes/onnx/graph_surgery/activations.py @@ -0,0 +1,164 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +import math + +from onnxscript.rewriter import pattern + +from olive.constants import MSFT_DOMAIN +from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon + +_SQRT_2 = math.sqrt(2.0) +_SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi) +_GELU_COEFFICIENT = 0.044715 + + +def _check_constant(value, expected: float, name: str) -> str | None: + if value.const_value is None: + return f"{name} is not a constant" + actual = float(value.const_value.numpy().flat[0]) + if not math.isclose(actual, expected, rel_tol=1e-3): + return f"{name} is {actual}, expected approximately {expected}" + return None + + +class _RemoveReplacedAdd: + def setup(self): + self._replaced_adds = [] + + def _record_replaced_add(self, add): + self._replaced_adds.append(add) + + def cleanup(self): + for add in self._replaced_adds: + if ( + add.graph is not None + and all(output not in add.graph.outputs for output in add.outputs) + and all(not list(output.uses()) for output in add.outputs) + ): + add.graph.remove(add, safe=True) + + +class _ExactGelu(pattern.RewriteRuleClassBase): + def pattern(self, op, x, sqrt_2, one, half): + divided = op.Div(x, sqrt_2) + erf = op.Erf(divided) + shifted = op.Add(erf, one) + scaled = op.Mul(x, shifted) + return op.Mul(scaled, half) + + def check(self, context, sqrt_2, one, half, **_): + result = pattern.MatchResult() + for value, expected, name in ( + (sqrt_2, _SQRT_2, "sqrt(2) divisor"), + (one, 1.0, "Add constant"), + (half, 0.5, "Mul half constant"), + ): + if error := _check_constant(value, expected, name): + return result.fail(error) + return result + + def rewrite(self, op, x, **_): + return op.Gelu(x, approximate="none") + + +class _ExactGeluHalfFirst(_ExactGelu): + def pattern(self, op, x, sqrt_2, one, half): + divided = op.Div(x, sqrt_2) + erf = op.Erf(divided) + shifted = op.Add(erf, one) + scaled = op.Mul(x, shifted) + return op.Mul(half, scaled) + + +class _ApproximateGelu(pattern.RewriteRuleClassBase): + def pattern(self, op, x, three, coefficient, sqrt_2_over_pi, one, half): + cubed = op.Pow(x, three) + scaled_cube = op.Mul(coefficient, cubed) + inner = op.Add(x, scaled_cube) + scaled = op.Mul(sqrt_2_over_pi, inner) + tanh = op.Tanh(scaled) + shifted = op.Add(tanh, one) + activated = op.Mul(x, shifted) + return op.Mul(activated, half) + + def check(self, context, three, coefficient, sqrt_2_over_pi, one, half, **_): + result = pattern.MatchResult() + for value, expected, name in ( + (three, 3.0, "Pow exponent"), + (coefficient, _GELU_COEFFICIENT, "Gelu coefficient"), + (sqrt_2_over_pi, _SQRT_2_OVER_PI, "sqrt(2/pi) constant"), + (one, 1.0, "Add constant"), + (half, 0.5, "Mul half constant"), + ): + if error := _check_constant(value, expected, name): + return result.fail(error) + return result + + def rewrite(self, op, x, **_): + return op.Gelu(x, approximate="tanh") + + +class _ApproximateGeluHalfFirst(_ApproximateGelu): + def pattern(self, op, x, three, coefficient, sqrt_2_over_pi, one, half): + cubed = op.Pow(x, three) + scaled_cube = op.Mul(coefficient, cubed) + inner = op.Add(x, scaled_cube) + scaled = op.Mul(sqrt_2_over_pi, inner) + tanh = op.Tanh(scaled) + shifted = op.Add(tanh, one) + activated = op.Mul(x, shifted) + return op.Mul(half, activated) + + +class FuseGelu(RewriteRuleSurgeon): + """Fuse exact and tanh-approximate decomposed Gelu subgraphs into ONNX Gelu.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet( + [ + _ExactGelu.rule(), + _ExactGeluHalfFirst.rule(), + _ApproximateGelu.rule(), + _ApproximateGeluHalfFirst.rule(), + ] + ) + + +class _AddGeluToBiasGelu(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): + def pattern(self, op, add_output): + return op.Gelu(add_output, _outputs=["gelu_output"]) + + def check(self, context, add_output, gelu_output, **_): + result = pattern.MatchResult() + + gelu = gelu_output.producer() + approximate = gelu.attributes.get("approximate", None) + approximate_value = approximate.value if approximate is not None else "none" + if approximate_value != "tanh": + return result.fail(f"Gelu uses approximate='{approximate_value}', BiasGelu requires 'tanh'") + + add = add_output.producer() + if add is None or add.op_type != "Add": + return result.fail("Input to Gelu is not produced by Add") + + uses = list(add_output.uses()) + if len(uses) != 1: + return result.fail(f"Add output has {len(uses)} consumers, expected exactly one") + + return result + + def rewrite(self, op, add_output, **_): + add = add_output.producer() + self._record_replaced_add(add) + return op.BiasGelu(add.inputs[0], add.inputs[1], _domain=MSFT_DOMAIN) + + +class FuseBiasGelu(RewriteRuleSurgeon): + """Fuse a single-use Add followed by tanh-approximate Gelu into BiasGelu.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_AddGeluToBiasGelu.rule()]) diff --git a/olive/passes/onnx/graph_surgery/normalization.py b/olive/passes/onnx/graph_surgery/normalization.py new file mode 100644 index 0000000000..50a1e47610 --- /dev/null +++ b/olive/passes/onnx/graph_surgery/normalization.py @@ -0,0 +1,257 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +import math + +from onnxscript.rewriter import pattern + +from olive.constants import MSFT_DOMAIN +from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon + + +def _check_scalar_constant(value, expected: float, name: str) -> str | None: + if value.const_value is None: + return f"{name} is not a constant" + actual = float(value.const_value.numpy().flat[0]) + if not math.isclose(actual, expected, rel_tol=1e-4): + return f"{name} is {actual}, expected {expected}" + return None + + +def _check_last_axis(value, name: str) -> str | None: + if value.const_value is None: + return f"ReduceMean {name} axes is not a constant" + axes = list(value.const_value.numpy().flat) + if axes != [-1]: + return f"ReduceMean {name} axes={axes}, expected [-1]" + return None + + +def _check_layer_normalization_constants(exponent, epsilon, first_axes, second_axes): + result = pattern.MatchResult() + if error := _check_scalar_constant(exponent, 2.0, "Pow exponent"): + return result.fail(error) + + if epsilon.const_value is None: + return result.fail("Epsilon is not a constant") + epsilon_value = float(epsilon.const_value.numpy().flat[0]) + if epsilon_value <= 0 or epsilon_value > 1.0: + return result.fail(f"Epsilon {epsilon_value} is outside the expected range (0, 1]") + + for name, axes in (("first", first_axes), ("second", second_axes)): + if error := _check_last_axis(axes, name): + return result.fail(error) + return result + + +class _RemoveReplacedAdd: + def setup(self): + self._replaced_adds = [] + + def _record_replaced_add(self, add): + self._replaced_adds.append(add) + + def cleanup(self): + for add in self._replaced_adds: + if ( + add.graph is not None + and all(output not in add.graph.outputs for output in add.outputs) + and all(not list(output.uses()) for output in add.outputs) + ): + add.graph.remove(add, safe=True) + + +class _LayerNormalization(pattern.RewriteRuleClassBase): + def pattern(self, op, x, first_axes, exponent, second_axes, epsilon, weight, bias): + mean = op.ReduceMean(x, first_axes, _allow_other_attributes=True) + difference = op.Sub(x, mean) + squared = op.Pow(difference, exponent) + variance = op.ReduceMean(squared, second_axes, _allow_other_attributes=True) + variance_epsilon = op.Add(variance, epsilon) + standard_deviation = op.Sqrt(variance_epsilon) + normalized = op.Div(difference, standard_deviation) + scaled = op.Mul(normalized, weight) + return op.Add(scaled, bias) + + def check(self, context, exponent, epsilon, first_axes, second_axes, **_): + return _check_layer_normalization_constants(exponent, epsilon, first_axes, second_axes) + + def rewrite(self, op, x, weight, bias, epsilon, **_): + epsilon_value = float(epsilon.const_value.numpy().flat[0]) + return op.LayerNormalization(x, weight, bias, axis=-1, epsilon=epsilon_value) + + +class _LayerNormalizationNoBias(pattern.RewriteRuleClassBase): + def pattern(self, op, x, first_axes, exponent, second_axes, epsilon, weight): + mean = op.ReduceMean(x, first_axes, _allow_other_attributes=True) + difference = op.Sub(x, mean) + squared = op.Pow(difference, exponent) + variance = op.ReduceMean(squared, second_axes, _allow_other_attributes=True) + variance_epsilon = op.Add(variance, epsilon) + standard_deviation = op.Sqrt(variance_epsilon) + normalized = op.Div(difference, standard_deviation) + return op.Mul(normalized, weight, _outputs=["norm_output"]) + + def check(self, context, exponent, epsilon, first_axes, second_axes, norm_output, **_): + result = _check_layer_normalization_constants(exponent, epsilon, first_axes, second_axes) + if not result: + return result + if list(norm_output.uses()): + return result.fail("Bias-free LayerNormalization output has another node consumer") + return result + + def rewrite(self, op, x, weight, epsilon, **_): + epsilon_value = float(epsilon.const_value.numpy().flat[0]) + return op.LayerNormalization(x, weight, axis=-1, epsilon=epsilon_value) + + +class FuseLayerNormalization(RewriteRuleSurgeon): + """Fuse the ReduceMean-based LayerNormalization decomposition.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_LayerNormalization.rule(), _LayerNormalizationNoBias.rule()]) + + +def _check_skip_input(add_output, norm_output, norm_op_type: str, *, check_axis: bool): + result = pattern.MatchResult() + add = add_output.producer() + if add is None or add.op_type != "Add": + return result.fail(f"Input to {norm_op_type} is not produced by Add") + + if norm_op_type == "LayerNormalization": + for index, value in enumerate(add.inputs): + if value is not None and value.shape is not None and len(value.shape) < 2: + return result.fail(f"Add input[{index}] has rank {len(value.shape)}, expected at least two") + else: + first_shape = add.inputs[0].shape + second_shape = add.inputs[1].shape + first_rank = len(first_shape) if first_shape is not None else None + second_rank = len(second_shape) if second_shape is not None else None + if first_rank is not None and second_rank is not None and first_rank != second_rank: + return result.fail( + f"Add inputs have different ranks ({first_rank} and {second_rank}); fused inputs must have the same shape" + ) + + graph = add.graph + if graph is not None and add_output in graph.outputs: + return result.fail("Add output is a graph output") + + norm = norm_output.producer() + if norm.attributes.get_float("epsilon", None) is None: + return result.fail(f"Missing epsilon attribute on {norm_op_type}") + if check_axis and norm.attributes.get_int("axis", -1) != -1: + return result.fail("LayerNormalization axis must be -1") + return result + + +class _AddLayerNormalizationToSkipLayerNormalization(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): + def pattern(self, op, add_output, weight, bias): + return op.LayerNormalization( + add_output, + weight, + bias, + _allow_other_attributes=True, + _outputs=["norm_output"], + ) + + def check(self, context, add_output, norm_output, **_): + return _check_skip_input(add_output, norm_output, "LayerNormalization", check_axis=True) + + def rewrite(self, op, add_output, weight, bias, norm_output, **_): + add = add_output.producer() + epsilon = norm_output.producer().attributes.get_float("epsilon") + outputs = op.SkipLayerNormalization( + add.inputs[0], + add.inputs[1], + weight, + bias, + _domain=MSFT_DOMAIN, + epsilon=epsilon, + _outputs=4, + ) + add_output.replace_all_uses_with(outputs[3]) + self._record_replaced_add(add) + return outputs[0] + + +class _AddLayerNormalizationNoBiasToSkipLayerNormalization(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): + def pattern(self, op, add_output, weight): + return op.LayerNormalization( + add_output, + weight, + _allow_other_attributes=True, + _outputs=["norm_output"], + ) + + def check(self, context, add_output, norm_output, **_): + result = _check_skip_input(add_output, norm_output, "LayerNormalization", check_axis=True) + if not result: + return result + if len(norm_output.producer().inputs) > 2: + return result.fail("LayerNormalization has a bias") + return result + + def rewrite(self, op, add_output, weight, norm_output, **_): + add = add_output.producer() + epsilon = norm_output.producer().attributes.get_float("epsilon") + outputs = op.SkipLayerNormalization( + add.inputs[0], + add.inputs[1], + weight, + _domain=MSFT_DOMAIN, + epsilon=epsilon, + _outputs=4, + ) + add_output.replace_all_uses_with(outputs[3]) + self._record_replaced_add(add) + return outputs[0] + + +class FuseSkipLayerNormalization(RewriteRuleSurgeon): + """Fuse residual Add and LayerNormalization into SkipLayerNormalization.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet( + [ + _AddLayerNormalizationToSkipLayerNormalization.rule(), + _AddLayerNormalizationNoBiasToSkipLayerNormalization.rule(), + ] + ) + + +class _AddRMSNormalizationToSkipRMSNormalization(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): + def pattern(self, op, add_output, weight): + return op.RMSNormalization( + add_output, + weight, + _allow_other_attributes=True, + _outputs=["norm_output"], + ) + + def check(self, context, add_output, norm_output, **_): + return _check_skip_input(add_output, norm_output, "RMSNormalization", check_axis=False) + + def rewrite(self, op, add_output, weight, norm_output, **_): + add = add_output.producer() + epsilon = norm_output.producer().attributes.get_float("epsilon") + outputs = op.SkipSimplifiedLayerNormalization( + add.inputs[0], + add.inputs[1], + weight, + _domain=MSFT_DOMAIN, + epsilon=epsilon, + _outputs=4, + ) + add_output.replace_all_uses_with(outputs[3]) + self._record_replaced_add(add) + return outputs[0] + + +class FuseSkipRMSNormalization(RewriteRuleSurgeon): + """Fuse residual Add and RMSNormalization into SkipSimplifiedLayerNormalization.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_AddRMSNormalizationToSkipRMSNormalization.rule()]) diff --git a/test/passes/onnx/test_graph_surgeries_activations.py b/test/passes/onnx/test_graph_surgeries_activations.py new file mode 100644 index 0000000000..c222a97b21 --- /dev/null +++ b/test/passes/onnx/test_graph_surgeries_activations.py @@ -0,0 +1,183 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +import math + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from olive.model import ONNXModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.onnx.graph_surgeries import GraphSurgeries +from olive.passes.onnx.graph_surgery.activations import FuseBiasGelu, FuseGelu +from olive.passes.onnx.graph_surgery.base import Surgeon + +_SQRT_2 = math.sqrt(2.0) + + +def _run_surgery(model, tmp_path, surgeon, name): + model_path = tmp_path / f"{name}.onnx" + onnx.save(model, model_path) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + {"surgeries": [{"surgeon": surgeon}], "remove_duplicate_initializers": False}, + disable_search=True, + ) + output_model = graph_surgeries.run(ONNXModelHandler(model_path=str(model_path)), tmp_path / f"{name}_output") + output = output_model.load_model() + onnx.checker.check_model(output) + return output + + +def _count_ops(model): + return { + op_type: sum(node.op_type == op_type for node in model.graph.node) + for op_type in {n.op_type for n in model.graph.node} + } + + +def _make_model(nodes, initializers, outputs): + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4, 8]) + graph_outputs = [helper.make_tensor_value_info(name, TensorProto.FLOAT, [1, 4, 8]) for name in outputs] + graph = helper.make_graph(nodes, "activation_test", [x], graph_outputs, initializers) + return helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid("", 21)]) + + +def _build_exact_gelu(*, half_first=False, sqrt_2=_SQRT_2): + initializers = [ + numpy_helper.from_array(np.array(sqrt_2, dtype=np.float32), name="sqrt_2"), + numpy_helper.from_array(np.array(1.0, dtype=np.float32), name="one"), + numpy_helper.from_array(np.array(0.5, dtype=np.float32), name="half"), + ] + nodes = [ + helper.make_node("Div", ["x", "sqrt_2"], ["divided"]), + helper.make_node("Erf", ["divided"], ["erf"]), + helper.make_node("Add", ["erf", "one"], ["shifted"]), + helper.make_node("Mul", ["x", "shifted"], ["scaled"]), + helper.make_node("Mul", ["half", "scaled"] if half_first else ["scaled", "half"], ["y"]), + ] + return _make_model(nodes, initializers, ["y"]) + + +def _build_approximate_gelu(*, half_first=False, coefficient=0.044715, input_name="x"): + initializers = [ + numpy_helper.from_array(np.array(3.0, dtype=np.float32), name="three"), + numpy_helper.from_array(np.array(coefficient, dtype=np.float32), name="coefficient"), + numpy_helper.from_array(np.array(math.sqrt(2.0 / math.pi), dtype=np.float32), name="sqrt_2_over_pi"), + numpy_helper.from_array(np.array(1.0, dtype=np.float32), name="one"), + numpy_helper.from_array(np.array(0.5, dtype=np.float32), name="half"), + ] + nodes = [ + helper.make_node("Pow", [input_name, "three"], ["cubed"]), + helper.make_node("Mul", ["coefficient", "cubed"], ["scaled_cube"]), + helper.make_node("Add", [input_name, "scaled_cube"], ["inner"]), + helper.make_node("Mul", ["sqrt_2_over_pi", "inner"], ["scaled"]), + helper.make_node("Tanh", ["scaled"], ["tanh"]), + helper.make_node("Add", ["tanh", "one"], ["shifted"]), + helper.make_node("Mul", [input_name, "shifted"], ["activated"]), + helper.make_node("Mul", ["half", "activated"] if half_first else ["activated", "half"], ["y"]), + ] + return nodes, initializers + + +def _build_bias_gelu(*, approximate="tanh", shared_add=False): + bias = numpy_helper.from_array(np.ones(8, dtype=np.float32), name="bias") + nodes = [helper.make_node("Add", ["x", "bias"], ["add_output"])] + gelu_attributes = {} if approximate is None else {"approximate": approximate} + nodes.append(helper.make_node("Gelu", ["add_output"], ["y"], **gelu_attributes)) + outputs = ["y"] + if shared_add: + nodes.append(helper.make_node("Identity", ["add_output"], ["residual"])) + outputs.append("residual") + return _make_model(nodes, [bias], outputs) + + +@pytest.mark.parametrize("surgeon_type", [FuseGelu, FuseBiasGelu]) +def test_activation_surgery_registers_on_module_import(surgeon_type): + assert Surgeon.registry[surgeon_type.__name__.lower()] is surgeon_type + + +@pytest.mark.parametrize("half_first", [False, True]) +def test_fuse_gelu_fuses_exact_variants(tmp_path, half_first): + model = _run_surgery(_build_exact_gelu(half_first=half_first), tmp_path, "FuseGelu", f"exact_{half_first}") + + assert _count_ops(model) == {"Gelu": 1} + gelu = model.graph.node[0] + assert helper.get_attribute_value(next(attr for attr in gelu.attribute if attr.name == "approximate")) == b"none" + + +@pytest.mark.parametrize("half_first", [False, True]) +def test_fuse_gelu_fuses_approximate_variants(tmp_path, half_first): + nodes, initializers = _build_approximate_gelu(half_first=half_first) + model = _run_surgery(_make_model(nodes, initializers, ["y"]), tmp_path, "FuseGelu", f"approx_{half_first}") + + assert _count_ops(model) == {"Gelu": 1} + gelu = model.graph.node[0] + assert helper.get_attribute_value(next(attr for attr in gelu.attribute if attr.name == "approximate")) == b"tanh" + + +@pytest.mark.parametrize( + ("model", "remaining_op"), + [ + (_build_exact_gelu(sqrt_2=2.0), "Erf"), + (_make_model(*_build_approximate_gelu(coefficient=0.05), ["y"]), "Tanh"), + ], +) +def test_fuse_gelu_preserves_non_matching_constants(tmp_path, model, remaining_op): + rewritten = _run_surgery(model, tmp_path, "FuseGelu", f"non_match_{remaining_op}") + + assert _count_ops(rewritten).get("Gelu", 0) == 0 + assert _count_ops(rewritten)[remaining_op] == 1 + + +def test_fuse_bias_gelu_fuses_single_use_tanh_gelu(tmp_path): + model = _run_surgery(_build_bias_gelu(), tmp_path, "FuseBiasGelu", "bias_gelu") + + assert _count_ops(model) == {"BiasGelu": 1} + assert model.graph.node[0].domain == "com.microsoft" + + +@pytest.mark.parametrize("approximate", [None, "none"]) +def test_fuse_bias_gelu_preserves_exact_gelu(tmp_path, approximate): + model = _run_surgery( + _build_bias_gelu(approximate=approximate), + tmp_path, + "FuseBiasGelu", + f"exact_bias_gelu_{approximate}", + ) + + assert _count_ops(model) == {"Add": 1, "Gelu": 1} + + +def test_fuse_bias_gelu_preserves_shared_add(tmp_path): + model = _run_surgery(_build_bias_gelu(shared_add=True), tmp_path, "FuseBiasGelu", "shared_bias_gelu") + + assert _count_ops(model) == {"Add": 1, "Gelu": 1, "Identity": 1} + + +def test_fuse_gelu_then_bias_gelu_through_public_pass(tmp_path): + bias = numpy_helper.from_array(np.ones(8, dtype=np.float32), name="bias") + nodes, initializers = _build_approximate_gelu(input_name="add_output") + model = _make_model([helper.make_node("Add", ["x", "bias"], ["add_output"]), *nodes], [bias, *initializers], ["y"]) + model_path = tmp_path / "combined.onnx" + onnx.save(model, model_path) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + { + "surgeries": [{"surgeon": "FuseGelu"}, {"surgeon": "FuseBiasGelu"}], + "remove_duplicate_initializers": False, + }, + disable_search=True, + ) + + rewritten = graph_surgeries.run( + ONNXModelHandler(model_path=str(model_path)), tmp_path / "combined_output" + ).load_model() + + onnx.checker.check_model(rewritten) + assert _count_ops(rewritten) == {"BiasGelu": 1} diff --git a/test/passes/onnx/test_graph_surgeries_normalization.py b/test/passes/onnx/test_graph_surgeries_normalization.py new file mode 100644 index 0000000000..5f79c88715 --- /dev/null +++ b/test/passes/onnx/test_graph_surgeries_normalization.py @@ -0,0 +1,250 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from olive.model import ONNXModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.onnx.graph_surgeries import GraphSurgeries +from olive.passes.onnx.graph_surgery.base import Surgeon +from olive.passes.onnx.graph_surgery.normalization import ( + FuseLayerNormalization, + FuseSkipLayerNormalization, + FuseSkipRMSNormalization, +) + + +def _run_surgery(model, tmp_path, surgeon, name): + model_path = tmp_path / f"{name}.onnx" + onnx.save(model, model_path) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + {"surgeries": [{"surgeon": surgeon}], "remove_duplicate_initializers": False}, + disable_search=True, + ) + output_model = graph_surgeries.run(ONNXModelHandler(model_path=str(model_path)), tmp_path / f"{name}_output") + output = output_model.load_model() + onnx.checker.check_model(output) + return output + + +def _count_ops(model): + return { + op_type: sum(node.op_type == op_type for node in model.graph.node) + for op_type in {n.op_type for n in model.graph.node} + } + + +def _build_decomposed_layer_normalization(*, include_bias=True, axes=-1, exponent=2.0, epsilon=1e-5): + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4, 8]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4, 8]) + initializers = [ + numpy_helper.from_array(np.array([axes], dtype=np.int64), name="axes"), + numpy_helper.from_array(np.array(exponent, dtype=np.float32), name="exponent"), + numpy_helper.from_array(np.array(epsilon, dtype=np.float32), name="epsilon"), + numpy_helper.from_array(np.ones(8, dtype=np.float32), name="weight"), + ] + nodes = [ + helper.make_node("ReduceMean", ["x", "axes"], ["mean"], keepdims=1), + helper.make_node("Sub", ["x", "mean"], ["difference"]), + helper.make_node("Pow", ["difference", "exponent"], ["squared"]), + helper.make_node("ReduceMean", ["squared", "axes"], ["variance"], keepdims=1), + helper.make_node("Add", ["variance", "epsilon"], ["variance_epsilon"]), + helper.make_node("Sqrt", ["variance_epsilon"], ["standard_deviation"]), + helper.make_node("Div", ["difference", "standard_deviation"], ["normalized"]), + helper.make_node("Mul", ["normalized", "weight"], ["scaled" if include_bias else "y"]), + ] + if include_bias: + initializers.append(numpy_helper.from_array(np.zeros(8, dtype=np.float32), name="bias")) + nodes.append(helper.make_node("Add", ["scaled", "bias"], ["y"])) + graph = helper.make_graph(nodes, "layer_normalization_test", [x], [y], initializers) + return helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid("", 21)]) + + +def _build_skip_normalization( + norm_op_type, + *, + include_bias=False, + shared_add=False, + epsilon=1e-5, + axis=-1, + skip_rank=3, + add_is_graph_output=False, +): + input_shape = [1, 4, 8] + skip_shape = input_shape if skip_rank == 3 else [8] + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape) + skip = helper.make_tensor_value_info("skip", TensorProto.FLOAT, skip_shape) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, input_shape) + initializers = [numpy_helper.from_array(np.ones(8, dtype=np.float32), name="weight")] + nodes = [helper.make_node("Add", ["x", "skip"], ["add_output"])] + norm_inputs = ["add_output", "weight"] + if include_bias: + initializers.append(numpy_helper.from_array(np.zeros(8, dtype=np.float32), name="bias")) + norm_inputs.append("bias") + attributes = {} if epsilon is None else {"epsilon": epsilon} + if norm_op_type == "LayerNormalization": + attributes["axis"] = axis + nodes.append(helper.make_node(norm_op_type, norm_inputs, ["y"], **attributes)) + + outputs = [y] + if shared_add: + nodes.append(helper.make_node("Identity", ["add_output"], ["residual"])) + outputs.append(helper.make_tensor_value_info("residual", TensorProto.FLOAT, input_shape)) + if add_is_graph_output: + outputs.append(helper.make_tensor_value_info("add_output", TensorProto.FLOAT, input_shape)) + + graph = helper.make_graph(nodes, "skip_normalization_test", [x, skip], outputs, initializers) + return helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid("", 23)]) + + +@pytest.mark.parametrize( + "surgeon_type", + [FuseLayerNormalization, FuseSkipLayerNormalization, FuseSkipRMSNormalization], +) +def test_normalization_surgery_registers_on_module_import(surgeon_type): + assert Surgeon.registry[surgeon_type.__name__.lower()] is surgeon_type + + +@pytest.mark.parametrize("include_bias", [False, True]) +def test_fuse_layer_normalization_fuses_bias_variants(tmp_path, include_bias): + model = _run_surgery( + _build_decomposed_layer_normalization(include_bias=include_bias), + tmp_path, + "FuseLayerNormalization", + f"layer_norm_{include_bias}", + ) + + assert _count_ops(model) == {"LayerNormalization": 1} + layer_norm = model.graph.node[0] + assert helper.get_attribute_value(next(attr for attr in layer_norm.attribute if attr.name == "axis")) == -1 + assert helper.get_attribute_value( + next(attr for attr in layer_norm.attribute if attr.name == "epsilon") + ) == pytest.approx(1e-5) + assert len(layer_norm.input) == (3 if include_bias else 2) + + +@pytest.mark.parametrize( + ("kwargs", "remaining_op"), + [ + ({"axes": -2}, "ReduceMean"), + ({"exponent": 3.0}, "Pow"), + ({"epsilon": 2.0}, "Sqrt"), + ], +) +def test_fuse_layer_normalization_preserves_non_matches(tmp_path, kwargs, remaining_op): + model = _run_surgery( + _build_decomposed_layer_normalization(**kwargs), + tmp_path, + "FuseLayerNormalization", + f"layer_norm_non_match_{remaining_op}", + ) + + counts = _count_ops(model) + assert counts.get("LayerNormalization", 0) == 0 + assert counts[remaining_op] >= 1 + + +@pytest.mark.parametrize("include_bias", [False, True]) +def test_fuse_skip_layer_normalization_rewires_shared_residual(tmp_path, include_bias): + model = _run_surgery( + _build_skip_normalization("LayerNormalization", include_bias=include_bias, shared_add=True), + tmp_path, + "FuseSkipLayerNormalization", + f"skip_layer_norm_{include_bias}", + ) + + assert _count_ops(model) == {"Identity": 1, "SkipLayerNormalization": 1} + fused = next(node for node in model.graph.node if node.op_type == "SkipLayerNormalization") + residual = next(node for node in model.graph.node if node.op_type == "Identity") + assert fused.domain == "com.microsoft" + assert residual.input[0] == fused.output[3] + assert len(fused.input) == (4 if include_bias else 3) + + +def test_fuse_skip_layer_normalization_fuses_single_consumer_add(tmp_path): + model = _run_surgery( + _build_skip_normalization("LayerNormalization", include_bias=True), + tmp_path, + "FuseSkipLayerNormalization", + "skip_layer_norm_single_consumer", + ) + + assert _count_ops(model) == {"SkipLayerNormalization": 1} + + +@pytest.mark.parametrize( + "kwargs", + [ + {"axis": 0}, + {"epsilon": None}, + {"skip_rank": 1}, + {"add_is_graph_output": True}, + ], +) +def test_fuse_skip_layer_normalization_preserves_non_matches(tmp_path, kwargs): + model = _run_surgery( + _build_skip_normalization("LayerNormalization", include_bias=True, **kwargs), + tmp_path, + "FuseSkipLayerNormalization", + f"skip_layer_norm_non_match_{next(iter(kwargs))}", + ) + + counts = _count_ops(model) + assert counts.get("SkipLayerNormalization", 0) == 0 + assert counts["Add"] == 1 + assert counts["LayerNormalization"] == 1 + + +def test_fuse_skip_rms_normalization_rewires_shared_residual(tmp_path): + model = _run_surgery( + _build_skip_normalization("RMSNormalization", shared_add=True), + tmp_path, + "FuseSkipRMSNormalization", + "skip_rms_norm_shared", + ) + + assert _count_ops(model) == {"Identity": 1, "SkipSimplifiedLayerNormalization": 1} + fused = next(node for node in model.graph.node if node.op_type == "SkipSimplifiedLayerNormalization") + residual = next(node for node in model.graph.node if node.op_type == "Identity") + assert fused.domain == "com.microsoft" + assert residual.input[0] == fused.output[3] + + +def test_fuse_skip_rms_normalization_fuses_single_consumer_add(tmp_path): + model = _run_surgery( + _build_skip_normalization("RMSNormalization"), + tmp_path, + "FuseSkipRMSNormalization", + "skip_rms_norm_single_consumer", + ) + + assert _count_ops(model) == {"SkipSimplifiedLayerNormalization": 1} + + +@pytest.mark.parametrize( + "kwargs", + [ + {"epsilon": None}, + {"skip_rank": 1}, + {"add_is_graph_output": True}, + ], +) +def test_fuse_skip_rms_normalization_preserves_non_matches(tmp_path, kwargs): + model = _run_surgery( + _build_skip_normalization("RMSNormalization", **kwargs), + tmp_path, + "FuseSkipRMSNormalization", + f"skip_rms_norm_non_match_{next(iter(kwargs))}", + ) + + counts = _count_ops(model) + assert counts.get("SkipSimplifiedLayerNormalization", 0) == 0 + assert counts["Add"] == 1 + assert counts["RMSNormalization"] == 1 From cdd86782d967a5619eefedea21611592cfd25bda Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:07:22 -0700 Subject: [PATCH 03/14] Add attention graph surgeries Port Mobius attention fusion, RoPE separation, QKV packing and unpacking, and block-diagonal packed attention rewrites into exporter-independent Olive surgeons. Add synthetic ONNX IR coverage through the public GraphSurgeries pass for rewrite variants, optional inputs, dtypes, shared inputs, outputs, and non-matches. Signed-off-by: Xiaoyu Zhang --- olive/passes/onnx/graph_surgery/attention.py | 827 ++++++++++++++++++ .../onnx/test_graph_surgeries_attention.py | 652 ++++++++++++++ 2 files changed, 1479 insertions(+) create mode 100644 olive/passes/onnx/graph_surgery/attention.py create mode 100644 test/passes/onnx/test_graph_surgeries_attention.py diff --git a/olive/passes/onnx/graph_surgery/attention.py b/olive/passes/onnx/graph_surgery/attention.py new file mode 100644 index 0000000000..02efada739 --- /dev/null +++ b/olive/passes/onnx/graph_surgery/attention.py @@ -0,0 +1,827 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Attention graph surgeries implemented with ONNXScript rewrite rules.""" + +from __future__ import annotations + +import numpy as np +from onnxscript import ir +from onnxscript.rewriter import pattern +from onnxscript.rewriter._basics import MatchFailureError, MatchResult +from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase + +from olive.constants import MSFT_DOMAIN +from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon + + +def _initializer_dtype(value: ir.Value) -> ir.DataType | None: + """Return an initializer's declared dtype, falling back to its tensor dtype.""" + declared = value.dtype + const_dtype = value.const_value.dtype if value.const_value is not None else None + if declared is not None and const_dtype is not None and declared != const_dtype: + raise ValueError( + f"Initializer {value.name!r} declares dtype {declared} but its const_value data is {const_dtype}." + ) + return declared if declared is not None else const_dtype + + +def _propagate_dtype(source: ir.Value, *targets: ir.Value) -> None: + dtype = _initializer_dtype(source) + if dtype is not None: + for target in targets: + target.dtype = dtype + + +def _skip_view_ops(value: ir.Value | None) -> ir.Value | None: + while value is not None: + producer = value.producer() + if producer is None or producer.op_type not in ("Cast", "Unsqueeze", "Identity"): + return value + value = producer.inputs[0] + return value + + +def _constant_int(value: ir.Value | None) -> int | None: + if value is None: + return None + tensor = value.const_value + if tensor is None: + producer = value.producer() + if producer is None or producer.op_type != "Constant": + return None + attr = producer.attributes.get("value") + tensor = getattr(attr, "value", None) + if tensor is None: + return None + try: + array = tensor.numpy() + except Exception: # pragma: no cover + return None + if array.size != 1: + return None + return int(array.reshape(-1)[0]) + + +class _MaskShape: + __slots__ = ("recognized", "window") + + def __init__(self, recognized: bool, window: int | None = None): + self.recognized = recognized + self.window = window + + +_MASK_WALK_LIMIT = 512 + + +def _sliding_window_from_less(node: ir.Node) -> int | None: + if len(node.inputs) < 2: + return None + distance = _skip_view_ops(node.inputs[0]) + distance_producer = distance.producer() if distance is not None else None + if distance_producer is None or distance_producer.op_type != "Sub": + return None + window = _constant_int(node.inputs[1]) + if window is None or window <= 0: + return None + return window + + +def _local_window_from_attention_bias(attention_bias: ir.Value | None) -> _MaskShape: + """Recover a sliding window and reject masks that GQA cannot represent.""" + window = None + seen = set() + stack = [attention_bias] + visited = 0 + while stack: + value = stack.pop() + if value is None or id(value) in seen: + continue + seen.add(id(value)) + producer = value.producer() + if producer is None: + continue + visited += 1 + if visited > _MASK_WALK_LIMIT: + return _MaskShape(False) + if producer.op_type == "Or": + return _MaskShape(False) + if producer.op_type == "Less": + found = _sliding_window_from_less(producer) + if found is not None: + if window is not None and window != found: + return _MaskShape(False) + window = found + continue + stack.extend(producer.inputs) + return _MaskShape(True, window) + + +def _has_unequal_kv_head_dimensions(k, v, past_key, past_value) -> bool: + def _static_last_dim(value): + if value is None or value.shape is None or len(value.shape) == 0: + return None + dim = value.shape[-1] + return dim if isinstance(dim, int) else None + + for key, value in ((k, v), (past_key, past_value)): + key_dim = _static_last_dim(key) + value_dim = _static_last_dim(value) + if key_dim is not None and value_dim is not None and key_dim != value_dim: + return True + return False + + +class _RotaryAttentionToGQA(RewriteRuleClassBase): + def __init__(self): + super().__init__() + self._seqlens_k = None + self._total_seq_len = None + self._cos_cache = None + self._sin_cache = None + + def pattern(self, op, q_pre, k_pre, v, attention_bias, past_key, past_value, cos, sin): + q_rot = op.RotaryEmbedding(q_pre, cos, sin, _allow_other_attributes=True) + k_rot = op.RotaryEmbedding(k_pre, cos, sin, _allow_other_attributes=True) + return op.Attention( + q_rot, + k_rot, + v, + attention_bias, + past_key, + past_value, + _allow_other_attributes=True, + _outputs=["attn_out", "present_key", "present_value"], + ) + + def check(self, context, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): + result = MatchResult() + if not _local_window_from_attention_bias(attention_bias).recognized: + return result.fail("Attention bias cannot be represented by GroupQueryAttention") + + attn = attn_out.producer() + if attn.attributes.get_float("scale", None) is None: + return result.fail("Missing scale attribute on Attention") + if attn.attributes.get_int("q_num_heads", None) is None: + return result.fail("Missing q_num_heads on Attention") + if attn.attributes.get_int("kv_num_heads", None) is None: + return result.fail("Missing kv_num_heads on Attention") + + cos_prod = cos.producer() + sin_prod = sin.producer() + if cos_prod is None or cos_prod.op_type != "Gather": + return result.fail("cos must be Gather-produced") + if sin_prod is None or sin_prod.op_type != "Gather": + return result.fail("sin must be Gather-produced") + + if past_key is None or past_value is None: + return result.fail("No KV cache inputs") + if past_key.producer() is not None: + return result.fail("past_key is not a graph input") + if past_value.producer() is not None: + return result.fail("past_value is not a graph input") + if _has_unequal_kv_head_dimensions(k_pre, v, past_key, past_value): + return result.fail("K and V head dimensions differ") + return result + + def rewrite( + self, + op, + q_pre, + k_pre, + v, + attention_bias, + past_key, + past_value, + cos, + sin, + attn_out, + present_key, + present_value, + **_, + ): + attn = attn_out.producer() + scale = attn.attributes.get_float("scale") + q_num_heads = attn.attributes.get_int("q_num_heads") + kv_num_heads = attn.attributes.get_int("kv_num_heads") + softcap = attn.attributes.get_float("softcap", 0.0) + + q_rope_node = attn.inputs[0].producer() + rotary_interleaved = q_rope_node.attributes.get_int("interleaved", 0) + rotary_embedding_dim = q_rope_node.attributes.get_int("rotary_embedding_dim", 0) + + if self._cos_cache is None: + self._cos_cache = cos.producer().inputs[0] + self._sin_cache = sin.producer().inputs[0] + + if self._seqlens_k is None: + attention_mask = next(gi for gi in attn.graph.inputs if gi.name == "attention_mask") + axis = op.Constant(value_ints=[1]) + reduce_sum = op.ReduceSum(attention_mask, axis) + one = op.Constant(value_ints=[1]) + self._seqlens_k = op.Cast(op.Sub(reduce_sum, one), to=6) + mask_shape = op.Shape(attention_mask) + idx_1 = op.Constant(value_int=1) + self._total_seq_len = op.Cast(op.Gather(mask_shape, idx_1), to=6) + + attrs = { + "num_heads": q_num_heads, + "kv_num_heads": kv_num_heads, + "scale": scale, + "do_rotary": 1, + "rotary_interleaved": rotary_interleaved, + } + if softcap: + attrs["softcap"] = softcap + if rotary_embedding_dim: + attrs["rotary_embedding_dim"] = rotary_embedding_dim + window = _local_window_from_attention_bias(attention_bias).window + if window is not None: + attrs["local_window_size"] = window + + outputs = op.GroupQueryAttention( + q_pre, + k_pre, + v, + past_key, + past_value, + self._seqlens_k, + self._total_seq_len, + self._cos_cache, + self._sin_cache, + _domain=MSFT_DOMAIN, + _outputs=3, + **attrs, + ) + return outputs[0], outputs[1], outputs[2] + + +class _AttentionToGQA(RewriteRuleClassBase): + def __init__(self): + super().__init__() + self._seqlens_k = None + self._total_seq_len = None + + def pattern(self, op, q, k, v, attention_bias, past_key, past_value): + return op.Attention( + q, + k, + v, + attention_bias, + past_key, + past_value, + _allow_other_attributes=True, + _outputs=["attn_out", "present_key", "present_value"], + ) + + def check(self, context, attn_out, k, v, attention_bias, past_key, past_value, **_): + result = MatchResult() + if not _local_window_from_attention_bias(attention_bias).recognized: + return result.fail("Attention bias cannot be represented by GroupQueryAttention") + + attn = attn_out.producer() + if attn.attributes.get_float("scale", None) is None: + return result.fail("Missing scale attribute on Attention") + if attn.attributes.get_int("q_num_heads", None) is None: + return result.fail("Missing q_num_heads on Attention") + if attn.attributes.get_int("kv_num_heads", None) is None: + return result.fail("Missing kv_num_heads on Attention") + + if past_key is None or past_value is None: + return result.fail("No KV cache inputs") + if past_key.producer() is not None: + return result.fail("past_key is not a graph input") + if past_value.producer() is not None: + return result.fail("past_value is not a graph input") + if not any(gi.name == "attention_mask" for gi in attn.graph.inputs): + return result.fail("No attention_mask graph input") + if _has_unequal_kv_head_dimensions(k, v, past_key, past_value): + return result.fail("K and V head dimensions differ") + return result + + def rewrite( + self, + op, + q, + k, + v, + attention_bias, + past_key, + past_value, + attn_out, + present_key, + present_value, + **_, + ): + attn = attn_out.producer() + scale = attn.attributes.get_float("scale") + q_num_heads = attn.attributes.get_int("q_num_heads") + kv_num_heads = attn.attributes.get_int("kv_num_heads") + softcap = attn.attributes.get_float("softcap", 0.0) + + if self._seqlens_k is None: + attention_mask = next(gi for gi in attn.graph.inputs if gi.name == "attention_mask") + axis = op.Constant(value_ints=[1]) + reduce_sum = op.ReduceSum(attention_mask, axis) + one = op.Constant(value_ints=[1]) + self._seqlens_k = op.Cast(op.Sub(reduce_sum, one), to=6) + mask_shape = op.Shape(attention_mask) + idx_1 = op.Constant(value_int=1) + self._total_seq_len = op.Cast(op.Gather(mask_shape, idx_1), to=6) + + attrs = { + "num_heads": q_num_heads, + "kv_num_heads": kv_num_heads, + "scale": scale, + "do_rotary": 0, + } + if softcap: + attrs["softcap"] = softcap + window = _local_window_from_attention_bias(attention_bias).window + if window is not None: + attrs["local_window_size"] = window + + outputs = op.GroupQueryAttention( + q, + k, + v, + past_key, + past_value, + self._seqlens_k, + self._total_seq_len, + _domain=MSFT_DOMAIN, + _outputs=3, + **attrs, + ) + return outputs[0], outputs[1], outputs[2] + + +class AttentionToGroupQueryAttention(RewriteRuleSurgeon): + """Fuse decoder Attention, and standard RoPE when possible, into GQA.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_RotaryAttentionToGQA.rule(), _AttentionToGQA.rule()]) + + +class _PackQKVForGQA(RewriteRuleClassBase): + def pattern(self, op, hidden, q_w, k_w, v_w): + q = op.MatMul(hidden, op.Transpose(q_w, perm=[1, 0])) + k = op.MatMul(hidden, op.Transpose(k_w, perm=[1, 0])) + v = op.MatMul(hidden, op.Transpose(v_w, perm=[1, 0])) + return op.GroupQueryAttention( + q, + k, + v, + _domain=MSFT_DOMAIN, + _allow_other_attributes=True, + _allow_other_inputs=True, + _outputs=["gqa_out", "present_key", "present_value"], + ) + + def check(self, context, q_w, k_w, v_w, **_): + for name, weight in (("q_w", q_w), ("k_w", k_w), ("v_w", v_w)): + if weight.producer() is not None: + raise MatchFailureError(f"{name} {weight.name!r} is not a graph parameter") + return True + + def rewrite(self, op, hidden, q_w, k_w, v_w, gqa_out, present_key, present_value, **_): + packed_w = op.Concat(q_w, k_w, v_w, axis=0) + packed_wt = op.Transpose(packed_w, perm=[1, 0]) + _propagate_dtype(q_w, packed_w, packed_wt) + packed_qkv = op.MatMul(hidden, packed_wt) + + gqa_node = gqa_out.producer() + attrs = {key: gqa_node.attributes[key].value for key in gqa_node.attributes} + outputs = op.GroupQueryAttention( + packed_qkv, + None, + None, + *gqa_node.inputs[3:], + _domain=MSFT_DOMAIN, + _outputs=3, + **attrs, + ) + return outputs[0], outputs[1], outputs[2] + + +class _PackQKVWithBiasForGQA(RewriteRuleClassBase): + def pattern(self, op, hidden, q_w, bias_q, k_w, bias_k, v_w, bias_v): + q = op.Add(op.MatMul(hidden, op.Transpose(q_w, perm=[1, 0])), bias_q) + k = op.Add(op.MatMul(hidden, op.Transpose(k_w, perm=[1, 0])), bias_k) + v = op.Add(op.MatMul(hidden, op.Transpose(v_w, perm=[1, 0])), bias_v) + return op.GroupQueryAttention( + q, + k, + v, + _domain=MSFT_DOMAIN, + _allow_other_attributes=True, + _allow_other_inputs=True, + _outputs=["gqa_out", "present_key", "present_value"], + ) + + def check(self, context, q_w, bias_q, k_w, bias_k, v_w, bias_v, **_): + for name, weight in (("q_w", q_w), ("k_w", k_w), ("v_w", v_w)): + if weight.producer() is not None: + raise MatchFailureError(f"{name} {weight.name!r} is not a graph parameter") + for bias in (bias_q, bias_k, bias_v): + if bias is None or bias.producer() is not None: + raise MatchFailureError(f"bias {bias!r} is not a graph parameter") + return True + + def rewrite( + self, + op, + hidden, + q_w, + bias_q, + k_w, + bias_k, + v_w, + bias_v, + gqa_out, + present_key, + present_value, + **_, + ): + packed_w = op.Concat(q_w, k_w, v_w, axis=0) + packed_wt = op.Transpose(packed_w, perm=[1, 0]) + _propagate_dtype(q_w, packed_w, packed_wt) + packed_mm = op.MatMul(hidden, packed_wt) + + packed_bias = op.Concat(bias_q, bias_k, bias_v, axis=0) + _propagate_dtype(bias_q, packed_bias) + packed_qkv = op.Add(packed_mm, packed_bias) + + gqa_node = gqa_out.producer() + attrs = {key: gqa_node.attributes[key].value for key in gqa_node.attributes} + outputs = op.GroupQueryAttention( + packed_qkv, + None, + None, + *gqa_node.inputs[3:], + _domain=MSFT_DOMAIN, + _outputs=3, + **attrs, + ) + return outputs[0], outputs[1], outputs[2] + + +class PackQKVForGroupQueryAttention(RewriteRuleSurgeon): + """Pack separate Q/K/V projections, with or without bias, for GQA.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_PackQKVForGQA.rule(), _PackQKVWithBiasForGQA.rule()]) + + +class _SeparateGroupQueryAttentionRoPE(RewriteRuleClassBase): + def pattern(self, op, q, k, v): + return op.GroupQueryAttention( + q, + k, + v, + _domain=MSFT_DOMAIN, + _allow_other_inputs=True, + _allow_other_attributes=True, + _outputs=["gqa_out", "present_key", "present_value"], + ) + + def check(self, context, q, k, v, gqa_out, **_): + result = MatchResult() + gqa_node = gqa_out.producer() + if gqa_node is None: + return result.fail("No GQA producer") + if gqa_node.attributes.get_int("do_rotary", 0) != 1: + return result.fail("do_rotary is not 1") + if len(gqa_node.inputs) < 9: + return result.fail("GQA has fewer than 9 inputs") + if gqa_node.inputs[7] is None or gqa_node.inputs[8] is None: + return result.fail("cos_cache or sin_cache is absent") + if not any(graph_input.name == "position_ids" for graph_input in gqa_node.graph.inputs): + return result.fail("No position_ids graph input") + return result + + def rewrite(self, op, q, k, v, gqa_out, present_key, present_value, **_): + gqa_node = gqa_out.producer() + attrs = {key: gqa_node.attributes[key].value for key in gqa_node.attributes} + attrs["do_rotary"] = 0 + num_heads = attrs.get("num_heads", 1) + kv_num_heads = attrs.get("kv_num_heads", 1) + + cos_cache = gqa_node.inputs[7] + sin_cache = gqa_node.inputs[8] + position_ids = next(graph_input for graph_input in gqa_node.graph.inputs if graph_input.name == "position_ids") + gathered_cos = op.Gather(cos_cache, position_ids) + gathered_sin = op.Gather(sin_cache, position_ids) + q_rot = op.RotaryEmbedding(q, gathered_cos, gathered_sin, num_heads=num_heads) + k_rot = op.RotaryEmbedding(k, gathered_cos, gathered_sin, num_heads=kv_num_heads) + + outputs = op.GroupQueryAttention( + q_rot, + k_rot, + v, + *gqa_node.inputs[3:7], + _domain=MSFT_DOMAIN, + _outputs=3, + **attrs, + ) + return outputs[0], outputs[1], outputs[2] + + +class SeparateGroupQueryAttentionRoPE(RewriteRuleSurgeon): + """Move fused rotary embedding out of GroupQueryAttention.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_SeparateGroupQueryAttentionRoPE.rule()]) + + +class _UnpackGroupQueryAttentionQKV(RewriteRuleClassBase): + def __init__(self): + super().__init__() + self._counter = 0 + self._split_weights = None + self._bias_concat_node = None + + def pattern(self, op, packed_qkv): + return op.GroupQueryAttention( + packed_qkv, + _domain=MSFT_DOMAIN, + _allow_other_inputs=True, + _allow_other_attributes=True, + _outputs=["gqa_out", "present_key", "present_value"], + ) + + def check(self, context, packed_qkv, gqa_out, **_): + result = MatchResult() + gqa_node = gqa_out.producer() + if gqa_node is None: + return result.fail("No GQA producer") + if len(gqa_node.inputs) < 3: + return result.fail("GQA has fewer than 3 inputs") + if gqa_node.inputs[1] is not None or gqa_node.inputs[2] is not None: + return result.fail("GQA is not in packed mode") + + self._split_weights = None + self._bias_concat_node = None + qkv_producer = packed_qkv.producer() + if qkv_producer is not None and qkv_producer.op_type == "Add": + if len(qkv_producer.inputs) < 2 or None in qkv_producer.inputs: + return result.fail("Add has missing inputs") + left, right = qkv_producer.inputs[0], qkv_producer.inputs[1] + left_prod = left.producer() + right_prod = right.producer() + if left_prod is not None and left_prod.op_type == "MatMul": + matmul, bias_value = left_prod, right + elif right_prod is not None and right_prod.op_type == "MatMul": + matmul, bias_value = right_prod, left + else: + return result.fail("Add inputs do not include a MatMul") + + bias_concat = bias_value.producer() if bias_value else None + if bias_concat is None or bias_concat.op_type != "Concat": + return result.fail("Bias is not produced by Concat") + if len(bias_concat.inputs) != 3: + return result.fail("Bias Concat does not have exactly 3 inputs") + for bias_input in bias_concat.inputs: + if bias_input is None or bias_input.producer() is not None: + return result.fail("Bias Concat input is not a graph parameter") + self._bias_concat_node = bias_concat + else: + matmul = qkv_producer + if matmul is None or matmul.op_type != "MatMul": + return result.fail("packed_qkv is not produced by MatMul") + + if len(matmul.inputs) < 2 or matmul.inputs[1] is None: + return result.fail("MatMul has no weight input") + transpose = matmul.inputs[1].producer() + if transpose is None or transpose.op_type != "Transpose": + return result.fail("MatMul weight is not produced by Transpose") + packed_weight = transpose.inputs[0] + if packed_weight is None: + return result.fail("Transpose input is absent") + + concat_node = packed_weight.producer() + if concat_node is not None and concat_node.op_type == "Concat": + if len(concat_node.inputs) != 3: + return result.fail("Concat does not have exactly 3 inputs") + for weight in concat_node.inputs: + if weight is None or weight.producer() is not None: + return result.fail("Concat input is not a graph parameter") + return result + + weight_tensor = ir.convenience.get_const_tensor(packed_weight) + if weight_tensor is None: + return result.fail("Packed weight is neither Concat nor a constant") + num_heads = gqa_node.attributes.get_int("num_heads", None) + kv_num_heads = gqa_node.attributes.get_int("kv_num_heads", None) + if num_heads is None or kv_num_heads is None: + return result.fail("Missing num_heads or kv_num_heads") + + weight_array = weight_tensor.numpy() + total_out = weight_array.shape[0] + total_heads = num_heads + 2 * kv_num_heads + if total_out % total_heads != 0: + return result.fail(f"Cannot determine head_dim from {total_out} outputs and {total_heads} heads") + + head_dim = total_out // total_heads + q_size = num_heads * head_dim + k_size = kv_num_heads * head_dim + self._split_weights = ( + weight_array[:q_size, :], + weight_array[q_size : q_size + k_size, :], + weight_array[q_size + k_size :, :], + ) + return result + + def rewrite(self, op, packed_qkv, gqa_out, present_key, present_value, **_): + gqa_node = gqa_out.producer() + qkv_producer = packed_qkv.producer() + if qkv_producer is not None and qkv_producer.op_type == "Add": + left_prod = qkv_producer.inputs[0].producer() + matmul = ( + left_prod + if left_prod is not None and left_prod.op_type == "MatMul" + else qkv_producer.inputs[1].producer() + ) + else: + matmul = qkv_producer + hidden_states = matmul.inputs[0] + packed_weight = matmul.inputs[1].producer().inputs[0] + concat_node = packed_weight.producer() + + self._counter += 1 + suffix = self._counter + if concat_node is not None and concat_node.op_type == "Concat": + w_q, w_k, w_v = concat_node.inputs + q_mm = op.MatMul(hidden_states, op.Transpose(w_q, perm=[1, 0])) + k_mm = op.MatMul(hidden_states, op.Transpose(w_k, perm=[1, 0])) + v_mm = op.MatMul(hidden_states, op.Transpose(w_v, perm=[1, 0])) + else: + w_q_array, w_k_array, w_v_array = self._split_weights + self._split_weights = None + + def _projection(weight: np.ndarray, name: str) -> ir.Value: + initializer = op.initializer(ir.tensor(weight, name=name), name=name) + return op.MatMul(hidden_states, op.Transpose(initializer, perm=[1, 0])) + + q_mm = _projection(w_q_array, f"unpack_q_weight_{suffix}") + k_mm = _projection(w_k_array, f"unpack_k_weight_{suffix}") + v_mm = _projection(w_v_array, f"unpack_v_weight_{suffix}") + + bias_concat = self._bias_concat_node + self._bias_concat_node = None + if bias_concat is not None: + bias_q, bias_k, bias_v = bias_concat.inputs + q = op.Add(q_mm, bias_q) + k = op.Add(k_mm, bias_k) + v = op.Add(v_mm, bias_v) + else: + q, k, v = q_mm, k_mm, v_mm + + attrs = {key: gqa_node.attributes[key].value for key in gqa_node.attributes} + outputs = op.GroupQueryAttention( + q, + k, + v, + *gqa_node.inputs[3:], + _domain=MSFT_DOMAIN, + _outputs=3, + **attrs, + ) + return outputs[0], outputs[1], outputs[2] + + +class UnpackGroupQueryAttentionQKV(RewriteRuleSurgeon): + """Split packed GQA QKV projections into separate projections.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_UnpackGroupQueryAttentionQKV.rule()]) + + +def _build_packed_token_offset(op, cu_seqlens): + axes_0 = op.Constant(value_ints=[0]) + axes_1 = op.Constant(value_ints=[1]) + neg_one = op.Constant(value_ints=[-1]) + start_1 = op.Constant(value_ints=[1]) + int_max = op.Constant(value_ints=[np.iinfo(np.int64).max]) + + cu_seqlens_i32 = op.Cast(cu_seqlens, to=ir.DataType.INT32) + batch_size = op.Cast(op.Sub(op.Size(cu_seqlens), op.Constant(value_int=1)), to=ir.DataType.INT32) + starts = op.Slice(cu_seqlens_i32, axes_0, neg_one, axes_0) + ends = op.Slice(cu_seqlens_i32, start_1, int_max, axes_0) + lengths = op.Sub(ends, starts) + max_len = op.Squeeze(op.ReduceMax(lengths), axes_0) + + zero_i32 = op.Cast(op.Constant(value_int=0), to=ir.DataType.INT32) + one_i32 = op.Cast(op.Constant(value_int=1), to=ir.DataType.INT32) + rows = op.Range(zero_i32, batch_size, one_i32) + cols = op.Range(zero_i32, max_len, one_i32) + pos_matrix = op.Add(op.Mul(op.Unsqueeze(rows, axes_1), max_len), op.Unsqueeze(cols, axes_0)) + pos_matrix_shape = op.Shape(pos_matrix) + + valid_mask = op.Less(op.Unsqueeze(cols, axes_0), op.Unsqueeze(lengths, axes_1)) + valid_mask_1d = op.Reshape(valid_mask, neg_one) + pos_1d = op.Reshape(pos_matrix, neg_one) + valid_indices = op.Compress(pos_1d, valid_mask_1d) + padding_indices = op.Compress(pos_1d, op.Not(valid_mask_1d)) + return op.Reshape(op.Concat(valid_indices, padding_indices, axis=0), pos_matrix_shape) + + +class _BlockDiagonalAttentionToPackedMHA(RewriteRuleClassBase): + def pattern(self, op, q, k, v, attn_bias_2d): + q_4d = op.Unsqueeze(q, [0]) + k_4d = op.Unsqueeze(k, [0]) + v_4d = op.Unsqueeze(v, [0]) + attn_bias_4d = op.Unsqueeze(attn_bias_2d, [0, 1]) + attn_out = op.Attention( + q_4d, + k_4d, + v_4d, + attn_bias_4d, + _allow_other_attributes=True, + _outputs=["attn_out"], + ) + return op.Squeeze(attn_out, [0]) + + def check(self, context, attn_bias_2d, attn_out, **_): + result = MatchResult() + where = attn_bias_2d.producer() + if where is None or where.op_type != "Where": + return result.fail("Expected Where producing attention bias") + equal = where.inputs[0].producer() + if equal is None or equal.op_type != "Equal": + return result.fail("Expected Equal in mask condition") + unsqueeze_row = equal.inputs[0].producer() + unsqueeze_col = equal.inputs[1].producer() + if not ( + unsqueeze_row + and unsqueeze_col + and unsqueeze_row.op_type == "Unsqueeze" + and unsqueeze_col.op_type == "Unsqueeze" + ): + return result.fail("Expected Unsqueeze ops feeding Equal") + if unsqueeze_row.inputs[0] is not unsqueeze_col.inputs[0]: + return result.fail("Unsqueezes do not share segment_ids") + + segment_ids = unsqueeze_row.inputs[0] + sub = segment_ids.producer() + if sub is None or sub.op_type != "Sub": + return result.fail("Expected Sub in segment_ids chain") + reduce_sum = sub.inputs[0].producer() + if reduce_sum is None or reduce_sum.op_type != "ReduceSum": + return result.fail("Expected ReduceSum in segment_ids chain") + cast = reduce_sum.inputs[0].producer() + if cast is None or cast.op_type != "Cast": + return result.fail("Expected Cast in segment_ids chain") + greater_equal = cast.inputs[0].producer() + if greater_equal is None or greater_equal.op_type != "GreaterOrEqual": + return result.fail("Expected GreaterOrEqual in segment_ids chain") + cu_unsqueeze = greater_equal.inputs[1].producer() + if cu_unsqueeze is None or cu_unsqueeze.op_type != "Unsqueeze": + return result.fail("Expected Unsqueeze for cu_seqlens") + + true_value = ir.convenience.get_const_tensor(where.inputs[1]) + false_value = ir.convenience.get_const_tensor(where.inputs[2]) + if true_value is None or not np.isclose(float(true_value.numpy()), 0.0, atol=1e-3): + return result.fail("Where true branch must be approximately zero") + if false_value is None or float(false_value.numpy()) > -100: + return result.fail("Where false branch must be large negative") + + attention = attn_out.producer() + if attention.attributes.get_float("scale", None) is None: + return result.fail("Missing scale attribute on Attention") + if attention.attributes.get_int("q_num_heads", None) is None: + return result.fail("Missing q_num_heads attribute on Attention") + return result + + @staticmethod + def _trace_cu_seqlens(attn_bias_2d): + where = attn_bias_2d.producer() + equal = where.inputs[0].producer() + segment_ids = equal.inputs[0].producer().inputs[0] + sub = segment_ids.producer() + reduce_sum = sub.inputs[0].producer() + cast = reduce_sum.inputs[0].producer() + greater_equal = cast.inputs[0].producer() + return greater_equal.inputs[1].producer().inputs[0] + + def rewrite(self, op, q, k, v, attn_bias_2d, attn_out, **_): + attention = attn_out.producer() + scale = attention.attributes.get_float("scale") + num_heads = attention.attributes.get_int("q_num_heads") + cu_seqlens = self._trace_cu_seqlens(attn_bias_2d) + cu_seqlens_i32 = op.Cast(cu_seqlens, to=6) + token_offset = _build_packed_token_offset(op, cu_seqlens) + return op.op( + "PackedMultiHeadAttention", + inputs=[q, k, v, None, token_offset, cu_seqlens_i32], + domain=MSFT_DOMAIN, + attributes={"scale": scale, "num_heads": num_heads}, + ) + + +class BlockDiagonalAttentionToPackedMHA(RewriteRuleSurgeon): + """Replace block-diagonal Attention with PackedMultiHeadAttention.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_BlockDiagonalAttentionToPackedMHA.rule()]) diff --git a/test/passes/onnx/test_graph_surgeries_attention.py b/test/passes/onnx/test_graph_surgeries_attention.py new file mode 100644 index 0000000000..a9b75d926f --- /dev/null +++ b/test/passes/onnx/test_graph_surgeries_attention.py @@ -0,0 +1,652 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +from collections import Counter + +import numpy as np +import pytest +from onnxscript import ir + +import olive.passes.onnx.graph_surgery.attention # noqa: F401 +from olive.model import ONNXModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.onnx.graph_surgeries import GraphSurgeries +from olive.passes.onnx.graph_surgery.base import Surgeon + +_MS_DOMAIN = "com.microsoft" + + +def _attribute(name, value): + if isinstance(value, float): + return ir.AttrFloat32(name, value) + if isinstance(value, int): + return ir.AttrInt64(name, value) + if isinstance(value, list): + return ir.AttrInt64s(name, value) + raise TypeError(f"Unsupported attribute value: {value!r}") + + +def _value(name, dtype=ir.DataType.FLOAT, shape=None): + return ir.val(name, dtype=dtype, shape=shape) + + +def _initializer(name, array, dtype=None): + array = np.asarray(array) + if dtype is not None: + array = array.astype(dtype.numpy()) + tensor = ir.tensor(array, dtype=dtype, name=name) + return ir.val(name, dtype=tensor.dtype, shape=array.shape, const_value=tensor) + + +def _node(op_type, inputs, *, attributes=None, domain="", num_outputs=1, output_names=None): + attrs = {name: _attribute(name, value) for name, value in (attributes or {}).items()} + node = ir.Node(domain, op_type, inputs, attributes=attrs, num_outputs=num_outputs) + if output_names: + for output, name in zip(node.outputs, output_names): + output.name = name + return node + + +def _constant(name, value, dtype=None): + array = np.asarray(value) + if dtype is not None: + array = array.astype(dtype.numpy()) + tensor = ir.tensor(array, dtype=dtype, name=f"{name}_value") + node = ir.Node("", "Constant", [], attributes={"value": ir.AttrTensor("value", tensor)}, num_outputs=1) + node.outputs[0].name = name + node.outputs[0].dtype = tensor.dtype + node.outputs[0].shape = ir.Shape(tensor.shape) + node.outputs[0].const_value = tensor + return node + + +def _model(inputs, outputs, nodes, initializers=()): + graph = ir.Graph( + inputs, + outputs, + nodes=nodes, + initializers=initializers, + name="attention_surgery_test", + opset_imports={"": 24, _MS_DOMAIN: 1}, + ) + return ir.Model(graph, ir_version=10) + + +def _count_ops(model): + return Counter(node.op_type for node in model.graph) + + +def _run_surgeries(tmp_path, model, *surgeons): + input_path = tmp_path / "input.onnx" + ir.save(model, input_path) + olive_model = ONNXModelHandler(model_path=str(input_path)) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + { + "surgeries": [{"surgeon": surgeon} for surgeon in surgeons], + "remove_duplicate_initializers": False, + }, + disable_search=True, + ) + output_model = graph_surgeries.run(olive_model, str(tmp_path / "output")) + return ir.load(output_model.model_path) + + +def _make_mask(nodes, *, prefix, sliding_window=None, bidirectional=False): + query_index = _value(f"{prefix}_query_index", ir.DataType.INT64, [1, 1]) + key_index = _value(f"{prefix}_key_index", ir.DataType.INT64, [1, 1]) + causal = _node("GreaterOrEqual", [query_index, key_index], output_names=[f"{prefix}_causal"]) + nodes.append(causal) + condition = causal.outputs[0] + + if sliding_window is not None: + distance = _node("Sub", [query_index, key_index], output_names=[f"{prefix}_distance"]) + window = _constant(f"{prefix}_window", sliding_window, ir.DataType.INT64) + within = _node("Less", [distance.outputs[0], window.outputs[0]], output_names=[f"{prefix}_within"]) + both = _node("And", [condition, within.outputs[0]], output_names=[f"{prefix}_visible"]) + nodes.extend([distance, window, within, both]) + condition = both.outputs[0] + + if bidirectional: + same_block = _node("Equal", [query_index, key_index], output_names=[f"{prefix}_same_block"]) + overlay = _node("Or", [condition, same_block.outputs[0]], output_names=[f"{prefix}_overlay"]) + nodes.extend([same_block, overlay]) + condition = overlay.outputs[0] + + zero = _constant(f"{prefix}_zero", np.float32(0)) + negative = _constant(f"{prefix}_negative", np.float32(-10000)) + bias = _node("Where", [condition, zero.outputs[0], negative.outputs[0]], output_names=[f"{prefix}_bias"]) + nodes.extend([zero, negative, bias]) + return [query_index, key_index], bias.outputs[0] + + +def _make_attention_model( + *, + rotary=False, + num_layers=1, + sliding_window=None, + bidirectional=False, + with_cache=True, + k_dim=32, + v_dim=32, +): + nodes = [] + inputs = [_value("attention_mask", ir.DataType.INT64, [1, "total_sequence"])] + mask_inputs, attention_bias = _make_mask( + nodes, + prefix="mask", + sliding_window=sliding_window, + bidirectional=bidirectional, + ) + inputs.extend(mask_inputs) + + cos_cache = _value("cos_cache", ir.DataType.FLOAT16, [128, 8]) + sin_cache = _value("sin_cache", ir.DataType.FLOAT16, [128, 8]) + position_ids = _value("position_ids", ir.DataType.INT64, [1, 2]) + if rotary: + inputs.extend([cos_cache, sin_cache, position_ids]) + + graph_outputs = [] + for layer in range(num_layers): + q = _value(f"q_{layer}", ir.DataType.FLOAT16, [1, 2, 64]) + k = _value(f"k_{layer}", ir.DataType.FLOAT16, [1, 2, k_dim]) + v = _value(f"v_{layer}", ir.DataType.FLOAT16, [1, 2, v_dim]) + past_key = _value(f"past_key_{layer}", ir.DataType.FLOAT16, [1, 2, 4, k_dim]) + past_value = _value(f"past_value_{layer}", ir.DataType.FLOAT16, [1, 2, 4, v_dim]) + inputs.extend([q, k, v]) + if with_cache: + inputs.extend([past_key, past_value]) + else: + past_key = None + past_value = None + + if rotary: + cos = _node("Gather", [cos_cache, position_ids], output_names=[f"cos_{layer}"]) + sin = _node("Gather", [sin_cache, position_ids], output_names=[f"sin_{layer}"]) + q_rope = _node( + "RotaryEmbedding", + [q, cos.outputs[0], sin.outputs[0]], + attributes={"interleaved": 1, "rotary_embedding_dim": 8}, + output_names=[f"q_rot_{layer}"], + ) + k_rope = _node( + "RotaryEmbedding", + [k, cos.outputs[0], sin.outputs[0]], + attributes={"interleaved": 1, "rotary_embedding_dim": 8}, + output_names=[f"k_rot_{layer}"], + ) + nodes.extend([cos, sin, q_rope, k_rope]) + q_input = q_rope.outputs[0] + k_input = k_rope.outputs[0] + else: + q_input = q + k_input = k + + attention = _node( + "Attention", + [q_input, k_input, v, attention_bias, past_key, past_value], + attributes={"scale": 0.125, "q_num_heads": 4, "kv_num_heads": 2, "softcap": 30.0}, + num_outputs=3, + output_names=[f"output_{layer}", f"present_key_{layer}", f"present_value_{layer}"], + ) + nodes.append(attention) + for output, shape in zip( + attention.outputs, + ([1, 2, 64], [1, 2, 6, k_dim], [1, 2, 6, v_dim]), + ): + output.dtype = ir.DataType.FLOAT16 + output.shape = ir.Shape(shape) + graph_outputs.extend(attention.outputs) + + return _model(inputs, graph_outputs, nodes) + + +def _make_gqa_projection_model(*, bias=False, computed_weight=False, dtype=ir.DataType.FLOAT16): + hidden = _value("hidden", dtype, [1, 2, 8]) + past_key = _value("past_key", dtype, [1, 2, 4, 4]) + past_value = _value("past_value", dtype, [1, 2, 4, 4]) + seqlens = _value("seqlens_k", ir.DataType.INT32, [1]) + total_seq_len = _value("total_sequence_length", ir.DataType.INT32, []) + inputs = [hidden, past_key, past_value, seqlens, total_seq_len] + initializers = [] + nodes = [] + projections = [] + + for name, output_size in (("q", 8), ("k", 4), ("v", 4)): + array = np.arange(output_size * 8, dtype=np.float16).reshape(output_size, 8) + weight = _initializer(f"{name}_weight", array, dtype) + initializers.append(weight) + weight_input = weight + if computed_weight and name == "q": + identity = _node("Identity", [weight], output_names=["computed_q_weight"]) + nodes.append(identity) + weight_input = identity.outputs[0] + transpose = _node("Transpose", [weight_input], attributes={"perm": [1, 0]}, output_names=[f"{name}_weight_t"]) + matmul = _node("MatMul", [hidden, transpose.outputs[0]], output_names=[f"{name}_matmul"]) + nodes.extend([transpose, matmul]) + projection = matmul.outputs[0] + if bias: + bias_value = _initializer(f"{name}_bias", np.arange(output_size, dtype=np.float16), dtype) + initializers.append(bias_value) + add = _node("Add", [projection, bias_value], output_names=[f"{name}_projection"]) + nodes.append(add) + projection = add.outputs[0] + projections.append(projection) + + gqa = _node( + "GroupQueryAttention", + [ + projections[0], + projections[1], + projections[2], + past_key, + past_value, + seqlens, + total_seq_len, + None, + None, + ], + domain=_MS_DOMAIN, + attributes={"num_heads": 4, "kv_num_heads": 2, "scale": 0.5, "do_rotary": 0}, + num_outputs=3, + output_names=["output", "present_key", "present_value"], + ) + nodes.append(gqa) + return _model(inputs, gqa.outputs, nodes, initializers) + + +def _make_packed_gqa_model(*, bias=False, legacy=False, invalid_size=False, dtype=ir.DataType.FLOAT16): + hidden = _value("hidden", dtype, [1, 2, 8]) + past_key = _value("past_key", dtype, [1, 2, 4, 4]) + past_value = _value("past_value", dtype, [1, 2, 4, 4]) + seqlens = _value("seqlens_k", ir.DataType.INT32, [1]) + total_seq_len = _value("total_sequence_length", ir.DataType.INT32, []) + inputs = [hidden, past_key, past_value, seqlens, total_seq_len] + nodes = [] + initializers = [] + + if legacy: + output_size = 15 if invalid_size else 16 + packed_array = np.arange(output_size * 8, dtype=np.float16).reshape(output_size, 8) + packed_weight = _initializer("packed_weight", packed_array, dtype) + initializers.append(packed_weight) + else: + weights = [] + for name, output_size in (("q", 8), ("k", 4), ("v", 4)): + weight = _initializer( + f"{name}_weight", + np.arange(output_size * 8, dtype=np.float16).reshape(output_size, 8), + dtype, + ) + initializers.append(weight) + weights.append(weight) + concat = _node("Concat", weights, attributes={"axis": 0}, output_names=["packed_weight"]) + nodes.append(concat) + packed_weight = concat.outputs[0] + + transpose = _node("Transpose", [packed_weight], attributes={"perm": [1, 0]}, output_names=["packed_weight_t"]) + matmul = _node("MatMul", [hidden, transpose.outputs[0]], output_names=["packed_matmul"]) + nodes.extend([transpose, matmul]) + packed_projection = matmul.outputs[0] + if bias: + biases = [] + for name, output_size in (("q", 8), ("k", 4), ("v", 4)): + bias_value = _initializer(f"{name}_bias", np.arange(output_size, dtype=np.float16), dtype) + initializers.append(bias_value) + biases.append(bias_value) + bias_concat = _node("Concat", biases, attributes={"axis": 0}, output_names=["packed_bias"]) + add = _node("Add", [packed_projection, bias_concat.outputs[0]], output_names=["packed_projection"]) + nodes.extend([bias_concat, add]) + packed_projection = add.outputs[0] + + gqa = _node( + "GroupQueryAttention", + [packed_projection, None, None, past_key, past_value, seqlens, total_seq_len, None, None], + domain=_MS_DOMAIN, + attributes={"num_heads": 4, "kv_num_heads": 2, "scale": 0.5, "do_rotary": 0}, + num_outputs=3, + output_names=["output", "present_key", "present_value"], + ) + nodes.append(gqa) + return _model(inputs, gqa.outputs, nodes, initializers) + + +def _make_fused_rope_gqa_model(*, do_rotary=1, include_caches=True, include_position_ids=True): + dtype = ir.DataType.FLOAT16 + q = _value("q", dtype, [1, 2, 8]) + k = _value("k", dtype, [1, 2, 4]) + v = _value("v", dtype, [1, 2, 4]) + past_key = _value("past_key", dtype, [1, 2, 4, 4]) + past_value = _value("past_value", dtype, [1, 2, 4, 4]) + seqlens = _value("seqlens_k", ir.DataType.INT32, [1]) + total_seq_len = _value("total_sequence_length", ir.DataType.INT32, []) + cos_cache = _value("cos_cache", dtype, [128, 2]) + sin_cache = _value("sin_cache", dtype, [128, 2]) + position_ids = _value("position_ids", ir.DataType.INT64, [1, 2]) + inputs = [q, k, v, past_key, past_value, seqlens, total_seq_len] + if include_caches: + inputs.extend([cos_cache, sin_cache]) + else: + cos_cache = None + sin_cache = None + if include_position_ids: + inputs.append(position_ids) + + gqa = _node( + "GroupQueryAttention", + [q, k, v, past_key, past_value, seqlens, total_seq_len, cos_cache, sin_cache], + domain=_MS_DOMAIN, + attributes={ + "num_heads": 4, + "kv_num_heads": 2, + "scale": 0.5, + "do_rotary": do_rotary, + "rotary_interleaved": 1, + "softcap": 20.0, + }, + num_outputs=3, + output_names=["output", "present_key", "present_value"], + ) + return _model(inputs, gqa.outputs, [gqa]) + + +def _make_block_diagonal_attention_model(*, false_bias=-10000.0, shared_segments=True): + dtype = ir.DataType.FLOAT16 + q = _value("q", dtype, [6, 8]) + k = _value("k", dtype, [6, 8]) + v = _value("v", dtype, [6, 8]) + positions = _value("positions", ir.DataType.INT64, [6, 1]) + cu_seqlens = _value("cu_seqlens", ir.DataType.INT64, [3]) + inputs = [q, k, v, positions, cu_seqlens] + nodes = [] + + axes_0 = _constant("axes_0", [0], ir.DataType.INT64) + axes_1 = _constant("axes_1", [1], ir.DataType.INT64) + cu_unsqueeze = _node("Unsqueeze", [cu_seqlens, axes_0.outputs[0]], output_names=["cu_unsqueeze"]) + greater_equal = _node("GreaterOrEqual", [positions, cu_unsqueeze.outputs[0]], output_names=["greater_equal"]) + cast = _node("Cast", [greater_equal.outputs[0]], attributes={"to": 7}, output_names=["cast"]) + reduce_sum = _node("ReduceSum", [cast.outputs[0], axes_1.outputs[0]], output_names=["reduce_sum"]) + one = _constant("one", np.int64(1), ir.DataType.INT64) + segment_ids = _node("Sub", [reduce_sum.outputs[0], one.outputs[0]], output_names=["segment_ids"]) + row = _node("Unsqueeze", [segment_ids.outputs[0], axes_1.outputs[0]], output_names=["segment_row"]) + column_source = segment_ids.outputs[0] if shared_segments else reduce_sum.outputs[0] + column = _node("Unsqueeze", [column_source, axes_0.outputs[0]], output_names=["segment_column"]) + equal = _node("Equal", [row.outputs[0], column.outputs[0]], output_names=["same_segment"]) + zero = _constant("zero", np.float16(0), dtype) + negative = _constant("negative", np.float16(false_bias), dtype) + bias = _node("Where", [equal.outputs[0], zero.outputs[0], negative.outputs[0]], output_names=["attention_bias"]) + q_4d = _node("Unsqueeze", [q, axes_0.outputs[0]], output_names=["q_4d"]) + k_4d = _node("Unsqueeze", [k, axes_0.outputs[0]], output_names=["k_4d"]) + v_4d = _node("Unsqueeze", [v, axes_0.outputs[0]], output_names=["v_4d"]) + axes_01 = _constant("axes_01", [0, 1], ir.DataType.INT64) + bias_4d = _node("Unsqueeze", [bias.outputs[0], axes_01.outputs[0]], output_names=["bias_4d"]) + attention = _node( + "Attention", + [q_4d.outputs[0], k_4d.outputs[0], v_4d.outputs[0], bias_4d.outputs[0]], + attributes={"scale": 0.5, "q_num_heads": 2}, + output_names=["attention_output"], + ) + squeeze = _node("Squeeze", [attention.outputs[0], axes_0.outputs[0]], output_names=["output"]) + nodes.extend( + [ + axes_0, + axes_1, + cu_unsqueeze, + greater_equal, + cast, + reduce_sum, + one, + segment_ids, + row, + column, + equal, + zero, + negative, + bias, + q_4d, + k_4d, + v_4d, + axes_01, + bias_4d, + attention, + squeeze, + ] + ) + squeeze.outputs[0].dtype = dtype + squeeze.outputs[0].shape = ir.Shape([6, 8]) + return _model(inputs, [squeeze.outputs[0]], nodes) + + +def test_attention_surgeries_register_on_module_import(): + expected = { + "attentiontogroupqueryattention", + "separategroupqueryattentionrope", + "unpackgroupqueryattentionqkv", + "blockdiagonalattentiontopackedmha", + "packqkvforgroupqueryattention", + } + assert expected <= Surgeon.registry.keys() + + +def test_attention_to_gqa_fuses_rotary_preserves_attributes_outputs_and_shared_inputs(tmp_path): + model = _make_attention_model(rotary=True, num_layers=2, sliding_window=64, k_dim=32, v_dim=32) + rewritten = _run_surgeries(tmp_path, model, "AttentionToGroupQueryAttention") + + counts = _count_ops(rewritten) + assert counts["Attention"] == 0 + assert counts["RotaryEmbedding"] == 0 + assert counts["GroupQueryAttention"] == 2 + assert counts["ReduceSum"] == 1 + assert counts["Shape"] == 1 + + gqa_nodes = [node for node in rewritten.graph if node.op_type == "GroupQueryAttention"] + assert len(rewritten.graph.outputs) == 6 + assert [len(node.outputs) for node in gqa_nodes] == [3, 3] + assert gqa_nodes[0].inputs[5].name == gqa_nodes[1].inputs[5].name + assert gqa_nodes[0].inputs[6].name == gqa_nodes[1].inputs[6].name + for node in gqa_nodes: + assert node.inputs[0].shape[-1] == 64 + assert node.inputs[1].shape[-1] == 32 + assert node.attributes.get_int("do_rotary") == 1 + assert node.attributes.get_int("rotary_interleaved") == 1 + assert node.attributes.get_int("rotary_embedding_dim") == 8 + assert node.attributes.get_int("local_window_size") == 64 + assert node.attributes.get_float("softcap") == pytest.approx(30.0) + assert node.inputs[7].name == "cos_cache" + assert node.inputs[8].name == "sin_cache" + + +def test_attention_to_gqa_fallback_uses_external_rope_and_preserves_three_outputs(tmp_path): + model = _make_attention_model(rotary=False, sliding_window=16) + rewritten = _run_surgeries(tmp_path, model, "AttentionToGroupQueryAttention") + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + + assert gqa.attributes.get_int("do_rotary") == 0 + assert gqa.attributes.get_int("local_window_size") == 16 + assert len(gqa.inputs) == 7 + assert len(gqa.outputs) == 3 + assert len(rewritten.graph.outputs) == 3 + + +@pytest.mark.parametrize( + "model", + [ + _make_attention_model(rotary=False, with_cache=False), + _make_attention_model(rotary=False, k_dim=32, v_dim=16), + _make_attention_model(rotary=False, bidirectional=True), + ], +) +def test_attention_to_gqa_preserves_non_matches(tmp_path, model): + rewritten = _run_surgeries(tmp_path, model, "AttentionToGroupQueryAttention") + assert _count_ops(rewritten)["Attention"] == 1 + assert _count_ops(rewritten)["GroupQueryAttention"] == 0 + + +@pytest.mark.parametrize("bias", [False, True]) +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16]) +def test_pack_qkv_packs_bias_variants_preserves_dtype_order_tail_and_outputs(tmp_path, bias, dtype): + model = _make_gqa_projection_model(bias=bias, dtype=dtype) + rewritten = _run_surgeries(tmp_path, model, "PackQKVForGroupQueryAttention") + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + + assert gqa.inputs[1] is None + assert gqa.inputs[2] is None + assert [value.name if value else None for value in gqa.inputs[3:]] == [ + "past_key", + "past_value", + "seqlens_k", + "total_sequence_length", + None, + None, + ] + assert len(gqa.outputs) == 3 + assert _count_ops(rewritten)["MatMul"] == 1 + + packed_projection = gqa.inputs[0] + if bias: + add = packed_projection.producer() + assert add.op_type == "Add" + matmul = next(value.producer() for value in add.inputs if value.producer().op_type == "MatMul") + bias_concat = next(value.producer() for value in add.inputs if value.producer().op_type == "Concat") + assert [value.name for value in bias_concat.inputs] == ["q_bias", "k_bias", "v_bias"] + assert bias_concat.outputs[0].dtype == dtype + else: + matmul = packed_projection.producer() + + transpose = matmul.inputs[1].producer() + weight_concat = transpose.inputs[0].producer() + assert [value.name for value in weight_concat.inputs] == ["q_weight", "k_weight", "v_weight"] + assert weight_concat.outputs[0].dtype == dtype + assert transpose.outputs[0].dtype == dtype + + +def test_pack_qkv_does_not_pack_computed_projection_weight(tmp_path): + model = _make_gqa_projection_model(computed_weight=True) + rewritten = _run_surgeries(tmp_path, model, "PackQKVForGroupQueryAttention") + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + assert gqa.inputs[1] is not None + assert gqa.inputs[2] is not None + assert _count_ops(rewritten)["MatMul"] == 3 + + +def test_separate_rope_adds_gathers_and_rotary_embeddings_and_preserves_attributes(tmp_path): + model = _make_fused_rope_gqa_model() + rewritten = _run_surgeries(tmp_path, model, "SeparateGroupQueryAttentionRoPE") + counts = _count_ops(rewritten) + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + + assert counts["Gather"] == 2 + assert counts["RotaryEmbedding"] == 2 + assert counts["GroupQueryAttention"] == 1 + assert gqa.attributes.get_int("do_rotary") == 0 + assert gqa.attributes.get_int("rotary_interleaved") == 1 + assert gqa.attributes.get_float("softcap") == pytest.approx(20.0) + assert len(gqa.inputs) == 7 + assert len(gqa.outputs) == 3 + q_rope = gqa.inputs[0].producer() + k_rope = gqa.inputs[1].producer() + assert q_rope.attributes.get_int("num_heads") == 4 + assert k_rope.attributes.get_int("num_heads") == 2 + + +@pytest.mark.parametrize( + "model", + [ + _make_fused_rope_gqa_model(do_rotary=0), + _make_fused_rope_gqa_model(include_caches=False), + _make_fused_rope_gqa_model(include_position_ids=False), + ], +) +def test_separate_rope_preserves_non_matches(tmp_path, model): + rewritten = _run_surgeries(tmp_path, model, "SeparateGroupQueryAttentionRoPE") + assert _count_ops(rewritten)["RotaryEmbedding"] == 0 + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + assert gqa.attributes.get_int("do_rotary") in (0, 1) + assert len(gqa.inputs) == 9 + + +@pytest.mark.parametrize("bias", [False, True]) +def test_unpack_qkv_unpacks_concat_variants_and_preserves_optional_tail(tmp_path, bias): + model = _make_packed_gqa_model(bias=bias) + rewritten = _run_surgeries(tmp_path, model, "UnpackGroupQueryAttentionQKV") + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + + assert all(gqa.inputs[index] is not None for index in range(3)) + assert [value.name if value else None for value in gqa.inputs[3:]] == [ + "past_key", + "past_value", + "seqlens_k", + "total_sequence_length", + None, + None, + ] + # The original packed projection was traced from the GQA input rather than + # included in the match pattern, so it remains dead until unused-node cleanup. + assert _count_ops(rewritten)["MatMul"] == 4 + expected_producer = "Add" if bias else "MatMul" + assert [gqa.inputs[index].producer().op_type for index in range(3)] == [expected_producer] * 3 + assert len(gqa.outputs) == 3 + + +def test_unpack_qkv_splits_legacy_initializer_with_unequal_q_and_kv_sizes(tmp_path): + model = _make_packed_gqa_model(legacy=True) + rewritten = _run_surgeries(tmp_path, model, "UnpackGroupQueryAttentionQKV") + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + + weights = [gqa.inputs[index].producer().inputs[1].producer().inputs[0] for index in range(3)] + assert [list(weight.shape) for weight in weights] == [[8, 8], [4, 8], [4, 8]] + assert [weight.dtype for weight in weights] == [ir.DataType.FLOAT16] * 3 + original = np.arange(16 * 8, dtype=np.float16).reshape(16, 8) + np.testing.assert_array_equal(weights[0].const_value.numpy(), original[:8]) + np.testing.assert_array_equal(weights[1].const_value.numpy(), original[8:12]) + np.testing.assert_array_equal(weights[2].const_value.numpy(), original[12:]) + + +@pytest.mark.parametrize( + "model", + [ + _make_gqa_projection_model(), + _make_packed_gqa_model(legacy=True, invalid_size=True), + ], +) +def test_unpack_qkv_preserves_non_matches(tmp_path, model): + rewritten = _run_surgeries(tmp_path, model, "UnpackGroupQueryAttentionQKV") + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + if gqa.inputs[1] is None: + assert _count_ops(rewritten)["MatMul"] == 1 + else: + assert _count_ops(rewritten)["MatMul"] == 3 + + +def test_block_diagonal_attention_rewrites_to_packed_mha_with_offsets(tmp_path): + model = _make_block_diagonal_attention_model() + rewritten = _run_surgeries(tmp_path, model, "BlockDiagonalAttentionToPackedMHA") + counts = _count_ops(rewritten) + packed = next(node for node in rewritten.graph if node.op_type == "PackedMultiHeadAttention") + + assert counts["Attention"] == 0 + assert counts["PackedMultiHeadAttention"] == 1 + assert packed.domain == _MS_DOMAIN + assert [value.name if value else None for value in packed.inputs[:4]] == ["q", "k", "v", None] + assert packed.inputs[4].producer().op_type == "Reshape" + assert packed.inputs[5].producer().op_type == "Cast" + assert packed.inputs[5].producer().inputs[0].name == "cu_seqlens" + assert packed.attributes.get_float("scale") == pytest.approx(0.5) + assert packed.attributes.get_int("num_heads") == 2 + assert len(rewritten.graph.outputs) == 1 + + +@pytest.mark.parametrize( + "model", + [ + _make_block_diagonal_attention_model(false_bias=-1.0), + _make_block_diagonal_attention_model(shared_segments=False), + ], +) +def test_block_diagonal_attention_preserves_non_matches(tmp_path, model): + rewritten = _run_surgeries(tmp_path, model, "BlockDiagonalAttentionToPackedMHA") + assert _count_ops(rewritten)["Attention"] == 1 + assert _count_ops(rewritten)["PackedMultiHeadAttention"] == 0 From d6b24cdb307173622a70df3b70001dc1ef90fe09 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:07:56 -0700 Subject: [PATCH 04/14] Port weight-aware MoE graph surgeries Add exporter-independent GraphSurgeries for native block-quantized and MatMulNBits dense MoE graphs. Preserve packed initializer bytes, routing semantics, fail-closed native-block behavior, and external-data support with focused synthetic tests. Signed-off-by: Xiaoyu Zhang --- olive/passes/onnx/graph_surgery/moe.py | 906 +++++++++++++++++++ test/passes/onnx/test_graph_surgeries_moe.py | 837 +++++++++++++++++ 2 files changed, 1743 insertions(+) create mode 100644 olive/passes/onnx/graph_surgery/moe.py create mode 100644 test/passes/onnx/test_graph_surgeries_moe.py diff --git a/olive/passes/onnx/graph_surgery/moe.py b/olive/passes/onnx/graph_surgery/moe.py new file mode 100644 index 0000000000..1c273ab804 --- /dev/null +++ b/olive/passes/onnx/graph_surgery/moe.py @@ -0,0 +1,906 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Weight-aware graph surgeries for dense-fallback Mixture-of-Experts graphs.""" + +from __future__ import annotations + +import dataclasses +import logging + +import numpy as np +from onnxscript import ir + +from olive.passes.onnx.graph_surgery import Surgeon + +logger = logging.getLogger(__name__) + +_MS_DOMAIN = "com.microsoft" +_NXRT_DOMAIN = "pkg.nxrt" +_BLOCK_MATMUL = "BlockQuantizedMatMul" +_BLOCK_MOE = "BlockQuantizedMoE" + +_NATIVE_BLOCK_FORMATS = { + "mxfp4": (32, 17), + "iq4_nl": (32, 18), + "iq4_xs": (256, 136), + "iq3_s": (256, 110), + "iq3_xxs": (256, 98), + "iq2_xxs": (256, 66), + "iq2_xs": (256, 74), + "iq2_s": (256, 82), + "iq1_s": (256, 50), + "iq1_m": (256, 56), +} + + +class MoEGraphSurgeryError(ValueError): + """A routed MoE graph cannot be represented by the requested sparse kernel.""" + + +class _UnfusableBlockMoEError(Exception): + """A native-block MoE layer cannot be represented by one fused node.""" + + +def _scalar_int(value: ir.Value | None) -> int | None: + if value is None: + return None + const = value.const_value + if const is None: + producer = value.producer() + if producer is None or producer.op_type != "Constant": + return None + for attr in ("value", "value_int", "value_ints"): + if attr in producer.attributes: + const = producer.attributes[attr].value + break + else: + return None + array = np.asarray(const.numpy() if hasattr(const, "numpy") else const).reshape(-1) + return int(array[0]) if array.size else None + + +def _consumers_of_type(value: ir.Value, op_type: str) -> list[ir.Node]: + return [node for node, _ in value.uses() if node.op_type == op_type] + + +def _single_consumer(value: ir.Value, *op_types: str) -> ir.Node | None: + for op_type in op_types: + matches = _consumers_of_type(value, op_type) + if matches: + return matches[0] + return None + + +def _require_producer(value: ir.Value, op_type: str) -> ir.Node | None: + producer = value.producer() + return producer if producer is not None and producer.op_type == op_type else None + + +def _array(value: ir.Value, *, block_moe: bool = False) -> np.ndarray: + const = value.const_value + if const is None: + message = f"expected an initializer for {value.name!r}" + if block_moe: + raise _UnfusableBlockMoEError(message) + raise ValueError(message) + return np.asarray(const.numpy()) + + +def _make_initializer(graph: ir.Graph, name: str, array: np.ndarray, dtype: ir.DataType) -> ir.Value: + tensor = ir.tensor(array, name=name, dtype=dtype) + value = ir.Value( + name=name, + shape=ir.Shape(array.shape), + type=ir.TensorType(dtype), + const_value=tensor, + ) + graph.register_initializer(value) + return value + + +def _remove_dead_nodes(graph: ir.Graph) -> None: + graph_outputs = set(graph.outputs) + changed = True + while changed: + changed = False + for node in reversed(list(graph)): + if any(output in graph_outputs for output in node.outputs): + continue + if all(not list(output.uses()) for output in node.outputs): + graph.remove(node, safe=True) + changed = True + + used = { + graph_input.name + for node in graph + for graph_input in node.inputs + if graph_input is not None and graph_input.name is not None + } + for name in list(graph.initializers): + if name not in used: + del graph.initializers[name] + + +class _ExpertProjections: + __slots__ = ("down", "gate", "up") + + def __init__(self, gate: ir.Node, up: ir.Node, down: ir.Node) -> None: + self.gate = gate + self.up = up + self.down = down + + +def _find_routed_output(contributions: list[ir.Value]) -> ir.Value | None: + if not contributions: + return None + contribution_set = set(contributions) + accumulation_outputs: set[ir.Value] = set() + current = contributions[0] + advanced = True + while advanced: + advanced = False + for node, index in current.uses(): + if node.op_type != "Add": + continue + other = node.inputs[1 - index] + if other in contribution_set or other in accumulation_outputs: + current = node.outputs[0] + accumulation_outputs.add(current) + advanced = True + break + return current + + +# MatMulNBits dense MoE to com.microsoft::QMoE. + + +def _trace_qmoe_expert(down: ir.Node) -> _ExpertProjections | None: + activation_mul = _require_producer(down.inputs[0], "Mul") + if activation_mul is None: + return None + first, second = activation_mul.inputs[:2] + if first.producer() is not None and first.producer().op_type == "MatMulNBits": + up_output, activation_output = first, second + elif second.producer() is not None and second.producer().op_type == "MatMulNBits": + up_output, activation_output = second, first + else: + return None + + up = up_output.producer() + activation = activation_output.producer() + if activation is None: + return None + if activation.op_type == "Swish": + gate = activation.inputs[0].producer() + if gate is not None and gate.op_type == "MatMulNBits": + return _ExpertProjections(gate, up, down) + return None + if activation.op_type != "Mul": + return None + + for gate_output, sigmoid_output in ( + (activation.inputs[0], activation.inputs[1]), + (activation.inputs[1], activation.inputs[0]), + ): + gate = gate_output.producer() + sigmoid = sigmoid_output.producer() + if ( + gate is not None + and gate.op_type == "MatMulNBits" + and sigmoid is not None + and sigmoid.op_type == "Sigmoid" + and sigmoid.inputs[0] is gate_output + ): + return _ExpertProjections(gate, up, down) + return None + + +def _router_matmul(logits: ir.Value) -> ir.Node | None: + producer = logits.producer() + if producer is not None and producer.op_type == "Cast": + source = producer.inputs[0] + producer = source.producer() if source is not None else None + return producer if producer is not None and producer.op_type == "MatMulNBits" else None + + +def _find_qmoe_anchors(graph: ir.Graph) -> list[ir.Node]: + anchors = [] + for node in graph: + if node.op_type != "TopK" or node.inputs[0] is None or _router_matmul(node.inputs[0]) is None: + continue + if _consumers_of_type(node.outputs[1], "Equal"): + anchors.append(node) + return anchors + + +class _DenseQMoELayer: + def __init__(self, topk: ir.Node) -> None: + self.topk = topk + self.gate_router = _router_matmul(topk.inputs[0]) + self.hidden = self.gate_router.inputs[0] + self.logits = topk.inputs[0] + self.k = _scalar_int(topk.inputs[1]) + softmax = _single_consumer(topk.outputs[0], "Softmax", "Cast") + if softmax is not None and softmax.op_type == "Cast": + softmax = _single_consumer(softmax.outputs[0], "Softmax") + self.softmax = softmax + self.experts: dict[int, _ExpertProjections] = {} + self.contributions: list[ir.Value] = [] + self._collect_experts() + self.routed_out = _find_routed_output(self.contributions) + + def _collect_experts(self) -> None: + for equal in _consumers_of_type(self.topk.outputs[1], "Equal"): + expert_id = _scalar_int(equal.inputs[1]) + if expert_id is None: + continue + cast = _single_consumer(equal.outputs[0], "CastLike", "Cast") + if cast is None: + continue + weight_mul = _single_consumer(cast.outputs[0], "Mul") + reduce_sum = _single_consumer(weight_mul.outputs[0], "ReduceSum") if weight_mul else None + contribution = _single_consumer(reduce_sum.outputs[0], "Mul") if reduce_sum else None + if contribution is None: + continue + weight_output = reduce_sum.outputs[0] + expert_output = ( + contribution.inputs[1] if contribution.inputs[0] is weight_output else contribution.inputs[0] + ) + down = _require_producer(expert_output, "MatMulNBits") + projections = _trace_qmoe_expert(down) if down is not None else None + if projections is not None: + self.experts[expert_id] = projections + self.contributions.append(contribution.outputs[0]) + + @property + def is_valid(self) -> bool: + if self.k is None or self.softmax is None or self.routed_out is None: + return False + expert_ids = sorted(self.experts) + return expert_ids == list(range(len(expert_ids))) and bool(expert_ids) + + +def _qmoe_geometry(node: ir.Node) -> tuple[int, int, bool]: + return ( + int(node.attributes["bits"].value), + int(node.attributes["block_size"].value), + len(node.inputs) > 3 and node.inputs[3] is not None, + ) + + +def _qmoe_abi_supported(bits: int, block_size: int) -> bool: + return bits == 4 and block_size >= 16 and not block_size & (block_size - 1) + + +def _pack_qmoe_projection(nodes: list[ir.Node], slot: int) -> np.ndarray: + arrays = [_array(node.inputs[slot]) for node in nodes] + return np.stack([array.reshape(array.shape[0], -1) for array in arrays], axis=0) + + +def _pack_qmoe_scales(nodes: list[ir.Node], slot: int, out_features: int, dtype: ir.DataType) -> np.ndarray: + return np.stack([_array(node.inputs[slot]).reshape(out_features, -1) for node in nodes], axis=0).astype( + dtype.numpy() + ) + + +def _pack_qmoe_zero_points(nodes: list[ir.Node], slot: int, out_features: int) -> np.ndarray: + return np.stack( + [_array(node.inputs[slot]).reshape(out_features, -1).astype(np.uint8) for node in nodes], + axis=0, + ) + + +def _emit_qmoe(graph: ir.Graph, layer: _DenseQMoELayer, index: int, activation_dtype: ir.DataType) -> None: + expert_ids = sorted(layer.experts) + gate_nodes = [layer.experts[expert_id].gate for expert_id in expert_ids] + up_nodes = [layer.experts[expert_id].up for expert_id in expert_ids] + down_nodes = [layer.experts[expert_id].down for expert_id in expert_ids] + bits, block_size, has_zero_point = _qmoe_geometry(down_nodes[0]) + + intermediate_size = _array(gate_nodes[0].inputs[1]).shape[0] + hidden_size = _array(down_nodes[0].inputs[1]).shape[0] + fc1_weights = np.concatenate([_pack_qmoe_projection(gate_nodes, 1), _pack_qmoe_projection(up_nodes, 1)], axis=1) + fc2_weights = _pack_qmoe_projection(down_nodes, 1) + fc1_scales = np.concatenate( + [ + _pack_qmoe_scales(gate_nodes, 2, intermediate_size, activation_dtype), + _pack_qmoe_scales(up_nodes, 2, intermediate_size, activation_dtype), + ], + axis=1, + ) + fc2_scales = _pack_qmoe_scales(down_nodes, 2, hidden_size, activation_dtype) + + fc1_zero_points = fc2_zero_points = None + if has_zero_point: + fc1_zero_points = np.concatenate( + [ + _pack_qmoe_zero_points(gate_nodes, 3, intermediate_size), + _pack_qmoe_zero_points(up_nodes, 3, intermediate_size), + ], + axis=1, + ) + fc2_zero_points = _pack_qmoe_zero_points(down_nodes, 3, hidden_size) + + prefix = f"moe.layer{index}" + fc1_weights_value = _make_initializer(graph, f"{prefix}.fc1_experts_weights", fc1_weights, ir.DataType.UINT8) + fc1_scales_value = _make_initializer(graph, f"{prefix}.fc1_scales", fc1_scales, activation_dtype) + fc2_weights_value = _make_initializer(graph, f"{prefix}.fc2_experts_weights", fc2_weights, ir.DataType.UINT8) + fc2_scales_value = _make_initializer(graph, f"{prefix}.fc2_scales", fc2_scales, activation_dtype) + fc1_zero_points_value = ( + _make_initializer(graph, f"{prefix}.fc1_zero_points", fc1_zero_points, ir.DataType.UINT8) + if fc1_zero_points is not None + else None + ) + fc2_zero_points_value = ( + _make_initializer(graph, f"{prefix}.fc2_zero_points", fc2_zero_points, ir.DataType.UINT8) + if fc2_zero_points is not None + else None + ) + + cast = ir.node( + "Cast", + inputs=[layer.logits], + attributes={"to": activation_dtype.value}, + num_outputs=1, + name=f"{prefix}.router_probs_cast", + ) + cast.outputs[0].name = f"{prefix}.router_probs" + qmoe = ir.node( + "QMoE", + inputs=[ + layer.hidden, + cast.outputs[0], + fc1_weights_value, + fc1_scales_value, + None, + fc2_weights_value, + fc2_scales_value, + None, + None, + None, + None, + fc1_zero_points_value, + fc2_zero_points_value, + None, + None, + ], + attributes={ + "activation_type": "swiglu", + "normalize_routing_weights": 1, + "k": layer.k, + "expert_weight_bits": bits, + "block_size": block_size, + "swiglu_fusion": 2, + "quant_type": "int", + "weights_prepacked": 0, + }, + domain=_MS_DOMAIN, + num_outputs=1, + name=f"{prefix}.qmoe", + ) + qmoe.outputs[0].name = f"{prefix}.qmoe_output" + qmoe.outputs[0].type = layer.hidden.type + qmoe.outputs[0].shape = layer.routed_out.shape + graph.insert_after(layer.topk, [cast, qmoe]) + layer.routed_out.replace_all_uses_with(qmoe.outputs[0]) + graph.opset_imports[_MS_DOMAIN] = 1 + + +def _fuse_dense_moe_to_qmoe(model: ir.Model) -> int: + graph = model.graph + fused = 0 + for index, layer in enumerate(_DenseQMoELayer(topk) for topk in _find_qmoe_anchors(graph)): + if not layer.is_valid: + logger.warning("Skipping MoE layer at %s: unrecognized dense-fallback structure", layer.topk.name) + continue + down = layer.experts[min(layer.experts)].down + bits, block_size, _ = _qmoe_geometry(down) + if not _qmoe_abi_supported(bits, block_size): + logger.warning( + "Skipping MoE layer at %s: QMoE ABI unsupported (bits=%d, block_size=%d)", + layer.topk.name, + bits, + block_size, + ) + continue + activation_dtype = layer.hidden.dtype or down.inputs[2].dtype + if activation_dtype not in {ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16}: + logger.warning( + "Skipping MoE layer at %s: activation dtype is missing or unsupported (%s)", + layer.topk.name, + activation_dtype, + ) + continue + _emit_qmoe(graph, layer, index, activation_dtype) + fused += 1 + if fused: + _remove_dead_nodes(graph) + graph.sort() + return fused + + +class FuseDenseMoEToQMoE(Surgeon): + """Fuse eligible dense MatMulNBits expert storms into com.microsoft::QMoE.""" + + def call_ir(self, model: ir.Model) -> ir.Model: + _fuse_dense_moe_to_qmoe(model) + return model + + +# Native BlockQuantizedMatMul dense MoE to pkg.nxrt::BlockQuantizedMoE. + + +def _skip_cast(value: ir.Value) -> ir.Value: + producer = value.producer() + if producer is not None and producer.op_type == "Cast" and producer.inputs[0] is not None: + return producer.inputs[0] + return value + + +def _producer_through_cast(value: ir.Value, op_type: str) -> ir.Node | None: + producer = _skip_cast(value).producer() + return producer if producer is not None and producer.op_type == op_type else None + + +def _trace_block_expert(down: ir.Node) -> _ExpertProjections | None: + activation_mul = _producer_through_cast(down.inputs[0], "Mul") + if activation_mul is None: + return None + first, second = activation_mul.inputs[:2] + if _producer_through_cast(first, _BLOCK_MATMUL) is not None: + up_output, activation_output = first, second + elif _producer_through_cast(second, _BLOCK_MATMUL) is not None: + up_output, activation_output = second, first + else: + return None + + up = _producer_through_cast(up_output, _BLOCK_MATMUL) + activation = _skip_cast(activation_output).producer() + if activation is None: + return None + if activation.op_type == "Swish": + gate = _producer_through_cast(activation.inputs[0], _BLOCK_MATMUL) + return _ExpertProjections(gate, up, down) if gate is not None else None + if activation.op_type != "Mul": + return None + + for gate_output, sigmoid_output in ( + (activation.inputs[0], activation.inputs[1]), + (activation.inputs[1], activation.inputs[0]), + ): + gate = _producer_through_cast(gate_output, _BLOCK_MATMUL) + sigmoid = _skip_cast(sigmoid_output).producer() + if ( + gate is not None + and sigmoid is not None + and sigmoid.op_type == "Sigmoid" + and _producer_through_cast(sigmoid.inputs[0], _BLOCK_MATMUL) is gate + ): + return _ExpertProjections(gate, up, down) + return None + + +def _find_moe_selectors(graph: ir.Graph) -> list[ir.Value]: + selectors = [] + seen: set[int] = set() + for node in graph: + if node.op_type != "Equal": + continue + selected = node.inputs[0] + if selected is None or _scalar_int(node.inputs[1]) is None: + continue + if len(_consumers_of_type(selected, "Equal")) < 2 or id(selected) in seen: + continue + seen.add(id(selected)) + selectors.append(selected) + return selectors + + +class _DenseBlockMoELayer: + def __init__(self, selected_experts: ir.Value) -> None: + self.selected_experts = selected_experts + self.routing_weights: ir.Value | None = None + self.declared_ids: set[int] = set() + self.experts: dict[int, _ExpertProjections] = {} + self.contributions: list[ir.Value] = [] + self.k: int | None = None + self._collect_experts() + self._resolve_k() + self.routed_out = _find_routed_output(self.contributions) + + def _collect_experts(self) -> None: + for equal in _consumers_of_type(self.selected_experts, "Equal"): + expert_id = _scalar_int(equal.inputs[1]) + if expert_id is None: + continue + self.declared_ids.add(expert_id) + cast = _single_consumer(equal.outputs[0], "CastLike", "Cast") + weight_mul = _single_consumer(cast.outputs[0], "Mul") if cast else None + reduce_sum = _single_consumer(weight_mul.outputs[0], "ReduceSum") if weight_mul else None + contribution = _single_consumer(reduce_sum.outputs[0], "Mul") if reduce_sum else None + if contribution is None: + continue + + routing_weights = weight_mul.inputs[1] if weight_mul.inputs[0] is cast.outputs[0] else weight_mul.inputs[0] + weighted_output = reduce_sum.outputs[0] + expert_output = ( + contribution.inputs[1] if contribution.inputs[0] is weighted_output else contribution.inputs[0] + ) + down = _producer_through_cast(expert_output, _BLOCK_MATMUL) + projections = _trace_block_expert(down) if down is not None else None + if projections is None: + continue + if self.routing_weights is None: + self.routing_weights = routing_weights + elif self.routing_weights is not routing_weights: + continue + self.experts[expert_id] = projections + self.contributions.append(contribution.outputs[0]) + + def _resolve_k(self) -> None: + shape = self.selected_experts.shape + if shape is not None and len(shape) >= 1 and isinstance(shape[-1], int): + self.k = shape[-1] + return + producer = self.selected_experts.producer() + if producer is not None and producer.op_type == "TopK" and len(producer.inputs) > 1: + self.k = _scalar_int(producer.inputs[1]) + + @property + def has_native_experts(self) -> bool: + return bool(self.experts) + + def check_fusable(self) -> None: + if self.routing_weights is None or self.routed_out is None: + raise _UnfusableBlockMoEError("could not trace the router mask or weighted-sum accumulation") + missing = sorted(self.declared_ids - set(self.experts)) + if missing: + raise _UnfusableBlockMoEError( + f"experts {missing} did not trace as native BlockQuantizedMatMul MLPs sharing the router" + ) + expert_ids = sorted(self.experts) + if expert_ids != list(range(len(expert_ids))): + raise _UnfusableBlockMoEError(f"routed expert ids are not contiguous 0..E-1: {expert_ids}") + if self.k is None or self.k < 1: + raise _UnfusableBlockMoEError("could not determine the router top-k") + + +def _projection_format(node: ir.Node, role: str) -> str: + attribute = node.attributes.get("format") + if attribute is None: + raise _UnfusableBlockMoEError(f"{role} projection {node.name!r} has no 'format' attribute") + fmt = attribute.value + if fmt not in _NATIVE_BLOCK_FORMATS: + raise _UnfusableBlockMoEError(f"{role} projection uses unknown native format {fmt!r}") + return fmt + + +def _uniform_format(nodes: list[ir.Node], role: str) -> str: + formats = {_projection_format(node, role) for node in nodes} + if len(formats) != 1: + raise _UnfusableBlockMoEError(f"{role} projections use mixed native formats {sorted(formats)}") + return next(iter(formats)) + + +def _require_no_bias(nodes: list[ir.Node], role: str) -> None: + for node in nodes: + if len(node.inputs) > 2 and node.inputs[2] is not None: + raise _UnfusableBlockMoEError(f"{role} projection {node.name!r} carries an unsupported bias") + + +def _stack_block_weights(nodes: list[ir.Node], out_features: int, role: str) -> np.ndarray: + arrays = [] + for node in nodes: + weight = _array(node.inputs[1], block_moe=True) + if weight.ndim != 3 or weight.shape[0] != out_features: + raise _UnfusableBlockMoEError( + f"{role} projection weight has shape {weight.shape}, expected [{out_features}, n_blocks, block_bytes]" + ) + arrays.append(weight) + shapes = {weight.shape for weight in arrays} + if len(shapes) != 1: + raise _UnfusableBlockMoEError(f"{role} projections have ragged packed shapes {sorted(shapes)}") + return np.stack(arrays, axis=0) + + +def _weight_out_features(node: ir.Node) -> int: + return int(_array(node.inputs[1], block_moe=True).shape[0]) + + +def _flat_last_dim_shape(source: ir.Value, prefix: str, negative_one: ir.Value) -> tuple[list[ir.Node], ir.Value]: + last_dimension = ir.node( + "Shape", + inputs=[source], + attributes={"start": -1}, + num_outputs=1, + name=f"{prefix}.last_dim", + ) + flat_shape = ir.node( + "Concat", + inputs=[negative_one, last_dimension.outputs[0]], + attributes={"axis": 0}, + num_outputs=1, + name=f"{prefix}.flat_shape", + ) + return [last_dimension, flat_shape], flat_shape.outputs[0] + + +def _build_block_routing( + graph: ir.Graph, layer: _DenseBlockMoELayer, experts: int, prefix: str +) -> tuple[list[ir.Node], ir.Value, ir.Value]: + negative_one = _make_initializer(graph, f"{prefix}.neg_one", np.array([-1], dtype=np.int64), ir.DataType.INT64) + expert_dimension = _make_initializer( + graph, f"{prefix}.experts", np.array([experts], dtype=np.int64), ir.DataType.INT64 + ) + zero = _make_initializer(graph, f"{prefix}.zero", np.array(0.0, dtype=np.float32), ir.DataType.FLOAT) + one = _make_initializer(graph, f"{prefix}.one", np.array(1.0, dtype=np.float32), ir.DataType.FLOAT) + nodes: list[ir.Node] = [] + + selected_shape_nodes, selected_flat_shape = _flat_last_dim_shape( + layer.selected_experts, f"{prefix}.sel", negative_one + ) + nodes.extend(selected_shape_nodes) + selected_2d = ir.node( + "Reshape", + inputs=[layer.selected_experts, selected_flat_shape], + num_outputs=1, + name=f"{prefix}.sel2d", + ) + nodes.append(selected_2d) + + weight_shape_nodes, weight_flat_shape = _flat_last_dim_shape(layer.routing_weights, f"{prefix}.rw", negative_one) + nodes.extend(weight_shape_nodes) + weights_2d = ir.node( + "Reshape", + inputs=[layer.routing_weights, weight_flat_shape], + num_outputs=1, + name=f"{prefix}.rw2d", + ) + nodes.append(weights_2d) + if layer.routing_weights.dtype == ir.DataType.FLOAT: + routing_weights = weights_2d.outputs[0] + else: + cast = ir.node( + "Cast", + inputs=[weights_2d.outputs[0]], + attributes={"to": ir.DataType.FLOAT.value}, + num_outputs=1, + name=f"{prefix}.rw_float", + ) + nodes.append(cast) + routing_weights = cast.outputs[0] + + rows = ir.node( + "Shape", + inputs=[selected_2d.outputs[0]], + attributes={"start": 0, "end": 1}, + num_outputs=1, + name=f"{prefix}.rows", + ) + dense_shape = ir.node( + "Concat", + inputs=[rows.outputs[0], expert_dimension], + attributes={"axis": 0}, + num_outputs=1, + name=f"{prefix}.dense_shape", + ) + zeros = ir.node("Expand", inputs=[zero, dense_shape.outputs[0]], num_outputs=1, name=f"{prefix}.zeros") + selected_shape = ir.node("Shape", inputs=[selected_2d.outputs[0]], num_outputs=1, name=f"{prefix}.sel_kshape") + ones = ir.node("Expand", inputs=[one, selected_shape.outputs[0]], num_outputs=1, name=f"{prefix}.ones") + nodes.extend([rows, dense_shape, zeros, selected_shape, ones]) + + logits = ir.node( + "ScatterElements", + inputs=[zeros.outputs[0], selected_2d.outputs[0], ones.outputs[0]], + attributes={"axis": 1}, + num_outputs=1, + name=f"{prefix}.router_logits", + ) + logits.outputs[0].name = f"{prefix}.router_logits" + weights = ir.node( + "ScatterElements", + inputs=[zeros.outputs[0], selected_2d.outputs[0], routing_weights], + attributes={"axis": 1}, + num_outputs=1, + name=f"{prefix}.router_weights", + ) + weights.outputs[0].name = f"{prefix}.router_weights" + nodes.extend([logits, weights]) + return nodes, logits.outputs[0], weights.outputs[0] + + +@dataclasses.dataclass +class _BlockFusionPlan: + layer: _DenseBlockMoELayer + prefix: str + moe_input: ir.Value + fc1_weights: np.ndarray + fc2_weights: np.ndarray + fc3_weights: np.ndarray | None + attributes: dict[str, object] + experts: int + + +def _plan_block_layer( + layer: _DenseBlockMoELayer, index: int, *, allow_per_projection_layout_v2: bool +) -> _BlockFusionPlan: + expert_ids = sorted(layer.experts) + gate_nodes = [layer.experts[expert_id].gate for expert_id in expert_ids] + up_nodes = [layer.experts[expert_id].up for expert_id in expert_ids] + down_nodes = [layer.experts[expert_id].down for expert_id in expert_ids] + + _require_no_bias(gate_nodes, "gate") + _require_no_bias(up_nodes, "up") + _require_no_bias(down_nodes, "down") + gate_format = _uniform_format(gate_nodes, "gate") + up_format = _uniform_format(up_nodes, "up") + down_format = _uniform_format(down_nodes, "down") + + intermediate_size = _weight_out_features(gate_nodes[0]) + hidden_size = _weight_out_features(down_nodes[0]) + if _weight_out_features(up_nodes[0]) != intermediate_size: + raise _UnfusableBlockMoEError("gate and up projections have different output widths") + + gate_weights = _stack_block_weights(gate_nodes, intermediate_size, "gate") + up_weights = _stack_block_weights(up_nodes, intermediate_size, "up") + down_weights = _stack_block_weights(down_nodes, hidden_size, "down") + + fc3_weights = None + projection_formats = {"fc2": down_format} + if gate_format == up_format: + swiglu_fusion = 2 + fc1_weights = np.concatenate([gate_weights, up_weights], axis=1) + projection_formats["fc1"] = gate_format + else: + swiglu_fusion = 0 + fc1_weights = gate_weights + fc3_weights = up_weights + projection_formats["fc1"] = gate_format + projection_formats["fc3"] = up_format + + attributes: dict[str, object] = { + "k": layer.k, + "activation_type": "swiglu", + "normalize_routing_weights": 0, + "swiglu_fusion": swiglu_fusion, + } + formats = set(projection_formats.values()) + if len(formats) == 1: + attributes["format"] = next(iter(formats)) + else: + if not allow_per_projection_layout_v2: + raise _UnfusableBlockMoEError( + "mixed per-projection native formats require block_layout_version=2, " + "which is not enabled for production graph surgery" + ) + attributes.update( + { + "block_layout_version": 2, + "format": projection_formats["fc1"], + "fc1_format": projection_formats["fc1"], + "fc2_format": projection_formats["fc2"], + } + ) + if "fc3" in projection_formats: + attributes["fc3_format"] = projection_formats["fc3"] + + return _BlockFusionPlan( + layer=layer, + prefix=f"bqmoe.layer{index}", + moe_input=gate_nodes[0].inputs[0], + fc1_weights=fc1_weights, + fc2_weights=down_weights, + fc3_weights=fc3_weights, + attributes=attributes, + experts=len(expert_ids), + ) + + +def _emit_block_layer(graph: ir.Graph, plan: _BlockFusionPlan) -> None: + fc1 = _make_initializer(graph, f"{plan.prefix}.fc1_experts_weights", plan.fc1_weights, ir.DataType.UINT8) + fc2 = _make_initializer(graph, f"{plan.prefix}.fc2_experts_weights", plan.fc2_weights, ir.DataType.UINT8) + fc3 = ( + _make_initializer(graph, f"{plan.prefix}.fc3_experts_weights", plan.fc3_weights, ir.DataType.UINT8) + if plan.fc3_weights is not None + else None + ) + routing_nodes, router_logits, router_weights = _build_block_routing(graph, plan.layer, plan.experts, plan.prefix) + moe = ir.node( + _BLOCK_MOE, + inputs=[plan.moe_input, router_logits, fc1, None, fc2, None, fc3, None, router_weights], + attributes=plan.attributes, + domain=_NXRT_DOMAIN, + num_outputs=1, + name=f"{plan.prefix}.bqmoe", + ) + moe.outputs[0].name = f"{plan.prefix}.bqmoe_output" + moe.outputs[0].type = ir.TensorType(ir.DataType.FLOAT) + moe.outputs[0].shape = plan.moe_input.shape + new_nodes = [*routing_nodes, moe] + + routed_dtype = plan.layer.routed_out.dtype + if routed_dtype not in (None, ir.DataType.FLOAT): + cast = ir.node( + "Cast", + inputs=[moe.outputs[0]], + attributes={"to": routed_dtype.value}, + num_outputs=1, + name=f"{plan.prefix}.bqmoe_cast", + ) + cast.outputs[0].name = f"{plan.prefix}.bqmoe_output_cast" + cast.outputs[0].type = ir.TensorType(routed_dtype) + cast.outputs[0].shape = plan.layer.routed_out.shape + new_nodes.append(cast) + final_output = cast.outputs[0] + else: + final_output = moe.outputs[0] + + graph.insert_after(plan.layer.routed_out.producer(), new_nodes) + plan.layer.routed_out.replace_all_uses_with(final_output) + graph.opset_imports[_NXRT_DOMAIN] = 1 + + +def _fail_block_moe_closed( + layer: _DenseBlockMoELayer, reason: _UnfusableBlockMoEError, *, allow_dense_moe: bool +) -> None: + detail = ( + f"routed native-block MoE at {layer.selected_experts.name!r} cannot be expressed as one sparse " + f"pkg.nxrt::BlockQuantizedMoE node: {reason}" + ) + if allow_dense_moe: + logger.warning( + "allow_dense_moe: %s. Keeping the dense BlockQuantizedMatMul fallback; every expert executes.", + detail, + ) + return + raise MoEGraphSurgeryError( + f"Sparse-MoE graph surgery blocker: {detail}. The surgery fails closed instead of emitting or preserving " + "a dense-all-expert performance path. Set allow_dense_moe=True only to retain the runnable dense fallback." + ) from reason + + +def _fuse_block_quantized_moe(model: ir.Model, *, allow_dense_moe: bool, allow_per_projection_layout_v2: bool) -> int: + graph = model.graph + plans: list[_BlockFusionPlan] = [] + for index, layer in enumerate(_DenseBlockMoELayer(value) for value in _find_moe_selectors(graph)): + if not layer.has_native_experts: + continue + try: + layer.check_fusable() + plans.append( + _plan_block_layer( + layer, + index, + allow_per_projection_layout_v2=allow_per_projection_layout_v2, + ) + ) + except _UnfusableBlockMoEError as reason: + _fail_block_moe_closed(layer, reason, allow_dense_moe=allow_dense_moe) + + for plan in plans: + _emit_block_layer(graph, plan) + if plans: + _remove_dead_nodes(graph) + graph.sort() + return len(plans) + + +class FuseBlockQuantizedMoE(Surgeon): + """Fuse eligible dense native-block expert storms into pkg.nxrt::BlockQuantizedMoE.""" + + def __init__(self, allow_dense_moe: bool = False, _allow_perproj_v2_schema: bool = False) -> None: + self.allow_dense_moe = allow_dense_moe + self._allow_perproj_v2_schema = _allow_perproj_v2_schema + + def call_ir(self, model: ir.Model) -> ir.Model: + _fuse_block_quantized_moe( + model, + allow_dense_moe=self.allow_dense_moe, + allow_per_projection_layout_v2=self._allow_perproj_v2_schema, + ) + return model + + +__all__ = ["FuseBlockQuantizedMoE", "FuseDenseMoEToQMoE", "MoEGraphSurgeryError"] diff --git a/test/passes/onnx/test_graph_surgeries_moe.py b/test/passes/onnx/test_graph_surgeries_moe.py new file mode 100644 index 0000000000..d2ff3211d2 --- /dev/null +++ b/test/passes/onnx/test_graph_surgeries_moe.py @@ -0,0 +1,837 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for weight-aware MoE graph surgeries.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import onnx +import pytest +from onnxscript import ir + +from olive.model import ONNXModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.onnx.graph_surgeries import GraphSurgeries +from olive.passes.onnx.graph_surgery import Surgeon +from olive.passes.onnx.graph_surgery import moe as moe_surgeries + +HIDDEN = 32 +INTERMEDIATE = 16 +EXPERTS = 4 +TOP_K = 2 +BLOCK_SIZE = 16 +BITS = 4 + + +def _initializer(graph: ir.Graph, name: str, array: np.ndarray, dtype: ir.DataType) -> ir.Value: + value = ir.Value( + name=name, + shape=ir.Shape(array.shape), + type=ir.TensorType(dtype), + const_value=ir.tensor(array, name=name, dtype=dtype), + ) + graph.register_initializer(value) + return value + + +def _constant_int(nodes: list[ir.Node], name: str, value: int) -> ir.Value: + node = ir.node("Constant", inputs=[], attributes={"value_int": value}, num_outputs=1, name=name) + node.outputs[0].name = f"{name}.out" + nodes.append(node) + return node.outputs[0] + + +def _count(model: ir.Model, op_type: str) -> int: + return sum(node.op_type == op_type for node in model.graph) + + +def _run_surgery( + model: ir.Model, + tmp_path: Path, + surgeon: str, + *, + input_external: bool = False, + output_external: bool = False, + **parameters, +) -> ir.Model: + input_path = tmp_path / "input.onnx" + proto = ir.to_proto(model) + if input_external: + onnx.save_model( + proto, + input_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="input.onnx.data", + size_threshold=0, + ) + else: + onnx.save_model(proto, input_path) + + config = { + "surgeries": [{"surgeon": surgeon, **parameters}], + "remove_duplicate_initializers": False, + "save_as_external_data": output_external, + "size_threshold": 0, + } + graph_surgeries = create_pass_from_dict(GraphSurgeries, config, disable_search=True) + output = graph_surgeries.run(ONNXModelHandler(model_path=str(input_path)), str(tmp_path / "output")) + if output_external: + assert (Path(output.model_path).parent / f"{Path(output.model_path).name}.data").exists() + return ir.from_proto(output.load_model()) + + +class _QuantizedWeight: + def __init__(self, rng: np.random.Generator, out_features: int, in_features: int, dtype: ir.DataType): + blocks = in_features // BLOCK_SIZE + codes = rng.integers(0, 16, size=(out_features, in_features), dtype=np.uint8) + weight = np.zeros((out_features, blocks, BLOCK_SIZE // 2), dtype=np.uint8) + for block in range(blocks): + values = codes[:, block * BLOCK_SIZE : (block + 1) * BLOCK_SIZE] + weight[:, block] = values[:, 0::2] | values[:, 1::2] << 4 + zero_point_codes = rng.integers(0, 16, size=(out_features, blocks), dtype=np.uint8) + zero_points = np.zeros((out_features, (blocks + 1) // 2), dtype=np.uint8) + for block in range(blocks): + zero_points[:, block // 2] |= zero_point_codes[:, block] << (4 * (block % 2)) + self.weight = weight + self.scales = rng.random((out_features, blocks), dtype=np.float32).astype(dtype.numpy()) + self.zero_points = zero_points + self.out_features = out_features + self.in_features = in_features + self.dtype = dtype + + +def _dequantize(weight: np.ndarray, scales: np.ndarray, zero_points: np.ndarray) -> np.ndarray: + out_features = weight.shape[0] + packed_block_size = BLOCK_SIZE // 2 + blocks = weight.shape[1] // packed_block_size if weight.ndim == 2 else weight.shape[1] + weight = weight.reshape(out_features, blocks, packed_block_size) + dense = np.empty((out_features, blocks * packed_block_size * 2), dtype=np.float32) + for block in range(blocks): + packed = weight[:, block].astype(np.int32) + codes = np.empty((out_features, BLOCK_SIZE), dtype=np.int32) + codes[:, 0::2] = packed & 0xF + codes[:, 1::2] = packed >> 4 + zero_point = (zero_points[:, block // 2].astype(np.int32) >> (4 * (block % 2))) & 0xF + dense[:, block * BLOCK_SIZE : (block + 1) * BLOCK_SIZE] = (codes - zero_point[:, None]) * scales[ + :, block : block + 1 + ].astype(np.float32) + return dense + + +def _matmul_nbits( + graph: ir.Graph, + nodes: list[ir.Node], + name: str, + value: ir.Value, + weight: _QuantizedWeight, + *, + bits: int = BITS, + block_size: int = BLOCK_SIZE, + include_zero_points: bool = True, +) -> ir.Value: + packed = _initializer(graph, f"{name}.weight", weight.weight, ir.DataType.UINT8) + scales = _initializer(graph, f"{name}.scales", weight.scales, weight.dtype) + zero_points = _initializer(graph, f"{name}.zero_points", weight.zero_points, ir.DataType.UINT8) + inputs = [value, packed, scales] + if include_zero_points: + inputs.append(zero_points) + node = ir.node( + "MatMulNBits", + inputs=inputs, + attributes={ + "K": weight.in_features, + "N": weight.out_features, + "bits": bits, + "block_size": block_size, + }, + domain="com.microsoft", + num_outputs=1, + name=name, + ) + node.outputs[0].name = f"{name}.out" + nodes.append(node) + return node.outputs[0] + + +def _build_qmoe_graph( + *, + activation_dtype: ir.DataType = ir.DataType.FLOAT16, + hidden_typed: bool = True, + legacy_activation: bool = False, + router_cast: bool = False, + topk_value_cast: bool = False, + block_size: int = BLOCK_SIZE, + include_zero_points: bool = True, + mask_cast_op: str = "CastLike", +) -> tuple[ir.Model, dict[str, _QuantizedWeight]]: + rng = np.random.default_rng(0) + weights: dict[str, _QuantizedWeight] = {} + nodes: list[ir.Node] = [] + hidden = ir.Value( + name="hidden", + shape=ir.Shape(["tokens", HIDDEN]), + type=ir.TensorType(activation_dtype) if hidden_typed else None, + ) + graph = ir.Graph([hidden], [], nodes=[], name="dense_qmoe") + + def quant(name: str, out_features: int, in_features: int) -> _QuantizedWeight: + weights[name] = _QuantizedWeight(rng, out_features, in_features, activation_dtype) + return weights[name] + + router = _matmul_nbits( + graph, + nodes, + "router", + hidden, + quant("router", EXPERTS, HIDDEN), + block_size=block_size, + include_zero_points=include_zero_points, + ) + logits = router + if router_cast: + cast = ir.node( + "Cast", + inputs=[router], + attributes={"to": activation_dtype.value}, + num_outputs=1, + name="router.cast", + ) + cast.outputs[0].name = "router.cast.out" + nodes.append(cast) + logits = cast.outputs[0] + + topk_k = _initializer(graph, "topk.k", np.array([TOP_K], dtype=np.int64), ir.DataType.INT64) + topk = ir.node("TopK", inputs=[logits, topk_k], attributes={"axis": -1}, num_outputs=2, name="topk") + topk.outputs[0].name = "topk.values" + topk.outputs[1].name = "topk.indices" + nodes.append(topk) + softmax_input = topk.outputs[0] + if topk_value_cast: + cast = ir.node( + "Cast", + inputs=[softmax_input], + attributes={"to": activation_dtype.value}, + num_outputs=1, + name="topk.cast", + ) + cast.outputs[0].name = "topk.cast.out" + nodes.append(cast) + softmax_input = cast.outputs[0] + softmax = ir.node("Softmax", inputs=[softmax_input], attributes={"axis": -1}, num_outputs=1, name="softmax") + softmax.outputs[0].name = "softmax.out" + nodes.append(softmax) + axes = _initializer(graph, "axes", np.array([-1], dtype=np.int64), ir.DataType.INT64) + + routed = None + for expert in range(EXPERTS): + expert_id = _constant_int(nodes, f"expert_id_{expert}", expert) + equal = ir.node("Equal", inputs=[topk.outputs[1], expert_id], num_outputs=1, name=f"equal_{expert}") + nodes.append(equal) + if mask_cast_op == "CastLike": + cast = ir.node( + "CastLike", + inputs=[equal.outputs[0], softmax.outputs[0]], + num_outputs=1, + name=f"mask_cast_{expert}", + ) + else: + cast = ir.node( + "Cast", + inputs=[equal.outputs[0]], + attributes={"to": activation_dtype.value}, + num_outputs=1, + name=f"mask_cast_{expert}", + ) + nodes.append(cast) + weighted_mask = ir.node( + "Mul", + inputs=[softmax.outputs[0], cast.outputs[0]], + num_outputs=1, + name=f"weighted_mask_{expert}", + ) + nodes.append(weighted_mask) + reduce_sum = ir.node( + "ReduceSum", + inputs=[weighted_mask.outputs[0], axes], + attributes={"keepdims": 1}, + num_outputs=1, + name=f"reduce_{expert}", + ) + nodes.append(reduce_sum) + + gate = _matmul_nbits( + graph, + nodes, + f"expert{expert}.gate", + hidden, + quant(f"gate{expert}", INTERMEDIATE, HIDDEN), + include_zero_points=include_zero_points, + ) + if legacy_activation: + sigmoid = ir.node("Sigmoid", inputs=[gate], num_outputs=1, name=f"expert{expert}.sigmoid") + nodes.append(sigmoid) + activation = ir.node("Mul", inputs=[gate, sigmoid.outputs[0]], num_outputs=1, name=f"expert{expert}.silu") + else: + activation = ir.node("Swish", inputs=[gate], num_outputs=1, name=f"expert{expert}.silu") + nodes.append(activation) + up = _matmul_nbits( + graph, + nodes, + f"expert{expert}.up", + hidden, + quant(f"up{expert}", INTERMEDIATE, HIDDEN), + include_zero_points=include_zero_points, + ) + product = ir.node("Mul", inputs=[activation.outputs[0], up], num_outputs=1, name=f"expert{expert}.product") + nodes.append(product) + down = _matmul_nbits( + graph, + nodes, + f"expert{expert}.down", + product.outputs[0], + quant(f"down{expert}", HIDDEN, INTERMEDIATE), + block_size=block_size, + include_zero_points=include_zero_points, + ) + contribution = ir.node( + "Mul", inputs=[down, reduce_sum.outputs[0]], num_outputs=1, name=f"expert{expert}.contribution" + ) + nodes.append(contribution) + if routed is None: + routed = contribution.outputs[0] + else: + add = ir.node("Add", inputs=[routed, contribution.outputs[0]], num_outputs=1, name=f"sum_{expert}") + nodes.append(add) + routed = add.outputs[0] + + shared = ir.node("Identity", inputs=[hidden], num_outputs=1, name="shared_expert") + nodes.append(shared) + final = ir.node("Add", inputs=[routed, shared.outputs[0]], num_outputs=1, name="final_add") + final.outputs[0].name = "output" + final.outputs[0].shape = hidden.shape + final.outputs[0].type = ir.TensorType(activation_dtype) + nodes.append(final) + for node in nodes: + graph.append(node) + graph.outputs.append(final.outputs[0]) + model = ir.Model(graph, ir_version=10, producer_name="test") + model.opset_imports[""] = 21 + model.opset_imports["com.microsoft"] = 1 + return model, weights + + +def _make_native_weight(rng: np.random.Generator, out_features: int, in_features: int, fmt: str) -> np.ndarray: + block_elements, block_bytes = moe_surgeries._NATIVE_BLOCK_FORMATS[fmt] + blocks = (in_features + block_elements - 1) // block_elements + return rng.integers(0, 256, size=(out_features, blocks, block_bytes), dtype=np.uint8) + + +def _block_matmul( + graph: ir.Graph, + nodes: list[ir.Node], + name: str, + value: ir.Value, + weight: np.ndarray, + fmt: str, + *, + bias: bool = False, +) -> ir.Value: + packed = _initializer(graph, f"{name}.weight", weight, ir.DataType.UINT8) + inputs = [value, packed] + if bias: + inputs.append(_initializer(graph, f"{name}.bias", np.zeros(weight.shape[0], np.float32), ir.DataType.FLOAT)) + node = ir.node( + "BlockQuantizedMatMul", + inputs=inputs, + attributes={"K": HIDDEN, "N": weight.shape[0], "format": fmt, "block_layout_version": 1}, + domain="pkg.nxrt", + num_outputs=1, + name=name, + ) + node.outputs[0].name = f"{name}.out" + node.outputs[0].type = ir.TensorType(ir.DataType.FLOAT) + nodes.append(node) + return node.outputs[0] + + +def _build_block_moe_graph( + *, + gate_format: str = "iq4_xs", + up_format: str = "iq4_xs", + down_format: str = "iq4_xs", + dtype: ir.DataType = ir.DataType.FLOAT, + legacy_activation: bool = False, + corrupt_gate_format: tuple[int, str] | None = None, + biased_expert: int | None = None, + broken_expert: int | None = None, + alien_routing_expert: int | None = None, + mask_cast_op: str = "CastLike", +) -> tuple[ir.Model, dict[str, np.ndarray]]: + rng = np.random.default_rng(1) + weights: dict[str, np.ndarray] = {} + nodes: list[ir.Node] = [] + hidden = ir.Value(name="hidden", shape=ir.Shape(["tokens", HIDDEN]), type=ir.TensorType(dtype)) + graph = ir.Graph([hidden], [], nodes=[], name="dense_block_moe") + + topk_k = _initializer(graph, "topk.k", np.array([TOP_K], dtype=np.int64), ir.DataType.INT64) + router_weight = _initializer( + graph, "router.weight", rng.standard_normal((HIDDEN, EXPERTS)).astype(np.float32), ir.DataType.FLOAT + ) + router = ir.node("MatMul", inputs=[hidden, router_weight], num_outputs=1, name="router") + nodes.append(router) + sigmoid = ir.node("Sigmoid", inputs=[router.outputs[0]], num_outputs=1, name="router.sigmoid") + nodes.append(sigmoid) + topk = ir.node("TopK", inputs=[sigmoid.outputs[0], topk_k], attributes={"axis": -1}, num_outputs=2, name="topk") + topk.outputs[0].shape = ir.Shape(["tokens", TOP_K]) + topk.outputs[1].shape = ir.Shape(["tokens", TOP_K]) + nodes.append(topk) + axes = _initializer(graph, "axes", np.array([-1], dtype=np.int64), ir.DataType.INT64) + total = ir.node( + "ReduceSum", + inputs=[topk.outputs[0], axes], + attributes={"keepdims": 1}, + num_outputs=1, + name="routing.total", + ) + nodes.append(total) + routing_weights = ir.node("Div", inputs=[topk.outputs[0], total.outputs[0]], num_outputs=1, name="routing") + routing_weights.outputs[0].shape = ir.Shape(["tokens", TOP_K]) + routing_weights.outputs[0].type = ir.TensorType(dtype) + nodes.append(routing_weights) + + def cast_in(value: ir.Value, name: str) -> ir.Value: + if dtype == ir.DataType.FLOAT: + return value + cast = ir.node( + "Cast", + inputs=[value], + attributes={"to": ir.DataType.FLOAT.value}, + num_outputs=1, + name=f"{name}.cast_in", + ) + nodes.append(cast) + return cast.outputs[0] + + def cast_out(value: ir.Value, name: str) -> ir.Value: + if dtype == ir.DataType.FLOAT: + return value + cast = ir.node("Cast", inputs=[value], attributes={"to": dtype.value}, num_outputs=1, name=f"{name}.cast_out") + cast.outputs[0].type = ir.TensorType(dtype) + nodes.append(cast) + return cast.outputs[0] + + routed = None + for expert in range(EXPERTS): + expert_id = _constant_int(nodes, f"expert_id_{expert}", expert) + equal = ir.node("Equal", inputs=[topk.outputs[1], expert_id], num_outputs=1, name=f"equal_{expert}") + nodes.append(equal) + if mask_cast_op == "CastLike": + cast = ir.node( + "CastLike", + inputs=[equal.outputs[0], routing_weights.outputs[0]], + num_outputs=1, + name=f"mask_cast_{expert}", + ) + else: + cast = ir.node( + "Cast", + inputs=[equal.outputs[0]], + attributes={"to": dtype.value}, + num_outputs=1, + name=f"mask_cast_{expert}", + ) + nodes.append(cast) + expert_routing_weights = routing_weights.outputs[0] + if alien_routing_expert == expert: + alien = ir.node( + "Identity", inputs=[routing_weights.outputs[0]], num_outputs=1, name=f"alien_routing_{expert}" + ) + nodes.append(alien) + expert_routing_weights = alien.outputs[0] + weighted_mask = ir.node( + "Mul", inputs=[expert_routing_weights, cast.outputs[0]], num_outputs=1, name=f"weighted_mask_{expert}" + ) + nodes.append(weighted_mask) + reduce_sum = ir.node( + "ReduceSum", + inputs=[weighted_mask.outputs[0], axes], + attributes={"keepdims": 1}, + num_outputs=1, + name=f"reduce_{expert}", + ) + nodes.append(reduce_sum) + + expert_gate_format = ( + corrupt_gate_format[1] + if corrupt_gate_format is not None and corrupt_gate_format[0] == expert + else gate_format + ) + gate_weight = _make_native_weight(rng, INTERMEDIATE, HIDDEN, expert_gate_format) + up_weight = _make_native_weight(rng, INTERMEDIATE, HIDDEN, up_format) + down_weight = _make_native_weight(rng, HIDDEN, INTERMEDIATE, down_format) + weights[f"gate{expert}"] = gate_weight + weights[f"up{expert}"] = up_weight + weights[f"down{expert}"] = down_weight + gate = _block_matmul( + graph, + nodes, + f"expert{expert}.gate", + cast_in(hidden, f"expert{expert}.gate"), + gate_weight, + expert_gate_format, + ) + gate = cast_out(gate, f"expert{expert}.gate") + if legacy_activation: + activation_sigmoid = ir.node("Sigmoid", inputs=[gate], num_outputs=1, name=f"expert{expert}.sigmoid") + nodes.append(activation_sigmoid) + activation = ir.node( + "Mul", + inputs=[gate, activation_sigmoid.outputs[0]], + num_outputs=1, + name=f"expert{expert}.silu", + ) + else: + activation = ir.node("Swish", inputs=[gate], num_outputs=1, name=f"expert{expert}.silu") + nodes.append(activation) + up = _block_matmul( + graph, nodes, f"expert{expert}.up", cast_in(hidden, f"expert{expert}.up"), up_weight, up_format + ) + up = cast_out(up, f"expert{expert}.up") + product = ir.node("Mul", inputs=[activation.outputs[0], up], num_outputs=1, name=f"expert{expert}.product") + nodes.append(product) + if broken_expert == expert: + down = ir.node("Identity", inputs=[cast_in(product.outputs[0], f"expert{expert}.down")], num_outputs=1) + nodes.append(down) + down_output = cast_out(down.outputs[0], f"expert{expert}.down") + else: + down_output = _block_matmul( + graph, + nodes, + f"expert{expert}.down", + cast_in(product.outputs[0], f"expert{expert}.down"), + down_weight, + down_format, + bias=biased_expert == expert, + ) + down_output = cast_out(down_output, f"expert{expert}.down") + contribution = ir.node( + "Mul", + inputs=[down_output, reduce_sum.outputs[0]], + num_outputs=1, + name=f"expert{expert}.contribution", + ) + contribution.outputs[0].type = ir.TensorType(dtype) + nodes.append(contribution) + if routed is None: + routed = contribution.outputs[0] + else: + add = ir.node("Add", inputs=[routed, contribution.outputs[0]], num_outputs=1, name=f"sum_{expert}") + add.outputs[0].type = ir.TensorType(dtype) + nodes.append(add) + routed = add.outputs[0] + + shared = ir.node("Identity", inputs=[hidden], num_outputs=1, name="shared_expert") + nodes.append(shared) + final = ir.node("Add", inputs=[routed, shared.outputs[0]], num_outputs=1, name="final_add") + final.outputs[0].name = "output" + final.outputs[0].shape = hidden.shape + final.outputs[0].type = ir.TensorType(dtype) + nodes.append(final) + for node in nodes: + graph.append(node) + graph.outputs.append(final.outputs[0]) + model = ir.Model(graph, ir_version=10, producer_name="test") + model.opset_imports[""] = 21 + model.opset_imports["pkg.nxrt"] = 1 + return model, weights + + +def test_moe_surgeries_register_on_explicit_module_import(): + assert Surgeon.registry["fusedensemoetoqmoe"] is moe_surgeries.FuseDenseMoEToQMoE + assert Surgeon.registry["fuseblockquantizedmoe"] is moe_surgeries.FuseBlockQuantizedMoE + + +def test_fuse_dense_moe_to_qmoe_preserves_bytes_attributes_and_shared_output(tmp_path): + model, weights = _build_qmoe_graph() + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + assert _count(rewritten, "QMoE") == 1 + assert _count(rewritten, "MatMulNBits") == 1 + assert _count(rewritten, "Equal") == 0 + qmoe = next(node for node in rewritten.graph if node.op_type == "QMoE") + assert qmoe.domain == "com.microsoft" + assert len(qmoe.inputs) == 15 + assert qmoe.inputs[4] is None + assert qmoe.inputs[7] is None + assert [qmoe.attributes[name].value for name in ("k", "expert_weight_bits", "block_size")] == [ + TOP_K, + BITS, + BLOCK_SIZE, + ] + assert qmoe.attributes["activation_type"].value == "swiglu" + assert qmoe.attributes["normalize_routing_weights"].value == 1 + assert qmoe.attributes["swiglu_fusion"].value == 2 + assert qmoe.attributes["quant_type"].value == "int" + assert qmoe.attributes["weights_prepacked"].value == 0 + fc1_weights = qmoe.inputs[2].const_value.numpy() + fc2_weights = qmoe.inputs[5].const_value.numpy() + fc1_zero_points = qmoe.inputs[11].const_value.numpy() + fc2_zero_points = qmoe.inputs[12].const_value.numpy() + for expert in range(EXPERTS): + np.testing.assert_array_equal( + fc1_weights[expert], + np.concatenate( + [ + weights[f"gate{expert}"].weight.reshape(INTERMEDIATE, -1), + weights[f"up{expert}"].weight.reshape(INTERMEDIATE, -1), + ] + ), + ) + np.testing.assert_array_equal(fc2_weights[expert], weights[f"down{expert}"].weight.reshape(HIDDEN, -1)) + np.testing.assert_array_equal( + fc1_zero_points[expert], + np.concatenate([weights[f"gate{expert}"].zero_points, weights[f"up{expert}"].zero_points]), + ) + np.testing.assert_array_equal(fc2_zero_points[expert], weights[f"down{expert}"].zero_points) + final_add = next(node for node in rewritten.graph if node.name == "final_add") + assert any(value.producer().op_type == "QMoE" for value in final_add.inputs) + assert any(value.producer().name == "shared_expert" for value in final_add.inputs) + + +def test_fuse_dense_moe_to_qmoe_preserves_dense_forward_semantics(tmp_path): + model, weights = _build_qmoe_graph(activation_dtype=ir.DataType.FLOAT) + hidden = np.random.default_rng(2).standard_normal((5, HIDDEN)).astype(np.float32) + + router_logits = ( + hidden + @ _dequantize( + weights["router"].weight, + weights["router"].scales, + weights["router"].zero_points, + ).T + ) + selected = np.argsort(-router_logits, axis=-1, kind="stable")[:, :TOP_K] + selected_logits = np.take_along_axis(router_logits, selected, axis=-1) + routing_weights = np.exp(selected_logits - selected_logits.max(axis=-1, keepdims=True)) + routing_weights /= routing_weights.sum(axis=-1, keepdims=True) + expected = hidden.copy() + for expert in range(EXPERTS): + gate = ( + hidden + @ _dequantize( + weights[f"gate{expert}"].weight, + weights[f"gate{expert}"].scales, + weights[f"gate{expert}"].zero_points, + ).T + ) + up = ( + hidden + @ _dequantize( + weights[f"up{expert}"].weight, + weights[f"up{expert}"].scales, + weights[f"up{expert}"].zero_points, + ).T + ) + activated = gate / (1 + np.exp(-gate)) * up + down = ( + activated + @ _dequantize( + weights[f"down{expert}"].weight, + weights[f"down{expert}"].scales, + weights[f"down{expert}"].zero_points, + ).T + ) + expected += down * np.where(selected == expert, routing_weights, 0).sum(axis=-1, keepdims=True) + + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + qmoe = next(node for node in rewritten.graph if node.op_type == "QMoE") + actual = hidden.copy() + for expert in range(EXPERTS): + fc1 = _dequantize( + qmoe.inputs[2].const_value.numpy()[expert], + qmoe.inputs[3].const_value.numpy()[expert], + qmoe.inputs[11].const_value.numpy()[expert], + ) + gate = hidden @ fc1[:INTERMEDIATE].T + up = hidden @ fc1[INTERMEDIATE:].T + activated = gate / (1 + np.exp(-gate)) * up + fc2 = _dequantize( + qmoe.inputs[5].const_value.numpy()[expert], + qmoe.inputs[6].const_value.numpy()[expert], + qmoe.inputs[12].const_value.numpy()[expert], + ) + down = activated @ fc2.T + actual += down * np.where(selected == expert, routing_weights, 0).sum(axis=-1, keepdims=True) + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +def test_fuse_dense_moe_to_qmoe_casts_scales_and_router_to_activation_dtype(tmp_path, dtype): + model, weights = _build_qmoe_graph(activation_dtype=dtype) + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + qmoe = next(node for node in rewritten.graph if node.op_type == "QMoE") + assert qmoe.inputs[3].dtype == dtype + assert qmoe.inputs[6].dtype == dtype + assert qmoe.inputs[1].producer().attributes["to"].value == dtype.value + expected = np.concatenate( + [weights["gate0"].scales.reshape(INTERMEDIATE, -1), weights["up0"].scales.reshape(INTERMEDIATE, -1)] + ).astype(dtype.numpy()) + np.testing.assert_array_equal(qmoe.inputs[3].const_value.numpy()[0], expected) + + +def test_fuse_dense_moe_to_qmoe_matches_all_cast_and_activation_variants(tmp_path): + model, _ = _build_qmoe_graph( + legacy_activation=True, + router_cast=True, + topk_value_cast=True, + mask_cast_op="Cast", + ) + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + assert _count(rewritten, "QMoE") == 1 + + +def test_fuse_dense_moe_to_qmoe_supports_symmetric_weights_without_zero_points(tmp_path): + model, _ = _build_qmoe_graph(include_zero_points=False) + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + qmoe = next(node for node in rewritten.graph if node.op_type == "QMoE") + assert qmoe.inputs[11] is None + assert qmoe.inputs[12] is None + + +def test_fuse_dense_moe_to_qmoe_uses_scale_dtype_when_hidden_is_untyped(tmp_path): + model, _ = _build_qmoe_graph(hidden_typed=False) + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + qmoe = next(node for node in rewritten.graph if node.op_type == "QMoE") + assert qmoe.inputs[3].dtype == ir.DataType.FLOAT16 + + +def test_fuse_dense_moe_to_qmoe_keeps_unsupported_geometry(tmp_path): + model, _ = _build_qmoe_graph(block_size=24) + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + assert _count(rewritten, "QMoE") == 0 + assert _count(rewritten, "TopK") == 1 + + +def test_fuse_dense_moe_to_qmoe_reads_and_writes_external_initializers(tmp_path): + model, weights = _build_qmoe_graph() + rewritten = _run_surgery( + model, + tmp_path, + "FuseDenseMoEToQMoE", + input_external=True, + output_external=True, + ) + qmoe = next(node for node in rewritten.graph if node.op_type == "QMoE") + np.testing.assert_array_equal( + qmoe.inputs[2].const_value.numpy()[0, :INTERMEDIATE], + weights["gate0"].weight.reshape(INTERMEDIATE, -1), + ) + + +def test_fuse_block_quantized_moe_preserves_native_bytes_and_v1_layout(tmp_path): + model, weights = _build_block_moe_graph() + rewritten = _run_surgery(model, tmp_path, "FuseBlockQuantizedMoE") + assert _count(rewritten, "BlockQuantizedMoE") == 1 + assert _count(rewritten, "BlockQuantizedMatMul") == 0 + assert _count(rewritten, "ScatterElements") == 2 + moe = next(node for node in rewritten.graph if node.op_type == "BlockQuantizedMoE") + assert moe.domain == "pkg.nxrt" + assert len(moe.inputs) == 9 + assert moe.inputs[3] is None + assert moe.inputs[5] is None + assert moe.inputs[6] is None + assert moe.inputs[7] is None + assert moe.attributes["format"].value == "iq4_xs" + assert "block_layout_version" not in moe.attributes + assert moe.attributes["k"].value == TOP_K + assert moe.attributes["normalize_routing_weights"].value == 0 + assert moe.attributes["swiglu_fusion"].value == 2 + for expert in range(EXPERTS): + np.testing.assert_array_equal( + moe.inputs[2].const_value.numpy()[expert], + np.concatenate([weights[f"gate{expert}"], weights[f"up{expert}"]], axis=0), + ) + np.testing.assert_array_equal(moe.inputs[4].const_value.numpy()[expert], weights[f"down{expert}"]) + final_add = next(node for node in rewritten.graph if node.name == "final_add") + assert any(value.producer().op_type == "BlockQuantizedMoE" for value in final_add.inputs) + + +def test_fuse_block_quantized_moe_supports_unfused_swiglu_layout_v2_schema(tmp_path): + model, weights = _build_block_moe_graph( + gate_format="iq1_s", up_format="iq3_xxs", down_format="iq4_xs", legacy_activation=True, mask_cast_op="Cast" + ) + rewritten = _run_surgery( + model, + tmp_path, + "FuseBlockQuantizedMoE", + _allow_perproj_v2_schema=True, + ) + moe = next(node for node in rewritten.graph if node.op_type == "BlockQuantizedMoE") + assert moe.attributes["block_layout_version"].value == 2 + assert moe.attributes["fc1_format"].value == "iq1_s" + assert moe.attributes["fc2_format"].value == "iq4_xs" + assert moe.attributes["fc3_format"].value == "iq3_xxs" + assert moe.attributes["swiglu_fusion"].value == 0 + for expert in range(EXPERTS): + np.testing.assert_array_equal(moe.inputs[2].const_value.numpy()[expert], weights[f"gate{expert}"]) + np.testing.assert_array_equal(moe.inputs[4].const_value.numpy()[expert], weights[f"down{expert}"]) + np.testing.assert_array_equal(moe.inputs[6].const_value.numpy()[expert], weights[f"up{expert}"]) + + +def test_fuse_block_quantized_moe_restores_non_float_output_dtype(tmp_path): + model, _ = _build_block_moe_graph(dtype=ir.DataType.FLOAT16) + rewritten = _run_surgery(model, tmp_path, "FuseBlockQuantizedMoE") + moe = next(node for node in rewritten.graph if node.op_type == "BlockQuantizedMoE") + consumer = next(node for node, _ in moe.outputs[0].uses()) + assert consumer.op_type == "Cast" + assert consumer.attributes["to"].value == ir.DataType.FLOAT16.value + + +@pytest.mark.parametrize( + ("builder_parameters", "message"), + [ + ({"gate_format": "iq1_s", "up_format": "iq1_s", "down_format": "iq4_xs"}, "block_layout_version=2"), + ({"corrupt_gate_format": (1, "iq1_m")}, "mixed native formats"), + ({"biased_expert": 1}, "unsupported bias"), + ({"broken_expert": EXPERTS - 1}, "did not trace"), + ({"alien_routing_expert": EXPERTS - 1}, "did not trace"), + ], +) +def test_fuse_block_quantized_moe_fails_closed_atomically(tmp_path, builder_parameters, message): + model, _ = _build_block_moe_graph(**builder_parameters) + input_path = tmp_path / "input.onnx" + onnx.save_model(ir.to_proto(model), input_path) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + { + "surgeries": [{"surgeon": "FuseBlockQuantizedMoE"}], + "remove_duplicate_initializers": False, + }, + disable_search=True, + ) + with pytest.raises(moe_surgeries.MoEGraphSurgeryError, match=message): + graph_surgeries.run(ONNXModelHandler(model_path=str(input_path)), str(tmp_path / "output")) + unchanged = ir.from_proto(onnx.load(input_path)) + assert _count(unchanged, "BlockQuantizedMoE") == 0 + assert _count(unchanged, "Equal") == EXPERTS + + +def test_fuse_block_quantized_moe_allow_dense_opt_in_preserves_fallback(tmp_path): + model, _ = _build_block_moe_graph(corrupt_gate_format=(1, "iq1_m")) + rewritten = _run_surgery(model, tmp_path, "FuseBlockQuantizedMoE", allow_dense_moe=True) + assert _count(rewritten, "BlockQuantizedMoE") == 0 + assert _count(rewritten, "Equal") == EXPERTS + + +def test_fuse_block_quantized_moe_reads_external_initializers(tmp_path): + model, weights = _build_block_moe_graph() + rewritten = _run_surgery(model, tmp_path, "FuseBlockQuantizedMoE", input_external=True) + moe = next(node for node in rewritten.graph if node.op_type == "BlockQuantizedMoE") + np.testing.assert_array_equal( + moe.inputs[2].const_value.numpy()[0, :INTERMEDIATE], + weights["gate0"], + ) From 99e8b09d4476b7a7a7815c67c848051c5771a956 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:18:38 -0700 Subject: [PATCH 05/14] Harden MoE graph surgery semantics Fail QMoE fusion closed for incomplete expert groups and inconsistent quantization metadata, support independent FC1 and FC2 zero-point banks, and replace direct graph outputs for both MoE surgeries. Add focused regression coverage for each review finding. Signed-off-by: Xiaoyu Zhang --- olive/passes/onnx/graph_surgery/moe.py | 205 +++++++++++++------ test/passes/onnx/test_graph_surgeries_moe.py | 180 +++++++++++++--- 2 files changed, 300 insertions(+), 85 deletions(-) diff --git a/olive/passes/onnx/graph_surgery/moe.py b/olive/passes/onnx/graph_surgery/moe.py index 1c273ab804..c672b12ef1 100644 --- a/olive/passes/onnx/graph_surgery/moe.py +++ b/olive/passes/onnx/graph_surgery/moe.py @@ -43,6 +43,10 @@ class _UnfusableBlockMoEError(Exception): """A native-block MoE layer cannot be represented by one fused node.""" +class _UnfusableQMoEError(Exception): + """A MatMulNBits MoE layer cannot be represented by one fused node.""" + + def _scalar_int(value: ir.Value | None) -> int | None: if value is None: return None @@ -226,6 +230,8 @@ def __init__(self, topk: ir.Node) -> None: if softmax is not None and softmax.op_type == "Cast": softmax = _single_consumer(softmax.outputs[0], "Softmax") self.softmax = softmax + self.declared_ids: set[int] = set() + self.declared_branch_counts: dict[int, int] = {} self.experts: dict[int, _ExpertProjections] = {} self.contributions: list[ir.Value] = [] self._collect_experts() @@ -236,6 +242,8 @@ def _collect_experts(self) -> None: expert_id = _scalar_int(equal.inputs[1]) if expert_id is None: continue + self.declared_ids.add(expert_id) + self.declared_branch_counts[expert_id] = self.declared_branch_counts.get(expert_id, 0) + 1 cast = _single_consumer(equal.outputs[0], "CastLike", "Cast") if cast is None: continue @@ -254,26 +262,55 @@ def _collect_experts(self) -> None: self.experts[expert_id] = projections self.contributions.append(contribution.outputs[0]) - @property - def is_valid(self) -> bool: + def check_fusable(self) -> None: if self.k is None or self.softmax is None or self.routed_out is None: - return False + raise _UnfusableQMoEError("could not trace the router or weighted-sum accumulation") + missing = sorted(self.declared_ids - set(self.experts)) + if missing: + raise _UnfusableQMoEError(f"experts {missing} did not trace as MatMulNBits MLPs") + duplicate_ids = sorted(expert_id for expert_id, count in self.declared_branch_counts.items() if count != 1) + if duplicate_ids: + raise _UnfusableQMoEError(f"experts {duplicate_ids} have duplicate routed branches") expert_ids = sorted(self.experts) - return expert_ids == list(range(len(expert_ids))) and bool(expert_ids) - - -def _qmoe_geometry(node: ir.Node) -> tuple[int, int, bool]: - return ( - int(node.attributes["bits"].value), - int(node.attributes["block_size"].value), - len(node.inputs) > 3 and node.inputs[3] is not None, - ) + if not expert_ids or expert_ids != list(range(len(expert_ids))): + raise _UnfusableQMoEError(f"routed expert ids are not contiguous 0..E-1: {expert_ids}") + + expected_counts = [] + output_width = self.gate_router.attributes.get("N") + if output_width is not None: + expected_counts.append(int(output_width.value)) + router_weight = self.gate_router.inputs[1] + if router_weight is not None and router_weight.const_value is not None: + expected_counts.append(int(_array(router_weight).shape[0])) + if expected_counts and len(set(expected_counts)) != 1: + raise _UnfusableQMoEError(f"router expert dimensions disagree: {expected_counts}") + if expected_counts and expert_ids != list(range(expected_counts[0])): + raise _UnfusableQMoEError( + f"router declares {expected_counts[0]} experts, but routed branches declare {expert_ids}" + ) def _qmoe_abi_supported(bits: int, block_size: int) -> bool: return bits == 4 and block_size >= 16 and not block_size & (block_size - 1) +def _uniform_qmoe_geometry(nodes: list[ir.Node]) -> tuple[int, int]: + try: + geometries = {(int(node.attributes["bits"].value), int(node.attributes["block_size"].value)) for node in nodes} + except KeyError as error: + raise _UnfusableQMoEError(f"MatMulNBits projection is missing {error.args[0]!r}") from error + if len(geometries) != 1: + raise _UnfusableQMoEError(f"MatMulNBits projections use mismatched quantization geometry: {sorted(geometries)}") + return next(iter(geometries)) + + +def _consistent_zero_points(nodes: list[ir.Node], bank: str) -> bool: + presence = {len(node.inputs) > 3 and node.inputs[3] is not None for node in nodes} + if len(presence) != 1: + raise _UnfusableQMoEError(f"{bank} projections disagree on zero-point presence") + return next(iter(presence)) + + def _pack_qmoe_projection(nodes: list[ir.Node], slot: int) -> np.ndarray: arrays = [_array(node.inputs[slot]) for node in nodes] return np.stack([array.reshape(array.shape[0], -1) for array in arrays], axis=0) @@ -292,57 +329,96 @@ def _pack_qmoe_zero_points(nodes: list[ir.Node], slot: int, out_features: int) - ) -def _emit_qmoe(graph: ir.Graph, layer: _DenseQMoELayer, index: int, activation_dtype: ir.DataType) -> None: +@dataclasses.dataclass +class _QMoEFusionPlan: + layer: _DenseQMoELayer + prefix: str + activation_dtype: ir.DataType + bits: int + block_size: int + fc1_weights: np.ndarray + fc1_scales: np.ndarray + fc1_zero_points: np.ndarray | None + fc2_weights: np.ndarray + fc2_scales: np.ndarray + fc2_zero_points: np.ndarray | None + + +def _plan_qmoe_layer(layer: _DenseQMoELayer, index: int, activation_dtype: ir.DataType) -> _QMoEFusionPlan: expert_ids = sorted(layer.experts) gate_nodes = [layer.experts[expert_id].gate for expert_id in expert_ids] up_nodes = [layer.experts[expert_id].up for expert_id in expert_ids] down_nodes = [layer.experts[expert_id].down for expert_id in expert_ids] - bits, block_size, has_zero_point = _qmoe_geometry(down_nodes[0]) - - intermediate_size = _array(gate_nodes[0].inputs[1]).shape[0] - hidden_size = _array(down_nodes[0].inputs[1]).shape[0] - fc1_weights = np.concatenate([_pack_qmoe_projection(gate_nodes, 1), _pack_qmoe_projection(up_nodes, 1)], axis=1) - fc2_weights = _pack_qmoe_projection(down_nodes, 1) - fc1_scales = np.concatenate( - [ - _pack_qmoe_scales(gate_nodes, 2, intermediate_size, activation_dtype), - _pack_qmoe_scales(up_nodes, 2, intermediate_size, activation_dtype), - ], - axis=1, - ) - fc2_scales = _pack_qmoe_scales(down_nodes, 2, hidden_size, activation_dtype) - - fc1_zero_points = fc2_zero_points = None - if has_zero_point: - fc1_zero_points = np.concatenate( + bits, block_size = _uniform_qmoe_geometry([*gate_nodes, *up_nodes, *down_nodes]) + fc1_has_zero_points = _consistent_zero_points([*gate_nodes, *up_nodes], "FC1") + fc2_has_zero_points = _consistent_zero_points(down_nodes, "FC2") + + try: + intermediate_size = _array(gate_nodes[0].inputs[1]).shape[0] + hidden_size = _array(down_nodes[0].inputs[1]).shape[0] + fc1_weights = np.concatenate([_pack_qmoe_projection(gate_nodes, 1), _pack_qmoe_projection(up_nodes, 1)], axis=1) + fc2_weights = _pack_qmoe_projection(down_nodes, 1) + fc1_scales = np.concatenate( [ - _pack_qmoe_zero_points(gate_nodes, 3, intermediate_size), - _pack_qmoe_zero_points(up_nodes, 3, intermediate_size), + _pack_qmoe_scales(gate_nodes, 2, intermediate_size, activation_dtype), + _pack_qmoe_scales(up_nodes, 2, intermediate_size, activation_dtype), ], axis=1, ) - fc2_zero_points = _pack_qmoe_zero_points(down_nodes, 3, hidden_size) + fc2_scales = _pack_qmoe_scales(down_nodes, 2, hidden_size, activation_dtype) + fc1_zero_points = ( + np.concatenate( + [ + _pack_qmoe_zero_points(gate_nodes, 3, intermediate_size), + _pack_qmoe_zero_points(up_nodes, 3, intermediate_size), + ], + axis=1, + ) + if fc1_has_zero_points + else None + ) + fc2_zero_points = _pack_qmoe_zero_points(down_nodes, 3, hidden_size) if fc2_has_zero_points else None + except ValueError as error: + raise _UnfusableQMoEError(str(error)) from error + + return _QMoEFusionPlan( + layer=layer, + prefix=f"moe.layer{index}", + activation_dtype=activation_dtype, + bits=bits, + block_size=block_size, + fc1_weights=fc1_weights, + fc1_scales=fc1_scales, + fc1_zero_points=fc1_zero_points, + fc2_weights=fc2_weights, + fc2_scales=fc2_scales, + fc2_zero_points=fc2_zero_points, + ) + + +def _emit_qmoe(graph: ir.Graph, plan: _QMoEFusionPlan) -> None: + layer = plan.layer + prefix = plan.prefix - prefix = f"moe.layer{index}" - fc1_weights_value = _make_initializer(graph, f"{prefix}.fc1_experts_weights", fc1_weights, ir.DataType.UINT8) - fc1_scales_value = _make_initializer(graph, f"{prefix}.fc1_scales", fc1_scales, activation_dtype) - fc2_weights_value = _make_initializer(graph, f"{prefix}.fc2_experts_weights", fc2_weights, ir.DataType.UINT8) - fc2_scales_value = _make_initializer(graph, f"{prefix}.fc2_scales", fc2_scales, activation_dtype) + fc1_weights_value = _make_initializer(graph, f"{prefix}.fc1_experts_weights", plan.fc1_weights, ir.DataType.UINT8) + fc1_scales_value = _make_initializer(graph, f"{prefix}.fc1_scales", plan.fc1_scales, plan.activation_dtype) + fc2_weights_value = _make_initializer(graph, f"{prefix}.fc2_experts_weights", plan.fc2_weights, ir.DataType.UINT8) + fc2_scales_value = _make_initializer(graph, f"{prefix}.fc2_scales", plan.fc2_scales, plan.activation_dtype) fc1_zero_points_value = ( - _make_initializer(graph, f"{prefix}.fc1_zero_points", fc1_zero_points, ir.DataType.UINT8) - if fc1_zero_points is not None + _make_initializer(graph, f"{prefix}.fc1_zero_points", plan.fc1_zero_points, ir.DataType.UINT8) + if plan.fc1_zero_points is not None else None ) fc2_zero_points_value = ( - _make_initializer(graph, f"{prefix}.fc2_zero_points", fc2_zero_points, ir.DataType.UINT8) - if fc2_zero_points is not None + _make_initializer(graph, f"{prefix}.fc2_zero_points", plan.fc2_zero_points, ir.DataType.UINT8) + if plan.fc2_zero_points is not None else None ) cast = ir.node( "Cast", inputs=[layer.logits], - attributes={"to": activation_dtype.value}, + attributes={"to": plan.activation_dtype.value}, num_outputs=1, name=f"{prefix}.router_probs_cast", ) @@ -370,8 +446,8 @@ def _emit_qmoe(graph: ir.Graph, layer: _DenseQMoELayer, index: int, activation_d "activation_type": "swiglu", "normalize_routing_weights": 1, "k": layer.k, - "expert_weight_bits": bits, - "block_size": block_size, + "expert_weight_bits": plan.bits, + "block_size": plan.block_size, "swiglu_fusion": 2, "quant_type": "int", "weights_prepacked": 0, @@ -384,19 +460,24 @@ def _emit_qmoe(graph: ir.Graph, layer: _DenseQMoELayer, index: int, activation_d qmoe.outputs[0].type = layer.hidden.type qmoe.outputs[0].shape = layer.routed_out.shape graph.insert_after(layer.topk, [cast, qmoe]) - layer.routed_out.replace_all_uses_with(qmoe.outputs[0]) + if layer.routed_out in graph.outputs: + qmoe.outputs[0].name = layer.routed_out.name + layer.routed_out.replace_all_uses_with(qmoe.outputs[0], replace_graph_outputs=True) graph.opset_imports[_MS_DOMAIN] = 1 def _fuse_dense_moe_to_qmoe(model: ir.Model) -> int: graph = model.graph - fused = 0 + plans: list[_QMoEFusionPlan] = [] for index, layer in enumerate(_DenseQMoELayer(topk) for topk in _find_qmoe_anchors(graph)): - if not layer.is_valid: - logger.warning("Skipping MoE layer at %s: unrecognized dense-fallback structure", layer.topk.name) - continue - down = layer.experts[min(layer.experts)].down - bits, block_size, _ = _qmoe_geometry(down) + try: + layer.check_fusable() + nodes = [ + projection for expert in layer.experts.values() for projection in (expert.gate, expert.up, expert.down) + ] + bits, block_size = _uniform_qmoe_geometry(nodes) + except _UnfusableQMoEError as reason: + raise MoEGraphSurgeryError(f"QMoE graph surgery blocker at {layer.topk.name!r}: {reason}") from reason if not _qmoe_abi_supported(bits, block_size): logger.warning( "Skipping MoE layer at %s: QMoE ABI unsupported (bits=%d, block_size=%d)", @@ -405,6 +486,7 @@ def _fuse_dense_moe_to_qmoe(model: ir.Model) -> int: block_size, ) continue + down = layer.experts[min(layer.experts)].down activation_dtype = layer.hidden.dtype or down.inputs[2].dtype if activation_dtype not in {ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16}: logger.warning( @@ -413,12 +495,17 @@ def _fuse_dense_moe_to_qmoe(model: ir.Model) -> int: activation_dtype, ) continue - _emit_qmoe(graph, layer, index, activation_dtype) - fused += 1 - if fused: + try: + plans.append(_plan_qmoe_layer(layer, index, activation_dtype)) + except _UnfusableQMoEError as reason: + raise MoEGraphSurgeryError(f"QMoE graph surgery blocker at {layer.topk.name!r}: {reason}") from reason + + for plan in plans: + _emit_qmoe(graph, plan) + if plans: _remove_dead_nodes(graph) graph.sort() - return fused + return len(plans) class FuseDenseMoEToQMoE(Surgeon): @@ -838,7 +925,9 @@ def _emit_block_layer(graph: ir.Graph, plan: _BlockFusionPlan) -> None: final_output = moe.outputs[0] graph.insert_after(plan.layer.routed_out.producer(), new_nodes) - plan.layer.routed_out.replace_all_uses_with(final_output) + if plan.layer.routed_out in graph.outputs: + final_output.name = plan.layer.routed_out.name + plan.layer.routed_out.replace_all_uses_with(final_output, replace_graph_outputs=True) graph.opset_imports[_NXRT_DOMAIN] = 1 diff --git a/test/passes/onnx/test_graph_surgeries_moe.py b/test/passes/onnx/test_graph_surgeries_moe.py index d2ff3211d2..5f2e15d2e6 100644 --- a/test/passes/onnx/test_graph_surgeries_moe.py +++ b/test/passes/onnx/test_graph_surgeries_moe.py @@ -167,7 +167,13 @@ def _build_qmoe_graph( topk_value_cast: bool = False, block_size: int = BLOCK_SIZE, include_zero_points: bool = True, + fc1_zero_points: bool | None = None, + fc2_zero_points: bool | None = None, mask_cast_op: str = "CastLike", + broken_expert: int | None = None, + geometry_override: tuple[int, str, int, int] | None = None, + zero_point_override: tuple[int, str, bool] | None = None, + direct_graph_output: bool = False, ) -> tuple[ir.Model, dict[str, _QuantizedWeight]]: rng = np.random.default_rng(0) weights: dict[str, _QuantizedWeight] = {} @@ -178,11 +184,25 @@ def _build_qmoe_graph( type=ir.TensorType(activation_dtype) if hidden_typed else None, ) graph = ir.Graph([hidden], [], nodes=[], name="dense_qmoe") + if fc1_zero_points is None: + fc1_zero_points = include_zero_points + if fc2_zero_points is None: + fc2_zero_points = include_zero_points def quant(name: str, out_features: int, in_features: int) -> _QuantizedWeight: weights[name] = _QuantizedWeight(rng, out_features, in_features, activation_dtype) return weights[name] + def geometry(expert: int, role: str) -> tuple[int, int]: + if geometry_override is not None and geometry_override[:2] == (expert, role): + return geometry_override[2:] + return BITS, block_size + + def has_zero_points(expert: int, role: str) -> bool: + if zero_point_override is not None and zero_point_override[:2] == (expert, role): + return zero_point_override[2] + return fc1_zero_points if role in {"gate", "up"} else fc2_zero_points + router = _matmul_nbits( graph, nodes, @@ -270,7 +290,9 @@ def quant(name: str, out_features: int, in_features: int) -> _QuantizedWeight: f"expert{expert}.gate", hidden, quant(f"gate{expert}", INTERMEDIATE, HIDDEN), - include_zero_points=include_zero_points, + bits=geometry(expert, "gate")[0], + block_size=geometry(expert, "gate")[1], + include_zero_points=has_zero_points(expert, "gate"), ) if legacy_activation: sigmoid = ir.node("Sigmoid", inputs=[gate], num_outputs=1, name=f"expert{expert}.sigmoid") @@ -285,19 +307,29 @@ def quant(name: str, out_features: int, in_features: int) -> _QuantizedWeight: f"expert{expert}.up", hidden, quant(f"up{expert}", INTERMEDIATE, HIDDEN), - include_zero_points=include_zero_points, + bits=geometry(expert, "up")[0], + block_size=geometry(expert, "up")[1], + include_zero_points=has_zero_points(expert, "up"), ) product = ir.node("Mul", inputs=[activation.outputs[0], up], num_outputs=1, name=f"expert{expert}.product") nodes.append(product) - down = _matmul_nbits( - graph, - nodes, - f"expert{expert}.down", - product.outputs[0], - quant(f"down{expert}", HIDDEN, INTERMEDIATE), - block_size=block_size, - include_zero_points=include_zero_points, - ) + if broken_expert == expert: + down_node = ir.node( + "Identity", inputs=[product.outputs[0]], num_outputs=1, name=f"expert{expert}.down_broken" + ) + nodes.append(down_node) + down = down_node.outputs[0] + else: + down = _matmul_nbits( + graph, + nodes, + f"expert{expert}.down", + product.outputs[0], + quant(f"down{expert}", HIDDEN, INTERMEDIATE), + bits=geometry(expert, "down")[0], + block_size=geometry(expert, "down")[1], + include_zero_points=has_zero_points(expert, "down"), + ) contribution = ir.node( "Mul", inputs=[down, reduce_sum.outputs[0]], num_outputs=1, name=f"expert{expert}.contribution" ) @@ -309,16 +341,23 @@ def quant(name: str, out_features: int, in_features: int) -> _QuantizedWeight: nodes.append(add) routed = add.outputs[0] - shared = ir.node("Identity", inputs=[hidden], num_outputs=1, name="shared_expert") - nodes.append(shared) - final = ir.node("Add", inputs=[routed, shared.outputs[0]], num_outputs=1, name="final_add") - final.outputs[0].name = "output" - final.outputs[0].shape = hidden.shape - final.outputs[0].type = ir.TensorType(activation_dtype) - nodes.append(final) + if direct_graph_output: + routed.name = "output" + routed.shape = hidden.shape + routed.type = ir.TensorType(activation_dtype) + graph_output = routed + else: + shared = ir.node("Identity", inputs=[hidden], num_outputs=1, name="shared_expert") + nodes.append(shared) + final = ir.node("Add", inputs=[routed, shared.outputs[0]], num_outputs=1, name="final_add") + final.outputs[0].name = "output" + final.outputs[0].shape = hidden.shape + final.outputs[0].type = ir.TensorType(activation_dtype) + nodes.append(final) + graph_output = final.outputs[0] for node in nodes: graph.append(node) - graph.outputs.append(final.outputs[0]) + graph.outputs.append(graph_output) model = ir.Model(graph, ir_version=10, producer_name="test") model.opset_imports[""] = 21 model.opset_imports["com.microsoft"] = 1 @@ -371,6 +410,7 @@ def _build_block_moe_graph( broken_expert: int | None = None, alien_routing_expert: int | None = None, mask_cast_op: str = "CastLike", + direct_graph_output: bool = False, ) -> tuple[ir.Model, dict[str, np.ndarray]]: rng = np.random.default_rng(1) weights: dict[str, np.ndarray] = {} @@ -535,16 +575,23 @@ def cast_out(value: ir.Value, name: str) -> ir.Value: nodes.append(add) routed = add.outputs[0] - shared = ir.node("Identity", inputs=[hidden], num_outputs=1, name="shared_expert") - nodes.append(shared) - final = ir.node("Add", inputs=[routed, shared.outputs[0]], num_outputs=1, name="final_add") - final.outputs[0].name = "output" - final.outputs[0].shape = hidden.shape - final.outputs[0].type = ir.TensorType(dtype) - nodes.append(final) + if direct_graph_output: + routed.name = "output" + routed.shape = hidden.shape + routed.type = ir.TensorType(dtype) + graph_output = routed + else: + shared = ir.node("Identity", inputs=[hidden], num_outputs=1, name="shared_expert") + nodes.append(shared) + final = ir.node("Add", inputs=[routed, shared.outputs[0]], num_outputs=1, name="final_add") + final.outputs[0].name = "output" + final.outputs[0].shape = hidden.shape + final.outputs[0].type = ir.TensorType(dtype) + nodes.append(final) + graph_output = final.outputs[0] for node in nodes: graph.append(node) - graph.outputs.append(final.outputs[0]) + graph.outputs.append(graph_output) model = ir.Model(graph, ir_version=10, producer_name="test") model.opset_imports[""] = 21 model.opset_imports["pkg.nxrt"] = 1 @@ -702,6 +749,18 @@ def test_fuse_dense_moe_to_qmoe_supports_symmetric_weights_without_zero_points(t assert qmoe.inputs[12] is None +@pytest.mark.parametrize(("fc1_zero_points", "fc2_zero_points"), [(True, False), (False, True)]) +def test_fuse_dense_moe_to_qmoe_wires_zero_point_banks_independently(tmp_path, fc1_zero_points, fc2_zero_points): + model, _ = _build_qmoe_graph( + fc1_zero_points=fc1_zero_points, + fc2_zero_points=fc2_zero_points, + ) + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + qmoe = next(node for node in rewritten.graph if node.op_type == "QMoE") + assert (qmoe.inputs[11] is not None) is fc1_zero_points + assert (qmoe.inputs[12] is not None) is fc2_zero_points + + def test_fuse_dense_moe_to_qmoe_uses_scale_dtype_when_hidden_is_untyped(tmp_path): model, _ = _build_qmoe_graph(hidden_typed=False) rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") @@ -716,6 +775,66 @@ def test_fuse_dense_moe_to_qmoe_keeps_unsupported_geometry(tmp_path): assert _count(rewritten, "TopK") == 1 +def test_fuse_dense_moe_to_qmoe_rejects_broken_trailing_expert_atomically(tmp_path): + model, _ = _build_qmoe_graph(broken_expert=EXPERTS - 1) + input_path = tmp_path / "input.onnx" + onnx.save_model(ir.to_proto(model), input_path) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + { + "surgeries": [{"surgeon": "FuseDenseMoEToQMoE"}], + "remove_duplicate_initializers": False, + }, + disable_search=True, + ) + with pytest.raises(moe_surgeries.MoEGraphSurgeryError, match=rf"experts \[{EXPERTS - 1}\]"): + graph_surgeries.run(ONNXModelHandler(model_path=str(input_path)), str(tmp_path / "output")) + unchanged = ir.from_proto(onnx.load(input_path)) + assert _count(unchanged, "QMoE") == 0 + assert _count(unchanged, "Equal") == EXPERTS + + +def test_fuse_dense_moe_to_qmoe_rejects_mismatched_projection_geometry(tmp_path): + model, _ = _build_qmoe_graph(geometry_override=(EXPERTS - 1, "up", 4, 32)) + input_path = tmp_path / "input.onnx" + onnx.save_model(ir.to_proto(model), input_path) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + { + "surgeries": [{"surgeon": "FuseDenseMoEToQMoE"}], + "remove_duplicate_initializers": False, + }, + disable_search=True, + ) + with pytest.raises(moe_surgeries.MoEGraphSurgeryError, match="mismatched quantization geometry"): + graph_surgeries.run(ONNXModelHandler(model_path=str(input_path)), str(tmp_path / "output")) + unchanged = ir.from_proto(onnx.load(input_path)) + assert _count(unchanged, "QMoE") == 0 + + +def test_fuse_dense_moe_to_qmoe_rejects_inconsistent_zero_points_within_bank(tmp_path): + model, _ = _build_qmoe_graph(zero_point_override=(EXPERTS - 1, "down", False)) + input_path = tmp_path / "input.onnx" + onnx.save_model(ir.to_proto(model), input_path) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + { + "surgeries": [{"surgeon": "FuseDenseMoEToQMoE"}], + "remove_duplicate_initializers": False, + }, + disable_search=True, + ) + with pytest.raises(moe_surgeries.MoEGraphSurgeryError, match="FC2 projections disagree"): + graph_surgeries.run(ONNXModelHandler(model_path=str(input_path)), str(tmp_path / "output")) + + +def test_fuse_dense_moe_to_qmoe_replaces_direct_graph_output(tmp_path): + model, _ = _build_qmoe_graph(direct_graph_output=True) + rewritten = _run_surgery(model, tmp_path, "FuseDenseMoEToQMoE") + assert rewritten.graph.outputs[0].producer().op_type == "QMoE" + assert rewritten.graph.outputs[0].name == "output" + + def test_fuse_dense_moe_to_qmoe_reads_and_writes_external_initializers(tmp_path): model, weights = _build_qmoe_graph() rewritten = _run_surgery( @@ -791,6 +910,13 @@ def test_fuse_block_quantized_moe_restores_non_float_output_dtype(tmp_path): assert consumer.attributes["to"].value == ir.DataType.FLOAT16.value +def test_fuse_block_quantized_moe_replaces_direct_graph_output(tmp_path): + model, _ = _build_block_moe_graph(direct_graph_output=True) + rewritten = _run_surgery(model, tmp_path, "FuseBlockQuantizedMoE") + assert rewritten.graph.outputs[0].producer().op_type == "BlockQuantizedMoE" + assert rewritten.graph.outputs[0].name == "output" + + @pytest.mark.parametrize( ("builder_parameters", "message"), [ From aabce342c379b100c5fdfb55aa7d2714fdc9371f Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:07:13 -0700 Subject: [PATCH 06/14] Add standard ONNX lowering graph surgeries Port six exporter-independent compatibility lowerings into the GraphSurgeries registry, including standard ONNX rotary embedding and Attention decomposition. Add synthetic behavior tests for numerical parity, non-matches, metadata preservation, and Microsoft rotary ABI separation. Signed-off-by: Xiaoyu Zhang --- olive/passes/onnx/graph_surgery/lowering.py | 486 ++++++++++++ .../onnx/test_graph_surgeries_lowering.py | 698 ++++++++++++++++++ 2 files changed, 1184 insertions(+) create mode 100644 olive/passes/onnx/graph_surgery/lowering.py create mode 100644 test/passes/onnx/test_graph_surgeries_lowering.py diff --git a/olive/passes/onnx/graph_surgery/lowering.py b/olive/passes/onnx/graph_surgery/lowering.py new file mode 100644 index 0000000000..6fb5ab325a --- /dev/null +++ b/olive/passes/onnx/graph_surgery/lowering.py @@ -0,0 +1,486 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Exporter-independent standard ONNX compatibility lowerings.""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +from onnx_ir import tape +from onnxscript.rewriter import pattern + +from olive.passes.onnx.graph_surgery import RewriteRuleSurgeon, Surgeon + +_MASK_NEG = -3.0e38 + + +class _BFloat16ClipRule(pattern.RewriteRuleClassBase): + def check(self, context, x, **_): + result = pattern.MatchResult() + if x.dtype != ir.DataType.BFLOAT16: + return result.fail("Clip input is not bfloat16") + return result + + +class _ClipBothToMinMax(_BFloat16ClipRule): + def pattern(self, op, x, lower, upper): + return op.Clip(x, lower, upper) + + def rewrite(self, op, x, lower, upper): + return op.Min(op.Max(x, lower), upper) + + +class _ClipMinToMax(_BFloat16ClipRule): + def pattern(self, op, x, lower): + return op.Clip(x, lower) + + def rewrite(self, op, x, lower): + return op.Max(x, lower) + + +class ClipToMinMax(RewriteRuleSurgeon): + """Lower bfloat16 Clip to Min and Max, which have ONNX Runtime kernels.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_ClipBothToMinMax().rule(), _ClipMinToMax().rule()]) + + +class _Rank4RMSNormRule(pattern.RewriteRuleClassBase): + def pattern(self, op, x, weight): + return op.RMSNormalization(x, weight, _allow_other_attributes=True, _outputs=["norm_out"]) + + def check(self, context, x, norm_out, **_): + result = pattern.MatchResult() + if x.shape is None or len(x.shape) != 4: + return result.fail("RMSNormalization input is not rank 4") + heads, head_dim = x.shape[2], x.shape[3] + if not isinstance(heads, int) or not isinstance(head_dim, int): + return result.fail("RMSNormalization head dimensions are not static") + + node = norm_out.producer() + if node.attributes.get_int("axis", -1) not in (-1, 3): + return result.fail("RMSNormalization does not reduce only the last axis") + if node.attributes.get_float("epsilon", None) is None: + return result.fail("RMSNormalization has no epsilon attribute") + return result + + def rewrite(self, op, x, weight, norm_out, **_): + node = norm_out.producer() + heads, head_dim = x.shape[2], x.shape[3] + attributes = {name: attribute.value for name, attribute in node.attributes.items()} + attributes["axis"] = -1 + + rank3 = op.Reshape(x, op.Constant(value_ints=[0, -1, head_dim])) + normalized = op.RMSNormalization(rank3, weight, **attributes) + return op.Reshape(normalized, op.Constant(value_ints=[0, -1, heads, head_dim])) + + +class Rank4RMSNormToRank3(RewriteRuleSurgeon): + """Reshape rank-4 RMSNormalization to an equivalent rank-3 normalization.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_Rank4RMSNormRule().rule()]) + + +class _DecomposeOnnxRotaryEmbeddingRule(pattern.RewriteRuleClassBase): + def pattern(self, op, x, cos, sin): + return op.RotaryEmbedding(x, cos, sin, _allow_other_attributes=True, _outputs=["rope_out"]) + + def check(self, context, x, cos, sin, rope_out, **_): + result = pattern.MatchResult() + node = rope_out.producer() + if node.domain not in ("", "ai.onnx"): + return result.fail("RotaryEmbedding is not in the standard ONNX domain") + if node.attributes.get_int("interleaved", 0) != 0: + return result.fail("interleaved RotaryEmbedding is not supported") + if node.attributes.get_int("rotary_embedding_dim", 0) != 0: + return result.fail("partial RotaryEmbedding is not supported") + if node.attributes.get_int("num_heads", 0) <= 0: + return result.fail("num_heads must be positive") + if x.shape is None or len(x.shape) != 3: + return result.fail("RotaryEmbedding input is not rank 3") + if cos.shape is None or len(cos.shape) != 3: + return result.fail("RotaryEmbedding cosine input is not rank 3") + if sin.shape is None or len(sin.shape) != 3: + return result.fail("RotaryEmbedding sine input is not rank 3") + return result + + def rewrite(self, op, x, cos, sin, rope_out, **_): + num_heads = rope_out.producer().attributes.get_int("num_heads") + + # (B, S, N*H) -> (B, S, N, H), with H inferred for symbolic hidden sizes. + x_rank4 = op.Reshape(x, op.Constant(value_ints=[0, 0, num_heads, -1])) + half = op.Shape(cos, start=2, end=3) + zero = op.Constant(value_ints=[0]) + last_axis = op.Constant(value_ints=[-1]) + int64_max = op.Constant(value_ints=[9223372036854775807]) + first_half = op.Slice(x_rank4, zero, half, last_axis) + second_half = op.Slice(x_rank4, half, int64_max, last_axis) + + head_axis = op.Constant(value_ints=[2]) + cos = op.Unsqueeze(cos, head_axis) + sin = op.Unsqueeze(sin, head_axis) + rotated_first = op.Sub(op.Mul(first_half, cos), op.Mul(second_half, sin)) + rotated_second = op.Add(op.Mul(second_half, cos), op.Mul(first_half, sin)) + rotated = op.Concat(rotated_first, rotated_second, axis=-1) + return op.Reshape(rotated, op.Constant(value_ints=[0, 0, -1])) + + +class DecomposeOnnxRotaryEmbedding(RewriteRuleSurgeon): + """Decompose the standard ONNX RotaryEmbedding without changing the Microsoft ABI surgery.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_DecomposeOnnxRotaryEmbeddingRule().rule()]) + + +class _TensorScatterToScatterNDRule(pattern.RewriteRuleClassBase): + def pattern(self, op, cache, update, write_indices): + return op.TensorScatter( + cache, + update, + write_indices, + _allow_other_attributes=True, + _outputs=["scatter_out"], + ) + + def check(self, context, cache, scatter_out, **_): + result = pattern.MatchResult() + node = scatter_out.producer() + if node.domain not in ("", "ai.onnx"): + return result.fail("TensorScatter is not in the standard ONNX domain") + if node.attributes.get_int("axis", 0) != 1: + return result.fail("TensorScatter axis is not 1") + if cache.shape is None or len(cache.shape) != 3: + return result.fail("TensorScatter cache is not rank 3") + if isinstance(cache.shape[0], int) and cache.shape[0] != 1: + return result.fail("TensorScatter batch size is not 1") + if not isinstance(cache.shape[1], int): + return result.fail("TensorScatter cache length is not static") + return result + + def rewrite(self, op, cache, update, write_indices, scatter_out, **_): + max_length = int(cache.shape[1]) + zero = op.Constant(value_ints=[0]) + last_axis = op.Constant(value_ints=[-1]) + + cache = op.Squeeze(cache, zero) + update_rank2 = op.Squeeze(update, zero) + full_range = op.Constant(value=ir.tensor(np.arange(max_length, dtype=np.int64))) + sequence_length = op.Shape(update, start=1, end=2) + offsets = op.Slice(full_range, zero, sequence_length, zero) + start = op.Squeeze(write_indices, zero) + positions = op.Add(offsets, start) + indices = op.Unsqueeze(positions, last_axis) + updated = op.ScatterND(cache, indices, update_rank2) + return op.Unsqueeze(updated, zero) + + +class TensorScatterToScatterND(RewriteRuleSurgeon): + """Lower batch-1 static-cache TensorScatter writes to ScatterND.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet([_TensorScatterToScatterNDRule().rule()]) + + +def _is_decomposable_attention(node: ir.Node) -> bool: + if node.op_type != "Attention" or node.domain not in ("", "ai.onnx"): + return False + query, key, value = node.inputs[:3] + if any( + input_value is None or input_value.shape is None or len(input_value.shape) != 3 + for input_value in (query, key, value) + ): + return False + + query_heads = node.attributes.get_int("q_num_heads", 0) + kv_heads = node.attributes.get_int("kv_num_heads", 0) + return ( + query_heads > 0 + and kv_heads > 0 + and query_heads % kv_heads == 0 + and node.attributes.get_int("qk_matmul_output_mode", 0) == 0 + ) + + +def _constant_ints(builder: tape.Tape, values) -> ir.Value: + return builder.op("Constant", [], {"value_ints": list(values)}) + + +def _constant_floats(builder: tape.Tape, values) -> ir.Value: + return builder.op("Constant", [], {"value_floats": [float(value) for value in values]}) + + +def _attention_indices(builder: tape.Tape, length, static_length: int | None): + if static_length is None: + zero = builder.op("Constant", [], {"value_int": 0}) + one = builder.op("Constant", [], {"value_int": 1}) + return builder.op("Range", [zero, length, one]) + full_range = _constant_ints(builder, range(static_length)) + return builder.op( + "Slice", + [ + full_range, + _constant_ints(builder, [0]), + builder.op("Unsqueeze", [length, _constant_ints(builder, [0])]), + ], + ) + + +def _static_kv_length(node: ir.Node) -> int | None: + key = node.inputs[1] + if key is None or key.shape is None or len(key.shape) != 3: + return None + key_length = key.shape[1] + + past_key = node.inputs[4] if len(node.inputs) > 4 else None + if past_key is None: + past_length = 0 + elif past_key.shape is None or len(past_key.shape) != 4 or not isinstance(past_key.shape[2], int): + return None + else: + past_length = past_key.shape[2] + return int(key_length) + int(past_length) if isinstance(key_length, int) else None + + +def _nonpadding_bias(builder: tape.Tape, nonpadding_length, key, scores, static_key_length): + key_length = builder.op( + "Squeeze", + [ + builder.op( + "Slice", [builder.op("Shape", [key]), _constant_ints(builder, [2]), _constant_ints(builder, [3])] + ), + _constant_ints(builder, [0]), + ], + ) + columns = _attention_indices(builder, key_length, static_key_length) + columns = builder.op("Unsqueeze", [columns, _constant_ints(builder, [0, 1, 2])]) + limit = builder.op( + "Unsqueeze", + [builder.op("CastLike", [nonpadding_length, columns]), _constant_ints(builder, [1, 2, 3])], + ) + allowed = builder.op("Less", [columns, limit]) + zero = builder.op("CastLike", [_constant_floats(builder, [0.0]), scores]) + masked = builder.op("CastLike", [_constant_floats(builder, [_MASK_NEG]), scores]) + return builder.op("Where", [allowed, zero, masked]) + + +def _causal_bias(builder: tape.Tape, query, key, scores, static_key_length): + query_length = builder.op( + "Squeeze", + [ + builder.op( + "Slice", [builder.op("Shape", [query]), _constant_ints(builder, [2]), _constant_ints(builder, [3])] + ), + _constant_ints(builder, [0]), + ], + ) + key_length = builder.op( + "Squeeze", + [ + builder.op( + "Slice", [builder.op("Shape", [key]), _constant_ints(builder, [2]), _constant_ints(builder, [3])] + ), + _constant_ints(builder, [0]), + ], + ) + rows = _attention_indices(builder, query_length, static_key_length) + columns = _attention_indices(builder, key_length, static_key_length) + offset = builder.op("Sub", [key_length, query_length]) + rows = builder.op("Add", [builder.op("Unsqueeze", [rows, _constant_ints(builder, [1])]), offset]) + columns = builder.op("Unsqueeze", [columns, _constant_ints(builder, [0])]) + allowed = builder.op("LessOrEqual", [columns, rows]) + zero = builder.op("CastLike", [_constant_floats(builder, [0.0]), scores]) + masked = builder.op("CastLike", [_constant_floats(builder, [_MASK_NEG]), scores]) + bias = builder.op("Where", [allowed, zero, masked]) + return builder.op("Unsqueeze", [bias, _constant_ints(builder, [0, 1])]) + + +def _build_attention_replacement(node: ir.Node) -> tuple[list[ir.Node], list[ir.Value]]: + builder = tape.Tape() + inputs = node.inputs + query, key, value = inputs[:3] + attention_mask = inputs[3] if len(inputs) > 3 else None + past_key = inputs[4] if len(inputs) > 4 else None + past_value = inputs[5] if len(inputs) > 5 else None + nonpadding_length = inputs[6] if len(inputs) > 6 else None + + query_heads = node.attributes.get_int("q_num_heads") + kv_heads = node.attributes.get_int("kv_num_heads") + group_size = query_heads // kv_heads + scale = node.attributes.get_float("scale", None) + softcap = node.attributes.get_float("softcap", 0.0) or 0.0 + is_causal = node.attributes.get_int("is_causal", 0) + + def split_heads(input_value, num_heads): + rank4 = builder.op("Reshape", [input_value, _constant_ints(builder, [0, 0, num_heads, -1])]) + return builder.op("Transpose", [rank4], {"perm": [0, 2, 1, 3]}) + + query = split_heads(query, query_heads) + key = split_heads(key, kv_heads) + value = split_heads(value, kv_heads) + key = builder.op("Concat", [past_key, key], {"axis": 2}) if past_key is not None else key + value = builder.op("Concat", [past_value, value], {"axis": 2}) if past_value is not None else value + present_key, present_value = key, value + + def repeat_kv(input_value): + if group_size == 1: + return input_value + rank5 = builder.op("Unsqueeze", [input_value, _constant_ints(builder, [2])]) + shape = builder.op("Shape", [input_value]) + batch = builder.op("Slice", [shape, _constant_ints(builder, [0]), _constant_ints(builder, [1])]) + heads = builder.op("Slice", [shape, _constant_ints(builder, [1]), _constant_ints(builder, [2])]) + sequence_and_head = builder.op("Slice", [shape, _constant_ints(builder, [2]), _constant_ints(builder, [4])]) + expand_shape = builder.op( + "Concat", + [batch, heads, _constant_ints(builder, [group_size]), sequence_and_head], + {"axis": 0}, + ) + expanded = builder.op("Expand", [rank5, expand_shape]) + target_shape = builder.op( + "Concat", [batch, _constant_ints(builder, [query_heads]), sequence_and_head], {"axis": 0} + ) + return builder.op("Reshape", [expanded, target_shape]) + + repeated_key = repeat_kv(key) + repeated_value = repeat_kv(value) + transposed_key = builder.op("Transpose", [repeated_key], {"perm": [0, 1, 3, 2]}) + scores = builder.op("MatMul", [query, transposed_key]) + if scale is not None: + scale_value = builder.op("CastLike", [_constant_floats(builder, [scale]), scores]) + scores = builder.op("Mul", [scores, scale_value]) + else: + head_dimension = builder.op( + "Squeeze", + [ + builder.op( + "Slice", + [builder.op("Shape", [query]), _constant_ints(builder, [3]), _constant_ints(builder, [4])], + ), + _constant_ints(builder, [0]), + ], + ) + head_dimension = builder.op("CastLike", [head_dimension, scores]) + scores = builder.op("Div", [scores, builder.op("Sqrt", [head_dimension])]) + + if softcap: + cap = builder.op("CastLike", [_constant_floats(builder, [softcap]), scores]) + scores = builder.op("Mul", [builder.op("Tanh", [builder.op("Div", [scores, cap])]), cap]) + if attention_mask is not None: + scores = builder.op("Add", [scores, builder.op("CastLike", [attention_mask, scores])]) + + static_key_length = _static_kv_length(node) + if nonpadding_length is not None: + scores = builder.op( + "Add", [scores, _nonpadding_bias(builder, nonpadding_length, key, scores, static_key_length)] + ) + if is_causal: + scores = builder.op("Add", [scores, _causal_bias(builder, query, key, scores, static_key_length)]) + + probabilities = builder.op("Softmax", [scores], {"axis": -1}) + output = builder.op("MatMul", [probabilities, repeated_value]) + output = builder.op("Transpose", [output], {"perm": [0, 2, 1, 3]}) + output = builder.op("Reshape", [output, _constant_ints(builder, [0, 0, -1])]) + return builder.nodes, [output, present_key, present_value] + + +class _DecomposeAttentionPass(ir.passes.InPlacePass): + def call(self, model: ir.Model) -> ir.passes.PassResult: + graph = model.graph + targets = [node for node in graph if _is_decomposable_attention(node)] + for node in targets: + new_nodes, new_values = _build_attention_replacement(node) + ir.convenience.replace_nodes_and_values( + graph, + insertion_point=node, + old_nodes=[node], + new_nodes=new_nodes, + old_values=list(node.outputs), + new_values=new_values[: len(node.outputs)], + ) + return ir.passes.PassResult(model, modified=bool(targets)) + + +class DecomposeAttention(Surgeon): + """Lower standard ONNX Attention to scaled dot-product attention primitives.""" + + def call_ir(self, model: ir.Model) -> ir.Model: + return _DecomposeAttentionPass()(model).model + + +def _shape_tail_size(shape_tail: ir.Value) -> int | None: + node = shape_tail.producer() + if node is None or node.op_type != "Constant": + return None + value_ints = node.attributes.get("value_ints") + if value_ints is None or len(value_ints.value) != 2 or value_ints.value[0] != 0: + return None + return int(value_ints.value[1]) + + +def _static_empty_kv(op, kv_hidden: int, dtype: ir.DataType): + return op.Constant(value=ir.tensor(np.zeros((1, 0, kv_hidden), dtype=dtype.numpy()))) + + +class _StaticEmptyKVCastLikeRule(pattern.RewriteRuleClassBase): + def pattern(self, op, query_states, shape_tail): + batch_dimension = op.Shape(query_states, start=0, end=1) + empty_shape = op.Concat(batch_dimension, shape_tail, axis=0) + empty = op.ConstantOfShape(empty_shape) + return op.CastLike(empty, query_states) + + def check(self, context, shape_tail, **_): + result = pattern.MatchResult() + if _shape_tail_size(shape_tail) is None: + return result.fail("shape tail is not Constant(value_ints=[0, kv_hidden])") + return result + + def rewrite(self, op, query_states, shape_tail, **_): + return _static_empty_kv( + op, + _shape_tail_size(shape_tail), + query_states.dtype or ir.DataType.FLOAT, + ) + + +class _StaticEmptyKVCastRule(pattern.RewriteRuleClassBase): + def pattern(self, op, query_states, shape_tail): + batch_dimension = op.Shape(query_states, start=0, end=1) + empty_shape = op.Concat(batch_dimension, shape_tail, axis=0) + empty = op.ConstantOfShape(empty_shape) + return op.Cast(empty, _allow_other_attributes=True, _outputs=["empty_kv"]) + + def check(self, context, shape_tail, empty_kv, **_): + result = pattern.MatchResult() + if _shape_tail_size(shape_tail) is None: + return result.fail("shape tail is not Constant(value_ints=[0, kv_hidden])") + if empty_kv.producer().attributes.get_int("to", None) is None: + return result.fail("Cast has no target dtype") + return result + + def rewrite(self, op, shape_tail, empty_kv, **_): + dtype = ir.DataType(empty_kv.producer().attributes.get_int("to")) + return _static_empty_kv(op, _shape_tail_size(shape_tail), dtype) + + +class StaticEmptyKV(RewriteRuleSurgeon): + """Replace dynamic empty-KV construction with a static batch-1 Constant.""" + + def rules(self) -> pattern.RewriteRuleSet: + return pattern.RewriteRuleSet( + [ + _StaticEmptyKVCastLikeRule().rule(), + _StaticEmptyKVCastRule().rule(), + ] + ) + + +__all__ = [ + "ClipToMinMax", + "DecomposeAttention", + "DecomposeOnnxRotaryEmbedding", + "Rank4RMSNormToRank3", + "StaticEmptyKV", + "TensorScatterToScatterND", +] diff --git a/test/passes/onnx/test_graph_surgeries_lowering.py b/test/passes/onnx/test_graph_surgeries_lowering.py new file mode 100644 index 0000000000..93ecaca2ef --- /dev/null +++ b/test/passes/onnx/test_graph_surgeries_lowering.py @@ -0,0 +1,698 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +from collections import Counter + +import ml_dtypes +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest + +from olive.model import ONNXModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.onnx.graph_surgeries import GraphSurgeries +from olive.passes.onnx.graph_surgery import Surgeon +from olive.passes.onnx.graph_surgery import lowering as lowering_surgeries + + +def _apply_surgery(tmp_path, model: ir.Model, surgeon: str, *, case: str = "model") -> ir.Model: + model_path = tmp_path / f"{case}.onnx" + ir.save(model, model_path) + olive_model = ONNXModelHandler(model_path=str(model_path)) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + { + "surgeries": [{"surgeon": surgeon}], + "remove_duplicate_initializers": False, + }, + disable_search=True, + ) + output_model = graph_surgeries.run(olive_model, str(tmp_path / f"{case}_output")) + return output_model.load_ir_model() + + +def _run(model: ir.Model, feeds: dict[str, np.ndarray]) -> list[np.ndarray]: + session = ort.InferenceSession( + ir.to_proto(model).SerializeToString(), + providers=["CPUExecutionProvider"], + ) + return session.run(None, feeds) + + +def _counts(model: ir.Model) -> Counter: + return Counter(node.op_type for node in model.graph.all_nodes()) + + +def _metadata(value: ir.Value): + return value.name, value.dtype, tuple(value.shape) if value.shape is not None else None + + +def test_lowering_module_registers_all_surgeons(): + expected = { + lowering_surgeries.ClipToMinMax, + lowering_surgeries.Rank4RMSNormToRank3, + lowering_surgeries.DecomposeOnnxRotaryEmbedding, + lowering_surgeries.TensorScatterToScatterND, + lowering_surgeries.DecomposeAttention, + lowering_surgeries.StaticEmptyKV, + } + assert {Surgeon.registry[surgeon.__name__.lower()] for surgeon in expected} == expected + + +def _clip_model(dtype: ir.DataType, *, both_bounds: bool = True) -> ir.Model: + numpy_dtype = ml_dtypes.bfloat16 if dtype == ir.DataType.BFLOAT16 else np.float32 + x = ir.Value(name="x", type=ir.TensorType(dtype), shape=ir.Shape(["batch", 4])) + lower = ir.Value( + name="lower", + type=ir.TensorType(dtype), + const_value=ir.tensor(np.array(-1.0, dtype=numpy_dtype)), + ) + inputs = [x, lower] + initializers = [lower] + if both_bounds: + upper = ir.Value( + name="upper", + type=ir.TensorType(dtype), + const_value=ir.tensor(np.array(1.0, dtype=numpy_dtype)), + ) + inputs.append(upper) + initializers.append(upper) + y = ir.Value(name="y", type=ir.TensorType(dtype), shape=ir.Shape(["batch", 4])) + graph = ir.Graph( + inputs=[x], + outputs=[y], + nodes=[ir.Node("", "Clip", inputs=inputs, outputs=[y], name="clip")], + initializers=initializers, + opset_imports={"": 24}, + name="clip", + ) + return ir.Model(graph, ir_version=10) + + +@pytest.mark.parametrize( + ("both_bounds", "expected"), + [(True, Counter({"Max": 1, "Min": 1})), (False, Counter({"Max": 1}))], +) +def test_clip_to_min_max_lowers_bfloat16_and_preserves_metadata(tmp_path, both_bounds, expected): + model = _clip_model(ir.DataType.BFLOAT16, both_bounds=both_bounds) + output_metadata = _metadata(model.graph.outputs[0]) + + lowered = _apply_surgery(tmp_path, model, "ClipToMinMax") + + assert _counts(lowered) == expected + assert _metadata(lowered.graph.outputs[0]) == output_metadata + + +def test_clip_to_min_max_leaves_other_dtypes_unchanged(tmp_path): + lowered = _apply_surgery(tmp_path, _clip_model(ir.DataType.FLOAT), "ClipToMinMax") + assert _counts(lowered) == Counter({"Clip": 1}) + + +def _rmsnorm_model(shape, *, axis=-1, epsilon=1e-6) -> ir.Model: + x = ir.Value(name="x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(shape)) + weight = ir.Value( + name="weight", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([shape[-1]]), + const_value=ir.tensor(np.linspace(0.5, 1.5, shape[-1], dtype=np.float32)), + ) + attributes = {"axis": axis} + if epsilon is not None: + attributes["epsilon"] = epsilon + y = ir.Value(name="y", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(shape)) + node = ir.Node( + "", + "RMSNormalization", + inputs=[x, weight], + outputs=[y], + attributes=ir.convenience.convert_attributes(attributes), + ) + graph = ir.Graph( + inputs=[x], + outputs=[y], + nodes=[node], + initializers=[weight], + opset_imports={"": 24}, + name="rmsnorm", + ) + return ir.Model(graph, ir_version=10) + + +def test_rank4_rmsnorm_to_rank3_is_exact_and_preserves_metadata(tmp_path): + model = _rmsnorm_model([1, 3, 4, 8]) + feeds = {"x": np.random.default_rng(0).standard_normal((1, 3, 4, 8)).astype(np.float32)} + reference = _run(model, feeds)[0] + output_metadata = _metadata(model.graph.outputs[0]) + + lowered = _apply_surgery(tmp_path, model, "Rank4RMSNormToRank3") + + assert _counts(lowered)["RMSNormalization"] == 1 + assert _counts(lowered)["Reshape"] == 2 + assert _metadata(lowered.graph.outputs[0]) == output_metadata + np.testing.assert_allclose(_run(lowered, feeds)[0], reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize( + ("shape", "axis", "epsilon"), + [ + ([1, 3, 8], -1, 1e-6), + ([1, 3, "heads", 8], -1, 1e-6), + ([1, 3, 4, 8], 2, 1e-6), + ([1, 3, 4, 8], -1, None), + ], +) +def test_rank4_rmsnorm_to_rank3_rejects_unsupported_forms(tmp_path, shape, axis, epsilon): + lowered = _apply_surgery( + tmp_path, + _rmsnorm_model(shape, axis=axis, epsilon=epsilon), + "Rank4RMSNormToRank3", + ) + counts = _counts(lowered) + assert counts["RMSNormalization"] == 1 + assert counts.get("Reshape", 0) == 0 + + +def _rope_model( + batch=1, + sequence=4, + num_heads=2, + head_dimension=8, + *, + domain="", + interleaved=0, + rotary_embedding_dim=0, + cos_rank=3, +) -> ir.Model: + hidden_size = num_heads * head_dimension + half_dimension = head_dimension // 2 + x = ir.Value( + name="x", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, sequence, hidden_size]), + ) + if domain: + position_ids = ir.Value( + name="position_ids", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape([batch, sequence]), + ) + cos = ir.Value( + name="cos", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([32, head_dimension]), + ) + sin = ir.Value( + name="sin", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([32, head_dimension]), + ) + inputs = [x, position_ids, cos, sin] + graph_inputs = inputs + else: + cos_shape = [batch, sequence, half_dimension] if cos_rank == 3 else [sequence, half_dimension] + cos = ir.Value(name="cos", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(cos_shape)) + sin = ir.Value(name="sin", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(cos_shape)) + inputs = [x, cos, sin] + graph_inputs = inputs + + y = ir.Value( + name="y", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, sequence, hidden_size]), + ) + node = ir.Node( + domain, + "RotaryEmbedding", + inputs=inputs, + outputs=[y], + attributes=ir.convenience.convert_attributes( + { + "num_heads": num_heads, + "rotary_embedding_dim": rotary_embedding_dim, + "interleaved": interleaved, + } + ), + ) + opsets = {"": 24} + if domain: + opsets[domain] = 1 + graph = ir.Graph( + inputs=graph_inputs, + outputs=[y], + nodes=[node], + opset_imports=opsets, + name="rope", + ) + return ir.Model(graph, ir_version=10) + + +def _rope_feeds(batch=1, sequence=4, num_heads=2, head_dimension=8): + rng = np.random.default_rng(1) + return { + "x": rng.standard_normal((batch, sequence, num_heads * head_dimension)).astype(np.float32), + "cos": rng.standard_normal((batch, sequence, head_dimension // 2)).astype(np.float32), + "sin": rng.standard_normal((batch, sequence, head_dimension // 2)).astype(np.float32), + } + + +def test_decompose_onnx_rotary_embedding_is_exact_and_preserves_metadata(tmp_path): + model = _rope_model(batch=2, sequence=3, num_heads=4, head_dimension=16) + feeds = _rope_feeds(batch=2, sequence=3, num_heads=4, head_dimension=16) + reference = _run(model, feeds)[0] + output_metadata = _metadata(model.graph.outputs[0]) + + lowered = _apply_surgery(tmp_path, model, "DecomposeOnnxRotaryEmbedding") + + counts = _counts(lowered) + assert counts.get("RotaryEmbedding", 0) == 0 + assert counts["Slice"] == 2 + assert counts["Concat"] == 1 + assert _metadata(lowered.graph.outputs[0]) == output_metadata + np.testing.assert_allclose(_run(lowered, feeds)[0], reference, rtol=0, atol=1e-5) + + +@pytest.mark.parametrize( + ("interleaved", "rotary_embedding_dim", "cos_rank"), + [(1, 0, 3), (0, 4, 3), (0, 0, 2)], +) +def test_decompose_onnx_rotary_embedding_rejects_unsupported_forms( + tmp_path, interleaved, rotary_embedding_dim, cos_rank +): + model = _rope_model( + interleaved=interleaved, + rotary_embedding_dim=rotary_embedding_dim, + cos_rank=cos_rank, + ) + lowered = _apply_surgery(tmp_path, model, "DecomposeOnnxRotaryEmbedding") + assert _counts(lowered)["RotaryEmbedding"] == 1 + + +def test_standard_and_microsoft_rotary_embedding_surgeries_are_distinct(tmp_path): + microsoft_model = _rope_model(domain="com.microsoft") + lowered_microsoft = _apply_surgery( + tmp_path, + microsoft_model, + "DecomposeOnnxRotaryEmbedding", + case="microsoft", + ) + assert _counts(lowered_microsoft)["RotaryEmbedding"] == 1 + + standard_model = _rope_model() + lowered_standard = _apply_surgery( + tmp_path, + standard_model, + "DecomposeRotaryEmbedding", + case="standard", + ) + assert _counts(lowered_standard)["RotaryEmbedding"] == 1 + + +def _tensor_scatter_model(max_length=16, dimension=4, *, batch=1, axis=1, symbolic_length=False): + cache_length = "max_length" if symbolic_length else max_length + cache = ir.Value( + name="cache", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, cache_length, dimension]), + ) + update = ir.Value( + name="update", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, "sequence", dimension]), + ) + write_indices = ir.Value( + name="write_indices", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape([1]), + ) + y = ir.Value( + name="y", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, cache_length, dimension]), + ) + node = ir.Node( + "", + "TensorScatter", + inputs=[cache, update, write_indices], + outputs=[y], + attributes=ir.convenience.convert_attributes({"axis": axis}), + ) + graph = ir.Graph( + inputs=[cache, update, write_indices], + outputs=[y], + nodes=[node], + opset_imports={"": 24}, + name="tensor_scatter", + ) + return ir.Model(graph, ir_version=10) + + +def _tensor_scatter_feeds(max_length, dimension, sequence, start): + rng = np.random.default_rng(sequence * 10 + start) + return { + "cache": rng.standard_normal((1, max_length, dimension)).astype(np.float32), + "update": rng.standard_normal((1, sequence, dimension)).astype(np.float32), + "write_indices": np.array([start], dtype=np.int64), + } + + +@pytest.mark.parametrize(("sequence", "start"), [(5, 0), (1, 7)]) +def test_tensor_scatter_to_scatternd_is_exact_and_preserves_metadata(tmp_path, sequence, start): + model = _tensor_scatter_model() + feeds = _tensor_scatter_feeds(16, 4, sequence, start) + reference = _run(model, feeds)[0] + output_metadata = _metadata(model.graph.outputs[0]) + + lowered = _apply_surgery(tmp_path, model, "TensorScatterToScatterND") + + counts = _counts(lowered) + assert counts.get("TensorScatter", 0) == 0 + assert counts["ScatterND"] == 1 + assert counts.get("Range", 0) == 0 + assert _metadata(lowered.graph.outputs[0]) == output_metadata + np.testing.assert_array_equal(_run(lowered, feeds)[0], reference) + + +@pytest.mark.parametrize( + ("batch", "axis", "symbolic_length"), + [(2, 1, False), (1, 0, False), (1, 1, True)], +) +def test_tensor_scatter_to_scatternd_rejects_unsupported_forms(tmp_path, batch, axis, symbolic_length): + model = _tensor_scatter_model(batch=batch, axis=axis, symbolic_length=symbolic_length) + lowered = _apply_surgery(tmp_path, model, "TensorScatterToScatterND") + assert _counts(lowered)["TensorScatter"] == 1 + assert _counts(lowered).get("ScatterND", 0) == 0 + + +def _attention_model( + *, + batch=1, + sequence=4, + past_sequence=0, + query_heads=8, + kv_heads=2, + head_dimension=16, + with_mask=True, + with_past=False, + with_nonpadding=False, + softcap=0.0, + is_causal=1, + rank=3, + qk_matmul_output_mode=0, + output_count=3, +) -> ir.Model: + query_hidden = query_heads * head_dimension + kv_hidden = kv_heads * head_dimension + if rank == 3: + query_shape = [batch, sequence, query_hidden] + kv_shape = [batch, sequence, kv_hidden] + else: + query_shape = [batch, query_heads, sequence, head_dimension] + kv_shape = [batch, kv_heads, sequence, head_dimension] + query = ir.Value(name="query", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(query_shape)) + key = ir.Value(name="key", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(kv_shape)) + value = ir.Value(name="value", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(kv_shape)) + inputs = [query, key, value] + graph_inputs = [query, key, value] + + if with_mask: + mask = ir.Value( + name="mask", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, 1, sequence, past_sequence + sequence]), + ) + inputs.append(mask) + graph_inputs.append(mask) + else: + inputs.append(None) + if with_past: + past_key = ir.Value( + name="past_key", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, kv_heads, past_sequence, head_dimension]), + ) + past_value = ir.Value( + name="past_value", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, kv_heads, past_sequence, head_dimension]), + ) + inputs.extend([past_key, past_value]) + graph_inputs.extend([past_key, past_value]) + elif with_nonpadding: + inputs.extend([None, None]) + if with_nonpadding: + nonpadding = ir.Value( + name="nonpadding", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape([batch]), + ) + inputs.append(nonpadding) + graph_inputs.append(nonpadding) + + outputs = [ + ir.Value( + name="y", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, sequence, query_hidden]), + ), + ir.Value( + name="present_key", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, kv_heads, past_sequence + sequence, head_dimension]), + ), + ir.Value( + name="present_value", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, kv_heads, past_sequence + sequence, head_dimension]), + ), + ][:output_count] + node = ir.Node( + "", + "Attention", + inputs=inputs, + outputs=outputs, + attributes=ir.convenience.convert_attributes( + { + "q_num_heads": query_heads, + "kv_num_heads": kv_heads, + "scale": 1.0, + "softcap": softcap, + "is_causal": is_causal, + "qk_matmul_output_mode": qk_matmul_output_mode, + } + ), + ) + graph = ir.Graph( + inputs=graph_inputs, + outputs=outputs, + nodes=[node], + opset_imports={"": 24}, + name="attention", + ) + return ir.Model(graph, ir_version=10) + + +def _attention_feeds( + *, + batch=1, + sequence=4, + past_sequence=0, + query_heads=8, + kv_heads=2, + head_dimension=16, + with_mask=True, + with_past=False, + with_nonpadding=False, + **_, +): + rng = np.random.default_rng(batch + sequence + past_sequence + kv_heads) + feeds = { + "query": rng.standard_normal((batch, sequence, query_heads * head_dimension)).astype(np.float32), + "key": rng.standard_normal((batch, sequence, kv_heads * head_dimension)).astype(np.float32), + "value": rng.standard_normal((batch, sequence, kv_heads * head_dimension)).astype(np.float32), + } + if with_mask: + feeds["mask"] = (rng.standard_normal((batch, 1, sequence, past_sequence + sequence)) * 0.1).astype(np.float32) + if with_past: + feeds["past_key"] = rng.standard_normal((batch, kv_heads, past_sequence, head_dimension)).astype(np.float32) + feeds["past_value"] = rng.standard_normal((batch, kv_heads, past_sequence, head_dimension)).astype(np.float32) + if with_nonpadding: + feeds["nonpadding"] = np.full((batch,), past_sequence + sequence - 1, dtype=np.int64) + return feeds + + +_ATTENTION_CASES = [ + pytest.param( + { + "batch": 1, + "sequence": 4, + "query_heads": 8, + "kv_heads": 8, + "with_mask": True, + "is_causal": 1, + }, + 0.0, + id="prefill-mha", + ), + pytest.param( + { + "batch": 2, + "sequence": 3, + "query_heads": 8, + "kv_heads": 2, + "with_mask": False, + "is_causal": 0, + }, + 0.0, + id="batch2-gqa", + ), + pytest.param( + { + "batch": 1, + "sequence": 1, + "past_sequence": 3, + "query_heads": 8, + "kv_heads": 1, + "with_mask": True, + "with_past": True, + "softcap": 30.0, + "is_causal": 1, + }, + 1e-5, + id="decode-gqa-past-softcap", + ), +] + + +@pytest.mark.parametrize(("configuration", "atol"), _ATTENTION_CASES) +def test_decompose_attention_is_exact_and_preserves_outputs(tmp_path, configuration, atol): + model = _attention_model(**configuration) + feeds = _attention_feeds(**configuration) + reference = _run(model, feeds) + output_metadata = [_metadata(output) for output in model.graph.outputs] + + lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") + + assert _counts(lowered).get("Attention", 0) == 0 + assert _counts(lowered)["Softmax"] == 1 + assert [_metadata(output) for output in lowered.graph.outputs] == output_metadata + actual = _run(lowered, feeds) + for expected, output in zip(reference, actual): + np.testing.assert_allclose(output, expected, rtol=0, atol=atol) + + +def test_decompose_attention_nonpadding_uses_static_indices(tmp_path): + configuration = { + "with_mask": False, + "with_nonpadding": True, + "is_causal": 0, + } + model = _attention_model(**configuration) + feeds = _attention_feeds(**configuration) + reference = _run(model, feeds) + + lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") + + counts = _counts(lowered) + assert counts.get("Attention", 0) == 0 + assert counts.get("Range", 0) == 0 + assert counts["Less"] >= 1 + np.testing.assert_allclose(_run(lowered, feeds)[0], reference[0], rtol=0, atol=0) + + +def test_decompose_attention_rewires_single_output(tmp_path): + model = _attention_model(output_count=1, with_mask=False, is_causal=0) + feeds = _attention_feeds(with_mask=False) + reference = _run(model, feeds)[0] + + lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") + + assert len(lowered.graph.outputs) == 1 + assert _metadata(lowered.graph.outputs[0]) == _metadata(model.graph.outputs[0]) + np.testing.assert_allclose(_run(lowered, feeds)[0], reference, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("rank", "qk_matmul_output_mode"), + [(4, 0), (3, 1)], +) +def test_decompose_attention_rejects_unsupported_forms(tmp_path, rank, qk_matmul_output_mode): + model = _attention_model( + rank=rank, + qk_matmul_output_mode=qk_matmul_output_mode, + output_count=1, + ) + lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") + assert _counts(lowered)["Attention"] == 1 + + +def _empty_kv_model( + kv_hidden=128, + *, + dtype=ir.DataType.FLOAT16, + cast_like=True, + shape_start=0, +) -> ir.Model: + query = ir.Value( + name="query", + type=ir.TensorType(dtype), + shape=ir.Shape(["batch", "sequence", 64]), + ) + builder = ir.tape.Tape() + batch_dimension = builder.op( + "Shape", + [query], + {"start": shape_start, "end": shape_start + 1}, + ) + shape_tail = builder.op("Constant", [], {"value_ints": [0, kv_hidden]}) + empty_shape = builder.op("Concat", [batch_dimension, shape_tail], {"axis": 0}) + empty = builder.op("ConstantOfShape", [empty_shape]) + output = builder.op("CastLike", [empty, query]) if cast_like else builder.op("Cast", [empty], {"to": int(dtype)}) + output.name = "empty_kv" + output.dtype = dtype + output.shape = ir.Shape([1, 0, kv_hidden]) + graph = ir.Graph( + inputs=[query], + outputs=[output], + nodes=builder.nodes, + opset_imports={"": 24}, + name="empty_kv", + ) + return ir.Model(graph, ir_version=10) + + +@pytest.mark.parametrize( + ("cast_like", "dtype"), + [ + (True, ir.DataType.FLOAT16), + (False, ir.DataType.FLOAT16), + (False, ir.DataType.BFLOAT16), + ], +) +def test_static_empty_kv_preserves_shape_dtype_and_output_metadata(tmp_path, cast_like, dtype): + model = _empty_kv_model(cast_like=cast_like, dtype=dtype) + output_metadata = _metadata(model.graph.outputs[0]) + + lowered = _apply_surgery(tmp_path, model, "StaticEmptyKV") + + counts = _counts(lowered) + assert counts.get("Shape", 0) == 0 + assert counts.get("ConstantOfShape", 0) == 0 + assert counts.get("CastLike", 0) == 0 + assert counts.get("Cast", 0) == 0 + assert _metadata(lowered.graph.outputs[0]) == output_metadata + constant = lowered.graph.outputs[0].producer().attributes["value"].value + assert tuple(constant.shape) == (1, 0, 128) + assert constant.dtype == dtype + + +def test_static_empty_kv_rejects_nonbatch_shape(tmp_path): + model = _empty_kv_model(shape_start=1) + lowered = _apply_surgery(tmp_path, model, "StaticEmptyKV") + counts = _counts(lowered) + assert counts["Shape"] == 1 + assert counts["ConstantOfShape"] == 1 + assert counts["CastLike"] == 1 From a0d9a10c6bd1abd73b1bb8f2b5e54edb4c65e113 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:25:09 -0700 Subject: [PATCH 07/14] Fix lowering surgery edge-case semantics Fail closed for unsupported Attention outputs, mask shapes, and softmax precision while preserving boolean-mask and softcap ordering. Correct external-cache causal alignment, tighten TensorScatter mode and batch handling, and add omitted-index, upper Clip bound, and default RMSNorm epsilon coverage. Signed-off-by: Xiaoyu Zhang --- olive/passes/onnx/graph_surgery/lowering.py | 224 ++++++++++--- .../onnx/test_graph_surgeries_lowering.py | 298 +++++++++++++----- 2 files changed, 401 insertions(+), 121 deletions(-) diff --git a/olive/passes/onnx/graph_surgery/lowering.py b/olive/passes/onnx/graph_surgery/lowering.py index 6fb5ab325a..e83a354235 100644 --- a/olive/passes/onnx/graph_surgery/lowering.py +++ b/olive/passes/onnx/graph_surgery/lowering.py @@ -13,7 +13,7 @@ from olive.passes.onnx.graph_surgery import RewriteRuleSurgeon, Surgeon -_MASK_NEG = -3.0e38 +_MASK_NEGATIVE_INFINITY = float("-inf") class _BFloat16ClipRule(pattern.RewriteRuleClassBase): @@ -40,11 +40,25 @@ def rewrite(self, op, x, lower): return op.Max(x, lower) +class _ClipMaxToMin(_BFloat16ClipRule): + def pattern(self, op, x, upper): + return op.Clip(x, None, upper) + + def rewrite(self, op, x, upper): + return op.Min(x, upper) + + class ClipToMinMax(RewriteRuleSurgeon): """Lower bfloat16 Clip to Min and Max, which have ONNX Runtime kernels.""" def rules(self) -> pattern.RewriteRuleSet: - return pattern.RewriteRuleSet([_ClipBothToMinMax().rule(), _ClipMinToMax().rule()]) + return pattern.RewriteRuleSet( + [ + _ClipBothToMinMax().rule(), + _ClipMinToMax().rule(), + _ClipMaxToMin().rule(), + ] + ) class _Rank4RMSNormRule(pattern.RewriteRuleClassBase): @@ -62,8 +76,6 @@ def check(self, context, x, norm_out, **_): node = norm_out.producer() if node.attributes.get_int("axis", -1) not in (-1, 3): return result.fail("RMSNormalization does not reduce only the last axis") - if node.attributes.get_float("epsilon", None) is None: - return result.fail("RMSNormalization has no epsilon attribute") return result def rewrite(self, op, x, weight, norm_out, **_): @@ -135,32 +147,25 @@ def rules(self) -> pattern.RewriteRuleSet: return pattern.RewriteRuleSet([_DecomposeOnnxRotaryEmbeddingRule().rule()]) -class _TensorScatterToScatterNDRule(pattern.RewriteRuleClassBase): - def pattern(self, op, cache, update, write_indices): - return op.TensorScatter( - cache, - update, - write_indices, - _allow_other_attributes=True, - _outputs=["scatter_out"], - ) - +class _TensorScatterToScatterNDBase(pattern.RewriteRuleClassBase): def check(self, context, cache, scatter_out, **_): result = pattern.MatchResult() node = scatter_out.producer() if node.domain not in ("", "ai.onnx"): return result.fail("TensorScatter is not in the standard ONNX domain") - if node.attributes.get_int("axis", 0) != 1: + if node.attributes.get_int("axis", -2) not in (-2, 1): return result.fail("TensorScatter axis is not 1") + if node.attributes.get_string("mode", "linear") != "linear": + return result.fail("TensorScatter mode is not linear") if cache.shape is None or len(cache.shape) != 3: return result.fail("TensorScatter cache is not rank 3") - if isinstance(cache.shape[0], int) and cache.shape[0] != 1: + if not isinstance(cache.shape[0], int) or cache.shape[0] != 1: return result.fail("TensorScatter batch size is not 1") if not isinstance(cache.shape[1], int): return result.fail("TensorScatter cache length is not static") return result - def rewrite(self, op, cache, update, write_indices, scatter_out, **_): + def _rewrite(self, op, cache, update, write_indices): max_length = int(cache.shape[1]) zero = op.Constant(value_ints=[0]) last_axis = op.Constant(value_ints=[-1]) @@ -170,23 +175,76 @@ def rewrite(self, op, cache, update, write_indices, scatter_out, **_): full_range = op.Constant(value=ir.tensor(np.arange(max_length, dtype=np.int64))) sequence_length = op.Shape(update, start=1, end=2) offsets = op.Slice(full_range, zero, sequence_length, zero) - start = op.Squeeze(write_indices, zero) + start = op.Constant(value_int=0) if write_indices is None else op.Squeeze(write_indices, zero) positions = op.Add(offsets, start) indices = op.Unsqueeze(positions, last_axis) updated = op.ScatterND(cache, indices, update_rank2) return op.Unsqueeze(updated, zero) +class _TensorScatterToScatterNDRule(_TensorScatterToScatterNDBase): + def pattern(self, op, cache, update, write_indices): + return op.TensorScatter( + cache, + update, + write_indices, + _allow_other_attributes=True, + _outputs=["scatter_out"], + ) + + def rewrite(self, op, cache, update, write_indices, **_): + return self._rewrite(op, cache, update, write_indices) + + +class _TensorScatterNoIndicesToScatterNDRule(_TensorScatterToScatterNDBase): + def pattern(self, op, cache, update): + return op.TensorScatter( + cache, + update, + _allow_other_attributes=True, + _outputs=["scatter_out"], + ) + + def rewrite(self, op, cache, update, **_): + return self._rewrite(op, cache, update, None) + + class TensorScatterToScatterND(RewriteRuleSurgeon): """Lower batch-1 static-cache TensorScatter writes to ScatterND.""" def rules(self) -> pattern.RewriteRuleSet: - return pattern.RewriteRuleSet([_TensorScatterToScatterNDRule().rule()]) + return pattern.RewriteRuleSet( + [ + _TensorScatterToScatterNDRule().rule(), + _TensorScatterNoIndicesToScatterNDRule().rule(), + ] + ) + + +def _attention_kv_dimension(node: ir.Node): + key = node.inputs[1] + if key is None or key.shape is None or len(key.shape) != 3: + return None + key_length = key.shape[1] + + past_key = node.inputs[4] if len(node.inputs) > 4 else None + if past_key is None: + return key_length + if past_key.shape is None or len(past_key.shape) != 4: + return None + past_length = past_key.shape[2] + if past_length == 0: + return key_length + if isinstance(key_length, int) and isinstance(past_length, int): + return key_length + past_length + return None def _is_decomposable_attention(node: ir.Node) -> bool: if node.op_type != "Attention" or node.domain not in ("", "ai.onnx"): return False + if len(node.outputs) not in (1, 3): + return False query, key, value = node.inputs[:3] if any( input_value is None or input_value.shape is None or len(input_value.shape) != 3 @@ -196,12 +254,32 @@ def _is_decomposable_attention(node: ir.Node) -> bool: query_heads = node.attributes.get_int("q_num_heads", 0) kv_heads = node.attributes.get_int("kv_num_heads", 0) - return ( - query_heads > 0 - and kv_heads > 0 - and query_heads % kv_heads == 0 - and node.attributes.get_int("qk_matmul_output_mode", 0) == 0 - ) + if query_heads <= 0 or kv_heads <= 0 or query_heads % kv_heads != 0: + return False + if node.attributes.get_int("qk_matmul_output_mode", 0) != 0: + return False + if node.attributes.get("softmax_precision") is not None: + return False + + past_key = node.inputs[4] if len(node.inputs) > 4 else None + past_value = node.inputs[5] if len(node.inputs) > 5 else None + if (past_key is None) != (past_value is None): + return False + nonpadding_length = node.inputs[6] if len(node.inputs) > 6 else None + if nonpadding_length is not None and (past_key is not None or len(node.outputs) != 1): + return False + + attention_mask = node.inputs[3] if len(node.inputs) > 3 else None + if attention_mask is None: + return True + if attention_mask.dtype == ir.DataType.BOOL: + pass + elif query.dtype is None or attention_mask.dtype != query.dtype: + return False + if attention_mask.shape is None or len(attention_mask.shape) == 0: + return False + kv_dimension = _attention_kv_dimension(node) + return kv_dimension is not None and attention_mask.shape[-1] == kv_dimension def _constant_ints(builder: tape.Tape, values) -> ir.Value: @@ -229,19 +307,8 @@ def _attention_indices(builder: tape.Tape, length, static_length: int | None): def _static_kv_length(node: ir.Node) -> int | None: - key = node.inputs[1] - if key is None or key.shape is None or len(key.shape) != 3: - return None - key_length = key.shape[1] - - past_key = node.inputs[4] if len(node.inputs) > 4 else None - if past_key is None: - past_length = 0 - elif past_key.shape is None or len(past_key.shape) != 4 or not isinstance(past_key.shape[2], int): - return None - else: - past_length = past_key.shape[2] - return int(key_length) + int(past_length) if isinstance(key_length, int) else None + kv_dimension = _attention_kv_dimension(node) + return kv_dimension if isinstance(kv_dimension, int) else None def _nonpadding_bias(builder: tape.Tape, nonpadding_length, key, scores, static_key_length): @@ -262,11 +329,19 @@ def _nonpadding_bias(builder: tape.Tape, nonpadding_length, key, scores, static_ ) allowed = builder.op("Less", [columns, limit]) zero = builder.op("CastLike", [_constant_floats(builder, [0.0]), scores]) - masked = builder.op("CastLike", [_constant_floats(builder, [_MASK_NEG]), scores]) + masked = builder.op("CastLike", [_constant_floats(builder, [_MASK_NEGATIVE_INFINITY]), scores]) return builder.op("Where", [allowed, zero, masked]) -def _causal_bias(builder: tape.Tape, query, key, scores, static_key_length): +def _causal_bias( + builder: tape.Tape, + query, + key, + scores, + static_key_length, + past_key, + nonpadding_length, +): query_length = builder.op( "Squeeze", [ @@ -287,14 +362,49 @@ def _causal_bias(builder: tape.Tape, query, key, scores, static_key_length): ) rows = _attention_indices(builder, query_length, static_key_length) columns = _attention_indices(builder, key_length, static_key_length) - offset = builder.op("Sub", [key_length, query_length]) - rows = builder.op("Add", [builder.op("Unsqueeze", [rows, _constant_ints(builder, [1])]), offset]) - columns = builder.op("Unsqueeze", [columns, _constant_ints(builder, [0])]) + if nonpadding_length is None: + if past_key is None: + offset = builder.op("Constant", [], {"value_int": 0}) + else: + offset = builder.op( + "Squeeze", + [ + builder.op( + "Slice", + [ + builder.op("Shape", [past_key]), + _constant_ints(builder, [2]), + _constant_ints(builder, [3]), + ], + ), + _constant_ints(builder, [0]), + ], + ) + rows = builder.op("Add", [builder.op("Unsqueeze", [rows, _constant_ints(builder, [1])]), offset]) + columns = builder.op("Unsqueeze", [columns, _constant_ints(builder, [0])]) + else: + offset = builder.op( + "Sub", + [builder.op("CastLike", [nonpadding_length, query_length]), query_length], + ) + rows = builder.op("Unsqueeze", [rows, _constant_ints(builder, [0, 2])]) + offset = builder.op("Unsqueeze", [offset, _constant_ints(builder, [1, 2])]) + rows = builder.op("Add", [rows, offset]) + columns = builder.op("Unsqueeze", [columns, _constant_ints(builder, [0, 1])]) allowed = builder.op("LessOrEqual", [columns, rows]) zero = builder.op("CastLike", [_constant_floats(builder, [0.0]), scores]) - masked = builder.op("CastLike", [_constant_floats(builder, [_MASK_NEG]), scores]) + masked = builder.op("CastLike", [_constant_floats(builder, [_MASK_NEGATIVE_INFINITY]), scores]) bias = builder.op("Where", [allowed, zero, masked]) - return builder.op("Unsqueeze", [bias, _constant_ints(builder, [0, 1])]) + axes = [1] if nonpadding_length is not None else [0, 1] + return builder.op("Unsqueeze", [bias, _constant_ints(builder, axes)]) + + +def _attention_mask_bias(builder: tape.Tape, attention_mask, scores): + if attention_mask.dtype != ir.DataType.BOOL: + return builder.op("CastLike", [attention_mask, scores]) + zero = builder.op("CastLike", [_constant_floats(builder, [0.0]), scores]) + masked = builder.op("CastLike", [_constant_floats(builder, [_MASK_NEGATIVE_INFINITY]), scores]) + return builder.op("Where", [attention_mask, zero, masked]) def _build_attention_replacement(node: ir.Node) -> tuple[list[ir.Node], list[ir.Value]]: @@ -364,11 +474,8 @@ def repeat_kv(input_value): head_dimension = builder.op("CastLike", [head_dimension, scores]) scores = builder.op("Div", [scores, builder.op("Sqrt", [head_dimension])]) - if softcap: - cap = builder.op("CastLike", [_constant_floats(builder, [softcap]), scores]) - scores = builder.op("Mul", [builder.op("Tanh", [builder.op("Div", [scores, cap])]), cap]) if attention_mask is not None: - scores = builder.op("Add", [scores, builder.op("CastLike", [attention_mask, scores])]) + scores = builder.op("Add", [scores, _attention_mask_bias(builder, attention_mask, scores)]) static_key_length = _static_kv_length(node) if nonpadding_length is not None: @@ -376,7 +483,24 @@ def repeat_kv(input_value): "Add", [scores, _nonpadding_bias(builder, nonpadding_length, key, scores, static_key_length)] ) if is_causal: - scores = builder.op("Add", [scores, _causal_bias(builder, query, key, scores, static_key_length)]) + scores = builder.op( + "Add", + [ + scores, + _causal_bias( + builder, + query, + key, + scores, + static_key_length, + past_key, + nonpadding_length, + ), + ], + ) + if softcap: + cap = builder.op("CastLike", [_constant_floats(builder, [softcap]), scores]) + scores = builder.op("Mul", [builder.op("Tanh", [builder.op("Div", [scores, cap])]), cap]) probabilities = builder.op("Softmax", [scores], {"axis": -1}) output = builder.op("MatMul", [probabilities, repeated_value]) diff --git a/test/passes/onnx/test_graph_surgeries_lowering.py b/test/passes/onnx/test_graph_surgeries_lowering.py index 93ecaca2ef..42e485e58d 100644 --- a/test/passes/onnx/test_graph_surgeries_lowering.py +++ b/test/passes/onnx/test_graph_surgeries_lowering.py @@ -11,6 +11,7 @@ import onnx_ir as ir import onnxruntime as ort import pytest +from onnx.reference import ReferenceEvaluator from olive.model import ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict @@ -43,6 +44,10 @@ def _run(model: ir.Model, feeds: dict[str, np.ndarray]) -> list[np.ndarray]: return session.run(None, feeds) +def _run_reference(model: ir.Model, feeds: dict[str, np.ndarray]) -> list[np.ndarray]: + return ReferenceEvaluator(ir.to_proto(model)).run(None, feeds) + + def _counts(model: ir.Model) -> Counter: return Counter(node.op_type for node in model.graph.all_nodes()) @@ -63,17 +68,27 @@ def test_lowering_module_registers_all_surgeons(): assert {Surgeon.registry[surgeon.__name__.lower()] for surgeon in expected} == expected -def _clip_model(dtype: ir.DataType, *, both_bounds: bool = True) -> ir.Model: +def _clip_model( + dtype: ir.DataType, + *, + lower_bound: bool = True, + upper_bound: bool = True, +) -> ir.Model: numpy_dtype = ml_dtypes.bfloat16 if dtype == ir.DataType.BFLOAT16 else np.float32 x = ir.Value(name="x", type=ir.TensorType(dtype), shape=ir.Shape(["batch", 4])) - lower = ir.Value( - name="lower", - type=ir.TensorType(dtype), - const_value=ir.tensor(np.array(-1.0, dtype=numpy_dtype)), - ) - inputs = [x, lower] - initializers = [lower] - if both_bounds: + inputs = [x] + initializers = [] + if lower_bound: + lower = ir.Value( + name="lower", + type=ir.TensorType(dtype), + const_value=ir.tensor(np.array(-1.0, dtype=numpy_dtype)), + ) + inputs.append(lower) + initializers.append(lower) + elif upper_bound: + inputs.append(None) + if upper_bound: upper = ir.Value( name="upper", type=ir.TensorType(dtype), @@ -94,11 +109,19 @@ def _clip_model(dtype: ir.DataType, *, both_bounds: bool = True) -> ir.Model: @pytest.mark.parametrize( - ("both_bounds", "expected"), - [(True, Counter({"Max": 1, "Min": 1})), (False, Counter({"Max": 1}))], + ("lower_bound", "upper_bound", "expected"), + [ + (True, True, Counter({"Max": 1, "Min": 1})), + (True, False, Counter({"Max": 1})), + (False, True, Counter({"Min": 1})), + ], ) -def test_clip_to_min_max_lowers_bfloat16_and_preserves_metadata(tmp_path, both_bounds, expected): - model = _clip_model(ir.DataType.BFLOAT16, both_bounds=both_bounds) +def test_clip_to_min_max_lowers_bfloat16_and_preserves_metadata(tmp_path, lower_bound, upper_bound, expected): + model = _clip_model( + ir.DataType.BFLOAT16, + lower_bound=lower_bound, + upper_bound=upper_bound, + ) output_metadata = _metadata(model.graph.outputs[0]) lowered = _apply_surgery(tmp_path, model, "ClipToMinMax") @@ -142,8 +165,9 @@ def _rmsnorm_model(shape, *, axis=-1, epsilon=1e-6) -> ir.Model: return ir.Model(graph, ir_version=10) -def test_rank4_rmsnorm_to_rank3_is_exact_and_preserves_metadata(tmp_path): - model = _rmsnorm_model([1, 3, 4, 8]) +@pytest.mark.parametrize("epsilon", [1e-6, None], ids=["explicit-epsilon", "default-epsilon"]) +def test_rank4_rmsnorm_to_rank3_is_exact_and_preserves_metadata(tmp_path, epsilon): + model = _rmsnorm_model([1, 3, 4, 8], epsilon=epsilon) feeds = {"x": np.random.default_rng(0).standard_normal((1, 3, 4, 8)).astype(np.float32)} reference = _run(model, feeds)[0] output_metadata = _metadata(model.graph.outputs[0]) @@ -162,7 +186,6 @@ def test_rank4_rmsnorm_to_rank3_is_exact_and_preserves_metadata(tmp_path): ([1, 3, 8], -1, 1e-6), ([1, 3, "heads", 8], -1, 1e-6), ([1, 3, 4, 8], 2, 1e-6), - ([1, 3, 4, 8], -1, None), ], ) def test_rank4_rmsnorm_to_rank3_rejects_unsupported_forms(tmp_path, shape, axis, epsilon): @@ -311,7 +334,16 @@ def test_standard_and_microsoft_rotary_embedding_surgeries_are_distinct(tmp_path assert _counts(lowered_standard)["RotaryEmbedding"] == 1 -def _tensor_scatter_model(max_length=16, dimension=4, *, batch=1, axis=1, symbolic_length=False): +def _tensor_scatter_model( + max_length=16, + dimension=4, + *, + batch=1, + axis=1, + mode="linear", + symbolic_length=False, + with_write_indices=True, +): cache_length = "max_length" if symbolic_length else max_length cache = ir.Value( name="cache", @@ -323,11 +355,16 @@ def _tensor_scatter_model(max_length=16, dimension=4, *, batch=1, axis=1, symbol type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([batch, "sequence", dimension]), ) - write_indices = ir.Value( - name="write_indices", - type=ir.TensorType(ir.DataType.INT64), - shape=ir.Shape([1]), - ) + inputs = [cache, update] + graph_inputs = [cache, update] + if with_write_indices: + write_indices = ir.Value( + name="write_indices", + type=ir.TensorType(ir.DataType.INT64), + shape=ir.Shape([1]), + ) + inputs.append(write_indices) + graph_inputs.append(write_indices) y = ir.Value( name="y", type=ir.TensorType(ir.DataType.FLOAT), @@ -336,12 +373,12 @@ def _tensor_scatter_model(max_length=16, dimension=4, *, batch=1, axis=1, symbol node = ir.Node( "", "TensorScatter", - inputs=[cache, update, write_indices], + inputs=inputs, outputs=[y], - attributes=ir.convenience.convert_attributes({"axis": axis}), + attributes=ir.convenience.convert_attributes({"axis": axis, "mode": mode}), ) graph = ir.Graph( - inputs=[cache, update, write_indices], + inputs=graph_inputs, outputs=[y], nodes=[node], opset_imports={"": 24}, @@ -350,19 +387,31 @@ def _tensor_scatter_model(max_length=16, dimension=4, *, batch=1, axis=1, symbol return ir.Model(graph, ir_version=10) -def _tensor_scatter_feeds(max_length, dimension, sequence, start): +def _tensor_scatter_feeds(max_length, dimension, sequence, start, *, with_write_indices=True): rng = np.random.default_rng(sequence * 10 + start) - return { + feeds = { "cache": rng.standard_normal((1, max_length, dimension)).astype(np.float32), "update": rng.standard_normal((1, sequence, dimension)).astype(np.float32), - "write_indices": np.array([start], dtype=np.int64), } + if with_write_indices: + feeds["write_indices"] = np.array([start], dtype=np.int64) + return feeds -@pytest.mark.parametrize(("sequence", "start"), [(5, 0), (1, 7)]) -def test_tensor_scatter_to_scatternd_is_exact_and_preserves_metadata(tmp_path, sequence, start): - model = _tensor_scatter_model() - feeds = _tensor_scatter_feeds(16, 4, sequence, start) +@pytest.mark.parametrize( + ("sequence", "start", "with_write_indices"), + [(5, 0, True), (1, 7, True), (5, 0, False)], + ids=["prefill-explicit-zero", "decode-offset", "prefill-implicit-zero"], +) +def test_tensor_scatter_to_scatternd_is_exact_and_preserves_metadata(tmp_path, sequence, start, with_write_indices): + model = _tensor_scatter_model(with_write_indices=with_write_indices) + feeds = _tensor_scatter_feeds( + 16, + 4, + sequence, + start, + with_write_indices=with_write_indices, + ) reference = _run(model, feeds)[0] output_metadata = _metadata(model.graph.outputs[0]) @@ -377,11 +426,22 @@ def test_tensor_scatter_to_scatternd_is_exact_and_preserves_metadata(tmp_path, s @pytest.mark.parametrize( - ("batch", "axis", "symbolic_length"), - [(2, 1, False), (1, 0, False), (1, 1, True)], + ("batch", "axis", "mode", "symbolic_length"), + [ + (2, 1, "linear", False), + ("batch", 1, "linear", False), + (1, 0, "linear", False), + (1, 1, "linear", True), + (1, 1, "circular", False), + ], ) -def test_tensor_scatter_to_scatternd_rejects_unsupported_forms(tmp_path, batch, axis, symbolic_length): - model = _tensor_scatter_model(batch=batch, axis=axis, symbolic_length=symbolic_length) +def test_tensor_scatter_to_scatternd_rejects_unsupported_forms(tmp_path, batch, axis, mode, symbolic_length): + model = _tensor_scatter_model( + batch=batch, + axis=axis, + mode=mode, + symbolic_length=symbolic_length, + ) lowered = _apply_surgery(tmp_path, model, "TensorScatterToScatterND") assert _counts(lowered)["TensorScatter"] == 1 assert _counts(lowered).get("ScatterND", 0) == 0 @@ -391,6 +451,7 @@ def _attention_model( *, batch=1, sequence=4, + kv_sequence=None, past_sequence=0, query_heads=8, kv_heads=2, @@ -399,19 +460,24 @@ def _attention_model( with_past=False, with_nonpadding=False, softcap=0.0, + softmax_precision=None, is_causal=1, rank=3, qk_matmul_output_mode=0, output_count=3, + mask_dtype=ir.DataType.FLOAT, + mask_last_dimension=None, + nonpadding_value=None, ) -> ir.Model: + kv_sequence = sequence if kv_sequence is None else kv_sequence query_hidden = query_heads * head_dimension kv_hidden = kv_heads * head_dimension if rank == 3: query_shape = [batch, sequence, query_hidden] - kv_shape = [batch, sequence, kv_hidden] + kv_shape = [batch, kv_sequence, kv_hidden] else: query_shape = [batch, query_heads, sequence, head_dimension] - kv_shape = [batch, kv_heads, sequence, head_dimension] + kv_shape = [batch, kv_heads, kv_sequence, head_dimension] query = ir.Value(name="query", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(query_shape)) key = ir.Value(name="key", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(kv_shape)) value = ir.Value(name="value", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(kv_shape)) @@ -419,10 +485,12 @@ def _attention_model( graph_inputs = [query, key, value] if with_mask: + if mask_last_dimension is None: + mask_last_dimension = past_sequence + kv_sequence mask = ir.Value( name="mask", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape([batch, 1, sequence, past_sequence + sequence]), + type=ir.TensorType(mask_dtype), + shape=ir.Shape([batch, 1, sequence, mask_last_dimension]), ) inputs.append(mask) graph_inputs.append(mask) @@ -461,29 +529,35 @@ def _attention_model( ir.Value( name="present_key", type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape([batch, kv_heads, past_sequence + sequence, head_dimension]), + shape=ir.Shape([batch, kv_heads, past_sequence + kv_sequence, head_dimension]), ), ir.Value( name="present_value", type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape([batch, kv_heads, past_sequence + sequence, head_dimension]), + shape=ir.Shape([batch, kv_heads, past_sequence + kv_sequence, head_dimension]), + ), + ir.Value( + name="qk_matmul_output", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([batch, query_heads, sequence, past_sequence + kv_sequence]), ), ][:output_count] + attributes = { + "q_num_heads": query_heads, + "kv_num_heads": kv_heads, + "scale": 1.0, + "softcap": softcap, + "is_causal": is_causal, + "qk_matmul_output_mode": qk_matmul_output_mode, + } + if softmax_precision is not None: + attributes["softmax_precision"] = int(softmax_precision) node = ir.Node( "", "Attention", inputs=inputs, outputs=outputs, - attributes=ir.convenience.convert_attributes( - { - "q_num_heads": query_heads, - "kv_num_heads": kv_heads, - "scale": 1.0, - "softcap": softcap, - "is_causal": is_causal, - "qk_matmul_output_mode": qk_matmul_output_mode, - } - ), + attributes=ir.convenience.convert_attributes(attributes), ) graph = ir.Graph( inputs=graph_inputs, @@ -499,6 +573,7 @@ def _attention_feeds( *, batch=1, sequence=4, + kv_sequence=None, past_sequence=0, query_heads=8, kv_heads=2, @@ -506,21 +581,32 @@ def _attention_feeds( with_mask=True, with_past=False, with_nonpadding=False, + mask_dtype=ir.DataType.FLOAT, + mask_last_dimension=None, + nonpadding_value=None, **_, ): + kv_sequence = sequence if kv_sequence is None else kv_sequence rng = np.random.default_rng(batch + sequence + past_sequence + kv_heads) feeds = { "query": rng.standard_normal((batch, sequence, query_heads * head_dimension)).astype(np.float32), - "key": rng.standard_normal((batch, sequence, kv_heads * head_dimension)).astype(np.float32), - "value": rng.standard_normal((batch, sequence, kv_heads * head_dimension)).astype(np.float32), + "key": rng.standard_normal((batch, kv_sequence, kv_heads * head_dimension)).astype(np.float32), + "value": rng.standard_normal((batch, kv_sequence, kv_heads * head_dimension)).astype(np.float32), } if with_mask: - feeds["mask"] = (rng.standard_normal((batch, 1, sequence, past_sequence + sequence)) * 0.1).astype(np.float32) + if mask_last_dimension is None: + mask_last_dimension = past_sequence + kv_sequence + if mask_dtype == ir.DataType.BOOL: + feeds["mask"] = rng.random((batch, 1, sequence, mask_last_dimension)) > 0.35 + else: + feeds["mask"] = (rng.standard_normal((batch, 1, sequence, mask_last_dimension)) * 0.1).astype(np.float32) if with_past: feeds["past_key"] = rng.standard_normal((batch, kv_heads, past_sequence, head_dimension)).astype(np.float32) feeds["past_value"] = rng.standard_normal((batch, kv_heads, past_sequence, head_dimension)).astype(np.float32) if with_nonpadding: - feeds["nonpadding"] = np.full((batch,), past_sequence + sequence - 1, dtype=np.int64) + if nonpadding_value is None: + nonpadding_value = past_sequence + kv_sequence - 1 + feeds["nonpadding"] = np.full((batch,), nonpadding_value, dtype=np.int64) return feeds @@ -534,7 +620,7 @@ def _attention_feeds( "with_mask": True, "is_causal": 1, }, - 0.0, + 1e-6, id="prefill-mha", ), pytest.param( @@ -546,7 +632,7 @@ def _attention_feeds( "with_mask": False, "is_causal": 0, }, - 0.0, + 1e-6, id="batch2-gqa", ), pytest.param( @@ -564,6 +650,34 @@ def _attention_feeds( 1e-5, id="decode-gqa-past-softcap", ), + pytest.param( + { + "batch": 1, + "sequence": 4, + "query_heads": 4, + "kv_heads": 4, + "with_mask": True, + "mask_dtype": ir.DataType.BOOL, + "softcap": 2.0, + "is_causal": 0, + }, + 1e-5, + id="boolean-mask-before-softcap", + ), + pytest.param( + { + "batch": 1, + "sequence": 2, + "kv_sequence": 4, + "query_heads": 4, + "kv_heads": 4, + "with_mask": False, + "is_causal": 1, + "output_count": 1, + }, + 1e-6, + id="causal-cross-attention-upper-left", + ), ] @@ -571,7 +685,7 @@ def _attention_feeds( def test_decompose_attention_is_exact_and_preserves_outputs(tmp_path, configuration, atol): model = _attention_model(**configuration) feeds = _attention_feeds(**configuration) - reference = _run(model, feeds) + reference = _run_reference(model, feeds) output_metadata = [_metadata(output) for output in model.graph.outputs] lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") @@ -589,10 +703,11 @@ def test_decompose_attention_nonpadding_uses_static_indices(tmp_path): "with_mask": False, "with_nonpadding": True, "is_causal": 0, + "output_count": 1, } model = _attention_model(**configuration) feeds = _attention_feeds(**configuration) - reference = _run(model, feeds) + reference = _run_reference(model, feeds) lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") @@ -600,31 +715,72 @@ def test_decompose_attention_nonpadding_uses_static_indices(tmp_path): assert counts.get("Attention", 0) == 0 assert counts.get("Range", 0) == 0 assert counts["Less"] >= 1 - np.testing.assert_allclose(_run(lowered, feeds)[0], reference[0], rtol=0, atol=0) + np.testing.assert_allclose(_run(lowered, feeds)[0], reference[0], rtol=0, atol=1e-6) + + +def test_decompose_attention_causal_mask_uses_valid_external_cache_prefix(tmp_path): + configuration = { + "sequence": 4, + "kv_sequence": 16, + "query_heads": 4, + "kv_heads": 2, + "head_dimension": 8, + "with_mask": False, + "with_nonpadding": True, + "nonpadding_value": 4, + "is_causal": 1, + "output_count": 1, + } + model = _attention_model(**configuration) + feeds = _attention_feeds(**configuration) + reference = _run_reference(model, feeds)[0] + + lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") + + assert _counts(lowered).get("Attention", 0) == 0 + np.testing.assert_allclose(_run(lowered, feeds)[0], reference, rtol=0, atol=1e-6) def test_decompose_attention_rewires_single_output(tmp_path): model = _attention_model(output_count=1, with_mask=False, is_causal=0) feeds = _attention_feeds(with_mask=False) - reference = _run(model, feeds)[0] + reference = _run_reference(model, feeds)[0] lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") assert len(lowered.graph.outputs) == 1 assert _metadata(lowered.graph.outputs[0]) == _metadata(model.graph.outputs[0]) - np.testing.assert_allclose(_run(lowered, feeds)[0], reference, rtol=0, atol=0) + np.testing.assert_allclose(_run(lowered, feeds)[0], reference, rtol=0, atol=1e-6) @pytest.mark.parametrize( - ("rank", "qk_matmul_output_mode"), - [(4, 0), (3, 1)], + "configuration", + [ + {"rank": 4, "output_count": 1}, + {"qk_matmul_output_mode": 1, "output_count": 1}, + {"qk_matmul_output_mode": 0, "output_count": 4}, + {"softmax_precision": ir.DataType.FLOAT, "output_count": 1}, + {"mask_last_dimension": 3, "output_count": 1}, + {"with_mask": False, "with_nonpadding": True, "output_count": 3}, + ], + ids=[ + "rank4-input", + "qk-output-mode", + "fourth-qk-output", + "softmax-precision", + "short-mask", + "nonpadding-with-cache-outputs", + ], ) -def test_decompose_attention_rejects_unsupported_forms(tmp_path, rank, qk_matmul_output_mode): - model = _attention_model( - rank=rank, - qk_matmul_output_mode=qk_matmul_output_mode, - output_count=1, - ) +def test_decompose_attention_rejects_unsupported_forms(tmp_path, configuration): + model = _attention_model(**configuration) + lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") + assert _counts(lowered)["Attention"] == 1 + + +def test_decompose_attention_rejects_unknown_mask_shape(tmp_path): + model = _attention_model(output_count=1) + next(iter(model.graph)).inputs[3].shape = None lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") assert _counts(lowered)["Attention"] == 1 From f81d002f790a5ed1c092c67924bca9c01b73d61e Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:23:59 -0700 Subject: [PATCH 08/14] Correct activation and normalization fusion semantics Fuse BiasGelu only for exact Gelu with a proven compatible 1-D bias, guard Gelu fusion by opset, and fail closed on unsafe residual shapes and axes. Preserve ONNX default epsilon explicitly and strengthen scalar, rejection, and numerical parity coverage. Signed-off-by: Xiaoyu Zhang --- .../passes/onnx/graph_surgery/activations.py | 53 ++++++++- .../onnx/graph_surgery/normalization.py | 53 ++++----- .../onnx/test_graph_surgeries_activations.py | 110 ++++++++++++++---- .../test_graph_surgeries_normalization.py | 64 +++++++--- 4 files changed, 216 insertions(+), 64 deletions(-) diff --git a/olive/passes/onnx/graph_surgery/activations.py b/olive/passes/onnx/graph_surgery/activations.py index 23550afd0d..4797230a8e 100644 --- a/olive/passes/onnx/graph_surgery/activations.py +++ b/olive/passes/onnx/graph_surgery/activations.py @@ -5,12 +5,16 @@ from __future__ import annotations import math +from typing import TYPE_CHECKING from onnxscript.rewriter import pattern from olive.constants import MSFT_DOMAIN from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon +if TYPE_CHECKING: + from onnxscript import ir + _SQRT_2 = math.sqrt(2.0) _SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi) _GELU_COEFFICIENT = 0.044715 @@ -19,7 +23,10 @@ def _check_constant(value, expected: float, name: str) -> str | None: if value.const_value is None: return f"{name} is not a constant" - actual = float(value.const_value.numpy().flat[0]) + array = value.const_value.numpy() + if array.size != 1: + return f"{name} must contain exactly one element" + actual = float(array.flat[0]) if not math.isclose(actual, expected, rel_tol=1e-3): return f"{name} is {actual}, expected approximately {expected}" return None @@ -117,6 +124,11 @@ def pattern(self, op, x, three, coefficient, sqrt_2_over_pi, one, half): class FuseGelu(RewriteRuleSurgeon): """Fuse exact and tanh-approximate decomposed Gelu subgraphs into ONNX Gelu.""" + def call_ir(self, model: ir.Model) -> ir.Model: + if model.opset_imports.get("", 0) < 20: + return model + return super().call_ir(model) + def rules(self) -> pattern.RewriteRuleSet: return pattern.RewriteRuleSet( [ @@ -128,6 +140,32 @@ def rules(self) -> pattern.RewriteRuleSet: ) +def _bias_gelu_inputs(add): + shapes = [value.shape for value in add.inputs] + if any(shape is None for shape in shapes): + return None, None, "Add input shapes must be known" + + bias_indices = [index for index, shape in enumerate(shapes) if len(shape) == 1] + if len(bias_indices) != 1: + return None, None, "Add must have exactly one rank-1 bias input" + + bias_index = bias_indices[0] + data_index = 1 - bias_index + data_shape = shapes[data_index] + bias_shape = shapes[bias_index] + if len(data_shape) not in (2, 3): + return None, None, f"BiasGelu data input rank must be 2 or 3, got {len(data_shape)}" + + hidden_size = data_shape[-1] + bias_size = bias_shape[0] + if not isinstance(hidden_size, int) or not isinstance(bias_size, int): + return None, None, "Bias and data last dimensions must be statically known" + if hidden_size != bias_size: + return None, None, f"Bias size {bias_size} does not match data last dimension {hidden_size}" + + return add.inputs[data_index], add.inputs[bias_index], None + + class _AddGeluToBiasGelu(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): def pattern(self, op, add_output): return op.Gelu(add_output, _outputs=["gelu_output"]) @@ -138,8 +176,8 @@ def check(self, context, add_output, gelu_output, **_): gelu = gelu_output.producer() approximate = gelu.attributes.get("approximate", None) approximate_value = approximate.value if approximate is not None else "none" - if approximate_value != "tanh": - return result.fail(f"Gelu uses approximate='{approximate_value}', BiasGelu requires 'tanh'") + if approximate_value != "none": + return result.fail(f"Gelu uses approximate='{approximate_value}', BiasGelu requires exact Gelu") add = add_output.producer() if add is None or add.op_type != "Add": @@ -149,16 +187,21 @@ def check(self, context, add_output, gelu_output, **_): if len(uses) != 1: return result.fail(f"Add output has {len(uses)} consumers, expected exactly one") + _, _, error = _bias_gelu_inputs(add) + if error: + return result.fail(error) + return result def rewrite(self, op, add_output, **_): add = add_output.producer() + data, bias, _ = _bias_gelu_inputs(add) self._record_replaced_add(add) - return op.BiasGelu(add.inputs[0], add.inputs[1], _domain=MSFT_DOMAIN) + return op.BiasGelu(data, bias, _domain=MSFT_DOMAIN) class FuseBiasGelu(RewriteRuleSurgeon): - """Fuse a single-use Add followed by tanh-approximate Gelu into BiasGelu.""" + """Fuse a compatible bias Add followed by exact Gelu into BiasGelu.""" def rules(self) -> pattern.RewriteRuleSet: return pattern.RewriteRuleSet([_AddGeluToBiasGelu.rule()]) diff --git a/olive/passes/onnx/graph_surgery/normalization.py b/olive/passes/onnx/graph_surgery/normalization.py index 50a1e47610..1bbce5b661 100644 --- a/olive/passes/onnx/graph_surgery/normalization.py +++ b/olive/passes/onnx/graph_surgery/normalization.py @@ -15,7 +15,10 @@ def _check_scalar_constant(value, expected: float, name: str) -> str | None: if value.const_value is None: return f"{name} is not a constant" - actual = float(value.const_value.numpy().flat[0]) + array = value.const_value.numpy() + if array.size != 1: + return f"{name} must contain exactly one element" + actual = float(array.flat[0]) if not math.isclose(actual, expected, rel_tol=1e-4): return f"{name} is {actual}, expected {expected}" return None @@ -37,7 +40,10 @@ def _check_layer_normalization_constants(exponent, epsilon, first_axes, second_a if epsilon.const_value is None: return result.fail("Epsilon is not a constant") - epsilon_value = float(epsilon.const_value.numpy().flat[0]) + epsilon_array = epsilon.const_value.numpy() + if epsilon_array.size != 1: + return result.fail("Epsilon must contain exactly one element") + epsilon_value = float(epsilon_array.flat[0]) if epsilon_value <= 0 or epsilon_value > 1.0: return result.fail(f"Epsilon {epsilon_value} is outside the expected range (0, 1]") @@ -115,35 +121,30 @@ def rules(self) -> pattern.RewriteRuleSet: return pattern.RewriteRuleSet([_LayerNormalization.rule(), _LayerNormalizationNoBias.rule()]) -def _check_skip_input(add_output, norm_output, norm_op_type: str, *, check_axis: bool): +def _check_skip_input(add_output, norm_output, norm_op_type: str): result = pattern.MatchResult() add = add_output.producer() if add is None or add.op_type != "Add": return result.fail(f"Input to {norm_op_type} is not produced by Add") - if norm_op_type == "LayerNormalization": - for index, value in enumerate(add.inputs): - if value is not None and value.shape is not None and len(value.shape) < 2: - return result.fail(f"Add input[{index}] has rank {len(value.shape)}, expected at least two") - else: - first_shape = add.inputs[0].shape - second_shape = add.inputs[1].shape - first_rank = len(first_shape) if first_shape is not None else None - second_rank = len(second_shape) if second_shape is not None else None - if first_rank is not None and second_rank is not None and first_rank != second_rank: - return result.fail( - f"Add inputs have different ranks ({first_rank} and {second_rank}); fused inputs must have the same shape" - ) + first_shape = add.inputs[0].shape + second_shape = add.inputs[1].shape + if first_shape is None or second_shape is None: + return result.fail("Add input shapes must be known") + if len(first_shape) not in (2, 3) or len(second_shape) not in (2, 3): + return result.fail("Add input ranks must both be 2 or 3") + if len(first_shape) != len(second_shape): + return result.fail("Add input ranks must match") + if any(first is None or second is None or first != second for first, second in zip(first_shape, second_shape)): + return result.fail("Add inputs must have provably identical shapes") graph = add.graph if graph is not None and add_output in graph.outputs: return result.fail("Add output is a graph output") norm = norm_output.producer() - if norm.attributes.get_float("epsilon", None) is None: - return result.fail(f"Missing epsilon attribute on {norm_op_type}") - if check_axis and norm.attributes.get_int("axis", -1) != -1: - return result.fail("LayerNormalization axis must be -1") + if norm.attributes.get_int("axis", -1) != -1: + return result.fail(f"{norm_op_type} axis must be -1") return result @@ -158,11 +159,11 @@ def pattern(self, op, add_output, weight, bias): ) def check(self, context, add_output, norm_output, **_): - return _check_skip_input(add_output, norm_output, "LayerNormalization", check_axis=True) + return _check_skip_input(add_output, norm_output, "LayerNormalization") def rewrite(self, op, add_output, weight, bias, norm_output, **_): add = add_output.producer() - epsilon = norm_output.producer().attributes.get_float("epsilon") + epsilon = norm_output.producer().attributes.get_float("epsilon", 1e-5) outputs = op.SkipLayerNormalization( add.inputs[0], add.inputs[1], @@ -187,7 +188,7 @@ def pattern(self, op, add_output, weight): ) def check(self, context, add_output, norm_output, **_): - result = _check_skip_input(add_output, norm_output, "LayerNormalization", check_axis=True) + result = _check_skip_input(add_output, norm_output, "LayerNormalization") if not result: return result if len(norm_output.producer().inputs) > 2: @@ -196,7 +197,7 @@ def check(self, context, add_output, norm_output, **_): def rewrite(self, op, add_output, weight, norm_output, **_): add = add_output.producer() - epsilon = norm_output.producer().attributes.get_float("epsilon") + epsilon = norm_output.producer().attributes.get_float("epsilon", 1e-5) outputs = op.SkipLayerNormalization( add.inputs[0], add.inputs[1], @@ -232,11 +233,11 @@ def pattern(self, op, add_output, weight): ) def check(self, context, add_output, norm_output, **_): - return _check_skip_input(add_output, norm_output, "RMSNormalization", check_axis=False) + return _check_skip_input(add_output, norm_output, "RMSNormalization") def rewrite(self, op, add_output, weight, norm_output, **_): add = add_output.producer() - epsilon = norm_output.producer().attributes.get_float("epsilon") + epsilon = norm_output.producer().attributes.get_float("epsilon", 1e-5) outputs = op.SkipSimplifiedLayerNormalization( add.inputs[0], add.inputs[1], diff --git a/test/passes/onnx/test_graph_surgeries_activations.py b/test/passes/onnx/test_graph_surgeries_activations.py index c222a97b21..d71933a68a 100644 --- a/test/passes/onnx/test_graph_surgeries_activations.py +++ b/test/passes/onnx/test_graph_surgeries_activations.py @@ -10,6 +10,7 @@ import onnx import pytest from onnx import TensorProto, helper, numpy_helper +from onnxruntime import InferenceSession from olive.model import ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict @@ -41,27 +42,32 @@ def _count_ops(model): } -def _make_model(nodes, initializers, outputs): - x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4, 8]) - graph_outputs = [helper.make_tensor_value_info(name, TensorProto.FLOAT, [1, 4, 8]) for name in outputs] +def _make_model(nodes, initializers, outputs, *, input_shape=(1, 4, 8), opset=21): + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape) + graph_outputs = [helper.make_tensor_value_info(name, TensorProto.FLOAT, input_shape) for name in outputs] graph = helper.make_graph(nodes, "activation_test", [x], graph_outputs, initializers) - return helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid("", 21)]) + return helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid("", opset)]) -def _build_exact_gelu(*, half_first=False, sqrt_2=_SQRT_2): +def _exact_gelu_parts(*, input_name="x", half_first=False, sqrt_2=_SQRT_2): initializers = [ numpy_helper.from_array(np.array(sqrt_2, dtype=np.float32), name="sqrt_2"), numpy_helper.from_array(np.array(1.0, dtype=np.float32), name="one"), numpy_helper.from_array(np.array(0.5, dtype=np.float32), name="half"), ] nodes = [ - helper.make_node("Div", ["x", "sqrt_2"], ["divided"]), + helper.make_node("Div", [input_name, "sqrt_2"], ["divided"]), helper.make_node("Erf", ["divided"], ["erf"]), helper.make_node("Add", ["erf", "one"], ["shifted"]), - helper.make_node("Mul", ["x", "shifted"], ["scaled"]), + helper.make_node("Mul", [input_name, "shifted"], ["scaled"]), helper.make_node("Mul", ["half", "scaled"] if half_first else ["scaled", "half"], ["y"]), ] - return _make_model(nodes, initializers, ["y"]) + return nodes, initializers + + +def _build_exact_gelu(*, half_first=False, sqrt_2=_SQRT_2, opset=21): + nodes, initializers = _exact_gelu_parts(half_first=half_first, sqrt_2=sqrt_2) + return _make_model(nodes, initializers, ["y"], opset=opset) def _build_approximate_gelu(*, half_first=False, coefficient=0.044715, input_name="x"): @@ -85,16 +91,24 @@ def _build_approximate_gelu(*, half_first=False, coefficient=0.044715, input_nam return nodes, initializers -def _build_bias_gelu(*, approximate="tanh", shared_add=False): - bias = numpy_helper.from_array(np.ones(8, dtype=np.float32), name="bias") - nodes = [helper.make_node("Add", ["x", "bias"], ["add_output"])] +def _build_bias_gelu( + *, + approximate=None, + shared_add=False, + bias_first=False, + data_shape=(1, 4, 8), + bias_shape=(8,), +): + bias = numpy_helper.from_array(np.ones(bias_shape, dtype=np.float32), name="bias") + add_inputs = ["bias", "x"] if bias_first else ["x", "bias"] + nodes = [helper.make_node("Add", add_inputs, ["add_output"])] gelu_attributes = {} if approximate is None else {"approximate": approximate} nodes.append(helper.make_node("Gelu", ["add_output"], ["y"], **gelu_attributes)) outputs = ["y"] if shared_add: nodes.append(helper.make_node("Identity", ["add_output"], ["residual"])) outputs.append("residual") - return _make_model(nodes, [bias], outputs) + return _make_model(nodes, [bias], outputs, input_shape=data_shape) @pytest.mark.parametrize("surgeon_type", [FuseGelu, FuseBiasGelu]) @@ -125,7 +139,9 @@ def test_fuse_gelu_fuses_approximate_variants(tmp_path, half_first): ("model", "remaining_op"), [ (_build_exact_gelu(sqrt_2=2.0), "Erf"), + (_build_exact_gelu(sqrt_2=[_SQRT_2, _SQRT_2]), "Erf"), (_make_model(*_build_approximate_gelu(coefficient=0.05), ["y"]), "Tanh"), + (_make_model(*_build_approximate_gelu(coefficient=[0.044715, 0.044715]), ["y"]), "Tanh"), ], ) def test_fuse_gelu_preserves_non_matching_constants(tmp_path, model, remaining_op): @@ -135,34 +151,88 @@ def test_fuse_gelu_preserves_non_matching_constants(tmp_path, model, remaining_o assert _count_ops(rewritten)[remaining_op] == 1 -def test_fuse_bias_gelu_fuses_single_use_tanh_gelu(tmp_path): - model = _run_surgery(_build_bias_gelu(), tmp_path, "FuseBiasGelu", "bias_gelu") +def test_fuse_gelu_preserves_decomposition_before_opset_20(tmp_path): + model = _run_surgery(_build_exact_gelu(opset=13), tmp_path, "FuseGelu", "gelu_opset_13") + + assert _count_ops(model).get("Gelu", 0) == 0 + assert _count_ops(model)["Erf"] == 1 + + +@pytest.mark.parametrize( + ("approximate", "bias_first", "data_shape"), + [(None, False, (1, 4, 8)), ("none", True, (1, 4, 8)), ("none", False, (4, 8))], +) +def test_fuse_bias_gelu_fuses_exact_gelu_and_orders_bias_last(tmp_path, approximate, bias_first, data_shape): + model = _run_surgery( + _build_bias_gelu(approximate=approximate, bias_first=bias_first, data_shape=data_shape), + tmp_path, + "FuseBiasGelu", + f"bias_gelu_{approximate}_{bias_first}_{len(data_shape)}", + ) assert _count_ops(model) == {"BiasGelu": 1} assert model.graph.node[0].domain == "com.microsoft" + assert list(model.graph.node[0].input) == ["x", "bias"] -@pytest.mark.parametrize("approximate", [None, "none"]) -def test_fuse_bias_gelu_preserves_exact_gelu(tmp_path, approximate): +def test_fuse_bias_gelu_preserves_tanh_gelu(tmp_path): model = _run_surgery( - _build_bias_gelu(approximate=approximate), + _build_bias_gelu(approximate="tanh"), tmp_path, "FuseBiasGelu", - f"exact_bias_gelu_{approximate}", + "tanh_bias_gelu", + ) + + assert _count_ops(model) == {"Add": 1, "Gelu": 1} + + +@pytest.mark.parametrize( + ("data_shape", "bias_shape"), + [ + ((8,), (8,)), + ((1, 4, 8), (1, 4, 8)), + ((1, 4, 8), (1, 8)), + ((1, 4, 8), (7,)), + ((1, 1, 4, 8), (8,)), + ((None, None, None), (8,)), + ], +) +def test_fuse_bias_gelu_rejects_unsafe_add_shapes(tmp_path, data_shape, bias_shape): + model = _run_surgery( + _build_bias_gelu(approximate="none", data_shape=data_shape, bias_shape=bias_shape), + tmp_path, + "FuseBiasGelu", + f"unsafe_bias_gelu_{bias_shape}", ) assert _count_ops(model) == {"Add": 1, "Gelu": 1} def test_fuse_bias_gelu_preserves_shared_add(tmp_path): - model = _run_surgery(_build_bias_gelu(shared_add=True), tmp_path, "FuseBiasGelu", "shared_bias_gelu") + model = _run_surgery( + _build_bias_gelu(approximate="none", shared_add=True), + tmp_path, + "FuseBiasGelu", + "shared_bias_gelu", + ) assert _count_ops(model) == {"Add": 1, "Gelu": 1, "Identity": 1} +def test_fuse_bias_gelu_matches_exact_gelu_numerically(tmp_path): + original = _build_bias_gelu(approximate="none", bias_first=True) + rewritten = _run_surgery(original, tmp_path, "FuseBiasGelu", "bias_gelu_parity") + inputs = {"x": np.random.default_rng(0).standard_normal((1, 4, 8), dtype=np.float32)} + + expected = InferenceSession(original.SerializeToString(), providers=["CPUExecutionProvider"]).run(None, inputs)[0] + actual = InferenceSession(rewritten.SerializeToString(), providers=["CPUExecutionProvider"]).run(None, inputs)[0] + + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + + def test_fuse_gelu_then_bias_gelu_through_public_pass(tmp_path): bias = numpy_helper.from_array(np.ones(8, dtype=np.float32), name="bias") - nodes, initializers = _build_approximate_gelu(input_name="add_output") + nodes, initializers = _exact_gelu_parts(input_name="add_output") model = _make_model([helper.make_node("Add", ["x", "bias"], ["add_output"]), *nodes], [bias, *initializers], ["y"]) model_path = tmp_path / "combined.onnx" onnx.save(model, model_path) diff --git a/test/passes/onnx/test_graph_surgeries_normalization.py b/test/passes/onnx/test_graph_surgeries_normalization.py index 5f79c88715..039f09b3fc 100644 --- a/test/passes/onnx/test_graph_surgeries_normalization.py +++ b/test/passes/onnx/test_graph_surgeries_normalization.py @@ -74,13 +74,16 @@ def _build_skip_normalization( shared_add=False, epsilon=1e-5, axis=-1, - skip_rank=3, + input_shape=(1, 4, 8), + skip_shape=None, + unknown_skip_shape=False, add_is_graph_output=False, ): - input_shape = [1, 4, 8] - skip_shape = input_shape if skip_rank == 3 else [8] x = helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape) - skip = helper.make_tensor_value_info("skip", TensorProto.FLOAT, skip_shape) + effective_skip_shape = ( + (None,) * len(input_shape) if unknown_skip_shape else (input_shape if skip_shape is None else skip_shape) + ) + skip = helper.make_tensor_value_info("skip", TensorProto.FLOAT, effective_skip_shape) y = helper.make_tensor_value_info("y", TensorProto.FLOAT, input_shape) initializers = [numpy_helper.from_array(np.ones(8, dtype=np.float32), name="weight")] nodes = [helper.make_node("Add", ["x", "skip"], ["add_output"])] @@ -89,7 +92,7 @@ def _build_skip_normalization( initializers.append(numpy_helper.from_array(np.zeros(8, dtype=np.float32), name="bias")) norm_inputs.append("bias") attributes = {} if epsilon is None else {"epsilon": epsilon} - if norm_op_type == "LayerNormalization": + if axis is not None: attributes["axis"] = axis nodes.append(helper.make_node(norm_op_type, norm_inputs, ["y"], **attributes)) @@ -135,7 +138,9 @@ def test_fuse_layer_normalization_fuses_bias_variants(tmp_path, include_bias): [ ({"axes": -2}, "ReduceMean"), ({"exponent": 3.0}, "Pow"), + ({"exponent": [2.0, 2.0]}, "Pow"), ({"epsilon": 2.0}, "Sqrt"), + ({"epsilon": [1e-5, 1e-5]}, "Sqrt"), ], ) def test_fuse_layer_normalization_preserves_non_matches(tmp_path, kwargs, remaining_op): @@ -151,13 +156,18 @@ def test_fuse_layer_normalization_preserves_non_matches(tmp_path, kwargs, remain assert counts[remaining_op] >= 1 -@pytest.mark.parametrize("include_bias", [False, True]) -def test_fuse_skip_layer_normalization_rewires_shared_residual(tmp_path, include_bias): +@pytest.mark.parametrize(("include_bias", "input_shape"), [(False, (4, 8)), (True, (1, 4, 8))]) +def test_fuse_skip_layer_normalization_rewires_shared_residual(tmp_path, include_bias, input_shape): model = _run_surgery( - _build_skip_normalization("LayerNormalization", include_bias=include_bias, shared_add=True), + _build_skip_normalization( + "LayerNormalization", + include_bias=include_bias, + shared_add=True, + input_shape=input_shape, + ), tmp_path, "FuseSkipLayerNormalization", - f"skip_layer_norm_{include_bias}", + f"skip_layer_norm_{include_bias}_{len(input_shape)}", ) assert _count_ops(model) == {"Identity": 1, "SkipLayerNormalization": 1} @@ -183,8 +193,11 @@ def test_fuse_skip_layer_normalization_fuses_single_consumer_add(tmp_path): "kwargs", [ {"axis": 0}, - {"epsilon": None}, - {"skip_rank": 1}, + {"input_shape": (1, 1, 4, 8)}, + {"skip_shape": (8,)}, + {"skip_shape": (4, 8)}, + {"skip_shape": (1, 2, 8)}, + {"unknown_skip_shape": True}, {"add_is_graph_output": True}, ], ) @@ -202,6 +215,27 @@ def test_fuse_skip_layer_normalization_preserves_non_matches(tmp_path, kwargs): assert counts["LayerNormalization"] == 1 +@pytest.mark.parametrize( + ("norm_op_type", "surgeon", "fused_op"), + [ + ("LayerNormalization", "FuseSkipLayerNormalization", "SkipLayerNormalization"), + ("RMSNormalization", "FuseSkipRMSNormalization", "SkipSimplifiedLayerNormalization"), + ], +) +def test_fuse_skip_normalization_emits_default_epsilon(tmp_path, norm_op_type, surgeon, fused_op): + model = _run_surgery( + _build_skip_normalization(norm_op_type, epsilon=None, axis=None), + tmp_path, + surgeon, + f"default_epsilon_{norm_op_type}", + ) + + assert _count_ops(model) == {fused_op: 1} + fused = model.graph.node[0] + epsilon = helper.get_attribute_value(next(attr for attr in fused.attribute if attr.name == "epsilon")) + assert epsilon == pytest.approx(1e-5) + + def test_fuse_skip_rms_normalization_rewires_shared_residual(tmp_path): model = _run_surgery( _build_skip_normalization("RMSNormalization", shared_add=True), @@ -231,8 +265,12 @@ def test_fuse_skip_rms_normalization_fuses_single_consumer_add(tmp_path): @pytest.mark.parametrize( "kwargs", [ - {"epsilon": None}, - {"skip_rank": 1}, + {"axis": 0}, + {"input_shape": (1, 1, 4, 8)}, + {"skip_shape": (8,)}, + {"skip_shape": (4, 8)}, + {"skip_shape": (1, 2, 8)}, + {"unknown_skip_shape": True}, {"add_is_graph_output": True}, ], ) From de65205cf4264e041e16bd6ff62ed6236542c505 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:39:18 -0700 Subject: [PATCH 09/14] Respect fused activation and normalization kernel contracts Preserve double normalization graphs and support the valid BiasGelu ranks and skip broadcasting forms. Reorder Add operands for the fused ABI and cover dtype and shape eligibility through GraphSurgeries. Signed-off-by: Xiaoyu Zhang --- .../passes/onnx/graph_surgery/activations.py | 15 +-- .../onnx/graph_surgery/normalization.py | 93 ++++++++++++++----- .../onnx/test_graph_surgeries_activations.py | 10 +- .../test_graph_surgeries_normalization.py | 74 +++++++++++++-- 4 files changed, 150 insertions(+), 42 deletions(-) diff --git a/olive/passes/onnx/graph_surgery/activations.py b/olive/passes/onnx/graph_surgery/activations.py index 4797230a8e..58a1d77558 100644 --- a/olive/passes/onnx/graph_surgery/activations.py +++ b/olive/passes/onnx/graph_surgery/activations.py @@ -146,15 +146,18 @@ def _bias_gelu_inputs(add): return None, None, "Add input shapes must be known" bias_indices = [index for index, shape in enumerate(shapes) if len(shape) == 1] - if len(bias_indices) != 1: - return None, None, "Add must have exactly one rank-1 bias input" + if len(bias_indices) == 2: + data_index, bias_index = 0, 1 + elif len(bias_indices) == 1: + bias_index = bias_indices[0] + data_index = 1 - bias_index + else: + return None, None, "Add must have a rank-1 bias input" - bias_index = bias_indices[0] - data_index = 1 - bias_index data_shape = shapes[data_index] bias_shape = shapes[bias_index] - if len(data_shape) not in (2, 3): - return None, None, f"BiasGelu data input rank must be 2 or 3, got {len(data_shape)}" + if len(data_shape) < 1: + return None, None, "BiasGelu data input rank must be at least one" hidden_size = data_shape[-1] bias_size = bias_shape[0] diff --git a/olive/passes/onnx/graph_surgery/normalization.py b/olive/passes/onnx/graph_surgery/normalization.py index 1bbce5b661..96314a49cb 100644 --- a/olive/passes/onnx/graph_surgery/normalization.py +++ b/olive/passes/onnx/graph_surgery/normalization.py @@ -6,11 +6,14 @@ import math +from onnxscript import ir from onnxscript.rewriter import pattern from olive.constants import MSFT_DOMAIN from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon +_SKIP_NORM_DTYPES = {ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16} + def _check_scalar_constant(value, expected: float, name: str) -> str | None: if value.const_value is None: @@ -121,22 +124,63 @@ def rules(self) -> pattern.RewriteRuleSet: return pattern.RewriteRuleSet([_LayerNormalization.rule(), _LayerNormalizationNoBias.rule()]) -def _check_skip_input(add_output, norm_output, norm_op_type: str): +def _same_dimensions(first_shape, second_shape) -> bool: + return len(first_shape) == len(second_shape) and all( + first is not None and second is not None and first == second for first, second in zip(first_shape, second_shape) + ) + + +def _skip_normalization_inputs(add): + first, second = add.inputs + first_shape = first.shape + second_shape = second.shape + if first_shape is None or second_shape is None: + return None, None, "Add input shapes must be known" + + first_rank = len(first_shape) + second_rank = len(second_shape) + if first_rank not in (2, 3) or second_rank not in (2, 3): + return None, None, "Add input ranks must both be 2 or 3" + + if first_rank == 2 and second_rank == 2: + if not _same_dimensions(first_shape, second_shape): + return None, None, "Rank-2 Add inputs must have identical shapes" + return first, second, None + + if first_rank != second_rank: + data, skip = (first, second) if first_rank == 3 else (second, first) + if not _same_dimensions(data.shape[1:], skip.shape): + return None, None, "Rank-2 skip shape must match rank-3 data sequence and hidden dimensions" + return data, skip, None + + if not _same_dimensions(first_shape[1:], second_shape[1:]): + return None, None, "Rank-3 Add inputs must have matching sequence and hidden dimensions" + if first_shape[0] == second_shape[0]: + return first, second, None + if first_shape[0] == 1: + return second, first, None + if second_shape[0] == 1: + return first, second, None + return None, None, "Rank-3 skip batch dimension must be one or match the data batch dimension" + + +def _check_skip_input(add_output, norm_output, norm_op_type: str, weight, bias=None): result = pattern.MatchResult() add = add_output.producer() if add is None or add.op_type != "Add": return result.fail(f"Input to {norm_op_type} is not produced by Add") - first_shape = add.inputs[0].shape - second_shape = add.inputs[1].shape - if first_shape is None or second_shape is None: - return result.fail("Add input shapes must be known") - if len(first_shape) not in (2, 3) or len(second_shape) not in (2, 3): - return result.fail("Add input ranks must both be 2 or 3") - if len(first_shape) != len(second_shape): - return result.fail("Add input ranks must match") - if any(first is None or second is None or first != second for first, second in zip(first_shape, second_shape)): - return result.fail("Add inputs must have provably identical shapes") + data, skip, error = _skip_normalization_inputs(add) + if error: + return result.fail(error) + + values = [data, skip, weight] + if bias is not None: + values.append(bias) + if any(value.dtype not in _SKIP_NORM_DTYPES for value in values): + return result.fail("Fused normalization inputs must use float, float16, or bfloat16") + if any(value.dtype != data.dtype for value in values[1:]): + return result.fail("Fused normalization inputs must use the same dtype") graph = add.graph if graph is not None and add_output in graph.outputs: @@ -158,15 +202,16 @@ def pattern(self, op, add_output, weight, bias): _outputs=["norm_output"], ) - def check(self, context, add_output, norm_output, **_): - return _check_skip_input(add_output, norm_output, "LayerNormalization") + def check(self, context, add_output, weight, bias, norm_output, **_): + return _check_skip_input(add_output, norm_output, "LayerNormalization", weight, bias) def rewrite(self, op, add_output, weight, bias, norm_output, **_): add = add_output.producer() + data, skip, _ = _skip_normalization_inputs(add) epsilon = norm_output.producer().attributes.get_float("epsilon", 1e-5) outputs = op.SkipLayerNormalization( - add.inputs[0], - add.inputs[1], + data, + skip, weight, bias, _domain=MSFT_DOMAIN, @@ -187,8 +232,8 @@ def pattern(self, op, add_output, weight): _outputs=["norm_output"], ) - def check(self, context, add_output, norm_output, **_): - result = _check_skip_input(add_output, norm_output, "LayerNormalization") + def check(self, context, add_output, weight, norm_output, **_): + result = _check_skip_input(add_output, norm_output, "LayerNormalization", weight) if not result: return result if len(norm_output.producer().inputs) > 2: @@ -197,10 +242,11 @@ def check(self, context, add_output, norm_output, **_): def rewrite(self, op, add_output, weight, norm_output, **_): add = add_output.producer() + data, skip, _ = _skip_normalization_inputs(add) epsilon = norm_output.producer().attributes.get_float("epsilon", 1e-5) outputs = op.SkipLayerNormalization( - add.inputs[0], - add.inputs[1], + data, + skip, weight, _domain=MSFT_DOMAIN, epsilon=epsilon, @@ -232,15 +278,16 @@ def pattern(self, op, add_output, weight): _outputs=["norm_output"], ) - def check(self, context, add_output, norm_output, **_): - return _check_skip_input(add_output, norm_output, "RMSNormalization") + def check(self, context, add_output, weight, norm_output, **_): + return _check_skip_input(add_output, norm_output, "RMSNormalization", weight) def rewrite(self, op, add_output, weight, norm_output, **_): add = add_output.producer() + data, skip, _ = _skip_normalization_inputs(add) epsilon = norm_output.producer().attributes.get_float("epsilon", 1e-5) outputs = op.SkipSimplifiedLayerNormalization( - add.inputs[0], - add.inputs[1], + data, + skip, weight, _domain=MSFT_DOMAIN, epsilon=epsilon, diff --git a/test/passes/onnx/test_graph_surgeries_activations.py b/test/passes/onnx/test_graph_surgeries_activations.py index d71933a68a..00dad053e9 100644 --- a/test/passes/onnx/test_graph_surgeries_activations.py +++ b/test/passes/onnx/test_graph_surgeries_activations.py @@ -160,7 +160,12 @@ def test_fuse_gelu_preserves_decomposition_before_opset_20(tmp_path): @pytest.mark.parametrize( ("approximate", "bias_first", "data_shape"), - [(None, False, (1, 4, 8)), ("none", True, (1, 4, 8)), ("none", False, (4, 8))], + [ + (None, False, (8,)), + ("none", False, (4, 8)), + ("none", True, (1, 4, 8)), + ("none", True, (1, 2, 4, 8)), + ], ) def test_fuse_bias_gelu_fuses_exact_gelu_and_orders_bias_last(tmp_path, approximate, bias_first, data_shape): model = _run_surgery( @@ -189,11 +194,10 @@ def test_fuse_bias_gelu_preserves_tanh_gelu(tmp_path): @pytest.mark.parametrize( ("data_shape", "bias_shape"), [ - ((8,), (8,)), ((1, 4, 8), (1, 4, 8)), ((1, 4, 8), (1, 8)), ((1, 4, 8), (7,)), - ((1, 1, 4, 8), (8,)), + ((), (8,)), ((None, None, None), (8,)), ], ) diff --git a/test/passes/onnx/test_graph_surgeries_normalization.py b/test/passes/onnx/test_graph_surgeries_normalization.py index 039f09b3fc..4f03a0ad59 100644 --- a/test/passes/onnx/test_graph_surgeries_normalization.py +++ b/test/passes/onnx/test_graph_surgeries_normalization.py @@ -77,19 +77,23 @@ def _build_skip_normalization( input_shape=(1, 4, 8), skip_shape=None, unknown_skip_shape=False, + skip_first=False, + dtype=TensorProto.FLOAT, add_is_graph_output=False, ): - x = helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape) + x = helper.make_tensor_value_info("x", dtype, input_shape) effective_skip_shape = ( (None,) * len(input_shape) if unknown_skip_shape else (input_shape if skip_shape is None else skip_shape) ) - skip = helper.make_tensor_value_info("skip", TensorProto.FLOAT, effective_skip_shape) - y = helper.make_tensor_value_info("y", TensorProto.FLOAT, input_shape) - initializers = [numpy_helper.from_array(np.ones(8, dtype=np.float32), name="weight")] - nodes = [helper.make_node("Add", ["x", "skip"], ["add_output"])] + skip = helper.make_tensor_value_info("skip", dtype, effective_skip_shape) + y = helper.make_tensor_value_info("y", dtype, input_shape) + numpy_dtype = np.float64 if dtype == TensorProto.DOUBLE else np.float32 + initializers = [numpy_helper.from_array(np.ones(8, dtype=numpy_dtype), name="weight")] + add_inputs = ["skip", "x"] if skip_first else ["x", "skip"] + nodes = [helper.make_node("Add", add_inputs, ["add_output"])] norm_inputs = ["add_output", "weight"] if include_bias: - initializers.append(numpy_helper.from_array(np.zeros(8, dtype=np.float32), name="bias")) + initializers.append(numpy_helper.from_array(np.zeros(8, dtype=numpy_dtype), name="bias")) norm_inputs.append("bias") attributes = {} if epsilon is None else {"epsilon": epsilon} if axis is not None: @@ -99,9 +103,9 @@ def _build_skip_normalization( outputs = [y] if shared_add: nodes.append(helper.make_node("Identity", ["add_output"], ["residual"])) - outputs.append(helper.make_tensor_value_info("residual", TensorProto.FLOAT, input_shape)) + outputs.append(helper.make_tensor_value_info("residual", dtype, input_shape)) if add_is_graph_output: - outputs.append(helper.make_tensor_value_info("add_output", TensorProto.FLOAT, input_shape)) + outputs.append(helper.make_tensor_value_info("add_output", dtype, input_shape)) graph = helper.make_graph(nodes, "skip_normalization_test", [x, skip], outputs, initializers) return helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid("", 23)]) @@ -195,8 +199,9 @@ def test_fuse_skip_layer_normalization_fuses_single_consumer_add(tmp_path): {"axis": 0}, {"input_shape": (1, 1, 4, 8)}, {"skip_shape": (8,)}, - {"skip_shape": (4, 8)}, + {"skip_shape": (5, 8)}, {"skip_shape": (1, 2, 8)}, + {"input_shape": (2, 4, 8), "skip_shape": (3, 4, 8)}, {"unknown_skip_shape": True}, {"add_is_graph_output": True}, ], @@ -215,6 +220,33 @@ def test_fuse_skip_layer_normalization_preserves_non_matches(tmp_path, kwargs): assert counts["LayerNormalization"] == 1 +@pytest.mark.parametrize( + ("norm_op_type", "surgeon", "fused_op"), + [ + ("LayerNormalization", "FuseSkipLayerNormalization", "SkipLayerNormalization"), + ("RMSNormalization", "FuseSkipRMSNormalization", "SkipSimplifiedLayerNormalization"), + ], +) +@pytest.mark.parametrize("skip_shape", [(4, 8), (1, 4, 8)]) +def test_fuse_skip_normalization_supports_ort_skip_broadcast_and_reorders_data( + tmp_path, norm_op_type, surgeon, fused_op, skip_shape +): + model = _run_surgery( + _build_skip_normalization( + norm_op_type, + input_shape=(2, 4, 8), + skip_shape=skip_shape, + skip_first=True, + ), + tmp_path, + surgeon, + f"broadcast_{norm_op_type}_{len(skip_shape)}", + ) + + assert _count_ops(model) == {fused_op: 1} + assert list(model.graph.node[0].input[:2]) == ["x", "skip"] + + @pytest.mark.parametrize( ("norm_op_type", "surgeon", "fused_op"), [ @@ -236,6 +268,27 @@ def test_fuse_skip_normalization_emits_default_epsilon(tmp_path, norm_op_type, s assert epsilon == pytest.approx(1e-5) +@pytest.mark.parametrize( + ("norm_op_type", "surgeon", "fused_op", "include_bias"), + [ + ("LayerNormalization", "FuseSkipLayerNormalization", "SkipLayerNormalization", True), + ("RMSNormalization", "FuseSkipRMSNormalization", "SkipSimplifiedLayerNormalization", False), + ], +) +def test_fuse_skip_normalization_preserves_double_norm(tmp_path, norm_op_type, surgeon, fused_op, include_bias): + model = _run_surgery( + _build_skip_normalization(norm_op_type, include_bias=include_bias, dtype=TensorProto.DOUBLE), + tmp_path, + surgeon, + f"double_{norm_op_type}", + ) + + counts = _count_ops(model) + assert counts.get(fused_op, 0) == 0 + assert counts["Add"] == 1 + assert counts[norm_op_type] == 1 + + def test_fuse_skip_rms_normalization_rewires_shared_residual(tmp_path): model = _run_surgery( _build_skip_normalization("RMSNormalization", shared_add=True), @@ -268,8 +321,9 @@ def test_fuse_skip_rms_normalization_fuses_single_consumer_add(tmp_path): {"axis": 0}, {"input_shape": (1, 1, 4, 8)}, {"skip_shape": (8,)}, - {"skip_shape": (4, 8)}, + {"skip_shape": (5, 8)}, {"skip_shape": (1, 2, 8)}, + {"input_shape": (2, 4, 8), "skip_shape": (3, 4, 8)}, {"unknown_skip_shape": True}, {"add_is_graph_output": True}, ], From a8ae1f86426421b9bcf7a0649fc2f9c07466c195 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 12:45:33 -0700 Subject: [PATCH 10/14] Register and document exporter-independent graph surgeries Load all 18 migrated surgeons through the normal GraphSurgeries entry point while preserving existing base-class imports. Share rule cleanup and scalar helpers, use onnx_ir directly, and surface tensor read errors. Cover cold-start registration without Mobius, ordered fusions, cache and weight preservation across QKV packing, and residual broadcast numerics. Document explicit configuration, runtime requirements, supported forms, and fail-closed MoE behavior. Signed-off-by: Xiaoyu Zhang --- docs/source/features/onnx-transformations.md | 111 +++++++ olive/passes/onnx/graph_surgery/__init__.py | 49 ++- olive/passes/onnx/graph_surgery/_common.py | 44 +++ .../passes/onnx/graph_surgery/activations.py | 41 +-- olive/passes/onnx/graph_surgery/attention.py | 10 +- olive/passes/onnx/graph_surgery/base.py | 2 +- olive/passes/onnx/graph_surgery/lowering.py | 5 +- olive/passes/onnx/graph_surgery/moe.py | 4 +- .../onnx/graph_surgery/normalization.py | 45 +-- test/passes/onnx/graph_surgery_test_utils.py | 32 ++ .../onnx/test_graph_surgeries_activations.py | 23 +- .../onnx/test_graph_surgeries_attention.py | 1 - test/passes/onnx/test_graph_surgeries_moe.py | 2 +- .../test_graph_surgeries_normalization.py | 27 +- .../onnx/test_graph_surgeries_pipeline.py | 305 ++++++++++++++++++ 15 files changed, 574 insertions(+), 127 deletions(-) create mode 100644 olive/passes/onnx/graph_surgery/_common.py create mode 100644 test/passes/onnx/graph_surgery_test_utils.py create mode 100644 test/passes/onnx/test_graph_surgeries_pipeline.py diff --git a/docs/source/features/onnx-transformations.md b/docs/source/features/onnx-transformations.md index c7a539755a..b5e4a6241f 100644 --- a/docs/source/features/onnx-transformations.md +++ b/docs/source/features/onnx-transformations.md @@ -119,6 +119,117 @@ Olive provides ability to apply many graph `surgeries` on the ONNX model. In the } ``` +### Transformer fusions and compatibility lowerings + +The following surgeons accept ONNX graphs directly and do not depend on Mobius +or another exporter. They are registered automatically when `GraphSurgeries` is +loaded; no implementation-module imports are required. + +These are explicit transformations, not an automatic execution-provider profile. +Choose the surgeries, their order, and the model component according to the +target runtime's operator and dtype support. A fusion that emits an ORT contrib +operator is not a portable standard-ONNX optimization. + +| Surgeon | Transformation | +| --- | --- | +| `FuseGelu` | Exact or tanh-approximate Gelu decomposition to standard ONNX `Gelu`; requires opset 20 or newer. | +| `FuseBiasGelu` | Compatible 1-D bias Add followed by **exact** Gelu to `com.microsoft::BiasGelu`. Tanh Gelu is not equivalent and is left unchanged. | +| `FuseLayerNormalization` | Last-axis ReduceMean-based normalization to standard `LayerNormalization`, with or without bias. | +| `FuseSkipLayerNormalization` | Compatible residual Add and last-axis `LayerNormalization` to `com.microsoft::SkipLayerNormalization`. | +| `FuseSkipRMSNormalization` | Compatible residual Add and last-axis `RMSNormalization` to `com.microsoft::SkipSimplifiedLayerNormalization`. | +| `AttentionToGroupQueryAttention` | Recognized causal standard `Attention`, optionally with RoPE, to `com.microsoft::GroupQueryAttention`. | +| `PackQKVForGroupQueryAttention` | Separate constant-weight Q/K/V projections to a packed GQA projection, including unequal Q/KV widths and bias. | +| `SeparateGroupQueryAttentionRoPE` | Move supported GQA-integrated RoPE into separate standard `RotaryEmbedding` nodes. | +| `UnpackGroupQueryAttentionQKV` | Split supported packed GQA projections into separate Q/K/V projections. | +| `BlockDiagonalAttentionToPackedMHA` | Recognized block-diagonal mask and standard `Attention` to `com.microsoft::PackedMultiHeadAttention`. | +| `ClipToMinMax` | BF16 `Clip` with both bounds, only a lower bound, or only an upper bound to `Max`/`Min`. | +| `Rank4RMSNormToRank3` | Rank-4 last-axis RMSNorm with static head dimensions to rank-3 RMSNorm surrounded by reshapes. | +| `DecomposeOnnxRotaryEmbedding` | Standard rank-3, full-width, non-interleaved `RotaryEmbedding` to primitive rotate-half operations. | +| `TensorScatterToScatterND` | Linear rank-3 static-cache writes with batch size 1 and known cache capacity to `ScatterND`; an omitted write index starts at zero. | +| `DecomposeAttention` | Supported standard rank-3 `Attention` to scaled dot-product primitives, including GQA, cache outputs, and causal/nonpadding masks. | +| `StaticEmptyKV` | Recognized dynamic empty-KV construction to a static empty tensor for graph-capture compatibility. | +| `FuseDenseMoEToQMoE` | Compatible `MatMulNBits` expert banks and routing to `com.microsoft::QMoE`. | +| `FuseBlockQuantizedMoE` | Compatible native block-quantized expert banks and routing to `pkg.nxrt::BlockQuantizedMoE`; requires a runtime implementing that operator. | + +The pattern surgeons preserve graphs that do not match their supported forms. +In particular, `DecomposeAttention` leaves a fourth QK output, an explicit +`softmax_precision`, and masks whose full KV width cannot be established +unchanged. It is not a general-purpose decomposition of every legal Attention +configuration. Check the resulting graph for unsupported operators before +deploying to a runtime without an Attention kernel. + +`DecomposeOnnxRotaryEmbedding` is distinct from the existing +[`DecomposeRotaryEmbedding`](#decomposerotaryembedding), which accepts the +Microsoft-domain four-input ABI. Likewise, `BlockDiagonalAttentionToPackedMHA` +accepts a standard Attention subgraph, not the `custom::PackedAttention` input +expected by `PackedAttentionToPackedMHA`. + +#### Ordering + +`GraphSurgeries` runs the list in order, once per surgeon. For example, on a +runtime that supports GQA and packed QKV: + +```json +{ + "type": "GraphSurgeries", + "surgeries": [ + {"surgeon": "AttentionToGroupQueryAttention"}, + {"surgeon": "PackQKVForGroupQueryAttention"}, + {"surgeon": "FuseSkipRMSNormalization"}, + {"surgeon": "FuseSkipLayerNormalization"}, + {"surgeon": "FuseGelu"}, + {"surgeon": "FuseBiasGelu"} + ] +} +``` + +`FuseGelu` must precede `FuseBiasGelu` when the source Gelu is decomposed. +Packing must follow GQA fusion. For a target requiring separate RoPE and Q/K/V, +apply `SeparateGroupQueryAttentionRoPE` followed by +`UnpackGroupQueryAttentionQKV` instead of packing. These lists are not universal +EP recipes: kernel availability also depends on the model dtype and component. +The surgeries do not implicitly run shape inference or a whole-model optimizer. +Projection packing may leave constant `Concat`/`Transpose` nodes, and traced +subgraphs may need unused-node cleanup; schedule an appropriate optimization +pass separately when required. + +#### Weight-aware MoE surgeries + +Run MoE surgeries **after** quantization or native-block import has supplied the +required weights. Both support loaded external initializers. `GraphSurgeries` +retains its normal external-data output options such as `save_as_external_data`. +The surgeries validate representable expert layouts before emitting replacements; +they do not silently drop an unrecognized expert. + +`FuseBlockQuantizedMoE` defaults to `allow_dense_moe=false`: an identified +native-block MoE layer that cannot be fused raises `MoEGraphSurgeryError`, rather +than quietly retaining an expensive all-expert path. Setting +`allow_dense_moe=true` explicitly retains that path with a warning: + +```json +{ + "type": "GraphSurgeries", + "surgeries": [ + {"surgeon": "FuseBlockQuantizedMoE", "allow_dense_moe": false} + ], + "save_as_external_data": true +} +``` + +`FuseDenseMoEToQMoE` requires compatible 4-bit geometry, complete expert banks, +and consistent zero-point presence within each FC bank. FC1 and FC2 may +independently have zero points. Malformed recognized groups raise +`MoEGraphSurgeryError`; unsupported QMoE ABI geometry is left unchanged with a +warning. Neither surgeon chooses an execution provider or quantizes weights. + +#### Adding a surgeon + +Implement new transformations in a focused module under +`olive/passes/onnx/graph_surgery/`, importing `Surgeon` or +`RewriteRuleSurgeon` from `graph_surgery.base`. Export the class from the +package's `__init__.py` to register it for normal `GraphSurgeries` use. +Existing imports of the base classes from `graph_surgeries` remain supported. + ### `RenameInputs` #### Description diff --git a/olive/passes/onnx/graph_surgery/__init__.py b/olive/passes/onnx/graph_surgery/__init__.py index ccd8d9b105..9503be8fdd 100644 --- a/olive/passes/onnx/graph_surgery/__init__.py +++ b/olive/passes/onnx/graph_surgery/__init__.py @@ -2,6 +2,53 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +"""Built-in surgeons, registered on import for the GraphSurgeries pass.""" + +from olive.passes.onnx.graph_surgery.activations import FuseBiasGelu, FuseGelu +from olive.passes.onnx.graph_surgery.attention import ( + AttentionToGroupQueryAttention, + BlockDiagonalAttentionToPackedMHA, + PackQKVForGroupQueryAttention, + SeparateGroupQueryAttentionRoPE, + UnpackGroupQueryAttentionQKV, +) from olive.passes.onnx.graph_surgery.base import ProtoSurgeon, RewriteRuleSurgeon, Surgeon +from olive.passes.onnx.graph_surgery.lowering import ( + ClipToMinMax, + DecomposeAttention, + DecomposeOnnxRotaryEmbedding, + Rank4RMSNormToRank3, + StaticEmptyKV, + TensorScatterToScatterND, +) +from olive.passes.onnx.graph_surgery.moe import FuseBlockQuantizedMoE, FuseDenseMoEToQMoE, MoEGraphSurgeryError +from olive.passes.onnx.graph_surgery.normalization import ( + FuseLayerNormalization, + FuseSkipLayerNormalization, + FuseSkipRMSNormalization, +) -__all__ = ["ProtoSurgeon", "RewriteRuleSurgeon", "Surgeon"] +__all__ = [ + "AttentionToGroupQueryAttention", + "BlockDiagonalAttentionToPackedMHA", + "ClipToMinMax", + "DecomposeAttention", + "DecomposeOnnxRotaryEmbedding", + "FuseBiasGelu", + "FuseBlockQuantizedMoE", + "FuseDenseMoEToQMoE", + "FuseGelu", + "FuseLayerNormalization", + "FuseSkipLayerNormalization", + "FuseSkipRMSNormalization", + "MoEGraphSurgeryError", + "PackQKVForGroupQueryAttention", + "ProtoSurgeon", + "Rank4RMSNormToRank3", + "RewriteRuleSurgeon", + "SeparateGroupQueryAttentionRoPE", + "StaticEmptyKV", + "Surgeon", + "TensorScatterToScatterND", + "UnpackGroupQueryAttentionQKV", +] diff --git a/olive/passes/onnx/graph_surgery/_common.py b/olive/passes/onnx/graph_surgery/_common.py new file mode 100644 index 0000000000..efa7ed3416 --- /dev/null +++ b/olive/passes/onnx/graph_surgery/_common.py @@ -0,0 +1,44 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import onnx_ir as ir + + +def check_scalar_constant(value: ir.Value, expected: float, name: str, *, rel_tol: float = 1e-4) -> str | None: + if value.const_value is None: + return f"{name} is not a constant" + array = value.const_value.numpy() + if array.size != 1: + return f"{name} must contain exactly one element" + actual = float(array.flat[0]) + if not math.isclose(actual, expected, rel_tol=rel_tol): + return f"{name} is {actual}, expected approximately {expected}" + return None + + +class ReplacedAddCleanupMixin: + """Clean up traced Adds outside the matched pattern after rewriting.""" + + _replaced_adds: list[ir.Node] + + def setup(self): + self._replaced_adds = [] + + def _record_replaced_add(self, add): + self._replaced_adds.append(add) + + def cleanup(self): + for add in self._replaced_adds: + if ( + add.graph is not None + and all(output not in add.graph.outputs for output in add.outputs) + and all(not list(output.uses()) for output in add.outputs) + ): + add.graph.remove(add, safe=True) diff --git a/olive/passes/onnx/graph_surgery/activations.py b/olive/passes/onnx/graph_surgery/activations.py index 58a1d77558..13a3f50fcd 100644 --- a/olive/passes/onnx/graph_surgery/activations.py +++ b/olive/passes/onnx/graph_surgery/activations.py @@ -10,45 +10,20 @@ from onnxscript.rewriter import pattern from olive.constants import MSFT_DOMAIN +from olive.passes.onnx.graph_surgery._common import ReplacedAddCleanupMixin, check_scalar_constant from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon if TYPE_CHECKING: - from onnxscript import ir + import onnx_ir as ir + +# ONNXScript binds each rule's named pattern operands to its callbacks. +# pylint: disable=arguments-differ _SQRT_2 = math.sqrt(2.0) _SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi) _GELU_COEFFICIENT = 0.044715 -def _check_constant(value, expected: float, name: str) -> str | None: - if value.const_value is None: - return f"{name} is not a constant" - array = value.const_value.numpy() - if array.size != 1: - return f"{name} must contain exactly one element" - actual = float(array.flat[0]) - if not math.isclose(actual, expected, rel_tol=1e-3): - return f"{name} is {actual}, expected approximately {expected}" - return None - - -class _RemoveReplacedAdd: - def setup(self): - self._replaced_adds = [] - - def _record_replaced_add(self, add): - self._replaced_adds.append(add) - - def cleanup(self): - for add in self._replaced_adds: - if ( - add.graph is not None - and all(output not in add.graph.outputs for output in add.outputs) - and all(not list(output.uses()) for output in add.outputs) - ): - add.graph.remove(add, safe=True) - - class _ExactGelu(pattern.RewriteRuleClassBase): def pattern(self, op, x, sqrt_2, one, half): divided = op.Div(x, sqrt_2) @@ -64,7 +39,7 @@ def check(self, context, sqrt_2, one, half, **_): (one, 1.0, "Add constant"), (half, 0.5, "Mul half constant"), ): - if error := _check_constant(value, expected, name): + if error := check_scalar_constant(value, expected, name, rel_tol=1e-3): return result.fail(error) return result @@ -101,7 +76,7 @@ def check(self, context, three, coefficient, sqrt_2_over_pi, one, half, **_): (one, 1.0, "Add constant"), (half, 0.5, "Mul half constant"), ): - if error := _check_constant(value, expected, name): + if error := check_scalar_constant(value, expected, name, rel_tol=1e-3): return result.fail(error) return result @@ -169,7 +144,7 @@ def _bias_gelu_inputs(add): return add.inputs[data_index], add.inputs[bias_index], None -class _AddGeluToBiasGelu(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): +class _AddGeluToBiasGelu(ReplacedAddCleanupMixin, pattern.RewriteRuleClassBase): def pattern(self, op, add_output): return op.Gelu(add_output, _outputs=["gelu_output"]) diff --git a/olive/passes/onnx/graph_surgery/attention.py b/olive/passes/onnx/graph_surgery/attention.py index 02efada739..d1cbd1fa10 100644 --- a/olive/passes/onnx/graph_surgery/attention.py +++ b/olive/passes/onnx/graph_surgery/attention.py @@ -7,7 +7,7 @@ from __future__ import annotations import numpy as np -from onnxscript import ir +import onnx_ir as ir from onnxscript.rewriter import pattern from onnxscript.rewriter._basics import MatchFailureError, MatchResult from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase @@ -15,6 +15,9 @@ from olive.constants import MSFT_DOMAIN from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon +# ONNXScript binds each rule's named pattern operands to its callbacks. +# pylint: disable=arguments-differ + def _initializer_dtype(value: ir.Value) -> ir.DataType | None: """Return an initializer's declared dtype, falling back to its tensor dtype.""" @@ -55,10 +58,7 @@ def _constant_int(value: ir.Value | None) -> int | None: tensor = getattr(attr, "value", None) if tensor is None: return None - try: - array = tensor.numpy() - except Exception: # pragma: no cover - return None + array = tensor.numpy() if array.size != 1: return None return int(array.reshape(-1)[0]) diff --git a/olive/passes/onnx/graph_surgery/base.py b/olive/passes/onnx/graph_surgery/base.py index 6acc780bd8..5f2a5b8787 100644 --- a/olive/passes/onnx/graph_surgery/base.py +++ b/olive/passes/onnx/graph_surgery/base.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, ClassVar -from onnxscript import ir +import onnx_ir as ir if TYPE_CHECKING: from onnx import ModelProto diff --git a/olive/passes/onnx/graph_surgery/lowering.py b/olive/passes/onnx/graph_surgery/lowering.py index e83a354235..a0913c8c03 100644 --- a/olive/passes/onnx/graph_surgery/lowering.py +++ b/olive/passes/onnx/graph_surgery/lowering.py @@ -11,7 +11,10 @@ from onnx_ir import tape from onnxscript.rewriter import pattern -from olive.passes.onnx.graph_surgery import RewriteRuleSurgeon, Surgeon +from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon, Surgeon + +# ONNXScript binds each rule's named pattern operands to its callbacks. +# pylint: disable=arguments-differ _MASK_NEGATIVE_INFINITY = float("-inf") diff --git a/olive/passes/onnx/graph_surgery/moe.py b/olive/passes/onnx/graph_surgery/moe.py index c672b12ef1..9bc5d674f6 100644 --- a/olive/passes/onnx/graph_surgery/moe.py +++ b/olive/passes/onnx/graph_surgery/moe.py @@ -10,9 +10,9 @@ import logging import numpy as np -from onnxscript import ir +import onnx_ir as ir -from olive.passes.onnx.graph_surgery import Surgeon +from olive.passes.onnx.graph_surgery.base import Surgeon logger = logging.getLogger(__name__) diff --git a/olive/passes/onnx/graph_surgery/normalization.py b/olive/passes/onnx/graph_surgery/normalization.py index 96314a49cb..6e3a440f8f 100644 --- a/olive/passes/onnx/graph_surgery/normalization.py +++ b/olive/passes/onnx/graph_surgery/normalization.py @@ -4,27 +4,17 @@ # -------------------------------------------------------------------------- from __future__ import annotations -import math - -from onnxscript import ir +import onnx_ir as ir from onnxscript.rewriter import pattern from olive.constants import MSFT_DOMAIN +from olive.passes.onnx.graph_surgery._common import ReplacedAddCleanupMixin, check_scalar_constant from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon -_SKIP_NORM_DTYPES = {ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16} - +# ONNXScript binds each rule's named pattern operands to its callbacks. +# pylint: disable=arguments-differ -def _check_scalar_constant(value, expected: float, name: str) -> str | None: - if value.const_value is None: - return f"{name} is not a constant" - array = value.const_value.numpy() - if array.size != 1: - return f"{name} must contain exactly one element" - actual = float(array.flat[0]) - if not math.isclose(actual, expected, rel_tol=1e-4): - return f"{name} is {actual}, expected {expected}" - return None +_SKIP_NORM_DTYPES = {ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16} def _check_last_axis(value, name: str) -> str | None: @@ -38,7 +28,7 @@ def _check_last_axis(value, name: str) -> str | None: def _check_layer_normalization_constants(exponent, epsilon, first_axes, second_axes): result = pattern.MatchResult() - if error := _check_scalar_constant(exponent, 2.0, "Pow exponent"): + if error := check_scalar_constant(exponent, 2.0, "Pow exponent"): return result.fail(error) if epsilon.const_value is None: @@ -56,23 +46,6 @@ def _check_layer_normalization_constants(exponent, epsilon, first_axes, second_a return result -class _RemoveReplacedAdd: - def setup(self): - self._replaced_adds = [] - - def _record_replaced_add(self, add): - self._replaced_adds.append(add) - - def cleanup(self): - for add in self._replaced_adds: - if ( - add.graph is not None - and all(output not in add.graph.outputs for output in add.outputs) - and all(not list(output.uses()) for output in add.outputs) - ): - add.graph.remove(add, safe=True) - - class _LayerNormalization(pattern.RewriteRuleClassBase): def pattern(self, op, x, first_axes, exponent, second_axes, epsilon, weight, bias): mean = op.ReduceMean(x, first_axes, _allow_other_attributes=True) @@ -192,7 +165,7 @@ def _check_skip_input(add_output, norm_output, norm_op_type: str, weight, bias=N return result -class _AddLayerNormalizationToSkipLayerNormalization(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): +class _AddLayerNormalizationToSkipLayerNormalization(ReplacedAddCleanupMixin, pattern.RewriteRuleClassBase): def pattern(self, op, add_output, weight, bias): return op.LayerNormalization( add_output, @@ -223,7 +196,7 @@ def rewrite(self, op, add_output, weight, bias, norm_output, **_): return outputs[0] -class _AddLayerNormalizationNoBiasToSkipLayerNormalization(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): +class _AddLayerNormalizationNoBiasToSkipLayerNormalization(ReplacedAddCleanupMixin, pattern.RewriteRuleClassBase): def pattern(self, op, add_output, weight): return op.LayerNormalization( add_output, @@ -269,7 +242,7 @@ def rules(self) -> pattern.RewriteRuleSet: ) -class _AddRMSNormalizationToSkipRMSNormalization(_RemoveReplacedAdd, pattern.RewriteRuleClassBase): +class _AddRMSNormalizationToSkipRMSNormalization(ReplacedAddCleanupMixin, pattern.RewriteRuleClassBase): def pattern(self, op, add_output, weight): return op.RMSNormalization( add_output, diff --git a/test/passes/onnx/graph_surgery_test_utils.py b/test/passes/onnx/graph_surgery_test_utils.py new file mode 100644 index 0000000000..87b86ba7c8 --- /dev/null +++ b/test/passes/onnx/graph_surgery_test_utils.py @@ -0,0 +1,32 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +from collections import Counter + +import onnx +import onnx_ir as ir + +from olive.model import ONNXModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.onnx.graph_surgeries import GraphSurgeries + + +def run_surgery(model, tmp_path, surgeon, name): + model_path = tmp_path / f"{name}.onnx" + onnx.save(model, model_path) + graph_surgeries = create_pass_from_dict( + GraphSurgeries, + {"surgeries": [{"surgeon": surgeon}], "remove_duplicate_initializers": False}, + disable_search=True, + ) + output_model = graph_surgeries.run(ONNXModelHandler(model_path=str(model_path)), tmp_path / f"{name}_output") + output = output_model.load_model() + onnx.checker.check_model(output) + return output + + +def count_ops(model): + return Counter(node.op_type for node in ir.from_proto(model).graph) diff --git a/test/passes/onnx/test_graph_surgeries_activations.py b/test/passes/onnx/test_graph_surgeries_activations.py index 00dad053e9..0bffbb8ced 100644 --- a/test/passes/onnx/test_graph_surgeries_activations.py +++ b/test/passes/onnx/test_graph_surgeries_activations.py @@ -17,31 +17,12 @@ from olive.passes.onnx.graph_surgeries import GraphSurgeries from olive.passes.onnx.graph_surgery.activations import FuseBiasGelu, FuseGelu from olive.passes.onnx.graph_surgery.base import Surgeon +from test.passes.onnx.graph_surgery_test_utils import count_ops as _count_ops +from test.passes.onnx.graph_surgery_test_utils import run_surgery as _run_surgery _SQRT_2 = math.sqrt(2.0) -def _run_surgery(model, tmp_path, surgeon, name): - model_path = tmp_path / f"{name}.onnx" - onnx.save(model, model_path) - graph_surgeries = create_pass_from_dict( - GraphSurgeries, - {"surgeries": [{"surgeon": surgeon}], "remove_duplicate_initializers": False}, - disable_search=True, - ) - output_model = graph_surgeries.run(ONNXModelHandler(model_path=str(model_path)), tmp_path / f"{name}_output") - output = output_model.load_model() - onnx.checker.check_model(output) - return output - - -def _count_ops(model): - return { - op_type: sum(node.op_type == op_type for node in model.graph.node) - for op_type in {n.op_type for n in model.graph.node} - } - - def _make_model(nodes, initializers, outputs, *, input_shape=(1, 4, 8), opset=21): x = helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape) graph_outputs = [helper.make_tensor_value_info(name, TensorProto.FLOAT, input_shape) for name in outputs] diff --git a/test/passes/onnx/test_graph_surgeries_attention.py b/test/passes/onnx/test_graph_surgeries_attention.py index a9b75d926f..862c49e98b 100644 --- a/test/passes/onnx/test_graph_surgeries_attention.py +++ b/test/passes/onnx/test_graph_surgeries_attention.py @@ -10,7 +10,6 @@ import pytest from onnxscript import ir -import olive.passes.onnx.graph_surgery.attention # noqa: F401 from olive.model import ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.graph_surgeries import GraphSurgeries diff --git a/test/passes/onnx/test_graph_surgeries_moe.py b/test/passes/onnx/test_graph_surgeries_moe.py index 5f2e15d2e6..2471d51b1f 100644 --- a/test/passes/onnx/test_graph_surgeries_moe.py +++ b/test/passes/onnx/test_graph_surgeries_moe.py @@ -365,7 +365,7 @@ def has_zero_points(expert: int, role: str) -> bool: def _make_native_weight(rng: np.random.Generator, out_features: int, in_features: int, fmt: str) -> np.ndarray: - block_elements, block_bytes = moe_surgeries._NATIVE_BLOCK_FORMATS[fmt] + block_elements, block_bytes = moe_surgeries._NATIVE_BLOCK_FORMATS[fmt] # pylint: disable=protected-access blocks = (in_features + block_elements - 1) // block_elements return rng.integers(0, 256, size=(out_features, blocks, block_bytes), dtype=np.uint8) diff --git a/test/passes/onnx/test_graph_surgeries_normalization.py b/test/passes/onnx/test_graph_surgeries_normalization.py index 4f03a0ad59..153795b913 100644 --- a/test/passes/onnx/test_graph_surgeries_normalization.py +++ b/test/passes/onnx/test_graph_surgeries_normalization.py @@ -5,40 +5,17 @@ from __future__ import annotations import numpy as np -import onnx import pytest from onnx import TensorProto, helper, numpy_helper -from olive.model import ONNXModelHandler -from olive.passes.olive_pass import create_pass_from_dict -from olive.passes.onnx.graph_surgeries import GraphSurgeries from olive.passes.onnx.graph_surgery.base import Surgeon from olive.passes.onnx.graph_surgery.normalization import ( FuseLayerNormalization, FuseSkipLayerNormalization, FuseSkipRMSNormalization, ) - - -def _run_surgery(model, tmp_path, surgeon, name): - model_path = tmp_path / f"{name}.onnx" - onnx.save(model, model_path) - graph_surgeries = create_pass_from_dict( - GraphSurgeries, - {"surgeries": [{"surgeon": surgeon}], "remove_duplicate_initializers": False}, - disable_search=True, - ) - output_model = graph_surgeries.run(ONNXModelHandler(model_path=str(model_path)), tmp_path / f"{name}_output") - output = output_model.load_model() - onnx.checker.check_model(output) - return output - - -def _count_ops(model): - return { - op_type: sum(node.op_type == op_type for node in model.graph.node) - for op_type in {n.op_type for n in model.graph.node} - } +from test.passes.onnx.graph_surgery_test_utils import count_ops as _count_ops +from test.passes.onnx.graph_surgery_test_utils import run_surgery as _run_surgery def _build_decomposed_layer_normalization(*, include_bias=True, axes=-1, exponent=2.0, epsilon=1e-5): diff --git a/test/passes/onnx/test_graph_surgeries_pipeline.py b/test/passes/onnx/test_graph_surgeries_pipeline.py new file mode 100644 index 0000000000..ea4820144b --- /dev/null +++ b/test/passes/onnx/test_graph_surgeries_pipeline.py @@ -0,0 +1,305 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import pytest +from onnxruntime import GraphOptimizationLevel as OptLevel +from onnxruntime import InferenceSession, SessionOptions + +from olive.model import ONNXModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.onnx.graph_surgeries import GraphSurgeries + +_MIGRATED_SURGEONS = ( + "FuseGelu", + "FuseBiasGelu", + "FuseLayerNormalization", + "FuseSkipLayerNormalization", + "FuseSkipRMSNormalization", + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "SeparateGroupQueryAttentionRoPE", + "UnpackGroupQueryAttentionQKV", + "BlockDiagonalAttentionToPackedMHA", + "ClipToMinMax", + "Rank4RMSNormToRank3", + "DecomposeOnnxRotaryEmbedding", + "TensorScatterToScatterND", + "DecomposeAttention", + "StaticEmptyKV", + "FuseDenseMoEToQMoE", + "FuseBlockQuantizedMoE", +) + + +def test_graph_surgeries_registers_builtins_without_exporter_imports(): + # A fresh interpreter prevents test collection from masking missing production imports. + script = """ +import sys +from importlib.abc import MetaPathFinder + +class BlockExporterImports(MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "mobius" or fullname.startswith("mobius."): + raise ImportError("Graph surgeries must not import Mobius") + return None + +sys.meta_path.insert(0, BlockExporterImports()) +from olive.passes.onnx.graph_surgeries import GraphSurgeries, Surgeon, RewriteRuleSurgeon, ProtoSurgeon +from olive.passes.onnx.graph_surgery.base import Surgeon as BaseSurgeon +from olive.passes.olive_pass import create_pass_from_dict + +assert Surgeon is BaseSurgeon +p = create_pass_from_dict(GraphSurgeries, {"surgeries": []}, disable_search=True) +for name in sys.argv[1:]: + instance = p.init_surgeon_instance({"surgeon": name.swapcase()}) + assert isinstance(instance, Surgeon), name + assert type(instance).__name__ == name +assert isinstance(p.init_surgeon_instance({"surgeon": "ReplaceErfWithTanh"}), RewriteRuleSurgeon) +assert isinstance(p.init_surgeon_instance({"surgeon": "DeduplicateNodes"}), ProtoSurgeon) +assert type(p.init_surgeon_instance({"surgeon": "DecomposeRotaryEmbedding"})).__module__.endswith("graph_surgeries") +assert not any(name == "mobius" or name.startswith("mobius.") for name in sys.modules) +""" + result = subprocess.run( + [sys.executable, "-c", script, *_MIGRATED_SURGEONS], + cwd=Path(__file__).resolve().parents[3], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def _run_pipeline(tmp_path, model, surgeons): + input_path = tmp_path / "input.onnx" + ir.save(model, input_path) + p = create_pass_from_dict( + GraphSurgeries, + {"surgeries": [{"surgeon": surgeon} for surgeon in surgeons]}, + disable_search=True, + ) + return p.run(ONNXModelHandler(model_path=str(input_path)), tmp_path / "output") + + +def _assert_pipeline_numerics(tmp_path, output, feeds): + options = SessionOptions() + options.graph_optimization_level = OptLevel.ORT_DISABLE_ALL + expected = InferenceSession( + str(tmp_path / "input.onnx"), sess_options=options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + actual = InferenceSession(output.model_path, sess_options=options, providers=["CPUExecutionProvider"]).run( + None, feeds + ) + for actual_output, expected_output in zip(actual, expected): + np.testing.assert_allclose(actual_output, expected_output, rtol=1e-6, atol=1e-6) + + +def test_graph_surgeries_preserves_nonmatching_model_and_metadata(tmp_path): + x = ir.val("x", dtype=ir.DataType.FLOAT, shape=["batch", 8]) + y = ir.val("y", dtype=ir.DataType.FLOAT, shape=["batch", 8]) + model = ir.Model( + ir.Graph([x], [y], nodes=[ir.Node("", "Relu", [x], outputs=[y])], opset_imports={"": 24}), + ir_version=10, + ) + model.metadata_props["pipeline.component"] = "encoder" + y.metadata_props["output.role"] = "hidden_states" + + output = ir.load(_run_pipeline(tmp_path, model, _MIGRATED_SURGEONS).model_path) + + assert [(node.domain, node.op_type) for node in output.graph] == [("", "Relu")] + assert output.metadata_props["pipeline.component"] == "encoder" + assert output.graph.outputs[0].metadata_props["output.role"] == "hidden_states" + assert [value.name for value in output.graph.inputs] == ["x"] + assert [value.name for value in output.graph.outputs] == ["y"] + assert output.graph.outputs[0].shape == y.shape + assert output.graph.outputs[0].dtype == y.dtype + + +def _exact_gelu_with_bias(): + x = ir.val("x", dtype=ir.DataType.FLOAT, shape=[2, 8]) + bias = ir.val("bias", const_value=ir.tensor(np.linspace(-0.5, 0.5, 8, dtype=np.float32))) + divisor = ir.val("divisor", const_value=ir.tensor(np.array(np.sqrt(2), dtype=np.float32))) + one = ir.val("one", const_value=ir.tensor(np.array(1, dtype=np.float32))) + half = ir.val("half", const_value=ir.tensor(np.array(0.5, dtype=np.float32))) + add = ir.Node("", "Add", [x, bias]) + add.outputs[0].shape = x.shape + add.outputs[0].dtype = x.dtype + div = ir.Node("", "Div", [add.outputs[0], divisor]) + erf = ir.Node("", "Erf", [div.outputs[0]]) + shifted = ir.Node("", "Add", [erf.outputs[0], one]) + scaled = ir.Node("", "Mul", [add.outputs[0], shifted.outputs[0]]) + y = ir.val("y", dtype=ir.DataType.FLOAT, shape=[2, 8]) + result = ir.Node("", "Mul", [scaled.outputs[0], half], outputs=[y]) + return ir.Model( + ir.Graph( + [x], + [y], + nodes=[add, div, erf, shifted, scaled, result], + initializers=[bias, divisor, one, half], + opset_imports={"": 21}, + ), + ir_version=10, + ) + + +@pytest.mark.parametrize( + ("surgeons", "expected_ops"), + [ + (["FuseGelu", "FuseBiasGelu"], ["BiasGelu"]), + (["FuseBiasGelu", "FuseGelu"], ["Add", "Gelu"]), + ], +) +def test_graph_surgeries_respects_fusion_order_and_preserves_numerics(tmp_path, surgeons, expected_ops): + output = _run_pipeline(tmp_path, _exact_gelu_with_bias(), surgeons) + model = ir.load(output.model_path) + assert [node.op_type for node in model.graph] == expected_ops + assert model.graph.outputs[0].name == "y" + feeds = {"x": np.random.default_rng(42).standard_normal((2, 8)).astype(np.float32)} + _assert_pipeline_numerics(tmp_path, output, feeds) + + +def _gqa_projections(dtype, with_bias): + hidden = ir.val("hidden", dtype=dtype, shape=[1, 2, 8]) + past_key = ir.val("past_key", dtype=dtype, shape=[1, 2, 6, 8]) + past_value = ir.val("past_value", dtype=dtype, shape=[1, 2, 6, 8]) + lengths = ir.val("seqlens", dtype=ir.DataType.INT32, shape=[1]) + total = ir.val("total", dtype=ir.DataType.INT32, shape=[]) + outputs = [ + ir.val("output", dtype=dtype, shape=[1, 2, 32]), + ir.val("present_key", dtype=dtype, shape=[1, 2, 6, 8]), + ir.val("present_value", dtype=dtype, shape=[1, 2, 6, 8]), + ] + rng = np.random.default_rng(7) + nodes, initializers, projections = [], [], [] + for name, width in (("q", 32), ("k", 16), ("v", 16)): + array = (rng.standard_normal((width, 8)) * 0.1).astype(dtype.numpy()) + weight = ir.val(f"{name}_weight", const_value=ir.tensor(array)) + initializers.append(weight) + transpose = ir.Node("", "Transpose", [weight], attributes=[ir.AttrInt64s("perm", [1, 0])]) + matmul = ir.Node("", "MatMul", [hidden, transpose.outputs[0]]) + nodes.extend([transpose, matmul]) + projection = matmul.outputs[0] + if with_bias: + bias = ir.val(f"{name}_bias", const_value=ir.tensor(rng.standard_normal(width).astype(dtype.numpy()))) + initializers.append(bias) + add = ir.Node("", "Add", [projection, bias]) + nodes.append(add) + projection = add.outputs[0] + projections.append(projection) + nodes.append( + ir.Node( + "com.microsoft", + "GroupQueryAttention", + [*projections, past_key, past_value, lengths, total, None, None], + outputs=outputs, + attributes=[ + ir.AttrInt64("num_heads", 4), + ir.AttrInt64("kv_num_heads", 2), + ir.AttrInt64("do_rotary", 0), + ir.AttrFloat32("scale", 0.5), + ], + ) + ) + return ir.Model( + ir.Graph( + [hidden, past_key, past_value, lengths, total], + outputs, + nodes=nodes, + initializers=initializers, + opset_imports={"": 24, "com.microsoft": 1}, + ), + ir_version=10, + ) + + +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_graph_surgeries_pack_unpack_preserves_weight_values_and_cache_contract(tmp_path, dtype, with_bias): + original = _gqa_projections(dtype, with_bias) + output = _run_pipeline(tmp_path, original, ["PackQKVForGroupQueryAttention", "UnpackGroupQueryAttentionQKV"]) + rewritten = ir.load(output.model_path) + gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + assert [value.name if value is not None else None for value in gqa.inputs[3:]] == [ + "past_key", + "past_value", + "seqlens", + "total", + None, + None, + ] + for actual, expected in zip(rewritten.graph.outputs, original.graph.outputs): + assert (actual.name, actual.dtype, actual.shape) == (expected.name, expected.dtype, expected.shape) + for index, name in enumerate(("q", "k", "v")): + projection = gqa.inputs[index].producer() + if with_bias: + assert projection.op_type == "Add" + actual_bias = projection.inputs[1].const_value + expected_bias = original.graph.initializers[f"{name}_bias"].const_value + assert actual_bias.tobytes() == expected_bias.tobytes() + projection = projection.inputs[0].producer() + assert projection.op_type == "MatMul" + weight = projection.inputs[1].producer().inputs[0] + assert weight.dtype == dtype + assert weight.const_value.tobytes() == original.graph.initializers[f"{name}_weight"].const_value.tobytes() + + if dtype == ir.DataType.FLOAT: + rng = np.random.default_rng(17) + feeds = { + "hidden": rng.standard_normal((1, 2, 8)).astype(np.float32), + "past_key": rng.standard_normal((1, 2, 6, 8)).astype(np.float32), + "past_value": rng.standard_normal((1, 2, 6, 8)).astype(np.float32), + "seqlens": np.array([5], dtype=np.int32), + "total": np.array(6, dtype=np.int32), + } + _assert_pipeline_numerics(tmp_path, output, feeds) + + +@pytest.mark.parametrize( + ("norm_op", "surgeon", "fused_op"), + [ + ("LayerNormalization", "FuseSkipLayerNormalization", "SkipLayerNormalization"), + ("RMSNormalization", "FuseSkipRMSNormalization", "SkipSimplifiedLayerNormalization"), + ], +) +@pytest.mark.parametrize("skip_shape", [(4, 8), (1, 4, 8)]) +def test_graph_surgeries_skip_broadcast_preserves_numerics_and_residual( + tmp_path, norm_op, surgeon, fused_op, skip_shape +): + x = ir.val("x", dtype=ir.DataType.FLOAT, shape=[2, 4, 8]) + skip = ir.val("skip", dtype=ir.DataType.FLOAT, shape=skip_shape) + weight = ir.val("weight", const_value=ir.tensor(np.linspace(0.5, 1.5, 8, dtype=np.float32))) + add = ir.Node("", "Add", [skip, x]) + y = ir.val("y", dtype=ir.DataType.FLOAT, shape=x.shape) + residual = ir.val("residual", dtype=ir.DataType.FLOAT, shape=x.shape) + model = ir.Model( + ir.Graph( + [x, skip], + [y, residual], + nodes=[ + add, + ir.Node("", norm_op, [add.outputs[0], weight], outputs=[y]), + ir.Node("", "Identity", [add.outputs[0]], outputs=[residual]), + ], + initializers=[weight], + opset_imports={"": 24}, + ), + ir_version=10, + ) + output = _run_pipeline(tmp_path, model, [surgeon]) + rewritten = ir.load(output.model_path) + assert sum(node.op_type == fused_op for node in rewritten.graph) == 1 + + rng = np.random.default_rng(23) + feeds = { + "x": rng.standard_normal((2, 4, 8)).astype(np.float32), + "skip": rng.standard_normal(skip_shape).astype(np.float32), + } + _assert_pipeline_numerics(tmp_path, output, feeds) From cf5f5691004cc0eda0231246e9187ba98fec0cbc Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 8 Sep 2026 13:14:19 -0700 Subject: [PATCH 11/14] Focus graph surgery unit tests on individual surgeons Remove redundant multi-surgery pipeline cases and retain single-surgeon contracts plus the registry smoke test. Split RoPE domain checks into independent parameterized cases and cover unrelated-node rejection directly, preserving the exact set of covered production lines. Signed-off-by: Xiaoyu Zhang --- .../onnx/test_graph_surgeries_activations.py | 27 -- .../onnx/test_graph_surgeries_contracts.py | 145 +++++++++ .../onnx/test_graph_surgeries_lowering.py | 44 +-- .../onnx/test_graph_surgeries_pipeline.py | 305 ------------------ 4 files changed, 171 insertions(+), 350 deletions(-) create mode 100644 test/passes/onnx/test_graph_surgeries_contracts.py delete mode 100644 test/passes/onnx/test_graph_surgeries_pipeline.py diff --git a/test/passes/onnx/test_graph_surgeries_activations.py b/test/passes/onnx/test_graph_surgeries_activations.py index 0bffbb8ced..9ee79a84e2 100644 --- a/test/passes/onnx/test_graph_surgeries_activations.py +++ b/test/passes/onnx/test_graph_surgeries_activations.py @@ -7,14 +7,10 @@ import math import numpy as np -import onnx import pytest from onnx import TensorProto, helper, numpy_helper from onnxruntime import InferenceSession -from olive.model import ONNXModelHandler -from olive.passes.olive_pass import create_pass_from_dict -from olive.passes.onnx.graph_surgeries import GraphSurgeries from olive.passes.onnx.graph_surgery.activations import FuseBiasGelu, FuseGelu from olive.passes.onnx.graph_surgery.base import Surgeon from test.passes.onnx.graph_surgery_test_utils import count_ops as _count_ops @@ -213,26 +209,3 @@ def test_fuse_bias_gelu_matches_exact_gelu_numerically(tmp_path): actual = InferenceSession(rewritten.SerializeToString(), providers=["CPUExecutionProvider"]).run(None, inputs)[0] np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) - - -def test_fuse_gelu_then_bias_gelu_through_public_pass(tmp_path): - bias = numpy_helper.from_array(np.ones(8, dtype=np.float32), name="bias") - nodes, initializers = _exact_gelu_parts(input_name="add_output") - model = _make_model([helper.make_node("Add", ["x", "bias"], ["add_output"]), *nodes], [bias, *initializers], ["y"]) - model_path = tmp_path / "combined.onnx" - onnx.save(model, model_path) - graph_surgeries = create_pass_from_dict( - GraphSurgeries, - { - "surgeries": [{"surgeon": "FuseGelu"}, {"surgeon": "FuseBiasGelu"}], - "remove_duplicate_initializers": False, - }, - disable_search=True, - ) - - rewritten = graph_surgeries.run( - ONNXModelHandler(model_path=str(model_path)), tmp_path / "combined_output" - ).load_model() - - onnx.checker.check_model(rewritten) - assert _count_ops(rewritten) == {"BiasGelu": 1} diff --git a/test/passes/onnx/test_graph_surgeries_contracts.py b/test/passes/onnx/test_graph_surgeries_contracts.py new file mode 100644 index 0000000000..344b07cba2 --- /dev/null +++ b/test/passes/onnx/test_graph_surgeries_contracts.py @@ -0,0 +1,145 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import pytest +from onnxruntime import GraphOptimizationLevel as OptLevel +from onnxruntime import InferenceSession, SessionOptions + +from olive.model import ONNXModelHandler +from olive.passes.olive_pass import create_pass_from_dict +from olive.passes.onnx.graph_surgeries import GraphSurgeries + +_MIGRATED_SURGEONS = ( + "FuseGelu", + "FuseBiasGelu", + "FuseLayerNormalization", + "FuseSkipLayerNormalization", + "FuseSkipRMSNormalization", + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "SeparateGroupQueryAttentionRoPE", + "UnpackGroupQueryAttentionQKV", + "BlockDiagonalAttentionToPackedMHA", + "ClipToMinMax", + "Rank4RMSNormToRank3", + "DecomposeOnnxRotaryEmbedding", + "TensorScatterToScatterND", + "DecomposeAttention", + "StaticEmptyKV", + "FuseDenseMoEToQMoE", + "FuseBlockQuantizedMoE", +) + + +def test_graph_surgeries_registers_builtins_without_exporter_imports(): + # A fresh interpreter prevents test collection from masking missing production imports. + script = """ +import sys +from importlib.abc import MetaPathFinder + +class BlockExporterImports(MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "mobius" or fullname.startswith("mobius."): + raise ImportError("Graph surgeries must not import Mobius") + return None + +sys.meta_path.insert(0, BlockExporterImports()) +from olive.passes.onnx.graph_surgeries import GraphSurgeries, Surgeon, RewriteRuleSurgeon, ProtoSurgeon +from olive.passes.onnx.graph_surgery.base import Surgeon as BaseSurgeon +from olive.passes.olive_pass import create_pass_from_dict + +assert Surgeon is BaseSurgeon +p = create_pass_from_dict(GraphSurgeries, {"surgeries": []}, disable_search=True) +for name in sys.argv[1:]: + instance = p.init_surgeon_instance({"surgeon": name.swapcase()}) + assert isinstance(instance, Surgeon), name + assert type(instance).__name__ == name +assert isinstance(p.init_surgeon_instance({"surgeon": "ReplaceErfWithTanh"}), RewriteRuleSurgeon) +assert isinstance(p.init_surgeon_instance({"surgeon": "DeduplicateNodes"}), ProtoSurgeon) +assert type(p.init_surgeon_instance({"surgeon": "DecomposeRotaryEmbedding"})).__module__.endswith("graph_surgeries") +assert not any(name == "mobius" or name.startswith("mobius.") for name in sys.modules) +""" + result = subprocess.run( + [sys.executable, "-c", script, *_MIGRATED_SURGEONS], + cwd=Path(__file__).resolve().parents[3], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def _run_surgery(tmp_path, model, surgeon): + input_path = tmp_path / "input.onnx" + ir.save(model, input_path) + p = create_pass_from_dict( + GraphSurgeries, + {"surgeries": [{"surgeon": surgeon}]}, + disable_search=True, + ) + return p.run(ONNXModelHandler(model_path=str(input_path)), tmp_path / "output") + + +def _assert_surgery_numerics(tmp_path, output, feeds): + options = SessionOptions() + options.graph_optimization_level = OptLevel.ORT_DISABLE_ALL + expected = InferenceSession( + str(tmp_path / "input.onnx"), sess_options=options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + actual = InferenceSession(output.model_path, sess_options=options, providers=["CPUExecutionProvider"]).run( + None, feeds + ) + for actual_output, expected_output in zip(actual, expected): + np.testing.assert_allclose(actual_output, expected_output, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize( + ("norm_op", "surgeon", "fused_op"), + [ + ("LayerNormalization", "FuseSkipLayerNormalization", "SkipLayerNormalization"), + ("RMSNormalization", "FuseSkipRMSNormalization", "SkipSimplifiedLayerNormalization"), + ], +) +@pytest.mark.parametrize("skip_shape", [(4, 8), (1, 4, 8)]) +def test_graph_surgeries_skip_broadcast_preserves_numerics_and_residual( + tmp_path, norm_op, surgeon, fused_op, skip_shape +): + x = ir.val("x", dtype=ir.DataType.FLOAT, shape=[2, 4, 8]) + skip = ir.val("skip", dtype=ir.DataType.FLOAT, shape=skip_shape) + weight = ir.val("weight", const_value=ir.tensor(np.linspace(0.5, 1.5, 8, dtype=np.float32))) + add = ir.Node("", "Add", [skip, x]) + y = ir.val("y", dtype=ir.DataType.FLOAT, shape=x.shape) + residual = ir.val("residual", dtype=ir.DataType.FLOAT, shape=x.shape) + model = ir.Model( + ir.Graph( + [x, skip], + [y, residual], + nodes=[ + add, + ir.Node("", norm_op, [add.outputs[0], weight], outputs=[y]), + ir.Node("", "Identity", [add.outputs[0]], outputs=[residual]), + ], + initializers=[weight], + opset_imports={"": 24}, + ), + ir_version=10, + ) + output = _run_surgery(tmp_path, model, surgeon) + rewritten = ir.load(output.model_path) + assert sum(node.op_type == fused_op for node in rewritten.graph) == 1 + + rng = np.random.default_rng(23) + feeds = { + "x": rng.standard_normal((2, 4, 8)).astype(np.float32), + "skip": rng.standard_normal(skip_shape).astype(np.float32), + } + _assert_surgery_numerics(tmp_path, output, feeds) diff --git a/test/passes/onnx/test_graph_surgeries_lowering.py b/test/passes/onnx/test_graph_surgeries_lowering.py index 42e485e58d..4c3d3a84bc 100644 --- a/test/passes/onnx/test_graph_surgeries_lowering.py +++ b/test/passes/onnx/test_graph_surgeries_lowering.py @@ -314,24 +314,17 @@ def test_decompose_onnx_rotary_embedding_rejects_unsupported_forms( assert _counts(lowered)["RotaryEmbedding"] == 1 -def test_standard_and_microsoft_rotary_embedding_surgeries_are_distinct(tmp_path): - microsoft_model = _rope_model(domain="com.microsoft") - lowered_microsoft = _apply_surgery( - tmp_path, - microsoft_model, - "DecomposeOnnxRotaryEmbedding", - case="microsoft", - ) - assert _counts(lowered_microsoft)["RotaryEmbedding"] == 1 - - standard_model = _rope_model() - lowered_standard = _apply_surgery( - tmp_path, - standard_model, - "DecomposeRotaryEmbedding", - case="standard", - ) - assert _counts(lowered_standard)["RotaryEmbedding"] == 1 +@pytest.mark.parametrize( + ("surgeon", "domain"), + [ + ("DecomposeOnnxRotaryEmbedding", "com.microsoft"), + ("DecomposeRotaryEmbedding", ""), + ], +) +def test_rotary_embedding_surgery_preserves_other_domain(tmp_path, surgeon, domain): + model = _rope_model(domain=domain) + lowered = _apply_surgery(tmp_path, model, surgeon) + assert _counts(lowered)["RotaryEmbedding"] == 1 def _tensor_scatter_model( @@ -785,6 +778,21 @@ def test_decompose_attention_rejects_unknown_mask_shape(tmp_path): assert _counts(lowered)["Attention"] == 1 +def test_decompose_attention_preserves_unrelated_nodes(tmp_path): + x = ir.val("x", dtype=ir.DataType.FLOAT, shape=["batch", 8]) + y = ir.val("y", dtype=ir.DataType.FLOAT, shape=["batch", 8]) + model = ir.Model( + ir.Graph([x], [y], nodes=[ir.Node("", "Relu", [x], outputs=[y])], opset_imports={"": 24}), + ir_version=10, + ) + + lowered = _apply_surgery(tmp_path, model, "DecomposeAttention") + + assert _counts(lowered) == {"Relu": 1} + assert _metadata(lowered.graph.inputs[0]) == _metadata(x) + assert _metadata(lowered.graph.outputs[0]) == _metadata(y) + + def _empty_kv_model( kv_hidden=128, *, diff --git a/test/passes/onnx/test_graph_surgeries_pipeline.py b/test/passes/onnx/test_graph_surgeries_pipeline.py deleted file mode 100644 index ea4820144b..0000000000 --- a/test/passes/onnx/test_graph_surgeries_pipeline.py +++ /dev/null @@ -1,305 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -from __future__ import annotations - -import subprocess -import sys -from pathlib import Path - -import numpy as np -import onnx_ir as ir -import pytest -from onnxruntime import GraphOptimizationLevel as OptLevel -from onnxruntime import InferenceSession, SessionOptions - -from olive.model import ONNXModelHandler -from olive.passes.olive_pass import create_pass_from_dict -from olive.passes.onnx.graph_surgeries import GraphSurgeries - -_MIGRATED_SURGEONS = ( - "FuseGelu", - "FuseBiasGelu", - "FuseLayerNormalization", - "FuseSkipLayerNormalization", - "FuseSkipRMSNormalization", - "AttentionToGroupQueryAttention", - "PackQKVForGroupQueryAttention", - "SeparateGroupQueryAttentionRoPE", - "UnpackGroupQueryAttentionQKV", - "BlockDiagonalAttentionToPackedMHA", - "ClipToMinMax", - "Rank4RMSNormToRank3", - "DecomposeOnnxRotaryEmbedding", - "TensorScatterToScatterND", - "DecomposeAttention", - "StaticEmptyKV", - "FuseDenseMoEToQMoE", - "FuseBlockQuantizedMoE", -) - - -def test_graph_surgeries_registers_builtins_without_exporter_imports(): - # A fresh interpreter prevents test collection from masking missing production imports. - script = """ -import sys -from importlib.abc import MetaPathFinder - -class BlockExporterImports(MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == "mobius" or fullname.startswith("mobius."): - raise ImportError("Graph surgeries must not import Mobius") - return None - -sys.meta_path.insert(0, BlockExporterImports()) -from olive.passes.onnx.graph_surgeries import GraphSurgeries, Surgeon, RewriteRuleSurgeon, ProtoSurgeon -from olive.passes.onnx.graph_surgery.base import Surgeon as BaseSurgeon -from olive.passes.olive_pass import create_pass_from_dict - -assert Surgeon is BaseSurgeon -p = create_pass_from_dict(GraphSurgeries, {"surgeries": []}, disable_search=True) -for name in sys.argv[1:]: - instance = p.init_surgeon_instance({"surgeon": name.swapcase()}) - assert isinstance(instance, Surgeon), name - assert type(instance).__name__ == name -assert isinstance(p.init_surgeon_instance({"surgeon": "ReplaceErfWithTanh"}), RewriteRuleSurgeon) -assert isinstance(p.init_surgeon_instance({"surgeon": "DeduplicateNodes"}), ProtoSurgeon) -assert type(p.init_surgeon_instance({"surgeon": "DecomposeRotaryEmbedding"})).__module__.endswith("graph_surgeries") -assert not any(name == "mobius" or name.startswith("mobius.") for name in sys.modules) -""" - result = subprocess.run( - [sys.executable, "-c", script, *_MIGRATED_SURGEONS], - cwd=Path(__file__).resolve().parents[3], - capture_output=True, - text=True, - check=False, - ) - assert result.returncode == 0, result.stdout + result.stderr - - -def _run_pipeline(tmp_path, model, surgeons): - input_path = tmp_path / "input.onnx" - ir.save(model, input_path) - p = create_pass_from_dict( - GraphSurgeries, - {"surgeries": [{"surgeon": surgeon} for surgeon in surgeons]}, - disable_search=True, - ) - return p.run(ONNXModelHandler(model_path=str(input_path)), tmp_path / "output") - - -def _assert_pipeline_numerics(tmp_path, output, feeds): - options = SessionOptions() - options.graph_optimization_level = OptLevel.ORT_DISABLE_ALL - expected = InferenceSession( - str(tmp_path / "input.onnx"), sess_options=options, providers=["CPUExecutionProvider"] - ).run(None, feeds) - actual = InferenceSession(output.model_path, sess_options=options, providers=["CPUExecutionProvider"]).run( - None, feeds - ) - for actual_output, expected_output in zip(actual, expected): - np.testing.assert_allclose(actual_output, expected_output, rtol=1e-6, atol=1e-6) - - -def test_graph_surgeries_preserves_nonmatching_model_and_metadata(tmp_path): - x = ir.val("x", dtype=ir.DataType.FLOAT, shape=["batch", 8]) - y = ir.val("y", dtype=ir.DataType.FLOAT, shape=["batch", 8]) - model = ir.Model( - ir.Graph([x], [y], nodes=[ir.Node("", "Relu", [x], outputs=[y])], opset_imports={"": 24}), - ir_version=10, - ) - model.metadata_props["pipeline.component"] = "encoder" - y.metadata_props["output.role"] = "hidden_states" - - output = ir.load(_run_pipeline(tmp_path, model, _MIGRATED_SURGEONS).model_path) - - assert [(node.domain, node.op_type) for node in output.graph] == [("", "Relu")] - assert output.metadata_props["pipeline.component"] == "encoder" - assert output.graph.outputs[0].metadata_props["output.role"] == "hidden_states" - assert [value.name for value in output.graph.inputs] == ["x"] - assert [value.name for value in output.graph.outputs] == ["y"] - assert output.graph.outputs[0].shape == y.shape - assert output.graph.outputs[0].dtype == y.dtype - - -def _exact_gelu_with_bias(): - x = ir.val("x", dtype=ir.DataType.FLOAT, shape=[2, 8]) - bias = ir.val("bias", const_value=ir.tensor(np.linspace(-0.5, 0.5, 8, dtype=np.float32))) - divisor = ir.val("divisor", const_value=ir.tensor(np.array(np.sqrt(2), dtype=np.float32))) - one = ir.val("one", const_value=ir.tensor(np.array(1, dtype=np.float32))) - half = ir.val("half", const_value=ir.tensor(np.array(0.5, dtype=np.float32))) - add = ir.Node("", "Add", [x, bias]) - add.outputs[0].shape = x.shape - add.outputs[0].dtype = x.dtype - div = ir.Node("", "Div", [add.outputs[0], divisor]) - erf = ir.Node("", "Erf", [div.outputs[0]]) - shifted = ir.Node("", "Add", [erf.outputs[0], one]) - scaled = ir.Node("", "Mul", [add.outputs[0], shifted.outputs[0]]) - y = ir.val("y", dtype=ir.DataType.FLOAT, shape=[2, 8]) - result = ir.Node("", "Mul", [scaled.outputs[0], half], outputs=[y]) - return ir.Model( - ir.Graph( - [x], - [y], - nodes=[add, div, erf, shifted, scaled, result], - initializers=[bias, divisor, one, half], - opset_imports={"": 21}, - ), - ir_version=10, - ) - - -@pytest.mark.parametrize( - ("surgeons", "expected_ops"), - [ - (["FuseGelu", "FuseBiasGelu"], ["BiasGelu"]), - (["FuseBiasGelu", "FuseGelu"], ["Add", "Gelu"]), - ], -) -def test_graph_surgeries_respects_fusion_order_and_preserves_numerics(tmp_path, surgeons, expected_ops): - output = _run_pipeline(tmp_path, _exact_gelu_with_bias(), surgeons) - model = ir.load(output.model_path) - assert [node.op_type for node in model.graph] == expected_ops - assert model.graph.outputs[0].name == "y" - feeds = {"x": np.random.default_rng(42).standard_normal((2, 8)).astype(np.float32)} - _assert_pipeline_numerics(tmp_path, output, feeds) - - -def _gqa_projections(dtype, with_bias): - hidden = ir.val("hidden", dtype=dtype, shape=[1, 2, 8]) - past_key = ir.val("past_key", dtype=dtype, shape=[1, 2, 6, 8]) - past_value = ir.val("past_value", dtype=dtype, shape=[1, 2, 6, 8]) - lengths = ir.val("seqlens", dtype=ir.DataType.INT32, shape=[1]) - total = ir.val("total", dtype=ir.DataType.INT32, shape=[]) - outputs = [ - ir.val("output", dtype=dtype, shape=[1, 2, 32]), - ir.val("present_key", dtype=dtype, shape=[1, 2, 6, 8]), - ir.val("present_value", dtype=dtype, shape=[1, 2, 6, 8]), - ] - rng = np.random.default_rng(7) - nodes, initializers, projections = [], [], [] - for name, width in (("q", 32), ("k", 16), ("v", 16)): - array = (rng.standard_normal((width, 8)) * 0.1).astype(dtype.numpy()) - weight = ir.val(f"{name}_weight", const_value=ir.tensor(array)) - initializers.append(weight) - transpose = ir.Node("", "Transpose", [weight], attributes=[ir.AttrInt64s("perm", [1, 0])]) - matmul = ir.Node("", "MatMul", [hidden, transpose.outputs[0]]) - nodes.extend([transpose, matmul]) - projection = matmul.outputs[0] - if with_bias: - bias = ir.val(f"{name}_bias", const_value=ir.tensor(rng.standard_normal(width).astype(dtype.numpy()))) - initializers.append(bias) - add = ir.Node("", "Add", [projection, bias]) - nodes.append(add) - projection = add.outputs[0] - projections.append(projection) - nodes.append( - ir.Node( - "com.microsoft", - "GroupQueryAttention", - [*projections, past_key, past_value, lengths, total, None, None], - outputs=outputs, - attributes=[ - ir.AttrInt64("num_heads", 4), - ir.AttrInt64("kv_num_heads", 2), - ir.AttrInt64("do_rotary", 0), - ir.AttrFloat32("scale", 0.5), - ], - ) - ) - return ir.Model( - ir.Graph( - [hidden, past_key, past_value, lengths, total], - outputs, - nodes=nodes, - initializers=initializers, - opset_imports={"": 24, "com.microsoft": 1}, - ), - ir_version=10, - ) - - -@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) -@pytest.mark.parametrize("with_bias", [False, True]) -def test_graph_surgeries_pack_unpack_preserves_weight_values_and_cache_contract(tmp_path, dtype, with_bias): - original = _gqa_projections(dtype, with_bias) - output = _run_pipeline(tmp_path, original, ["PackQKVForGroupQueryAttention", "UnpackGroupQueryAttentionQKV"]) - rewritten = ir.load(output.model_path) - gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") - assert [value.name if value is not None else None for value in gqa.inputs[3:]] == [ - "past_key", - "past_value", - "seqlens", - "total", - None, - None, - ] - for actual, expected in zip(rewritten.graph.outputs, original.graph.outputs): - assert (actual.name, actual.dtype, actual.shape) == (expected.name, expected.dtype, expected.shape) - for index, name in enumerate(("q", "k", "v")): - projection = gqa.inputs[index].producer() - if with_bias: - assert projection.op_type == "Add" - actual_bias = projection.inputs[1].const_value - expected_bias = original.graph.initializers[f"{name}_bias"].const_value - assert actual_bias.tobytes() == expected_bias.tobytes() - projection = projection.inputs[0].producer() - assert projection.op_type == "MatMul" - weight = projection.inputs[1].producer().inputs[0] - assert weight.dtype == dtype - assert weight.const_value.tobytes() == original.graph.initializers[f"{name}_weight"].const_value.tobytes() - - if dtype == ir.DataType.FLOAT: - rng = np.random.default_rng(17) - feeds = { - "hidden": rng.standard_normal((1, 2, 8)).astype(np.float32), - "past_key": rng.standard_normal((1, 2, 6, 8)).astype(np.float32), - "past_value": rng.standard_normal((1, 2, 6, 8)).astype(np.float32), - "seqlens": np.array([5], dtype=np.int32), - "total": np.array(6, dtype=np.int32), - } - _assert_pipeline_numerics(tmp_path, output, feeds) - - -@pytest.mark.parametrize( - ("norm_op", "surgeon", "fused_op"), - [ - ("LayerNormalization", "FuseSkipLayerNormalization", "SkipLayerNormalization"), - ("RMSNormalization", "FuseSkipRMSNormalization", "SkipSimplifiedLayerNormalization"), - ], -) -@pytest.mark.parametrize("skip_shape", [(4, 8), (1, 4, 8)]) -def test_graph_surgeries_skip_broadcast_preserves_numerics_and_residual( - tmp_path, norm_op, surgeon, fused_op, skip_shape -): - x = ir.val("x", dtype=ir.DataType.FLOAT, shape=[2, 4, 8]) - skip = ir.val("skip", dtype=ir.DataType.FLOAT, shape=skip_shape) - weight = ir.val("weight", const_value=ir.tensor(np.linspace(0.5, 1.5, 8, dtype=np.float32))) - add = ir.Node("", "Add", [skip, x]) - y = ir.val("y", dtype=ir.DataType.FLOAT, shape=x.shape) - residual = ir.val("residual", dtype=ir.DataType.FLOAT, shape=x.shape) - model = ir.Model( - ir.Graph( - [x, skip], - [y, residual], - nodes=[ - add, - ir.Node("", norm_op, [add.outputs[0], weight], outputs=[y]), - ir.Node("", "Identity", [add.outputs[0]], outputs=[residual]), - ], - initializers=[weight], - opset_imports={"": 24}, - ), - ir_version=10, - ) - output = _run_pipeline(tmp_path, model, [surgeon]) - rewritten = ir.load(output.model_path) - assert sum(node.op_type == fused_op for node in rewritten.graph) == 1 - - rng = np.random.default_rng(23) - feeds = { - "x": rng.standard_normal((2, 4, 8)).astype(np.float32), - "skip": rng.standard_normal(skip_shape).astype(np.float32), - } - _assert_pipeline_numerics(tmp_path, output, feeds) From df3cbfda2c1d34ea69e17adf645518c4258a57c0 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 9 Sep 2026 13:54:04 -0700 Subject: [PATCH 12/14] Fix Mobius graph surgery review feedback --- olive/passes/onnx/graph_surgery/base.py | 2 +- olive/passes/onnx/graph_surgery/lowering.py | 2 +- .../onnx/graph_surgery/normalization.py | 6 +++-- .../test_graph_surgeries_normalization.py | 22 +++++++++++++++++-- 4 files changed, 26 insertions(+), 6 deletions(-) diff --git a/olive/passes/onnx/graph_surgery/base.py b/olive/passes/onnx/graph_surgery/base.py index 5f2a5b8787..a90933742c 100644 --- a/olive/passes/onnx/graph_surgery/base.py +++ b/olive/passes/onnx/graph_surgery/base.py @@ -44,7 +44,7 @@ def __call__(self, model: ModelProto) -> ModelProto: raise NotImplementedError def call_ir(self, model: ir.Model) -> ir.Model: - raise RuntimeError("Implement __call__ method instead of operator on onnx.ModelProto directly.") + raise RuntimeError("Implement __call__ to operate directly on onnx.ModelProto.") @staticmethod def get_node_by_name(model, name: str, match_output: bool = False): diff --git a/olive/passes/onnx/graph_surgery/lowering.py b/olive/passes/onnx/graph_surgery/lowering.py index a0913c8c03..35addbbeb4 100644 --- a/olive/passes/onnx/graph_surgery/lowering.py +++ b/olive/passes/onnx/graph_surgery/lowering.py @@ -157,7 +157,7 @@ def check(self, context, cache, scatter_out, **_): if node.domain not in ("", "ai.onnx"): return result.fail("TensorScatter is not in the standard ONNX domain") if node.attributes.get_int("axis", -2) not in (-2, 1): - return result.fail("TensorScatter axis is not 1") + return result.fail("TensorScatter axis must be 1 or -2") if node.attributes.get_string("mode", "linear") != "linear": return result.fail("TensorScatter mode is not linear") if cache.shape is None or len(cache.shape) != 3: diff --git a/olive/passes/onnx/graph_surgery/normalization.py b/olive/passes/onnx/graph_surgery/normalization.py index 6e3a440f8f..b44b843f1d 100644 --- a/olive/passes/onnx/graph_surgery/normalization.py +++ b/olive/passes/onnx/graph_surgery/normalization.py @@ -81,8 +81,10 @@ def check(self, context, exponent, epsilon, first_axes, second_axes, norm_output result = _check_layer_normalization_constants(exponent, epsilon, first_axes, second_axes) if not result: return result - if list(norm_output.uses()): - return result.fail("Bias-free LayerNormalization output has another node consumer") + + uses = list(norm_output.uses()) + if len(uses) == 1 and uses[0].node.domain in ("", "ai.onnx") and uses[0].node.op_type == "Add": + return result.fail("Bias-free LayerNormalization output is consumed by Add") return result def rewrite(self, op, x, weight, epsilon, **_): diff --git a/test/passes/onnx/test_graph_surgeries_normalization.py b/test/passes/onnx/test_graph_surgeries_normalization.py index 153795b913..f85e41facf 100644 --- a/test/passes/onnx/test_graph_surgeries_normalization.py +++ b/test/passes/onnx/test_graph_surgeries_normalization.py @@ -18,7 +18,9 @@ from test.passes.onnx.graph_surgery_test_utils import run_surgery as _run_surgery -def _build_decomposed_layer_normalization(*, include_bias=True, axes=-1, exponent=2.0, epsilon=1e-5): +def _build_decomposed_layer_normalization( + *, include_bias=True, output_consumer=False, axes=-1, exponent=2.0, epsilon=1e-5 +): x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4, 8]) y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4, 8]) initializers = [ @@ -35,11 +37,13 @@ def _build_decomposed_layer_normalization(*, include_bias=True, axes=-1, exponen helper.make_node("Add", ["variance", "epsilon"], ["variance_epsilon"]), helper.make_node("Sqrt", ["variance_epsilon"], ["standard_deviation"]), helper.make_node("Div", ["difference", "standard_deviation"], ["normalized"]), - helper.make_node("Mul", ["normalized", "weight"], ["scaled" if include_bias else "y"]), + helper.make_node("Mul", ["normalized", "weight"], ["scaled" if include_bias or output_consumer else "y"]), ] if include_bias: initializers.append(numpy_helper.from_array(np.zeros(8, dtype=np.float32), name="bias")) nodes.append(helper.make_node("Add", ["scaled", "bias"], ["y"])) + elif output_consumer: + nodes.append(helper.make_node("Identity", ["scaled"], ["y"])) graph = helper.make_graph(nodes, "layer_normalization_test", [x], [y], initializers) return helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid("", 21)]) @@ -114,6 +118,20 @@ def test_fuse_layer_normalization_fuses_bias_variants(tmp_path, include_bias): assert len(layer_norm.input) == (3 if include_bias else 2) +def test_fuse_layer_normalization_fuses_bias_free_output_with_consumer(tmp_path): + model = _run_surgery( + _build_decomposed_layer_normalization(include_bias=False, output_consumer=True), + tmp_path, + "FuseLayerNormalization", + "layer_norm_bias_free_consumer", + ) + + assert _count_ops(model) == {"Identity": 1, "LayerNormalization": 1} + layer_norm = next(node for node in model.graph.node if node.op_type == "LayerNormalization") + identity = next(node for node in model.graph.node if node.op_type == "Identity") + assert identity.input[0] == layer_norm.output[0] + + @pytest.mark.parametrize( ("kwargs", "remaining_op"), [ From fdb45f116d7cec996541032f576ac48b1911ff8c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Wed, 9 Sep 2026 15:56:00 -0700 Subject: [PATCH 13/14] Remove redundant graph surgery import test --- .../onnx/test_graph_surgeries_contracts.py | 63 ------------------- 1 file changed, 63 deletions(-) diff --git a/test/passes/onnx/test_graph_surgeries_contracts.py b/test/passes/onnx/test_graph_surgeries_contracts.py index 344b07cba2..4f5fd6bd96 100644 --- a/test/passes/onnx/test_graph_surgeries_contracts.py +++ b/test/passes/onnx/test_graph_surgeries_contracts.py @@ -4,10 +4,6 @@ # -------------------------------------------------------------------------- from __future__ import annotations -import subprocess -import sys -from pathlib import Path - import numpy as np import onnx_ir as ir import pytest @@ -18,65 +14,6 @@ from olive.passes.olive_pass import create_pass_from_dict from olive.passes.onnx.graph_surgeries import GraphSurgeries -_MIGRATED_SURGEONS = ( - "FuseGelu", - "FuseBiasGelu", - "FuseLayerNormalization", - "FuseSkipLayerNormalization", - "FuseSkipRMSNormalization", - "AttentionToGroupQueryAttention", - "PackQKVForGroupQueryAttention", - "SeparateGroupQueryAttentionRoPE", - "UnpackGroupQueryAttentionQKV", - "BlockDiagonalAttentionToPackedMHA", - "ClipToMinMax", - "Rank4RMSNormToRank3", - "DecomposeOnnxRotaryEmbedding", - "TensorScatterToScatterND", - "DecomposeAttention", - "StaticEmptyKV", - "FuseDenseMoEToQMoE", - "FuseBlockQuantizedMoE", -) - - -def test_graph_surgeries_registers_builtins_without_exporter_imports(): - # A fresh interpreter prevents test collection from masking missing production imports. - script = """ -import sys -from importlib.abc import MetaPathFinder - -class BlockExporterImports(MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == "mobius" or fullname.startswith("mobius."): - raise ImportError("Graph surgeries must not import Mobius") - return None - -sys.meta_path.insert(0, BlockExporterImports()) -from olive.passes.onnx.graph_surgeries import GraphSurgeries, Surgeon, RewriteRuleSurgeon, ProtoSurgeon -from olive.passes.onnx.graph_surgery.base import Surgeon as BaseSurgeon -from olive.passes.olive_pass import create_pass_from_dict - -assert Surgeon is BaseSurgeon -p = create_pass_from_dict(GraphSurgeries, {"surgeries": []}, disable_search=True) -for name in sys.argv[1:]: - instance = p.init_surgeon_instance({"surgeon": name.swapcase()}) - assert isinstance(instance, Surgeon), name - assert type(instance).__name__ == name -assert isinstance(p.init_surgeon_instance({"surgeon": "ReplaceErfWithTanh"}), RewriteRuleSurgeon) -assert isinstance(p.init_surgeon_instance({"surgeon": "DeduplicateNodes"}), ProtoSurgeon) -assert type(p.init_surgeon_instance({"surgeon": "DecomposeRotaryEmbedding"})).__module__.endswith("graph_surgeries") -assert not any(name == "mobius" or name.startswith("mobius.") for name in sys.modules) -""" - result = subprocess.run( - [sys.executable, "-c", script, *_MIGRATED_SURGEONS], - cwd=Path(__file__).resolve().parents[3], - capture_output=True, - text=True, - check=False, - ) - assert result.returncode == 0, result.stdout + result.stderr - def _run_surgery(tmp_path, model, surgeon): input_path = tmp_path / "input.onnx" From afe4049b075c679b6e956fc01d856fc7367ee236 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Mon, 14 Sep 2026 14:17:44 -0700 Subject: [PATCH 14/14] address comments --- docs/source/features/onnx-transformations.md | 2 +- olive/passes/onnx/graph_surgery/base.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/docs/source/features/onnx-transformations.md b/docs/source/features/onnx-transformations.md index b5e4a6241f..7d713615c9 100644 --- a/docs/source/features/onnx-transformations.md +++ b/docs/source/features/onnx-transformations.md @@ -149,7 +149,7 @@ operator is not a portable standard-ONNX optimization. | `DecomposeAttention` | Supported standard rank-3 `Attention` to scaled dot-product primitives, including GQA, cache outputs, and causal/nonpadding masks. | | `StaticEmptyKV` | Recognized dynamic empty-KV construction to a static empty tensor for graph-capture compatibility. | | `FuseDenseMoEToQMoE` | Compatible `MatMulNBits` expert banks and routing to `com.microsoft::QMoE`. | -| `FuseBlockQuantizedMoE` | Compatible native block-quantized expert banks and routing to `pkg.nxrt::BlockQuantizedMoE`; requires a runtime implementing that operator. | +| `FuseBlockQuantizedMoE` | Compatible native block-quantized expert banks; requires a runtime implementing that operator. | The pattern surgeons preserve graphs that do not match their supported forms. In particular, `DecomposeAttention` leaves a fourth QK output, an explicit diff --git a/olive/passes/onnx/graph_surgery/base.py b/olive/passes/onnx/graph_surgery/base.py index a90933742c..9a922ed840 100644 --- a/olive/passes/onnx/graph_surgery/base.py +++ b/olive/passes/onnx/graph_surgery/base.py @@ -4,6 +4,7 @@ # -------------------------------------------------------------------------- from __future__ import annotations +from abc import ABC, abstractmethod from typing import TYPE_CHECKING, ClassVar import onnx_ir as ir @@ -13,10 +14,10 @@ from onnxscript.rewriter import pattern -class Surgeon: +class Surgeon(ABC): """Base class for surgeons that operate on the ONNX IR model.""" - # Refer to https://microsoft.github.io/onnxscript/intermediate_representation/ir_api.html#onnxscript.ir.Model + # Refer to https://onnx.ai/ir-py/api/generated/onnx_ir.Model.html#onnx_ir.Model # for the IR model API. registry: ClassVar[dict[str, type[Surgeon]]] = {} @@ -26,22 +27,20 @@ def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) Surgeon.registry[cls.__name__.lower()] = cls - def __init__(self): - pass - def __call__(self, model: ModelProto) -> ModelProto: return ir.to_proto(self.call_ir(ir.from_proto(model))) + @abstractmethod def call_ir(self, model: ir.Model) -> ir.Model: - # Implement this method in subclasses to operate on the IR model. - raise NotImplementedError + """Operate on and return the ONNX IR model.""" class ProtoSurgeon(Surgeon): """Base class for surgeons that operate on the ONNX model proto directly.""" + @abstractmethod def __call__(self, model: ModelProto) -> ModelProto: - raise NotImplementedError + """Operate on and return the ONNX model proto.""" def call_ir(self, model: ir.Model) -> ir.Model: raise RuntimeError("Implement __call__ to operate directly on onnx.ModelProto.") @@ -92,8 +91,9 @@ class RewriteRuleSurgeon(Surgeon): operand commutativity, use-count bookkeeping, and dead-node cleanup. """ + @abstractmethod def rules(self) -> pattern.RewriteRuleSet: - raise NotImplementedError + """Return the rewrite rules to apply.""" def call_ir(self, model: ir.Model) -> ir.Model: self.rules().apply_to_model(model)