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

docs/source/features/onnx-transformations.md

@@ -142,6 +142,8 @@ operator is not a portable standard-ONNX optimization.
 | `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`. |
+| `ConvertGroupQueryAttentionKVCacheToFp8` | Retype GQA past/present KV-cache I/O to FP8 E4M3 and attach per-layer scale inputs. |
+| `InlineModelLocalFunctions` | Inline all model-local function calls and reject any remaining non-standard-domain node without a standard fallback body. |
 | `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. |

olive/cli/capture_onnx.py

@@ -19,6 +19,8 @@
     update_shared_cache_options,
 )
 from olive.common.utils import set_nested_dict_value
+from olive.hardware.constants import DEVICE_TO_EXECUTION_PROVIDERS, ExecutionProvider
+from olive.model import ModelConfig
 from olive.model.utils.diffusers_utils import is_valid_diffusers_model
 from olive.telemetry import action
 
@@ -30,6 +32,56 @@ class ModelBuilderAccuracyLevel(IntEnum):
     int8 = 4
 
 
+_EP_ALIASES = {
+    "cuda": ExecutionProvider.CUDAExecutionProvider,
+    "openvino": ExecutionProvider.OpenVINOExecutionProvider,
+    "qnn": ExecutionProvider.QNNExecutionProvider,
+    "trt-rtx": ExecutionProvider.NvTensorRTRTXExecutionProvider,
+    "vitisai": ExecutionProvider.VitisAIExecutionProvider,
+}
+
+
+def _surgery(name: str, **kwargs) -> dict:
+    return {"surgeon": name, **kwargs}
+
+
+def _resolve_recipe_ep_profile(ep: str, device: str) -> tuple[ExecutionProvider, list[dict]]:
+    provider = _EP_ALIASES[ep]
+    if provider not in DEVICE_TO_EXECUTION_PROVIDERS[device]:
+        raise ValueError(f"Execution provider {ep!r} does not support device {device!r}.")
+
+    if ep == "cuda":
+        return provider, [
+            _surgery(
+                "AttentionToGroupQueryAttention",
+                supported_dtypes=["FLOAT16", "BFLOAT16"],
+            ),
+            _surgery("PackQKVForGroupQueryAttention"),
+            _surgery("FuseSkipRMSNormalization"),
+            _surgery("FuseSkipLayerNormalization"),
+        ]
+    if ep == "qnn":
+        surgeries = [
+            _surgery("AttentionToGroupQueryAttention"),
+            _surgery("PackQKVForGroupQueryAttention"),
+            _surgery("FuseSkipRMSNormalization"),
+            _surgery("AttentionMaskToSequenceLengths"),
+        ]
+        if device == "npu":
+            surgeries.append(_surgery("SimplifiedLayerNormToL2Norm"))
+        return provider, surgeries
+    if ep == "trt-rtx":
+        return provider, [
+            _surgery(
+                "AttentionToGroupQueryAttention",
+                supported_dtypes=["FLOAT16", "BFLOAT16"],
+            ),
+            _surgery("PackQKVForGroupQueryAttention"),
+        ]
+
+    return provider, []
+
+
 def parse_dim_dict(s):
     try:
         return {k: int(v) if v.isdigit() else v for k, v in (item.split("=") for item in s.split(","))}
@@ -84,10 +136,32 @@ def register_subcommand(parser: ArgumentParser):
             help=(
                 "Whether to use MobiusBuilder (mobius-onnx) to capture ONNX model. "
                 "Supports multi-component multimodal models (VLMs). "
+                "Preserves model precision and exports Mobius' canonical graph before applying "
+                "optional strict-ONNX expansion or execution-provider graph surgeries. "
                 "Requires 'pip install mobius-onnx'."
             ),
         )
 
+        mobius_ep_group = sub_parser.add_argument_group("Mobius Builder graph surgery options")
+        mobius_ep_group.add_argument(
+            "--execution_provider",
+            choices=sorted(_EP_ALIASES),
+            help="Target execution provider used to select post-export graph surgeries.",
+        )
+        mobius_ep_group.add_argument(
+            "--device",
+            choices=["cpu", "gpu", "npu"],
+            help="Target device used with --execution_provider to select post-export graph surgeries.",
+        )
+        mobius_ep_group.add_argument(
+            "--onnx_standard",
+            action="store_true",
+            help=(
+                "Expand all model-local non-standard functions after Mobius export. "
+                "May be combined with --execution_provider and --device to retain the target build contract."
+            ),
+        )
+
         # PyTorch Exporter options
         pte_group = sub_parser.add_argument_group("PyTorch Exporter options")
         pte_group.add_argument(
@@ -130,7 +204,7 @@ def register_subcommand(parser: ArgumentParser):
             type=str,
             default="fp16",
             choices=["fp16", "fp32", "int4", "bf16"],
-            help="The precision of the ONNX model. Used by Model Builder and Mobius Builder.",
+            help="The precision of the ONNX model. Used by Model Builder.",
         )
         mb_group.add_argument(
             "--int4_block_size",
@@ -194,7 +268,32 @@ def register_subcommand(parser: ArgumentParser):
 
     @action
     def run(self):
-        return self._run_workflow()
+        workflow_output = self._run_workflow()
+        if workflow_output is None or not self.args.use_mobius_builder or not workflow_output.has_output_model():
+            return workflow_output
+
+        if self.args.onnx_standard:
+            surgeries = [_surgery("InlineModelLocalFunctions")]
+            decoder_only = False
+        elif self.args.execution_provider:
+            _, surgeries = _resolve_recipe_ep_profile(self.args.execution_provider, self.args.device)
+            decoder_only = True
+        else:
+            return workflow_output
+        if not surgeries:
+            return workflow_output
+
+        from olive.workflows import run as olive_run
+
+        exported_model = workflow_output.get_best_candidate()
+        post_export_config = self._get_post_export_run_config(
+            exported_model.olive_model_config,
+            surgeries,
+            decoder_only=decoder_only,
+        )
+        if self.args.save_config_file:
+            self._save_config_file(post_export_config, file_name="graph_surgery_config.json")
+        return olive_run(post_export_config)
 
     def _get_run_config(self, tempdir: str) -> dict:
         config = deepcopy(TEMPLATE)
@@ -217,36 +316,41 @@ def _get_run_config(self, tempdir: str) -> dict:
 
         is_diffusers_model = input_model_config["type"].lower() == "diffusersmodel"
 
+        if bool(self.args.execution_provider) != bool(self.args.device):
+            raise ValueError("--execution_provider and --device must be provided together.")
+        if self.args.execution_provider and not self.args.use_mobius_builder:
+            raise ValueError("--execution_provider and --device graph-surgery profiles require --use_mobius_builder.")
+        if self.args.onnx_standard and not self.args.use_mobius_builder:
+            raise ValueError("--onnx_standard requires --use_mobius_builder.")
+
         # whether model is in fp16 or bf16 (currently not supported by CPU EP)
         is_fp16_or_bf16 = (
-            (
-                not self.args.use_model_builder
-                and not self.args.use_mobius_builder
-                and self.args.torch_dtype == "float16"
+            not self.args.use_model_builder and not self.args.use_mobius_builder and self.args.torch_dtype == "float16"
+        ) or (self.args.use_model_builder and self.args.precision in ("fp16", "bf16"))
+
+        if self.args.use_mobius_builder and self.args.execution_provider:
+            provider, _ = _resolve_recipe_ep_profile(self.args.execution_provider, self.args.device)
+            device = self.args.device
+        elif self.args.use_mobius_builder:
+            provider = ExecutionProvider.CPUExecutionProvider
+            device = "cpu"
+        else:
+            provider = (
+                ExecutionProvider.CUDAExecutionProvider if is_fp16_or_bf16 else ExecutionProvider.CPUExecutionProvider
             )
-            or (self.args.use_model_builder and self.args.precision in ("fp16", "bf16"))
-            or (self.args.use_mobius_builder and self.args.precision in ("fp16", "bf16"))
-        )
+            device = "gpu" if is_fp16_or_bf16 else "cpu"
+
         to_replace = [
             ("input_model", input_model_config),
             ("output_dir", self.args.output_path),
             ("log_severity_level", self.args.log_level),
-            (("systems", "local_system", "accelerators", 0, "device"), "gpu" if is_fp16_or_bf16 else "cpu"),
-            (
-                ("systems", "local_system", "accelerators", 0, "execution_providers"),
-                [("CUDAExecutionProvider" if is_fp16_or_bf16 else "CPUExecutionProvider")],
-            ),
+            (("systems", "local_system", "accelerators", 0, "device"), device),
+            (("systems", "local_system", "accelerators", 0, "execution_providers"), [provider.value]),
         ]
 
         if self.args.use_mobius_builder:
-            if self.args.precision not in ("fp32", "fp16", "bf16"):
-                raise ValueError(
-                    f"MobiusBuilder supports precisions fp32/fp16/bf16; got '{self.args.precision}'. "
-                    "For INT4, capture in fp32/fp16/bf16 first and run a quantization pass afterwards."
-                )
             del config["passes"]["c"]
             del config["passes"]["m"]
-            to_replace.append((("passes", "b", "precision"), self.args.precision))
         elif is_diffusers_model:
             del config["passes"]["m"]
             del config["passes"]["b"]
@@ -325,12 +429,43 @@ def _get_run_config(self, tempdir: str) -> dict:
 
         return config
 
+    def _get_post_export_run_config(
+        self,
+        input_model_config: dict,
+        surgeries: list[dict],
+        *,
+        decoder_only: bool,
+    ) -> dict:
+        model_config = ModelConfig.model_validate(input_model_config)
+        components = model_config.get_components()
+        is_multimodal = components is not None
+        if decoder_only and is_multimodal and "decoder" not in components:
+            raise ValueError(
+                "Execution-provider graph surgeries require a 'decoder' component, "
+                f"but the exported model contains {components}."
+            )
+
+        config = deepcopy(MULTIMODAL_TEMPLATE)
+        if not decoder_only or not is_multimodal:
+            config.pop("builds")
+
+        config["input_model"] = model_config.model_dump()
+        config["output_dir"] = self.args.output_path
+        config["log_severity_level"] = self.args.log_level
+        accelerator = config["systems"]["local_system"]["accelerators"][0]
+        accelerator["device"] = self.args.device or "cpu"
+        accelerator["execution_providers"] = [
+            _EP_ALIASES.get(self.args.execution_provider, ExecutionProvider.CPUExecutionProvider).value
+        ]
+        config["passes"]["g"]["surgeries"] = surgeries
+        update_shared_cache_options(config, self.args)
+        return config
+
 
 TEMPLATE = {
     "systems": {
         "local_system": {
             "type": "LocalSystem",
-            # might need an ep option to set for model builder, it is sensitive to ep
             "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}],
         }
     },
@@ -346,3 +481,20 @@ def _get_run_config(self, tempdir: str) -> dict:
     "target": "local_system",
     "no_artifacts": True,
 }
+
+
+MULTIMODAL_TEMPLATE = {
+    "systems": {
+        "local_system": {
+            "type": "LocalSystem",
+            "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}],
+        }
+    },
+    "passes": {"g": {"type": "GraphSurgeries"}},
+    "builds": {
+        "decoder": {"components": ["decoder"], "pipeline": ["g"]},
+    },
+    "host": "local_system",
+    "target": "local_system",
+    "no_artifacts": True,
+}

olive/olive_config.json

@@ -215,7 +215,7 @@
             "module_path": "olive.passes.onnx.mobius_model_builder.MobiusBuilder",
             "supported_providers": [ "*" ],
             "supported_accelerators": [ "*" ],
-            "supported_precisions": [ "fp32", "fp16", "bf16" ],
+            "supported_precisions": [ "*" ],
             "supported_algorithms": [  ],
             "supported_quantization_encodings": [  ]
         },

olive/passes/onnx/graph_surgery/__init__.py

@@ -8,6 +8,7 @@
 from olive.passes.onnx.graph_surgery.attention import (
     AttentionToGroupQueryAttention,
     BlockDiagonalAttentionToPackedMHA,
+    ConvertGroupQueryAttentionKVCacheToFp8,
     PackQKVForGroupQueryAttention,
     SeparateGroupQueryAttentionRoPE,
     UnpackGroupQueryAttentionQKV,
@@ -17,6 +18,7 @@
     ClipToMinMax,
     DecomposeAttention,
     DecomposeOnnxRotaryEmbedding,
+    InlineModelLocalFunctions,
     Rank4RMSNormToRank3,
     StaticEmptyKV,
     TensorScatterToScatterND,
@@ -32,6 +34,7 @@
     "AttentionToGroupQueryAttention",
     "BlockDiagonalAttentionToPackedMHA",
     "ClipToMinMax",
+    "ConvertGroupQueryAttentionKVCacheToFp8",
     "DecomposeAttention",
     "DecomposeOnnxRotaryEmbedding",
     "FuseBiasGelu",
@@ -41,6 +44,7 @@
     "FuseLayerNormalization",
     "FuseSkipLayerNormalization",
     "FuseSkipRMSNormalization",
+    "InlineModelLocalFunctions",
     "MoEGraphSurgeryError",
     "PackQKVForGroupQueryAttention",
     "ProtoSurgeon",