[6287717][ONNX][Quantization] Preserve trt.plugins custom-op value_info in clear_stale_value_info

modelopt/onnx/utils.py

@@ -1868,14 +1868,15 @@ def clear_stale_value_info(model: onnx.ModelProto) -> int:
     Walks every ``Cast`` node and forces the ``elem_type`` of any
     ``graph.output`` entry produced by that Cast to match the Cast's ``to``
     attribute (the spec-defined contract for a Cast's output dtype). Then
-    clears ``value_info`` wholesale so ORT/shape-inference re-derives
-    intermediate-tensor types from the operator graph during session setup.
+    clears ``value_info`` so ORT/shape-inference re-derives intermediate-tensor
+    types from the operator graph during session setup -- except entries for
+    outputs of ``trt.plugins`` custom-op nodes, whose types ORT cannot infer.
 
     Args:
         model: Loaded in-memory onnx ModelProto.
 
     Returns:
-        Total number of entries reconciled or cleared.
+        Number of Cast outputs reconciled plus value_info entries cleared.
     """
     cast_to_by_output = {
         node.output[0]: get_cast_to_type(node)
@@ -1890,7 +1891,14 @@ def clear_stale_value_info(model: onnx.ModelProto) -> int:
             o.type.tensor_type.elem_type = to_attr
             fixed_outputs += 1
 
-    n_vi = len(model.graph.value_info)
-    if n_vi:
+    # Outputs of TensorRT-plugin nodes carry types ORT cannot infer so they must survive the
+    # value_info clear, otherwise ORT fails output type inference for the custom op.
+    preserve_names = {
+        out for node in model.graph.node if node.domain == "trt.plugins" for out in node.output
+    }
+    preserved = [vi for vi in model.graph.value_info if vi.name in preserve_names]
+    n_cleared = len(model.graph.value_info) - len(preserved)
+    if n_cleared:
         del model.graph.value_info[:]
-    return fixed_outputs + n_vi
+        model.graph.value_info.extend(preserved)
+    return fixed_outputs + n_cleared

tests/gpu/onnx/quantization/test_plugin.py

@@ -127,6 +127,29 @@ def test_trt_plugin_quantization(tmp_path):
         assert assert_nodes_are_quantized(quantizable_nodes)
 
 
+def test_trt_plugin_quantization_int4_awq(tmp_path):
+    model = _create_test_model_trt()
+    with open(os.path.join(tmp_path, "model_with_trt_plugin_int4.onnx"), "w") as f:
+        onnx.save_model(model, f.name)
+
+        # Quantize at int4 with awq_clip (the path that forces opset >= 21).
+        quantize(
+            f.name,
+            quantize_mode="int4",
+            calibration_method="awq_clip",
+            calibration_eps=["trt", "cuda:0", "cpu"],
+        )
+
+        # The regression was a hard failure at calibration-session load; reaching a
+        # written output model means the custom op's type survived the value_info clear.
+        output_onnx_path = f.name.replace(".onnx", ".quant.onnx")
+        assert os.path.isfile(output_onnx_path)
+
+        # The custom op must still be present (not dropped) in the quantized model.
+        graph = gs.import_onnx(onnx.load(output_onnx_path))
+        assert any(n.op == "CustomSkipLayerNormPluginDynamic" for n in graph.nodes)
+
+
 def test_trt_plugin_autocast(tmp_path):
     model = _create_test_model_trt()
     with open(os.path.join(tmp_path, "model_with_trt_plugin_autocast.onnx"), "w") as f: