NVIDIA · ajrasane · Sep 4, 2026
@@ -64,7 +64,7 @@ jobs:
 
           # nvcr.io/nvidia/tensorrt:26.05-py3 ships cuDNN 9.22 with no preinstalled torch, and
           # torch 2.14 pins cuDNN 9.24: mixing them fails with CUDNN_SUBLIBRARY_LOADING_FAILED.
-          if [[ "${{ inputs.docker_image }}" == *"/tensorrt:"* ]]; then
+          if [[ "${{ inputs.docker_image }}" == *"/tensorrt:26.05-"* ]]; then
             echo "torch<2.14" > /tmp/pip-constraints.txt
             export PIP_CONSTRAINT=/tmp/pip-constraints.txt
           fi
@@ -79,6 +79,22 @@ jobs:
           fi
 
           find examples/${{ inputs.example }} -name "requirements.txt" | while read req_file; do python -m pip install -r "$req_file" || exit 1; done
+
+          if [[ "${{ inputs.example }}" == "torch_onnx" ]]; then
+            # Prefer the CUDA libraries bundled with PyPI torch to avoid mixing cuDNN sublibraries.
+            torch_lib_path=$(python - <<'PY'
+          from pathlib import Path
+
+          import torch
+
+          root = Path(torch.__file__).resolve().parent.parent / "nvidia"
+          paths = sorted(str(path) for path in root.glob("*/lib") if path.is_dir())
+          assert paths
+          print(":".join(paths))
+          PY
+            )
+            echo "LD_LIBRARY_PATH=${torch_lib_path}:${LD_LIBRARY_PATH}" >> "$GITHUB_ENV"
+          fi
       - name: Run tests
         id: run_tests
         continue-on-error: ${{ inputs.allow_failure }}

@@ -168,18 +168,26 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        example: [diffusers, torch_onnx, torch_trt]
+        include:
+          - example: diffusers
+            docker_image: "nvcr.io/nvidia/tensorrt:26.05-py3"
+            timeout_minutes: 45
+          - example: torch_onnx
+            docker_image: "nvcr.io/nvidia/tensorrt:26.06-py3"
+            timeout_minutes: 90
+          - example: torch_trt
+            docker_image: "nvcr.io/nvidia/tensorrt:26.05-py3"
+            timeout_minutes: 45
     uses: ./.github/workflows/_example_tests_runner.yml
     permissions:
       contents: read
     secrets: inherit
     with:
-      # Pinned to 26.05 (TensorRT 10): torch-tensorrt is capped at <2.13 (== 2.12.1),
-      # which needs libnvinfer.so.10; newer tensorrt containers drop it. Bump only once
-      # a torch-tensorrt build for the newer TensorRT is available.
-      docker_image: "nvcr.io/nvidia/tensorrt:26.05-py3"
+      # torch_trt stays on TensorRT 10 because torch-tensorrt needs libnvinfer.so.10;
+      # torch_onnx uses TensorRT 11 for W4A4 NVFP4 engine coverage.
+      docker_image: ${{ matrix.docker_image }}
       example: ${{ matrix.example }}
-      timeout_minutes: 45
+      timeout_minutes: ${{ matrix.timeout_minutes }}
       pip_install_extras: "[onnx,hf,dev-test]"
       runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}

@@ -0,0 +1,58 @@
+# 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.
+
+from pathlib import Path
+
+import onnx
+
+from modelopt.onnx.quantization.ort_utils import _check_for_trtexec
+from modelopt.onnx.utils import has_node_op_type
+
+DYNAMIC_NVFP4_OP = "TRT_FP4DynamicQuantize"
+DYNAMIC_NVFP4_MIN_TRT_VERSION = "11.0"
+W4A16_NVFP4_RECIPE = "w4a16_nvfp4"
+
+_DYNAMIC_NVFP4_FORMATS = {"nvfp4", "nvfp4_awq_lite"}
+_DYNAMIC_NVFP4_TRT_ERROR = (
+    "Dynamic NVFP4 (W4A4) TensorRT engine builds require TensorRT 11.0 or newer. "
+    "Upgrade TensorRT, or re-export with "
+    "`--quantize_mode=nvfp4 --recipe=w4a16_nvfp4` to use the weight-only NVFP4 "
+    "recipe on TensorRT 10.16."
+)
+
+
+def request_uses_dynamic_nvfp4(
+    quantize_mode: str, recipe: str | None, auto_quantization_formats: list[str]
+) -> bool:
+    """Return whether the requested quantization can emit dynamic NVFP4 activations."""
+    if quantize_mode == "nvfp4":
+        return recipe != W4A16_NVFP4_RECIPE
+    return quantize_mode == "auto" and bool(
+        _DYNAMIC_NVFP4_FORMATS.intersection(auto_quantization_formats)
+    )
+
+
+def onnx_uses_dynamic_nvfp4(onnx_path: str | Path) -> bool:
+    """Return whether an ONNX graph contains dynamic NVFP4 activation quantization."""
+    model = onnx.load(str(onnx_path), load_external_data=False)
+    return has_node_op_type(model.graph, DYNAMIC_NVFP4_OP)
+
+
+def check_dynamic_nvfp4_trt_support() -> None:
+    """Require a TensorRT release that reliably compiles dynamic NVFP4 graphs."""
+    try:
+        _check_for_trtexec(min_version=DYNAMIC_NVFP4_MIN_TRT_VERSION)
+    except ImportError as e:
+        raise ImportError(f"{_DYNAMIC_NVFP4_TRT_ERROR} ({e})") from e
@@ -17,6 +17,7 @@
 import csv
 
 import timm
+from _trt_compat import check_dynamic_nvfp4_trt_support, onnx_uses_dynamic_nvfp4
 from evaluation import evaluate
 
 from modelopt.torch._deploy._runtime import RuntimeRegistry
@@ -25,7 +26,13 @@
 
 
 def main():
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(
+        description=(
+            "Dynamic NVFP4 (W4A4) models require TensorRT 11.0 or newer. "
+            "For TensorRT 10.16, re-export with "
+            "--quantize_mode=nvfp4 --recipe=w4a16_nvfp4."
+        )
+    )
     parser.add_argument(
         "--onnx_path",
         type=str,
@@ -80,6 +87,12 @@ def main():
     )
 
     args = parser.parse_args()
+    if onnx_uses_dynamic_nvfp4(args.onnx_path):
+        try:
+            check_dynamic_nvfp4_trt_support()
+        except ImportError as e:
+            parser.error(str(e))
+
     deployment = {
         "runtime": "TRT",
         "precision": args.engine_precision,

@@ -74,6 +74,25 @@ Quantization configs are loaded from the YAML preset recipes under
 `--recipe=<preset basename or path to a QuantizeConfig YAML>` to use a different
 recipe (e.g. `--recipe=nvfp4_awq_lite` or `--recipe=/path/to/my_quant_cfg.yaml`).
 
+### TensorRT Compatibility
+
+Building a dynamic W4A4 NVFP4 vision model requires TensorRT 11.0 or later.
+TensorRT 10.16 users can explicitly select the validated weight-only W4A16
+NVFP4 recipe instead:
+
+```bash
+python torch_quant_to_onnx.py \
+    --timm_model_name=vit_small_patch16_224 \
+    --quantize_mode=nvfp4 \
+    --recipe=w4a16_nvfp4 \
+    --onnx_save_path=vit_small_patch16_224.w4a16_nvfp4.onnx \
+    --trt_build
+```
+
+The fallback changes the quantization behavior: Linear weights use NVFP4 while
+their activations remain in higher precision. The existing FP8 Conv2d override
+still applies.
+
 ### Conv2d Quantization Override
 
 TensorRT only supports FP8 and INT8 for convolution operations. When quantizing models with Conv2d layers (like SwinTransformer), the script automatically applies the following overrides:
@@ -88,7 +107,8 @@ TensorRT only supports FP8 and INT8 for convolution operations. When quantizing
 
 If the input model is of type image classification, use the following script to evaluate it. The script automatically downloads and uses the [ILSVRC/imagenet-1k](https://huggingface.co/datasets/ILSVRC/imagenet-1k) dataset from Hugging Face. This gated repository requires authentication via Hugging Face access token. See <https://huggingface.co/docs/hub/en/security-tokens> for details.
 
-> *Note: TensorRT 10.11 or later is required to evaluate the MXFP8 or NVFP4 ONNX models.*
+> *Note: TensorRT 10.11 or later is required to evaluate MXFP8 ONNX models. W4A4
+> NVFP4 vision models require TensorRT 11.0 or later.*
 
 ```bash
 python ../onnx_ptq/evaluate.py \

@@ -29,6 +29,7 @@
 import torch
 import torch.multiprocessing as mp
 import torch.nn.functional as F
+from _trt_compat import check_dynamic_nvfp4_trt_support, request_uses_dynamic_nvfp4
 from datasets import load_dataset
 from download_example_onnx import export_to_onnx
 from evaluation import evaluate
@@ -525,7 +526,11 @@ def main():
     parser.add_argument(
         "--trt_build",
         action="store_true",
-        help="Build a TensorRT engine from the exported ONNX model using trtexec.",
+        help=(
+            "Build a TensorRT engine from the exported ONNX model using trtexec. "
+            "Dynamic NVFP4 (W4A4) builds require TensorRT 11.0 or newer. "
+            "For TensorRT 10.16, use --quantize_mode=nvfp4 --recipe=w4a16_nvfp4."
+        ),
     )
     parser.add_argument(
         "--no_pretrained",
@@ -546,6 +551,15 @@ def main():
             "--recipe is not supported with --quantize_mode=auto; "
             "use --auto_quantization_formats instead."
         )
+    if args.trt_build and request_uses_dynamic_nvfp4(
+        args.quantize_mode,
+        args.recipe,
+        args.auto_quantization_formats,
+    ):
+        try:
+            check_dynamic_nvfp4_trt_support()
+        except ImportError as e:
+            parser.error(str(e))
 
     # Create model and move to appropriate device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

@@ -105,7 +105,7 @@ def _add_initializer(initializer):
         if initializer.name not in initializer_indices:
             graph.initializer.append(initializer)
 
-    def _add_input_value_info(graph, tensor_proto):
+    def _add_initializer_value_info(graph, tensor_proto):
         assert tensor_proto.name not in graph_inputs, (
             f"{tensor_proto.name} already in graph inputs."
         )
@@ -116,7 +116,7 @@ def _add_input_value_info(graph, tensor_proto):
         value_info = onnx.helper.make_tensor_value_info(
             tensor_proto.name, tensor_proto.data_type, tensor_proto.dims
         )
-        graph.input.append(value_info)
+        graph.value_info.append(value_info)
 
     # Remove the original node from the graph
     graph.node.remove(node)
@@ -148,9 +148,9 @@ def _add_input_value_info(graph, tensor_proto):
     )
 
     # Add ValueInfo for the initializers if not present
-    _add_input_value_info(graph, w_f4_proto)
-    _add_input_value_info(graph, sw_f32_per_tensor_proto)
-    _add_input_value_info(graph, sw_f8_per_block_proto)
+    _add_initializer_value_info(graph, w_f4_proto)
+    _add_initializer_value_info(graph, sw_f32_per_tensor_proto)
+    _add_initializer_value_info(graph, sw_f8_per_block_proto)
 
     # Add the initializers to the graph
     _add_initializer(w_f4_proto)

@@ -21,6 +21,7 @@
 import tempfile
 import uuid
 from collections import defaultdict
+from collections.abc import Collection
 from typing import Any
 
 import numpy as np
@@ -128,6 +129,21 @@ def get_node_names(model: onnx.ModelProto) -> list[str]:
     return [node.name for node in model.graph.node]
 
 
+def has_node_op_type(graph: onnx.GraphProto, op_type: str) -> bool:
+    """Return whether a graph or any nested subgraph contains an operator type."""
+    for node in graph.node:
+        if node.op_type == op_type:
+            return True
+        for attr in node.attribute:
+            if attr.type == onnx.AttributeProto.GRAPH:
+                if has_node_op_type(attr.g, op_type):
+                    return True
+            elif attr.type == onnx.AttributeProto.GRAPHS:
+                if any(has_node_op_type(subgraph, op_type) for subgraph in attr.graphs):
+                    return True
+    return False
+
+
 def _get_tensor_shape(tensor: onnx.ValueInfoProto) -> list[int]:
     """This function returns the shape of the input onnx tensor.
 
@@ -1859,18 +1875,25 @@ def remove_node_training_mode(onnx_model: onnx.ModelProto, node_op_type: str) ->
     return onnx_model
 
 
-def change_casts_to_fp16(model: onnx.ModelProto, target_op_types: list[str]) -> onnx.ModelProto:
-    """Change FP16-to-FP32 Cast nodes whose entire fanout feeds target ops to cast to FP16 instead.
+def change_casts_to_fp16(
+    model: onnx.ModelProto,
+    target_op_types: list[str],
+    source_types: Collection[int] | None = None,
+) -> onnx.ModelProto:
+    """Retarget eligible Cast nodes whose entire fanout feeds target ops to FP16.
 
     Args:
         model: The ONNX model to modify.
         target_op_types: List of op types to check for. Cast nodes feeding exclusively into
             these will be changed from FP32 to FP16.
+        source_types: Source element types eligible for retargeting. Defaults to FP16.
 
     Returns:
         The modified ONNX model with Cast nodes updated.
     """
     type_map = _build_tensor_type_map(model)
+    if source_types is None:
+        source_types = {onnx.TensorProto.FLOAT16}
 
     # Build a map of tensor name -> consumer nodes
     tensor_to_consumers: dict[str, list[onnx.NodeProto]] = {}
@@ -1879,17 +1902,16 @@ def change_casts_to_fp16(model: onnx.ModelProto, target_op_types: list[str]) ->
             if inp:
                 tensor_to_consumers.setdefault(inp, []).append(node)
 
-    # Find Cast nodes that feed into target ops and change FP16->FP32 to FP16->FP16
+    # Find Cast nodes that feed into target ops and change their destination to FP16
     for node in model.graph.node:
         if node.op_type != "Cast":
             continue
 
-        # Only retarget FP16->FP32 casts; leave other casts (e.g. FP64->FP32) alone
         cast_to = get_cast_to_type(node)
         if cast_to != onnx.TensorProto.FLOAT:
             continue
         source_type = type_map.get(node.input[0])
-        if source_type != onnx.TensorProto.FLOAT16:
+        if source_type not in source_types:
             continue
 
         # Only change when ALL consumers are target ops to avoid breaking non-target branches

@@ -180,10 +180,6 @@ def initialize_input_output_tensors(self, engine):
                         tensor_name,
                         output_tensors[idx - len(input_tensors)].data_ptr(),
                     )
-            assert self.execution_context.all_shape_inputs_specified, (
-                "Not all shape inputs are specified."
-            )
-
             # Set selected profile idx
             self.execution_context.set_optimization_profile_async(0, self.stream.cuda_stream)