Migrate mobius rewrites

[Xiaoyu Z (xiaoyu-work)]

This pull request introduces a new framework for ONNX graph transformations called "surgeons," which are modular, composable graph transformation classes. The changes add a registry and base classes for surgeons, implement built-in surgeon modules for common graph fusions and lowerings, and update documentation to describe these new capabilities. This allows users to apply a sequence of explicit, configurable ONNX graph transformations (such as fusions and lowerings) in a more maintainable and extensible way.

The most important changes are:

Core framework and API refactor:

• Moves the Surgeon, ProtoSurgeon, and RewriteRuleSurgeon base classes into a new dedicated module (base.py), and updates their implementation to support registration and modular import. (olive/passes/onnx/graph_surgery/base.py, olive/passes/onnx/graph_surgeries.py) [1] [2] [3]

• Adds a new module (__init__.py) to automatically register a set of built-in surgeons on import, making them available for the GraphSurgeries pass without manual imports. (olive/passes/onnx/graph_surgery/__init__.py)

New and refactored graph transformation logic:

• Implements new surgeon modules for common ONNX graph transformations, including GELU and BiasGELU fusion, normalization fusions, attention pattern lowerings, MoE fusions, and more. These are organized by theme in their own files (e.g., activations.py, attention.py). (olive/passes/onnx/graph_surgery/activations.py, etc.)

• Adds a shared utility module for common pattern-matching and cleanup logic used by surgeons. (olive/passes/onnx/graph_surgery/_common.py)

Documentation and user guidance:

• Updates the ONNX transformations documentation to describe the new "surgeons," their usage, available built-in transformations, ordering considerations, and guidance for adding new surgeons. (docs/source/features/onnx-transformations.md)

[Copilot]

🟡 Changes recommended

The bias-free LayerNormalization fusion currently rejects matches when the output is consumed, which will prevent expected fusions in real models and should be corrected before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR introduces a new “surgeons” framework for ONNX graph transformations in Olive, providing modular, composable transformation classes with automatic registration and an updated GraphSurgeries workflow to apply ordered rewrite sequences. It also adds extensive test coverage across activation fusions, normalization fusions, attention/GQA packing and lowering, and weight-aware MoE fusions, plus documentation describing the new capabilities and recommended ordering.

Changes:

• Added a new olive.passes.onnx.graph_surgery package with surgeon base classes, built-in surgeon modules, and auto-registration via package import.
• Implemented multiple built-in surgeons for common fusions/lowerings (activations, normalization, attention/GQA, standard ONNX lowerings, MoE fusions).
• Added comprehensive new tests validating registry behavior, ordering, metadata preservation, and numerical parity; updated docs to describe surgeons and ordering guidance.

File summaries

File	Description
olive/passes/onnx/graph_surgery/base.py	New surgeon base classes and registration mechanism.
olive/passes/onnx/graph_surgery/__init__.py	Auto-imports built-in surgeons to ensure registration on load.
olive/passes/onnx/graph_surgeries.py	Refactors to use the new surgeon package types/registry.
olive/passes/onnx/graph_surgery/activations.py	GELU and BiasGELU fusion rewrite rules.
olive/passes/onnx/graph_surgery/normalization.py	LayerNorm and Skip{Layer,RMS}Norm fusion rewrite rules.
olive/passes/onnx/graph_surgery/attention.py	Attention→GQA, packing/unpacking QKV, RoPE separation, packed MHA rewrite rules.
olive/passes/onnx/graph_surgery/lowering.py	Standard-ONNX compatibility lowerings (Clip, RotaryEmbedding, TensorScatter, Attention decomposition, etc.).
olive/passes/onnx/graph_surgery/moe.py	Weight-aware MoE fusions to QMoE / BlockQuantizedMoE with validation and fail-closed behavior.
olive/passes/onnx/graph_surgery/_common.py	Shared helpers and cleanup mixin for rewrite rules.
docs/source/features/onnx-transformations.md	Documents built-in surgeons, ordering considerations, and MoE guidance.
test/passes/onnx/graph_surgery_test_utils.py	Shared helpers for surgery tests.
test/passes/onnx/test_graph_surgeries_activations.py	Activation surgeon registration + fusion behavior tests.
test/passes/onnx/test_graph_surgeries_normalization.py	Normalization surgeon registration + fusion behavior tests.
test/passes/onnx/test_graph_surgeries_attention.py	Attention/GQA surgeon behavior and packing/unpacking tests.
test/passes/onnx/test_graph_surgeries_lowering.py	Lowering surgeon behavior + numerical equivalence tests.
test/passes/onnx/test_graph_surgeries_moe.py	MoE fusion validation, parity, and atomic failure-mode tests.
test/passes/onnx/test_graph_surgeries_pipeline.py	End-to-end pass pipeline tests (registration, ordering, metadata, numerics).

Review details

• Files reviewed: 17/17 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.