Skip to content

Separate target build contracts from graph rewrites - #743

Draft
xiaoyu-work wants to merge 3 commits into
mainfrom
xiaoyu-work/onnx-standard-build-contract
Draft

xiaoyu-work wants to merge 3 commits into
mainfrom
xiaoyu-work/onnx-standard-build-contract

Conversation

@xiaoyu-work

@xiaoyu-work xiaoyu-work commented Sep 21, 2026

Copy link
Copy Markdown
Member

Summary

Remove graph rewrite ownership from Mobius and move target graph resolution to Olive.

Mobius now exports a canonical graph without fusion or lowering rewrites. execution_provider and device select only structural BuildContract requirements and runtime metadata. Function-backed custom operators retain their model-local standard ONNX fallback bodies for downstream resolution.

Removed from Mobius

  • The complete src/mobius/rewrite_rules package and its tests.
  • optimize_model and EP rewrite orchestration.
  • CLI --optimize.
  • FP8 KV-cache graph rewriting and CLI feature flags.
  • Automatic GGUF block-quantized MoE fusion.
  • Rewrite-specific examples, skills, and documentation.

This removes approximately 14,000 lines of duplicated rewrite implementation and tests.

Retained in Mobius

  • Canonical model/task graph construction.
  • Standard fallback function bodies for canonical custom operators.
  • Structural BuildContract fields such as OpenVINO rank-4 per_layer_inputs, target buffer limits, and static-range materialization.
  • Symbolic shape analysis, debug metadata stripping, and post-weight initializer materialization.

Olive ownership

microsoft/Olive#2670 owns EP fusion/lowering, strict-ONNX function expansion, and target profiles. All former Mobius rewrite families already have Olive surgeons; this change additionally migrates the FP8 GQA KV-cache transformation as ConvertGroupQueryAttentionKVCacheToFp8.

Validation

  • 1,713 Mobius L1/build-graph tests passed after rewrite removal.
  • 474 canonical builder/workflow core tests passed before final deletion; the final focused core/CLI suite passed with rewrite APIs absent.
  • Mobius test collection succeeds with no mobius.rewrite_rules imports.
  • Ruff lint and format checks passed.

@github-actions

Copy link
Copy Markdown

Performance Comparison

Comparing 6b27a3f009f14a77

Model Metric Baseline Current Delta
bert (feature-extraction) model_size_bytes 359 KB 359 KB +0.0%
bert (feature-extraction) num_nodes 68 68 +0.0%
falcon model_size_bytes 364 KB 364 KB +0.0%
falcon num_nodes 66 66 +0.0%
gemma2 model_size_bytes 428 KB 428 KB +0.0%
gemma2 num_nodes 105 105 +0.0%
gpt2 model_size_bytes 324 KB 324 KB +0.0%
gpt2 num_nodes 54 54 +0.0%
llama model_size_bytes 425 KB 425 KB +0.0%
llama num_nodes 60 60 +0.0%
llama (static-cache) model_size_bytes 425 KB 425 KB +0.0%
llama (static-cache) num_nodes 56 56 +0.0%
mamba (ssm-text-generation) model_size_bytes 296 KB 296 KB +0.0%
mamba (ssm-text-generation) num_nodes 94 94 +0.0%
phi3 model_size_bytes 421 KB 421 KB +0.0%
phi3 num_nodes 58 58 +0.0%
phi3 (static-cache) model_size_bytes 421 KB 421 KB +0.0%
phi3 (static-cache) num_nodes 54 54 +0.0%
qwen2 model_size_bytes 425 KB 425 KB +0.0%
qwen2 num_nodes 60 60 +0.0%
qwen2 (static-cache) model_size_bytes 425 KB 425 KB +0.0%
qwen2 (static-cache) num_nodes 56 56 +0.0%
qwen3_5_moe (hybrid-text-generation) model_size_bytes 506 KB 506 KB +0.0%
qwen3_5_moe (hybrid-text-generation) num_nodes 265 265 +0.0%
qwen3_5_text (hybrid-text-generation) model_size_bytes 458 KB 458 KB +0.0%
qwen3_5_text (hybrid-text-generation) num_nodes 127 127 +0.0%
qwen3_5_vl (hybrid-qwen-vl) model_size_bytes 977 KB 977 KB +0.0%
qwen3_5_vl (hybrid-qwen-vl) num_nodes 450 450 +0.0%
t5 (seq2seq) model_size_bytes 836 KB 836 KB +0.0%
t5 (seq2seq) num_nodes 176 176 +0.0%
whisper (speech-to-text) model_size_bytes 1008 KB 1008 KB +0.0%
whisper (speech-to-text) num_nodes 128 128 +0.0%

No performance regressions.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟡 Changes recommended

Several Transformers dispatch paths silently drop target parameters before nested builds.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1 Medium severity

Open (1)
What changed in this PR

Separates graph optimization providers from target-specific structural build contracts.

Changes:

  • Adds thread-safe build contract handling.
  • Applies target requirements to Gemma4 layouts and static Range support.
  • Forwards target options through builders and documents the API.
File Description
src/​mobius/​tasks/​_gemma4.py Uses structural contract requirements.
src/​mobius/​models/​gemma4.py Supports target-specific layouts.
src/​mobius/​models/​gemma4_test.py Tests independent target layouts.
src/​mobius/​integrations/​transformers/​_builder.py Adds target parameters; some dispatch paths drop them before nested builds.
src/​mobius/​integrations/​transformers/​_builder_test.py Verifies parameter forwarding.
src/​mobius/​components/​_common.py Uses contract-based static Range support.
src/​mobius/​_builder.py Constructs and scopes build contracts.
src/​mobius/​_build_context.py Defines thread-safe contract state.
src/​mobius/​_build_context_test.py Tests contract projection and scoping.
src/​mobius/​__init__.py Exposes the contract API.
docs/​execution_providers.md Documents independent target contracts.
CHANGELOG.md Records the new capability.

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

Comment on lines +338 to +339
target_execution_provider: str | None = None,
target_device: str | None = None,
@github-actions

github-actions Bot commented Sep 21, 2026

Copy link
Copy Markdown

🏗️ Architecture Diff

Comparing 6b27a3f00d435731

Model Sub-model Changes Status
bert (feature-extraction) model 0
falcon model 0
gemma2 model 0
gemma4 (gemma4) decoder 0
gemma4 (gemma4) embedding 0
gemma4 (gemma4) vision_encoder 0
gemma4_text model 0
gpt2 model 0
llama model 0
llama (static-cache) model 0
mamba (ssm-text-generation) model 33 🟡
phi3 model 0
phi3 (static-cache) model 0
qwen model 0
qwen (static-cache) model 0
qwen2 model 0
qwen2 (static-cache) model 0
qwen2_moe model 0
qwen2_moe (static-cache) model 0
qwen3 model 0
qwen3 (static-cache) model 0
qwen3_5_moe (hybrid-text-generation) model 0
qwen3_5_text (hybrid-text-generation) model 0
qwen3_5_vl (hybrid-qwen-vl) decoder 0
qwen3_5_vl (hybrid-qwen-vl) embedding 0
qwen3_5_vl (hybrid-qwen-vl) vision_encoder 0
qwen3_moe model 0
qwen3_moe (static-cache) model 0
qwen3_next (hybrid-text-generation) model 0
t5 (seq2seq) decoder 0
t5 (seq2seq) encoder 0
whisper (speech-to-text) decoder 32 🟡
whisper (speech-to-text) encoder 19 🟡
mamba (ssm-text-generation) / model — 33 change(s)

Op summary: 94 → 116 nodes

--- base
+++ head
@@ -17,34 +17,45 @@
 Transpose
 MatMul
 Add
+Cast
 Softplus
+Cast
 Exp
 Neg
 Unsqueeze
 Unsqueeze
 Mul
 Exp
+Cast
 Unsqueeze
 Mul
+Cast
 Unsqueeze
 Mul
 Squeeze
 Squeeze
+Cast
 Mul
 Add
+Cast
 Squeeze
 Unsqueeze
 Mul
 ReduceSum
+Cast
 Squeeze
+Cast
 Mul
 Add
 Unsqueeze
+CastLike
+CastLike
 Swish
 Mul
 Transpose
 MatMul
-SkipSimplifiedLayerNormalization
+Add
+RMSNormalization
 Transpose
 MatMul
 Split
@@ -62,33 +73,44 @@
 Transpose
 MatMul
 Add
+Cast
 Softplus
+Cast
 Exp
 Neg
 Unsqueeze
 Unsqueeze
 Mul
 Exp
+Cast
 Unsqueeze
 Mul
+Cast
 Unsqueeze
 Mul
 Squeeze
 Squeeze
+Cast
 Mul
 Add
+Cast
 Squeeze
 Unsqueeze
 Mul
 ReduceSum
+Cast
 Squeeze
+Cast
 Mul
 Add
 Unsqueeze
+CastLike
+CastLike
 Swish
 Mul
 Transpose
 MatMul
-SkipSimplifiedLayerNormalization
+Add
+RMSNormalization
 Transpose
 MatMul

Added nodes:

  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + CastLike
  • + CastLike
  • + Add
  • + RMSNormalization
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + Cast
  • + CastLike
  • + CastLike
  • + Add
  • + RMSNormalization

Removed nodes:

  • - SkipSimplifiedLayerNormalization
  • - SkipSimplifiedLayerNormalization

Modified attributes:

  • node[61] Transpose: perm: [1, 0] → [0, 2, 1]

Connectivity changes:

  • node[34] Squeeze: input_ids [53, 9] → [65, 9]
  • node[58] Transpose: input_ids [30] → [25]
  • node[59] MatMul: input_ids [99, 100] → [95, 96]
  • node[60] Split: input_ids [101, 18] → [97, 8]
  • node[61] Transpose: input_ids [31] → [98]
  • node[80] Unsqueeze: input_ids [123, 9] → [117, 21]
whisper (speech-to-text) / decoder — 32 change(s)

Op summary: 78 → 85 nodes

--- base
+++ head
@@ -1,7 +1,8 @@
 Gather
 Mul
 Gather
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -15,7 +16,8 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -29,7 +31,8 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -37,7 +40,8 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -51,7 +55,8 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -65,7 +70,8 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -73,6 +79,7 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul

Added nodes:

  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization

Removed nodes:

  • - SkipLayerNormalization
  • - SkipLayerNormalization
  • - SkipLayerNormalization
  • - SkipLayerNormalization
  • - SkipLayerNormalization
  • - SkipLayerNormalization
  • - SkipLayerNormalization

Connectivity changes:

  • node[24] Transpose: input_ids [25] → [24]
  • node[25] MatMul: input_ids [1, 94] → [1, 88]
  • node[42] Add: input_ids [117, 38] → [99, 107]
  • node[44] Transpose: input_ids [39] → [37]
  • node[45] MatMul: input_ids [112, 120] → [109, 110]
  • node[50] Transpose: input_ids [42] → [40]
  • node[51] MatMul: input_ids [125, 128] → [109, 116]
  • node[52] Add: input_ids [129, 43] → [117, 41]
  • node[54] Transpose: input_ids [46] → [42]
  • node[55] MatMul: input_ids [131, 135] → [119, 122]
  • node[56] Add: input_ids [136, 47] → [123, 43]
whisper (speech-to-text) / encoder — 19 change(s)

Op summary: 50 → 55 nodes

--- base
+++ head
@@ -3,7 +3,8 @@
 Conv
 Gelu
 Transpose
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -17,7 +18,8 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -25,7 +27,8 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -39,7 +42,8 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization
 Transpose
 MatMul
 Add
@@ -47,4 +51,5 @@
 Transpose
 MatMul
 Add
-SkipLayerNormalization
+Add
+LayerNormalization

Added nodes:

  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization
  • + Add
  • + LayerNormalization

Removed nodes:

  • - SkipLayerNormalization
  • - SkipLayerNormalization
  • - SkipLayerNormalization
  • - SkipLayerNormalization
  • - SkipLayerNormalization

Connectivity changes:

  • node[44] Add: input_ids [94, 34] → [70, 86]
  • node[46] Transpose: input_ids [35] → [33]
  • node[47] MatMul: input_ids [96, 97] → [88, 89]
  • node[48] Add: input_ids [98, 36] → [90, 34]

Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

@xiaoyu-work
xiaoyu-work marked this pull request as draft September 21, 2026 22:51

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants