[6421642][ONNX][Quantization] Fix NVFP4 exporter node ordering

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
@@ -421,4 +426,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")
 
+        utils.topologically_sort_graph_nodes(graph)
+
         return onnx_model

modelopt/onnx/utils.py

@@ -2020,3 +2020,89 @@ 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.
+    """
+
+    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):
+        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):
+        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
+            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)

tests/unit/torch/quantization/test_onnx_export_cpu.py

@@ -15,13 +15,25 @@
 
 """Unit tests for ONNX export for CPU quantization."""
 
+import inspect
+import io
+
+import numpy as np
 import pytest
 import torch
 
+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
+from modelopt.onnx import utils
+from modelopt.onnx.export import NVFP4QuantExporter
+from modelopt.torch.quantization.utils import is_quantized_linear
 
 
 @pytest.mark.parametrize("model_cls", TEST_MODELS)
@@ -42,3 +54,157 @@ 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)
+
+
+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)
+
+
+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)