From c245ed0bee512956858195bc6a493489cc7a8e3c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 17 Sep 2026 14:17:06 -0700 Subject: [PATCH 01/10] Preserve model precision in MobiusBuilder --- olive/cli/capture_onnx.py | 18 ++---------- olive/passes/onnx/mobius_model_builder.py | 29 +------------------ test/cli/test_cli.py | 24 +++++---------- test/passes/onnx/test_mobius_model_builder.py | 22 +++++++------- 4 files changed, 23 insertions(+), 70 deletions(-) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 3966b798c..63dc0e0a0 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -130,7 +130,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", @@ -219,14 +219,8 @@ def _get_run_config(self, tempdir: str) -> dict: # 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" - ) - 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")) - ) + 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")) to_replace = [ ("input_model", input_model_config), ("output_dir", self.args.output_path), @@ -239,14 +233,8 @@ def _get_run_config(self, tempdir: str) -> dict: ] 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"] diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index d52b6d26f..0446cf41e 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Any, ClassVar from olive.common.utils import StrEnumBase -from olive.constants import Precision from olive.hardware.constants import EXECUTION_PROVIDER_TO_MOBIUS_EP, ExecutionProvider from olive.model import HfModelHandler, ONNXModelHandler from olive.model.handler.composite import CompositeModelHandler @@ -24,18 +23,6 @@ logger = logging.getLogger(__name__) -# Maps Olive Precision values to mobius dtype strings. -# "f32" = 32-bit float (torch.float32), standard full precision. -# "f16" = 16-bit float (torch.float16), half precision — good for GPU inference. -# "bf16" = bfloat16 (torch.bfloat16), brain float — preferred over f16 on newer hardware. -# For INT4/INT8 quantization, use a downstream Olive quantization pass (e.g. OnnxMatMulNBits) -# after this pass rather than setting precision here. -_PRECISION_TO_DTYPE: dict[str, str] = { - Precision.FP32: "f32", - Precision.FP16: "f16", - Precision.BF16: "bf16", -} - class MobiusBuilder(Pass): """Olive pass that uses mobius to build ONNX models from HuggingFace model IDs. @@ -89,7 +76,6 @@ class MobiusEP(StrEnumBase): TRT_RTX = "trt-rtx" ONNX_STANDARD = "onnx-standard" - # Maps Olive ExecutionProvider enum values to mobius EP names. EP_MAP: ClassVar[dict[ExecutionProvider, str]] = { ExecutionProvider.CPUExecutionProvider: "cpu", ExecutionProvider.CUDAExecutionProvider: "cuda", @@ -106,16 +92,6 @@ def is_accelerator_agnostic(cls, accelerator_spec: AcceleratorSpec) -> bool: @classmethod def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]: return { - "precision": PassConfigParam( - type_=Precision, - required=False, - default_value=Precision.FP32, - description=( - "Model weight / compute precision. One of: fp32, fp16, bf16. " - "Defaults to fp32. For INT4 quantization, run an Olive " - "quantization pass (e.g. OnnxMatMulNBits) after this pass." - ), - ), "text_only": PassConfigParam( type_=bool, required=False, @@ -170,7 +146,6 @@ def _run_for_config( self.accelerator_spec.accelerator_type, ) - dtype_str: str = _PRECISION_TO_DTYPE.get(config.precision, "f32") model_id: str = model.model_name_or_path load_kwargs = model.get_load_kwargs() @@ -178,10 +153,9 @@ def _run_for_config( trust_remote_code: bool = load_kwargs.get("trust_remote_code", False) logger.info( - "MobiusBuilder: building '%s' (ep=%s, dtype=%s)", + "MobiusBuilder: building '%s' (ep=%s)", model_id, ep_str, - dtype_str, ) if trust_remote_code: @@ -201,7 +175,6 @@ def _run_for_config( pkg = build( model_id, revision=revision, - dtype=dtype_str, execution_provider=ep_str, load_weights=True, trust_remote_code=trust_remote_code, diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index f12ab8359..97cb933bd 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -539,15 +539,8 @@ def test_capture_onnx_command_fix_shape(_, mock_run, use_model_builder, tmp_path @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) -@pytest.mark.parametrize( - ("precision", "use_ort_genai"), - [ - ("fp16", True), - ("fp32", False), - ("bf16", True), - ], -) -def test_capture_onnx_command_use_mobius_builder(_, mock_run, precision, use_ort_genai, tmp_path): +@pytest.mark.parametrize("use_ort_genai", [True, False]) +def test_capture_onnx_command_use_mobius_builder(_, mock_run, use_ort_genai, tmp_path): # setup output_dir = tmp_path / "output_dir" model_id = "dummy-model-id" @@ -558,8 +551,6 @@ def test_capture_onnx_command_use_mobius_builder(_, mock_run, precision, use_ort "-o", str(output_dir), "--use_mobius_builder", - "--precision", - precision, ] if use_ort_genai: command_args.append("--use_ort_genai") @@ -574,13 +565,13 @@ def test_capture_onnx_command_use_mobius_builder(_, mock_run, precision, use_ort assert "c" not in config["passes"] assert "m" not in config["passes"] assert config["passes"]["b"]["type"] == "MobiusBuilder" - assert config["passes"]["b"]["precision"] == precision + assert "precision" not in config["passes"]["b"] assert mock_run.call_count == 1 @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) -def test_capture_onnx_command_use_mobius_builder_rejects_int4(_, __, tmp_path): +def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision(_, mock_run, tmp_path): # setup output_dir = tmp_path / "output_dir" command_args = [ @@ -594,9 +585,10 @@ def test_capture_onnx_command_use_mobius_builder_rejects_int4(_, __, tmp_path): "int4", ] - # execute / verify - with pytest.raises(ValueError, match="MobiusBuilder supports precisions fp32/fp16/bf16"): - cli_main(command_args) + cli_main(command_args) + + config = mock_run.call_args[0][0] + assert "precision" not in config["passes"]["b"] @patch("olive.workflows.run") diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index fe43864ca..0c8d91ee1 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -79,7 +79,7 @@ def _make_hf_model(model_path: str, load_kwargs: dict | None = None, task: str | def _make_pass(ep: str = ExecutionProvider.CPUExecutionProvider, text_only: bool | None = None) -> MobiusBuilder: accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ep) - pass_config = {"precision": "fp32"} + pass_config = {} if text_only is not None: pass_config["text_only"] = text_only return create_pass_from_dict( @@ -198,12 +198,12 @@ def __exit__(self, *args): def test_default_config_params(): - """MobiusBuilder must declare precision, and must not declare execution_provider or trust_remote_code.""" + """MobiusBuilder preserves model precision and does not expose build environment options.""" accelerator_spec = AcceleratorSpec( accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider ) config = MobiusBuilder._default_config(accelerator_spec) # pylint: disable=protected-access - assert "precision" in config + assert "precision" not in config assert config["text_only"].default_value is False assert config["text_only"].required is False assert "execution_provider" not in config @@ -257,7 +257,7 @@ def test_single_component_returns_onnx_handler(tmp_path): mock_build.assert_called_once() call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == "cpu" - assert call_kwargs["dtype"] == "f32" + assert "dtype" not in call_kwargs def test_text_only_default_omits_mobius_build_kwarg(tmp_path): @@ -416,7 +416,7 @@ def test_ep_auto_detected_from_accelerator(tmp_path): ) p = create_pass_from_dict( MobiusBuilder, - {"precision": "fp16"}, + {}, disable_search=True, accelerator_spec=accelerator_spec, ) @@ -426,7 +426,7 @@ def test_ep_auto_detected_from_accelerator(tmp_path): call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == "cuda" - assert call_kwargs["dtype"] == "f16" + assert "dtype" not in call_kwargs def test_hf_load_options_forwarded_to_build_and_genai_config(tmp_path): @@ -521,7 +521,7 @@ def test_unsupported_ep_falls_back_to_default(tmp_path): ) p = create_pass_from_dict( MobiusBuilder, - {"precision": "fp32"}, + {}, disable_search=True, accelerator_spec=accelerator_spec, ) @@ -542,7 +542,7 @@ def test_none_execution_provider_falls_back_to_default(tmp_path): accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=None) p = create_pass_from_dict( MobiusBuilder, - {"precision": "fp32"}, + {}, disable_search=True, accelerator_spec=accelerator_spec, ) @@ -637,7 +637,7 @@ def test_trust_remote_code_warning_logged(tmp_path): pkg = _fake_pkg(["model"], out) p = create_pass_from_dict( MobiusBuilder, - {"precision": "fp32"}, + {}, disable_search=True, accelerator_spec=AcceleratorSpec( accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider @@ -672,13 +672,13 @@ def test_no_warning_when_trust_remote_code_false(tmp_path): # --------------------------------------------------------------------------- -def _make_filtered_pass(components_to_export, precision: str = "fp16") -> MobiusBuilder: +def _make_filtered_pass(components_to_export) -> MobiusBuilder: accelerator_spec = AcceleratorSpec( accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider ) return create_pass_from_dict( MobiusBuilder, - {"precision": precision, "components_to_export": components_to_export}, + {"components_to_export": components_to_export}, disable_search=True, accelerator_spec=accelerator_spec, ) From 7ce7525cee4c7a125c7159fe2b0ab7f948f93f33 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 17 Sep 2026 14:41:38 -0700 Subject: [PATCH 02/10] Add EP graph surgeries to ONNX capture --- olive/cli/capture_onnx.py | 133 +++++++++- olive/olive_config.json | 2 +- olive/passes/onnx/graph_surgery/attention.py | 33 ++- olive/passes/onnx/mobius_model_builder.py | 29 ++- test/cli/test_cli.py | 236 ++++++++++++++++++ .../onnx/test_graph_surgeries_attention.py | 39 ++- test/passes/onnx/test_mobius_model_builder.py | 41 ++- 7 files changed, 470 insertions(+), 43 deletions(-) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 63dc0e0a0..d2239ce23 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -19,6 +19,7 @@ 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.utils.diffusers_utils import is_valid_diffusers_model from olive.telemetry import action @@ -30,6 +31,89 @@ class ModelBuilderAccuracyLevel(IntEnum): int8 = 4 +_EP_ALIASES = { + "cpu": ExecutionProvider.CPUExecutionProvider, + "cuda": ExecutionProvider.CUDAExecutionProvider, + "default": ExecutionProvider.CPUExecutionProvider, + "dml": ExecutionProvider.DmlExecutionProvider, + "migraphx": ExecutionProvider.MIGraphXExecutionProvider, + "onnx-standard": ExecutionProvider.CPUExecutionProvider, + "openvino": ExecutionProvider.OpenVINOExecutionProvider, + "qnn": ExecutionProvider.QNNExecutionProvider, + "rocm": ExecutionProvider.ROCMExecutionProvider, + "trt-rtx": ExecutionProvider.NvTensorRTRTXExecutionProvider, + "vitisai": ExecutionProvider.VitisAIExecutionProvider, + "webgpu": ExecutionProvider.WebGpuExecutionProvider, +} + + +def _surgery(name: str, **kwargs) -> dict: + return {"surgeon": name, **kwargs} + + +def _resolve_mobius_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 == "qnn" and device == "gpu": + raise ValueError("QNN GPU has no generic Mobius graph-surgery profile; use --device npu for QNN HTP.") + + clip = _surgery("ClipToMinMax") + gelu = _surgery("FuseGelu") + skip_norms = [ + _surgery("FuseSkipRMSNormalization"), + _surgery("FuseSkipLayerNormalization"), + ] + + if ep == "default": + return provider, [*skip_norms, gelu, clip] + if ep in ("onnx-standard", "migraphx", "rocm", "vitisai"): + return provider, [] + if ep == "openvino": + return provider, [gelu, clip] + if ep == "qnn": + return provider, [ + gelu, + _surgery("SeparateGroupQueryAttentionRoPE"), + _surgery("UnpackGroupQueryAttentionQKV"), + _surgery("Rank4RMSNormToRank3"), + _surgery("DecomposeOnnxRotaryEmbedding"), + _surgery("TensorScatterToScatterND"), + _surgery("DecomposeAttention"), + clip, + ] + + supported_dtypes = { + "cpu": ["FLOAT"], + "cuda": ["FLOAT16", "BFLOAT16"], + "dml": ["FLOAT16"], + "trt-rtx": ["FLOAT16", "BFLOAT16"], + "webgpu": ["FLOAT", "FLOAT16"], + }[ep] + attention = _surgery("AttentionToGroupQueryAttention", supported_dtypes=supported_dtypes) + + if ep == "dml": + return provider, [ + attention, + *skip_norms, + gelu, + _surgery("SeparateGroupQueryAttentionRoPE"), + _surgery("UnpackGroupQueryAttentionQKV"), + clip, + ] + + surgeries = [attention] + if ep in ("cpu", "cuda", "trt-rtx", "webgpu"): + surgeries.append(_surgery("PackQKVForGroupQueryAttention")) + if ep != "trt-rtx": + surgeries.extend(skip_norms) + surgeries.append(gelu) + if ep == "webgpu": + surgeries.append(_surgery("StaticEmptyKV")) + surgeries.append(clip) + return provider, surgeries + + 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 +168,24 @@ def register_subcommand(parser: ArgumentParser): help=( "Whether to use MobiusBuilder (mobius-onnx) to capture ONNX model. " "Supports multi-component multimodal models (VLMs). " + "Preserves the model precision and exports standard ONNX before applying " + "optional 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.", + ) + # PyTorch Exporter options pte_group = sub_parser.add_argument_group("PyTorch Exporter options") pte_group.add_argument( @@ -217,27 +315,48 @@ 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.") + # 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" ) or (self.args.use_model_builder and self.args.precision in ("fp16", "bf16")) + + mobius_surgeries = [] + if self.args.use_mobius_builder and self.args.execution_provider: + provider, mobius_surgeries = _resolve_mobius_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 + ) + 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: del config["passes"]["c"] del config["passes"]["m"] + if mobius_surgeries: + to_replace.append((("passes", "g", "surgeries"), mobius_surgeries)) + else: + del config["passes"]["g"] elif is_diffusers_model: del config["passes"]["m"] del config["passes"]["b"] + del config["passes"]["g"] to_replace.extend( [ ( @@ -251,6 +370,7 @@ def _get_run_config(self, tempdir: str) -> dict: elif self.args.use_model_builder: del config["passes"]["c"] del config["passes"]["b"] + del config["passes"]["g"] to_replace.extend( [ (("passes", "m", "precision"), self.args.precision), @@ -272,6 +392,7 @@ def _get_run_config(self, tempdir: str) -> dict: to_replace.append((("passes", "m", "int4_accuracy_level"), self.args.int4_accuracy_level)) else: del config["passes"]["b"] + del config["passes"]["g"] to_replace.extend( [ ( @@ -318,7 +439,6 @@ def _get_run_config(self, tempdir: str) -> dict: "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"]}], } }, @@ -328,6 +448,7 @@ def _get_run_config(self, tempdir: str) -> dict: }, "m": {"type": "ModelBuilder", "metadata_only": False}, "b": {"type": "MobiusBuilder"}, + "g": {"type": "GraphSurgeries"}, "f": {"type": "DynamicToFixedShape"}, }, "host": "local_system", diff --git a/olive/olive_config.json b/olive/olive_config.json index 85ab4233e..dc7ecaeb6 100644 --- a/olive/olive_config.json +++ b/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": [ ] }, diff --git a/olive/passes/onnx/graph_surgery/attention.py b/olive/passes/onnx/graph_surgery/attention.py index d1cbd1fa1..4a32c87e0 100644 --- a/olive/passes/onnx/graph_surgery/attention.py +++ b/olive/passes/onnx/graph_surgery/attention.py @@ -133,9 +133,14 @@ def _static_last_dim(value): return False +def _is_supported_dtype(value: ir.Value, supported_dtypes: frozenset[ir.DataType] | None) -> bool: + return supported_dtypes is None or value.dtype in supported_dtypes + + class _RotaryAttentionToGQA(RewriteRuleClassBase): - def __init__(self): + def __init__(self, supported_dtypes: frozenset[ir.DataType] | None = None): super().__init__() + self._supported_dtypes = supported_dtypes self._seqlens_k = None self._total_seq_len = None self._cos_cache = None @@ -155,8 +160,10 @@ def pattern(self, op, q_pre, k_pre, v, attention_bias, past_key, past_value, cos _outputs=["attn_out", "present_key", "present_value"], ) - def check(self, context, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): + def check(self, context, q_pre, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): result = MatchResult() + if not _is_supported_dtype(q_pre, self._supported_dtypes): + return result.fail("Attention dtype is unsupported by the target execution provider") if not _local_window_from_attention_bias(attention_bias).recognized: return result.fail("Attention bias cannot be represented by GroupQueryAttention") @@ -258,8 +265,9 @@ def rewrite( class _AttentionToGQA(RewriteRuleClassBase): - def __init__(self): + def __init__(self, supported_dtypes: frozenset[ir.DataType] | None = None): super().__init__() + self._supported_dtypes = supported_dtypes self._seqlens_k = None self._total_seq_len = None @@ -275,8 +283,10 @@ def pattern(self, op, q, k, v, attention_bias, past_key, past_value): _outputs=["attn_out", "present_key", "present_value"], ) - def check(self, context, attn_out, k, v, attention_bias, past_key, past_value, **_): + def check(self, context, q, attn_out, k, v, attention_bias, past_key, past_value, **_): result = MatchResult() + if not _is_supported_dtype(q, self._supported_dtypes): + return result.fail("Attention dtype is unsupported by the target execution provider") if not _local_window_from_attention_bias(attention_bias).recognized: return result.fail("Attention bias cannot be represented by GroupQueryAttention") @@ -360,8 +370,21 @@ def rewrite( class AttentionToGroupQueryAttention(RewriteRuleSurgeon): """Fuse decoder Attention, and standard RoPE when possible, into GQA.""" + def __init__(self, supported_dtypes: list[str] | None = None): + try: + self.supported_dtypes = ( + frozenset(ir.DataType[dtype.upper()] for dtype in supported_dtypes) if supported_dtypes else None + ) + except KeyError as exc: + raise ValueError(f"Unsupported ONNX dtype {exc.args[0]!r}") from exc + def rules(self) -> pattern.RewriteRuleSet: - return pattern.RewriteRuleSet([_RotaryAttentionToGQA.rule(), _AttentionToGQA.rule()]) + return pattern.RewriteRuleSet( + [ + _RotaryAttentionToGQA.rule(self.supported_dtypes), + _AttentionToGQA.rule(self.supported_dtypes), + ] + ) class _PackQKVForGQA(RewriteRuleClassBase): diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index 0446cf41e..62e7b752a 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -33,6 +33,8 @@ class MobiusBuilder(Pass): pass returns a :class:`~olive.model.handler.composite.CompositeModelHandler` whose components are individual :class:`~olive.model.ONNXModelHandler` objects. Single-component models return a plain :class:`~olive.model.ONNXModelHandler`. + Mobius preserves the source model precision and exports a standard-ONNX graph; + target-specific graph transformations run in downstream Olive passes. Use ``components_to_export`` to export only a subset of components. This is useful when some components (e.g. a text decoder) are already exported and @@ -80,13 +82,16 @@ class MobiusEP(StrEnumBase): ExecutionProvider.CPUExecutionProvider: "cpu", ExecutionProvider.CUDAExecutionProvider: "cuda", ExecutionProvider.DmlExecutionProvider: "dml", + ExecutionProvider.NvTensorRTRTXExecutionProvider: "trt-rtx", + ExecutionProvider.OpenVINOExecutionProvider: "openvino", + ExecutionProvider.QNNExecutionProvider: "qnn", ExecutionProvider.WebGpuExecutionProvider: "webgpu", } @classmethod def is_accelerator_agnostic(cls, accelerator_spec: AcceleratorSpec) -> bool: - # EP selection determines which fused ops are emitted, so this pass is - # EP-specific. + # The graph is accelerator-agnostic, but ORT GenAI runtime packaging + # records the target EP. return False @classmethod @@ -135,13 +140,17 @@ def _run_for_config( if not isinstance(model, HfModelHandler): raise ValueError(f"MobiusBuilder requires an HfModelHandler input, got {type(model).__name__}.") - # Map Olive EP to mobius EP. If unsupported/unknown, fall back to mobius default EP. + # The graph is always exported through Mobius' standard-ONNX path. + # The requested EP is retained only for ORT GenAI runtime packaging. requested_ep = self.accelerator_spec.execution_provider - ep_str: str = EXECUTION_PROVIDER_TO_MOBIUS_EP.get(requested_ep, self.MobiusEP.DEFAULT) - if ep_str == self.MobiusEP.DEFAULT: + runtime_ep: str = self.EP_MAP.get( + requested_ep, + EXECUTION_PROVIDER_TO_MOBIUS_EP.get(requested_ep, self.MobiusEP.DEFAULT), + ) + if runtime_ep == self.MobiusEP.DEFAULT: logger.warning( "MobiusBuilder: execution provider '%s' on accelerator '%s' is not explicitly supported; " - "falling back to mobius default EP.", + "using the default Mobius runtime configuration.", requested_ep, self.accelerator_spec.accelerator_type, ) @@ -153,9 +162,9 @@ def _run_for_config( trust_remote_code: bool = load_kwargs.get("trust_remote_code", False) logger.info( - "MobiusBuilder: building '%s' (ep=%s)", + "MobiusBuilder: building '%s' as standard ONNX (runtime ep=%s)", model_id, - ep_str, + runtime_ep, ) if trust_remote_code: @@ -175,7 +184,7 @@ def _run_for_config( pkg = build( model_id, revision=revision, - execution_provider=ep_str, + execution_provider=self.MobiusEP.ONNX_STANDARD, load_weights=True, trust_remote_code=trust_remote_code, **text_only_kwargs, @@ -231,7 +240,7 @@ def components_filter(name: str) -> bool: pkg, str(output_dir), model_id, - ep_str, + runtime_ep, revision=revision, trust_remote_code=trust_remote_code, ) diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index 97cb933bd..b261ab1fc 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -564,8 +564,12 @@ def test_capture_onnx_command_use_mobius_builder(_, mock_run, use_ort_genai, tmp assert "b" in config["passes"] assert "c" not in config["passes"] assert "m" not in config["passes"] + assert "g" not in config["passes"] assert config["passes"]["b"]["type"] == "MobiusBuilder" assert "precision" not in config["passes"]["b"] + accelerator = config["systems"]["local_system"]["accelerators"][0] + assert accelerator["device"] == "cpu" + assert accelerator["execution_providers"] == ["CPUExecutionProvider"] assert mock_run.call_count == 1 @@ -591,6 +595,238 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision assert "precision" not in config["passes"]["b"] +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +@pytest.mark.parametrize( + ("ep", "device", "provider", "surgeons"), + [ + ( + "cuda", + "gpu", + "CUDAExecutionProvider", + [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + "FuseGelu", + "ClipToMinMax", + ], + ), + ( + "trt-rtx", + "gpu", + "NvTensorRTRTXExecutionProvider", + [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseGelu", + "ClipToMinMax", + ], + ), + ( + "openvino", + "npu", + "OpenVINOExecutionProvider", + ["FuseGelu", "ClipToMinMax"], + ), + ( + "qnn", + "npu", + "QNNExecutionProvider", + [ + "FuseGelu", + "SeparateGroupQueryAttentionRoPE", + "UnpackGroupQueryAttentionQKV", + "Rank4RMSNormToRank3", + "DecomposeOnnxRotaryEmbedding", + "TensorScatterToScatterND", + "DecomposeAttention", + "ClipToMinMax", + ], + ), + ], +) +def test_capture_onnx_command_adds_mobius_ep_surgeries(_, mock_run, ep, device, provider, surgeons, tmp_path): + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--use_mobius_builder", + "--execution_provider", + ep, + "--device", + device, + ] + ) + + config = mock_run.call_args[0][0] + accelerator = config["systems"]["local_system"]["accelerators"][0] + assert accelerator == {"device": device, "execution_providers": [provider]} + assert list(config["passes"]) == ["b", "g"] + assert [surgery["surgeon"] for surgery in config["passes"]["g"]["surgeries"]] == surgeons + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_mobius_ep_without_surgeries(_, mock_run, tmp_path): + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--use_mobius_builder", + "--execution_provider", + "vitisai", + "--device", + "npu", + ] + ) + + config = mock_run.call_args[0][0] + assert list(config["passes"]) == ["b"] + assert config["systems"]["local_system"]["accelerators"][0] == { + "device": "npu", + "execution_providers": ["VitisAIExecutionProvider"], + } + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +@pytest.mark.parametrize( + ("extra_args", "message"), + [ + (["--execution_provider", "cuda"], "must be provided together"), + (["--device", "gpu"], "must be provided together"), + (["--execution_provider", "cuda", "--device", "npu"], "does not support device"), + (["--execution_provider", "qnn", "--device", "gpu"], "has no generic Mobius"), + ], +) +def test_capture_onnx_command_rejects_invalid_mobius_ep_device(_, __, extra_args, message, tmp_path): + command_args = [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--use_mobius_builder", + *extra_args, + ] + + with pytest.raises(ValueError, match=message): + cli_main(command_args) + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_rejects_ep_device_without_mobius(_, __, tmp_path): + with pytest.raises(ValueError, match="require --use_mobius_builder"): + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--execution_provider", + "cuda", + "--device", + "gpu", + ] + ) + + +@pytest.mark.parametrize( + ("ep", "device", "expected"), + [ + ( + "cpu", + "cpu", + [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + "FuseGelu", + "ClipToMinMax", + ], + ), + ( + "dml", + "gpu", + [ + "AttentionToGroupQueryAttention", + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + "FuseGelu", + "SeparateGroupQueryAttentionRoPE", + "UnpackGroupQueryAttentionQKV", + "ClipToMinMax", + ], + ), + ( + "webgpu", + "gpu", + [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + "FuseGelu", + "StaticEmptyKV", + "ClipToMinMax", + ], + ), + ( + "default", + "cpu", + [ + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + "FuseGelu", + "ClipToMinMax", + ], + ), + ("onnx-standard", "cpu", []), + ("migraphx", "gpu", []), + ("rocm", "gpu", []), + ("vitisai", "npu", []), + ], +) +def test_mobius_ep_surgery_profiles(ep, device, expected): + from olive.cli.capture_onnx import _resolve_mobius_ep_profile + + _, surgeries = _resolve_mobius_ep_profile(ep, device) + + assert [surgery["surgeon"] for surgery in surgeries] == expected + + +@pytest.mark.parametrize( + ("ep", "device", "expected"), + [ + ("cpu", "cpu", ["FLOAT"]), + ("cuda", "gpu", ["FLOAT16", "BFLOAT16"]), + ("dml", "gpu", ["FLOAT16"]), + ("trt-rtx", "gpu", ["FLOAT16", "BFLOAT16"]), + ("webgpu", "gpu", ["FLOAT", "FLOAT16"]), + ], +) +def test_mobius_gqa_profiles_filter_using_exported_graph_dtype(ep, device, expected): + from olive.cli.capture_onnx import _resolve_mobius_ep_profile + + _, surgeries = _resolve_mobius_ep_profile(ep, device) + + assert surgeries[0] == { + "surgeon": "AttentionToGroupQueryAttention", + "supported_dtypes": expected, + } + + @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) @pytest.mark.parametrize("conflicting_flag", ["--use_model_builder", "--use_dynamo_exporter"]) diff --git a/test/passes/onnx/test_graph_surgeries_attention.py b/test/passes/onnx/test_graph_surgeries_attention.py index 862c49e98..36f152491 100644 --- a/test/passes/onnx/test_graph_surgeries_attention.py +++ b/test/passes/onnx/test_graph_surgeries_attention.py @@ -85,7 +85,7 @@ def _run_surgeries(tmp_path, model, *surgeons): graph_surgeries = create_pass_from_dict( GraphSurgeries, { - "surgeries": [{"surgeon": surgeon} for surgeon in surgeons], + "surgeries": [{"surgeon": surgeon} if isinstance(surgeon, str) else surgeon for surgeon in surgeons], "remove_duplicate_initializers": False, }, disable_search=True, @@ -131,6 +131,7 @@ def _make_attention_model( with_cache=True, k_dim=32, v_dim=32, + dtype=ir.DataType.FLOAT16, ): nodes = [] inputs = [_value("attention_mask", ir.DataType.INT64, [1, "total_sequence"])] @@ -142,19 +143,19 @@ def _make_attention_model( ) inputs.extend(mask_inputs) - cos_cache = _value("cos_cache", ir.DataType.FLOAT16, [128, 8]) - sin_cache = _value("sin_cache", ir.DataType.FLOAT16, [128, 8]) + cos_cache = _value("cos_cache", dtype, [128, 8]) + sin_cache = _value("sin_cache", dtype, [128, 8]) position_ids = _value("position_ids", ir.DataType.INT64, [1, 2]) if rotary: inputs.extend([cos_cache, sin_cache, position_ids]) graph_outputs = [] for layer in range(num_layers): - q = _value(f"q_{layer}", ir.DataType.FLOAT16, [1, 2, 64]) - k = _value(f"k_{layer}", ir.DataType.FLOAT16, [1, 2, k_dim]) - v = _value(f"v_{layer}", ir.DataType.FLOAT16, [1, 2, v_dim]) - past_key = _value(f"past_key_{layer}", ir.DataType.FLOAT16, [1, 2, 4, k_dim]) - past_value = _value(f"past_value_{layer}", ir.DataType.FLOAT16, [1, 2, 4, v_dim]) + q = _value(f"q_{layer}", dtype, [1, 2, 64]) + k = _value(f"k_{layer}", dtype, [1, 2, k_dim]) + v = _value(f"v_{layer}", dtype, [1, 2, v_dim]) + past_key = _value(f"past_key_{layer}", dtype, [1, 2, 4, k_dim]) + past_value = _value(f"past_value_{layer}", dtype, [1, 2, 4, v_dim]) inputs.extend([q, k, v]) if with_cache: inputs.extend([past_key, past_value]) @@ -196,7 +197,7 @@ def _make_attention_model( attention.outputs, ([1, 2, 64], [1, 2, 6, k_dim], [1, 2, 6, v_dim]), ): - output.dtype = ir.DataType.FLOAT16 + output.dtype = dtype output.shape = ir.Shape(shape) graph_outputs.extend(attention.outputs) @@ -470,6 +471,26 @@ def test_attention_to_gqa_fallback_uses_external_rope_and_preserves_three_output assert len(rewritten.graph.outputs) == 3 +@pytest.mark.parametrize( + ("dtype", "expected_gqa"), + [(ir.DataType.FLOAT16, 1), (ir.DataType.FLOAT, 0)], +) +@pytest.mark.parametrize("rotary", [False, True]) +def test_attention_to_gqa_filters_target_unsupported_dtypes(tmp_path, dtype, expected_gqa, rotary): + model = _make_attention_model(dtype=dtype, rotary=rotary) + rewritten = _run_surgeries( + tmp_path, + model, + { + "surgeon": "AttentionToGroupQueryAttention", + "supported_dtypes": ["FLOAT16", "BFLOAT16"], + }, + ) + + assert _count_ops(rewritten)["GroupQueryAttention"] == expected_gqa + assert _count_ops(rewritten)["Attention"] == 1 - expected_gqa + + @pytest.mark.parametrize( "model", [ diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index 0c8d91ee1..250c7b654 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -211,7 +211,7 @@ def test_default_config_params(): def test_is_not_accelerator_agnostic(): - """Pass must be EP-specific because it chooses fused ops based on the EP.""" + """Pass remains EP-specific because runtime packaging records the target EP.""" accelerator_spec = AcceleratorSpec( accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider ) @@ -222,10 +222,16 @@ def test_ep_map_covers_common_providers(): assert ExecutionProvider.CPUExecutionProvider in MobiusBuilder.EP_MAP assert ExecutionProvider.CUDAExecutionProvider in MobiusBuilder.EP_MAP assert ExecutionProvider.DmlExecutionProvider in MobiusBuilder.EP_MAP + assert ExecutionProvider.NvTensorRTRTXExecutionProvider in MobiusBuilder.EP_MAP + assert ExecutionProvider.OpenVINOExecutionProvider in MobiusBuilder.EP_MAP + assert ExecutionProvider.QNNExecutionProvider in MobiusBuilder.EP_MAP assert ExecutionProvider.WebGpuExecutionProvider in MobiusBuilder.EP_MAP assert MobiusBuilder.EP_MAP[ExecutionProvider.CPUExecutionProvider] == "cpu" assert MobiusBuilder.EP_MAP[ExecutionProvider.CUDAExecutionProvider] == "cuda" assert MobiusBuilder.EP_MAP[ExecutionProvider.DmlExecutionProvider] == "dml" + assert MobiusBuilder.EP_MAP[ExecutionProvider.NvTensorRTRTXExecutionProvider] == "trt-rtx" + assert MobiusBuilder.EP_MAP[ExecutionProvider.OpenVINOExecutionProvider] == "openvino" + assert MobiusBuilder.EP_MAP[ExecutionProvider.QNNExecutionProvider] == "qnn" assert MobiusBuilder.EP_MAP[ExecutionProvider.WebGpuExecutionProvider] == "webgpu" @@ -256,7 +262,7 @@ def test_single_component_returns_onnx_handler(tmp_path): assert Path(result.model_path).exists() mock_build.assert_called_once() call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == "cpu" + assert call_kwargs["execution_provider"] == "onnx-standard" assert "dtype" not in call_kwargs @@ -406,8 +412,8 @@ def test_multi_component_returns_composite_handler(tmp_path): # --------------------------------------------------------------------------- -def test_ep_auto_detected_from_accelerator(tmp_path): - """Execution provider is determined by the Olive accelerator spec.""" +def test_export_is_standard_onnx_for_target_accelerator(tmp_path): + """The target accelerator affects runtime packaging, not Mobius graph construction.""" out = tmp_path / "out" pkg = _fake_pkg(["model"], out) @@ -421,12 +427,23 @@ def test_ep_auto_detected_from_accelerator(tmp_path): accelerator_spec=accelerator_spec, ) - with _patch_build(pkg) as mock_build: + with ( + patch("mobius.build", return_value=pkg) as mock_build, + patch.object(MobiusBuilder, "_write_genai_config", return_value={}) as mock_write, + ): p.run(_make_hf_model("org/model"), out) call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == "cuda" + assert call_kwargs["execution_provider"] == "onnx-standard" assert "dtype" not in call_kwargs + mock_write.assert_called_once_with( + pkg, + str(out), + "org/model", + "cuda", + revision=None, + trust_remote_code=False, + ) def test_hf_load_options_forwarded_to_build_and_genai_config(tmp_path): @@ -509,8 +526,8 @@ def test_write_genai_config_defaults_hf_load_options(tmp_path): ) -def test_unsupported_ep_falls_back_to_default(tmp_path): - """If accelerator EP is unsupported, pass should fall back to mobius default EP.""" +def test_unsupported_ep_uses_standard_export(tmp_path): + """An unsupported runtime EP does not change the standard ONNX export path.""" out = tmp_path / "out" pkg = _fake_pkg(["model"], out) @@ -530,11 +547,11 @@ def test_unsupported_ep_falls_back_to_default(tmp_path): p.run(_make_hf_model("org/model"), out) call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT + assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.ONNX_STANDARD -def test_none_execution_provider_falls_back_to_default(tmp_path): - """If execution_provider is None, pass should fall back to mobius default EP.""" +def test_none_execution_provider_uses_standard_export(tmp_path): + """An unspecified runtime EP does not change the standard ONNX export path.""" out = tmp_path / "out" pkg = _fake_pkg(["model"], out) @@ -551,7 +568,7 @@ def test_none_execution_provider_falls_back_to_default(tmp_path): p.run(_make_hf_model("org/model"), out) call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT + assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.ONNX_STANDARD @pytest.mark.skipif(not _HAS_REAL_MOBIUS, reason="mobius-onnx is not publicly available in CI yet") From 1410069da1d671fcd585193d35409e01ba161f06 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 17 Sep 2026 14:52:47 -0700 Subject: [PATCH 03/10] Align capture surgeries with Olive recipes --- olive/cli/capture_onnx.py | 86 ++------- olive/passes/onnx/graph_surgery/attention.py | 33 +--- test/cli/test_cli.py | 171 +++++------------- .../onnx/test_graph_surgeries_attention.py | 39 +--- 4 files changed, 83 insertions(+), 246 deletions(-) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index d2239ce23..5f7d969af 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -32,18 +32,11 @@ class ModelBuilderAccuracyLevel(IntEnum): _EP_ALIASES = { - "cpu": ExecutionProvider.CPUExecutionProvider, "cuda": ExecutionProvider.CUDAExecutionProvider, - "default": ExecutionProvider.CPUExecutionProvider, - "dml": ExecutionProvider.DmlExecutionProvider, - "migraphx": ExecutionProvider.MIGraphXExecutionProvider, - "onnx-standard": ExecutionProvider.CPUExecutionProvider, "openvino": ExecutionProvider.OpenVINOExecutionProvider, "qnn": ExecutionProvider.QNNExecutionProvider, - "rocm": ExecutionProvider.ROCMExecutionProvider, "trt-rtx": ExecutionProvider.NvTensorRTRTXExecutionProvider, "vitisai": ExecutionProvider.VitisAIExecutionProvider, - "webgpu": ExecutionProvider.WebGpuExecutionProvider, } @@ -51,67 +44,28 @@ def _surgery(name: str, **kwargs) -> dict: return {"surgeon": name, **kwargs} -def _resolve_mobius_ep_profile(ep: str, device: str) -> tuple[ExecutionProvider, list[dict]]: +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 == "qnn" and device == "gpu": - raise ValueError("QNN GPU has no generic Mobius graph-surgery profile; use --device npu for QNN HTP.") - - clip = _surgery("ClipToMinMax") - gelu = _surgery("FuseGelu") - skip_norms = [ - _surgery("FuseSkipRMSNormalization"), - _surgery("FuseSkipLayerNormalization"), - ] - - if ep == "default": - return provider, [*skip_norms, gelu, clip] - if ep in ("onnx-standard", "migraphx", "rocm", "vitisai"): - return provider, [] - if ep == "openvino": - return provider, [gelu, clip] - if ep == "qnn": - return provider, [ - gelu, - _surgery("SeparateGroupQueryAttentionRoPE"), - _surgery("UnpackGroupQueryAttentionQKV"), - _surgery("Rank4RMSNormToRank3"), - _surgery("DecomposeOnnxRotaryEmbedding"), - _surgery("TensorScatterToScatterND"), - _surgery("DecomposeAttention"), - clip, - ] - supported_dtypes = { - "cpu": ["FLOAT"], - "cuda": ["FLOAT16", "BFLOAT16"], - "dml": ["FLOAT16"], - "trt-rtx": ["FLOAT16", "BFLOAT16"], - "webgpu": ["FLOAT", "FLOAT16"], - }[ep] - attention = _surgery("AttentionToGroupQueryAttention", supported_dtypes=supported_dtypes) - - if ep == "dml": - return provider, [ - attention, - *skip_norms, - gelu, - _surgery("SeparateGroupQueryAttentionRoPE"), - _surgery("UnpackGroupQueryAttentionQKV"), - clip, + if ep == "cuda": + return provider, [_surgery("TieWordEmbeddings")] + if ep == "qnn": + surgeries = [ + _surgery("RemoveRopeMultiCache"), + _surgery("AttentionMaskToSequenceLengths"), ] + if device == "npu": + surgeries.extend( + [ + _surgery("RemoveGidxFromMatMulNBits"), + _surgery("SimplifiedLayerNormToL2Norm"), + ] + ) + return provider, surgeries - surgeries = [attention] - if ep in ("cpu", "cuda", "trt-rtx", "webgpu"): - surgeries.append(_surgery("PackQKVForGroupQueryAttention")) - if ep != "trt-rtx": - surgeries.extend(skip_norms) - surgeries.append(gelu) - if ep == "webgpu": - surgeries.append(_surgery("StaticEmptyKV")) - surgeries.append(clip) - return provider, surgeries + return provider, [] def parse_dim_dict(s): @@ -325,9 +279,9 @@ def _get_run_config(self, tempdir: str) -> dict: 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")) - mobius_surgeries = [] + recipe_surgeries = [] if self.args.use_mobius_builder and self.args.execution_provider: - provider, mobius_surgeries = _resolve_mobius_ep_profile(self.args.execution_provider, self.args.device) + provider, recipe_surgeries = _resolve_recipe_ep_profile(self.args.execution_provider, self.args.device) device = self.args.device elif self.args.use_mobius_builder: provider = ExecutionProvider.CPUExecutionProvider @@ -349,8 +303,8 @@ def _get_run_config(self, tempdir: str) -> dict: if self.args.use_mobius_builder: del config["passes"]["c"] del config["passes"]["m"] - if mobius_surgeries: - to_replace.append((("passes", "g", "surgeries"), mobius_surgeries)) + if recipe_surgeries: + to_replace.append((("passes", "g", "surgeries"), recipe_surgeries)) else: del config["passes"]["g"] elif is_diffusers_model: diff --git a/olive/passes/onnx/graph_surgery/attention.py b/olive/passes/onnx/graph_surgery/attention.py index 4a32c87e0..d1cbd1fa1 100644 --- a/olive/passes/onnx/graph_surgery/attention.py +++ b/olive/passes/onnx/graph_surgery/attention.py @@ -133,14 +133,9 @@ def _static_last_dim(value): return False -def _is_supported_dtype(value: ir.Value, supported_dtypes: frozenset[ir.DataType] | None) -> bool: - return supported_dtypes is None or value.dtype in supported_dtypes - - class _RotaryAttentionToGQA(RewriteRuleClassBase): - def __init__(self, supported_dtypes: frozenset[ir.DataType] | None = None): + def __init__(self): super().__init__() - self._supported_dtypes = supported_dtypes self._seqlens_k = None self._total_seq_len = None self._cos_cache = None @@ -160,10 +155,8 @@ def pattern(self, op, q_pre, k_pre, v, attention_bias, past_key, past_value, cos _outputs=["attn_out", "present_key", "present_value"], ) - def check(self, context, q_pre, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): + def check(self, context, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): result = MatchResult() - if not _is_supported_dtype(q_pre, self._supported_dtypes): - return result.fail("Attention dtype is unsupported by the target execution provider") if not _local_window_from_attention_bias(attention_bias).recognized: return result.fail("Attention bias cannot be represented by GroupQueryAttention") @@ -265,9 +258,8 @@ def rewrite( class _AttentionToGQA(RewriteRuleClassBase): - def __init__(self, supported_dtypes: frozenset[ir.DataType] | None = None): + def __init__(self): super().__init__() - self._supported_dtypes = supported_dtypes self._seqlens_k = None self._total_seq_len = None @@ -283,10 +275,8 @@ def pattern(self, op, q, k, v, attention_bias, past_key, past_value): _outputs=["attn_out", "present_key", "present_value"], ) - def check(self, context, q, attn_out, k, v, attention_bias, past_key, past_value, **_): + def check(self, context, attn_out, k, v, attention_bias, past_key, past_value, **_): result = MatchResult() - if not _is_supported_dtype(q, self._supported_dtypes): - return result.fail("Attention dtype is unsupported by the target execution provider") if not _local_window_from_attention_bias(attention_bias).recognized: return result.fail("Attention bias cannot be represented by GroupQueryAttention") @@ -370,21 +360,8 @@ def rewrite( class AttentionToGroupQueryAttention(RewriteRuleSurgeon): """Fuse decoder Attention, and standard RoPE when possible, into GQA.""" - def __init__(self, supported_dtypes: list[str] | None = None): - try: - self.supported_dtypes = ( - frozenset(ir.DataType[dtype.upper()] for dtype in supported_dtypes) if supported_dtypes else None - ) - except KeyError as exc: - raise ValueError(f"Unsupported ONNX dtype {exc.args[0]!r}") from exc - def rules(self) -> pattern.RewriteRuleSet: - return pattern.RewriteRuleSet( - [ - _RotaryAttentionToGQA.rule(self.supported_dtypes), - _AttentionToGQA.rule(self.supported_dtypes), - ] - ) + return pattern.RewriteRuleSet([_RotaryAttentionToGQA.rule(), _AttentionToGQA.rule()]) class _PackQKVForGQA(RewriteRuleClassBase): diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index b261ab1fc..b821120c1 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -604,50 +604,31 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision "cuda", "gpu", "CUDAExecutionProvider", - [ - "AttentionToGroupQueryAttention", - "PackQKVForGroupQueryAttention", - "FuseSkipRMSNormalization", - "FuseSkipLayerNormalization", - "FuseGelu", - "ClipToMinMax", - ], + ["TieWordEmbeddings"], ), ( - "trt-rtx", + "qnn", "gpu", - "NvTensorRTRTXExecutionProvider", + "QNNExecutionProvider", [ - "AttentionToGroupQueryAttention", - "PackQKVForGroupQueryAttention", - "FuseGelu", - "ClipToMinMax", + "RemoveRopeMultiCache", + "AttentionMaskToSequenceLengths", ], ), - ( - "openvino", - "npu", - "OpenVINOExecutionProvider", - ["FuseGelu", "ClipToMinMax"], - ), ( "qnn", "npu", "QNNExecutionProvider", [ - "FuseGelu", - "SeparateGroupQueryAttentionRoPE", - "UnpackGroupQueryAttentionQKV", - "Rank4RMSNormToRank3", - "DecomposeOnnxRotaryEmbedding", - "TensorScatterToScatterND", - "DecomposeAttention", - "ClipToMinMax", + "RemoveRopeMultiCache", + "AttentionMaskToSequenceLengths", + "RemoveGidxFromMatMulNBits", + "SimplifiedLayerNormToL2Norm", ], ), ], ) -def test_capture_onnx_command_adds_mobius_ep_surgeries(_, mock_run, ep, device, provider, surgeons, tmp_path): +def test_capture_onnx_command_adds_recipe_ep_surgeries(_, mock_run, ep, device, provider, surgeons, tmp_path): cli_main( [ "capture-onnx-graph", @@ -672,7 +653,17 @@ def test_capture_onnx_command_adds_mobius_ep_surgeries(_, mock_run, ep, device, @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) -def test_capture_onnx_command_mobius_ep_without_surgeries(_, mock_run, tmp_path): +@pytest.mark.parametrize( + ("ep", "device", "provider"), + [ + ("openvino", "cpu", "OpenVINOExecutionProvider"), + ("openvino", "gpu", "OpenVINOExecutionProvider"), + ("openvino", "npu", "OpenVINOExecutionProvider"), + ("trt-rtx", "gpu", "NvTensorRTRTXExecutionProvider"), + ("vitisai", "npu", "VitisAIExecutionProvider"), + ], +) +def test_capture_onnx_command_recipe_ep_without_surgeries(_, mock_run, ep, device, provider, tmp_path): cli_main( [ "capture-onnx-graph", @@ -682,17 +673,17 @@ def test_capture_onnx_command_mobius_ep_without_surgeries(_, mock_run, tmp_path) str(tmp_path / "output"), "--use_mobius_builder", "--execution_provider", - "vitisai", + ep, "--device", - "npu", + device, ] ) config = mock_run.call_args[0][0] assert list(config["passes"]) == ["b"] assert config["systems"]["local_system"]["accelerators"][0] == { - "device": "npu", - "execution_providers": ["VitisAIExecutionProvider"], + "device": device, + "execution_providers": [provider], } @@ -704,10 +695,10 @@ def test_capture_onnx_command_mobius_ep_without_surgeries(_, mock_run, tmp_path) (["--execution_provider", "cuda"], "must be provided together"), (["--device", "gpu"], "must be provided together"), (["--execution_provider", "cuda", "--device", "npu"], "does not support device"), - (["--execution_provider", "qnn", "--device", "gpu"], "has no generic Mobius"), + (["--execution_provider", "trt-rtx", "--device", "npu"], "does not support device"), ], ) -def test_capture_onnx_command_rejects_invalid_mobius_ep_device(_, __, extra_args, message, tmp_path): +def test_capture_onnx_command_rejects_invalid_recipe_ep_device(_, __, extra_args, message, tmp_path): command_args = [ "capture-onnx-graph", "-m", @@ -724,8 +715,8 @@ def test_capture_onnx_command_rejects_invalid_mobius_ep_device(_, __, extra_args @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) -def test_capture_onnx_command_rejects_ep_device_without_mobius(_, __, tmp_path): - with pytest.raises(ValueError, match="require --use_mobius_builder"): +def test_capture_onnx_command_rejects_execution_provider_without_recipe_profile(_, __, tmp_path): + with pytest.raises(SystemExit) as exc_info: cli_main( [ "capture-onnx-graph", @@ -733,98 +724,34 @@ def test_capture_onnx_command_rejects_ep_device_without_mobius(_, __, tmp_path): "dummy-model-id", "-o", str(tmp_path / "output"), + "--use_mobius_builder", "--execution_provider", - "cuda", + "dml", "--device", "gpu", ] ) + assert exc_info.value.code == 2 -@pytest.mark.parametrize( - ("ep", "device", "expected"), - [ - ( - "cpu", - "cpu", - [ - "AttentionToGroupQueryAttention", - "PackQKVForGroupQueryAttention", - "FuseSkipRMSNormalization", - "FuseSkipLayerNormalization", - "FuseGelu", - "ClipToMinMax", - ], - ), - ( - "dml", - "gpu", - [ - "AttentionToGroupQueryAttention", - "FuseSkipRMSNormalization", - "FuseSkipLayerNormalization", - "FuseGelu", - "SeparateGroupQueryAttentionRoPE", - "UnpackGroupQueryAttentionQKV", - "ClipToMinMax", - ], - ), - ( - "webgpu", - "gpu", - [ - "AttentionToGroupQueryAttention", - "PackQKVForGroupQueryAttention", - "FuseSkipRMSNormalization", - "FuseSkipLayerNormalization", - "FuseGelu", - "StaticEmptyKV", - "ClipToMinMax", - ], - ), - ( - "default", - "cpu", - [ - "FuseSkipRMSNormalization", - "FuseSkipLayerNormalization", - "FuseGelu", - "ClipToMinMax", - ], - ), - ("onnx-standard", "cpu", []), - ("migraphx", "gpu", []), - ("rocm", "gpu", []), - ("vitisai", "npu", []), - ], -) -def test_mobius_ep_surgery_profiles(ep, device, expected): - from olive.cli.capture_onnx import _resolve_mobius_ep_profile - - _, surgeries = _resolve_mobius_ep_profile(ep, device) - - assert [surgery["surgeon"] for surgery in surgeries] == expected - - -@pytest.mark.parametrize( - ("ep", "device", "expected"), - [ - ("cpu", "cpu", ["FLOAT"]), - ("cuda", "gpu", ["FLOAT16", "BFLOAT16"]), - ("dml", "gpu", ["FLOAT16"]), - ("trt-rtx", "gpu", ["FLOAT16", "BFLOAT16"]), - ("webgpu", "gpu", ["FLOAT", "FLOAT16"]), - ], -) -def test_mobius_gqa_profiles_filter_using_exported_graph_dtype(ep, device, expected): - from olive.cli.capture_onnx import _resolve_mobius_ep_profile - - _, surgeries = _resolve_mobius_ep_profile(ep, device) - assert surgeries[0] == { - "surgeon": "AttentionToGroupQueryAttention", - "supported_dtypes": expected, - } +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_rejects_ep_device_without_mobius(_, __, tmp_path): + with pytest.raises(ValueError, match="require --use_mobius_builder"): + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--execution_provider", + "cuda", + "--device", + "gpu", + ] + ) @patch("olive.workflows.run") diff --git a/test/passes/onnx/test_graph_surgeries_attention.py b/test/passes/onnx/test_graph_surgeries_attention.py index 36f152491..862c49e98 100644 --- a/test/passes/onnx/test_graph_surgeries_attention.py +++ b/test/passes/onnx/test_graph_surgeries_attention.py @@ -85,7 +85,7 @@ def _run_surgeries(tmp_path, model, *surgeons): graph_surgeries = create_pass_from_dict( GraphSurgeries, { - "surgeries": [{"surgeon": surgeon} if isinstance(surgeon, str) else surgeon for surgeon in surgeons], + "surgeries": [{"surgeon": surgeon} for surgeon in surgeons], "remove_duplicate_initializers": False, }, disable_search=True, @@ -131,7 +131,6 @@ def _make_attention_model( with_cache=True, k_dim=32, v_dim=32, - dtype=ir.DataType.FLOAT16, ): nodes = [] inputs = [_value("attention_mask", ir.DataType.INT64, [1, "total_sequence"])] @@ -143,19 +142,19 @@ def _make_attention_model( ) inputs.extend(mask_inputs) - cos_cache = _value("cos_cache", dtype, [128, 8]) - sin_cache = _value("sin_cache", dtype, [128, 8]) + cos_cache = _value("cos_cache", ir.DataType.FLOAT16, [128, 8]) + sin_cache = _value("sin_cache", ir.DataType.FLOAT16, [128, 8]) position_ids = _value("position_ids", ir.DataType.INT64, [1, 2]) if rotary: inputs.extend([cos_cache, sin_cache, position_ids]) graph_outputs = [] for layer in range(num_layers): - q = _value(f"q_{layer}", dtype, [1, 2, 64]) - k = _value(f"k_{layer}", dtype, [1, 2, k_dim]) - v = _value(f"v_{layer}", dtype, [1, 2, v_dim]) - past_key = _value(f"past_key_{layer}", dtype, [1, 2, 4, k_dim]) - past_value = _value(f"past_value_{layer}", dtype, [1, 2, 4, v_dim]) + q = _value(f"q_{layer}", ir.DataType.FLOAT16, [1, 2, 64]) + k = _value(f"k_{layer}", ir.DataType.FLOAT16, [1, 2, k_dim]) + v = _value(f"v_{layer}", ir.DataType.FLOAT16, [1, 2, v_dim]) + past_key = _value(f"past_key_{layer}", ir.DataType.FLOAT16, [1, 2, 4, k_dim]) + past_value = _value(f"past_value_{layer}", ir.DataType.FLOAT16, [1, 2, 4, v_dim]) inputs.extend([q, k, v]) if with_cache: inputs.extend([past_key, past_value]) @@ -197,7 +196,7 @@ def _make_attention_model( attention.outputs, ([1, 2, 64], [1, 2, 6, k_dim], [1, 2, 6, v_dim]), ): - output.dtype = dtype + output.dtype = ir.DataType.FLOAT16 output.shape = ir.Shape(shape) graph_outputs.extend(attention.outputs) @@ -471,26 +470,6 @@ def test_attention_to_gqa_fallback_uses_external_rope_and_preserves_three_output assert len(rewritten.graph.outputs) == 3 -@pytest.mark.parametrize( - ("dtype", "expected_gqa"), - [(ir.DataType.FLOAT16, 1), (ir.DataType.FLOAT, 0)], -) -@pytest.mark.parametrize("rotary", [False, True]) -def test_attention_to_gqa_filters_target_unsupported_dtypes(tmp_path, dtype, expected_gqa, rotary): - model = _make_attention_model(dtype=dtype, rotary=rotary) - rewritten = _run_surgeries( - tmp_path, - model, - { - "surgeon": "AttentionToGroupQueryAttention", - "supported_dtypes": ["FLOAT16", "BFLOAT16"], - }, - ) - - assert _count_ops(rewritten)["GroupQueryAttention"] == expected_gqa - assert _count_ops(rewritten)["Attention"] == 1 - expected_gqa - - @pytest.mark.parametrize( "model", [ From f0bdb444fb7cab5156502f8f84acdc8f93a64e73 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 17 Sep 2026 15:27:24 -0700 Subject: [PATCH 04/10] Add GQA to QNN surgery profiles --- olive/cli/capture_onnx.py | 1 + test/cli/test_cli.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 5f7d969af..cd09dea11 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -53,6 +53,7 @@ def _resolve_recipe_ep_profile(ep: str, device: str) -> tuple[ExecutionProvider, return provider, [_surgery("TieWordEmbeddings")] if ep == "qnn": surgeries = [ + _surgery("AttentionToGroupQueryAttention"), _surgery("RemoveRopeMultiCache"), _surgery("AttentionMaskToSequenceLengths"), ] diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index b821120c1..78c5cc8f5 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -611,6 +611,7 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision "gpu", "QNNExecutionProvider", [ + "AttentionToGroupQueryAttention", "RemoveRopeMultiCache", "AttentionMaskToSequenceLengths", ], @@ -620,6 +621,7 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision "npu", "QNNExecutionProvider", [ + "AttentionToGroupQueryAttention", "RemoveRopeMultiCache", "AttentionMaskToSequenceLengths", "RemoveGidxFromMatMulNBits", From d5b7307f5baa70f70630c3d6acbc781477720d62 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 17 Sep 2026 15:38:11 -0700 Subject: [PATCH 05/10] Align surgery profiles with ModelBuilder graphs --- olive/cli/capture_onnx.py | 24 +++++++++++- olive/passes/onnx/graph_surgery/attention.py | 33 +++++++++++++--- test/cli/test_cli.py | 21 +++++++++- .../onnx/test_graph_surgeries_attention.py | 39 ++++++++++++++----- 4 files changed, 100 insertions(+), 17 deletions(-) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index cd09dea11..67274b257 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -50,10 +50,20 @@ def _resolve_recipe_ep_profile(ep: str, device: str) -> tuple[ExecutionProvider, raise ValueError(f"Execution provider {ep!r} does not support device {device!r}.") if ep == "cuda": - return provider, [_surgery("TieWordEmbeddings")] + return provider, [ + _surgery( + "AttentionToGroupQueryAttention", + supported_dtypes=["FLOAT16", "BFLOAT16"], + ), + _surgery("PackQKVForGroupQueryAttention"), + _surgery("FuseSkipRMSNormalization"), + _surgery("FuseSkipLayerNormalization"), + _surgery("TieWordEmbeddings"), + ] if ep == "qnn": surgeries = [ _surgery("AttentionToGroupQueryAttention"), + _surgery("PackQKVForGroupQueryAttention"), _surgery("RemoveRopeMultiCache"), _surgery("AttentionMaskToSequenceLengths"), ] @@ -62,9 +72,21 @@ def _resolve_recipe_ep_profile(ep: str, device: str) -> tuple[ExecutionProvider, [ _surgery("RemoveGidxFromMatMulNBits"), _surgery("SimplifiedLayerNormToL2Norm"), + _surgery("Rank4RMSNormToRank3"), + _surgery("DecomposeOnnxRotaryEmbedding"), + _surgery("TensorScatterToScatterND"), + _surgery("DecomposeAttention"), ] ) return provider, surgeries + if ep == "trt-rtx": + return provider, [ + _surgery( + "AttentionToGroupQueryAttention", + supported_dtypes=["FLOAT16", "BFLOAT16"], + ), + _surgery("PackQKVForGroupQueryAttention"), + ] return provider, [] diff --git a/olive/passes/onnx/graph_surgery/attention.py b/olive/passes/onnx/graph_surgery/attention.py index d1cbd1fa1..4a32c87e0 100644 --- a/olive/passes/onnx/graph_surgery/attention.py +++ b/olive/passes/onnx/graph_surgery/attention.py @@ -133,9 +133,14 @@ def _static_last_dim(value): return False +def _is_supported_dtype(value: ir.Value, supported_dtypes: frozenset[ir.DataType] | None) -> bool: + return supported_dtypes is None or value.dtype in supported_dtypes + + class _RotaryAttentionToGQA(RewriteRuleClassBase): - def __init__(self): + def __init__(self, supported_dtypes: frozenset[ir.DataType] | None = None): super().__init__() + self._supported_dtypes = supported_dtypes self._seqlens_k = None self._total_seq_len = None self._cos_cache = None @@ -155,8 +160,10 @@ def pattern(self, op, q_pre, k_pre, v, attention_bias, past_key, past_value, cos _outputs=["attn_out", "present_key", "present_value"], ) - def check(self, context, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): + def check(self, context, q_pre, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): result = MatchResult() + if not _is_supported_dtype(q_pre, self._supported_dtypes): + return result.fail("Attention dtype is unsupported by the target execution provider") if not _local_window_from_attention_bias(attention_bias).recognized: return result.fail("Attention bias cannot be represented by GroupQueryAttention") @@ -258,8 +265,9 @@ def rewrite( class _AttentionToGQA(RewriteRuleClassBase): - def __init__(self): + def __init__(self, supported_dtypes: frozenset[ir.DataType] | None = None): super().__init__() + self._supported_dtypes = supported_dtypes self._seqlens_k = None self._total_seq_len = None @@ -275,8 +283,10 @@ def pattern(self, op, q, k, v, attention_bias, past_key, past_value): _outputs=["attn_out", "present_key", "present_value"], ) - def check(self, context, attn_out, k, v, attention_bias, past_key, past_value, **_): + def check(self, context, q, attn_out, k, v, attention_bias, past_key, past_value, **_): result = MatchResult() + if not _is_supported_dtype(q, self._supported_dtypes): + return result.fail("Attention dtype is unsupported by the target execution provider") if not _local_window_from_attention_bias(attention_bias).recognized: return result.fail("Attention bias cannot be represented by GroupQueryAttention") @@ -360,8 +370,21 @@ def rewrite( class AttentionToGroupQueryAttention(RewriteRuleSurgeon): """Fuse decoder Attention, and standard RoPE when possible, into GQA.""" + def __init__(self, supported_dtypes: list[str] | None = None): + try: + self.supported_dtypes = ( + frozenset(ir.DataType[dtype.upper()] for dtype in supported_dtypes) if supported_dtypes else None + ) + except KeyError as exc: + raise ValueError(f"Unsupported ONNX dtype {exc.args[0]!r}") from exc + def rules(self) -> pattern.RewriteRuleSet: - return pattern.RewriteRuleSet([_RotaryAttentionToGQA.rule(), _AttentionToGQA.rule()]) + return pattern.RewriteRuleSet( + [ + _RotaryAttentionToGQA.rule(self.supported_dtypes), + _AttentionToGQA.rule(self.supported_dtypes), + ] + ) class _PackQKVForGQA(RewriteRuleClassBase): diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index 78c5cc8f5..8dc359454 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -604,7 +604,19 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision "cuda", "gpu", "CUDAExecutionProvider", - ["TieWordEmbeddings"], + [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + "TieWordEmbeddings", + ], + ), + ( + "trt-rtx", + "gpu", + "NvTensorRTRTXExecutionProvider", + ["AttentionToGroupQueryAttention", "PackQKVForGroupQueryAttention"], ), ( "qnn", @@ -612,6 +624,7 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision "QNNExecutionProvider", [ "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", "RemoveRopeMultiCache", "AttentionMaskToSequenceLengths", ], @@ -622,10 +635,15 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision "QNNExecutionProvider", [ "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", "RemoveRopeMultiCache", "AttentionMaskToSequenceLengths", "RemoveGidxFromMatMulNBits", "SimplifiedLayerNormToL2Norm", + "Rank4RMSNormToRank3", + "DecomposeOnnxRotaryEmbedding", + "TensorScatterToScatterND", + "DecomposeAttention", ], ), ], @@ -661,7 +679,6 @@ def test_capture_onnx_command_adds_recipe_ep_surgeries(_, mock_run, ep, device, ("openvino", "cpu", "OpenVINOExecutionProvider"), ("openvino", "gpu", "OpenVINOExecutionProvider"), ("openvino", "npu", "OpenVINOExecutionProvider"), - ("trt-rtx", "gpu", "NvTensorRTRTXExecutionProvider"), ("vitisai", "npu", "VitisAIExecutionProvider"), ], ) diff --git a/test/passes/onnx/test_graph_surgeries_attention.py b/test/passes/onnx/test_graph_surgeries_attention.py index 862c49e98..36f152491 100644 --- a/test/passes/onnx/test_graph_surgeries_attention.py +++ b/test/passes/onnx/test_graph_surgeries_attention.py @@ -85,7 +85,7 @@ def _run_surgeries(tmp_path, model, *surgeons): graph_surgeries = create_pass_from_dict( GraphSurgeries, { - "surgeries": [{"surgeon": surgeon} for surgeon in surgeons], + "surgeries": [{"surgeon": surgeon} if isinstance(surgeon, str) else surgeon for surgeon in surgeons], "remove_duplicate_initializers": False, }, disable_search=True, @@ -131,6 +131,7 @@ def _make_attention_model( with_cache=True, k_dim=32, v_dim=32, + dtype=ir.DataType.FLOAT16, ): nodes = [] inputs = [_value("attention_mask", ir.DataType.INT64, [1, "total_sequence"])] @@ -142,19 +143,19 @@ def _make_attention_model( ) inputs.extend(mask_inputs) - cos_cache = _value("cos_cache", ir.DataType.FLOAT16, [128, 8]) - sin_cache = _value("sin_cache", ir.DataType.FLOAT16, [128, 8]) + cos_cache = _value("cos_cache", dtype, [128, 8]) + sin_cache = _value("sin_cache", dtype, [128, 8]) position_ids = _value("position_ids", ir.DataType.INT64, [1, 2]) if rotary: inputs.extend([cos_cache, sin_cache, position_ids]) graph_outputs = [] for layer in range(num_layers): - q = _value(f"q_{layer}", ir.DataType.FLOAT16, [1, 2, 64]) - k = _value(f"k_{layer}", ir.DataType.FLOAT16, [1, 2, k_dim]) - v = _value(f"v_{layer}", ir.DataType.FLOAT16, [1, 2, v_dim]) - past_key = _value(f"past_key_{layer}", ir.DataType.FLOAT16, [1, 2, 4, k_dim]) - past_value = _value(f"past_value_{layer}", ir.DataType.FLOAT16, [1, 2, 4, v_dim]) + q = _value(f"q_{layer}", dtype, [1, 2, 64]) + k = _value(f"k_{layer}", dtype, [1, 2, k_dim]) + v = _value(f"v_{layer}", dtype, [1, 2, v_dim]) + past_key = _value(f"past_key_{layer}", dtype, [1, 2, 4, k_dim]) + past_value = _value(f"past_value_{layer}", dtype, [1, 2, 4, v_dim]) inputs.extend([q, k, v]) if with_cache: inputs.extend([past_key, past_value]) @@ -196,7 +197,7 @@ def _make_attention_model( attention.outputs, ([1, 2, 64], [1, 2, 6, k_dim], [1, 2, 6, v_dim]), ): - output.dtype = ir.DataType.FLOAT16 + output.dtype = dtype output.shape = ir.Shape(shape) graph_outputs.extend(attention.outputs) @@ -470,6 +471,26 @@ def test_attention_to_gqa_fallback_uses_external_rope_and_preserves_three_output assert len(rewritten.graph.outputs) == 3 +@pytest.mark.parametrize( + ("dtype", "expected_gqa"), + [(ir.DataType.FLOAT16, 1), (ir.DataType.FLOAT, 0)], +) +@pytest.mark.parametrize("rotary", [False, True]) +def test_attention_to_gqa_filters_target_unsupported_dtypes(tmp_path, dtype, expected_gqa, rotary): + model = _make_attention_model(dtype=dtype, rotary=rotary) + rewritten = _run_surgeries( + tmp_path, + model, + { + "surgeon": "AttentionToGroupQueryAttention", + "supported_dtypes": ["FLOAT16", "BFLOAT16"], + }, + ) + + assert _count_ops(rewritten)["GroupQueryAttention"] == expected_gqa + assert _count_ops(rewritten)["Attention"] == 1 - expected_gqa + + @pytest.mark.parametrize( "model", [ From 852c92e0e629828dd2a56bc1e802da50bc78dbea Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Thu, 17 Sep 2026 15:45:06 -0700 Subject: [PATCH 06/10] Limit EP surgeries to generic decoder paths --- olive/cli/capture_onnx.py | 14 ++------------ test/cli/test_cli.py | 10 ++-------- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 67274b257..682b5a96f 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -58,26 +58,16 @@ def _resolve_recipe_ep_profile(ep: str, device: str) -> tuple[ExecutionProvider, _surgery("PackQKVForGroupQueryAttention"), _surgery("FuseSkipRMSNormalization"), _surgery("FuseSkipLayerNormalization"), - _surgery("TieWordEmbeddings"), ] if ep == "qnn": surgeries = [ _surgery("AttentionToGroupQueryAttention"), _surgery("PackQKVForGroupQueryAttention"), - _surgery("RemoveRopeMultiCache"), + _surgery("FuseSkipRMSNormalization"), _surgery("AttentionMaskToSequenceLengths"), ] if device == "npu": - surgeries.extend( - [ - _surgery("RemoveGidxFromMatMulNBits"), - _surgery("SimplifiedLayerNormToL2Norm"), - _surgery("Rank4RMSNormToRank3"), - _surgery("DecomposeOnnxRotaryEmbedding"), - _surgery("TensorScatterToScatterND"), - _surgery("DecomposeAttention"), - ] - ) + surgeries.append(_surgery("SimplifiedLayerNormToL2Norm")) return provider, surgeries if ep == "trt-rtx": return provider, [ diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index 8dc359454..c971458f9 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -609,7 +609,6 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision "PackQKVForGroupQueryAttention", "FuseSkipRMSNormalization", "FuseSkipLayerNormalization", - "TieWordEmbeddings", ], ), ( @@ -625,7 +624,7 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision [ "AttentionToGroupQueryAttention", "PackQKVForGroupQueryAttention", - "RemoveRopeMultiCache", + "FuseSkipRMSNormalization", "AttentionMaskToSequenceLengths", ], ), @@ -636,14 +635,9 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision [ "AttentionToGroupQueryAttention", "PackQKVForGroupQueryAttention", - "RemoveRopeMultiCache", + "FuseSkipRMSNormalization", "AttentionMaskToSequenceLengths", - "RemoveGidxFromMatMulNBits", "SimplifiedLayerNormToL2Norm", - "Rank4RMSNormToRank3", - "DecomposeOnnxRotaryEmbedding", - "TensorScatterToScatterND", - "DecomposeAttention", ], ), ], From cb38a08fc6a3e597fc87b635e19afe99e3e3a63e Mon Sep 17 00:00:00 2001 From: Xiaoyu Date: Mon, 21 Sep 2026 13:59:00 -0700 Subject: [PATCH 07/10] Apply EP surgeries after Mobius export --- olive/cli/capture_onnx.py | 74 +++++++++++++++++++++++++++++++++------ test/cli/test_cli.py | 74 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 133 insertions(+), 15 deletions(-) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 682b5a96f..6290ad306 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -20,6 +20,7 @@ ) 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 @@ -259,7 +260,26 @@ 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 self.args.execution_provider + or not workflow_output.has_output_model() + ): + return workflow_output + + _, surgeries = _resolve_recipe_ep_profile(self.args.execution_provider, self.args.device) + 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) + 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) @@ -292,9 +312,8 @@ def _get_run_config(self, tempdir: str) -> dict: 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")) - recipe_surgeries = [] if self.args.use_mobius_builder and self.args.execution_provider: - provider, recipe_surgeries = _resolve_recipe_ep_profile(self.args.execution_provider, self.args.device) + 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 @@ -316,14 +335,9 @@ def _get_run_config(self, tempdir: str) -> dict: if self.args.use_mobius_builder: del config["passes"]["c"] del config["passes"]["m"] - if recipe_surgeries: - to_replace.append((("passes", "g", "surgeries"), recipe_surgeries)) - else: - del config["passes"]["g"] elif is_diffusers_model: del config["passes"]["m"] del config["passes"]["b"] - del config["passes"]["g"] to_replace.extend( [ ( @@ -337,7 +351,6 @@ def _get_run_config(self, tempdir: str) -> dict: elif self.args.use_model_builder: del config["passes"]["c"] del config["passes"]["b"] - del config["passes"]["g"] to_replace.extend( [ (("passes", "m", "precision"), self.args.precision), @@ -359,7 +372,6 @@ def _get_run_config(self, tempdir: str) -> dict: to_replace.append((("passes", "m", "int4_accuracy_level"), self.args.int4_accuracy_level)) else: del config["passes"]["b"] - del config["passes"]["g"] to_replace.extend( [ ( @@ -401,6 +413,30 @@ def _get_run_config(self, tempdir: str) -> dict: return config + def _get_post_export_run_config(self, input_model_config: dict, surgeries: list[dict]) -> dict: + model_config = ModelConfig.model_validate(input_model_config) + components = model_config.get_components() + is_multimodal = components is not None + if 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 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 + accelerator["execution_providers"] = [_EP_ALIASES[self.args.execution_provider].value] + config["passes"]["g"]["surgeries"] = surgeries + update_shared_cache_options(config, self.args) + return config + TEMPLATE = { "systems": { @@ -415,10 +451,26 @@ def _get_run_config(self, tempdir: str) -> dict: }, "m": {"type": "ModelBuilder", "metadata_only": False}, "b": {"type": "MobiusBuilder"}, - "g": {"type": "GraphSurgeries"}, "f": {"type": "DynamicToFixedShape"}, }, "host": "local_system", "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, +} diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index c971458f9..9b46ef508 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -12,6 +12,7 @@ from olive.cli.base import TEST_OUTPUT_MARKER_FILE from olive.cli.launcher import main as cli_main +from olive.workflows.run.builds import parse_run_config @pytest.mark.parametrize("console_script", [True, False]) @@ -643,6 +644,14 @@ def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision ], ) def test_capture_onnx_command_adds_recipe_ep_surgeries(_, mock_run, ep, device, provider, surgeons, tmp_path): + exported_output = MagicMock() + exported_output.has_output_model.return_value = True + exported_output.get_best_candidate.return_value.olive_model_config = { + "type": "ONNXModel", + "config": {"model_path": str(tmp_path / "output"), "onnx_file_name": "model.onnx"}, + } + mock_run.side_effect = [exported_output, MagicMock()] + cli_main( [ "capture-onnx-graph", @@ -658,11 +667,68 @@ def test_capture_onnx_command_adds_recipe_ep_surgeries(_, mock_run, ep, device, ] ) - config = mock_run.call_args[0][0] - accelerator = config["systems"]["local_system"]["accelerators"][0] + export_config, surgery_config = [call.args[0] for call in mock_run.call_args_list] + assert list(export_config["passes"]) == ["b"] + assert "builds" not in export_config + accelerator = surgery_config["systems"]["local_system"]["accelerators"][0] assert accelerator == {"device": device, "execution_providers": [provider]} - assert list(config["passes"]) == ["b", "g"] - assert [surgery["surgeon"] for surgery in config["passes"]["g"]["surgeries"]] == surgeons + assert list(surgery_config["passes"]) == ["g"] + assert "builds" not in surgery_config + assert [surgery["surgeon"] for surgery in surgery_config["passes"]["g"]["surgeries"]] == surgeons + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_limits_recipe_ep_surgeries_to_multimodal_decoder(_, mock_run, tmp_path): + output_dir = tmp_path / "output" + exported_output = MagicMock() + exported_output.has_output_model.return_value = True + exported_output.get_best_candidate.return_value.olive_model_config = { + "type": "CompositeModel", + "config": { + "model_path": str(output_dir), + "model_components": [ + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "decoder")}}, + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "vision_encoder")}}, + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "embedding")}}, + ], + "model_component_names": ["decoder", "vision_encoder", "embedding"], + }, + } + mock_run.side_effect = [exported_output, {"decoder": MagicMock()}] + + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(output_dir), + "--use_mobius_builder", + "--execution_provider", + "cuda", + "--device", + "gpu", + ] + ) + + export_config, surgery_config = [call.args[0] for call in mock_run.call_args_list] + assert list(export_config["passes"]) == ["b"] + assert "builds" not in export_config + assert surgery_config["builds"] == { + "decoder": {"components": ["decoder"], "pipeline": ["g"]}, + } + assert "b" not in surgery_config["passes"] + assert [surgery["surgeon"] for surgery in surgery_config["passes"]["g"]["surgeries"]] == [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + ] + parsed = parse_run_config(surgery_config) + assert list(parsed) == ["decoder"] + assert list(parsed["decoder"].passes) == ["g"] + assert parsed["decoder"].engine.output_dir == (output_dir / "decoder").resolve() @patch("olive.workflows.run") From 8e41b9fa22051ccd3ba64eb8c8a59f4e596f9bd6 Mon Sep 17 00:00:00 2001 From: Xiaoyu Date: Mon, 21 Sep 2026 14:38:42 -0700 Subject: [PATCH 08/10] Forward target build contracts to Mobius --- olive/passes/onnx/mobius_model_builder.py | 8 ++++- test/passes/onnx/test_mobius_model_builder.py | 32 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index 62e7b752a..e4442473d 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -35,6 +35,9 @@ class MobiusBuilder(Pass): Single-component models return a plain :class:`~olive.model.ONNXModelHandler`. Mobius preserves the source model precision and exports a standard-ONNX graph; target-specific graph transformations run in downstream Olive passes. + The target accelerator is forwarded only as a structural build contract, so + requirements such as OpenVINO's rank-4 Gemma4 component interface are preserved + without enabling Mobius graph rewrites or introducing non-standard operators. Use ``components_to_export`` to export only a subset of components. This is useful when some components (e.g. a text decoder) are already exported and @@ -141,7 +144,8 @@ def _run_for_config( raise ValueError(f"MobiusBuilder requires an HfModelHandler input, got {type(model).__name__}.") # The graph is always exported through Mobius' standard-ONNX path. - # The requested EP is retained only for ORT GenAI runtime packaging. + # The requested EP/device provide structural build requirements and + # ORT GenAI runtime packaging without enabling Mobius graph rewrites. requested_ep = self.accelerator_spec.execution_provider runtime_ep: str = self.EP_MAP.get( requested_ep, @@ -185,6 +189,8 @@ def _run_for_config( model_id, revision=revision, execution_provider=self.MobiusEP.ONNX_STANDARD, + target_execution_provider=runtime_ep, + target_device=str(self.accelerator_spec.accelerator_type), load_weights=True, trust_remote_code=trust_remote_code, **text_only_kwargs, diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index 250c7b654..37c26490f 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -263,6 +263,8 @@ def test_single_component_returns_onnx_handler(tmp_path): mock_build.assert_called_once() call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == "onnx-standard" + assert call_kwargs["target_execution_provider"] == "cpu" + assert call_kwargs["target_device"] == "cpu" assert "dtype" not in call_kwargs @@ -435,6 +437,8 @@ def test_export_is_standard_onnx_for_target_accelerator(tmp_path): call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == "onnx-standard" + assert call_kwargs["target_execution_provider"] == "cuda" + assert call_kwargs["target_device"] == "gpu" assert "dtype" not in call_kwargs mock_write.assert_called_once_with( pkg, @@ -446,6 +450,30 @@ def test_export_is_standard_onnx_for_target_accelerator(tmp_path): ) +def test_openvino_target_contract_keeps_standard_export(tmp_path): + """OpenVINO structural requirements are independent of graph rewrites.""" + out = tmp_path / "out" + pkg = _fake_pkg(["model"], out) + accelerator_spec = AcceleratorSpec( + accelerator_type=Device.NPU, + execution_provider=ExecutionProvider.OpenVINOExecutionProvider, + ) + p = create_pass_from_dict( + MobiusBuilder, + {}, + disable_search=True, + accelerator_spec=accelerator_spec, + ) + + with _patch_build(pkg) as mock_build: + p.run(_make_hf_model("org/model"), out) + + call_kwargs = mock_build.call_args.kwargs + assert call_kwargs["execution_provider"] == "onnx-standard" + assert call_kwargs["target_execution_provider"] == "openvino" + assert call_kwargs["target_device"] == "npu" + + def test_hf_load_options_forwarded_to_build_and_genai_config(tmp_path): """Build and GenAI config generation resolve assets with the same Hugging Face options.""" out = tmp_path / "out" @@ -548,6 +576,8 @@ def test_unsupported_ep_uses_standard_export(tmp_path): call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.ONNX_STANDARD + assert call_kwargs["target_execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT + assert call_kwargs["target_device"] == "npu" def test_none_execution_provider_uses_standard_export(tmp_path): @@ -569,6 +599,8 @@ def test_none_execution_provider_uses_standard_export(tmp_path): call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.ONNX_STANDARD + assert call_kwargs["target_execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT + assert call_kwargs["target_device"] == "cpu" @pytest.mark.skipif(not _HAS_REAL_MOBIUS, reason="mobius-onnx is not publicly available in CI yet") From d798bf2007fb2cc2d3772e5017992298e1af0a77 Mon Sep 17 00:00:00 2001 From: Xiaoyu Date: Mon, 21 Sep 2026 16:02:27 -0700 Subject: [PATCH 09/10] Add explicit ONNX standard expansion --- docs/source/features/onnx-transformations.md | 1 + olive/cli/capture_onnx.py | 54 ++++++++--- olive/passes/onnx/graph_surgery/__init__.py | 2 + olive/passes/onnx/graph_surgery/lowering.py | 24 +++++ olive/passes/onnx/mobius_model_builder.py | 13 ++- test/cli/test_cli.py | 97 +++++++++++++++++++ .../onnx/test_graph_surgeries_lowering.py | 49 ++++++++++ test/passes/onnx/test_mobius_model_builder.py | 29 +++--- 8 files changed, 230 insertions(+), 39 deletions(-) diff --git a/docs/source/features/onnx-transformations.md b/docs/source/features/onnx-transformations.md index 7d713615c..e5368cbc7 100644 --- a/docs/source/features/onnx-transformations.md +++ b/docs/source/features/onnx-transformations.md @@ -142,6 +142,7 @@ 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`. | +| `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. | diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 6290ad306..cdd02d189 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -136,8 +136,8 @@ def register_subcommand(parser: ArgumentParser): help=( "Whether to use MobiusBuilder (mobius-onnx) to capture ONNX model. " "Supports multi-component multimodal models (VLMs). " - "Preserves the model precision and exports standard ONNX before applying " - "optional execution-provider graph surgeries. " + "Preserves model precision and exports Mobius' canonical graph before applying " + "optional strict-ONNX expansion or execution-provider graph surgeries. " "Requires 'pip install mobius-onnx'." ), ) @@ -153,6 +153,14 @@ def register_subcommand(parser: ArgumentParser): 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") @@ -261,22 +269,28 @@ def register_subcommand(parser: ArgumentParser): @action def run(self): workflow_output = self._run_workflow() - if ( - workflow_output is None - or not self.args.use_mobius_builder - or not self.args.execution_provider - or not workflow_output.has_output_model() - ): + if workflow_output is None or not self.args.use_mobius_builder or not workflow_output.has_output_model(): return workflow_output - _, surgeries = _resolve_recipe_ep_profile(self.args.execution_provider, self.args.device) + 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) + 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) @@ -306,6 +320,8 @@ def _get_run_config(self, tempdir: str) -> dict: 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 = ( @@ -413,26 +429,34 @@ def _get_run_config(self, tempdir: str) -> dict: return config - def _get_post_export_run_config(self, input_model_config: dict, surgeries: list[dict]) -> dict: + 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 is_multimodal and "decoder" not in components: + 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 is_multimodal: + 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 - accelerator["execution_providers"] = [_EP_ALIASES[self.args.execution_provider].value] + 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 diff --git a/olive/passes/onnx/graph_surgery/__init__.py b/olive/passes/onnx/graph_surgery/__init__.py index 9503be8fd..fde204258 100644 --- a/olive/passes/onnx/graph_surgery/__init__.py +++ b/olive/passes/onnx/graph_surgery/__init__.py @@ -17,6 +17,7 @@ ClipToMinMax, DecomposeAttention, DecomposeOnnxRotaryEmbedding, + InlineModelLocalFunctions, Rank4RMSNormToRank3, StaticEmptyKV, TensorScatterToScatterND, @@ -41,6 +42,7 @@ "FuseLayerNormalization", "FuseSkipLayerNormalization", "FuseSkipRMSNormalization", + "InlineModelLocalFunctions", "MoEGraphSurgeryError", "PackQKVForGroupQueryAttention", "ProtoSurgeon", diff --git a/olive/passes/onnx/graph_surgery/lowering.py b/olive/passes/onnx/graph_surgery/lowering.py index 35addbbeb..65652fa14 100644 --- a/olive/passes/onnx/graph_surgery/lowering.py +++ b/olive/passes/onnx/graph_surgery/lowering.py @@ -9,6 +9,7 @@ import numpy as np import onnx_ir as ir from onnx_ir import tape +from onnx_ir.passes.common import InlinePass, RemoveUnusedOpsetsPass from onnxscript.rewriter import pattern from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon, Surgeon @@ -17,6 +18,29 @@ # pylint: disable=arguments-differ _MASK_NEGATIVE_INFINITY = float("-inf") +_STANDARD_ONNX_DOMAINS = frozenset({"", "ai.onnx"}) + + +class InlineModelLocalFunctions(Surgeon): + """Inline model-local functions and require the resulting graph to use standard ONNX domains.""" + + def call_ir(self, model: ir.Model) -> ir.Model: + InlinePass()(model) + model.functions.clear() + RemoveUnusedOpsetsPass()(model) + non_standard_ops = sorted( + { + f"{node.domain}::{node.op_type}" + for node in model.graph.all_nodes() + if node.domain not in _STANDARD_ONNX_DOMAINS + } + ) + if non_standard_ops: + raise ValueError( + "Cannot produce strict ONNX because these operators have no model-local " + f"standard function body: {non_standard_ops}." + ) + return model class _BFloat16ClipRule(pattern.RewriteRuleClassBase): diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index e4442473d..87f595df0 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -33,11 +33,12 @@ class MobiusBuilder(Pass): pass returns a :class:`~olive.model.handler.composite.CompositeModelHandler` whose components are individual :class:`~olive.model.ONNXModelHandler` objects. Single-component models return a plain :class:`~olive.model.ONNXModelHandler`. - Mobius preserves the source model precision and exports a standard-ONNX graph; - target-specific graph transformations run in downstream Olive passes. + Mobius preserves the source model precision and exports its canonical graph + without automatic rewrites; strict-ONNX expansion and target-specific graph + transformations run in downstream Olive passes. The target accelerator is forwarded only as a structural build contract, so requirements such as OpenVINO's rank-4 Gemma4 component interface are preserved - without enabling Mobius graph rewrites or introducing non-standard operators. + without enabling Mobius graph rewrites. Use ``components_to_export`` to export only a subset of components. This is useful when some components (e.g. a text decoder) are already exported and @@ -143,7 +144,6 @@ def _run_for_config( if not isinstance(model, HfModelHandler): raise ValueError(f"MobiusBuilder requires an HfModelHandler input, got {type(model).__name__}.") - # The graph is always exported through Mobius' standard-ONNX path. # The requested EP/device provide structural build requirements and # ORT GenAI runtime packaging without enabling Mobius graph rewrites. requested_ep = self.accelerator_spec.execution_provider @@ -188,9 +188,8 @@ def _run_for_config( pkg = build( model_id, revision=revision, - execution_provider=self.MobiusEP.ONNX_STANDARD, - target_execution_provider=runtime_ep, - target_device=str(self.accelerator_spec.accelerator_type), + execution_provider=runtime_ep, + device=str(self.accelerator_spec.accelerator_type), load_weights=True, trust_remote_code=trust_remote_code, **text_only_kwargs, diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index 9b46ef508..d29b57d74 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -731,6 +731,87 @@ def test_capture_onnx_command_limits_recipe_ep_surgeries_to_multimodal_decoder(_ assert parsed["decoder"].engine.output_dir == (output_dir / "decoder").resolve() +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_onnx_standard_inlines_all_multimodal_components(_, mock_run, tmp_path): + output_dir = tmp_path / "output" + exported_output = MagicMock() + exported_output.has_output_model.return_value = True + exported_output.get_best_candidate.return_value.olive_model_config = { + "type": "CompositeModel", + "config": { + "model_path": str(output_dir), + "model_components": [ + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "decoder")}}, + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "vision_encoder")}}, + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "embedding")}}, + ], + "model_component_names": ["decoder", "vision_encoder", "embedding"], + }, + } + mock_run.side_effect = [exported_output, MagicMock()] + + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(output_dir), + "--use_mobius_builder", + "--execution_provider", + "openvino", + "--device", + "npu", + "--onnx_standard", + ] + ) + + _, standard_config = [call.args[0] for call in mock_run.call_args_list] + assert "builds" not in standard_config + assert standard_config["passes"]["g"]["surgeries"] == [ + {"surgeon": "InlineModelLocalFunctions"}, + ] + assert standard_config["systems"]["local_system"]["accelerators"][0] == { + "device": "npu", + "execution_providers": ["OpenVINOExecutionProvider"], + } + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_onnx_standard_without_target_ep(_, mock_run, tmp_path): + output_dir = tmp_path / "output" + exported_output = MagicMock() + exported_output.has_output_model.return_value = True + exported_output.get_best_candidate.return_value.olive_model_config = { + "type": "ONNXModel", + "config": {"model_path": str(output_dir), "onnx_file_name": "model.onnx"}, + } + mock_run.side_effect = [exported_output, MagicMock()] + + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(output_dir), + "--use_mobius_builder", + "--onnx_standard", + ] + ) + + _, standard_config = [call.args[0] for call in mock_run.call_args_list] + assert standard_config["passes"]["g"]["surgeries"] == [ + {"surgeon": "InlineModelLocalFunctions"}, + ] + assert standard_config["systems"]["local_system"]["accelerators"][0] == { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"], + } + + @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) @pytest.mark.parametrize( @@ -833,6 +914,22 @@ def test_capture_onnx_command_rejects_ep_device_without_mobius(_, __, tmp_path): ) +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_rejects_onnx_standard_without_mobius(_, __, tmp_path): + with pytest.raises(ValueError, match="requires --use_mobius_builder"): + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--onnx_standard", + ] + ) + + @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) @pytest.mark.parametrize("conflicting_flag", ["--use_model_builder", "--use_dynamo_exporter"]) diff --git a/test/passes/onnx/test_graph_surgeries_lowering.py b/test/passes/onnx/test_graph_surgeries_lowering.py index 4c3d3a84b..8a8dd0cd7 100644 --- a/test/passes/onnx/test_graph_surgeries_lowering.py +++ b/test/passes/onnx/test_graph_surgeries_lowering.py @@ -11,6 +11,7 @@ import onnx_ir as ir import onnxruntime as ort import pytest +from onnx import TensorProto, helper from onnx.reference import ReferenceEvaluator from olive.model import ONNXModelHandler @@ -59,6 +60,7 @@ def _metadata(value: ir.Value): def test_lowering_module_registers_all_surgeons(): expected = { lowering_surgeries.ClipToMinMax, + lowering_surgeries.InlineModelLocalFunctions, lowering_surgeries.Rank4RMSNormToRank3, lowering_surgeries.DecomposeOnnxRotaryEmbedding, lowering_surgeries.TensorScatterToScatterND, @@ -68,6 +70,53 @@ def test_lowering_module_registers_all_surgeons(): assert {Surgeon.registry[surgeon.__name__.lower()] for surgeon in expected} == expected +def _model_with_local_identity_function() -> ir.Model: + function = helper.make_function( + "com.microsoft", + "CustomIdentity", + ["x"], + ["y"], + [helper.make_node("Identity", ["x"], ["y"])], + [helper.make_opsetid("", 24)], + ) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1]) + graph = helper.make_graph( + [helper.make_node("CustomIdentity", ["x"], ["y"], domain="com.microsoft")], + "custom_identity", + [x], + [y], + ) + model = helper.make_model( + graph, + functions=[function], + opset_imports=[ + helper.make_opsetid("", 24), + helper.make_opsetid("com.microsoft", 1), + ], + ) + return ir.from_proto(model) + + +def test_inline_model_local_functions_produces_standard_onnx(tmp_path): + output = _apply_surgery( + tmp_path, + _model_with_local_identity_function(), + "InlineModelLocalFunctions", + ) + + assert _counts(output) == Counter({"Identity": 1}) + assert not output.functions + + +def test_inline_model_local_functions_rejects_custom_op_without_function(tmp_path): + model = _model_with_local_identity_function() + model.functions.clear() + + with pytest.raises(ValueError, match="no model-local standard function body"): + _apply_surgery(tmp_path, model, "InlineModelLocalFunctions") + + def _clip_model( dtype: ir.DataType, *, diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index 37c26490f..0ffee8199 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -262,9 +262,8 @@ def test_single_component_returns_onnx_handler(tmp_path): assert Path(result.model_path).exists() mock_build.assert_called_once() call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == "onnx-standard" - assert call_kwargs["target_execution_provider"] == "cpu" - assert call_kwargs["target_device"] == "cpu" + assert call_kwargs["execution_provider"] == "cpu" + assert call_kwargs["device"] == "cpu" assert "dtype" not in call_kwargs @@ -414,8 +413,8 @@ def test_multi_component_returns_composite_handler(tmp_path): # --------------------------------------------------------------------------- -def test_export_is_standard_onnx_for_target_accelerator(tmp_path): - """The target accelerator affects runtime packaging, not Mobius graph construction.""" +def test_target_accelerator_is_forwarded_without_enabling_mobius_rewrites(tmp_path): + """The target accelerator controls the build contract and runtime packaging.""" out = tmp_path / "out" pkg = _fake_pkg(["model"], out) @@ -436,9 +435,8 @@ def test_export_is_standard_onnx_for_target_accelerator(tmp_path): p.run(_make_hf_model("org/model"), out) call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == "onnx-standard" - assert call_kwargs["target_execution_provider"] == "cuda" - assert call_kwargs["target_device"] == "gpu" + assert call_kwargs["execution_provider"] == "cuda" + assert call_kwargs["device"] == "gpu" assert "dtype" not in call_kwargs mock_write.assert_called_once_with( pkg, @@ -469,9 +467,8 @@ def test_openvino_target_contract_keeps_standard_export(tmp_path): p.run(_make_hf_model("org/model"), out) call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == "onnx-standard" - assert call_kwargs["target_execution_provider"] == "openvino" - assert call_kwargs["target_device"] == "npu" + assert call_kwargs["execution_provider"] == "openvino" + assert call_kwargs["device"] == "npu" def test_hf_load_options_forwarded_to_build_and_genai_config(tmp_path): @@ -575,9 +572,8 @@ def test_unsupported_ep_uses_standard_export(tmp_path): p.run(_make_hf_model("org/model"), out) call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.ONNX_STANDARD - assert call_kwargs["target_execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT - assert call_kwargs["target_device"] == "npu" + assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT + assert call_kwargs["device"] == "npu" def test_none_execution_provider_uses_standard_export(tmp_path): @@ -598,9 +594,8 @@ def test_none_execution_provider_uses_standard_export(tmp_path): p.run(_make_hf_model("org/model"), out) call_kwargs = mock_build.call_args.kwargs - assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.ONNX_STANDARD - assert call_kwargs["target_execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT - assert call_kwargs["target_device"] == "cpu" + assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT + assert call_kwargs["device"] == "cpu" @pytest.mark.skipif(not _HAS_REAL_MOBIUS, reason="mobius-onnx is not publicly available in CI yet") From b6f23004763a61bff13c235fce555251d830e138 Mon Sep 17 00:00:00 2001 From: Xiaoyu Date: Tue, 22 Sep 2026 11:58:00 -0700 Subject: [PATCH 10/10] Migrate FP8 KV cache surgery to Olive --- docs/source/features/onnx-transformations.md | 1 + olive/passes/onnx/graph_surgery/__init__.py | 2 + olive/passes/onnx/graph_surgery/attention.py | 98 ++++++++++++++++++- .../onnx/test_graph_surgeries_attention.py | 48 +++++++++ 4 files changed, 148 insertions(+), 1 deletion(-) diff --git a/docs/source/features/onnx-transformations.md b/docs/source/features/onnx-transformations.md index e5368cbc7..558009008 100644 --- a/docs/source/features/onnx-transformations.md +++ b/docs/source/features/onnx-transformations.md @@ -142,6 +142,7 @@ 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. | diff --git a/olive/passes/onnx/graph_surgery/__init__.py b/olive/passes/onnx/graph_surgery/__init__.py index fde204258..07e5fdce0 100644 --- a/olive/passes/onnx/graph_surgery/__init__.py +++ b/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, @@ -33,6 +34,7 @@ "AttentionToGroupQueryAttention", "BlockDiagonalAttentionToPackedMHA", "ClipToMinMax", + "ConvertGroupQueryAttentionKVCacheToFp8", "DecomposeAttention", "DecomposeOnnxRotaryEmbedding", "FuseBiasGelu", diff --git a/olive/passes/onnx/graph_surgery/attention.py b/olive/passes/onnx/graph_surgery/attention.py index 4a32c87e0..aa5a1ff42 100644 --- a/olive/passes/onnx/graph_surgery/attention.py +++ b/olive/passes/onnx/graph_surgery/attention.py @@ -6,6 +6,10 @@ from __future__ import annotations +import math +import re +import warnings + import numpy as np import onnx_ir as ir from onnxscript.rewriter import pattern @@ -13,11 +17,103 @@ from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase from olive.constants import MSFT_DOMAIN -from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon +from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon, Surgeon # ONNXScript binds each rule's named pattern operands to its callbacks. # pylint: disable=arguments-differ +_FP8 = ir.DataType.FLOAT8E4M3FN +_K_SCALE_INDEX = 12 +_V_SCALE_INDEX = 13 +_MIN_GQA_INPUTS_WITH_SCALES = 14 +_LAYER_ID_RE = re.compile(r"\.(\d+)\.") + + +def _validate_fp8_scales(scales: dict[int, tuple[float, float]]) -> dict[int, tuple[float, float]]: + validated = {} + for layer_id, pair in scales.items(): + try: + k_scale, v_scale = pair + k_scale, v_scale = float(k_scale), float(v_scale) + except (TypeError, ValueError) as error: + raise ValueError(f"Layer {layer_id!r} must provide numeric (k_scale, v_scale) values.") from error + if not (math.isfinite(k_scale) and k_scale > 0 and math.isfinite(v_scale) and v_scale > 0): + raise ValueError(f"Layer {layer_id!r} FP8 scales must be finite and greater than zero.") + validated[int(layer_id)] = (k_scale, v_scale) + return validated + + +def _retype_fp8(value: ir.Value | None) -> None: + if value is None: + return + value.type = ir.TensorType(_FP8) + if value.const_value is not None and value.const_value.size == 0: + value.const_value = ir.tensor( + np.zeros(tuple(value.const_value.shape), dtype=_FP8.numpy()), + name=value.name, + ) + + +class ConvertGroupQueryAttentionKVCacheToFp8(Surgeon): + """Convert GroupQueryAttention past/present KV-cache tensors to FP8 E4M3.""" + + def __init__(self, scales: dict[int, tuple[float, float]] | None = None): + self.scales = _validate_fp8_scales(scales or {}) + + @staticmethod + def _scale_initializer(graph: ir.Graph, name: str, value: float) -> ir.Value: + existing = graph.initializers.get(name) + if existing is not None: + return existing + scale = ir.Value( + name=name, + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([1]), + const_value=ir.tensor(np.array([value], dtype=np.float32), name=name), + ) + graph.initializers[name] = scale + return scale + + def call_ir(self, model: ir.Model) -> ir.Model: + converted = 0 + for node in model.graph: + if node.domain != MSFT_DOMAIN or node.op_type != "GroupQueryAttention": + continue + past_key = node.inputs[3] if len(node.inputs) > 4 else None + past_value = node.inputs[4] if len(node.inputs) > 4 else None + if past_key is None or past_value is None: + continue + if any(value.const_value is not None and value.const_value.size > 0 for value in (past_key, past_value)): + warnings.warn( + f"Skipping {node.name!r}: only graph-input or empty-placeholder KV caches can become FP8.", + stacklevel=2, + ) + continue + + layer_match = _LAYER_ID_RE.search(past_key.name or "") + layer_id = int(layer_match.group(1)) if layer_match else -1 + k_value, v_value = self.scales.get(layer_id, (1.0, 1.0)) + + _retype_fp8(past_key) + _retype_fp8(past_value) + _retype_fp8(node.outputs[1] if len(node.outputs) > 1 else None) + _retype_fp8(node.outputs[2] if len(node.outputs) > 2 else None) + + k_scale = self._scale_initializer(model.graph, f"{past_key.name}.fp8_scale", k_value) + v_scale = self._scale_initializer(model.graph, f"{past_value.name}.fp8_scale", v_value) + if len(node.inputs) < _MIN_GQA_INPUTS_WITH_SCALES: + node.resize_inputs(_MIN_GQA_INPUTS_WITH_SCALES) + node.replace_input_with(_K_SCALE_INDEX, k_scale) + node.replace_input_with(_V_SCALE_INDEX, v_scale) + node.attributes.add(ir.AttrString("k_quant_type", "PER_TENSOR")) + node.attributes.add(ir.AttrString("v_quant_type", "PER_TENSOR")) + node.attributes.add(ir.AttrInt64("kv_cache_bit_width", 8)) + converted += 1 + + if converted == 0: + raise ValueError("No retypable GroupQueryAttention KV cache was found.") + return model + def _initializer_dtype(value: ir.Value) -> ir.DataType | None: """Return an initializer's declared dtype, falling back to its tensor dtype.""" diff --git a/test/passes/onnx/test_graph_surgeries_attention.py b/test/passes/onnx/test_graph_surgeries_attention.py index 36f152491..466620fe2 100644 --- a/test/passes/onnx/test_graph_surgeries_attention.py +++ b/test/passes/onnx/test_graph_surgeries_attention.py @@ -423,6 +423,7 @@ def _make_block_diagonal_attention_model(*, false_bias=-10000.0, shared_segments def test_attention_surgeries_register_on_module_import(): expected = { "attentiontogroupqueryattention", + "convertgroupqueryattentionkvcachetofp8", "separategroupqueryattentionrope", "unpackgroupqueryattentionqkv", "blockdiagonalattentiontopackedmha", @@ -431,6 +432,53 @@ def test_attention_surgeries_register_on_module_import(): assert expected <= Surgeon.registry.keys() +def test_convert_gqa_kv_cache_to_fp8_adds_scales_and_retypes_io(tmp_path): + query = _value("query", ir.DataType.FLOAT16, [1, 1, 32]) + key = _value("key", ir.DataType.FLOAT16, [1, 1, 16]) + value = _value("value", ir.DataType.FLOAT16, [1, 1, 16]) + past_key = _value("past_key_values.0.key", ir.DataType.FLOAT16, [1, 1, 0, 16]) + past_value = _value("past_key_values.0.value", ir.DataType.FLOAT16, [1, 1, 0, 16]) + seqlens = _value("seqlens", ir.DataType.INT32, [1]) + total_sequence = _value("total_sequence", ir.DataType.INT32, []) + gqa = _node( + "GroupQueryAttention", + [query, key, value, past_key, past_value, seqlens, total_sequence], + domain=_MS_DOMAIN, + num_outputs=3, + output_names=["output", "present.0.key", "present.0.value"], + ) + model = _model( + [query, key, value, past_key, past_value, seqlens, total_sequence], + list(gqa.outputs), + [gqa], + ) + + rewritten = _run_surgeries( + tmp_path, + model, + { + "surgeon": "ConvertGroupQueryAttentionKVCacheToFp8", + "scales": {0: [0.25, 0.5]}, + }, + ) + rewritten_gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + + assert rewritten_gqa.inputs[3].dtype == ir.DataType.FLOAT8E4M3FN + assert rewritten_gqa.inputs[4].dtype == ir.DataType.FLOAT8E4M3FN + assert rewritten_gqa.outputs[1].dtype == ir.DataType.FLOAT8E4M3FN + assert rewritten_gqa.outputs[2].dtype == ir.DataType.FLOAT8E4M3FN + assert len(rewritten_gqa.inputs) == 14 + np.testing.assert_array_equal( + rewritten_gqa.inputs[12].const_value.numpy(), + np.array([0.25], dtype=np.float32), + ) + np.testing.assert_array_equal( + rewritten_gqa.inputs[13].const_value.numpy(), + np.array([0.5], dtype=np.float32), + ) + assert rewritten_gqa.attributes.get_int("kv_cache_bit_width") == 8 + + def test_attention_to_gqa_fuses_rotary_preserves_attributes_outputs_and_shared_inputs(tmp_path): model = _make_attention_model(rotary=True, num_layers=2, sliding_window=64, k_dim=32, v_dim=32) rewritten = _run_surgeries(tmp_path, model, "AttentionToGroupQueryAttention")