Separate target build contracts from graph rewrites - #743
xiaoyu-work wants to merge 3 commits into
Conversation
Performance Comparison
|
There was a problem hiding this comment.
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
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
Rangesupport. - 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.
| target_execution_provider: str | None = None, | ||
| target_device: str | None = None, |
🏗️ Architecture Diff
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
MatMulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
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
MatMulAdded nodes:
Removed nodes:
Connectivity changes:
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
+LayerNormalizationAdded nodes:
Removed nodes:
Connectivity changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |

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_provideranddeviceselect only structuralBuildContractrequirements and runtime metadata. Function-backed custom operators retain their model-local standard ONNX fallback bodies for downstream resolution.Removed from Mobius
src/mobius/rewrite_rulespackage and its tests.optimize_modeland EP rewrite orchestration.--optimize.This removes approximately 14,000 lines of duplicated rewrite implementation and tests.
Retained in Mobius
BuildContractfields such as OpenVINO rank-4per_layer_inputs, target buffer limits, and static-range 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
mobius.rewrite_rulesimports.