From ce6352cfab6d64b8b9d131f1c91341615b50763d Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Mon, 6 Jul 2026 14:26:24 -0700
Subject: [PATCH 01/15] Migrate bnb4, io_datatype_converter, and
 dynamic_to_fixed_shape passes to onnx_ir

- OnnxBnb4Quantization: traverse/replace MatMul->MatMulBnb4 nodes via onnx_ir
  instead of the ORT MatMulBnb4Quantizer proto wrapper, keeping the native
  quantize_matmul_bnb4 kernel for FP4/NF4 packing. Preserves graph output names.
- OnnxIODataTypeConverter: reimplement input/output Cast insertion and rewiring
  on onnx_ir.
- DynamicToFixedShape: add IR-native fix_dim_params_ir/fix_input_shapes_ir helpers
  in common.py (native reimpl of ORT make_dim_param_fixed/make_input_shape_fixed/
  remove_invalid_dim_values over ir.Graph incl. subgraphs); output shape inference
  still uses ORT SymbolicShapeInference. Proto helpers kept for static_llm.
- Add OpType.MatMulBnb4 and OpType.Cast.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/constants.py                          |   2 +
 olive/passes/onnx/bnb_quantization.py       | 193 +++++++++++++++-----
 olive/passes/onnx/common.py                 | 160 ++++++++++++++++
 olive/passes/onnx/dynamic_to_fixed_shape.py |  16 +-
 olive/passes/onnx/io_datatype_converter.py  | 141 ++++++--------
 test/passes/onnx/test_bnb_quantization.py   |   2 +-
 6 files changed, 384 insertions(+), 130 deletions(-)

diff --git a/olive/constants.py b/olive/constants.py
index 968e970053..5df2b4e26a 100644
--- a/olive/constants.py
+++ b/olive/constants.py
@@ -96,6 +96,7 @@ class OpType(StrEnumBase):
     Gather = "Gather"
     GatherBlockQuantized = "GatherBlockQuantized"
     MatMulNBits = "MatMulNBits"
+    MatMulBnb4 = "MatMulBnb4"
     MatMul = "MatMul"
     QuickGelu = "QuickGelu"
     Sigmoid = "Sigmoid"
@@ -114,6 +115,7 @@ class OpType(StrEnumBase):
     PackedMultiHeadAttention = "PackedMultiHeadAttention"
     MultiHeadAttention = "MultiHeadAttention"
     Loop = "Loop"
+    Cast = "Cast"
 
 
 class AccuracyLevel(IntEnum):
diff --git a/olive/passes/onnx/bnb_quantization.py b/olive/passes/onnx/bnb_quantization.py
index 5894aa1a98..c5db8c9e98 100644
--- a/olive/passes/onnx/bnb_quantization.py
+++ b/olive/passes/onnx/bnb_quantization.py
@@ -7,19 +7,30 @@
 from pathlib import Path
 from typing import Optional
 
-import onnx
-from packaging import version
+import numpy as np
+import numpy.typing as npt
+import onnx_ir as ir
 
-from olive.constants import Precision
+from olive.constants import MSFT_DOMAIN, OpType, Precision
 from olive.hardware import AcceleratorSpec
 from olive.model import ONNXModelHandler
 from olive.model.utils import resolve_onnx_path
 from olive.passes import Pass
-from olive.passes.onnx.common import get_external_data_config, model_proto_to_olive_model
+from olive.passes.onnx.common import get_external_data_config, ir_model_to_olive_model
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
 
+# 4-bit quantization types, must be consistent with the native ORT MatMulBnb4 kernel
+# (Bnb_DataType_t defined in blockwise_quant_block_bnb4.h).
+_BNB4_QUANT_TYPES = {
+    # 4b floating point with bias of 3
+    "fp4": 0,
+    # 4b NormalFloat
+    "nf4": 1,
+}
+_BNB4_BLOCK_SIZE = 64
+
 
 class OnnxBnb4Quantization(Pass):
     """Quantize MatMul nodes in ONNX model using 4bit FP4/NF4 quantization."""
@@ -63,14 +74,6 @@ def validate_config(
     def _run_for_config(
         self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: str
     ) -> ONNXModelHandler:
-        from onnxruntime import __version__ as OrtVersion
-
-        assert version.parse(OrtVersion) >= version.parse("1.16.2"), (
-            "MatMulBnb4Quantizer is only supported in onnxruntime >= 1.16.2"
-        )
-
-        from onnxruntime.quantization.matmul_bnb4_quantizer import MatMulBnb4Quantizer
-
         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
 
         precision = config.precision.value if config.precision else None
@@ -93,50 +96,158 @@ def _run_for_config(
                     logger.info(
                         "bnb_4bit_use_double_quant is set to True but double quantization is not supported. Ignoring."
                     )
-        assert precision in {"fp4", "nf4"}, f"quant_type must be one of 'fp4' or 'nf4'. Got {precision}."
-        quant_type_enum = getattr(MatMulBnb4Quantizer, precision.upper())
+        assert precision in _BNB4_QUANT_TYPES, f"quant_type must be one of 'fp4' or 'nf4'. Got {precision}."
+        quant_type = _BNB4_QUANT_TYPES[precision]
 
         # load the model
-        onnx_model = model.load_model()
+        ir_model = model.load_ir_model()
+        ir.external_data.load_to_model(ir_model)
+        ir_model.graph.opset_imports[MSFT_DOMAIN] = 1
 
         # get nodes to exclude from quantization
         nodes_to_exclude = config.nodes_to_exclude or []
 
         # find all MatMul nodes in the graph
-        matmul_nodes = self._find_matmul_nodes(onnx_model.graph)
+        matmul_nodes = self._find_matmul_nodes(ir_model.graph)
         # filter based on quantized_modules
         quantized_modules = set(quantized_modules or [])
-        nodes_to_exclude = nodes_to_exclude + [
+        nodes_to_exclude = set(nodes_to_exclude) | {
             node
             for node in matmul_nodes
             if quantized_modules and not any(re.match(f".*[./]{key}[./]MatMul$", node) for key in quantized_modules)
-        ]
+        }
 
         # quantize the model
-        quantizer = MatMulBnb4Quantizer(
-            onnx_model, quant_type=quant_type_enum, block_size=64, nodes_to_exclude=nodes_to_exclude
-        )
-        quantizer.process()
-        # topologically sort the graph at the end since previous optimizations may have broken it
-        quantizer.model.topological_sort()
+        self._quantize_model(ir_model, quant_type, nodes_to_exclude)
 
         # save the model to the output path and return the model
-        return model_proto_to_olive_model(onnx_model, output_model_path, config)
+        return ir_model_to_olive_model(ir_model, output_model_path, config)
+
+    def _quantize_model(self, ir_model: ir.Model, quant_type: int, nodes_to_exclude: set[str]) -> None:
+        """Replace eligible MatMul nodes with MatMulBnb4 nodes carrying 4-bit quantized weights."""
+        ir_model.graph.sort()
+        for node in ir_model.graph.all_nodes():
+            if node.op_type != str(OpType.MatMul):
+                continue
+
+            if node.name in nodes_to_exclude:
+                logger.debug("exclude to quantize %s as specified by nodes_to_exclude...", node.name)
+                continue
+
+            quantized_node = self._quantize_matmul(node, quant_type)
+            if quantized_node is node:
+                # nothing changed (non-const or non-2D weight)
+                continue
+
+            weight_graph = node.inputs[1].graph
+            for input_value in quantized_node.inputs:
+                if input_value is not None and input_value.const_value is not None:
+                    weight_graph.register_initializer(input_value)
+
+            ir.convenience.replace_nodes_and_values(
+                node.graph, node, [node], [quantized_node], node.outputs, quantized_node.outputs
+            )
+
+        self._remove_unused_initializers(ir_model)
+
+    def _quantize_matmul(self, node: ir.Node, quant_type: int) -> ir.Node:
+        """Quantize weight B of a MatMul node to 4-bit and return the new MatMulBnb4 node.
+
+        Returns the original node unchanged if the weight is not a 2D constant.
+        """
+        logger.debug("start to quantize %s ...", node.name)
+
+        node_initializer = node.inputs[1]
+        if node_initializer is None or node_initializer.const_value is None:
+            logger.debug("MatMul doesn't have const weight. Skip to quantize")
+            return node
+
+        b_ndarray = node_initializer.const_value.numpy()
+        if len(b_ndarray.shape) != 2:
+            logger.debug("MatMul weight is not 2D. Skip to quantize")
+            return node
+
+        packed, absmax = self._bnb4_block_quant(b_ndarray, quant_type)
+
+        b_quant = ir.Value(name=node_initializer.name + "_Bnb4", const_value=ir.tensor(packed))
+        absmax_value = ir.Value(name=node_initializer.name + "_absmax", const_value=ir.tensor(absmax))
+
+        rows, cols = b_ndarray.shape
+        kwargs = {
+            "K": rows,
+            "N": cols,
+            "block_size": _BNB4_BLOCK_SIZE,
+            "quant_type": quant_type,
+        }
+
+        self._rename_output_unless_graph_output(node)
+
+        logger.debug("complete quantization of %s ...", node.name)
+
+        return ir.node(
+            domain=MSFT_DOMAIN,
+            op_type=str(OpType.MatMulBnb4),
+            inputs=[node.inputs[0], b_quant, absmax_value],
+            name=node.name + "_Bnb4" if node.name else "",
+            attributes=kwargs,
+        )
+
+    @staticmethod
+    def _bnb4_block_quant(fpweight: npt.ArrayLike, quant_type: int) -> tuple[np.ndarray, np.ndarray]:
+        """4b quantize fp32/fp16 weight using the native ORT bnb4 kernel."""
+        from onnxruntime.capi._pybind_state import quantize_matmul_bnb4
+
+        if len(fpweight.shape) != 2:
+            raise ValueError("Current bnb4 block quantization only supports 2D tensors!")
+        # need to copy since the transposed weight still has the original memory layout
+        # Linear4bit quantizes its weight data which is the transposed weight
+        fpweight_t = fpweight.transpose().copy()
+
+        rows, cols = fpweight.shape
+        numel = rows * cols
+        block_size = _BNB4_BLOCK_SIZE
+        num_blocks = (numel + block_size - 1) // block_size
+        quantized_numel = (numel + 1) // 2
+
+        packed = np.zeros(quantized_numel, dtype="uint8")
+        absmax = np.zeros(num_blocks, dtype=fpweight.dtype)
+        # block wise quantization, fpweight_t is flattened and divided into blocks
+        quantize_matmul_bnb4(packed, fpweight_t, absmax, block_size, quant_type, cols, rows)
+
+        return packed, absmax
+
+    @staticmethod
+    def _rename_output_unless_graph_output(node: ir.Node) -> None:
+        """Append a `_Bnb4` suffix to the node's output name for renaming after quantization.
+
+        Skips values that are graph outputs so external consumers relying on those names are
+        not broken. Internal tensors are safe to rename because their consumers are rewired by
+        replace_nodes_and_values.
+        """
+        graph = node.graph
+        is_graph_output = graph is not None and node.outputs[0] in graph.outputs
+        if not is_graph_output:
+            node.outputs[0].name = node.outputs[0].name + "_Bnb4"
+
+    @staticmethod
+    def _remove_unused_initializers(ir_model: ir.Model) -> None:
+        """Remove initializers that are no longer referenced by any node after quantization."""
+        used_names: set[str] = set()
+        for node in ir_model.graph.all_nodes():
+            for inp in node.inputs:
+                if inp is not None and inp.name:
+                    used_names.add(inp.name)
+        for out in ir_model.graph.outputs:
+            if out is not None and out.name:
+                used_names.add(out.name)
+
+        unused = [name for name in ir_model.graph.initializers if name not in used_names]
+        for name in unused:
+            del ir_model.graph.initializers[name]
+        if unused:
+            logger.debug("Removed %d unused initializers after quantization.", len(unused))
 
     @classmethod
-    def _find_matmul_nodes(cls, graph: onnx.GraphProto) -> list[str]:
-        """Find all MatMul nodes in the graph and return their names."""
-        matmul_nodes = []
-        for node in graph.node:
-            for attr in node.attribute:
-                if attr.type == onnx.AttributeProto.GRAPH:
-                    # recursive call to take care of sub-graph
-                    matmul_nodes += cls._find_matmul_nodes(attr.g)
-                elif attr.type == onnx.AttributeProto.GRAPHS:
-                    for subgraph in attr.graphs:
-                        # recursive call to take care of sub-graph
-                        matmul_nodes += cls._find_matmul_nodes(subgraph)
-            if node.op_type == "MatMul":
-                matmul_nodes.append(node.name)
-
-        return matmul_nodes
+    def _find_matmul_nodes(cls, graph: ir.Graph) -> list[str]:
+        """Find all MatMul nodes in the graph (including subgraphs) and return their names."""
+        return [node.name for node in graph.all_nodes() if node.op_type == str(OpType.MatMul)]
diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index b2df6c8705..c4221f6ccf 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -629,6 +629,166 @@ def fix_input_shapes(model_proto: onnx.ModelProto, input_names: list[str], input
     _fix_output_shapes(model_proto)
 
 
+def _ir_shape_is_fixed(value: Optional[ir.Value]) -> bool:
+    """Return True if the value has a shape where every dimension is a fixed positive integer."""
+    if value is None or value.shape is None:
+        return False
+    return all(isinstance(dim, int) and dim > 0 for dim in value.shape)
+
+
+def _iter_graphs_ir(graph: ir.Graph):
+    """Yield the graph and all of its subgraphs recursively (for control-flow nodes)."""
+    yield graph
+    for node in graph:
+        for attr in node.attributes.values():
+            if not isinstance(attr, ir.Attr):
+                continue
+            if attr.type == ir.AttributeType.GRAPH:
+                yield from _iter_graphs_ir(attr.value)
+            elif attr.type == ir.AttributeType.GRAPHS:
+                for subgraph in attr.value:
+                    yield from _iter_graphs_ir(subgraph)
+
+
+def _iter_shaped_values_ir(graph: ir.Graph):
+    """Yield all value objects in a single graph that may carry shape information."""
+    yield from graph.inputs
+    for node in graph:
+        yield from node.outputs
+    yield from graph.outputs
+
+
+def _make_dim_param_fixed_ir(graph: ir.Graph, param_name: str, value: int) -> None:
+    """Replace every occurrence of the symbolic dim ``param_name`` with ``value`` across the graph.
+
+    Mirrors onnxruntime.tools.onnx_model_utils.make_dim_param_fixed but operates on an ir.Graph,
+    including subgraphs.
+    """
+    for subgraph in _iter_graphs_ir(graph):
+        for val in _iter_shaped_values_ir(subgraph):
+            if val is None or val.shape is None:
+                continue
+            dims = list(val.shape)
+            changed = False
+            for idx, dim in enumerate(dims):
+                if isinstance(dim, ir.SymbolicDim) and dim.value == param_name:
+                    dims[idx] = value
+                    changed = True
+            if changed:
+                val.shape = ir.Shape(dims)
+
+
+def _remove_invalid_dim_values_ir(graph: ir.Graph) -> None:
+    """Unset any fixed dim values that are less than 1 (typically -1 placeholders for dynamic dims)."""
+    for subgraph in _iter_graphs_ir(graph):
+        for val in _iter_shaped_values_ir(subgraph):
+            if val is None or val.shape is None:
+                continue
+            dims = list(val.shape)
+            changed = False
+            for idx, dim in enumerate(dims):
+                if isinstance(dim, int) and dim < 1:
+                    dims[idx] = None
+                    changed = True
+            if changed:
+                val.shape = ir.Shape(dims)
+
+
+def _make_input_shape_fixed_ir(graph: ir.Graph, input_name: str, fixed_shape: list[int]) -> None:
+    """Set the shape of the named graph input to ``fixed_shape``.
+
+    Mirrors onnxruntime.tools.onnx_model_utils.make_input_shape_fixed but operates on an ir.Graph.
+    """
+    # remove any invalid dim values first. typically this is a dim_value of -1.
+    _remove_invalid_dim_values_ir(graph)
+
+    for graph_input in graph.inputs:
+        if graph_input.name != input_name:
+            continue
+
+        # graph inputs are required to have a shape to provide the rank
+        if graph_input.shape is None:
+            raise ValueError(f"Input {input_name} does not have a shape")
+
+        dims = list(graph_input.shape)
+        if len(dims) != len(fixed_shape):
+            raise ValueError(f"Rank mismatch. Existing:{len(dims)} Replacement:{len(fixed_shape)}")
+
+        new_dims = list(dims)
+        for idx, dim in enumerate(dims):
+            if isinstance(dim, int):
+                # check any existing fixed dims match
+                if dim != fixed_shape[idx]:
+                    raise ValueError(
+                        f"Can't replace existing fixed size of {dim} with {fixed_shape[idx]} for dimension {idx + 1}"
+                    )
+            elif isinstance(dim, ir.SymbolicDim) and dim.value is not None:
+                # replacing a dim_param so have to do that through the entire graph
+                _make_dim_param_fixed_ir(graph, dim.value, fixed_shape[idx])
+                new_dims[idx] = fixed_shape[idx]
+            else:
+                # replacing an unknown dim
+                new_dims[idx] = fixed_shape[idx]
+
+        graph_input.shape = ir.Shape(new_dims)
+        return
+
+    valid_names = ",".join(i.name for i in graph.inputs if i.name)
+    raise ValueError(f"Input {input_name} was not found in graph inputs. Valid input names are: {valid_names}")
+
+
+def _fix_output_shapes_ir(ir_model: ir.Model) -> None:
+    """Run shape inference on the model and update graph output shapes to make them fixed."""
+    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
+    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+
+    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
+    model_proto = ir.to_proto(ir_model)
+    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
+    inferred_outputs = {o.name: o for o in inferred_proto.graph.output}
+
+    for output in ir_model.graph.outputs:
+        if output is None or output.name is None or _ir_shape_is_fixed(output):
+            continue
+        new_o = inferred_outputs.get(output.name)
+        if new_o is not None and is_fixed_size_tensor(new_o):
+            output.shape = ir.Shape([dim.dim_value for dim in new_o.type.tensor_type.shape.dim])
+
+
+def fix_dim_params_ir(ir_model: ir.Model, dim_params: list[str], dim_values: list[int]) -> None:
+    """Fix the dimension parameters in an ir.Model.
+
+    :param dim_params: The dimension parameters to fix.
+    :param dim_values: The values to set for the dimension parameters.
+    """
+    dim_params = list(dim_params)
+    dim_values = list(dim_values)
+    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
+    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
+
+    for param, value in zip(dim_params, dim_values):
+        _make_dim_param_fixed_ir(ir_model.graph, param, value)
+
+    # update the output shapes to make them fixed
+    _fix_output_shapes_ir(ir_model)
+
+
+def fix_input_shapes_ir(ir_model: ir.Model, input_names: list[str], input_shapes: list[list[int]]) -> None:
+    """Fix the input shapes in an ir.Model.
+
+    :param input_names: The input names to fix.
+    :param input_shapes: The shapes to set for the inputs.
+    """
+    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
+    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
+
+    for name, shape in zip(input_names, input_shapes):
+        _make_input_shape_fixed_ir(ir_model.graph, name, shape)
+
+    # update the output shapes to make them fixed
+    _fix_output_shapes_ir(ir_model)
+
+
 def process_llm_pipeline(
     model: CompositeModelHandler,
     llm_pipeline: list,
diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py
index a019482605..2f73c341ad 100644
--- a/olive/passes/onnx/dynamic_to_fixed_shape.py
+++ b/olive/passes/onnx/dynamic_to_fixed_shape.py
@@ -6,6 +6,7 @@
 import logging
 from typing import Any, Callable
 
+import onnx_ir as ir
 from pydantic import model_validator
 
 from olive.hardware import AcceleratorSpec
@@ -13,10 +14,10 @@
 from olive.model.utils import resolve_onnx_path
 from olive.passes.olive_pass import Pass
 from olive.passes.onnx.common import (
-    fix_dim_params,
-    fix_input_shapes,
+    fix_dim_params_ir,
+    fix_input_shapes_ir,
     get_external_data_config,
-    model_proto_to_olive_model,
+    ir_model_to_olive_model,
 )
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
@@ -72,15 +73,16 @@ def _run_for_config(
         config: type[BasePassConfig],
         output_model_path: str,
     ) -> ONNXModelHandler:
-        onnx_model = model.load_model()
+        ir_model = model.load_ir_model()
+        ir.external_data.load_to_model(ir_model)
         output_model_path = resolve_onnx_path(output_model_path)
 
         if config.dim_param:
-            fix_dim_params(onnx_model, config.dim_param, config.dim_value)
+            fix_dim_params_ir(ir_model, config.dim_param, config.dim_value)
         elif config.input_name:
-            fix_input_shapes(onnx_model, config.input_name, config.input_shape)
+            fix_input_shapes_ir(ir_model, config.input_name, config.input_shape)
 
-        return model_proto_to_olive_model(onnx_model, output_model_path, config)
+        return ir_model_to_olive_model(ir_model, output_model_path, config)
 
 
 def _jointly_validate_configs(cls, values):
diff --git a/olive/passes/onnx/io_datatype_converter.py b/olive/passes/onnx/io_datatype_converter.py
index 3dd7258933..7f05e5a926 100644
--- a/olive/passes/onnx/io_datatype_converter.py
+++ b/olive/passes/onnx/io_datatype_converter.py
@@ -4,17 +4,18 @@
 # --------------------------------------------------------------------------
 import logging
 import re
-from collections import defaultdict
 from pathlib import Path
 from typing import Optional
 
 import onnx
+import onnx_ir as ir
 
+from olive.constants import OpType
 from olive.hardware.accelerator import AcceleratorSpec
 from olive.model import ONNXModelHandler
 from olive.model.utils import resolve_onnx_path
 from olive.passes import Pass
-from olive.passes.onnx.common import get_external_data_config, model_proto_to_olive_model
+from olive.passes.onnx.common import get_external_data_config, ir_model_to_olive_model
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
@@ -51,67 +52,63 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
         config.update(get_external_data_config())
         return config
 
-    def create_io_mapping(self, graph, i_map, o_map):
-        for n in graph.node:
-            for i in n.input:
-                i_map[i].append(n)
-        for n in graph.node:
-            for o in n.output:
-                assert o not in o_map[o]
-                o_map[o] = [n]
-
-    def wrap_inputs(self, graph, i_map, names, source_dtype, target_dtype) -> int:
-        # 1. find source_dtype inputs
-        # 2. rewrite all consumers
-        # 3. insert cast
-        # 4. rewrite graph inputs
-        inputs = [n for n in graph.input if n.type.tensor_type.elem_type == source_dtype]
+    def _wrap_inputs(
+        self, graph: ir.Graph, names: Optional[re.Pattern], source_dtype: ir.DataType, target_dtype: ir.DataType
+    ) -> int:
+        # 1. find source_dtype graph inputs
+        # 2. rewrite all consumers to read from a Cast output
+        # 3. insert Cast that converts the (now target_dtype) input back to source_dtype
+        # 4. rewrite the graph input dtype to target_dtype
         converted_count = 0
-        for i in inputs:
-            if not self._is_name_matched(i.name, names):
+        for graph_input in list(graph.inputs):
+            if graph_input.dtype != source_dtype:
                 continue
-            logger.debug("Converting input %s from %s to %s", i.name, source_dtype, target_dtype)
-            for n in i_map[i.name]:
-                for j, o in enumerate(n.input):
-                    if o == i.name:
-                        n.input[j] = i.name + "_converted"
+            if not self._is_name_matched(graph_input.name, names):
+                continue
+            logger.debug("Converting input %s from %s to %s", graph_input.name, source_dtype, target_dtype)
 
-            cast = onnx.helper.make_node("Cast", inputs=[i.name], outputs=[i.name + "_converted"], to=source_dtype)
+            cast_out = ir.Value(
+                name=graph_input.name + "_converted", shape=graph_input.shape, type=ir.TensorType(source_dtype)
+            )
+            # redirect existing consumers before the Cast node exists so it is not itself redirected
+            graph_input.replace_all_uses_with(cast_out)
 
-            graph.node.insert(0, cast)
-            i.type.tensor_type.elem_type = target_dtype
+            cast_node = ir.node(
+                str(OpType.Cast), inputs=[graph_input], attributes={"to": int(source_dtype)}, outputs=[cast_out]
+            )
+            graph_input.type = ir.TensorType(target_dtype)
+            graph.append(cast_node)
             converted_count += 1
 
         return converted_count
 
-    def wrap_outputs(self, graph, i_map, o_map, names, source_dtype, target_dtype) -> int:
-        # 1. find source dtype outputs
-        # 2. rewrite all providers
-        # 3. append cast
-        # 4. rewrite graph outputs
-        outputs = [n for n in graph.output if n.type.tensor_type.elem_type == source_dtype]
+    def _wrap_outputs(
+        self, graph: ir.Graph, names: Optional[re.Pattern], source_dtype: ir.DataType, target_dtype: ir.DataType
+    ) -> int:
+        # 1. find source_dtype graph outputs
+        # 2. rename the internal tensor (keeping its source_dtype consumers intact)
+        # 3. append a Cast that converts the internal tensor to target_dtype under the original output name
+        # 4. rewrite the graph output slot to the Cast output
         converted_count = 0
-        for o in outputs:
-            if not self._is_name_matched(o.name, names):
+        for idx, graph_output in list(enumerate(graph.outputs)):
+            if graph_output is None or graph_output.dtype != source_dtype:
                 continue
-            logger.debug("Converting output %s from %s to %s", o.name, source_dtype, target_dtype)
-            for n in o_map[o.name]:
-                for j, i_ in enumerate(n.output):
-                    if i_ == o.name:
-                        n.output[j] = o.name + "_converted"
-            for n in i_map[o.name]:
-                for j, i_ in enumerate(n.input):
-                    if i_ == o.name:
-                        n.input[j] = o.name + "_converted"
-
-            cast = onnx.helper.make_node(
-                "Cast",
-                inputs=[o.name + "_converted"],
-                outputs=[o.name],
-                to=target_dtype,
-            )
-            graph.node.append(cast)
-            o.type.tensor_type.elem_type = target_dtype
+            if not self._is_name_matched(graph_output.name, names):
+                continue
+            logger.debug("Converting output %s from %s to %s", graph_output.name, source_dtype, target_dtype)
+
+            original_name = graph_output.name
+            # keep the internal tensor in source_dtype; all its existing consumers follow the rename
+            graph_output.name = original_name + "_converted"
+
+            cast_node = ir.node(str(OpType.Cast), inputs=[graph_output], attributes={"to": int(target_dtype)})
+            cast_out = cast_node.outputs[0]
+            cast_out.name = original_name
+            cast_out.shape = graph_output.shape
+            cast_out.type = ir.TensorType(target_dtype)
+
+            graph.outputs[idx] = cast_out
+            graph.append(cast_node)
             converted_count += 1
 
         return converted_count
@@ -119,13 +116,6 @@ def wrap_outputs(self, graph, i_map, o_map, names, source_dtype, target_dtype) -
     def _is_name_matched(self, name: str, names: Optional[re.Pattern]) -> bool:
         return not names or bool(names.search(name))
 
-    @staticmethod
-    def get_elem_type_from_number(num):
-        for value in vars(onnx.TensorProto).values():
-            if isinstance(value, int) and value == num:
-                return value
-        raise ValueError(f"Invalid elem_type number: {num}")
-
     def _get_available_elem_types(self):
         return onnx.TensorProto.DataType.values()
 
@@ -141,32 +131,20 @@ def _verify_elem_type(self, elem_type):
     def _run_for_config(
         self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: str
     ) -> ONNXModelHandler:
-        from onnxruntime.transformers.onnx_model import OnnxModel
-
         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
 
-        ort_onnx_model = OnnxModel(model.load_model())
-
-        i_map = defaultdict(list)
-        o_map = defaultdict(list)
-
-        self.create_io_mapping(ort_onnx_model.model.graph, i_map, o_map)
-
-        pat = None
-        if config.name_pattern:
-            pat = re.compile(config.name_pattern)
+        self._verify_elem_type(config.source_dtype)
+        self._verify_elem_type(config.target_dtype)
 
-        source_dtype = config.source_dtype
-        target_dtype = config.target_dtype
+        source_dtype = ir.DataType(config.source_dtype)
+        target_dtype = ir.DataType(config.target_dtype)
 
-        self._verify_elem_type(source_dtype)
-        self._verify_elem_type(target_dtype)
+        ir_model = model.load_ir_model()
 
-        source_dtype = self.get_elem_type_from_number(source_dtype)
-        target_dtype = self.get_elem_type_from_number(target_dtype)
+        pat = re.compile(config.name_pattern) if config.name_pattern else None
 
-        wrapped_inputs = self.wrap_inputs(ort_onnx_model.model.graph, i_map, pat, source_dtype, target_dtype)
-        wrapped_outputs = self.wrap_outputs(ort_onnx_model.model.graph, i_map, o_map, pat, source_dtype, target_dtype)
+        wrapped_inputs = self._wrap_inputs(ir_model.graph, pat, source_dtype, target_dtype)
+        wrapped_outputs = self._wrap_outputs(ir_model.graph, pat, source_dtype, target_dtype)
         if wrapped_inputs + wrapped_outputs == 0:
             logger.info("No inputs/outputs found with source_dtype=%s. Skip conversion.", source_dtype)
             return model
@@ -178,4 +156,5 @@ def _run_for_config(
             target_dtype,
         )
 
-        return model_proto_to_olive_model(ort_onnx_model.model, output_model_path, config)
+        ir_model.graph.sort()
+        return ir_model_to_olive_model(ir_model, output_model_path, config)
diff --git a/test/passes/onnx/test_bnb_quantization.py b/test/passes/onnx/test_bnb_quantization.py
index e1ddd4eb3c..059902da08 100644
--- a/test/passes/onnx/test_bnb_quantization.py
+++ b/test/passes/onnx/test_bnb_quantization.py
@@ -94,7 +94,7 @@ def test_validate_precision(pass_config, model_attributes, expected_error, tmp_p
 @pytest.mark.parametrize(("model_generator", "expected_count"), [(get_onnx_matmul_model, 1), (get_onnx_gemm_model, 0)])
 def test__find_matmul_nodes(tmp_path, model_generator, expected_count):
     onnx_model = model_generator(str(tmp_path / "model.onnx"))
-    matmul_nodes = OnnxBnb4Quantization._find_matmul_nodes(onnx_model.load_model().graph)
+    matmul_nodes = OnnxBnb4Quantization._find_matmul_nodes(onnx_model.load_ir_model().graph)
     assert len(matmul_nodes) == expected_count
     if expected_count:
         assert "fc1" in matmul_nodes[0]

From 9f7df525bc5de41dbfedd2bef3778ec2dd434023 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Mon, 6 Jul 2026 14:29:28 -0700
Subject: [PATCH 02/15] Move dynamic_to_fixed_shape IR helpers into the pass
 and use model.graphs()

The IR shape-fixing helpers are only used by DynamicToFixedShape, so move them
out of common.py into the pass module. Replace the custom subgraph iterator with
ir.Model.graphs() for traversal.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 olive/passes/onnx/common.py                 | 160 --------------------
 olive/passes/onnx/dynamic_to_fixed_shape.py | 152 ++++++++++++++++++-
 2 files changed, 148 insertions(+), 164 deletions(-)

diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index c4221f6ccf..b2df6c8705 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -629,166 +629,6 @@ def fix_input_shapes(model_proto: onnx.ModelProto, input_names: list[str], input
     _fix_output_shapes(model_proto)
 
 
-def _ir_shape_is_fixed(value: Optional[ir.Value]) -> bool:
-    """Return True if the value has a shape where every dimension is a fixed positive integer."""
-    if value is None or value.shape is None:
-        return False
-    return all(isinstance(dim, int) and dim > 0 for dim in value.shape)
-
-
-def _iter_graphs_ir(graph: ir.Graph):
-    """Yield the graph and all of its subgraphs recursively (for control-flow nodes)."""
-    yield graph
-    for node in graph:
-        for attr in node.attributes.values():
-            if not isinstance(attr, ir.Attr):
-                continue
-            if attr.type == ir.AttributeType.GRAPH:
-                yield from _iter_graphs_ir(attr.value)
-            elif attr.type == ir.AttributeType.GRAPHS:
-                for subgraph in attr.value:
-                    yield from _iter_graphs_ir(subgraph)
-
-
-def _iter_shaped_values_ir(graph: ir.Graph):
-    """Yield all value objects in a single graph that may carry shape information."""
-    yield from graph.inputs
-    for node in graph:
-        yield from node.outputs
-    yield from graph.outputs
-
-
-def _make_dim_param_fixed_ir(graph: ir.Graph, param_name: str, value: int) -> None:
-    """Replace every occurrence of the symbolic dim ``param_name`` with ``value`` across the graph.
-
-    Mirrors onnxruntime.tools.onnx_model_utils.make_dim_param_fixed but operates on an ir.Graph,
-    including subgraphs.
-    """
-    for subgraph in _iter_graphs_ir(graph):
-        for val in _iter_shaped_values_ir(subgraph):
-            if val is None or val.shape is None:
-                continue
-            dims = list(val.shape)
-            changed = False
-            for idx, dim in enumerate(dims):
-                if isinstance(dim, ir.SymbolicDim) and dim.value == param_name:
-                    dims[idx] = value
-                    changed = True
-            if changed:
-                val.shape = ir.Shape(dims)
-
-
-def _remove_invalid_dim_values_ir(graph: ir.Graph) -> None:
-    """Unset any fixed dim values that are less than 1 (typically -1 placeholders for dynamic dims)."""
-    for subgraph in _iter_graphs_ir(graph):
-        for val in _iter_shaped_values_ir(subgraph):
-            if val is None or val.shape is None:
-                continue
-            dims = list(val.shape)
-            changed = False
-            for idx, dim in enumerate(dims):
-                if isinstance(dim, int) and dim < 1:
-                    dims[idx] = None
-                    changed = True
-            if changed:
-                val.shape = ir.Shape(dims)
-
-
-def _make_input_shape_fixed_ir(graph: ir.Graph, input_name: str, fixed_shape: list[int]) -> None:
-    """Set the shape of the named graph input to ``fixed_shape``.
-
-    Mirrors onnxruntime.tools.onnx_model_utils.make_input_shape_fixed but operates on an ir.Graph.
-    """
-    # remove any invalid dim values first. typically this is a dim_value of -1.
-    _remove_invalid_dim_values_ir(graph)
-
-    for graph_input in graph.inputs:
-        if graph_input.name != input_name:
-            continue
-
-        # graph inputs are required to have a shape to provide the rank
-        if graph_input.shape is None:
-            raise ValueError(f"Input {input_name} does not have a shape")
-
-        dims = list(graph_input.shape)
-        if len(dims) != len(fixed_shape):
-            raise ValueError(f"Rank mismatch. Existing:{len(dims)} Replacement:{len(fixed_shape)}")
-
-        new_dims = list(dims)
-        for idx, dim in enumerate(dims):
-            if isinstance(dim, int):
-                # check any existing fixed dims match
-                if dim != fixed_shape[idx]:
-                    raise ValueError(
-                        f"Can't replace existing fixed size of {dim} with {fixed_shape[idx]} for dimension {idx + 1}"
-                    )
-            elif isinstance(dim, ir.SymbolicDim) and dim.value is not None:
-                # replacing a dim_param so have to do that through the entire graph
-                _make_dim_param_fixed_ir(graph, dim.value, fixed_shape[idx])
-                new_dims[idx] = fixed_shape[idx]
-            else:
-                # replacing an unknown dim
-                new_dims[idx] = fixed_shape[idx]
-
-        graph_input.shape = ir.Shape(new_dims)
-        return
-
-    valid_names = ",".join(i.name for i in graph.inputs if i.name)
-    raise ValueError(f"Input {input_name} was not found in graph inputs. Valid input names are: {valid_names}")
-
-
-def _fix_output_shapes_ir(ir_model: ir.Model) -> None:
-    """Run shape inference on the model and update graph output shapes to make them fixed."""
-    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
-    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
-
-    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
-    model_proto = ir.to_proto(ir_model)
-    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
-    inferred_outputs = {o.name: o for o in inferred_proto.graph.output}
-
-    for output in ir_model.graph.outputs:
-        if output is None or output.name is None or _ir_shape_is_fixed(output):
-            continue
-        new_o = inferred_outputs.get(output.name)
-        if new_o is not None and is_fixed_size_tensor(new_o):
-            output.shape = ir.Shape([dim.dim_value for dim in new_o.type.tensor_type.shape.dim])
-
-
-def fix_dim_params_ir(ir_model: ir.Model, dim_params: list[str], dim_values: list[int]) -> None:
-    """Fix the dimension parameters in an ir.Model.
-
-    :param dim_params: The dimension parameters to fix.
-    :param dim_values: The values to set for the dimension parameters.
-    """
-    dim_params = list(dim_params)
-    dim_values = list(dim_values)
-    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
-    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
-
-    for param, value in zip(dim_params, dim_values):
-        _make_dim_param_fixed_ir(ir_model.graph, param, value)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes_ir(ir_model)
-
-
-def fix_input_shapes_ir(ir_model: ir.Model, input_names: list[str], input_shapes: list[list[int]]) -> None:
-    """Fix the input shapes in an ir.Model.
-
-    :param input_names: The input names to fix.
-    :param input_shapes: The shapes to set for the inputs.
-    """
-    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
-    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
-
-    for name, shape in zip(input_names, input_shapes):
-        _make_input_shape_fixed_ir(ir_model.graph, name, shape)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes_ir(ir_model)
-
-
 def process_llm_pipeline(
     model: CompositeModelHandler,
     llm_pipeline: list,
diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py
index 2f73c341ad..1f49995b1c 100644
--- a/olive/passes/onnx/dynamic_to_fixed_shape.py
+++ b/olive/passes/onnx/dynamic_to_fixed_shape.py
@@ -14,8 +14,6 @@
 from olive.model.utils import resolve_onnx_path
 from olive.passes.olive_pass import Pass
 from olive.passes.onnx.common import (
-    fix_dim_params_ir,
-    fix_input_shapes_ir,
     get_external_data_config,
     ir_model_to_olive_model,
 )
@@ -24,6 +22,152 @@
 logger = logging.getLogger(__name__)
 
 
+def _iter_shaped_values(graph: ir.Graph):
+    """Yield all value objects in a single graph that may carry shape information."""
+    yield from graph.inputs
+    for node in graph:
+        yield from node.outputs
+    yield from graph.outputs
+
+
+def _make_dim_param_fixed(ir_model: ir.Model, param_name: str, value: int) -> None:
+    """Replace every occurrence of the symbolic dim ``param_name`` with ``value`` across the model.
+
+    Mirrors onnxruntime.tools.onnx_model_utils.make_dim_param_fixed but operates on an ir.Model,
+    including subgraphs.
+    """
+    for graph in ir_model.graphs():
+        for val in _iter_shaped_values(graph):
+            if val is None or val.shape is None:
+                continue
+            dims = list(val.shape)
+            changed = False
+            for idx, dim in enumerate(dims):
+                if isinstance(dim, ir.SymbolicDim) and dim.value == param_name:
+                    dims[idx] = value
+                    changed = True
+            if changed:
+                val.shape = ir.Shape(dims)
+
+
+def _remove_invalid_dim_values(ir_model: ir.Model) -> None:
+    """Unset any fixed dim values that are less than 1 (typically -1 placeholders for dynamic dims)."""
+    for graph in ir_model.graphs():
+        for val in _iter_shaped_values(graph):
+            if val is None or val.shape is None:
+                continue
+            dims = list(val.shape)
+            changed = False
+            for idx, dim in enumerate(dims):
+                if isinstance(dim, int) and dim < 1:
+                    dims[idx] = None
+                    changed = True
+            if changed:
+                val.shape = ir.Shape(dims)
+
+
+def _make_input_shape_fixed(ir_model: ir.Model, input_name: str, fixed_shape: list[int]) -> None:
+    """Set the shape of the named graph input to ``fixed_shape``.
+
+    Mirrors onnxruntime.tools.onnx_model_utils.make_input_shape_fixed but operates on an ir.Model.
+    """
+    # remove any invalid dim values first. typically this is a dim_value of -1.
+    _remove_invalid_dim_values(ir_model)
+
+    for graph_input in ir_model.graph.inputs:
+        if graph_input.name != input_name:
+            continue
+
+        # graph inputs are required to have a shape to provide the rank
+        if graph_input.shape is None:
+            raise ValueError(f"Input {input_name} does not have a shape")
+
+        dims = list(graph_input.shape)
+        if len(dims) != len(fixed_shape):
+            raise ValueError(f"Rank mismatch. Existing:{len(dims)} Replacement:{len(fixed_shape)}")
+
+        new_dims = list(dims)
+        for idx, dim in enumerate(dims):
+            if isinstance(dim, int):
+                # check any existing fixed dims match
+                if dim != fixed_shape[idx]:
+                    raise ValueError(
+                        f"Can't replace existing fixed size of {dim} with {fixed_shape[idx]} for dimension {idx + 1}"
+                    )
+            elif isinstance(dim, ir.SymbolicDim) and dim.value is not None:
+                # replacing a dim_param so have to do that through the entire model
+                _make_dim_param_fixed(ir_model, dim.value, fixed_shape[idx])
+                new_dims[idx] = fixed_shape[idx]
+            else:
+                # replacing an unknown dim
+                new_dims[idx] = fixed_shape[idx]
+
+        graph_input.shape = ir.Shape(new_dims)
+        return
+
+    valid_names = ",".join(i.name for i in ir_model.graph.inputs if i.name)
+    raise ValueError(f"Input {input_name} was not found in graph inputs. Valid input names are: {valid_names}")
+
+
+def _fix_output_shapes(ir_model: ir.Model) -> None:
+    """Run shape inference on the model and update graph output shapes to make them fixed."""
+    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
+    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+
+    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
+    model_proto = ir.to_proto(ir_model)
+    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
+    inferred_outputs = {o.name: o for o in inferred_proto.graph.output}
+
+    for output in ir_model.graph.outputs:
+        if output is None or output.name is None or _shape_is_fixed(output):
+            continue
+        new_o = inferred_outputs.get(output.name)
+        if new_o is not None and is_fixed_size_tensor(new_o):
+            output.shape = ir.Shape([dim.dim_value for dim in new_o.type.tensor_type.shape.dim])
+
+
+def _shape_is_fixed(value: ir.Value) -> bool:
+    """Return True if the value has a shape where every dimension is a fixed positive integer."""
+    if value is None or value.shape is None:
+        return False
+    return all(isinstance(dim, int) and dim > 0 for dim in value.shape)
+
+
+def fix_dim_params(ir_model: ir.Model, dim_params: list[str], dim_values: list[int]) -> None:
+    """Fix the dimension parameters in an ir.Model.
+
+    :param dim_params: The dimension parameters to fix.
+    :param dim_values: The values to set for the dimension parameters.
+    """
+    dim_params = list(dim_params)
+    dim_values = list(dim_values)
+    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
+    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
+
+    for param, value in zip(dim_params, dim_values):
+        _make_dim_param_fixed(ir_model, param, value)
+
+    # update the output shapes to make them fixed
+    _fix_output_shapes(ir_model)
+
+
+def fix_input_shapes(ir_model: ir.Model, input_names: list[str], input_shapes: list[list[int]]) -> None:
+    """Fix the input shapes in an ir.Model.
+
+    :param input_names: The input names to fix.
+    :param input_shapes: The shapes to set for the inputs.
+    """
+    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
+    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
+
+    for name, shape in zip(input_names, input_shapes):
+        _make_input_shape_fixed(ir_model, name, shape)
+
+    # update the output shapes to make them fixed
+    _fix_output_shapes(ir_model)
+
+
 class DynamicToFixedShape(Pass):
     """Convert dynamic shape to fixed shape for ONNX model."""
 
@@ -78,9 +222,9 @@ def _run_for_config(
         output_model_path = resolve_onnx_path(output_model_path)
 
         if config.dim_param:
-            fix_dim_params_ir(ir_model, config.dim_param, config.dim_value)
+            fix_dim_params(ir_model, config.dim_param, config.dim_value)
         elif config.input_name:
-            fix_input_shapes_ir(ir_model, config.input_name, config.input_shape)
+            fix_input_shapes(ir_model, config.input_name, config.input_shape)
 
         return ir_model_to_olive_model(ir_model, output_model_path, config)
 

From 432d48c99a9a1e0f73662bf01348b5307b71ffcb Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 6 Jul 2026 21:45:44 +0000
Subject: [PATCH 03/15] Migrate static_llm to onnx-ir, drop onnxruntime shape
 helpers

---
 olive/passes/onnx/common.py     | 65 ++++++----------------------
 olive/passes/onnx/static_llm.py | 75 +++++++++++++++------------------
 2 files changed, 47 insertions(+), 93 deletions(-)

diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index b2df6c8705..4f929c2b29 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -103,6 +103,20 @@ def add_version_metadata_to_model_proto(model: onnx.ModelProto) -> onnx.ModelPro
     return model
 
 
+def add_version_metadata_to_ir_model(model: ir.Model) -> ir.Model:
+    olive_version = None
+    try:
+        import olive
+
+        olive_version = getattr(olive, "__version__", "unknown")
+    except Exception:
+        olive_version = "unknown"
+
+    model.metadata_props["olive_version"] = olive_version
+
+    return model
+
+
 def model_proto_to_file(
     model: onnx.ModelProto,
     output_path: Union[str, Path],
@@ -578,57 +592,6 @@ def model_has_adapters(model_path: Union[str, Path], adapter_type: AdapterType =
     )
 
 
-def _fix_output_shapes(model_proto: onnx.ModelProto):
-    """Run shape inference on the model and update the output shapes to make them fixed."""
-    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
-    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
-
-    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
-    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
-
-    for idx, o in enumerate(model_proto.graph.output):
-        if not is_fixed_size_tensor(o):
-            new_o = inferred_proto.graph.output[idx]
-            if is_fixed_size_tensor(new_o):
-                o.type.tensor_type.shape.CopyFrom(new_o.type.tensor_type.shape)
-
-
-def fix_dim_params(model_proto: onnx.ModelProto, dim_params: list[str], dim_values: list[int]):
-    """Fix the dimension parameters in the model.
-
-    :param dim_params: The dimension parameters to fix.
-    :param dim_values: The values to set for the dimension parameters.
-    """
-    from onnxruntime.tools.onnx_model_utils import make_dim_param_fixed
-
-    assert len(dim_params) == len(dim_values), "dim_params and dim_values must have the same number of elements."
-    assert all(i >= 0 for i in dim_values), "dim_values must be all >= 0"
-
-    for param, value in zip(dim_params, dim_values):
-        make_dim_param_fixed(model_proto.graph, param, value)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes(model_proto)
-
-
-def fix_input_shapes(model_proto: onnx.ModelProto, input_names: list[str], input_shapes: list[list[int]]):
-    """Fix the input shapes in the model.
-
-    :param input_names: The input names to fix.
-    :param input_shapes: The shapes to set for the inputs.
-    """
-    from onnxruntime.tools.onnx_model_utils import make_input_shape_fixed
-
-    assert len(input_names) == len(input_shapes), "input_names and input_shapes must have the same number of elements."
-    assert all(all(i > 0 for i in shape) for shape in input_shapes), "input_shapes must be all > 0"
-
-    for name, shape in zip(input_names, input_shapes):
-        make_input_shape_fixed(model_proto.graph, name, shape)
-
-    # update the output shapes to make them fixed
-    _fix_output_shapes(model_proto)
-
-
 def process_llm_pipeline(
     model: CompositeModelHandler,
     llm_pipeline: list,
diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 5b76f8967f..2c61deaf14 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -5,7 +5,6 @@
 import logging
 from pathlib import Path
 
-import onnx
 import onnx_ir as ir
 
 from olive.hardware import Device
@@ -14,12 +13,12 @@
 from olive.model import CompositeModelHandler, ONNXModelHandler
 from olive.passes import Pass
 from olive.passes.onnx.common import (
-    add_version_metadata_to_model_proto,
-    fix_dim_params,
+    add_version_metadata_to_ir_model,
     process_llm_pipeline,
     resave_model,
     update_llm_pipeline_genai_config_gpu,
 )
+from olive.passes.onnx.dynamic_to_fixed_shape import fix_dim_params
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
@@ -32,12 +31,12 @@ def _ir_io_shape(value: ir.Value) -> list:
     return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
 
 
-def _proto_io_shape(model_proto: onnx.ModelProto, name: str) -> list:
-    """Return the shape of a graph input/output as a list of ints and symbolic dim names."""
-    for value_info in list(model_proto.graph.input) + list(model_proto.graph.output):
-        if value_info.name == name:
-            return [dim.dim_param if dim.dim_param else dim.dim_value for dim in value_info.type.tensor_type.shape.dim]
-    return None
+def _get_ir_input(ir_model: ir.Model, name: str) -> ir.Value:
+    """Return the named graph input value of the model."""
+    for graph_input in ir_model.graph.inputs:
+        if graph_input.name == name:
+            return graph_input
+    raise ValueError(f"Input {name} was not found in graph inputs.")
 
 
 class StaticLLM(Pass):
@@ -99,15 +98,13 @@ def _run_generic(self, model: CompositeModelHandler, config: type[BasePassConfig
         )
 
         # only gqa models are supported for now
-        transformer_model = ir.from_proto(onnx.load(model_components[1].model_path, load_external_data=False))
+        transformer_model = ir.load(model_components[1].model_path)
         assert any(node.op_type == "GroupQueryAttention" for node in transformer_model.graph.all_nodes()), (
             "Only GQA models are supported for now."
         )
         # get dimension params from embeddings model
-        embedding_model = ir.from_proto(onnx.load(model_components[0].model_path, load_external_data=False))
-        input_ids = embedding_model.graph.inputs[
-            [value.name for value in embedding_model.graph.inputs].index("input_ids")
-        ]
+        embedding_model = ir.load(model_components[0].model_path)
+        input_ids = _get_ir_input(embedding_model, "input_ids")
         batch_size, sequence_length = _ir_io_shape(input_ids)
         assert isinstance(batch_size, str), "Batch size must be a symbolic dimension"
         assert isinstance(sequence_length, str), "Sequence length must be a symbolic dimension"
@@ -126,7 +123,7 @@ def _run_generic(self, model: CompositeModelHandler, config: type[BasePassConfig
         # update the param mapping with the new shapes from the embeddings model
         for param_mapping in param_mapping_dict.values():
             self.fix_shape(
-                onnx.load(model_components[0].model_path, load_external_data=False),
+                ir.load(model_components[0].model_path),
                 param_mapping,
             )
 
@@ -148,14 +145,15 @@ def process_context_iterator(component_models, llm_pipeline, output_dir):
                 for key, param_mapping in param_mapping_dict.items():
                     new_component_name = f"{key}{suffix}"
 
-                    component_proto = onnx.load(intermediate_model_path, load_external_data=False)
-                    self.fix_shape(component_proto, param_mapping)
+                    # load lazily so the fixed-shape models share the intermediate external data file
+                    component_ir = ir.load(intermediate_model_path)
+                    self.fix_shape(component_ir, param_mapping)
 
                     # save the model with fixed shapes
                     component_model_path = output_dir / f"{new_component_name}.onnx"
                     # Add olive version to metadata
-                    add_version_metadata_to_model_proto(component_proto)
-                    onnx.save_model(component_proto, component_model_path)
+                    add_version_metadata_to_ir_model(component_ir)
+                    ir.save(component_ir, component_model_path)
                     new_groups[key][new_component_name] = ONNXModelHandler(
                         model_path=output_dir, onnx_file_name=component_model_path.name
                     )
@@ -200,34 +198,28 @@ def process_context_iterator(component_models, llm_pipeline, output_dir):
 
     def _run_qnn_gpu(self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: Path):
         output_model_dir = Path(output_model_path).with_suffix("")
-        model_path = Path(model.model_path)
 
         # --- Step 1: Load model (handle both single and external data) ---
         try:
-            model_proto = onnx.load(model_path, load_external_data=True)
+            ir_model = model.load_ir_model()
+            ir.external_data.load_to_model(ir_model)
         except Exception as e:
             raise RuntimeError(f"Failed to load ONNX model: {e}") from e
 
         # --- Step 2: Fix symbolic dimensions ---
-        batch_size, sequence_length = _proto_io_shape(model_proto, "input_ids")
+        batch_size, sequence_length = _ir_io_shape(_get_ir_input(ir_model, "input_ids"))
         if not (isinstance(batch_size, str) and isinstance(sequence_length, str)):
             raise ValueError("Input dimensions must be symbolic before static shape fixing.")
 
         param_mapping = {batch_size: config.batch_size, sequence_length: config.context_length}
-        self.fix_shape(model_proto, param_mapping)
+        self.fix_shape(ir_model, param_mapping)
 
         # --- Step 3: Save model as external-data format ---
-        output_model_file = Path(output_model_dir) / "model.onnx"
-        external_data_file = Path(output_model_dir) / "model.onnx.data"
-
-        onnx.save(
-            model_proto,
-            str(output_model_file),
-            save_as_external_data=True,
-            all_tensors_to_one_file=True,
-            location=external_data_file.name,
-            convert_attribute=False,
-        )
+        output_model_dir.mkdir(parents=True, exist_ok=True)
+        output_model_file = output_model_dir / "model.onnx"
+        external_data_file = output_model_dir / "model.onnx.data"
+
+        ir.save(ir_model, output_model_file, external_data=external_data_file.name)
 
         decoder_config_extra = {
             "inputs": {
@@ -253,23 +245,22 @@ def _run_qnn_gpu(self, model: ONNXModelHandler, config: type[BasePassConfig], ou
         )
 
     @staticmethod
-    def fix_shape(model_proto: onnx.ModelProto, param_mapping: dict[str, int]):
+    def fix_shape(ir_model: ir.Model, param_mapping: dict[str, int]):
         """Fix the shape of the model based on the param mapping.
 
-        :param model_path: Path to the model.
+        :param ir_model: The ONNX IR model to fix in place.
         :param param_mapping: Mapping from params to fixed values. This gets updated with the output shapes of the new
             model.
         """
-        original_shapes = {}
-        for output in model_proto.graph.output:
-            original_shapes[output.name] = _proto_io_shape(model_proto, output.name)
+        original_shapes = {output.name: _ir_io_shape(output) for output in ir_model.graph.outputs}
 
         # fix dim params
-        fix_dim_params(model_proto, param_mapping.keys(), param_mapping.values())
+        fix_dim_params(ir_model, list(param_mapping.keys()), list(param_mapping.values()))
 
         # update the param mapping with the new shapes
-        for output_name, original_shape in original_shapes.items():
-            new_shape = _proto_io_shape(model_proto, output_name)
+        for output in ir_model.graph.outputs:
+            original_shape = original_shapes[output.name]
+            new_shape = _ir_io_shape(output)
 
             for old_dim, new_dim in zip(original_shape, new_shape):
                 if isinstance(old_dim, str) and isinstance(new_dim, int):

From 60eeb884e2e73f58bc4fcecb4b14b60867eb7eba Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 6 Jul 2026 21:48:09 +0000
Subject: [PATCH 04/15] Address review: simplify version helper and clarify
 external data load

---
 olive/passes/onnx/common.py     | 1 -
 olive/passes/onnx/static_llm.py | 2 ++
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/olive/passes/onnx/common.py b/olive/passes/onnx/common.py
index 4f929c2b29..9be91f9a2f 100644
--- a/olive/passes/onnx/common.py
+++ b/olive/passes/onnx/common.py
@@ -104,7 +104,6 @@ def add_version_metadata_to_model_proto(model: onnx.ModelProto) -> onnx.ModelPro
 
 
 def add_version_metadata_to_ir_model(model: ir.Model) -> ir.Model:
-    olive_version = None
     try:
         import olive
 
diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 2c61deaf14..06f1a86cec 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -202,6 +202,8 @@ def _run_qnn_gpu(self, model: ONNXModelHandler, config: type[BasePassConfig], ou
         # --- Step 1: Load model (handle both single and external data) ---
         try:
             ir_model = model.load_ir_model()
+            # load_ir_model() references external data lazily; materialize it so the model can be
+            # re-saved into a fresh external data file under the output directory
             ir.external_data.load_to_model(ir_model)
         except Exception as e:
             raise RuntimeError(f"Failed to load ONNX model: {e}") from e

From 8ad62fcd1c3d700666917ff1ffe45ada11ce13b5 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Mon, 6 Jul 2026 22:46:29 +0000
Subject: [PATCH 05/15] Use onnx-shape-inference package instead of onnxruntime
 symbolic_shape_infer

---
 olive/passes/onnx/dynamic_to_fixed_shape.py | 26 ++++-----------------
 requirements.txt                            |  1 +
 2 files changed, 6 insertions(+), 21 deletions(-)

diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py
index 1f49995b1c..9bae6ffbbd 100644
--- a/olive/passes/onnx/dynamic_to_fixed_shape.py
+++ b/olive/passes/onnx/dynamic_to_fixed_shape.py
@@ -110,28 +110,12 @@ def _make_input_shape_fixed(ir_model: ir.Model, input_name: str, fixed_shape: li
 
 
 def _fix_output_shapes(ir_model: ir.Model) -> None:
-    """Run shape inference on the model and update graph output shapes to make them fixed."""
-    from onnxruntime.tools.onnx_model_utils import is_fixed_size_tensor
-    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
+    """Run symbolic shape inference to propagate fixed shapes to the graph outputs."""
+    from onnx_shape_inference import infer_symbolic_shapes
 
-    # use the onnxruntime shape inference tool since it can handle large models as well as contrib ops
-    model_proto = ir.to_proto(ir_model)
-    inferred_proto = SymbolicShapeInference.infer_shapes(model_proto, auto_merge=True, guess_output_rank=True)
-    inferred_outputs = {o.name: o for o in inferred_proto.graph.output}
-
-    for output in ir_model.graph.outputs:
-        if output is None or output.name is None or _shape_is_fixed(output):
-            continue
-        new_o = inferred_outputs.get(output.name)
-        if new_o is not None and is_fixed_size_tensor(new_o):
-            output.shape = ir.Shape([dim.dim_value for dim in new_o.type.tensor_type.shape.dim])
-
-
-def _shape_is_fixed(value: ir.Value) -> bool:
-    """Return True if the value has a shape where every dimension is a fixed positive integer."""
-    if value is None or value.shape is None:
-        return False
-    return all(isinstance(dim, int) and dim > 0 for dim in value.shape)
+    # infer_symbolic_shapes operates directly on the ir.Model (no proto round-trip) and refines
+    # existing shapes in place, propagating the now-fixed dimensions to the model outputs.
+    infer_symbolic_shapes(ir_model)
 
 
 def fix_dim_params(ir_model: ir.Model, dim_params: list[str], dim_values: list[int]) -> None:
diff --git a/requirements.txt b/requirements.txt
index 1032c72f97..106a2115e3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
 hf-xet
 numpy
 onnx
+onnx-shape-inference>=0.2.0
 onnx_ir>=0.1.2
 onnxscript>=0.5.3
 opentelemetry-sdk>=1.39.1

From 7fb78db453409ec4a28719d79fea8beee714eaee Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 7 Jul 2026 02:38:46 +0000
Subject: [PATCH 06/15] Remove unnecessary load_to_model in DynamicToFixedShape

---
 olive/passes/onnx/dynamic_to_fixed_shape.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/olive/passes/onnx/dynamic_to_fixed_shape.py b/olive/passes/onnx/dynamic_to_fixed_shape.py
index 9bae6ffbbd..2debf60960 100644
--- a/olive/passes/onnx/dynamic_to_fixed_shape.py
+++ b/olive/passes/onnx/dynamic_to_fixed_shape.py
@@ -202,7 +202,6 @@ def _run_for_config(
         output_model_path: str,
     ) -> ONNXModelHandler:
         ir_model = model.load_ir_model()
-        ir.external_data.load_to_model(ir_model)
         output_model_path = resolve_onnx_path(output_model_path)
 
         if config.dim_param:

From 1655286012eacf0fb270170139636da6dc0b22c6 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 15:46:01 +0000
Subject: [PATCH 07/15] Fix StaticLLM fix_shape for untyped outputs

---
 olive/passes/onnx/static_llm.py     |  2 ++
 test/passes/onnx/test_static_llm.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 06f1a86cec..52b3f98a5a 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -263,6 +263,8 @@ def fix_shape(ir_model: ir.Model, param_mapping: dict[str, int]):
         for output in ir_model.graph.outputs:
             original_shape = original_shapes[output.name]
             new_shape = _ir_io_shape(output)
+            if original_shape is None or new_shape is None:
+                continue
 
             for old_dim, new_dim in zip(original_shape, new_shape):
                 if isinstance(old_dim, str) and isinstance(new_dim, int):
diff --git a/test/passes/onnx/test_static_llm.py b/test/passes/onnx/test_static_llm.py
index d111aaab04..095f92f3ac 100644
--- a/test/passes/onnx/test_static_llm.py
+++ b/test/passes/onnx/test_static_llm.py
@@ -4,6 +4,9 @@
 # --------------------------------------------------------------------------
 import json
 
+import onnx
+import onnx_ir as ir
+
 from olive.model import CompositeModelHandler, ONNXModelHandler
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx.static_llm import StaticLLM
@@ -64,3 +67,28 @@ def test_static_llm(tmp_path):
     assert set(genai_config["model"]["decoder"]["pipeline"][0].keys()) == set(output_model.model_component_names)
     assert not genai_config["model"]["decoder"]["pipeline"][0]["context_0"]["run_on_token_gen"]
     assert not genai_config["model"]["decoder"]["pipeline"][0]["iterator_0"]["run_on_prompt"]
+
+
+def test_static_llm_fix_shape_handles_outputs_without_shape_metadata(tmp_path):
+    model_path = tmp_path / "model.onnx"
+    model = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [onnx.helper.make_node("Identity", ["input_ids"], ["output"])],
+            "test_graph",
+            [
+                onnx.helper.make_tensor_value_info(
+                    "input_ids", onnx.TensorProto.FLOAT, ["batch_size", "sequence_length"]
+                )
+            ],
+            [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, None)],
+        ),
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+    )
+    onnx.save(model, model_path)
+
+    param_mapping = {"batch_size": 1, "sequence_length": 64}
+    ir_model = ir.load(model_path)
+    assert ir_model.graph.outputs[0].shape is None
+    StaticLLM.fix_shape(ir_model, param_mapping)
+
+    assert param_mapping == {"batch_size": 1, "sequence_length": 64}

From 5070b275f839fde9de91975d66fd4b0fcbb6cc1c Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 16:42:12 +0000
Subject: [PATCH 08/15] Fix StaticLLM output shape fallback for compose
 compatibility

---
 olive/passes/onnx/static_llm.py     |  8 +++++++
 test/passes/onnx/test_static_llm.py | 34 +++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 52b3f98a5a..7dbb224eca 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -263,6 +263,14 @@ def fix_shape(ir_model: ir.Model, param_mapping: dict[str, int]):
         for output in ir_model.graph.outputs:
             original_shape = original_shapes[output.name]
             new_shape = _ir_io_shape(output)
+            if original_shape is not None and new_shape is None:
+                # keep output shapes stable even if symbolic shape inference cannot infer this output.
+                # this preserves inter-model interface metadata used by compose.
+                fallback_shape = [
+                    param_mapping.get(dim, dim) if isinstance(dim, str) else dim for dim in original_shape
+                ]
+                output.shape = ir.Shape(fallback_shape)
+                new_shape = fallback_shape
             if original_shape is None or new_shape is None:
                 continue
 
diff --git a/test/passes/onnx/test_static_llm.py b/test/passes/onnx/test_static_llm.py
index 095f92f3ac..a6cad73cc2 100644
--- a/test/passes/onnx/test_static_llm.py
+++ b/test/passes/onnx/test_static_llm.py
@@ -9,6 +9,7 @@
 
 from olive.model import CompositeModelHandler, ONNXModelHandler
 from olive.passes.olive_pass import create_pass_from_dict
+from olive.passes.onnx import static_llm as static_llm_module
 from olive.passes.onnx.static_llm import StaticLLM
 from test.utils import make_local_tiny_llama
 
@@ -92,3 +93,36 @@ def test_static_llm_fix_shape_handles_outputs_without_shape_metadata(tmp_path):
     StaticLLM.fix_shape(ir_model, param_mapping)
 
     assert param_mapping == {"batch_size": 1, "sequence_length": 64}
+
+
+def test_static_llm_fix_shape_restores_output_shape_when_shape_inference_drops_it(tmp_path, monkeypatch):
+    model_path = tmp_path / "model.onnx"
+    model = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [onnx.helper.make_node("Identity", ["input_ids"], ["output"])],
+            "test_graph",
+            [
+                onnx.helper.make_tensor_value_info(
+                    "input_ids", onnx.TensorProto.FLOAT, ["batch_size", "sequence_length"]
+                )
+            ],
+            [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, ["batch_size", "sequence_length"])],
+        ),
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+    )
+    onnx.save(model, model_path)
+
+    def fake_fix_dim_params(ir_model, dim_params, dim_values):
+        for graph_input in ir_model.graph.inputs:
+            if graph_input.name == "input_ids":
+                graph_input.shape = ir.Shape(dim_values)
+        ir_model.graph.outputs[0].shape = None
+
+    monkeypatch.setattr(static_llm_module, "fix_dim_params", fake_fix_dim_params)
+
+    param_mapping = {"batch_size": 1, "sequence_length": 64}
+    ir_model = ir.load(model_path)
+    StaticLLM.fix_shape(ir_model, param_mapping)
+
+    assert list(ir_model.graph.outputs[0].shape) == [1, 64]
+    assert param_mapping == {"batch_size": 1, "sequence_length": 64}

From c8f623ee3327583e2fd0369992e4daabd5157e04 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:21:10 +0000
Subject: [PATCH 09/15] Handle compose shape metadata gaps across split model
 boundaries

---
 olive/passes/onnx/compose.py     | 30 +++++++++++++++++++++++++++++-
 test/passes/onnx/test_compose.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/olive/passes/onnx/compose.py b/olive/passes/onnx/compose.py
index d57ce24423..7741dec2b4 100644
--- a/olive/passes/onnx/compose.py
+++ b/olive/passes/onnx/compose.py
@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import logging
+from numbers import Integral
 from pathlib import Path
 from typing import Optional, Union
 
@@ -117,6 +118,33 @@ def shape_list(value: ir.Value):
                 return None
             return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
 
+        def merge_value_shapes(existing: ir.Value, new_value: ir.Value, name: str):
+            existing_shape = shape_list(existing)
+            new_shape = shape_list(new_value)
+
+            if existing_shape is None:
+                existing.shape = new_value.shape
+                return
+            if new_shape is None:
+                return
+
+            assert len(existing_shape) == len(new_shape), f"Input rank mismatch: {name}"
+
+            merged_shape = []
+            for existing_dim, new_dim in zip(existing_shape, new_shape):
+                if isinstance(existing_dim, Integral) and isinstance(new_dim, Integral):
+                    assert existing_dim == new_dim, f"Input shape mismatch: {name}"
+                    merged_shape.append(existing_dim)
+                elif isinstance(existing_dim, Integral):
+                    merged_shape.append(existing_dim)
+                elif isinstance(new_dim, Integral) or existing_dim is None:
+                    merged_shape.append(new_dim)
+                else:
+                    merged_shape.append(existing_dim)
+
+            if merged_shape != existing_shape:
+                existing.shape = ir.Shape(merged_shape)
+
         ir_models = []
         for path in onnx_model_paths:
             ir_model = ir.load(path)
@@ -170,7 +198,7 @@ def get_value(name: str) -> ir.Value:
                 if name in composed_input_names or name in produced_names:
                     # already a graph input or an internal connection from a previous model
                     existing = composed_values[name]
-                    assert shape_list(inp) == shape_list(existing), f"Input shape mismatch: {name}"
+                    merge_value_shapes(existing, inp, name)
                     assert inp.dtype == existing.dtype, f"Input dtype mismatch: {name}"
                     continue
 
diff --git a/test/passes/onnx/test_compose.py b/test/passes/onnx/test_compose.py
index 43edba9bac..92db221cce 100644
--- a/test/passes/onnx/test_compose.py
+++ b/test/passes/onnx/test_compose.py
@@ -5,6 +5,7 @@
 import json
 from pathlib import Path
 
+import onnx
 import pytest
 
 from olive.model import CompositeModelHandler, ONNXModelHandler
@@ -123,3 +124,34 @@ def test_compose_onnx_models_llm_pipeline(tmp_path):
         genai_config = json.load(f)
     assert genai_config["model"]["decoder"]["pipeline"][0]["context"]["session_options"] == session_options
     assert genai_config["model"]["decoder"]["pipeline"][0]["iterator"]["session_options"] == session_options
+
+
+def test_compose_onnx_models_merges_missing_input_shape_metadata(tmp_path):
+    model_1_path = tmp_path / "model_1.onnx"
+    model_2_path = tmp_path / "model_2.onnx"
+
+    model_1 = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [onnx.helper.make_node("Identity", ["x"], ["y"])],
+            "model_1",
+            [onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 2])],
+            [onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, None)],
+        ),
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+    )
+    onnx.save(model_1, model_1_path)
+
+    model_2 = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [onnx.helper.make_node("Identity", ["y"], ["z"])],
+            "model_2",
+            [onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1, 2])],
+            [onnx.helper.make_tensor_value_info("z", onnx.TensorProto.FLOAT, [1, 2])],
+        ),
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+    )
+    onnx.save(model_2, model_2_path)
+
+    output_model = ComposeOnnxModels._get_composed_model([model_1_path, model_2_path], tmp_path / "output.onnx", {})
+
+    assert isinstance(output_model, ONNXModelHandler)

From 1df9cf3d556b0a86d28b216e6ced9efc8b932573 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:23:16 +0000
Subject: [PATCH 10/15] Merge partial shape metadata when composing linked ONNX
 values

---
 olive/passes/onnx/compose.py     | 8 +++++---
 test/passes/onnx/test_compose.py | 2 +-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/olive/passes/onnx/compose.py b/olive/passes/onnx/compose.py
index 7741dec2b4..57f8dc06e8 100644
--- a/olive/passes/onnx/compose.py
+++ b/olive/passes/onnx/compose.py
@@ -118,12 +118,12 @@ def shape_list(value: ir.Value):
                 return None
             return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
 
-        def merge_value_shapes(existing: ir.Value, new_value: ir.Value, name: str):
+        def merge_value_shapes(existing: ir.Value, consumer_input: ir.Value, name: str):
             existing_shape = shape_list(existing)
-            new_shape = shape_list(new_value)
+            new_shape = shape_list(consumer_input)
 
             if existing_shape is None:
-                existing.shape = new_value.shape
+                existing.shape = consumer_input.shape
                 return
             if new_shape is None:
                 return
@@ -137,6 +137,8 @@ def merge_value_shapes(existing: ir.Value, new_value: ir.Value, name: str):
                     merged_shape.append(existing_dim)
                 elif isinstance(existing_dim, Integral):
                     merged_shape.append(existing_dim)
+                # Prefer known dimension metadata from the consumer side when available, and
+                # fill missing producer metadata (`None`) from the consumer metadata.
                 elif isinstance(new_dim, Integral) or existing_dim is None:
                     merged_shape.append(new_dim)
                 else:
diff --git a/test/passes/onnx/test_compose.py b/test/passes/onnx/test_compose.py
index 92db221cce..9087e01da4 100644
--- a/test/passes/onnx/test_compose.py
+++ b/test/passes/onnx/test_compose.py
@@ -126,7 +126,7 @@ def test_compose_onnx_models_llm_pipeline(tmp_path):
     assert genai_config["model"]["decoder"]["pipeline"][0]["iterator"]["session_options"] == session_options
 
 
-def test_compose_onnx_models_merges_missing_input_shape_metadata(tmp_path):
+def test_compose_onnx_models_merges_partial_shape_metadata(tmp_path):
     model_1_path = tmp_path / "model_1.onnx"
     model_2_path = tmp_path / "model_2.onnx"
 

From 9ed6851ec6f7ac2adc1791add6dcfdc5dbbf802f Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:25:19 +0000
Subject: [PATCH 11/15] Assert merged intermediate shape metadata in compose
 regression test

---
 test/passes/onnx/test_compose.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/passes/onnx/test_compose.py b/test/passes/onnx/test_compose.py
index 9087e01da4..beeb49cef6 100644
--- a/test/passes/onnx/test_compose.py
+++ b/test/passes/onnx/test_compose.py
@@ -155,3 +155,6 @@ def test_compose_onnx_models_merges_partial_shape_metadata(tmp_path):
     output_model = ComposeOnnxModels._get_composed_model([model_1_path, model_2_path], tmp_path / "output.onnx", {})
 
     assert isinstance(output_model, ONNXModelHandler)
+    composed_model = onnx.load(output_model.model_path)
+    y_info = next(value_info for value_info in composed_model.graph.value_info if value_info.name == "y")
+    assert [dim.dim_value for dim in y_info.type.tensor_type.shape.dim] == [1, 2]

From 0f7a41309e0999c977c8a1fadce4195507775b59 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:33:27 +0000
Subject: [PATCH 12/15] Use Value.merge_shapes in compose and relax
 dynamic-shape error assertion

---
 olive/passes/onnx/compose.py                  | 37 ++-----------------
 .../passes/onnx/test_session_params_tuning.py |  7 +++-
 2 files changed, 10 insertions(+), 34 deletions(-)

diff --git a/olive/passes/onnx/compose.py b/olive/passes/onnx/compose.py
index 57f8dc06e8..3e2251572a 100644
--- a/olive/passes/onnx/compose.py
+++ b/olive/passes/onnx/compose.py
@@ -3,7 +3,6 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import logging
-from numbers import Integral
 from pathlib import Path
 from typing import Optional, Union
 
@@ -113,39 +112,11 @@ def _get_composed_model(
         :return: Composed ONNX model.
         """
 
-        def shape_list(value: ir.Value):
-            if value.shape is None:
-                return None
-            return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
-
         def merge_value_shapes(existing: ir.Value, consumer_input: ir.Value, name: str):
-            existing_shape = shape_list(existing)
-            new_shape = shape_list(consumer_input)
-
-            if existing_shape is None:
-                existing.shape = consumer_input.shape
-                return
-            if new_shape is None:
-                return
-
-            assert len(existing_shape) == len(new_shape), f"Input rank mismatch: {name}"
-
-            merged_shape = []
-            for existing_dim, new_dim in zip(existing_shape, new_shape):
-                if isinstance(existing_dim, Integral) and isinstance(new_dim, Integral):
-                    assert existing_dim == new_dim, f"Input shape mismatch: {name}"
-                    merged_shape.append(existing_dim)
-                elif isinstance(existing_dim, Integral):
-                    merged_shape.append(existing_dim)
-                # Prefer known dimension metadata from the consumer side when available, and
-                # fill missing producer metadata (`None`) from the consumer metadata.
-                elif isinstance(new_dim, Integral) or existing_dim is None:
-                    merged_shape.append(new_dim)
-                else:
-                    merged_shape.append(existing_dim)
-
-            if merged_shape != existing_shape:
-                existing.shape = ir.Shape(merged_shape)
+            try:
+                existing.merge_shapes(consumer_input.shape)
+            except ValueError as e:
+                raise AssertionError(f"Input shape mismatch: {name}") from e
 
         ir_models = []
         for path in onnx_model_paths:
diff --git a/test/passes/onnx/test_session_params_tuning.py b/test/passes/onnx/test_session_params_tuning.py
index 4eed0653a9..bf230f4118 100644
--- a/test/passes/onnx/test_session_params_tuning.py
+++ b/test/passes/onnx/test_session_params_tuning.py
@@ -171,4 +171,9 @@ def test_ort_session_params_tuning_pass_with_dynamic_shapes(mock_get_io_config,
     with pytest.raises(TypeError) as e:
         # execute
         p.run(input_model, output_folder)
-    assert "ones() received an invalid combination of arguments" in str(e.value)
+    error = str(e.value)
+    assert "ones()" in error
+    assert (
+        "received an invalid combination of arguments" in error
+        or "must be tuple of ints, but found element of type str" in error
+    )

From 6306163ad1596f6fdb16d168467f13910e3fa132 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Thu, 9 Jul 2026 23:50:57 +0000
Subject: [PATCH 13/15] Add StaticLLM test asserting context/iterator share
 external data

Verify that each split's context_i and iterator_i models reference the
same shared transformer_i.onnx.data external file (mmapped once for both
prompt and token-gen) instead of writing per-component duplicates. Guards
against a regression where routing the pipeline save through a helper that
rewrites the external-data path would fork the shared weights into
context_i.onnx.data / iterator_i.onnx.data.

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 test/passes/onnx/test_static_llm.py | 37 +++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/test/passes/onnx/test_static_llm.py b/test/passes/onnx/test_static_llm.py
index a6cad73cc2..a407a3e753 100644
--- a/test/passes/onnx/test_static_llm.py
+++ b/test/passes/onnx/test_static_llm.py
@@ -69,6 +69,43 @@ def test_static_llm(tmp_path):
     assert not genai_config["model"]["decoder"]["pipeline"][0]["context_0"]["run_on_token_gen"]
     assert not genai_config["model"]["decoder"]["pipeline"][0]["iterator_0"]["run_on_prompt"]
 
+    # The context and iterator components for a given split share a single external
+    # data file (the intermediate transformer_{i}.onnx.data) rather than each
+    # writing its own duplicate copy. This is what lets the runtime memory-map the
+    # weights once for both the prompt (context) and token-gen (iterator) models.
+    def _external_data_locations(onnx_path):
+        # Return the set of external data file names referenced by the model's
+        # initializers, without materializing the tensor data.
+        model_proto = onnx.load(onnx_path, load_external_data=False)
+        locations = set()
+        for initializer in model_proto.graph.initializer:
+            if initializer.data_location == onnx.TensorProto.EXTERNAL:
+                for entry in initializer.external_data:
+                    if entry.key == "location":
+                        locations.add(entry.value)
+        return locations
+
+    output_files = {f.name for f in output_model_path.iterdir()}
+    for idx in range(2):  # two split transformer components -> context_i / iterator_i
+        context_locations = _external_data_locations(output_model_path / f"context_{idx}.onnx")
+        iterator_locations = _external_data_locations(output_model_path / f"iterator_{idx}.onnx")
+
+        # Both components must reference exactly the same shared external data file(s)...
+        assert context_locations, f"context_{idx}.onnx should reference external data"
+        assert context_locations == iterator_locations, (
+            f"context_{idx} and iterator_{idx} must share external data, got "
+            f"{context_locations} vs {iterator_locations}"
+        )
+        # ...which is the intermediate transformer_{idx}.onnx.data that survives on disk.
+        assert context_locations == {f"transformer_{idx}.onnx.data"}, (
+            f"expected shared transformer_{idx}.onnx.data, got {context_locations}"
+        )
+        assert f"transformer_{idx}.onnx.data" in output_files
+
+        # No per-component duplicate external data files should have been created.
+        assert f"context_{idx}.onnx.data" not in output_files
+        assert f"iterator_{idx}.onnx.data" not in output_files
+
 
 def test_static_llm_fix_shape_handles_outputs_without_shape_metadata(tmp_path):
     model_path = tmp_path / "model.onnx"

From aed279c29cc69f6dd643b72829614445170198af Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Thu, 9 Jul 2026 23:55:53 +0000
Subject: [PATCH 14/15] Drop unnecessary load_to_model in bnb and static_llm
 (qnn-gpu)

ir.external_data.load_to_model eagerly materializes every external tensor
into memory. Neither of these two paths needs it:

* OnnxBnb4Quantization: reads each MatMul weight lazily via const_value.numpy()
  (ExternalTensor reads on demand) and writes new quantized initializers; the
  remaining lazy tensors are read from the still-present source file when the
  output is saved to a different path.
* StaticLLM._run_qnn_gpu: only edits shape metadata (fix_shape) and saves to a
  fresh model.onnx.data under a distinct output directory, so the lazy
  references stay valid at save time.

Removing the eager materialization lowers peak memory for large models with no
behavior change. Mirrors the earlier DynamicToFixedShape removal. Existing bnb
and static_llm tests pass (19).

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 olive/passes/onnx/bnb_quantization.py | 1 -
 olive/passes/onnx/static_llm.py       | 6 +++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/olive/passes/onnx/bnb_quantization.py b/olive/passes/onnx/bnb_quantization.py
index c5db8c9e98..17da628b2a 100644
--- a/olive/passes/onnx/bnb_quantization.py
+++ b/olive/passes/onnx/bnb_quantization.py
@@ -101,7 +101,6 @@ def _run_for_config(
 
         # load the model
         ir_model = model.load_ir_model()
-        ir.external_data.load_to_model(ir_model)
         ir_model.graph.opset_imports[MSFT_DOMAIN] = 1
 
         # get nodes to exclude from quantization
diff --git a/olive/passes/onnx/static_llm.py b/olive/passes/onnx/static_llm.py
index 7dbb224eca..da1f592056 100644
--- a/olive/passes/onnx/static_llm.py
+++ b/olive/passes/onnx/static_llm.py
@@ -201,10 +201,10 @@ def _run_qnn_gpu(self, model: ONNXModelHandler, config: type[BasePassConfig], ou
 
         # --- Step 1: Load model (handle both single and external data) ---
         try:
+            # External data is referenced lazily. We only edit shape metadata below (no tensor
+            # reads/writes) and save to a fresh file under a distinct output directory, so the
+            # lazy references remain valid at save time — no need to materialize the weights.
             ir_model = model.load_ir_model()
-            # load_ir_model() references external data lazily; materialize it so the model can be
-            # re-saved into a fresh external data file under the output directory
-            ir.external_data.load_to_model(ir_model)
         except Exception as e:
             raise RuntimeError(f"Failed to load ONNX model: {e}") from e
 

From 27655a4d8383b5f8930ee4b4c415227dbd32435c Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 10 Jul 2026 04:05:39 +0000
Subject: [PATCH 15/15] Drop unnecessary load_to_model in rtn/kquant/hqq
 quantizers

Same rationale as bnb: each pass reads MatMul weights lazily via
const_value.numpy() and writes new quantized initializers, saving to a
distinct output path. The remaining lazy tensors are read from the still-
present source at save time, so eagerly materializing the whole model up
front is unnecessary and only inflates peak memory.

onnx_ir tensors are read-only whether or not load_to_model is called
(verified), so this changes no behavior; the quant kernels already copy
(.float()/.copy()/new buffers) before use. rtn/kquant/hqq tests pass (21).

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
---
 olive/passes/onnx/hqq_quantization.py    | 1 -
 olive/passes/onnx/kquant_quantization.py | 1 -
 olive/passes/onnx/rtn_quantization.py    | 1 -
 3 files changed, 3 deletions(-)

diff --git a/olive/passes/onnx/hqq_quantization.py b/olive/passes/onnx/hqq_quantization.py
index 992e430c79..c55a177a85 100644
--- a/olive/passes/onnx/hqq_quantization.py
+++ b/olive/passes/onnx/hqq_quantization.py
@@ -72,7 +72,6 @@ def _run_for_config(
             return model
         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
         ir_model = model.load_ir_model()
-        ir.external_data.load_to_model(ir_model)
         ir_model.graph.opset_imports[MSFT_DOMAIN] = 1
         self._quantize_model(
             ir_model,
diff --git a/olive/passes/onnx/kquant_quantization.py b/olive/passes/onnx/kquant_quantization.py
index 84e93434cf..e848a17cde 100644
--- a/olive/passes/onnx/kquant_quantization.py
+++ b/olive/passes/onnx/kquant_quantization.py
@@ -273,7 +273,6 @@ def _run_for_config(
             output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
 
         ir_model = model.load_ir_model()
-        ir.external_data.load_to_model(ir_model)
         ir_model.graph.opset_imports[MSFT_DOMAIN] = 1
         self._quantize_model(
             ir_model,
diff --git a/olive/passes/onnx/rtn_quantization.py b/olive/passes/onnx/rtn_quantization.py
index 0665a3afc5..4e541aa86c 100644
--- a/olive/passes/onnx/rtn_quantization.py
+++ b/olive/passes/onnx/rtn_quantization.py
@@ -77,7 +77,6 @@ def _run_for_config(
     ) -> ONNXModelHandler:
         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
         ir_model = model.load_ir_model()
-        ir.external_data.load_to_model(ir_model)
         ir_model.graph.opset_imports[MSFT_DOMAIN] = 1
         self._quantize_model(
             ir_model,