Skip to content

Add EP- and device-specific graph surgery support to capture-onnx-graph CLI - #2670

Draft
Xiaoyu Z (xiaoyu-work) wants to merge 10 commits into
mainfrom
xiaoyu-work/capture-onnx-graph-ep-surgeries
Draft

Xiaoyu Z (xiaoyu-work) wants to merge 10 commits into
mainfrom
xiaoyu-work/capture-onnx-graph-ep-surgeries

Conversation

@xiaoyu-work

@xiaoyu-work Xiaoyu Z (xiaoyu-work) commented Sep 17, 2026

Copy link
Copy Markdown
Member

Describe your changes

This PR changes olive capture-onnx-graph --use_mobius_builder to use a two-stage graph flow:

  1. Mobius exports its canonical graph without automatically applying graph rewrites. Function-backed custom operators and their standard ONNX fallback bodies are preserved.
  2. Olive decides the final graph representation after export.

When --execution_provider and --device are supplied, Olive applies the selected EP surgery profile. For a multi-component model, builds selects only the exported decoder, leaving vision, embedding, and other components unchanged.

When --onnx_standard is supplied, Olive runs InlineModelLocalFunctions across every exported component. All model-local functions are inlined, function definitions are removed, and Olive rejects any remaining non-standard-domain operator that has no standard fallback body. --onnx_standard may be combined with EP/device so structural requirements such as OpenVINO Gemma4 rank-4 per_layer_inputs and runtime configuration are retained.

Olive now owns every graph rewrite removed from Mobius. In addition to the migrated fusion/lowering surgeons, ConvertGroupQueryAttentionKVCacheToFp8 carries the former Mobius FP8 GQA KV-cache transformation, including calibrated per-layer scale inputs.

The accelerator EP/device is forwarded to Mobius only as a structural build contract; it never enables Mobius graph rewrites. This depends on onnxruntime/mobius#743.

Execution provider \ Device cpu gpu npu
cuda Invalid AttentionToGroupQueryAttention (FP16/BF16)
PackQKVForGroupQueryAttention
FuseSkipRMSNormalization
FuseSkipLayerNormalization
Invalid
openvino None None None
qnn Invalid AttentionToGroupQueryAttention
PackQKVForGroupQueryAttention
FuseSkipRMSNormalization
AttentionMaskToSequenceLengths
AttentionToGroupQueryAttention
PackQKVForGroupQueryAttention
FuseSkipRMSNormalization
AttentionMaskToSequenceLengths
SimplifiedLayerNormToL2Norm
trt-rtx Invalid AttentionToGroupQueryAttention (FP16/BF16)
PackQKVForGroupQueryAttention
Invalid
vitisai Invalid Invalid None

Invalid denotes an unsupported provider/device combination. None denotes a valid combination with no automatic EP surgeries; users may still request strict expansion with --onnx_standard.

The EP profiles intentionally cover ordinary decoder graphs. Model-specific paths remain explicit rather than silently applying an incompatible generic rewrite.

Checklist before requesting a review

  • Add unit tests for this change.
  • Make sure all relevant tests pass.
  • Update directly related documentation.
  • Run lintrunner on changed files.
  • Include a user-facing release note.

Release note: capture-onnx-graph now keeps Mobius export canonical and lets Olive explicitly apply EP graph surgeries or expand all function-backed custom operators with --onnx_standard.

(Optional) Issue link

N/A

Copilot AI lite review requested due to automatic review settings September 17, 2026 21:43

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.

🟡 Changes recommended

The QNN mapping passes an unsupported Mobius runtime EP instead of retaining the existing standard-export fallback.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This PR updates Mobius ONNX capture to preserve source precision and optionally apply execution-provider-specific graph surgeries.

Changes:

  • Standardizes Mobius export and removes precision overrides.
  • Adds CLI execution-provider/device surgery profiles.
  • Adds dtype-aware GQA surgery filtering and tests.
File summaries
File Description
olive/cli/capture_onnx.py Adds EP profiles and CLI validation.
olive/passes/onnx/mobius_model_builder.py Uses standard ONNX export and runtime packaging.
olive/passes/onnx/graph_surgery/attention.py Filters GQA rewrites by dtype.
olive/olive_config.json Allows all source precisions.
test/cli/test_cli.py Tests CLI profiles and validation.
test/passes/onnx/test_mobius_model_builder.py Tests standard Mobius export behavior.
test/passes/onnx/test_graph_surgeries_attention.py Tests dtype filtering.
Review details
  • Files reviewed: 7/7 changed files
  • Comments generated: 1
  • Review effort level: Lite

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

Comment on lines +85 to 88
ExecutionProvider.NvTensorRTRTXExecutionProvider: "trt-rtx",
ExecutionProvider.OpenVINOExecutionProvider: "openvino",
ExecutionProvider.QNNExecutionProvider: "qnn",
ExecutionProvider.WebGpuExecutionProvider: "webgpu",
@xiaoyu-work Xiaoyu Z (xiaoyu-work) changed the title Add execution provider graph surgeries to ONNX capture Add recipe-aligned graph surgeries to ONNX capture Sep 17, 2026
@xiaoyu-work Xiaoyu Z (xiaoyu-work) changed the title Add recipe-aligned graph surgeries to ONNX capture Add ModelBuilder-aligned graph surgeries to ONNX capture Sep 17, 2026
@xiaoyu-work Xiaoyu Z (xiaoyu-work) changed the title Add ModelBuilder-aligned graph surgeries to ONNX capture Add generic decoder graph surgeries to ONNX capture Sep 17, 2026
@xiaoyu-work Xiaoyu Z (xiaoyu-work) changed the title Add generic decoder graph surgeries to ONNX capture Add EP- and device-specific graph surgery support to capture-onnx-graph CLI Sep 17, 2026
@xiaoyu-work
Xiaoyu Z (xiaoyu-work) marked this pull request as draft September 21, 2026 20:24

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