From 19b2e731d85c557afe1ceaedd372d4ca1b1e475c Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Thu, 15 Oct 2026 12:46:24 +0000
Subject: [PATCH 01/17] [DLA_transforms_17] [17/17] Opt-in FP16 strong typing
 for DLA

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 CHANGELOG.rst                                 |   3 +-
 modelopt/onnx/graph_surgery/__main__.py       |   6 +
 .../dla_force_fp16_activations.py             | 396 ++++++++++++++++++
 .../onnx/graph_surgery/make_dla_compatible.py |  17 +
 .../test_dla_force_fp16_activations.py        | 216 ++++++++++
 5 files changed, 637 insertions(+), 1 deletion(-)
 create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
 create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 3906cc2ff84..ff846f92e2f 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -35,7 +35,8 @@ Changelog
 
 - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs.
 - Add the ordered DLA compatibility pipeline through the ``modelopt.onnx.graph_surgery.make_dla_compatible``
-  API and ``make-dla-compatible`` CLI command. Use ``--remove-all-qdq`` to remove Q/DQ pairs of every dtype.
+  API and ``make-dla-compatible`` CLI command. Use ``--remove-all-qdq`` to remove Q/DQ pairs of every dtype,
+  or ``--force-fp16-activations`` for FP16 internal computation with FP32 model boundaries.
 
 **Backward Breaking Changes**
 
diff --git a/modelopt/onnx/graph_surgery/__main__.py b/modelopt/onnx/graph_surgery/__main__.py
index 0c130960de4..62135c974ac 100644
--- a/modelopt/onnx/graph_surgery/__main__.py
+++ b/modelopt/onnx/graph_surgery/__main__.py
@@ -284,6 +284,11 @@ def main():
         action="store_true",
         help="Remove QDQ pairs for every dtype, matching the C++ remove-all mode",
     )
+    dla_parser.add_argument(
+        "--force-fp16-activations",
+        action="store_true",
+        help="Use FP16 internal activations and weights with FP32 model boundaries",
+    )
 
     # Analyze subcommand
     analyze_parser = subparsers.add_parser(
@@ -367,6 +372,7 @@ def main():
             verbose=not args.quiet,
             qdq_quantized_dtypes=args.qdq_quantized_dtypes,
             remove_all_qdq=args.remove_all_qdq,
+            force_fp16_activations=args.force_fp16_activations,
         )
 
     elif args.command == "analyze":
diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
new file mode 100644
index 00000000000..08f76ab13e3
--- /dev/null
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -0,0 +1,396 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Apply the C++ FP16 activation policy while retaining FP32 public boundaries."""
+
+import copy
+import math
+
+import numpy as np
+from onnx import TensorProto, helper, numpy_helper
+
+from ..utils.graph_utils import GraphCache, _attribute, convert_initializers_to_dtype
+from .dla_rank_conversion import _apply_non4d_lpnorm
+
+__all__ = []
+
+
+def _cast(source, output, dtype):
+    name = source + "_to_fp16" if dtype == TensorProto.FLOAT16 else output + "_to_fp32"
+    return helper.make_node("Cast", [source], [output], name=name, to=dtype)
+
+
+def _is_cast(node, dtype):
+    target = _attribute(node, "to") if node is not None else None
+    return node is not None and node.op_type == "Cast" and target is not None and target.i == dtype
+
+
+def _apply_fp16_lpnorm(model):
+    """Decompose last-axis L2 normalization followed by a static gamma multiply."""
+    graph = model.graph
+    cache = GraphCache(graph)
+    static = cache.init_names | {n.output[0] for n in graph.node if n.op_type == "Constant"}
+    opset = next((o.version for o in model.opset_import if o.domain in ("", "ai.onnx")), 1)
+    nodes = []
+    for node in graph.node:
+        if node.op_type != "LpNormalization":
+            nodes.append(node)
+            continue
+        axis_attr, p_attr = _attribute(node, "axis"), _attribute(node, "p")
+        axis = axis_attr.i if axis_attr is not None else -1
+        shape = cache.get_shape(node.input[0])
+        consumers = cache.get_consumers(node.output[0])
+        scaled = (
+            len(consumers) == 1
+            and consumers[0][0].op_type == "Mul"
+            and any(name in static for name in consumers[0][0].input if name != node.output[0])
+        )
+        if (
+            (p_attr is not None and p_attr.i != 2)
+            or not scaled
+            or not (axis == -1 or (shape is not None and axis == len(shape) - 1))
+        ):
+            nodes.append(node)
+            continue
+        base = node.name or node.output[0] + "_lpnorm"
+        axes = cache.unique_name(base + "_axes")
+        epsilon = cache.unique_name(base + "_epsilon")
+        if opset >= 13:
+            graph.initializer.append(numpy_helper.from_array(np.array([axis], np.int64), axes))
+        graph.initializer.append(numpy_helper.from_array(np.float32(2**-14), epsilon))
+        power, total, norm, safe = [
+            cache.unique_name(base + suffix)
+            for suffix in ("_power_output", "_sum_output", "_norm_output", "_safe_norm_output")
+        ]
+        nodes.extend(
+            [
+                helper.make_node(
+                    "Mul", [node.input[0], node.input[0]], [power], name=base + "_square"
+                ),
+                helper.make_node(
+                    "ReduceSum",
+                    [power, axes] if opset >= 13 else [power],
+                    [total],
+                    name=base + "_sum",
+                    **({} if opset >= 13 else {"axes": [axis]}),
+                    keepdims=1,
+                ),
+                helper.make_node("Sqrt", [total], [norm], name=base + "_sqrt"),
+                helper.make_node("Max", [norm, epsilon], [safe], name=base + "_safe_norm"),
+                helper.make_node(
+                    "Div", [node.input[0], safe], list(node.output), name=base + "_divide"
+                ),
+            ]
+        )
+    del graph.node[:]
+    graph.node.extend(nodes)
+    return _apply_non4d_lpnorm(model)
+
+
+def _apply_force_fp16_activations(model):
+    """Run after the ordered DLA passes and their final symbolic shape inference."""
+    graph = model.graph
+    cache = GraphCache(graph)
+    fp32, fp16 = TensorProto.FLOAT, TensorProto.FLOAT16
+    preserved_inits = {
+        name
+        for name in cache.init_names
+        if any(n.op_type == "Resize" and port in (1, 2) for n, port in cache.get_consumers(name))
+    }
+    excluded_inits = [init for init in graph.initializer if init.name in preserved_inits]
+    for init in excluded_inits:
+        graph.initializer.remove(init)
+    try:
+        convert_initializers_to_dtype(graph, "float16")
+    finally:
+        graph.initializer.extend(excluded_inits)
+    # Conversion replaces initializer objects, invalidating the cached references.
+    cache = GraphCache(graph)
+    preserved = set()
+    terminal_concats = set()
+    for node in graph.node:
+        if not _is_cast(node, fp32) or cache.get_dtype(node.input[0]) in (fp32, fp16, None):
+            continue
+        for concat, _ in cache.get_consumers(node.output[0]):
+            if concat.op_type != "Concat" or len(concat.output) != 1:
+                continue
+            uses = cache.get_consumers(concat.output[0])
+            if len(uses) == 1:
+                boundary, port = uses[0]
+                if (
+                    port == 0
+                    and boundary.op_type == "Squeeze"
+                    and len(boundary.output) == 1
+                    and (boundary.output[0] in cache.graph_output_names)
+                ):
+                    preserved.add(concat.output[0])
+                    terminal_concats.add(id(concat))
+    for node in graph.node:
+        if _is_cast(node, fp32) and node.output[0] not in cache.graph_output_names | preserved:
+            _attribute(node, "to").i = fp16
+    for value in graph.value_info:
+        if (
+            value.type.HasField("tensor_type")
+            and value.type.tensor_type.elem_type == fp32
+            and (value.name not in cache.graph_output_names | preserved | preserved_inits)
+        ):
+            value.type.tensor_type.elem_type = fp16
+    for value in graph.input:
+        if (
+            value.name in cache.init_names - preserved_inits
+            and value.type.tensor_type.elem_type == fp32
+        ):
+            value.type.tensor_type.elem_type = fp16
+    before, after = {}, {}
+    front, back = [], []
+    for node in graph.node:
+        if id(node) not in terminal_concats:
+            continue
+        for port, source in enumerate(node.input):
+            if source not in preserved:
+                output = cache.unique_name(source + "_concat_fp32")
+                before.setdefault(id(node), []).append(_cast(source, output, fp32))
+                node.input[port] = output
+                preserved.add(output)
+    for value in graph.input:
+        if value.name in cache.init_names or value.type.tensor_type.elem_type != fp32:
+            continue
+        direct = None
+        rank = len(value.type.tensor_type.shape.dim)
+        for consumer, port in cache.get_consumers(value.name):
+            if _is_cast(consumer, fp16):
+                continue
+            shape_info = cache.vi_map.get(consumer.output[0])
+            adapter = (
+                rank < 4
+                and port == 0
+                and consumer.op_type == "Unsqueeze"
+                and (shape_info is not None and len(shape_info.type.tensor_type.shape.dim) == 4)
+            )
+            if adapter:
+                assert shape_info is not None
+                source = consumer.output[0]
+                output = cache.unique_name(source + "_fp16")
+                after.setdefault(id(consumer), []).append(_cast(source, output, fp16))
+                for downstream, slot in cache.get_consumers(source):
+                    if not _is_cast(downstream, fp16):
+                        downstream.input[slot] = output
+                shape_info.type.tensor_type.elem_type = fp32
+            else:
+                if direct is None:
+                    direct = cache.unique_name(value.name + "_fp16")
+                    front.append(_cast(value.name, direct, fp16))
+                consumer.input[port] = direct
+    if any(
+        n.op_type in {"QuantizeLinear", "DequantizeLinear"} and n.domain in ("", "ai.onnx")
+        for n in graph.node
+    ):
+        opsets = [o for o in model.opset_import if o.domain in ("", "ai.onnx")]
+        if opsets:
+            opsets[0].version = max(19, opsets[0].version)
+        else:
+            model.opset_import.append(helper.make_opsetid("", 19))
+    for value in graph.output:
+        if value.type.tensor_type.elem_type != fp32:
+            continue
+        producer = cache.get_producer(value.name)
+        if producer is not None and producer.op_type == "Squeeze":
+            source = producer.input[0]
+            if source not in preserved:
+                output = cache.unique_name(source + "_to_fp32_fp16")
+                before.setdefault(id(producer), []).append(_cast(source, output, fp32))
+                producer.input[0] = output
+            continue
+        if _is_cast(producer, fp32):
+            continue
+        tensor = cache.get_init(value.name) if producer is None else None
+        if producer is None and (tensor is None or tensor.data_type != fp16):
+            continue
+        internal = cache.unique_name(value.name + "_fp16")
+        if producer is None:
+            assert tensor is not None
+            tensor.name = internal
+        else:
+            for port, name in enumerate(producer.output):
+                if name == value.name:
+                    producer.output[port] = internal
+                    break
+        for annotation in graph.value_info:
+            if annotation.name == value.name:
+                annotation.name = internal
+                annotation.type.tensor_type.elem_type = fp16
+        for consumer, port in cache.get_consumers(value.name):
+            consumer.input[port] = internal
+        back.append(_cast(internal, value.name, fp32))
+    nodes = list(front)
+    for node in graph.node:
+        nodes.extend(before.get(id(node), []))
+        nodes.append(node)
+        nodes.extend(after.get(id(node), []))
+    nodes.extend(back)
+    del graph.node[:]
+    graph.node.extend(nodes)
+    _rewrite_terminal_projections(graph)
+    return model
+
+
+def _rewrite_terminal_projections(graph):
+    """Bound terminal projection channels to DLA's 8192-element dimension limit."""
+    cache = GraphCache(graph)
+    removed, removed_values, removed_inits = set(), set(), set()
+    replacements, additions = {}, []
+
+    def attributes(node, key, expected):
+        value = _attribute(node, key)
+        return (
+            value is not None
+            and (list(value.ints) if isinstance(expected, list) else value.i) == expected
+        )
+
+    def axes(node, expected):
+        value = cache.get_init_array(node.input[1]) if len(node.input) == 2 else None
+        return value is not None and value.reshape(-1).tolist() == expected
+
+    for output in graph.output:
+        if output.type.tensor_type.elem_type != TensorProto.FLOAT:
+            continue
+        chain = []
+        current = output.name
+        for op in ("Squeeze", "Cast", "Concat", "ReduceSum", "Concat", "Transpose", "Conv"):
+            node = cache.get_producer(current)
+            if node is None or node.op_type != op:
+                break
+            chain.append(node)
+            current = node.input[0]
+        if len(chain) != 7:
+            continue
+        squeeze, cast, final, reduce, initial, transpose, conv = chain
+        if not (
+            axes(squeeze, [0])
+            and _is_cast(cast, TensorProto.FLOAT)
+            and len(final.input) == len(initial.input) == 1
+            and axes(reduce, [0])
+            and attributes(reduce, "keepdims", 1)
+            and attributes(transpose, "perm", [0, 3, 2, 1])
+            and len(conv.input) == 2
+            and len(conv.output) == 1
+        ):
+            continue
+        weight = cache.get_init(conv.input[1])
+        dq = cache.get_producer(conv.input[1]) if weight is None else None
+        if dq is not None:
+            parameters = [cache.get_init(name) for name in dq.input[1:]]
+            if (
+                dq.op_type != "DequantizeLinear"
+                or len(dq.input) < 2
+                or any(
+                    parameter is None or math.prod(parameter.dims) != 1 for parameter in parameters
+                )
+            ):
+                continue
+            weight = cache.get_init(dq.input[0])
+            chain.append(dq)
+        if (
+            weight is None
+            or len(weight.dims) != 4
+            or weight.dims[0] <= 8192
+            or not weight.HasField("raw_data")
+        ):
+            continue
+        if weight.data_type != (TensorProto.UINT8 if dq is not None else TensorProto.FLOAT16):
+            continue
+        chain_ids = {id(n) for n in chain}
+        if any(
+            id(consumer) not in chain_ids
+            for name in (weight.name, squeeze.input[1], reduce.input[1])
+            for consumer, _ in cache.get_consumers(name)
+        ):
+            continue
+        if any(len(cache.get_consumers(n.output[0])) != 1 for n in chain[1:]):
+            continue
+        channels = weight.dims[0]
+        if cache.get_shape(output.name) != [1, 1, channels]:
+            continue
+        bytes_per_channel = math.prod(weight.dims[1:]) * (1 if dq is not None else 2)
+        if bytes_per_channel <= 0 or len(weight.raw_data) != channels * bytes_per_channel:
+            continue
+        nodes, outputs = [], []
+        for index, start in enumerate(range(0, channels, 8192)):
+            count = min(8192, channels - start)
+            suffix = f"_dla_chunk_{index}"
+            chunk = TensorProto(
+                name=cache.unique_name(weight.name + suffix),
+                data_type=weight.data_type,
+                dims=[count, *weight.dims[1:]],
+            )
+            chunk.raw_data = weight.raw_data[
+                start * bytes_per_channel : (start + count) * bytes_per_channel
+            ]
+            additions.append(chunk)
+            cloned_conv = copy.deepcopy(conv)
+            cloned_conv.name = cache.unique_name(conv.name + suffix)
+            if dq is not None:
+                cloned_dq = copy.deepcopy(dq)
+                cloned_dq.name = cache.unique_name(dq.name + suffix)
+                cloned_dq.input[0] = chunk.name
+                cloned_dq.output[0] = cache.unique_name(dq.output[0] + suffix)
+                nodes.append(cloned_dq)
+                cloned_conv.input[1] = cloned_dq.output[0]
+            else:
+                cloned_conv.input[1] = chunk.name
+            cloned_conv.output[0] = cache.unique_name(conv.output[0] + suffix)
+            nodes.append(cloned_conv)
+            shape = cache.unique_name(conv.output[0] + suffix + "_bounded_shape")
+            additions.append(numpy_helper.from_array(np.array([1, 1, count], np.int64), shape))
+            bounded = cache.unique_name(conv.output[0] + suffix + "_bounded")
+            nodes.append(
+                helper.make_node(
+                    "Reshape",
+                    [cloned_conv.output[0], shape],
+                    [bounded],
+                    name=cache.unique_name(conv.name + suffix + "_bounded_reshape"),
+                )
+            )
+            result = cache.unique_name(conv.output[0] + suffix + "_fp32")
+            nodes.append(_cast(bounded, result, TensorProto.FLOAT))
+            outputs.append(result)
+        nodes.append(
+            helper.make_node(
+                "Concat",
+                outputs,
+                [output.name],
+                axis=2,
+                name=cache.unique_name(output.name + "_dla_terminal_concat"),
+            )
+        )
+        replacements[id(conv)] = nodes
+        removed.update(id(n) for n in chain)
+        removed_values.update(name for n in chain for name in n.output if name != output.name)
+        removed_inits.update([weight.name, squeeze.input[1], reduce.input[1]])
+    if not replacements:
+        return
+    nodes = []
+    for node in graph.node:
+        nodes.extend(replacements.get(id(node), []))
+        if id(node) not in removed:
+            nodes.append(node)
+    retained_inits = [t for t in graph.initializer if t.name not in removed_inits]
+    retained_info = [v for v in graph.value_info if v.name not in removed_values]
+    del graph.node[:]
+    graph.node.extend(nodes)
+    del graph.initializer[:]
+    graph.initializer.extend([*retained_inits, *additions])
+    del graph.value_info[:]
+    graph.value_info.extend(retained_info)
diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py
index 2700f54ce00..71e4ba95875 100644
--- a/modelopt/onnx/graph_surgery/make_dla_compatible.py
+++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py
@@ -35,6 +35,10 @@
 from .dla_transforms.dla_cast_to_fp32 import _apply_cast_to_fp32
 from .dla_transforms.dla_constants_to_initializers import _constants_to_initializers_from_model
 from .dla_transforms.dla_decompose_lstm import _apply_decompose_lstm
+from .dla_transforms.dla_force_fp16_activations import (
+    _apply_force_fp16_activations,
+    _apply_fp16_lpnorm,
+)
 from .dla_transforms.dla_graph_cleanup import _remove_unused_initializers
 from .dla_transforms.dla_greater import _apply_greater
 from .dla_transforms.dla_indexing import (
@@ -242,6 +246,7 @@ def _transform_make_dla_compatible(
     verbose: bool = True,
     qdq_quantized_dtypes: Iterable[int | str] | None = None,
     remove_all_qdq: bool = False,
+    force_fp16_activations: bool = False,
 ) -> onnx.ModelProto:
     """Apply the logged DLA sequence in memory, then infer shapes and validate."""
     started = time.perf_counter()
@@ -311,12 +316,21 @@ def _transform_make_dla_compatible(
         options = _boundary_stage_options(
             name, quant_types, verbose, prepared_input_ops, prepared_output_ops
         )
+        if name == "Non4DLpNorm" and force_fp16_activations:
+            transform = _apply_fp16_lpnorm
         result = transform(model, **options)
         if result is not None:
             model = result
     model = infer_shapes(model)
     model = _apply_reorder_nodes(model)
     onnx.checker.check_model(model)
+    if force_fp16_activations:
+        if verbose:
+            logger.info("[DLA pipeline] Applying ForceFp16Activations")
+        model = _apply_force_fp16_activations(model)
+        # Match C++: symbolic inference stays before precision conversion.
+        model = onnx.shape_inference.infer_shapes(model)
+        onnx.checker.check_model(model)
     if verbose:
         logger.info("[DLA pipeline] Done (%.1fs).", time.perf_counter() - started)
     return model
@@ -331,6 +345,7 @@ def dla_make_dla_compatible(
     verbose: bool = True,
     qdq_quantized_dtypes: Iterable[int | str] | None = None,
     remove_all_qdq: bool = False,
+    force_fp16_activations: bool = False,
 ) -> onnx.ModelProto:
     """Load an ONNX model, apply the ordered DLA transforms, and save the result.
 
@@ -348,6 +363,7 @@ def dla_make_dla_compatible(
         qdq_quantized_dtypes: Zero-point dtypes whose QDQ pairs are removed.
             Defaults to UINT16 and INT16. Accepts ONNX dtype IDs or names.
         remove_all_qdq: Remove QDQ pairs for every dtype, matching the C++ flag.
+        force_fp16_activations: Use FP16 internals while preserving FP32 public boundaries.
             Cannot be combined with qdq_quantized_dtypes.
 
     Returns:
@@ -358,6 +374,7 @@ def dla_make_dla_compatible(
         verbose=verbose,
         qdq_quantized_dtypes=qdq_quantized_dtypes,
         remove_all_qdq=remove_all_qdq,
+        force_fp16_activations=force_fp16_activations,
     )
     output_dir = os.path.dirname(output_path)
     if output_dir:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
new file mode 100644
index 00000000000..40a1dc5b570
--- /dev/null
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -0,0 +1,216 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Focused FP16 payload, boundary, and DLA projection assertions on small fixtures."""
+
+import numpy as np
+import onnx
+import onnxruntime as ort
+import pytest
+from onnx import TensorProto as T
+from onnx import helper as h
+from onnx import numpy_helper as nh
+from onnx.reference import ReferenceEvaluator
+
+from modelopt.onnx.graph_surgery import make_dla_compatible
+from modelopt.onnx.graph_surgery.dla_transforms.dla_force_fp16_activations import (
+    _apply_force_fp16_activations,
+    _apply_fp16_lpnorm,
+)
+
+
+def _value(name, shape, dtype=T.FLOAT):
+    return h.make_tensor_value_info(name, dtype, shape)
+
+
+def _model(nodes, inputs, outputs, initializers, opset=19):
+    model = h.make_model(
+        h.make_graph(nodes, "fp16", inputs, outputs, initializers),
+        opset_imports=[h.make_opsetid("", opset)],
+    )
+    model.ir_version = 9
+    return onnx.shape_inference.infer_shapes(model)
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+def test_public_option_preserves_boundaries_and_selects_weight_precision(tmp_path, enabled):
+    values = np.array([1.0003, 0.12345], np.float32)
+    model = _model(
+        [h.make_node("Mul", ["x", "w"], ["y"])],
+        [_value("x", [1, 1, 1, 2])],
+        [_value("y", [1, 1, 1, 2])],
+        [nh.from_array(values.reshape(1, 1, 1, 2), "w")],
+    )
+    source, output = tmp_path / "source.onnx", tmp_path / "output.onnx"
+    onnx.save(model, source)
+    result = make_dla_compatible(
+        str(source),
+        str(output),
+        verbose=False,
+        use_external_data=False,
+        force_fp16_activations=enabled,
+    )
+    assert result.graph.input[0].type.tensor_type.elem_type == T.FLOAT
+    assert result.graph.output[0].name == "y"
+    assert result.graph.output[0].type.tensor_type.elem_type == T.FLOAT
+    assert result.graph.initializer[0].data_type == (T.FLOAT16 if enabled else T.FLOAT)
+    data = np.array([[[[0.3333, 2.0003]]]], np.float32)
+    dtype = np.float16 if enabled else np.float32
+    expected = (data.astype(dtype) * values.astype(dtype)).astype(np.float32)
+    actual = ReferenceEvaluator(onnx.load(output)).run(None, {"x": data})[0]
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_initializer_rounding_preserves_subnormals_signed_zero_and_nan():
+    values = np.array([0, -0.0, 2**-24, 2**-25, 2**-14, 65504, 65520, 0, 0], np.float32)
+    values.view(np.uint32)[-2:] = [0x7F812345, 0xFF812345]
+    model = _model([], [], [_value("weight", [len(values)])], [nh.from_array(values, "weight")])
+    result = _apply_force_fp16_activations(model)
+    actual = nh.to_array(result.graph.initializer[0]).view(np.uint16)
+    np.testing.assert_array_equal(actual[:-2], [0, 0x8000, 1, 0, 0x400, 0x7BFF, 0x7C00])
+    assert np.isnan(actual[-2:].view(np.float16)).all()
+    assert result.graph.node[0].op_type == "Cast"
+    assert result.graph.node[0].output[0] == "weight"
+
+
+def test_boundary_adapters_and_direct_input_consumers_get_their_own_casts():
+    model = _model(
+        [
+            h.make_node("Unsqueeze", ["x", "axes"], ["x4"]),
+            h.make_node("Relu", ["x4"], ["r4"]),
+            h.make_node("Squeeze", ["r4", "axes"], ["y"]),
+            h.make_node("Neg", ["x"], ["direct"]),
+        ],
+        [_value("x", [1, 2])],
+        [_value("y", [1, 2]), _value("direct", [1, 2])],
+        [nh.from_array(np.array([0, 1], np.int64), "axes")],
+    )
+    result = _apply_force_fp16_activations(model)
+    producers = {name: n for n in result.graph.node for name in n.output}
+    relu = next(n for n in result.graph.node if n.op_type == "Relu")
+    assert producers[relu.input[0]].op_type == "Cast"
+    assert producers[producers[relu.input[0]].input[0]].op_type == "Unsqueeze"
+    squeeze = producers["y"]
+    assert squeeze.op_type == "Squeeze"
+    assert producers[squeeze.input[0]].op_type == "Cast"
+    actual = ReferenceEvaluator(result).run(None, {"x": np.array([[-2, 3]], np.float32)})
+    np.testing.assert_array_equal(actual[0], [[0, 3]])
+    np.testing.assert_array_equal(actual[1], [[2, -3]])
+
+
+@pytest.mark.parametrize("operation", ["Resize", "QuantizeLinear"])
+def test_schema_specific_float_parameters_and_qdq_opset(operation):
+    if operation == "Resize":
+        nodes = [
+            h.make_node("Mul", ["x", "weight"], ["weighted"]),
+            h.make_node("Resize", ["weighted", "roi", "scales"], ["y"]),
+        ]
+        initializers = [
+            nh.from_array(np.full((1, 1, 1, 2), 0.5, np.float32), "weight"),
+            nh.from_array(np.array([], np.float32), "roi"),
+            nh.from_array(np.array([1, 1, 1, 2], np.float32), "scales"),
+        ]
+        output_shape = [1, 1, 1, 4]
+    else:
+        nodes = [
+            h.make_node("QuantizeLinear", ["x", "scale", "zero"], ["q"]),
+            h.make_node("DequantizeLinear", ["q", "scale", "zero"], ["y"]),
+        ]
+        initializers = [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")]
+        output_shape = [1, 1, 1, 2]
+    model = _model(
+        nodes, [_value("x", [1, 1, 1, 2])], [_value("y", output_shape)], initializers, opset=17
+    )
+    result = _apply_force_fp16_activations(model)
+    result = onnx.shape_inference.infer_shapes(result, check_type=True)
+    onnx.checker.check_model(result)
+    types = {t.name: t.data_type for t in result.graph.initializer}
+    if operation == "Resize":
+        assert types == {"weight": T.FLOAT16, "roi": T.FLOAT, "scales": T.FLOAT}
+        assert result.opset_import[0].version == 17
+    else:
+        assert types == {"scale": T.FLOAT16, "zero": T.INT8}
+        assert result.opset_import[0].version == 19
+
+
+def test_scaled_l2_normalization_keeps_zero_slices_finite():
+    model = _model(
+        [
+            h.make_node("LpNormalization", ["x"], ["norm"], p=2, axis=-1),
+            h.make_node("Mul", ["norm", "gamma"], ["y"]),
+        ],
+        [_value("x", [1, 1, 2, 2])],
+        [_value("y", [1, 1, 2, 2])],
+        [nh.from_array(np.ones((1, 1, 1, 2), np.float32), "gamma")],
+    )
+    result = _apply_force_fp16_activations(_apply_fp16_lpnorm(model))
+    assert not any(n.op_type == "LpNormalization" for n in result.graph.node)
+    assert any(n.op_type == "Max" for n in result.graph.node)
+    data = np.array([[[[0, 0], [3, 4]]]], np.float32)
+    actual = ReferenceEvaluator(result).run(None, {"x": data})[0]
+    assert np.isfinite(actual).all()
+    np.testing.assert_allclose(actual, [[[[0, 0], [0.6, 0.8]]]], atol=3e-4)
+
+
+def test_terminal_projection_splits_channels_and_preserves_values():
+    channels = 8193
+    model = _model(
+        [
+            h.make_node("Conv", ["x", "weight"], ["conv"]),
+            h.make_node("Transpose", ["conv"], ["transposed"], perm=[0, 3, 2, 1]),
+            h.make_node("Concat", ["transposed"], ["concat1"], axis=0),
+            h.make_node("ReduceSum", ["concat1", "axes"], ["reduced"], keepdims=1),
+            h.make_node("Concat", ["reduced"], ["concat2"], axis=2),
+            h.make_node("Squeeze", ["concat2", "axes"], ["y"]),
+        ],
+        [_value("x", [1, 1, 1, 1])],
+        [_value("y", [1, 1, channels])],
+        [
+            nh.from_array(np.full((channels, 1, 1, 1), np.float32(0.25)), "weight"),
+            nh.from_array(np.array([0], np.int64), "axes"),
+        ],
+    )
+    result = _apply_force_fp16_activations(model)
+    assert sum(n.op_type == "Conv" for n in result.graph.node) == 2
+    weights = [t for t in result.graph.initializer if t.data_type == T.FLOAT16]
+    assert [list(t.dims) for t in weights] == [[8192, 1, 1, 1], [1, 1, 1, 1]]
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.inter_op_num_threads = 1
+    actual = ort.InferenceSession(
+        result.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, {"x": np.array([[[[2]]]], np.float32)})[0]
+    np.testing.assert_array_equal(actual, np.full((1, 1, channels), np.float32(0.5)))
+
+
+def test_excluded_initializers_are_restored_when_conversion_fails():
+    """Retain excluded weight payloads when the original converter rejects an included tensor."""
+    roi = nh.from_array(np.array([], np.float32), "roi")
+    scales = nh.from_array(np.array([1, 1, 1, 2], np.float32), "scales")
+    invalid = T(name="invalid", data_type=T.FLOAT, dims=[2], float_data=[1])
+    model = h.make_model(
+        h.make_graph(
+            [h.make_node("Resize", ["x", "roi", "scales"], ["y"])],
+            "conversion_failure",
+            [_value("x", [1, 1, 1, 2])],
+            [_value("y", [1, 1, 1, 4])],
+            [roi, invalid, scales],
+        )
+    )
+    expected = {value.name: value.SerializeToString() for value in (roi, scales)}
+    with pytest.raises(ValueError):
+        _apply_force_fp16_activations(model)
+    actual = {value.name: value.SerializeToString() for value in model.graph.initializer}
+    assert all(actual[name] == payload for name, payload in expected.items())

From 3bd0ac9a8bb7d070bafbf6b1cd481fd241257fb9 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 06:22:44 +0000
Subject: [PATCH 02/17] Preserve FP16 bindings and terminal projection owners

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../dla_force_fp16_activations.py             | 379 +++++++++--
 .../onnx/graph_surgery/make_dla_compatible.py |   3 +-
 .../test_dla_force_fp16_activations.py        | 615 ++++++++++++++++++
 3 files changed, 931 insertions(+), 66 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
index 08f76ab13e3..40e375338a4 100644
--- a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -19,33 +19,63 @@
 import math
 
 import numpy as np
-from onnx import TensorProto, helper, numpy_helper
+from onnx import AttributeProto, TensorProto, defs, helper, numpy_helper
 
-from ..utils.graph_utils import GraphCache, _attribute, convert_initializers_to_dtype
-from .dla_rank_conversion import _apply_non4d_lpnorm
+from ..utils.graph_utils import (
+    GraphCache,
+    _attribute,
+    batch_replace_nodes,
+    convert_initializers_to_dtype,
+)
+from .dla_constants_to_initializers import _add_legacy_initializer_inputs
+from .dla_graph_cleanup import _subgraph_captures
+from .dla_indexing import _insert_support_nodes, _operator_nodes
+from .dla_rank_conversion import _apply_non4d_lpnorm, _runtime_input_names
 
 __all__ = []
 
 
-def _cast(source, output, dtype):
+def _cast(source, output, dtype, cache):
+    """Allocate a boundary cast without colliding with existing graph owners."""
     name = source + "_to_fp16" if dtype == TensorProto.FLOAT16 else output + "_to_fp32"
-    return helper.make_node("Cast", [source], [output], name=name, to=dtype)
+    return helper.make_node("Cast", [source], [output], name=cache.unique_name(name), to=dtype)
+
+
+def _fp16_operator_nodes(model, operations):
+    """Select standard operations and supported contrib schemas for FP16 policy."""
+    nodes = _operator_nodes(model, operations)
+    if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import):
+        nodes += _operator_nodes(
+            model, operations & {"QuantizeLinear", "DequantizeLinear", "Gelu"}, ("com.microsoft",)
+        )
+    return nodes
 
 
 def _is_cast(node, dtype):
+    """Match a standard Cast with the requested destination type."""
     target = _attribute(node, "to") if node is not None else None
-    return node is not None and node.op_type == "Cast" and target is not None and target.i == dtype
+    return (
+        node is not None
+        and node.domain in {"", "ai.onnx"}
+        and not node.overload
+        and node.op_type == "Cast"
+        and target is not None
+        and target.i == dtype
+    )
 
 
 def _apply_fp16_lpnorm(model):
     """Decompose last-axis L2 normalization followed by a static gamma multiply."""
     graph = model.graph
     cache = GraphCache(graph)
-    static = cache.init_names | {n.output[0] for n in graph.node if n.op_type == "Constant"}
+    standard = {id(node) for node in _operator_nodes(model, {"LpNormalization", "Mul", "Constant"})}
+    static = cache.init_names - _runtime_input_names(model) | {
+        node.output[0] for node in graph.node if node.op_type == "Constant" and id(node) in standard
+    }
     opset = next((o.version for o in model.opset_import if o.domain in ("", "ai.onnx")), 1)
     nodes = []
     for node in graph.node:
-        if node.op_type != "LpNormalization":
+        if node.op_type != "LpNormalization" or id(node) not in standard:
             nodes.append(node)
             continue
         axis_attr, p_attr = _attribute(node, "axis"), _attribute(node, "p")
@@ -55,6 +85,7 @@ def _apply_fp16_lpnorm(model):
         scaled = (
             len(consumers) == 1
             and consumers[0][0].op_type == "Mul"
+            and id(consumers[0][0]) in standard
             and any(name in static for name in consumers[0][0].input if name != node.output[0])
         )
         if (
@@ -77,48 +108,219 @@ def _apply_fp16_lpnorm(model):
         nodes.extend(
             [
                 helper.make_node(
-                    "Mul", [node.input[0], node.input[0]], [power], name=base + "_square"
+                    "Mul",
+                    [node.input[0], node.input[0]],
+                    [power],
+                    name=cache.unique_name(base + "_square"),
                 ),
                 helper.make_node(
                     "ReduceSum",
                     [power, axes] if opset >= 13 else [power],
                     [total],
-                    name=base + "_sum",
+                    name=cache.unique_name(base + "_sum"),
                     **({} if opset >= 13 else {"axes": [axis]}),
                     keepdims=1,
                 ),
-                helper.make_node("Sqrt", [total], [norm], name=base + "_sqrt"),
-                helper.make_node("Max", [norm, epsilon], [safe], name=base + "_safe_norm"),
+                helper.make_node("Sqrt", [total], [norm], name=cache.unique_name(base + "_sqrt")),
                 helper.make_node(
-                    "Div", [node.input[0], safe], list(node.output), name=base + "_divide"
+                    "Max", [norm, epsilon], [safe], name=cache.unique_name(base + "_safe_norm")
+                ),
+                helper.make_node(
+                    "Div",
+                    [node.input[0], safe],
+                    list(node.output),
+                    name=cache.unique_name(base + "_divide"),
                 ),
             ]
         )
-    del graph.node[:]
-    graph.node.extend(nodes)
+    retained = {id(node) for node in nodes}
+    batch_replace_nodes(
+        graph,
+        {index: ([node], []) for index, node in enumerate(graph.node) if id(node) not in retained},
+    )
+    _insert_support_nodes(graph, nodes)
+    _add_legacy_initializer_inputs(model)
     return _apply_non4d_lpnorm(model)
 
 
+def _validate_float_weights(graph):
+    """Validate all FLOAT payloads before calling the existing weight converter."""
+    for tensor in graph.initializer:
+        if tensor.data_type != TensorProto.FLOAT:
+            continue
+        count = math.prod(tensor.dims)
+        raw = tensor.HasField("raw_data")
+        if (
+            count < 0
+            or (raw and len(tensor.raw_data) != count * 4)
+            or (not raw and len(tensor.float_data) != count)
+            or (tensor.data_location == TensorProto.EXTERNAL and not raw)
+        ):
+            raise ValueError(f"Invalid or unloaded FLOAT initializer {tensor.name!r}")
+
+
+def _fp32_control(node, port, native, required):
+    """Retain schema-mandated FLOAT inputs and inputs of native FP32 regions."""
+    return id(node) in native or (id(node), port) in required
+
+
+def _fp32_regions(model, cache):
+    """Preserve native floating dependencies and schema-required FLOAT bindings."""
+    supported = _fp16_operator_nodes(model, {node.op_type for node in model.graph.node})
+    contrib = {id(node) for node in supported if node.domain == "com.microsoft"}
+    standard = {id(node) for node in supported}
+    native = {
+        id(node)
+        for node in model.graph.node
+        if id(node) not in standard
+        or any(
+            attr.type in {AttributeProto.GRAPH, AttributeProto.GRAPHS} for attr in node.attribute
+        )
+        or any(
+            (info := cache.vi_map.get(name)) is not None
+            and info.type.WhichOneof("value") not in {None, "tensor_type"}
+            for name in [*node.input, *node.output]
+        )
+    }
+    version = next((op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1)
+    if _operator_nodes(model, {"QuantizeLinear", "DequantizeLinear"}):
+        version = max(version, 19)
+    pending = list(_subgraph_captures(model.graph))
+    schemas, required, outputs = {}, set(), set()
+    for node in model.graph.node:
+        if id(node) in native:
+            pending.extend([*node.input, *node.output])
+        if id(node) not in standard:
+            continue
+        schema_version = version
+        if id(node) in contrib:
+            # Contrib v1 supports HALF without upgrading the standard opset.
+            schema_version = 20 if node.op_type == "Gelu" else 19
+        key = node.op_type, schema_version
+        if key not in schemas:
+            schema = defs.get_schema(node.op_type, schema_version)
+            only_fp32 = {
+                item.type_param_str
+                for item in schema.type_constraints
+                if "tensor(float)" in item.allowed_type_strs
+                and "tensor(float16)" not in item.allowed_type_strs
+            }
+            schemas[key] = schema, only_fp32 | {"tensor(float)"}
+        schema, only_fp32 = schemas[key]
+        for port, name in enumerate(node.input):
+            if name and (
+                schema.inputs[min(port, len(schema.inputs) - 1)].type_str in only_fp32
+                or (node.op_type == "Resize" and port in (1, 2))
+            ):
+                required.add((id(node), port))
+                pending.append(name)
+        float_outputs = {name for name in node.output if cache.get_dtype(name) == TensorProto.FLOAT}
+        dtype = _attribute(node, "dtype")
+        if (dtype is not None and dtype.i == TensorProto.FLOAT) or (
+            node.op_type != "Cast"
+            and not any(cache.get_dtype(name) == TensorProto.FLOAT for name in node.input)
+        ):
+            outputs.update(float_outputs)
+        outputs.update(
+            name
+            for port, name in enumerate(node.output)
+            if name in float_outputs
+            and schema.outputs[min(port, len(schema.outputs) - 1)].type_str in only_fp32
+        )
+    values = set()
+    while pending:
+        name = pending.pop()
+        if name in values or cache.get_dtype(name) != TensorProto.FLOAT:
+            continue
+        values.add(name)
+        producer = cache.get_producer(name)
+        if producer is not None:
+            native.add(id(producer))
+            pending.extend([*producer.input, *producer.output])
+    return values | outputs, native, required
+
+
+def _upgrade_fp16_quantizers(model, cache):
+    """Migrate inherited reduction schemas when FP16 quantizers require opset 19."""
+    if not _operator_nodes(model, {"QuantizeLinear", "DequantizeLinear"}):
+        return
+    opset = next(op for op in model.opset_import if op.domain in {"", "ai.onnx"})
+    if opset.version >= 19:
+        return
+    functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions}
+    pending = [model.graph]
+    while pending:
+        graph = pending.pop()
+        allocator = cache if graph is model.graph else GraphCache(graph)
+        allocator._names.update(_subgraph_captures(graph))
+        for node in graph.node:
+            for attr in node.attribute:
+                if attr.type == AttributeProto.GRAPH:
+                    pending.append(attr.g)
+                elif attr.type == AttributeProto.GRAPHS:
+                    pending.extend(attr.graphs)
+            if (
+                node.domain not in {"", "ai.onnx"}
+                or node.overload
+                or (node.domain, node.op_type, node.overload) in functions
+            ):
+                continue
+            axes = _attribute(node, "axes")
+            if axes is None or not node.op_type.startswith("Reduce"):
+                continue
+            if any(value.name == "axes" for value in defs.get_schema(node.op_type, 19).inputs):
+                control = allocator.unique_name((node.name or node.output[0]) + "_axes")
+                length = len(axes.ints)
+                graph.initializer.append(
+                    numpy_helper.from_array(np.asarray(axes.ints, np.int64), control)
+                )
+                node.input.append(control)
+                node.attribute.remove(axes)
+                if model.ir_version < 4:
+                    graph.input.append(
+                        helper.make_tensor_value_info(control, TensorProto.INT64, [length])
+                    )
+    opset.version = 19
+
+
 def _apply_force_fp16_activations(model):
     """Run after the ordered DLA passes and their final symbolic shape inference."""
     graph = model.graph
     cache = GraphCache(graph)
+    opset = next((op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1)
+    if opset < 13 and _operator_nodes(model, {"QuantizeLinear", "DequantizeLinear"}):
+        raise ValueError("FP16 quantizer migration requires ONNX opset 13 or newer")
     fp32, fp16 = TensorProto.FLOAT, TensorProto.FLOAT16
+    runtime = _runtime_input_names(model)
+    native_values, native_nodes, fp32_inputs = _fp32_regions(model, cache)
     preserved_inits = {
         name
         for name in cache.init_names
-        if any(n.op_type == "Resize" and port in (1, 2) for n, port in cache.get_consumers(name))
+        if name in runtime
+        or name in native_values
+        or any(
+            _fp32_control(node, port, native_nodes, fp32_inputs)
+            for node, port in cache.get_consumers(name)
+        )
     }
-    excluded_inits = [init for init in graph.initializer if init.name in preserved_inits]
-    for init in excluded_inits:
-        graph.initializer.remove(init)
-    try:
-        convert_initializers_to_dtype(graph, "float16")
-    finally:
-        graph.initializer.extend(excluded_inits)
-    # Conversion replaces initializer objects, invalidating the cached references.
-    cache = GraphCache(graph)
-    preserved = set()
+    _validate_float_weights(graph)
+    if any(
+        tensor.data_type == fp32 and tensor.name not in preserved_inits
+        for tensor in graph.initializer
+    ):
+        excluded = [
+            tensor
+            for tensor in graph.initializer
+            if tensor.data_type == fp32 and tensor.name in preserved_inits
+        ]
+        for tensor in excluded:
+            graph.initializer.remove(tensor)
+        try:
+            convert_initializers_to_dtype(graph, "float16")
+        finally:
+            graph.initializer.extend(excluded)
+        cache = GraphCache(graph)
+    preserved = set(native_values)
     terminal_concats = set()
     for node in graph.node:
         if not _is_cast(node, fp32) or cache.get_dtype(node.input[0]) in (fp32, fp16, None):
@@ -155,22 +357,50 @@ def _apply_force_fp16_activations(model):
             value.type.tensor_type.elem_type = fp16
     before, after = {}, {}
     front, back = [], []
+    for tensor in graph.initializer:
+        if tensor.name not in preserved_inits - runtime or tensor.data_type != fp32:
+            continue
+        consumers = [
+            (node, port)
+            for node, port in cache.get_consumers(tensor.name)
+            if not _fp32_control(node, port, native_nodes, fp32_inputs) and not _is_cast(node, fp16)
+        ]
+        if consumers:
+            output = cache.unique_name(tensor.name + "_fp16")
+            front.append(_cast(tensor.name, output, fp16, cache))
+            for node, port in consumers:
+                node.input[port] = output
+    for node in graph.node:
+        for name in node.output:
+            if name not in native_values:
+                continue
+            consumers = [
+                (consumer, port)
+                for consumer, port in cache.get_consumers(name)
+                if not _fp32_control(consumer, port, native_nodes, fp32_inputs)
+                and not _is_cast(consumer, fp16)
+            ]
+            if consumers:
+                output = cache.unique_name(name + "_fp16")
+                after.setdefault(id(node), []).append(_cast(name, output, fp16, cache))
+                for consumer, port in consumers:
+                    consumer.input[port] = output
     for node in graph.node:
         if id(node) not in terminal_concats:
             continue
         for port, source in enumerate(node.input):
             if source not in preserved:
                 output = cache.unique_name(source + "_concat_fp32")
-                before.setdefault(id(node), []).append(_cast(source, output, fp32))
+                before.setdefault(id(node), []).append(_cast(source, output, fp32, cache))
                 node.input[port] = output
                 preserved.add(output)
     for value in graph.input:
-        if value.name in cache.init_names or value.type.tensor_type.elem_type != fp32:
+        if value.name in cache.init_names - runtime or value.type.tensor_type.elem_type != fp32:
             continue
         direct = None
         rank = len(value.type.tensor_type.shape.dim)
         for consumer, port in cache.get_consumers(value.name):
-            if _is_cast(consumer, fp16):
+            if _fp32_control(consumer, port, native_nodes, fp32_inputs) or _is_cast(consumer, fp16):
                 continue
             shape_info = cache.vi_map.get(consumer.output[0])
             adapter = (
@@ -183,7 +413,7 @@ def _apply_force_fp16_activations(model):
                 assert shape_info is not None
                 source = consumer.output[0]
                 output = cache.unique_name(source + "_fp16")
-                after.setdefault(id(consumer), []).append(_cast(source, output, fp16))
+                after.setdefault(id(consumer), []).append(_cast(source, output, fp16, cache))
                 for downstream, slot in cache.get_consumers(source):
                     if not _is_cast(downstream, fp16):
                         downstream.input[slot] = output
@@ -191,26 +421,18 @@ def _apply_force_fp16_activations(model):
             else:
                 if direct is None:
                     direct = cache.unique_name(value.name + "_fp16")
-                    front.append(_cast(value.name, direct, fp16))
+                    front.append(_cast(value.name, direct, fp16, cache))
                 consumer.input[port] = direct
-    if any(
-        n.op_type in {"QuantizeLinear", "DequantizeLinear"} and n.domain in ("", "ai.onnx")
-        for n in graph.node
-    ):
-        opsets = [o for o in model.opset_import if o.domain in ("", "ai.onnx")]
-        if opsets:
-            opsets[0].version = max(19, opsets[0].version)
-        else:
-            model.opset_import.append(helper.make_opsetid("", 19))
+    _upgrade_fp16_quantizers(model, cache)
     for value in graph.output:
-        if value.type.tensor_type.elem_type != fp32:
+        if value.type.tensor_type.elem_type != fp32 or value.name in native_values:
             continue
         producer = cache.get_producer(value.name)
         if producer is not None and producer.op_type == "Squeeze":
             source = producer.input[0]
             if source not in preserved:
                 output = cache.unique_name(source + "_to_fp32_fp16")
-                before.setdefault(id(producer), []).append(_cast(source, output, fp32))
+                before.setdefault(id(producer), []).append(_cast(source, output, fp32, cache))
                 producer.input[0] = output
             continue
         if _is_cast(producer, fp32):
@@ -222,6 +444,10 @@ def _apply_force_fp16_activations(model):
         if producer is None:
             assert tensor is not None
             tensor.name = internal
+            if model.ir_version < 4:
+                for declaration in graph.input:
+                    if declaration.name == value.name:
+                        declaration.name = internal
         else:
             for port, name in enumerate(producer.output):
                 if name == value.name:
@@ -233,26 +459,33 @@ def _apply_force_fp16_activations(model):
                 annotation.type.tensor_type.elem_type = fp16
         for consumer, port in cache.get_consumers(value.name):
             consumer.input[port] = internal
-        back.append(_cast(internal, value.name, fp32))
+        back.append(_cast(internal, value.name, fp32, cache))
     nodes = list(front)
     for node in graph.node:
         nodes.extend(before.get(id(node), []))
         nodes.append(node)
         nodes.extend(after.get(id(node), []))
     nodes.extend(back)
-    del graph.node[:]
-    graph.node.extend(nodes)
-    _rewrite_terminal_projections(graph)
+    _insert_support_nodes(graph, nodes)
+    _rewrite_terminal_projections(model)
+    _add_legacy_initializer_inputs(model)
     return model
 
 
-def _rewrite_terminal_projections(graph):
-    """Bound terminal projection channels to DLA's 8192-element dimension limit."""
+def _rewrite_terminal_projections(model):
+    """Bound private fixed projection channels while preserving native owners."""
+    graph = model.graph
     cache = GraphCache(graph)
+    standard = {
+        id(node) for node in _fp16_operator_nodes(model, {node.op_type for node in graph.node})
+    }
+    protected = cache.graph_output_names | _subgraph_captures(graph)
+    runtime = _runtime_input_names(model)
     removed, removed_values, removed_inits = set(), set(), set()
     replacements, additions = {}, []
 
     def attributes(node, key, expected):
+        """Match one fixed operator attribute used by the projection pattern."""
         value = _attribute(node, key)
         return (
             value is not None
@@ -260,6 +493,7 @@ def attributes(node, key, expected):
         )
 
     def axes(node, expected):
+        """Match fixed axes without changing their initializer owner."""
         value = cache.get_init_array(node.input[1]) if len(node.input) == 2 else None
         return value is not None and value.reshape(-1).tolist() == expected
 
@@ -270,7 +504,7 @@ def axes(node, expected):
         current = output.name
         for op in ("Squeeze", "Cast", "Concat", "ReduceSum", "Concat", "Transpose", "Conv"):
             node = cache.get_producer(current)
-            if node is None or node.op_type != op:
+            if node is None or id(node) not in standard or node.op_type != op:
                 break
             chain.append(node)
             current = node.input[0]
@@ -286,6 +520,7 @@ def axes(node, expected):
             and attributes(transpose, "perm", [0, 3, 2, 1])
             and len(conv.input) == 2
             and len(conv.output) == 1
+            and ((_attribute(conv, "group") is None) or _attribute(conv, "group").i == 1)
         ):
             continue
         weight = cache.get_init(conv.input[1])
@@ -294,6 +529,7 @@ def axes(node, expected):
             parameters = [cache.get_init(name) for name in dq.input[1:]]
             if (
                 dq.op_type != "DequantizeLinear"
+                or id(dq) not in standard
                 or len(dq.input) < 2
                 or any(
                     parameter is None or math.prod(parameter.dims) != 1 for parameter in parameters
@@ -312,6 +548,12 @@ def axes(node, expected):
         if weight.data_type != (TensorProto.UINT8 if dq is not None else TensorProto.FLOAT16):
             continue
         chain_ids = {id(n) for n in chain}
+        if any(
+            name in protected | runtime for name in (weight.name, squeeze.input[1], reduce.input[1])
+        ):
+            continue
+        if any(name in protected for item in chain[1:] for name in item.output):
+            continue
         if any(
             id(consumer) not in chain_ids
             for name in (weight.name, squeeze.input[1], reduce.input[1])
@@ -324,7 +566,8 @@ def axes(node, expected):
         if cache.get_shape(output.name) != [1, 1, channels]:
             continue
         bytes_per_channel = math.prod(weight.dims[1:]) * (1 if dq is not None else 2)
-        if bytes_per_channel <= 0 or len(weight.raw_data) != channels * bytes_per_channel:
+        payload = weight.raw_data
+        if bytes_per_channel <= 0 or len(payload) != channels * bytes_per_channel:
             continue
         nodes, outputs = [], []
         for index, start in enumerate(range(0, channels, 8192)):
@@ -335,7 +578,7 @@ def axes(node, expected):
                 data_type=weight.data_type,
                 dims=[count, *weight.dims[1:]],
             )
-            chunk.raw_data = weight.raw_data[
+            chunk.raw_data = payload[
                 start * bytes_per_channel : (start + count) * bytes_per_channel
             ]
             additions.append(chunk)
@@ -364,7 +607,7 @@ def axes(node, expected):
                 )
             )
             result = cache.unique_name(conv.output[0] + suffix + "_fp32")
-            nodes.append(_cast(bounded, result, TensorProto.FLOAT))
+            nodes.append(_cast(bounded, result, TensorProto.FLOAT, cache))
             outputs.append(result)
         nodes.append(
             helper.make_node(
@@ -381,16 +624,22 @@ def axes(node, expected):
         removed_inits.update([weight.name, squeeze.input[1], reduce.input[1]])
     if not replacements:
         return
-    nodes = []
-    for node in graph.node:
-        nodes.extend(replacements.get(id(node), []))
-        if id(node) not in removed:
-            nodes.append(node)
-    retained_inits = [t for t in graph.initializer if t.name not in removed_inits]
-    retained_info = [v for v in graph.value_info if v.name not in removed_values]
-    del graph.node[:]
-    graph.node.extend(nodes)
-    del graph.initializer[:]
-    graph.initializer.extend([*retained_inits, *additions])
-    del graph.value_info[:]
-    graph.value_info.extend(retained_info)
+    batch_replace_nodes(
+        graph,
+        {
+            index: ([node], replacements.get(id(node), []))
+            for index, node in enumerate(graph.node)
+            if id(node) in removed
+        },
+    )
+    for index in range(len(graph.initializer) - 1, -1, -1):
+        if graph.initializer[index].name in removed_inits:
+            del graph.initializer[index]
+    graph.initializer.extend(additions)
+    for index in range(len(graph.value_info) - 1, -1, -1):
+        if graph.value_info[index].name in removed_values:
+            del graph.value_info[index]
+    if model.ir_version < 4:
+        for index in range(len(graph.input) - 1, -1, -1):
+            if graph.input[index].name in removed_inits:
+                del graph.input[index]
diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py
index 71e4ba95875..920c7286246 100644
--- a/modelopt/onnx/graph_surgery/make_dla_compatible.py
+++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py
@@ -363,7 +363,8 @@ def dla_make_dla_compatible(
         qdq_quantized_dtypes: Zero-point dtypes whose QDQ pairs are removed.
             Defaults to UINT16 and INT16. Accepts ONNX dtype IDs or names.
         remove_all_qdq: Remove QDQ pairs for every dtype, matching the C++ flag.
-        force_fp16_activations: Use FP16 internals while preserving FP32 public boundaries.
+        force_fp16_activations: Use FP16 where supported, keeping FP32 public
+            boundaries and operations or controls that require FP32.
             Cannot be combined with qdq_quantized_dtypes.
 
     Returns:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index 40a1dc5b570..37d1d34ad54 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -32,10 +32,12 @@
 
 
 def _value(name, shape, dtype=T.FLOAT):
+    """Declare a typed tensor for a checked precision fixture."""
     return h.make_tensor_value_info(name, dtype, shape)
 
 
 def _model(nodes, inputs, outputs, initializers, opset=19):
+    """Build a model with inferred intermediate types for a precision fixture."""
     model = h.make_model(
         h.make_graph(nodes, "fp16", inputs, outputs, initializers),
         opset_imports=[h.make_opsetid("", opset)],
@@ -46,6 +48,7 @@ def _model(nodes, inputs, outputs, initializers, opset=19):
 
 @pytest.mark.parametrize("enabled", [False, True])
 def test_public_option_preserves_boundaries_and_selects_weight_precision(tmp_path, enabled):
+    """Keep public FLOAT interfaces while selecting internal weight precision."""
     values = np.array([1.0003, 0.12345], np.float32)
     model = _model(
         [h.make_node("Mul", ["x", "w"], ["y"])],
@@ -74,6 +77,7 @@ def test_public_option_preserves_boundaries_and_selects_weight_precision(tmp_pat
 
 
 def test_initializer_rounding_preserves_subnormals_signed_zero_and_nan():
+    """Retain the original converter rounding, signed zero and NaN behavior."""
     values = np.array([0, -0.0, 2**-24, 2**-25, 2**-14, 65504, 65520, 0, 0], np.float32)
     values.view(np.uint32)[-2:] = [0x7F812345, 0xFF812345]
     model = _model([], [], [_value("weight", [len(values)])], [nh.from_array(values, "weight")])
@@ -86,6 +90,7 @@ def test_initializer_rounding_preserves_subnormals_signed_zero_and_nan():
 
 
 def test_boundary_adapters_and_direct_input_consumers_get_their_own_casts():
+    """Give boundary adapters and direct input uses their own typed views."""
     model = _model(
         [
             h.make_node("Unsqueeze", ["x", "axes"], ["x4"]),
@@ -112,6 +117,7 @@ def test_boundary_adapters_and_direct_input_consumers_get_their_own_casts():
 
 @pytest.mark.parametrize("operation", ["Resize", "QuantizeLinear"])
 def test_schema_specific_float_parameters_and_qdq_opset(operation):
+    """Keep required FLOAT controls and migrate FP16 quantizer schemas."""
     if operation == "Resize":
         nodes = [
             h.make_node("Mul", ["x", "weight"], ["weighted"]),
@@ -146,6 +152,7 @@ def test_schema_specific_float_parameters_and_qdq_opset(operation):
 
 
 def test_scaled_l2_normalization_keeps_zero_slices_finite():
+    """Stabilize zero-vector normalization before applying FP16 precision."""
     model = _model(
         [
             h.make_node("LpNormalization", ["x"], ["norm"], p=2, axis=-1),
@@ -165,6 +172,7 @@ def test_scaled_l2_normalization_keeps_zero_slices_finite():
 
 
 def test_terminal_projection_splits_channels_and_preserves_values():
+    """Bound private projection channels while preserving real output values."""
     channels = 8193
     model = _model(
         [
@@ -214,3 +222,610 @@ def test_excluded_initializers_are_restored_when_conversion_fails():
         _apply_force_fp16_activations(model)
     actual = {value.name: value.SerializeToString() for value in model.graph.initializer}
     assert all(actual[name] == payload for name, payload in expected.items())
+
+
+def _checked_cpu_outputs(model, feeds):
+    """Execute a fully checked graph with bounded CPU runtime threading."""
+    onnx.checker.check_model(model, full_check=True)
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
+    return ort.InferenceSession(
+        model.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, feeds)
+
+
+def test_fp16_runtime_initializer_default_retains_fp32_public_contract():
+    """Keep FP32 default-input declarations and accept FP32 runtime overrides."""
+    model = _model(
+        [h.make_node("Mul", ["x", "w"], ["y"])],
+        [_value("x", [2]), _value("w", [2])],
+        [_value("y", [2])],
+        [nh.from_array(np.asarray([0.5, 2], np.float32), "w")],
+    )
+    before = model.graph.initializer[0].SerializeToString()
+    cases = [
+        {"x": np.asarray([2, 4], np.float32)},
+        {"x": np.asarray([2, 4], np.float32), "w": np.asarray([2, 0.25], np.float32)},
+    ]
+    expected = [_checked_cpu_outputs(model, feeds)[0] for feeds in cases]
+    changed = _apply_force_fp16_activations(model)
+    assert (
+        next(value for value in changed.graph.input if value.name == "w").type.tensor_type.elem_type
+        == T.FLOAT
+    )
+    assert (
+        next(value for value in changed.graph.initializer if value.name == "w").SerializeToString()
+        == before
+    )
+    for feeds, wanted in zip(cases, expected):
+        np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], wanted)
+
+
+def test_fp16_resize_parameter_shared_with_arithmetic_uses_private_precision_view():
+    """Keep required FP32 scales while providing a separate FP16 arithmetic operand."""
+    model = _model(
+        [
+            h.make_node("Mul", ["x", "scales"], ["weighted"]),
+            h.make_node("Resize", ["weighted", "roi", "scales"], ["y"]),
+        ],
+        [_value("x", [1, 1, 1, 4])],
+        [_value("y", [1, 1, 1, 8])],
+        [
+            nh.from_array(np.asarray([1, 1, 1, 2], np.float32), "scales"),
+            nh.from_array(np.asarray([], np.float32), "roi"),
+        ],
+    )
+    feeds = {"x": np.asarray([[[[2, 4, 6, 8]]]], np.float32)}
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    changed = _apply_force_fp16_activations(model)
+    np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected)
+
+
+def test_fp16_conversion_preserves_unrelated_constant_object():
+    """Avoid copying an unrelated integer Constant while inserting floating boundaries."""
+    constant = h.make_node(
+        "Constant", [], ["integer_payload"], value=nh.from_array(np.arange(64, dtype=np.uint8))
+    )
+    model = _model(
+        [constant, h.make_node("Relu", ["x"], ["y"])], [_value("x", [2])], [_value("y", [2])], []
+    )
+    constant = model.graph.node[0]
+    before = constant.SerializeToString()
+    changed = _apply_force_fp16_activations(model)
+    assert next(node for node in changed.graph.node if node.op_type == "Constant") is constant
+    assert constant.SerializeToString() == before
+    _checked_cpu_outputs(changed, {"x": np.asarray([-1, 2], np.float32)})
+
+
+def test_fp16_malformed_late_weight_refuses_before_any_mutation():
+    """Validate every floating payload before converting an earlier fixed initializer."""
+    model = h.make_model(
+        h.make_graph(
+            [h.make_node("Mul", ["x", "valid"], ["y"])],
+            "invalid_late_weight",
+            [_value("x", [2])],
+            [_value("y", [2])],
+            [
+                nh.from_array(np.asarray([1, 2], np.float32), "valid"),
+                T(name="invalid", data_type=T.FLOAT, dims=[2], float_data=[1]),
+            ],
+        ),
+        opset_imports=[h.make_opsetid("", 19)],
+    )
+    model.ir_version = 9
+    before = model.SerializeToString()
+    with pytest.raises(ValueError):
+        _apply_force_fp16_activations(model)
+    assert model.SerializeToString() == before
+
+
+def test_fp16_qdq_upgrade_keeps_legacy_reduce_axes_valid():
+    """Migrate schema-sensitive reduction axes when enabling FP16 quantizers."""
+    model = _model(
+        [
+            h.make_node("ReduceMean", ["x"], ["mean"], axes=[1], keepdims=1),
+            h.make_node("QuantizeLinear", ["mean", "scale", "zero"], ["q"]),
+            h.make_node("DequantizeLinear", ["q", "scale", "zero"], ["y"]),
+        ],
+        [_value("x", [2, 2])],
+        [_value("y", [2, 1])],
+        [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")],
+        opset=17,
+    )
+    onnx.checker.check_model(model, full_check=True)
+    feeds = {"x": np.asarray([[0.2, 0.2], [0.4, 0.4]], np.float32)}
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    changed = _apply_force_fp16_activations(model)
+    onnx.checker.check_model(changed, full_check=True)
+    np.testing.assert_allclose(
+        ReferenceEvaluator(changed).run(None, feeds)[0], expected, rtol=2e-3, atol=1e-4
+    )
+
+
+def test_fp16_lpnorm_keeps_opaque_function_native():
+    """Do not interpret an opaque LpNormalization name as standard normalization."""
+    model = _model(
+        [
+            h.make_node("LpNormalization", ["x"], ["norm"], p=2, axis=-1),
+            h.make_node("Mul", ["norm", "gamma"], ["y"]),
+        ],
+        [_value("x", [1, 1, 1, 2])],
+        [_value("y", [1, 1, 1, 2])],
+        [nh.from_array(np.ones([1, 1, 1, 2], np.float32), "gamma")],
+    )
+    model.graph.node[0].domain = "opaque"
+    model.functions.append(
+        h.make_function(
+            "opaque",
+            "LpNormalization",
+            ["x"],
+            ["y"],
+            [h.make_node("Neg", ["x"], ["y"])],
+            [h.make_opsetid("", 19)],
+            attributes=["p", "axis"],
+        )
+    )
+    model.opset_import.append(h.make_opsetid("opaque", 1))
+    before = model.SerializeToString()
+    changed = _apply_fp16_lpnorm(model)
+    assert changed.SerializeToString() == before
+    _checked_cpu_outputs(changed, {"x": np.asarray([[[[3, 4]]]], np.float32)})
+
+
+def test_fp16_boundary_cast_names_are_allocated():
+    """Keep generated Cast node names distinct from existing nodes and tensor names."""
+    model = _model(
+        [
+            h.make_node(
+                "Constant",
+                [],
+                ["integer"],
+                name="x_to_fp16",
+                value=nh.from_array(np.asarray([1], np.int8)),
+            ),
+            h.make_node("Relu", ["x"], ["y"]),
+        ],
+        [_value("x", [2])],
+        [_value("y", [2])],
+        [],
+    )
+    changed = _apply_force_fp16_activations(model)
+    names = [node.name for node in changed.graph.node if node.name]
+    assert len(names) == len(set(names))
+    _checked_cpu_outputs(changed, {"x": np.asarray([-1, 2], np.float32)})
+
+
+def test_fp16_lpnorm_keeps_unrelated_constant_object():
+    """Retain a large-payload owner while replacing a private normalization node."""
+    model = _model(
+        [
+            h.make_node(
+                "Constant", [], ["integer"], value=nh.from_array(np.arange(64, dtype=np.uint8))
+            ),
+            h.make_node("LpNormalization", ["x"], ["norm"], p=2, axis=-1),
+            h.make_node("Mul", ["norm", "gamma"], ["y"]),
+        ],
+        [_value("x", [1, 1, 1, 2])],
+        [_value("y", [1, 1, 1, 2])],
+        [nh.from_array(np.ones([1, 1, 1, 2], np.float32), "gamma")],
+    )
+    constant = model.graph.node[0]
+    changed = _apply_fp16_lpnorm(model)
+    assert next(node for node in changed.graph.node if node.op_type == "Constant") is constant
+    onnx.checker.check_model(changed, full_check=True)
+
+
+def test_fp16_keeps_opaque_float_function_type_contract():
+    """Retain a local function's FP32 region and cast only its downstream standard use."""
+    model = _model(
+        [h.make_node("Identity", ["x"], ["middle"]), h.make_node("Relu", ["middle"], ["y"])],
+        [_value("x", [2])],
+        [_value("y", [2])],
+        [],
+    )
+    model.graph.node[0].domain = "opaque"
+    model.functions.append(
+        h.make_function(
+            "opaque",
+            "Identity",
+            ["x"],
+            ["out"],
+            [
+                h.make_node("Constant", [], ["bias"], value=nh.from_array(np.float32(0.25))),
+                h.make_node("Add", ["x", "bias"], ["out"]),
+            ],
+            [h.make_opsetid("", 19)],
+        )
+    )
+    model.opset_import.append(h.make_opsetid("opaque", 1))
+    feeds = {"x": np.asarray([0.5, -1], np.float32)}
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    changed = _apply_force_fp16_activations(model)
+    np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected)
+
+
+def test_fp16_retains_captured_values_and_control_flow_types():
+    """Keep a captured FP32 result available to unchanged nested graph bodies."""
+    branch = h.make_graph(
+        [h.make_node("Identity", ["middle"], ["branch_value"])],
+        "capture",
+        [],
+        [_value("branch_value", [2])],
+    )
+    model = _model(
+        [
+            h.make_node("Relu", ["x"], ["middle"]),
+            h.make_node("Neg", ["middle"], ["y"]),
+            h.make_node("If", ["condition"], ["captured"], then_branch=branch, else_branch=branch),
+        ],
+        [_value("x", [2]), _value("condition", [], T.BOOL)],
+        [_value("y", [2]), _value("captured", [2])],
+        [],
+    )
+    feeds = {"x": np.asarray([1, -2], np.float32), "condition": np.asarray(True)}
+    expected = _checked_cpu_outputs(model, feeds)
+    changed = _apply_force_fp16_activations(model)
+    for actual, wanted in zip(_checked_cpu_outputs(changed, feeds), expected):
+        np.testing.assert_array_equal(actual, wanted)
+
+
+def test_fp16_float_constant_keeps_valid_consumer_types():
+    """Give an unfurled FLOAT Constant a private FP16 data view."""
+    model = _model(
+        [
+            h.make_node(
+                "Constant", [], ["bias"], value=nh.from_array(np.asarray([0.5, 0.25], np.float32))
+            ),
+            h.make_node("Add", ["x", "bias"], ["y"]),
+        ],
+        [_value("x", [2])],
+        [_value("y", [2])],
+        [],
+    )
+    feeds = {"x": np.asarray([1, 2], np.float32)}
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    changed = _apply_force_fp16_activations(model)
+    np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected)
+
+
+def _terminal_model(channels=8193, *, groups=1, quantized=False):
+    """Build the exact terminal projection pattern with independently checkable values."""
+    data_channels = groups
+    weights = nh.from_array(
+        np.ones((channels, 1, 1, 1), np.uint8)
+        if quantized
+        else np.full((channels, 1, 1, 1), 0.25, np.float32),
+        "weight",
+    )
+    nodes = []
+    initializers = [weights, nh.from_array(np.asarray([0], np.int64), "axes")]
+    if quantized:
+        initializers.extend(
+            [nh.from_array(np.float32(0.25), "scale"), nh.from_array(np.uint8(0), "zero")]
+        )
+        nodes.append(h.make_node("DequantizeLinear", ["weight", "scale", "zero"], ["dequantized"]))
+    nodes.extend(
+        [
+            h.make_node(
+                "Conv", ["x", "dequantized" if quantized else "weight"], ["conv"], group=groups
+            ),
+            h.make_node("Transpose", ["conv"], ["transposed"], perm=[0, 3, 2, 1]),
+            h.make_node("Concat", ["transposed"], ["concat1"], axis=0),
+            h.make_node("ReduceSum", ["concat1", "axes"], ["reduced"], keepdims=1),
+            h.make_node("Concat", ["reduced"], ["concat2"], axis=2),
+            h.make_node("Squeeze", ["concat2", "axes"], ["y"]),
+        ]
+    )
+    return _model(
+        nodes,
+        [_value("x", [1, data_channels, 1, 1])],
+        [_value("y", [1, 1, channels])],
+        initializers,
+    )
+
+
+@pytest.mark.parametrize(
+    "owner", ["grouped", "runtime_weight", "public_axes", "legacy", "constant"]
+)
+def test_terminal_projection_preserves_native_owners_and_group_layout(owner):
+    """Chunk only private fixed projections and preserve control owners and grouped channels."""
+    grouped = owner == "grouped"
+    model = _terminal_model(
+        channels=8194 if grouped else 8193,
+        groups=2 if grouped else 1,
+        quantized=owner == "runtime_weight",
+    )
+    if owner == "runtime_weight":
+        model.graph.input.append(_value("weight", [8193, 1, 1, 1], T.UINT8))
+    if owner == "public_axes":
+        model.graph.output.append(_value("axes", [1], T.INT64))
+    if owner == "legacy":
+        model.ir_version = 3
+        model.graph.input.extend(
+            _value(t.name, list(t.dims), t.data_type) for t in model.graph.initializer
+        )
+    constant = None
+    if owner == "constant":
+        model.graph.node.insert(
+            0,
+            h.make_node(
+                "Constant", [], ["integer"], value=nh.from_array(np.arange(64, dtype=np.uint8))
+            ),
+        )
+        constant = model.graph.node[0]
+    data = (
+        np.asarray([[[[2]], [[8]]]], np.float32) if grouped else np.asarray([[[[2]]]], np.float32)
+    )
+    feeds = {"x": data}
+    onnx.checker.check_model(model, full_check=True)
+    expected = ReferenceEvaluator(model).run(None, feeds)
+    override = (
+        {"x": data, "weight": np.full((8193, 1, 1, 1), 2, np.uint8)}
+        if owner == "runtime_weight"
+        else None
+    )
+    override_expected = (
+        ReferenceEvaluator(model).run(None, override) if override is not None else None
+    )
+    changed = _apply_force_fp16_activations(model)
+    if owner in {"grouped", "runtime_weight", "public_axes"}:
+        assert sum(node.op_type == "Conv" for node in changed.graph.node) == 1
+    if owner == "legacy":
+        assert not {"weight", "axes"} & {value.name for value in changed.graph.input}
+    if constant is not None:
+        assert next(node for node in changed.graph.node if node.op_type == "Constant") is constant
+    onnx.checker.check_model(changed, full_check=True)
+    for actual, wanted in zip(ReferenceEvaluator(changed).run(None, feeds), expected):
+        np.testing.assert_array_equal(actual, wanted)
+    if override is not None:
+        for actual, wanted in zip(
+            ReferenceEvaluator(changed).run(None, override), override_expected
+        ):
+            np.testing.assert_array_equal(actual, wanted)
+
+
+def test_fp16_public_legacy_initializer_uses_private_declaration():
+    """Rename the legacy initializer declaration together with its private FP16 payload."""
+    model = _model(
+        [],
+        [_value("weight", [2])],
+        [_value("weight", [2])],
+        [nh.from_array(np.asarray([1, 2], np.float32), "weight")],
+    )
+    model.ir_version = 3
+    changed = _apply_force_fp16_activations(model)
+    onnx.checker.check_model(changed, full_check=True)
+    np.testing.assert_array_equal(
+        _checked_cpu_outputs(changed, {})[0], np.asarray([1, 2], np.float32)
+    )
+
+
+@pytest.mark.parametrize("legacy", [False, True])
+def test_fp16_qdq_upgrade_migrates_nested_reduce_axes(legacy):
+    """Migrate axes in nested graphs that inherit the upgraded main operator set."""
+    branch = h.make_graph(
+        [h.make_node("ReduceMean", ["x"], ["branch_value"], axes=[1], keepdims=1)],
+        "nested_reduce",
+        [],
+        [_value("branch_value", [2, 1])],
+    )
+    model = _model(
+        [
+            h.make_node("QuantizeLinear", ["x", "scale", "zero"], ["q"]),
+            h.make_node("DequantizeLinear", ["q", "scale", "zero"], ["y"]),
+            h.make_node("If", ["condition"], ["reduced"], then_branch=branch, else_branch=branch),
+        ],
+        [_value("x", [2, 2]), _value("condition", [], T.BOOL)],
+        [_value("y", [2, 2]), _value("reduced", [2, 1])],
+        [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")],
+        opset=17,
+    )
+    if legacy:
+        model.ir_version = 3
+        model.graph.input.extend(
+            _value(t.name, list(t.dims), t.data_type) for t in model.graph.initializer
+        )
+    feeds = {"x": np.asarray([[0.2, 0.2], [0.4, 0.4]], np.float32), "condition": np.asarray(True)}
+    expected = _checked_cpu_outputs(model, feeds)
+    changed = _apply_force_fp16_activations(model)
+    onnx.checker.check_model(changed, full_check=True)
+    for got, wanted in zip(ReferenceEvaluator(changed).run(None, feeds), expected):
+        np.testing.assert_allclose(got, wanted, rtol=2e-3, atol=1e-4)
+
+
+def test_fp16_legacy_qdq_schema_refuses_before_mutation():
+    """Refuse legacy operator-set migration that needs adapters beyond the supported range."""
+    model = _model(
+        [
+            h.make_node("Clip", ["x"], ["clipped"], min=0.0, max=1.0),
+            h.make_node("QuantizeLinear", ["clipped", "scale", "zero"], ["y"]),
+        ],
+        [_value("x", [2])],
+        [_value("y", [2], T.INT8)],
+        [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")],
+        opset=10,
+    )
+    onnx.checker.check_model(model, full_check=True)
+    before = model.SerializeToString()
+    with pytest.raises(ValueError, match="opset"):
+        _apply_force_fp16_activations(model)
+    assert model.SerializeToString() == before
+
+
+@pytest.mark.parametrize("operation", ["Range", "NonMaxSuppression"])
+def test_fp16_preserves_operator_schemas_without_half_inputs(operation):
+    """Keep mandatory FP32 schema bindings valid instead of forcing unsupported HALF inputs."""
+    if operation == "Range":
+        model = _model(
+            [h.make_node("Range", ["start", "limit", "delta"], ["y"])],
+            [_value("start", [])],
+            [_value("y", [4])],
+            [nh.from_array(np.float32(4), "limit"), nh.from_array(np.float32(1), "delta")],
+        )
+        feeds = {"start": np.asarray(0, np.float32)}
+    else:
+        model = _model(
+            [h.make_node("NonMaxSuppression", ["boxes", "scores", "count", "iou"], ["y"])],
+            [_value("boxes", [1, 2, 4]), _value("scores", [1, 1, 2])],
+            [_value("y", [None, 3], T.INT64)],
+            [nh.from_array(np.int64(1), "count"), nh.from_array(np.float32(0.5), "iou")],
+        )
+        feeds = {
+            "boxes": np.asarray([[[0, 0, 1, 1], [2, 2, 3, 3]]], np.float32),
+            "scores": np.asarray([[[0.8, 0.5]]], np.float32),
+        }
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    changed = _apply_force_fp16_activations(model)
+    np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected)
+
+
+@pytest.mark.parametrize("operation", ["ConstantOfShape", "RandomNormal", "EyeLike"])
+def test_fp16_preserves_explicit_generated_float_types(operation):
+    """Keep generated FLOAT values typed correctly and provide private HALF consumer views."""
+    inputs = []
+    initializers = []
+    feeds = {}
+    shape = [2, 2]
+    if operation == "ConstantOfShape":
+        initializers = [nh.from_array(np.asarray(shape, np.int64), "shape")]
+        generated = h.make_node(
+            operation, ["shape"], ["generated"], value=nh.from_array(np.asarray([0.25], np.float32))
+        )
+    elif operation == "RandomNormal":
+        generated = h.make_node(operation, [], ["generated"], shape=shape, dtype=T.FLOAT, seed=42.0)
+    else:
+        inputs = [_value("x", shape)]
+        feeds = {"x": np.zeros(shape, np.float32)}
+        generated = h.make_node(operation, ["x"], ["generated"], dtype=T.FLOAT)
+    model = _model(
+        [generated, h.make_node("Relu", ["generated"], ["y"])],
+        inputs,
+        [_value("y", shape)],
+        initializers,
+    )
+    expected_cpu = _checked_cpu_outputs(model, feeds)[0]
+    expected = ReferenceEvaluator(model).run(None, feeds)[0].astype(np.float16).astype(np.float32)
+    changed = _apply_force_fp16_activations(model)
+    onnx.checker.check_model(changed, full_check=True)
+    np.testing.assert_array_equal(ReferenceEvaluator(changed).run(None, feeds)[0], expected)
+    # CPU fallback can retain FP32 arithmetic even when the graph requests HALF.
+    np.testing.assert_allclose(
+        _checked_cpu_outputs(changed, feeds)[0], expected_cpu, rtol=1e-3, atol=1e-3
+    )
+
+
+def test_fp16_lpnorm_allocates_support_node_names():
+    """Reserve normalization support-node names across existing graph owners."""
+    model = _model(
+        [
+            h.make_node(
+                "Constant",
+                [],
+                ["integer"],
+                name="normalize_square",
+                value=nh.from_array(np.asarray([1], np.int8)),
+            ),
+            h.make_node("LpNormalization", ["x"], ["norm"], name="normalize", p=2, axis=-1),
+            h.make_node("Mul", ["norm", "gamma"], ["y"]),
+        ],
+        [_value("x", [1, 1, 1, 2])],
+        [_value("y", [1, 1, 1, 2])],
+        [nh.from_array(np.ones([1, 1, 1, 2], np.float32), "gamma")],
+    )
+    changed = _apply_fp16_lpnorm(model)
+    names = [node.name for node in changed.graph.node if node.name]
+    assert len(names) == len(set(names))
+    onnx.checker.check_model(changed, full_check=True)
+
+
+@pytest.mark.parametrize("dtype", [T.UINT8, T.FLOAT])
+def test_fp16_skips_conversion_for_preserved_runtime_defaults_only(dtype):
+    """Skip weight conversion when the model contains only preserved defaults."""
+    numpy_dtype = np.uint8 if dtype == T.UINT8 else np.float32
+    model = _model(
+        [h.make_node("Cast", ["weight"], ["y"], to=T.FLOAT)],
+        [_value("weight", [2], dtype)],
+        [_value("y", [2])],
+        [nh.from_array(np.asarray([1, 2], numpy_dtype), "weight")],
+    )
+    weight = model.graph.initializer[0]
+    changed = _apply_force_fp16_activations(model)
+    assert next(t for t in changed.graph.initializer if t.name == "weight") is weight
+    np.testing.assert_array_equal(
+        _checked_cpu_outputs(changed, {})[0], np.asarray([1, 2], np.float32)
+    )
+    np.testing.assert_array_equal(
+        _checked_cpu_outputs(changed, {"weight": np.asarray([2, 3], numpy_dtype)})[0],
+        np.asarray([2, 3], np.float32),
+    )
+
+
+@pytest.mark.parametrize("operation", ["QuantizeLinear", "DequantizeLinear"])
+def test_fp16_supported_contrib_quantizers_keep_half_interface(operation):
+    """Retain the supported contrib quantizers' HALF scales without treating them as opaque."""
+    initializers = [nh.from_array(np.float32(0.5), "scale"), nh.from_array(np.int8(0), "zero")]
+    if operation == "QuantizeLinear":
+        nodes = [h.make_node(operation, ["x", "scale", "zero"], ["y"])]
+        inputs = [_value("x", [2])]
+        outputs = [_value("y", [2], T.INT8)]
+        feeds = {"x": np.asarray([-1, 2], np.float32)}
+    else:
+        initializers.append(nh.from_array(np.asarray([-2, 4], np.int8), "quantized"))
+        nodes = [
+            h.make_node(operation, ["quantized", "scale", "zero"], ["middle"]),
+            h.make_node("Relu", ["middle"], ["y"]),
+        ]
+        inputs = []
+        outputs = [_value("y", [2])]
+        feeds = {}
+    model = _model(nodes, inputs, outputs, initializers, opset=17)
+    model.graph.node[0].domain = "com.microsoft"
+    model.opset_import.append(h.make_opsetid("com.microsoft", 1))
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    changed = _apply_force_fp16_activations(model)
+    assert next(t for t in changed.graph.initializer if t.name == "scale").data_type == T.FLOAT16
+    assert next(op for op in changed.opset_import if not op.domain).version == 17
+    np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected)
+
+
+def test_terminal_projection_chunks_supported_contrib_dequantization():
+    """Keep the supported contrib dequantizer on each private bounded weight chunk."""
+    model = _terminal_model(quantized=True)
+    model.graph.node[0].domain = "com.microsoft"
+    model.opset_import.append(h.make_opsetid("com.microsoft", 1))
+    feeds = {"x": np.asarray([[[[2]]]], np.float32)}
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    changed = _apply_force_fp16_activations(model)
+    assert sum(node.op_type == "Conv" for node in changed.graph.node) == 2
+    np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected)
+
+
+def test_fp16_supported_contrib_gelu_keeps_upstream_half_precision():
+    """Keep a supported contrib activation and its standard normalization in HALF."""
+    model = h.make_model(
+        h.make_graph(
+            [
+                h.make_node("LayerNormalization", ["x", "scale", "bias"], ["normalized"], axis=-1),
+                h.make_node("Gelu", ["normalized"], ["y"], domain="com.microsoft"),
+            ],
+            "contrib_gelu",
+            [_value("x", [2, 2])],
+            [_value("y", [2, 2])],
+            [
+                nh.from_array(np.ones(2, np.float32), "scale"),
+                nh.from_array(np.zeros(2, np.float32), "bias"),
+            ],
+        ),
+        opset_imports=[h.make_opsetid("", 17), h.make_opsetid("com.microsoft", 1)],
+    )
+    model.ir_version = 9
+    model = onnx.shape_inference.infer_shapes(model)
+    feeds = {"x": np.asarray([[1, 3], [2, 4]], np.float32)}
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    changed = _apply_force_fp16_activations(model)
+    assert all(t.data_type == T.FLOAT16 for t in changed.graph.initializer)
+    assert next(op for op in changed.opset_import if not op.domain).version == 17
+    np.testing.assert_allclose(
+        _checked_cpu_outputs(changed, feeds)[0], expected, rtol=2e-3, atol=1e-4
+    )

From cc5df0b933144e37e49dc3b53d03c3ea75f4e60d Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 07:02:07 +0000
Subject: [PATCH 03/17] Retain terminal projection batch reduction and correct
 option docs

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../dla_transforms/dla_force_fp16_activations.py  |  7 ++++++-
 .../onnx/graph_surgery/make_dla_compatible.py     |  2 +-
 .../test_dla_force_fp16_activations.py            | 15 +++++++++++++--
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
index 40e375338a4..e90751630e4 100644
--- a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -563,7 +563,12 @@ def axes(node, expected):
         if any(len(cache.get_consumers(n.output[0])) != 1 for n in chain[1:]):
             continue
         channels = weight.dims[0]
-        if cache.get_shape(output.name) != [1, 1, channels]:
+        if cache.get_shape(output.name) != [1, 1, channels] or cache.get_shape(conv.output[0]) != [
+            1,
+            channels,
+            1,
+            1,
+        ]:
             continue
         bytes_per_channel = math.prod(weight.dims[1:]) * (1 if dq is not None else 2)
         payload = weight.raw_data
diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py
index 920c7286246..4b07e55bec3 100644
--- a/modelopt/onnx/graph_surgery/make_dla_compatible.py
+++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py
@@ -363,9 +363,9 @@ def dla_make_dla_compatible(
         qdq_quantized_dtypes: Zero-point dtypes whose QDQ pairs are removed.
             Defaults to UINT16 and INT16. Accepts ONNX dtype IDs or names.
         remove_all_qdq: Remove QDQ pairs for every dtype, matching the C++ flag.
+            Cannot be combined with qdq_quantized_dtypes.
         force_fp16_activations: Use FP16 where supported, keeping FP32 public
             boundaries and operations or controls that require FP32.
-            Cannot be combined with qdq_quantized_dtypes.
 
     Returns:
         The transformed model. Invalid models and failed transformations raise.
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index 37d1d34ad54..ede6c66e53b 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -489,7 +489,7 @@ def test_fp16_float_constant_keeps_valid_consumer_types():
     np.testing.assert_array_equal(_checked_cpu_outputs(changed, feeds)[0], expected)
 
 
-def _terminal_model(channels=8193, *, groups=1, quantized=False):
+def _terminal_model(channels=8193, *, groups=1, quantized=False, batch=1):
     """Build the exact terminal projection pattern with independently checkable values."""
     data_channels = groups
     weights = nh.from_array(
@@ -519,12 +519,23 @@ def _terminal_model(channels=8193, *, groups=1, quantized=False):
     )
     return _model(
         nodes,
-        [_value("x", [1, data_channels, 1, 1])],
+        [_value("x", [batch, data_channels, 1, 1])],
         [_value("y", [1, 1, channels])],
         initializers,
     )
 
 
+@pytest.mark.parametrize("quantized", [False, True])
+def test_terminal_projection_preserves_nontrivial_batch_reduction(quantized):
+    """Retain batch summation instead of replacing it with singleton projection chunks."""
+    model = _terminal_model(batch=2, quantized=quantized)
+    feeds = {"x": np.asarray([2, 6], np.float32).reshape(2, 1, 1, 1)}
+    expected = _checked_cpu_outputs(model, feeds)[0]
+    converted = _apply_force_fp16_activations(model)
+    assert sum(node.op_type == "Conv" for node in converted.graph.node) == 1
+    np.testing.assert_array_equal(_checked_cpu_outputs(converted, feeds)[0], expected)
+
+
 @pytest.mark.parametrize(
     "owner", ["grouped", "runtime_weight", "public_axes", "legacy", "constant"]
 )

From 96fb3c77e6f9e68b3cbeebe09864a127b3eeea6a Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 09:03:04 +0000
Subject: [PATCH 04/17] Map collapsed normalization axes before FP16
 decomposition

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../dla_force_fp16_activations.py             | 19 ++++-
 .../test_dla_force_fp16_activations.py        | 70 +++++++++++++++++++
 2 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
index e90751630e4..37feafc6f9b 100644
--- a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -24,13 +24,18 @@
 from ..utils.graph_utils import (
     GraphCache,
     _attribute,
+    _shape_before_boundary_promotion,
     batch_replace_nodes,
     convert_initializers_to_dtype,
 )
 from .dla_constants_to_initializers import _add_legacy_initializer_inputs
 from .dla_graph_cleanup import _subgraph_captures
 from .dla_indexing import _insert_support_nodes, _operator_nodes
-from .dla_rank_conversion import _apply_non4d_lpnorm, _runtime_input_names
+from .dla_rank_conversion import (
+    _apply_non4d_lpnorm,
+    _runtime_input_names,
+    _specialized_collapsed_axis,
+)
 
 __all__ = []
 
@@ -73,6 +78,17 @@ def _apply_fp16_lpnorm(model):
         node.output[0] for node in graph.node if node.op_type == "Constant" and id(node) in standard
     }
     opset = next((o.version for o in model.opset_import if o.domain in ("", "ai.onnx")), 1)
+    reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})}
+    mapped_axes = {}
+    for node in _operator_nodes(model, {"LpNormalization"}):
+        attribute = _attribute(node, "axis")
+        axis = attribute.i if attribute is not None else -1
+        shape = _shape_before_boundary_promotion(cache, node.input[0])
+        mapped_axes[id(node)] = (
+            _specialized_collapsed_axis(cache, node, shape, axis, reshapes)
+            if shape and len(shape) == 5
+            else axis
+        )
     nodes = []
     for node in graph.node:
         if node.op_type != "LpNormalization" or id(node) not in standard:
@@ -96,6 +112,7 @@ def _apply_fp16_lpnorm(model):
             nodes.append(node)
             continue
         base = node.name or node.output[0] + "_lpnorm"
+        axis = mapped_axes[id(node)]
         axes = cache.unique_name(base + "_axes")
         epsilon = cache.unique_name(base + "_epsilon")
         if opset >= 13:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index ede6c66e53b..4e4360f237b 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -29,6 +29,8 @@
     _apply_force_fp16_activations,
     _apply_fp16_lpnorm,
 )
+from modelopt.onnx.graph_surgery.dla_transforms.dla_rank_conversion import _apply_non4d_reshape
+from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible
 
 
 def _value(name, shape, dtype=T.FLOAT):
@@ -171,6 +173,74 @@ def test_scaled_l2_normalization_keeps_zero_slices_finite():
     np.testing.assert_allclose(actual, [[[[0, 0], [0.6, 0.8]]]], atol=3e-4)
 
 
+def _scaled_collapsed_lpnorm_model(shape, axis, prefix=False):
+    """Build checked native normalization around a real singleton Reshape collapse."""
+    collapsed = shape.copy()
+    collapsed.remove(1)
+    nodes = [
+        h.make_node("Reshape", ["x", "native_shape"], ["native"]),
+        h.make_node("LpNormalization", ["native"], ["norm"], axis=axis, p=2),
+        h.make_node("Mul", ["norm", "gamma"], ["scaled"]),
+        h.make_node("Reshape", ["scaled", "output_shape"], ["y"]),
+    ]
+    outputs = [_value("y", collapsed)]
+    initializers = [
+        nh.from_array(np.asarray(shape, np.int64), "native_shape"),
+        nh.from_array(np.asarray(collapsed, np.int64), "output_shape"),
+        nh.from_array(np.ones(shape[-1], np.float32), "gamma"),
+    ]
+    if prefix:
+        nodes = [
+            h.make_node("LpNormalization", ["x"], ["first_norm"], axis=-1, p=2),
+            h.make_node("Mul", ["first_norm", "first_gamma"], ["first"]),
+            *nodes,
+        ]
+        outputs.insert(0, _value("first", collapsed))
+        initializers.append(nh.from_array(np.ones(collapsed[-1], np.float32), "first_gamma"))
+    return _model(nodes, [_value("x", collapsed)], outputs, initializers, opset=17)
+
+
+@pytest.mark.parametrize("shape", [[1, 2, 3, 4, 5], [2, 3, 1, 4, 5]])
+@pytest.mark.parametrize("axis", [-1, 4])
+def test_fp16_scaled_lpnorm_uses_actual_collapsed_axis(shape, axis):
+    """Use the mapped last axis after the registered five-dimensional collapse."""
+    model = _scaled_collapsed_lpnorm_model(shape, axis)
+    collapsed = shape.copy()
+    collapsed.remove(1)
+    feeds = {"x": np.linspace(0.25, 2, np.prod(collapsed), dtype=np.float32).reshape(collapsed)}
+    onnx.checker.check_model(model, full_check=True)
+    expected = ReferenceEvaluator(model).run(None, feeds)[0]
+    converted = _transform_make_dla_compatible(model, verbose=False, force_fp16_activations=True)
+    onnx.checker.check_model(converted, full_check=True)
+    actual = ReferenceEvaluator(converted).run(None, feeds)[0]
+    np.testing.assert_allclose(actual, expected, rtol=2e-3, atol=3e-4)
+
+
+@pytest.mark.parametrize("axis", [-1, 4])
+@pytest.mark.parametrize("enabled", [False, True])
+def test_scaled_lpnorm_refuses_normalizing_over_the_removed_axis(axis, enabled):
+    """Give FP16 and FP32 the same refusal for a normalization axis removed by collapse."""
+    model = _scaled_collapsed_lpnorm_model([2, 3, 4, 5, 1], axis)
+    feeds = {"x": np.ones([2, 3, 4, 5], np.float32)}
+    onnx.checker.check_model(model, full_check=True)
+    np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], feeds["x"])
+    with pytest.raises(ValueError, match="cannot map axis"):
+        _transform_make_dla_compatible(model, verbose=False, force_fp16_activations=enabled)
+
+
+@pytest.mark.parametrize("axis", [-1, 4])
+def test_fp16_lpnorm_validates_late_collapsed_axis_before_decomposition(axis):
+    """Refuse unsupported normalization before mutating an earlier eligible FP16 branch."""
+    model = _scaled_collapsed_lpnorm_model([2, 3, 4, 5, 1], axis, prefix=True)
+    onnx.checker.check_model(model, full_check=True)
+    ReferenceEvaluator(model).run(None, {"x": np.ones([2, 3, 4, 5], np.float32)})
+    _apply_non4d_reshape(model)
+    before = model.SerializeToString()
+    with pytest.raises(ValueError, match="cannot map axis"):
+        _apply_fp16_lpnorm(model)
+    assert model.SerializeToString() == before
+
+
 def test_terminal_projection_splits_channels_and_preserves_values():
     """Bound private projection channels while preserving real output values."""
     channels = 8193

From 4f6e5f69046aee4d49fa1455a8ec2f45ff667c0d Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 12:51:02 +0000
Subject: [PATCH 05/17] Preserve last-axis FP16 normalization after promotion

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../dla_force_fp16_activations.py             |  2 +-
 .../test_dla_force_fp16_activations.py        | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
index 37feafc6f9b..df2a9aff63a 100644
--- a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -87,7 +87,7 @@ def _apply_fp16_lpnorm(model):
         mapped_axes[id(node)] = (
             _specialized_collapsed_axis(cache, node, shape, axis, reshapes)
             if shape and len(shape) == 5
-            else axis
+            else -1
         )
     nodes = []
     for node in graph.node:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index 4e4360f237b..56394839790 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -173,6 +173,36 @@ def test_scaled_l2_normalization_keeps_zero_slices_finite():
     np.testing.assert_allclose(actual, [[[[0, 0], [0.6, 0.8]]]], atol=3e-4)
 
 
+@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]])
+@pytest.mark.parametrize("positive_axis", [False, True])
+@pytest.mark.parametrize("enabled", [False, True])
+@pytest.mark.parametrize("unary", [False, True])
+def test_scaled_low_rank_lpnorm_preserves_last_axis(shape, positive_axis, enabled, unary):
+    """Normalize the last runtime dimension after promoting a native positive axis."""
+    axis = len(shape) - 1 if positive_axis else -1
+    nodes = [h.make_node("Relu", ["x"], ["native"])] if unary else []
+    nodes.extend(
+        [
+            h.make_node("LpNormalization", ["native" if unary else "x"], ["norm"], axis=axis, p=2),
+            h.make_node("Mul", ["norm", "gamma"], ["y"]),
+        ]
+    )
+    model = _model(
+        nodes,
+        [_value("x", shape)],
+        [_value("y", shape)],
+        [nh.from_array(np.linspace(0.5, 1.5, shape[-1], dtype=np.float32), "gamma")],
+        opset=17,
+    )
+    feeds = {"x": np.linspace(0.25, 2, np.prod(shape), dtype=np.float32).reshape(shape)}
+    onnx.checker.check_model(model, full_check=True)
+    expected = ReferenceEvaluator(model).run(None, feeds)[0]
+    converted = _transform_make_dla_compatible(model, verbose=False, force_fp16_activations=enabled)
+    onnx.checker.check_model(converted, full_check=True)
+    actual = ReferenceEvaluator(converted).run(None, feeds)[0]
+    np.testing.assert_allclose(actual, expected, rtol=2e-3, atol=3e-4)
+
+
 def _scaled_collapsed_lpnorm_model(shape, axis, prefix=False):
     """Build checked native normalization around a real singleton Reshape collapse."""
     collapsed = shape.copy()

From feb6b5e6406d79e1dad9507410e9c58c7ab3698c Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 17:15:15 +0000
Subject: [PATCH 06/17] Preserve legacy equal Split semantics during FP16 opset
 migration

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../dla_force_fp16_activations.py             | 11 ++-
 .../test_dla_force_fp16_activations.py        | 77 +++++++++++++++----
 2 files changed, 71 insertions(+), 17 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
index df2a9aff63a..7351853ed68 100644
--- a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -258,7 +258,7 @@ def _fp32_regions(model, cache):
 
 
 def _upgrade_fp16_quantizers(model, cache):
-    """Migrate inherited reduction schemas when FP16 quantizers require opset 19."""
+    """Migrate inherited reduction and Split schemas when FP16 quantizers require opset 19."""
     if not _operator_nodes(model, {"QuantizeLinear", "DequantizeLinear"}):
         return
     opset = next(op for op in model.opset_import if op.domain in {"", "ai.onnx"})
@@ -282,6 +282,15 @@ def _upgrade_fp16_quantizers(model, cache):
                 or (node.domain, node.op_type, node.overload) in functions
             ):
                 continue
+            if (
+                node.op_type == "Split"
+                and (len(node.input) < 2 or not node.input[1])
+                and _attribute(node, "num_outputs") is None
+            ):
+                # Split-18 requires an explicit choice; legacy equal splits were even.
+                del node.input[1:]
+                node.attribute.append(helper.make_attribute("num_outputs", len(node.output)))
+                continue
             axes = _attribute(node, "axes")
             if axes is None or not node.op_type.startswith("Reduce"):
                 continue
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index 56394839790..248e26ec6bd 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -420,27 +420,72 @@ def test_fp16_malformed_late_weight_refuses_before_any_mutation():
     assert model.SerializeToString() == before
 
 
-def test_fp16_qdq_upgrade_keeps_legacy_reduce_axes_valid():
-    """Migrate schema-sensitive reduction axes when enabling FP16 quantizers."""
+@pytest.mark.parametrize(
+    ("opset", "split"),
+    [(opset, split) for opset in [13, 17, 18, 19] for split in ["none", "equal", "sizes", "nested"]]
+    + [(13, "empty"), (17, "empty")],
+)
+def test_fp16_qdq_upgrade_keeps_legacy_reduce_axes_valid(opset, split):
+    """Migrate reduction axes and live equal splits while retaining explicit Split configurations."""
+    initializers = [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")]
+    mean_inputs, mean_attributes = ["x"], {"keepdims": 1}
+    if opset < 18:
+        mean_attributes["axes"] = [1]
+    else:
+        mean_inputs.append("axes")
+        initializers.append(nh.from_array(np.array([1], np.int64), "axes"))
+    nodes = [
+        h.make_node("ReduceMean", mean_inputs, ["mean"], **mean_attributes),
+        h.make_node("QuantizeLinear", ["mean", "scale", "zero"], ["q"]),
+        h.make_node(
+            "DequantizeLinear", ["q", "scale", "zero"], ["y" if split == "none" else "decoded"]
+        ),
+    ]
+    outputs = [_value("y", [4, 1])]
+    if split != "none":
+        split_inputs = ["x" if split == "nested" else "decoded"]
+        attrs = {"axis": 0}
+        if split == "sizes":
+            split_inputs.append("sizes")
+            initializers.append(nh.from_array(np.array([1, 3], np.int64), "sizes"))
+        else:
+            if split == "empty":
+                split_inputs.append("")
+            if opset >= 18:
+                attrs["num_outputs"] = 2
+        splitting = [
+            h.make_node("Split", split_inputs, ["left", "right"], **attrs),
+            h.make_node(
+                "Concat", ["left", "right"], ["y" if split != "nested" else "branch_y"], axis=0
+            ),
+        ]
+        if split == "nested":
+            branch = h.make_graph(splitting, "split_branch", [], [_value("branch_y", [4, 2])])
+            initializers.append(nh.from_array(np.asarray(True), "condition"))
+            nodes.append(
+                h.make_node("If", ["condition"], ["y"], then_branch=branch, else_branch=branch)
+            )
+            outputs = [_value("y", [4, 2]), _value("decoded", [4, 1])]
+        else:
+            nodes.extend(splitting)
     model = _model(
-        [
-            h.make_node("ReduceMean", ["x"], ["mean"], axes=[1], keepdims=1),
-            h.make_node("QuantizeLinear", ["mean", "scale", "zero"], ["q"]),
-            h.make_node("DequantizeLinear", ["q", "scale", "zero"], ["y"]),
-        ],
-        [_value("x", [2, 2])],
-        [_value("y", [2, 1])],
-        [nh.from_array(np.float32(0.1), "scale"), nh.from_array(np.int8(0), "zero")],
-        opset=17,
+        nodes,
+        [_value("x", [4, 2])],
+        outputs,
+        initializers,
+        opset=opset,
     )
     onnx.checker.check_model(model, full_check=True)
-    feeds = {"x": np.asarray([[0.2, 0.2], [0.4, 0.4]], np.float32)}
-    expected = _checked_cpu_outputs(model, feeds)[0]
+    feeds = {"x": np.repeat(np.asarray([0.2, 0.4, 0.6, 0.8], np.float32), 2).reshape(4, 2)}
+    expected = _checked_cpu_outputs(model, feeds)
     changed = _apply_force_fp16_activations(model)
     onnx.checker.check_model(changed, full_check=True)
-    np.testing.assert_allclose(
-        ReferenceEvaluator(changed).run(None, feeds)[0], expected, rtol=2e-3, atol=1e-4
-    )
+    for actuals in [
+        ReferenceEvaluator(changed).run(None, feeds),
+        _checked_cpu_outputs(changed, feeds),
+    ]:
+        for actual, reference in zip(actuals, expected):
+            np.testing.assert_allclose(actual, reference, rtol=2e-3, atol=1e-4)
 
 
 def test_fp16_lpnorm_keeps_opaque_function_native():

From 0e16914aca55f5e7c4019377625ca7e369f72435 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 18:09:14 +0000
Subject: [PATCH 07/17] Document FP16 L2 intermediate range limitations

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 modelopt/onnx/graph_surgery/__main__.py            | 5 ++++-
 modelopt/onnx/graph_surgery/make_dla_compatible.py | 7 +++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/modelopt/onnx/graph_surgery/__main__.py b/modelopt/onnx/graph_surgery/__main__.py
index 62135c974ac..de97becee5e 100644
--- a/modelopt/onnx/graph_surgery/__main__.py
+++ b/modelopt/onnx/graph_surgery/__main__.py
@@ -287,7 +287,10 @@ def main():
     dla_parser.add_argument(
         "--force-fp16-activations",
         action="store_true",
-        help="Use FP16 internal activations and weights with FP32 model boundaries",
+        help=(
+            "Use FP16 internal activations and weights with FP32 model boundaries; "
+            "L2 normalization intermediates can overflow or underflow"
+        ),
     )
 
     # Analyze subcommand
diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py
index 4b07e55bec3..866ac4ef843 100644
--- a/modelopt/onnx/graph_surgery/make_dla_compatible.py
+++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py
@@ -354,6 +354,13 @@ def dla_make_dla_compatible(
     Masked multiplication can change the sign of zero even for finite inputs; validate
     sign-sensitive downstream operations before using this conversion.
 
+    With ``force_fp16_activations=True``, last-axis L2 normalization followed by a
+    static gamma multiply uses FP16 arithmetic with a denominator floor of ``2**-14``.
+    This decomposition does not rescale inputs: squares or their sum can overflow,
+    and small squares can underflow, even for finite FP32 inputs. Validate numerical
+    accuracy on the target runtime; leave this option disabled when its intermediate
+    computations need a wider floating-point range.
+
     Args:
         model_path: Input file; external weights are loaded from its directory.
         output_path: Destination ONNX file.

From 86a59fd8092293c6d38e8c0d44522d1c61856667 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 18:29:12 +0000
Subject: [PATCH 08/17] Remove redundant FP16 validation test with misleading
 rollback claim

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../test_dla_force_fp16_activations.py        | 21 -------------------
 1 file changed, 21 deletions(-)

diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index 248e26ec6bd..b0cef6c1ca2 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -303,27 +303,6 @@ def test_terminal_projection_splits_channels_and_preserves_values():
     np.testing.assert_array_equal(actual, np.full((1, 1, channels), np.float32(0.5)))
 
 
-def test_excluded_initializers_are_restored_when_conversion_fails():
-    """Retain excluded weight payloads when the original converter rejects an included tensor."""
-    roi = nh.from_array(np.array([], np.float32), "roi")
-    scales = nh.from_array(np.array([1, 1, 1, 2], np.float32), "scales")
-    invalid = T(name="invalid", data_type=T.FLOAT, dims=[2], float_data=[1])
-    model = h.make_model(
-        h.make_graph(
-            [h.make_node("Resize", ["x", "roi", "scales"], ["y"])],
-            "conversion_failure",
-            [_value("x", [1, 1, 1, 2])],
-            [_value("y", [1, 1, 1, 4])],
-            [roi, invalid, scales],
-        )
-    )
-    expected = {value.name: value.SerializeToString() for value in (roi, scales)}
-    with pytest.raises(ValueError):
-        _apply_force_fp16_activations(model)
-    actual = {value.name: value.SerializeToString() for value in model.graph.initializer}
-    assert all(actual[name] == payload for name, payload in expected.items())
-
-
 def _checked_cpu_outputs(model, feeds):
     """Execute a fully checked graph with bounded CPU runtime threading."""
     onnx.checker.check_model(model, full_check=True)

From 18af7ec3789fd334d0f08589d0e4acb546f4c696 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 18:53:01 +0000
Subject: [PATCH 09/17] Forward registered rank context through FP16
 normalization override

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../dla_transforms/dla_force_fp16_activations.py              | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
index 7351853ed68..1062128d496 100644
--- a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -69,7 +69,7 @@ def _is_cast(node, dtype):
     )
 
 
-def _apply_fp16_lpnorm(model):
+def _apply_fp16_lpnorm(model, *, prepared_data_ops=None):
     """Decompose last-axis L2 normalization followed by a static gamma multiply."""
     graph = model.graph
     cache = GraphCache(graph)
@@ -157,7 +157,7 @@ def _apply_fp16_lpnorm(model):
     )
     _insert_support_nodes(graph, nodes)
     _add_legacy_initializer_inputs(model)
-    return _apply_non4d_lpnorm(model)
+    return _apply_non4d_lpnorm(model, prepared_data_ops=prepared_data_ops)
 
 
 def _validate_float_weights(graph):

From 36f5164db202eb1b9bf6e0ab3a00d9767a42be58 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Tue, 27 Oct 2026 20:53:59 +0000
Subject: [PATCH 10/17] Give public Squeeze outputs private FP16 reader views

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../dla_force_fp16_activations.py             | 11 ++++
 .../test_dla_force_fp16_activations.py        | 52 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
index 1062128d496..c89c8cdf38d 100644
--- a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -460,6 +460,17 @@ def _apply_force_fp16_activations(model):
                 output = cache.unique_name(source + "_to_fp32_fp16")
                 before.setdefault(id(producer), []).append(_cast(source, output, fp32, cache))
                 producer.input[0] = output
+            internal = [
+                (consumer, port)
+                for consumer, port in cache.get_consumers(value.name)
+                if not _fp32_control(consumer, port, native_nodes, fp32_inputs)
+                and not _is_cast(consumer, fp16)
+            ]
+            if internal:
+                view = cache.unique_name(value.name + "_fp16")
+                after.setdefault(id(producer), []).append(_cast(value.name, view, fp16, cache))
+                for consumer, port in internal:
+                    consumer.input[port] = view
             continue
         if _is_cast(producer, fp32):
             continue
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index b0cef6c1ca2..9ac5fab9f4d 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -964,3 +964,55 @@ def test_fp16_supported_contrib_gelu_keeps_upstream_half_precision():
     np.testing.assert_allclose(
         _checked_cpu_outputs(changed, feeds)[0], expected, rtol=2e-3, atol=1e-4
     )
+
+
+@pytest.mark.parametrize("operation", ["Mul", "Add"])
+@pytest.mark.parametrize("control_reader", [False, True])
+def test_public_squeeze_has_private_fp16_readers_and_native_controls(operation, control_reader):
+    """Preserve public FP32 Squeeze values while keeping internal arithmetic well typed."""
+    nodes = [
+        h.make_node("Relu", ["x"], ["r4"]),
+        h.make_node("Squeeze", ["r4", "axes"], ["public"]),
+        h.make_node(operation, ["public", "weight"], ["result"]),
+        h.make_node("Cast", ["public"], ["half"], to=T.FLOAT16),
+        h.make_node("Cast", ["public"], ["integer"], to=T.INT32),
+    ]
+    inputs = [_value("x", [1, 1, 1, 2])]
+    outputs = [
+        _value("public", [2]),
+        _value("result", [2]),
+        _value("half", [2], T.FLOAT16),
+        _value("integer", [2], T.INT32),
+    ]
+    if control_reader:
+        nodes.append(h.make_node("Resize", ["image", "", "public"], ["resized"], mode="nearest"))
+        inputs.append(_value("image", [1, 2]))
+        outputs.append(_value("resized", [1, 4]))
+    model = _model(
+        nodes,
+        inputs,
+        outputs,
+        [
+            nh.from_array(np.array([0, 1, 2], np.int64), "axes"),
+            nh.from_array(np.array([2, 3], np.float32), "weight"),
+        ],
+    )
+    feeds = {"x": np.array([[[[1, 2]]]], np.float32)}
+    if control_reader:
+        feeds["image"] = np.array([[4, 5]], np.float32)
+    expected = ReferenceEvaluator(model).run(None, feeds)
+    changed = _apply_force_fp16_activations(model)
+    onnx.checker.check_model(changed, full_check=True)
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    actual = ort.InferenceSession(
+        changed.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, feeds)
+    for left, right in zip(actual, expected):
+        np.testing.assert_array_equal(left, right)
+    assert actual[0].dtype == np.float32
+    casts = {node.output[0]: node for node in changed.graph.node if node.op_type == "Cast"}
+    assert casts["half"].input[0] == "public"
+    if control_reader:
+        resize = next(node for node in changed.graph.node if node.op_type == "Resize")
+        assert resize.input[2] == "public"

From b7c167265baeed76851cb0c920558d55108ae75c Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Wed, 28 Oct 2026 03:37:10 +0000
Subject: [PATCH 11/17] Validate FP16 tensor types before saving converted
 models

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../onnx/graph_surgery/make_dla_compatible.py |  4 +-
 .../test_dla_force_fp16_activations.py        | 37 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py
index 866ac4ef843..1da734f7522 100644
--- a/modelopt/onnx/graph_surgery/make_dla_compatible.py
+++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py
@@ -329,8 +329,8 @@ def _transform_make_dla_compatible(
             logger.info("[DLA pipeline] Applying ForceFp16Activations")
         model = _apply_force_fp16_activations(model)
         # Match C++: symbolic inference stays before precision conversion.
-        model = onnx.shape_inference.infer_shapes(model)
-        onnx.checker.check_model(model)
+        model = onnx.shape_inference.infer_shapes(model, check_type=True)
+        onnx.checker.check_model(model, full_check=True)
     if verbose:
         logger.info("[DLA pipeline] Done (%.1fs).", time.perf_counter() - started)
     return model
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index 9ac5fab9f4d..4e9ba073cec 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -1016,3 +1016,40 @@ def test_public_squeeze_has_private_fp16_readers_and_native_controls(operation,
     if control_reader:
         resize = next(node for node in changed.graph.node if node.op_type == "Resize")
         assert resize.input[2] == "public"
+
+
+@pytest.mark.parametrize("existing_output", [False, True])
+def test_fp16_public_conversion_refuses_invalid_control_types_before_saving(
+    tmp_path, existing_output
+):
+    """Reject half-precision Resize scales without creating or replacing an output file."""
+    model = h.make_model(
+        h.make_graph(
+            [h.make_node("Resize", ["x", "", "scales"], ["y"], mode="nearest")],
+            "invalid_resize_scales",
+            [_value("x", [1, 1, 1, 2])],
+            [_value("y", [1, 1, 1, 4])],
+            [nh.from_array(np.array([1, 1, 1, 2], np.float16), "scales")],
+        ),
+        opset_imports=[h.make_opsetid("", 19)],
+        ir_version=9,
+    )
+    onnx.checker.check_model(model)
+    with pytest.raises((onnx.shape_inference.InferenceError, onnx.checker.ValidationError)):
+        onnx.checker.check_model(model, full_check=True)
+    source, output = tmp_path / "source.onnx", tmp_path / "output.onnx"
+    onnx.save(model, source)
+    original = source.read_bytes()
+    previous = b"existing output must survive validation failure"
+    if existing_output:
+        output.write_bytes(previous)
+    with pytest.raises((onnx.shape_inference.InferenceError, onnx.checker.ValidationError)):
+        make_dla_compatible(
+            str(source),
+            str(output),
+            use_external_data=False,
+            verbose=False,
+            force_fp16_activations=True,
+        )
+    assert source.read_bytes() == original
+    assert output.read_bytes() == previous if existing_output else not output.exists()

From 48c8ab1a15f6dc46f92448c0174eaac594ff11fc Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Wed, 28 Oct 2026 03:54:53 +0000
Subject: [PATCH 12/17] Reuse strict typed inference without copying payloads
 twice

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 modelopt/onnx/graph_surgery/make_dla_compatible.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py
index 1da734f7522..92c72f81d7e 100644
--- a/modelopt/onnx/graph_surgery/make_dla_compatible.py
+++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py
@@ -329,8 +329,9 @@ def _transform_make_dla_compatible(
             logger.info("[DLA pipeline] Applying ForceFp16Activations")
         model = _apply_force_fp16_activations(model)
         # Match C++: symbolic inference stays before precision conversion.
-        model = onnx.shape_inference.infer_shapes(model, check_type=True)
-        onnx.checker.check_model(model, full_check=True)
+        model = onnx.shape_inference.infer_shapes(model, check_type=True, strict_mode=True)
+        # Strict typed inference already covers full_check's extra model-copy pass.
+        onnx.checker.check_model(model)
     if verbose:
         logger.info("[DLA pipeline] Done (%.1fs).", time.perf_counter() - started)
     return model

From 4dfa2eae64f2c94a5052a3ed9b726fd2280645f1 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Wed, 28 Oct 2026 04:26:53 +0000
Subject: [PATCH 13/17] Preserve public FLOAT Cast outputs with FP16 arithmetic
 readers

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../dla_force_fp16_activations.py             | 17 ++--
 .../test_dla_force_fp16_activations.py        | 95 +++++++++++++++++++
 2 files changed, 104 insertions(+), 8 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
index c89c8cdf38d..ce45ed3d6a7 100644
--- a/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_force_fp16_activations.py
@@ -454,17 +454,20 @@ def _apply_force_fp16_activations(model):
         if value.type.tensor_type.elem_type != fp32 or value.name in native_values:
             continue
         producer = cache.get_producer(value.name)
-        if producer is not None and producer.op_type == "Squeeze":
-            source = producer.input[0]
-            if source not in preserved:
-                output = cache.unique_name(source + "_to_fp32_fp16")
-                before.setdefault(id(producer), []).append(_cast(source, output, fp32, cache))
-                producer.input[0] = output
+        public_cast = _is_cast(producer, fp32)
+        if producer is not None and (producer.op_type == "Squeeze" or public_cast):
+            if not public_cast:
+                source = producer.input[0]
+                if source not in preserved:
+                    output = cache.unique_name(source + "_to_fp32_fp16")
+                    before.setdefault(id(producer), []).append(_cast(source, output, fp32, cache))
+                    producer.input[0] = output
             internal = [
                 (consumer, port)
                 for consumer, port in cache.get_consumers(value.name)
                 if not _fp32_control(consumer, port, native_nodes, fp32_inputs)
                 and not _is_cast(consumer, fp16)
+                and (not public_cast or consumer.op_type != "Cast")
             ]
             if internal:
                 view = cache.unique_name(value.name + "_fp16")
@@ -472,8 +475,6 @@ def _apply_force_fp16_activations(model):
                 for consumer, port in internal:
                     consumer.input[port] = view
             continue
-        if _is_cast(producer, fp32):
-            continue
         tensor = cache.get_init(value.name) if producer is None else None
         if producer is None and (tensor is None or tensor.data_type != fp16):
             continue
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index 4e9ba073cec..837010069a0 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -1053,3 +1053,98 @@ def test_fp16_public_conversion_refuses_invalid_control_types_before_saving(
         )
     assert source.read_bytes() == original
     assert output.read_bytes() == previous if existing_output else not output.exists()
+
+
+@pytest.mark.parametrize("source_type", [T.INT32, T.INT64])
+@pytest.mark.parametrize("operation", ["Mul", "Add"])
+@pytest.mark.parametrize("control_reader", [False, True])
+def test_public_float_cast_keeps_boundaries_and_typed_internal_readers(
+    source_type, operation, control_reader
+):
+    """Share a half view for arithmetic while preserving public values and explicit Cast readers."""
+    nodes = [
+        h.make_node("Cast", ["x"], ["public"], to=T.FLOAT),
+        h.make_node(operation, ["public", "weight"], ["result"]),
+        h.make_node("Cast", ["public"], ["half"], to=T.FLOAT16),
+        h.make_node("Cast", ["public"], ["alias"], to=T.FLOAT),
+        h.make_node("Cast", ["public"], ["integer"], to=T.INT64),
+    ]
+    inputs = [_value("x", [2], source_type)]
+    outputs = [
+        _value("public", [2]),
+        _value("result", [2]),
+        _value("half", [2], T.FLOAT16),
+        _value("alias", [2]),
+        _value("integer", [2], T.INT64),
+    ]
+    dtype = np.int32 if source_type == T.INT32 else np.int64
+    feeds = {"x": np.array([1, 2] if control_reader else [2049, 2051], dtype)}
+    if control_reader:
+        nodes.append(h.make_node("Resize", ["image", "", "public"], ["resized"], mode="nearest"))
+        inputs.append(_value("image", [1, 2]))
+        outputs.append(_value("resized", [1, 4]))
+        feeds["image"] = np.array([[4, 5]], np.float32)
+    weight = np.array([2, 3], np.float32)
+    model = _model(nodes, inputs, outputs, [nh.from_array(weight, "weight")])
+    expected = ReferenceEvaluator(model).run(None, feeds)
+    expected[1] = (
+        (expected[2] * weight.astype(np.float16))
+        if operation == "Mul"
+        else (expected[2] + weight.astype(np.float16))
+    ).astype(np.float32)
+    changed = _apply_force_fp16_activations(model)
+    onnx.checker.check_model(changed, full_check=True)
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    actual = ort.InferenceSession(
+        changed.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, feeds)
+    for left, right in zip(actual, expected):
+        np.testing.assert_array_equal(left, right)
+    assert actual[0].dtype == np.float32
+    casts = {node.output[0]: node for node in changed.graph.node if node.op_type == "Cast"}
+    assert casts["half"].input[0] == "public"
+    if not control_reader:
+        assert casts["alias"].input[0] == "public"
+        assert casts["integer"].input[0] == "public"
+    else:
+        resize = next(node for node in changed.graph.node if node.op_type == "Resize")
+        assert resize.input[2] == "public"
+
+
+@pytest.mark.parametrize("source_type", [T.INT32, T.INT64])
+def test_public_fp16_conversion_supports_shared_float_cast_outputs(tmp_path, source_type):
+    """Convert and reload a valid shared integer-to-FLOAT graph through the real public API."""
+    shape = [1, 1, 1, 2]
+    model = _model(
+        [
+            h.make_node("Cast", ["x"], ["public"], to=T.FLOAT),
+            h.make_node("Mul", ["public", "weight"], ["result"]),
+        ],
+        [_value("x", shape, source_type)],
+        [_value("public", shape), _value("result", shape)],
+        [nh.from_array(np.array([2, 3], np.float32).reshape(shape), "weight")],
+    )
+    source, output = tmp_path / "source.onnx", tmp_path / "output.onnx"
+    onnx.save(model, source)
+    make_dla_compatible(
+        str(source),
+        str(output),
+        use_external_data=False,
+        verbose=False,
+        force_fp16_activations=True,
+    )
+    changed = onnx.load(output)
+    onnx.checker.check_model(changed, full_check=True)
+    dtype = np.int32 if source_type == T.INT32 else np.int64
+    values = np.array([2049, 2051], dtype).reshape(shape)
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    actual = ort.InferenceSession(
+        changed.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, {"x": values})
+    np.testing.assert_array_equal(actual[0], values.astype(np.float32))
+    expected = (values.astype(np.float16) * np.array([2, 3], np.float16).reshape(shape)).astype(
+        np.float32
+    )
+    np.testing.assert_array_equal(actual[1], expected)

From ba8c85ebc3c900300884fa58525876911b0620ff Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Wed, 28 Oct 2026 04:28:10 +0000
Subject: [PATCH 14/17] Check strict FP16 rounding with the reference evaluator

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../test_dla_force_fp16_activations.py        | 25 +++++++++++--------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
index 837010069a0..826567d77ef 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_force_fp16_activations.py
@@ -1094,11 +1094,14 @@ def test_public_float_cast_keeps_boundaries_and_typed_internal_readers(
     ).astype(np.float32)
     changed = _apply_force_fp16_activations(model)
     onnx.checker.check_model(changed, full_check=True)
-    options = ort.SessionOptions()
-    options.intra_op_num_threads = 1
-    actual = ort.InferenceSession(
-        changed.SerializeToString(), options, providers=["CPUExecutionProvider"]
-    ).run(None, feeds)
+    actual = ReferenceEvaluator(changed).run(None, feeds)
+    # ORT CPU promotes FP16 arithmetic; use exact values for the runtime smoke check.
+    cpu_feeds = dict(feeds, x=np.array([1, 2], dtype))
+    for left, right in zip(
+        _checked_cpu_outputs(changed, cpu_feeds),
+        ReferenceEvaluator(changed).run(None, cpu_feeds),
+    ):
+        np.testing.assert_array_equal(left, right)
     for left, right in zip(actual, expected):
         np.testing.assert_array_equal(left, right)
     assert actual[0].dtype == np.float32
@@ -1138,11 +1141,13 @@ def test_public_fp16_conversion_supports_shared_float_cast_outputs(tmp_path, sou
     onnx.checker.check_model(changed, full_check=True)
     dtype = np.int32 if source_type == T.INT32 else np.int64
     values = np.array([2049, 2051], dtype).reshape(shape)
-    options = ort.SessionOptions()
-    options.intra_op_num_threads = 1
-    actual = ort.InferenceSession(
-        changed.SerializeToString(), options, providers=["CPUExecutionProvider"]
-    ).run(None, {"x": values})
+    actual = ReferenceEvaluator(changed).run(None, {"x": values})
+    cpu_feeds = {"x": np.array([1, 2], dtype).reshape(shape)}
+    for left, right in zip(
+        _checked_cpu_outputs(changed, cpu_feeds),
+        ReferenceEvaluator(changed).run(None, cpu_feeds),
+    ):
+        np.testing.assert_array_equal(left, right)
     np.testing.assert_array_equal(actual[0], values.astype(np.float32))
     expected = (values.astype(np.float16) * np.array([2, 3], np.float16).reshape(shape)).astype(
         np.float32

From ddae0506aa98bb6463ef3b6993619bd6271b8a12 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Thu, 29 Oct 2026 20:17:51 +0000
Subject: [PATCH 15/17] Reduce redundant DLA test matrices and CPU execution

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../test_dla_5d_reshape_to_4d.py              |  10 +-
 .../dla_transforms/test_dla_cast_to_fp32.py   |   8 +-
 .../test_dla_constants_to_initializers.py     |   8 +-
 .../test_dla_convert_ops_to_4d.py             |   8 +-
 .../dla_transforms/test_dla_decompose_lstm.py |  10 +-
 ...t_dla_fix_instancenorm_channel_mismatch.py |   8 +-
 .../dla_transforms/test_dla_graph_cleanup.py  |   8 +-
 .../dla_transforms/test_dla_handle_qdq.py     |   4 +-
 .../dla_transforms/test_dla_indexing.py       | 121 ++++++++++++++----
 ..._dla_matmul_to_transpose_conv_transpose.py |   4 +-
 .../test_dla_model_boundaries.py              | 114 ++++++++++++++---
 .../test_dla_rank_conversion.py               | 118 ++++++++++++++---
 .../dla_transforms/test_dla_remove_deqlin.py  |  11 +-
 .../dla_transforms/test_dla_remove_qdq.py     |   8 +-
 .../test_dla_remove_reshapes.py               |   8 +-
 .../test_dla_reshape_patterns.py              |   6 +-
 16 files changed, 350 insertions(+), 104 deletions(-)

diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py
index 866c8c3a4f7..f3f3d778be1 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py
@@ -56,14 +56,12 @@ def _run_ort(model: onnx.ModelProto, feeds: dict[str, np.ndarray]) -> dict[str,
     """Execute all graph outputs with the CPU runtime."""
     m = copy.deepcopy(model)
     _clamp_ir(m)
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
     sess = ort.InferenceSession(
-        m.SerializeToString(),
-        providers=["CPUExecutionProvider"],
+        m.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
     )
-    out = {}
-    for o in m.graph.output:
-        out[o.name] = sess.run([o.name], feeds)[0]
-    return out
+    return dict(zip([output.name for output in m.graph.output], sess.run(None, feeds)))
 
 
 def _make_reduce_min_axis4_model() -> onnx.ModelProto:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py
index fbe61e881e3..6f315057feb 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py
@@ -66,8 +66,12 @@ def _prepare_model(model: onnx.ModelProto) -> None:
 def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute a small test graph on the CPU and return named outputs."""
     _clamp_ir(model)
-    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    sess = ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    )
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 def _cast_to_attr(model: onnx.ModelProto, cast_name: str) -> int | None:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_constants_to_initializers.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_constants_to_initializers.py
index d584f216f5d..0d03fe24b2e 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_constants_to_initializers.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_constants_to_initializers.py
@@ -147,8 +147,12 @@ def _prepare_model(model: onnx.ModelProto) -> None:
 def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute the test model on the CPU and return named outputs."""
     _clamp_ir(model)
-    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    sess = ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    )
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 def _initializer_names(model: onnx.ModelProto) -> set[str]:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py
index 4558c962581..0b777d4e45c 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py
@@ -142,8 +142,12 @@ def _make_model(nodes, name, inputs, outputs, inits=None, *, infer_symbolic_shap
 def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute all public outputs on the CPU runtime."""
     _clamp_ir(model)
-    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    sess = ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    )
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 @pytest.mark.parametrize("operation", ["Concat", "TopK"])
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py
index 7830f7fcd3e..66aba3f55bc 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py
@@ -155,14 +155,12 @@ def _expect_unsupported_lstm(model: onnx.ModelProto, expected: str) -> None:
 def _run_ort(model: onnx.ModelProto, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
     """Run ort."""
     _clamp_ir_version(model)
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
     sess = ort.InferenceSession(
-        model.SerializeToString(),
-        providers=["CPUExecutionProvider"],
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
     )
-    out: dict[str, np.ndarray] = {}
-    for o in model.graph.output:
-        out[o.name] = sess.run([o.name], feeds)[0]
-    return out
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py
index e8ee9b92900..d2b69adf4a7 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py
@@ -110,8 +110,12 @@ def _make_instancenorm_model(input_shape: list, epsilon: float = 1e-5) -> onnx.M
 def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute all public outputs on the CPU runtime."""
     _clamp_ir(model)
-    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    sess = ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    )
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py
index 640b3c1c0f9..79962162595 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py
@@ -134,8 +134,12 @@ def _prepare_model(model: onnx.ModelProto) -> None:
 def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Run all graph outputs with ORT's CPU execution provider."""
     _clamp_ir(model)
-    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    sess = ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    )
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 _RNG = np.random.default_rng(2024)
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_handle_qdq.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_handle_qdq.py
index c2df30040fe..969a79f746e 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_handle_qdq.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_handle_qdq.py
@@ -87,14 +87,14 @@ def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute the test model on the CPU and return named outputs."""
     _clamp_ir(model)
     options = ort.SessionOptions()
-    # Preserve the wrappers under test; ORT's QDQ propagation can introduce invalid signed Q nodes.
+    options.intra_op_num_threads = 1
     options.add_session_config_entry(
         "optimization.disable_specified_optimizers", "QDQPropagationTransformer"
     )
     sess = ort.InferenceSession(
         model.SerializeToString(), options, providers=["CPUExecutionProvider"]
     )
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 def _assert_wrapped(tr, op_type: str, orig_name: str) -> None:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py
index 08c977706c0..571eefee7f1 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py
@@ -40,10 +40,16 @@
 )
 
 
-@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]])
-@pytest.mark.parametrize("pipeline", [False, True])
-@pytest.mark.parametrize("producer", ["Reshape", "Squeeze"])
-@pytest.mark.parametrize("metadata", ["inferred", "omitted"])
+@pytest.mark.parametrize(
+    ("shape", "pipeline", "producer", "metadata"),
+    [
+        ([2, 3], True, "Reshape", "inferred"),
+        ([2, 3], False, "Squeeze", "inferred"),
+        ([2, 3, 4], True, "Squeeze", "inferred"),
+        ([2, 3, 4], False, "Reshape", "omitted"),
+        ([2, 3], True, "Squeeze", "omitted"),
+    ],
+)
 def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline, producer, metadata):
     """A native internal index view cannot follow data into a prepared four-dimensional layout."""
     source_shape = [1] * (4 - len(shape)) + shape
@@ -340,18 +346,45 @@ def _model(nodes, inputs, outputs, initializers=(), extra_opsets=()):
 
 
 @pytest.mark.parametrize(
-    "operation", ["GatherElements", "ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"]
-)
-@pytest.mark.parametrize(
-    ("producer", "shape"),
+    ("operation", "producer", "shape", "unary"),
     [
-        (producer, shape)
-        for producer in ["Reshape", "Squeeze", "opaque"]
-        for shape in [[2, 3], [2, 3, 4]]
-    ]
-    + [("Gemm", [2, 3])],
+        ("GatherElements", "Reshape", [2, 3], False),
+        ("ReduceMean", "Reshape", [2, 3], False),
+        ("ReduceMax", "Reshape", [2, 3], False),
+        ("ReduceMin", "Reshape", [2, 3], False),
+        ("ReduceSum", "Reshape", [2, 3, 4], False),
+        ("ReduceMean", "Reshape", [2, 3, 4], False),
+        ("ReduceMax", "Reshape", [2, 3, 4], False),
+        ("ReduceMin", "Reshape", [2, 3, 4], False),
+        ("GatherElements", "Squeeze", [2, 3], False),
+        ("ReduceSum", "Squeeze", [2, 3], False),
+        ("ReduceMean", "Squeeze", [2, 3], False),
+        ("ReduceMin", "Squeeze", [2, 3], False),
+        ("GatherElements", "Squeeze", [2, 3, 4], False),
+        ("ReduceSum", "Squeeze", [2, 3, 4], False),
+        ("ReduceMax", "Squeeze", [2, 3, 4], False),
+        ("ReduceMin", "Squeeze", [2, 3, 4], False),
+        ("ReduceSum", "opaque", [2, 3], False),
+        ("ReduceMean", "opaque", [2, 3], False),
+        ("ReduceMax", "opaque", [2, 3], False),
+        ("ReduceMin", "opaque", [2, 3], False),
+        ("GatherElements", "opaque", [2, 3, 4], False),
+        ("ReduceSum", "opaque", [2, 3, 4], False),
+        ("ReduceMean", "opaque", [2, 3, 4], False),
+        ("ReduceMax", "opaque", [2, 3, 4], False),
+        ("GatherElements", "Gemm", [2, 3], False),
+        ("ReduceMean", "Gemm", [2, 3], False),
+        ("ReduceMax", "Gemm", [2, 3], False),
+        ("ReduceMin", "Gemm", [2, 3], False),
+        ("ReduceSum", "Reshape", [2, 3], True),
+        ("GatherElements", "Reshape", [2, 3, 4], True),
+        ("ReduceMax", "Squeeze", [2, 3], True),
+        ("ReduceMean", "Squeeze", [2, 3, 4], True),
+        ("GatherElements", "opaque", [2, 3], True),
+        ("ReduceMin", "opaque", [2, 3, 4], True),
+        ("ReduceSum", "Gemm", [2, 3], True),
+    ],
 )
-@pytest.mark.parametrize("unary", [False, True])
 def test_internal_native_reshape_refuses_before_index_or_reduction_mutation(
     operation, producer, shape, unary
 ):
@@ -490,9 +523,17 @@ def test_symbolic_internal_reshape_retains_native_index_and_reduction_layout(ope
     np.testing.assert_array_equal(_execute(converted, feeds)[0], expected[0])
 
 
-@pytest.mark.parametrize("source", ["input", "initializer", "shape"])
-@pytest.mark.parametrize("scalar", [False, True])
-@pytest.mark.parametrize("entrypoint", ["stage", "pipeline"])
+@pytest.mark.parametrize(
+    ("source", "scalar", "entrypoint"),
+    [
+        ("input", False, "stage"),
+        ("shape", False, "stage"),
+        ("initializer", True, "stage"),
+        ("initializer", False, "pipeline"),
+        ("input", True, "pipeline"),
+        ("shape", True, "pipeline"),
+    ],
+)
 def test_rank_one_gather_control_paths_remain_native(source, scalar, entrypoint):
     """Preserve scalar/vector selection on native vectors, including Shape control paths."""
     dtype = TensorProto.INT64 if source == "shape" else TensorProto.FLOAT
@@ -537,12 +578,24 @@ def test_rank_one_gather_control_paths_remain_native(source, scalar, entrypoint)
 
 
 @pytest.mark.parametrize(
-    ("producer", "shape"),
-    [(name, shape) for name in ("Reshape", "Squeeze", "opaque") for shape in ([2, 3], [2, 3, 4])]
-    + [("Gemm", [2, 3])],
+    ("producer", "shape", "unary", "positive_axis"),
+    [
+        ("Reshape", [2, 3], False, False),
+        ("Squeeze", [2, 3], False, False),
+        ("Squeeze", [2, 3, 4], False, False),
+        ("opaque", [2, 3], False, False),
+        ("opaque", [2, 3, 4], False, False),
+        ("Gemm", [2, 3], False, False),
+        ("Reshape", [2, 3, 4], True, False),
+        ("Reshape", [2, 3, 4], False, True),
+        ("Reshape", [2, 3], True, True),
+        ("Squeeze", [2, 3], True, True),
+        ("Squeeze", [2, 3, 4], True, True),
+        ("opaque", [2, 3], True, True),
+        ("opaque", [2, 3, 4], True, True),
+        ("Gemm", [2, 3], True, True),
+    ],
 )
-@pytest.mark.parametrize("unary", [False, True])
-@pytest.mark.parametrize("positive_axis", [False, True])
 def test_gather_refuses_unprepared_internal_data_before_mutation(
     producer, shape, unary, positive_axis
 ):
@@ -752,10 +805,16 @@ def test_internal_index_data_uses_registered_rank_preparation(operation, produce
 
 
 @pytest.mark.parametrize(
-    ("operation", "shape"), [("MatMul", [2, 3]), ("MatMul", [2, 2, 3]), ("Gemm", [2, 3])]
+    ("operation", "shape", "owner", "entrypoint"),
+    [
+        ("MatMul", [2, 3], "output", "stage"),
+        ("Gemm", [2, 3], "output", "stage"),
+        ("MatMul", [2, 2, 3], "capture", "stage"),
+        ("MatMul", [2, 2, 3], "output", "pipeline"),
+        ("MatMul", [2, 3], "capture", "pipeline"),
+        ("Gemm", [2, 3], "capture", "pipeline"),
+    ],
 )
-@pytest.mark.parametrize("owner", ["output", "capture"])
-@pytest.mark.parametrize("entrypoint", ["stage", "pipeline"])
 def test_gather_refuses_native_matrix_outputs_with_other_rank_owners(
     operation, shape, owner, entrypoint
 ):
@@ -1227,10 +1286,16 @@ def test_gather_preserves_overridable_scalar_and_vector_indices(indices, produce
 
 
 @pytest.mark.parametrize(
-    ("shape", "unsqueeze_axes"), [([2, 3], [0]), ([2, 3], [0, 1]), ([2, 3, 4], [0])]
+    ("shape", "unsqueeze_axes", "negative", "keepdims"),
+    [
+        ([2, 3], [0], False, 0),
+        ([2, 3, 4], [0], False, 0),
+        ([2, 3], [0, 1], True, 0),
+        ([2, 3], [0, 1], False, 1),
+        ([2, 3], [0], True, 1),
+        ([2, 3, 4], [0], True, 1),
+    ],
 )
-@pytest.mark.parametrize("negative", [False, True])
-@pytest.mark.parametrize("keepdims", [0, 1])
 def test_reduction_preserves_rank_of_replaced_boundary_unsqueeze(
     shape, unsqueeze_axes, negative, keepdims
 ):
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py
index d8b941b7a3c..5b75969e53d 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py
@@ -84,12 +84,12 @@ def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute all graph outputs with the optional CPU runtime."""
     _clamp_ir_version(model)
     options = ort.SessionOptions()
-    # Compare graph-surgery semantics before runtime quantization fusions.
+    options.intra_op_num_threads = 1
     options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
     sess = ort.InferenceSession(
         model.SerializeToString(), options, providers=["CPUExecutionProvider"]
     )
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py
index 1c10edaa960..3a0dc3390c9 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py
@@ -323,10 +323,27 @@ def test_unnamed_input_adapters_keep_distinct_targets():
         np.testing.assert_array_equal(actual, reference)
 
 
-@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax"])
-@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]])
-@pytest.mark.parametrize("axis", [0, 1, -1, None])
-@pytest.mark.parametrize("opset", [13, 17])
+@pytest.mark.parametrize(
+    ("operation", "shape", "axis", "opset"),
+    [
+        ("Softmax", [2, 3], 0, 13),
+        ("LogSoftmax", [2, 3], 0, 13),
+        ("Softmax", [2, 3, 4], 0, 13),
+        ("LogSoftmax", [2, 3], 1, 13),
+        ("Softmax", [2, 3, 4], 1, 13),
+        ("LogSoftmax", [2, 3, 4], 1, 13),
+        ("Softmax", [2, 3], -1, 13),
+        ("Softmax", [2, 3, 4], -1, 13),
+        ("LogSoftmax", [2, 3, 4], -1, 13),
+        ("Softmax", [2, 3], None, 13),
+        ("LogSoftmax", [2, 3], None, 13),
+        ("LogSoftmax", [2, 3, 4], None, 13),
+        ("LogSoftmax", [2, 3, 4], 0, 17),
+        ("Softmax", [2, 3], 1, 17),
+        ("LogSoftmax", [2, 3], -1, 17),
+        ("Softmax", [2, 3, 4], None, 17),
+    ],
+)
 def test_boundary_softmax_preserves_native_axes(operation, shape, axis, opset):
     """Preserve explicit and schema-default normalization axes across input promotion."""
     model = _make_model(
@@ -516,10 +533,23 @@ def test_normalization_probe_accepts_unknown_sequence_element_rank():
     np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected[0])
 
 
-@pytest.mark.parametrize("batch", [1, 2, 3])
-@pytest.mark.parametrize("operation", ["Add", "Mul"])
-@pytest.mark.parametrize("tail", [False, True])
-@pytest.mark.parametrize("reverse", [False, True])
+@pytest.mark.parametrize(
+    ("batch", "operation", "tail", "reverse"),
+    [
+        (1, "Add", False, False),
+        (2, "Add", False, False),
+        (1, "Mul", False, False),
+        (3, "Mul", False, False),
+        (3, "Add", True, False),
+        (2, "Mul", True, False),
+        (2, "Add", False, True),
+        (3, "Add", False, True),
+        (2, "Mul", False, True),
+        (3, "Mul", False, True),
+        (1, "Add", True, True),
+        (1, "Mul", True, True),
+    ],
+)
 def test_normalization_binary_branch_preserves_broadcast_layout(batch, operation, tail, reverse):
     """Keep shared normalized public values and binary residual inputs in compatible layouts."""
     shape = [batch, 3, 4]
@@ -553,10 +583,23 @@ def test_normalization_binary_branch_preserves_broadcast_layout(batch, operation
         np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6)
 
 
-@pytest.mark.parametrize("width", [8, 12])
-@pytest.mark.parametrize("axis", [-1, 4])
-@pytest.mark.parametrize("quantized", [False, True])
-@pytest.mark.parametrize("preparation", ["direct", "intermediate", "constant"])
+@pytest.mark.parametrize(
+    ("width", "axis", "quantized", "preparation"),
+    [
+        (8, -1, False, "direct"),
+        (8, 4, False, "direct"),
+        (8, -1, True, "direct"),
+        (12, 4, True, "direct"),
+        (12, -1, False, "intermediate"),
+        (8, 4, False, "intermediate"),
+        (8, -1, True, "intermediate"),
+        (8, 4, True, "intermediate"),
+        (8, -1, False, "constant"),
+        (8, 4, False, "constant"),
+        (12, -1, True, "constant"),
+        (8, 4, True, "constant"),
+    ],
+)
 def test_early_registered_reshape_reduction_is_not_rejected(
     width, axis, quantized, preparation, monkeypatch
 ):
@@ -683,9 +726,17 @@ def test_ineligible_gemm_keeps_native_boundary_and_weights(tail):
     assert [list(tensor.dims) for tensor in converted.graph.initializer] == [[3, 4], [4]]
 
 
-@pytest.mark.parametrize("source", ["initializer", "constant", "runtime"])
-@pytest.mark.parametrize("axis", [0, 1])
-@pytest.mark.parametrize("tail", [False, True])
+@pytest.mark.parametrize(
+    ("source", "axis", "tail"),
+    [
+        ("initializer", 0, False),
+        ("runtime", 0, False),
+        ("constant", 1, False),
+        ("constant", 0, True),
+        ("initializer", 1, True),
+        ("runtime", 1, True),
+    ],
+)
 def test_scalar_gather_requires_registered_rank_conversion(source, axis, tail, monkeypatch):
     """Refuse unprepared scalar Gather before the actual pipeline changes caller state."""
     pipeline = import_module(_transform_make_dla_compatible.__module__)
@@ -1343,10 +1394,19 @@ def _rank_sensitive_adapter_model(operation, unary, metadata):
     return model, feeds
 
 
-@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax", "Transpose"])
-@pytest.mark.parametrize("unary", [False, True])
 @pytest.mark.parametrize(
-    ("entrypoint", "metadata"), [("stage", True), ("pipeline", True), ("pipeline", False)]
+    ("operation", "unary", "entrypoint", "metadata"),
+    [
+        ("Softmax", False, "stage", True),
+        ("LogSoftmax", True, "stage", True),
+        ("Transpose", True, "stage", True),
+        ("LogSoftmax", False, "pipeline", True),
+        ("Transpose", False, "pipeline", True),
+        ("Softmax", True, "pipeline", True),
+        ("LogSoftmax", False, "pipeline", False),
+        ("Transpose", False, "pipeline", False),
+        ("Softmax", True, "pipeline", False),
+    ],
 )
 def test_rank_changing_adapter_requires_actual_consumer_preparation(
     operation, unary, entrypoint, metadata, monkeypatch
@@ -1554,9 +1614,21 @@ def test_output_adapter_refuses_native_rank_five_squeeze_before_mutation(axes):
     assert model.SerializeToString() == before
 
 
-@pytest.mark.parametrize("operation", ["Flatten", "ArgMax", "ArgMin", "ReduceProd", "ReduceL2"])
-@pytest.mark.parametrize("tail", [False, True])
-@pytest.mark.parametrize("entrypoint", ["deferred", "pipeline"])
+@pytest.mark.parametrize(
+    ("operation", "tail", "entrypoint"),
+    [
+        ("Flatten", False, "deferred"),
+        ("ArgMin", False, "deferred"),
+        ("ReduceProd", False, "deferred"),
+        ("ReduceL2", False, "deferred"),
+        ("ArgMax", True, "deferred"),
+        ("ArgMax", False, "pipeline"),
+        ("Flatten", True, "pipeline"),
+        ("ArgMin", True, "pipeline"),
+        ("ReduceProd", True, "pipeline"),
+        ("ReduceL2", True, "pipeline"),
+    ],
+)
 def test_native_rank_reducing_outputs_require_registered_conversion(
     operation, tail, entrypoint, monkeypatch
 ):
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py
index 1f138505cd5..6955eef5505 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py
@@ -238,10 +238,22 @@ def test_symbolic_native_transpose_keeps_its_permutation(explicit, batch):
 
 
 @pytest.mark.parametrize(
-    "operation", ["Softmax", "LogSoftmax", "LpNormalization", "ArgMax", "Slice", "Transpose"]
+    ("operation", "negative", "tail"),
+    [
+        ("Softmax", False, False),
+        ("LpNormalization", False, False),
+        ("ArgMax", False, False),
+        ("Slice", False, False),
+        ("Transpose", False, False),
+        ("LogSoftmax", True, False),
+        ("LogSoftmax", False, True),
+        ("Softmax", True, True),
+        ("LpNormalization", True, True),
+        ("ArgMax", True, True),
+        ("Slice", True, True),
+        ("Transpose", True, True),
+    ],
 )
-@pytest.mark.parametrize("negative", [False, True])
-@pytest.mark.parametrize("tail", [False, True])
 def test_symbolic_native_axis_chain_preserves_concrete_intermediate_rank(operation, negative, tail):
     """Retain native axes after symbolic data acquires concrete Reshape metadata."""
     axis = -1 if negative else 1
@@ -398,10 +410,35 @@ def test_reshape_resolves_copied_dimensions_across_output_aliases(target, collis
         np.testing.assert_array_equal(actual, reference)
 
 
-@pytest.mark.parametrize("axis", [0, 1, 2, -1])
-@pytest.mark.parametrize("value_info", [False, True, "symbolic"])
-@pytest.mark.parametrize("ir_version", [3, 9])
-@pytest.mark.parametrize("entrypoint", ["stages", "pipeline"])
+@pytest.mark.parametrize(
+    ("axis", "value_info", "ir_version", "entrypoint"),
+    [
+        (0, False, 3, "stages"),
+        (2, False, 3, "stages"),
+        (-1, False, 3, "stages"),
+        (0, True, 3, "stages"),
+        (1, True, 3, "stages"),
+        (-1, True, 3, "stages"),
+        (0, "symbolic", 3, "stages"),
+        (1, "symbolic", 3, "stages"),
+        (-1, "symbolic", 3, "stages"),
+        (0, False, 9, "stages"),
+        (1, False, 9, "stages"),
+        (2, False, 9, "stages"),
+        (0, True, 9, "stages"),
+        (1, True, 9, "stages"),
+        (2, True, 9, "stages"),
+        (1, "symbolic", 9, "stages"),
+        (2, "symbolic", 9, "stages"),
+        (-1, "symbolic", 9, "stages"),
+        (1, False, 3, "pipeline"),
+        (2, True, 3, "pipeline"),
+        (2, "symbolic", 3, "pipeline"),
+        (-1, False, 9, "pipeline"),
+        (-1, True, 9, "pipeline"),
+        (0, "symbolic", 9, "pipeline"),
+    ],
+)
 def test_registered_axis_update_reads_native_initializer_rank(
     axis, value_info, ir_version, entrypoint
 ):
@@ -1293,10 +1330,17 @@ def test_native_index_views_survive_rank_preparation(producer, copied):
     np.testing.assert_array_equal(_execute(converted, feeds)[0], expected)
 
 
-@pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"])
-@pytest.mark.parametrize("batch", [1, 2])
-@pytest.mark.parametrize("static_output", [False, True])
-@pytest.mark.parametrize("tail", [False, True])
+@pytest.mark.parametrize(
+    ("operation", "batch", "static_output", "tail"),
+    [
+        ("Tile", 2, False, False),
+        ("Reshape", 1, True, False),
+        ("Expand", 2, True, False),
+        ("Expand", 1, False, True),
+        ("Reshape", 2, False, True),
+        ("Tile", 1, True, True),
+    ],
+)
 def test_symbolic_structural_controls_keep_native_data_rank(operation, batch, static_output, tail):
     """Retain symbolic native layouts, including fixed public outputs and unary readers."""
     target = {
@@ -1996,10 +2040,23 @@ def test_legacy_initializer_controls_remain_fixed(operation):
     np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0])
 
 
-@pytest.mark.parametrize("operation", ["QuantizeLinear", "DequantizeLinear"])
-@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]])
-@pytest.mark.parametrize("per_channel", [False, True, "singleton"])
-@pytest.mark.parametrize("domain", ["", "com.microsoft"])
+@pytest.mark.parametrize(
+    ("operation", "shape", "per_channel", "domain"),
+    [
+        ("QuantizeLinear", [2, 3], False, ""),
+        ("DequantizeLinear", [2, 3], False, ""),
+        ("QuantizeLinear", [2, 3, 4], True, ""),
+        ("DequantizeLinear", [2, 3, 4], True, ""),
+        ("QuantizeLinear", [2, 3, 4], "singleton", ""),
+        ("DequantizeLinear", [2, 3, 4], "singleton", ""),
+        ("QuantizeLinear", [2, 3, 4], False, "com.microsoft"),
+        ("DequantizeLinear", [2, 3, 4], False, "com.microsoft"),
+        ("QuantizeLinear", [2, 3], True, "com.microsoft"),
+        ("DequantizeLinear", [2, 3], True, "com.microsoft"),
+        ("QuantizeLinear", [2, 3], "singleton", "com.microsoft"),
+        ("DequantizeLinear", [2, 3], "singleton", "com.microsoft"),
+    ],
+)
 def test_registered_implicit_quantizer_axis_preserves_channel_values(
     operation, shape, per_channel, domain
 ):
@@ -2129,11 +2186,34 @@ def test_native_four_dimensional_runtime_controls_need_no_padding(operation, sou
 
 
 @pytest.mark.parametrize(
-    "operation",
-    ["Slice", "Transpose", "Softmax", "LogSoftmax", "LpNormalization", "ArgMax", "Expand", "Tile"],
+    ("operation", "rank", "source"),
+    [
+        ("Slice", 2, "unsqueeze"),
+        ("LogSoftmax", 2, "unsqueeze"),
+        ("LpNormalization", 2, "unsqueeze"),
+        ("ArgMax", 2, "unsqueeze"),
+        ("Expand", 2, "unsqueeze"),
+        ("Tile", 2, "unsqueeze"),
+        ("Transpose", 3, "unsqueeze"),
+        ("Softmax", 3, "unsqueeze"),
+        ("Transpose", 2, "unary"),
+        ("Softmax", 2, "unary"),
+        ("LogSoftmax", 2, "unary"),
+        ("Slice", 3, "unary"),
+        ("LpNormalization", 3, "unary"),
+        ("ArgMax", 3, "unary"),
+        ("Expand", 3, "unary"),
+        ("Tile", 3, "unary"),
+        ("Slice", 2, "reshape"),
+        ("Transpose", 2, "reshape"),
+        ("Softmax", 2, "reshape"),
+        ("LpNormalization", 2, "reshape"),
+        ("ArgMax", 2, "reshape"),
+        ("Expand", 2, "reshape"),
+        ("Tile", 2, "reshape"),
+        ("LogSoftmax", 3, "reshape"),
+    ],
 )
-@pytest.mark.parametrize("rank", [2, 3])
-@pytest.mark.parametrize("source", ["unsqueeze", "unary", "reshape"])
 def test_rank_one_adapters_keep_native_controls_unless_reshape_promotes(operation, rank, source):
     """Distinguish native rank-one input adapters from an explicitly padded Reshape."""
     shape = [1] * (rank - 1) + [6]
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
index 539a2e2e7f7..3b4471e3238 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
@@ -100,12 +100,15 @@ def _prepare_model(model: onnx.ModelProto) -> None:
 
 def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute the graph outputs with the CPU runtime."""
-    # ONNX Runtime is optional for these execution tests.
-    import onnxruntime as ort  # type: ignore[import-untyped]
+    import onnxruntime as ort
 
     _clamp_ir(model)
-    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    sess = ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    )
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 def _has_dq(model: onnx.ModelProto) -> bool:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_qdq.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_qdq.py
index f160e1a7a81..4e633d670db 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_qdq.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_qdq.py
@@ -93,8 +93,12 @@ def _prepare_model(model: onnx.ModelProto) -> onnx.ModelProto:
 def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute the test model on the CPU and return named outputs."""
     _clamp_ir(model)
-    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
-    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    sess = ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    )
+    return dict(zip([output.name for output in model.graph.output], sess.run(None, feeds)))
 
 
 def _apply(model, quant_types=None, **kw):
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py
index b488aa052a9..cc0dd9a60a2 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py
@@ -30,9 +30,11 @@
 def _execute(model, feeds):
     """Check and execute the actual CPU graph with native public outputs."""
     onnx.checker.check_model(model, full_check=True)
-    return ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]).run(
-        None, feeds
-    )
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    return ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    ).run(None, feeds)
 
 
 @pytest.mark.parametrize("ir_version", [3, 9])
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py
index a7fb53af5bc..a00d70a112a 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py
@@ -146,7 +146,11 @@ def _execute(model, feeds, runtime):
     """Execute all graph outputs with the selected independent engine."""
     if runtime == "reference":
         return ReferenceEvaluator(model).run(None, feeds)
-    session = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    session = ort.InferenceSession(
+        model.SerializeToString(), providers=["CPUExecutionProvider"], sess_options=options
+    )
     return session.run(None, feeds)
 
 

From 0f574b4a934e45d4a0be33a92d7a74858497e35d Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 30 Oct 2026 06:02:33 +0000
Subject: [PATCH 16/17] Move optional ORT test import to module scope

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../graph_surgery/dla_transforms/test_dla_remove_deqlin.py    | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
index 3b4471e3238..084cb4a5892 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
@@ -60,7 +60,7 @@
 
 from modelopt.onnx.graph_surgery.dla_transforms import dla_remove_deqlin as _mod
 
-pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks")
+ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks")
 
 _OPSET_VERSION = 21
 _MAX_IR_VERSION_FOR_ORT = 10
@@ -100,8 +100,6 @@ def _prepare_model(model: onnx.ModelProto) -> None:
 
 def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
     """Execute the graph outputs with the CPU runtime."""
-    import onnxruntime as ort
-
     _clamp_ir(model)
     options = ort.SessionOptions()
     options.intra_op_num_threads = 1

From 27d51e74015f32b3a91b0a1ed3f71301f8163145 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 30 Oct 2026 06:06:48 +0000
Subject: [PATCH 17/17] Keep optional test imports stable across Ruff versions

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../graph_surgery/dla_transforms/test_dla_cast_to_fp32.py     | 4 +---
 .../dla_transforms/test_dla_convert_ops_to_4d.py              | 4 +---
 .../test_dla_matmul_to_transpose_conv_transpose.py            | 4 +---
 .../graph_surgery/dla_transforms/test_dla_remove_deqlin.py    | 4 +---
 4 files changed, 4 insertions(+), 12 deletions(-)

diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py
index 6f315057feb..a3b916889bd 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py
@@ -51,9 +51,7 @@ def _prepare_model(model: onnx.ModelProto) -> None:
     """Populate intermediate shapes when the optional ORT tool is available."""
     try:
         # The tool submodule is optional; model execution only requires ORT itself.
-        from onnxruntime.tools.symbolic_shape_infer import (
-            SymbolicShapeInference,  # type: ignore[import-untyped]
-        )
+        from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
 
         inferred = SymbolicShapeInference.infer_shapes(model)
         if inferred is not None:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py
index 0b777d4e45c..11bf73b0c27 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py
@@ -127,9 +127,7 @@ def _make_model(nodes, name, inputs, outputs, inits=None, *, infer_symbolic_shap
     if infer_symbolic_shapes:
         try:
             # Shape inference is optional in partial runtime installations.
-            from onnxruntime.tools.symbolic_shape_infer import (
-                SymbolicShapeInference,  # type: ignore[import-untyped]
-            )
+            from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
 
         except ImportError:
             pass
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py
index 5b75969e53d..a58491000ec 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py
@@ -122,9 +122,7 @@ def _make_model(nodes, graph_name, inputs, outputs, initializers=None):
     _clamp_ir_version(model)
     try:
         # Optional runtime tools may be absent without preventing test collection.
-        from onnxruntime.tools.symbolic_shape_infer import (
-            SymbolicShapeInference,  # type: ignore[import-untyped]
-        )
+        from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
 
         model = SymbolicShapeInference.infer_shapes(model, auto_merge=True) or model
     except ImportError:
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
index 084cb4a5892..0406cf7144a 100644
--- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py
@@ -86,9 +86,7 @@ def _prepare_model(model: onnx.ModelProto) -> None:
     """Infer intermediate metadata when the optional runtime supports the graph."""
     try:
         # ONNX Runtime is optional for these execution tests.
-        from onnxruntime.tools.symbolic_shape_infer import (
-            SymbolicShapeInference,  # type: ignore[import-untyped]
-        )
+        from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
 
         inferred = SymbolicShapeInference.infer_shapes(model)
         if inferred is not None: