From 5f99679037a57b806bf772c9d7764f2bc2686c82 Mon Sep 17 00:00:00 2001
From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Date: Fri, 17 Jul 2026 15:17:28 +0000
Subject: [PATCH 1/7] Fix NVFP4 exporter node ordering

Ensure NVFP4 ONNX post-processing returns graphs with producers before consumers after TRT_FP4QDQ lowering and Cast insertion.

Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/export/nvfp4_exporter.py        |  5 ++
 .../test_nvfp4_onnx_export_cpu.py             | 77 +++++++++++++++++++
 2 files changed, 82 insertions(+)
 create mode 100644 tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py

diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py
index e8bdfa2db1f..ceda1d0b48b 100644
--- a/modelopt/onnx/export/nvfp4_exporter.py
+++ b/modelopt/onnx/export/nvfp4_exporter.py
@@ -17,6 +17,7 @@
 
 import numpy as np
 import onnx
+import onnx_graphsurgeon as gs
 import torch
 from onnx import numpy_helper
 
@@ -421,4 +422,8 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
         graph.initializer.extend(new_initializers)
         logger.info(f"Removed {len(initializers_to_delete)} initializers")
 
+        gs_graph = gs.import_onnx(onnx_model)
+        gs_graph.toposort()
+        onnx_model = gs.export_onnx(gs_graph)
+
         return onnx_model
diff --git a/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py b/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py
new file mode 100644
index 00000000000..68d6c1c2a24
--- /dev/null
+++ b/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py
@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""CPU tests for NVFP4 Torch quantization ONNX export."""
+
+import inspect
+import io
+
+import onnx
+import pytest
+import torch
+from _test_utils.torch.quantization.models import SimpleLinear
+
+import modelopt.torch.quantization as mtq
+import modelopt.torch.quantization.tensor_quant as tensor_quant
+from modelopt.onnx.export import NVFP4QuantExporter
+from modelopt.torch.quantization.utils import is_quantized_linear
+
+pytest.importorskip("onnxruntime")
+
+
+def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch):
+    def forward_loop(model):
+        model(sample_input)
+
+    def cpu_dynamic_block_quantize(inputs, *args):
+        return inputs
+
+    monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_dynamic_block_quantize)
+
+    model = SimpleLinear().eval()
+    sample_input = model.get_input()
+    model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop)
+
+    for module in model.modules():
+        assert not isinstance(module, torch.nn.Linear) or is_quantized_linear(module)
+        if isinstance(module, torch.nn.Linear):
+            module.input_quantizer.disable()
+            module.weight_quantizer._onnx_quantizer_type = "static"
+
+    buffer = io.BytesIO()
+    if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
+        kwargs = {"enable_onnx_checker": False}
+    else:
+        kwargs = {}
+
+    torch.onnx.export(
+        model,
+        sample_input,
+        buffer,
+        input_names=["input"],
+        output_names=["output"],
+        export_params=True,
+        opset_version=21,
+        dynamo=False,
+        **kwargs,
+    )
+
+    buffer.seek(0)
+    exported_model = onnx.load_model_from_string(buffer.read())
+    assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node)
+
+    converted_model = NVFP4QuantExporter.process_model(exported_model)
+    assert not any(node.op_type == "TRT_FP4QDQ" for node in converted_model.graph.node)
+    onnx.checker.check_model(converted_model)

From 84f2e702e92eb9f267d6f6bf15d61d0334803d26 Mon Sep 17 00:00:00 2001
From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Date: Fri, 17 Jul 2026 15:17:29 +0000
Subject: [PATCH 2/7] Move NVFP4 ONNX regression into CPU export tests

Keep the NVFP4 ONNX exporter regression with the existing CPU Torch ONNX export coverage instead of a standalone test file.

Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
---
 .../test_nvfp4_onnx_export_cpu.py             | 77 -------------------
 .../quantization/test_onnx_export_cpu.py      | 56 ++++++++++++++
 2 files changed, 56 insertions(+), 77 deletions(-)
 delete mode 100644 tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py

diff --git a/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py b/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py
deleted file mode 100644
index 68d6c1c2a24..00000000000
--- a/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py
+++ /dev/null
@@ -1,77 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 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.
-
-"""CPU tests for NVFP4 Torch quantization ONNX export."""
-
-import inspect
-import io
-
-import onnx
-import pytest
-import torch
-from _test_utils.torch.quantization.models import SimpleLinear
-
-import modelopt.torch.quantization as mtq
-import modelopt.torch.quantization.tensor_quant as tensor_quant
-from modelopt.onnx.export import NVFP4QuantExporter
-from modelopt.torch.quantization.utils import is_quantized_linear
-
-pytest.importorskip("onnxruntime")
-
-
-def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch):
-    def forward_loop(model):
-        model(sample_input)
-
-    def cpu_dynamic_block_quantize(inputs, *args):
-        return inputs
-
-    monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_dynamic_block_quantize)
-
-    model = SimpleLinear().eval()
-    sample_input = model.get_input()
-    model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop)
-
-    for module in model.modules():
-        assert not isinstance(module, torch.nn.Linear) or is_quantized_linear(module)
-        if isinstance(module, torch.nn.Linear):
-            module.input_quantizer.disable()
-            module.weight_quantizer._onnx_quantizer_type = "static"
-
-    buffer = io.BytesIO()
-    if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
-        kwargs = {"enable_onnx_checker": False}
-    else:
-        kwargs = {}
-
-    torch.onnx.export(
-        model,
-        sample_input,
-        buffer,
-        input_names=["input"],
-        output_names=["output"],
-        export_params=True,
-        opset_version=21,
-        dynamo=False,
-        **kwargs,
-    )
-
-    buffer.seek(0)
-    exported_model = onnx.load_model_from_string(buffer.read())
-    assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node)
-
-    converted_model = NVFP4QuantExporter.process_model(exported_model)
-    assert not any(node.op_type == "TRT_FP4QDQ" for node in converted_model.graph.node)
-    onnx.checker.check_model(converted_model)
diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py
index ed062df8c9b..dcfc78b58e4 100644
--- a/tests/unit/torch/quantization/test_onnx_export_cpu.py
+++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py
@@ -15,14 +15,24 @@
 
 """Unit tests for ONNX export for CPU quantization."""
 
+import inspect
+import io
+
+import onnx
 import pytest
 import torch
 
 pytest.importorskip("onnxruntime")
 
 from _test_utils.torch.misc import set_seed
+from _test_utils.torch.quantization.models import SimpleLinear
 from _test_utils.torch.quantization.onnx_export import TEST_MODELS, onnx_export_tester
 
+import modelopt.torch.quantization as mtq
+import modelopt.torch.quantization.tensor_quant as tensor_quant
+from modelopt.onnx.export import NVFP4QuantExporter
+from modelopt.torch.quantization.utils import is_quantized_linear
+
 
 @pytest.mark.parametrize("model_cls", TEST_MODELS)
 @pytest.mark.parametrize(
@@ -42,3 +52,49 @@ def test_onnx_export_cpu(model_cls, num_bits, per_channel_quantization, constant
     onnx_export_tester(
         model_cls(), "cpu", num_bits, per_channel_quantization, constant_folding, dtype
     )
+
+
+def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch):
+    def forward_loop(model):
+        model(sample_input)
+
+    def cpu_dynamic_block_quantize(inputs, *args):
+        return inputs
+
+    monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_dynamic_block_quantize)
+
+    model = SimpleLinear().eval()
+    sample_input = model.get_input()
+    model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop)
+
+    for module in model.modules():
+        assert not isinstance(module, torch.nn.Linear) or is_quantized_linear(module)
+        if isinstance(module, torch.nn.Linear):
+            module.input_quantizer.disable()
+            module.weight_quantizer._onnx_quantizer_type = "static"
+
+    buffer = io.BytesIO()
+    if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
+        kwargs = {"enable_onnx_checker": False}
+    else:
+        kwargs = {}
+
+    torch.onnx.export(
+        model,
+        sample_input,
+        buffer,
+        input_names=["input"],
+        output_names=["output"],
+        export_params=True,
+        opset_version=21,
+        dynamo=False,
+        **kwargs,
+    )
+
+    buffer.seek(0)
+    exported_model = onnx.load_model_from_string(buffer.read())
+    assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node)
+
+    converted_model = NVFP4QuantExporter.process_model(exported_model)
+    assert not any(node.op_type == "TRT_FP4QDQ" for node in converted_model.graph.node)
+    onnx.checker.check_model(converted_model)

From 9063c957d46ad4de978d467bbbb6b30e3cf16ef4 Mon Sep 17 00:00:00 2001
From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Date: Fri, 17 Jul 2026 15:17:30 +0000
Subject: [PATCH 3/7] Avoid GraphSurgeon round-trip for NVFP4 BF16 export

Use a protobuf-only graph node topological sort so NVFP4 post-processing preserves BF16 ONNX metadata while still returning producer-before-consumer node order.

Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/export/nvfp4_exporter.py | 48 +++++++++++++++++++++++---
 1 file changed, 44 insertions(+), 4 deletions(-)

diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py
index ceda1d0b48b..707550312dc 100644
--- a/modelopt/onnx/export/nvfp4_exporter.py
+++ b/modelopt/onnx/export/nvfp4_exporter.py
@@ -17,7 +17,6 @@
 
 import numpy as np
 import onnx
-import onnx_graphsurgeon as gs
 import torch
 from onnx import numpy_helper
 
@@ -35,6 +34,49 @@
 from .base_exporter import ONNXQuantExporter
 
 
+def _topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None:
+    """Stable-sort graph nodes so tensor producers precede consumers."""
+    nodes = list(graph.node)
+    producer_by_tensor: dict[str, int] = {}
+    for node_index, node in enumerate(nodes):
+        for output_name in node.output:
+            if not output_name:
+                continue
+            if output_name in producer_by_tensor:
+                raise ValueError(f"Duplicate producer for tensor {output_name!r}.")
+            producer_by_tensor[output_name] = node_index
+
+    dependencies: list[set[int]] = [set() for _ in nodes]
+    dependents: list[set[int]] = [set() for _ in nodes]
+
+    for consumer_index, node in enumerate(nodes):
+        for input_name in node.input:
+            producer_index = producer_by_tensor.get(input_name)
+            if producer_index is None:
+                continue
+            if producer_index == consumer_index:
+                raise ValueError(f"Node {node.name!r} consumes its own output {input_name!r}.")
+            dependencies[consumer_index].add(producer_index)
+            dependents[producer_index].add(consumer_index)
+
+    ready = [node_index for node_index, dependency in enumerate(dependencies) if not dependency]
+    sorted_nodes: list[onnx.NodeProto] = []
+    while ready:
+        node_index = ready.pop(0)
+        sorted_nodes.append(nodes[node_index])
+        for dependent_index in sorted(dependents[node_index]):
+            dependencies[dependent_index].remove(node_index)
+            if not dependencies[dependent_index]:
+                ready.append(dependent_index)
+        ready.sort()
+
+    if len(sorted_nodes) != len(nodes):
+        raise ValueError("Cycle detected while sorting NVFP4 ONNX graph nodes.")
+
+    graph.ClearField("node")
+    graph.node.extend(sorted_nodes)
+
+
 def _cast_fp4(array: np.ndarray) -> np.ndarray:
     """Cast a numpy array to FLOAT4E2M1 using PyTorch.
 
@@ -422,8 +464,6 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
         graph.initializer.extend(new_initializers)
         logger.info(f"Removed {len(initializers_to_delete)} initializers")
 
-        gs_graph = gs.import_onnx(onnx_model)
-        gs_graph.toposort()
-        onnx_model = gs.export_onnx(gs_graph)
+        _topologically_sort_graph_nodes(graph)
 
         return onnx_model

From 0370f27bf95948cbf8dca9a082336ea5a083c0e6 Mon Sep 17 00:00:00 2001
From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Date: Fri, 17 Jul 2026 15:17:31 +0000
Subject: [PATCH 4/7] Move ONNX graph toposort helper to utils

Share the protobuf-only ONNX graph node sorter from modelopt.onnx.utils and document why it avoids GraphSurgeon dtype round-tripping for enum dtypes such as BF16.

Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/export/nvfp4_exporter.py | 45 +----------------------
 modelopt/onnx/utils.py                 | 49 ++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 44 deletions(-)

diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py
index 707550312dc..b86ade13de5 100644
--- a/modelopt/onnx/export/nvfp4_exporter.py
+++ b/modelopt/onnx/export/nvfp4_exporter.py
@@ -34,49 +34,6 @@
 from .base_exporter import ONNXQuantExporter
 
 
-def _topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None:
-    """Stable-sort graph nodes so tensor producers precede consumers."""
-    nodes = list(graph.node)
-    producer_by_tensor: dict[str, int] = {}
-    for node_index, node in enumerate(nodes):
-        for output_name in node.output:
-            if not output_name:
-                continue
-            if output_name in producer_by_tensor:
-                raise ValueError(f"Duplicate producer for tensor {output_name!r}.")
-            producer_by_tensor[output_name] = node_index
-
-    dependencies: list[set[int]] = [set() for _ in nodes]
-    dependents: list[set[int]] = [set() for _ in nodes]
-
-    for consumer_index, node in enumerate(nodes):
-        for input_name in node.input:
-            producer_index = producer_by_tensor.get(input_name)
-            if producer_index is None:
-                continue
-            if producer_index == consumer_index:
-                raise ValueError(f"Node {node.name!r} consumes its own output {input_name!r}.")
-            dependencies[consumer_index].add(producer_index)
-            dependents[producer_index].add(consumer_index)
-
-    ready = [node_index for node_index, dependency in enumerate(dependencies) if not dependency]
-    sorted_nodes: list[onnx.NodeProto] = []
-    while ready:
-        node_index = ready.pop(0)
-        sorted_nodes.append(nodes[node_index])
-        for dependent_index in sorted(dependents[node_index]):
-            dependencies[dependent_index].remove(node_index)
-            if not dependencies[dependent_index]:
-                ready.append(dependent_index)
-        ready.sort()
-
-    if len(sorted_nodes) != len(nodes):
-        raise ValueError("Cycle detected while sorting NVFP4 ONNX graph nodes.")
-
-    graph.ClearField("node")
-    graph.node.extend(sorted_nodes)
-
-
 def _cast_fp4(array: np.ndarray) -> np.ndarray:
     """Cast a numpy array to FLOAT4E2M1 using PyTorch.
 
@@ -464,6 +421,6 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
         graph.initializer.extend(new_initializers)
         logger.info(f"Removed {len(initializers_to_delete)} initializers")
 
-        _topologically_sort_graph_nodes(graph)
+        utils.topologically_sort_graph_nodes(graph)
 
         return onnx_model
diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py
index 3b8edf76a84..8f5591f636a 100644
--- a/modelopt/onnx/utils.py
+++ b/modelopt/onnx/utils.py
@@ -35,6 +35,55 @@
 BASE_MIN_OPSET = 19
 
 
+def topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None:
+    """Stable-sort graph nodes so tensor producers precede consumers.
+
+    Unlike GraphSurgeon topological sorting, this operates directly on GraphProto
+    nodes and does not import/export tensor metadata that may include ONNX enum
+    dtypes such as BF16, which can be misinterpreted as NumPy dtypes during
+    conversion.
+    """
+    nodes = list(graph.node)
+    producer_by_tensor: dict[str, int] = {}
+    for node_index, node in enumerate(nodes):
+        for output_name in node.output:
+            if not output_name:
+                continue
+            if output_name in producer_by_tensor:
+                raise ValueError(f"Duplicate producer for tensor {output_name!r}.")
+            producer_by_tensor[output_name] = node_index
+
+    dependencies: list[set[int]] = [set() for _ in nodes]
+    dependents: list[set[int]] = [set() for _ in nodes]
+
+    for consumer_index, node in enumerate(nodes):
+        for input_name in node.input:
+            producer_index = producer_by_tensor.get(input_name)
+            if producer_index is None:
+                continue
+            if producer_index == consumer_index:
+                raise ValueError(f"Node {node.name!r} consumes its own output {input_name!r}.")
+            dependencies[consumer_index].add(producer_index)
+            dependents[producer_index].add(consumer_index)
+
+    ready = [node_index for node_index, dependency in enumerate(dependencies) if not dependency]
+    sorted_nodes: list[onnx.NodeProto] = []
+    while ready:
+        node_index = ready.pop(0)
+        sorted_nodes.append(nodes[node_index])
+        for dependent_index in sorted(dependents[node_index]):
+            dependencies[dependent_index].remove(node_index)
+            if not dependencies[dependent_index]:
+                ready.append(dependent_index)
+        ready.sort()
+
+    if len(sorted_nodes) != len(nodes):
+        raise ValueError("Cycle detected while sorting ONNX graph nodes.")
+
+    graph.ClearField("node")
+    graph.node.extend(sorted_nodes)
+
+
 def get_input_names_from_bytes(model_bytes: bytes, external_inputs_only: bool = True) -> list[str]:
     """This function returns the inputs names of the given onnx model in bytes.
 

From 7c88dbd394a8c1ce20fb694a8eaea0a56ef0fb7f Mon Sep 17 00:00:00 2001
From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Date: Fri, 17 Jul 2026 15:17:32 +0000
Subject: [PATCH 5/7] Reuse NVFP4 Cast nodes for shared activations

Cache Cast outputs per input and precision dtype so multiple NVFP4-lowered MatMul consumers reuse a single activation Cast instead of producing duplicate tensor names.

Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/export/nvfp4_exporter.py        | 29 +++---
 modelopt/onnx/utils.py                        | 98 +++++++++----------
 .../quantization/test_onnx_export_cpu.py      | 67 +++++++++++++
 3 files changed, 133 insertions(+), 61 deletions(-)

diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py
index b86ade13de5..ee5ee39aec2 100644
--- a/modelopt/onnx/export/nvfp4_exporter.py
+++ b/modelopt/onnx/export/nvfp4_exporter.py
@@ -331,6 +331,7 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
         }
         value_info_map = {vi.name: vi for vi in graph.value_info}
         graph_inputs = {inp.name for inp in graph.input}
+        cast_output_cache: dict[tuple[str, str], str] = {}
 
         def _get_precision_dtype() -> str:
             # Check initializers to determine the precision of the weights
@@ -350,18 +351,22 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str):
 
             # Create Cast nodes for each input of the target node except bias
             for i, input_name in enumerate(node.input[:2]):
-                cast_output_name = input_name + "_f16"  # Unique name for the cast output
-
-                # Create a Cast node to convert the input to FP16/BF16
-                cast_node = onnx.helper.make_node(
-                    "Cast",
-                    inputs=[input_name],  # Original input of the target node
-                    outputs=[cast_output_name],
-                    to=onnx_dtype_map[precision_dtype],  # Cast to FP16/BF16
-                )
-
-                # Insert the Cast node into the graph
-                graph.node.extend([cast_node])
+                cast_output_name = cast_output_cache.get((input_name, precision_dtype))
+                if cast_output_name is None:
+                    cast_output_suffix = "bf16" if precision_dtype == "BFloat16" else "f16"
+                    cast_output_name = f"{input_name}_{cast_output_suffix}"
+                    cast_output_cache[(input_name, precision_dtype)] = cast_output_name
+
+                    # Create a Cast node to convert the input to FP16/BF16
+                    cast_node = onnx.helper.make_node(
+                        "Cast",
+                        inputs=[input_name],  # Original input of the target node
+                        outputs=[cast_output_name],
+                        to=onnx_dtype_map[precision_dtype],  # Cast to FP16/BF16
+                    )
+
+                    # Insert the Cast node into the graph
+                    graph.node.extend([cast_node])
 
                 # Update the target node input to use the cast node output
                 node.input[i] = cast_output_name
diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py
index 8f5591f636a..eaab43179b5 100644
--- a/modelopt/onnx/utils.py
+++ b/modelopt/onnx/utils.py
@@ -35,55 +35,6 @@
 BASE_MIN_OPSET = 19
 
 
-def topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None:
-    """Stable-sort graph nodes so tensor producers precede consumers.
-
-    Unlike GraphSurgeon topological sorting, this operates directly on GraphProto
-    nodes and does not import/export tensor metadata that may include ONNX enum
-    dtypes such as BF16, which can be misinterpreted as NumPy dtypes during
-    conversion.
-    """
-    nodes = list(graph.node)
-    producer_by_tensor: dict[str, int] = {}
-    for node_index, node in enumerate(nodes):
-        for output_name in node.output:
-            if not output_name:
-                continue
-            if output_name in producer_by_tensor:
-                raise ValueError(f"Duplicate producer for tensor {output_name!r}.")
-            producer_by_tensor[output_name] = node_index
-
-    dependencies: list[set[int]] = [set() for _ in nodes]
-    dependents: list[set[int]] = [set() for _ in nodes]
-
-    for consumer_index, node in enumerate(nodes):
-        for input_name in node.input:
-            producer_index = producer_by_tensor.get(input_name)
-            if producer_index is None:
-                continue
-            if producer_index == consumer_index:
-                raise ValueError(f"Node {node.name!r} consumes its own output {input_name!r}.")
-            dependencies[consumer_index].add(producer_index)
-            dependents[producer_index].add(consumer_index)
-
-    ready = [node_index for node_index, dependency in enumerate(dependencies) if not dependency]
-    sorted_nodes: list[onnx.NodeProto] = []
-    while ready:
-        node_index = ready.pop(0)
-        sorted_nodes.append(nodes[node_index])
-        for dependent_index in sorted(dependents[node_index]):
-            dependencies[dependent_index].remove(node_index)
-            if not dependencies[dependent_index]:
-                ready.append(dependent_index)
-        ready.sort()
-
-    if len(sorted_nodes) != len(nodes):
-        raise ValueError("Cycle detected while sorting ONNX graph nodes.")
-
-    graph.ClearField("node")
-    graph.node.extend(sorted_nodes)
-
-
 def get_input_names_from_bytes(model_bytes: bytes, external_inputs_only: bool = True) -> list[str]:
     """This function returns the inputs names of the given onnx model in bytes.
 
@@ -2069,3 +2020,52 @@ def clear_stale_value_info(model: onnx.ModelProto) -> int:
     fixed_shapes = _reconcile_stale_output_shapes(model)
 
     return fixed_outputs + fixed_shapes + n_cleared
+
+
+def topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None:
+    """Stable-sort graph nodes so tensor producers precede consumers.
+
+    Unlike GraphSurgeon topological sorting, this operates directly on GraphProto
+    nodes and does not import/export tensor metadata that may include ONNX enum
+    dtypes such as BF16, which can be misinterpreted as NumPy dtypes during
+    conversion.
+    """
+    nodes = list(graph.node)
+    producer_by_tensor: dict[str, int] = {}
+    for node_index, node in enumerate(nodes):
+        for output_name in node.output:
+            if not output_name:
+                continue
+            if output_name in producer_by_tensor:
+                raise ValueError(f"Duplicate producer for tensor {output_name!r}.")
+            producer_by_tensor[output_name] = node_index
+
+    dependencies: list[set[int]] = [set() for _ in nodes]
+    dependents: list[set[int]] = [set() for _ in nodes]
+
+    for consumer_index, node in enumerate(nodes):
+        for input_name in node.input:
+            producer_index = producer_by_tensor.get(input_name)
+            if producer_index is None:
+                continue
+            if producer_index == consumer_index:
+                raise ValueError(f"Node {node.name!r} consumes its own output {input_name!r}.")
+            dependencies[consumer_index].add(producer_index)
+            dependents[producer_index].add(consumer_index)
+
+    ready = [node_index for node_index, dependency in enumerate(dependencies) if not dependency]
+    sorted_nodes: list[onnx.NodeProto] = []
+    while ready:
+        node_index = ready.pop(0)
+        sorted_nodes.append(nodes[node_index])
+        for dependent_index in sorted(dependents[node_index]):
+            dependencies[dependent_index].remove(node_index)
+            if not dependencies[dependent_index]:
+                ready.append(dependent_index)
+        ready.sort()
+
+    if len(sorted_nodes) != len(nodes):
+        raise ValueError("Cycle detected while sorting ONNX graph nodes.")
+
+    graph.ClearField("node")
+    graph.node.extend(sorted_nodes)
diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py
index dcfc78b58e4..455ea5a8816 100644
--- a/tests/unit/torch/quantization/test_onnx_export_cpu.py
+++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py
@@ -18,9 +18,11 @@
 import inspect
 import io
 
+import numpy as np
 import onnx
 import pytest
 import torch
+from onnx import TensorProto, helper, numpy_helper
 
 pytest.importorskip("onnxruntime")
 
@@ -98,3 +100,68 @@ def cpu_dynamic_block_quantize(inputs, *args):
     converted_model = NVFP4QuantExporter.process_model(exported_model)
     assert not any(node.op_type == "TRT_FP4QDQ" for node in converted_model.graph.node)
     onnx.checker.check_model(converted_model)
+
+
+def test_nvfp4_shared_activation_reuses_cast():
+    input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [4, 32])
+    outputs = [
+        helper.make_tensor_value_info("output0", TensorProto.FLOAT, [4, 64]),
+        helper.make_tensor_value_info("output1", TensorProto.FLOAT, [4, 64]),
+    ]
+    nodes = []
+    initializers = []
+    value_info = []
+
+    for index in range(2):
+        weight_name = f"linear{index}.weight"
+        fp4qdq_output = f"fp4qdq_output{index}"
+        initializers.append(
+            numpy_helper.from_array(
+                np.linspace(-1.0, 1.0, num=32 * 64, dtype=np.float32).reshape(32, 64),
+                weight_name,
+            )
+        )
+        value_info.append(helper.make_tensor_value_info(fp4qdq_output, TensorProto.FLOAT, [32, 64]))
+        nodes.extend(
+            [
+                helper.make_node(
+                    "TRT_FP4QDQ",
+                    inputs=[weight_name],
+                    outputs=[fp4qdq_output],
+                    name=f"weight{index}_fp4qdq",
+                    block_size=16,
+                ),
+                helper.make_node(
+                    "MatMul",
+                    inputs=["input", fp4qdq_output],
+                    outputs=[f"output{index}"],
+                    name=f"matmul{index}",
+                ),
+            ]
+        )
+
+    model = helper.make_model(
+        helper.make_graph(
+            nodes,
+            "shared_activation_nvfp4",
+            [input_tensor],
+            outputs,
+            initializers,
+            value_info=value_info,
+        )
+    )
+
+    converted_model = NVFP4QuantExporter.process_model(model)
+    activation_casts = [
+        node
+        for node in converted_model.graph.node
+        if node.op_type == "Cast" and node.input == ["input"]
+    ]
+    assert len(activation_casts) == 1
+    assert activation_casts[0].output == ["input_f16"]
+    assert all(
+        node.input[0] == "input_f16"
+        for node in converted_model.graph.node
+        if node.op_type == "MatMul"
+    )
+    onnx.checker.check_model(converted_model)

From e4039402ad7b314507bf67c6989ee842f0ada22b Mon Sep 17 00:00:00 2001
From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Date: Fri, 17 Jul 2026 15:17:33 +0000
Subject: [PATCH 6/7] Account for subgraph captures in ONNX graph sorting

Include implicit outer-scope tensor references from nested ONNX graph attributes when sorting graph nodes so control-flow nodes are ordered after captured tensor producers.

Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/utils.py                        | 39 +++++++++++++++-
 .../quantization/test_onnx_export_cpu.py      | 44 +++++++++++++++++++
 2 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py
index eaab43179b5..bc37c4a9333 100644
--- a/modelopt/onnx/utils.py
+++ b/modelopt/onnx/utils.py
@@ -2030,6 +2030,35 @@ def topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None:
     dtypes such as BF16, which can be misinterpreted as NumPy dtypes during
     conversion.
     """
+
+    def get_graph_outer_scope_inputs(subgraph: onnx.GraphProto) -> set[str]:
+        local_names = {
+            value.name for value in (*subgraph.input, *subgraph.initializer) if value.name
+        }
+        local_names.update(output_name for node in subgraph.node for output_name in node.output)
+
+        outer_scope_inputs = set()
+        for node in subgraph.node:
+            outer_scope_inputs.update(
+                input_name
+                for input_name in node.input
+                if input_name and input_name not in local_names
+            )
+            for attr in node.attribute:
+                graphs = []
+                if attr.type == onnx.AttributeProto.GRAPH:
+                    graphs.append(attr.g)
+                elif attr.type == onnx.AttributeProto.GRAPHS:
+                    graphs.extend(attr.graphs)
+
+                for nested_graph in graphs:
+                    outer_scope_inputs.update(
+                        input_name
+                        for input_name in get_graph_outer_scope_inputs(nested_graph)
+                        if input_name not in local_names
+                    )
+        return outer_scope_inputs
+
     nodes = list(graph.node)
     producer_by_tensor: dict[str, int] = {}
     for node_index, node in enumerate(nodes):
@@ -2044,7 +2073,15 @@ def topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None:
     dependents: list[set[int]] = [set() for _ in nodes]
 
     for consumer_index, node in enumerate(nodes):
-        for input_name in node.input:
+        input_names = set(node.input)
+        for attr in node.attribute:
+            if attr.type == onnx.AttributeProto.GRAPH:
+                input_names.update(get_graph_outer_scope_inputs(attr.g))
+            elif attr.type == onnx.AttributeProto.GRAPHS:
+                for subgraph in attr.graphs:
+                    input_names.update(get_graph_outer_scope_inputs(subgraph))
+
+        for input_name in input_names:
             producer_index = producer_by_tensor.get(input_name)
             if producer_index is None:
                 continue
diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py
index 455ea5a8816..f158b1c540a 100644
--- a/tests/unit/torch/quantization/test_onnx_export_cpu.py
+++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py
@@ -32,6 +32,7 @@
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.tensor_quant as tensor_quant
+from modelopt.onnx import utils
 from modelopt.onnx.export import NVFP4QuantExporter
 from modelopt.torch.quantization.utils import is_quantized_linear
 
@@ -165,3 +166,46 @@ def test_nvfp4_shared_activation_reuses_cast():
         if node.op_type == "MatMul"
     )
     onnx.checker.check_model(converted_model)
+
+
+def test_topologically_sort_graph_nodes_accounts_for_subgraph_captures():
+    input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])
+    cond_tensor = helper.make_tensor_value_info("cond", TensorProto.BOOL, [])
+    output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])
+    then_output = helper.make_tensor_value_info("then_output", TensorProto.FLOAT, [1])
+    else_output = helper.make_tensor_value_info("else_output", TensorProto.FLOAT, [1])
+
+    then_graph = helper.make_graph(
+        [helper.make_node("Identity", ["captured"], ["then_output"], name="then_use_captured")],
+        "then_branch",
+        [],
+        [then_output],
+    )
+    else_graph = helper.make_graph(
+        [helper.make_node("Identity", ["captured"], ["else_output"], name="else_use_captured")],
+        "else_branch",
+        [],
+        [else_output],
+    )
+    if_node = helper.make_node(
+        "If",
+        ["cond"],
+        ["output"],
+        name="if_uses_captured",
+        then_branch=then_graph,
+        else_branch=else_graph,
+    )
+    producer = helper.make_node("Identity", ["input"], ["captured"], name="producer")
+    model = helper.make_model(
+        helper.make_graph(
+            [if_node, producer],
+            "outer_scope_capture",
+            [input_tensor, cond_tensor],
+            [output_tensor],
+        )
+    )
+
+    utils.topologically_sort_graph_nodes(model.graph)
+
+    assert [node.name for node in model.graph.node] == ["producer", "if_uses_captured"]
+    onnx.checker.check_model(model)

From 1afe0deca27ccce9db84a1a378ead8ba814b3979 Mon Sep 17 00:00:00 2001
From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Date: Fri, 17 Jul 2026 15:17:34 +0000
Subject: [PATCH 7/7] Skip ONNX CPU export tests when ONNX is unavailable

Import ONNX through pytest.importorskip before using ONNX helper symbols so partial torch unit sessions skip this optional coverage instead of failing during collection.

Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
---
 tests/unit/torch/quantization/test_onnx_export_cpu.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py
index f158b1c540a..bb5a6dcb44c 100644
--- a/tests/unit/torch/quantization/test_onnx_export_cpu.py
+++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py
@@ -19,16 +19,15 @@
 import io
 
 import numpy as np
-import onnx
 import pytest
 import torch
-from onnx import TensorProto, helper, numpy_helper
 
+onnx = pytest.importorskip("onnx")
 pytest.importorskip("onnxruntime")
-
 from _test_utils.torch.misc import set_seed
 from _test_utils.torch.quantization.models import SimpleLinear
 from _test_utils.torch.quantization.onnx_export import TEST_MODELS, onnx_export_tester
+from onnx import TensorProto, helper, numpy_helper
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.tensor_quant as tensor_quant