Add TensorRT-RTX ABI EP support for ONNX quantization

CHANGELOG.rst

@@ -8,6 +8,7 @@ Changelog
 
 *Quantization*
 
+- Add opt-in TensorRT-RTX ABI Execution Provider support for ONNX calibration on Windows x64. Select it with ``--calibration_eps=NvTensorRtRtx --trt_rtx_backend=abi``; the legacy backend remains the default.
 - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint.
 
 *Megatron Framework (M-LM / M-Bridge)*

docs/source/getting_started/windows/_installation_standalone.rst

@@ -44,16 +44,26 @@ If you install ModelOpt-Windows without the extra ``[onnx]`` option, only the mi
 The Post-Training Quantization (PTQ) process for ONNX models usually involves running the base model with user-supplied inputs, a process called calibration. The user-supplied model inputs are referred to as calibration data. To perform calibration, the base model must be run using a suitable ONNX Execution Provider (EP), such as *DmlExecutionProvider* (DirectML EP) or *CUDAExecutionProvider* (CUDA EP). There are different ONNX Runtime packages for each EP:
 
 - *onnxruntime-directml* provides the DirectML EP.
-- *onnxruntime-trt-rtx* provides TensorRT-RTX EP.
+- *onnxruntime-ep-nv-tensorrt-rtx-cu13* provides the standalone TensorRT-RTX EP ABI plugin.
 - *onnxruntime-gpu* provides the CUDA EP.
 - *onnxruntime* provides the CPU EP.
 
-By default, ModelOpt-Windows installs *onnxruntime-gpu*. The default CUDA version needed for *onnxruntime-gpu* since v1.19.0 is 12.x. The *onnxruntime-gpu* package (i.e. CUDA EP) has CUDA and cuDNN dependencies:
+By default, ModelOpt-Windows on Windows x64 installs *onnxruntime-gpu* together with the
+standalone TensorRT-RTX EP ABI plugin. Keeping *onnxruntime-gpu* allows calibration to
+switch between CUDA EP and TensorRT-RTX. Select the TensorRT-RTX ABI implementation with
+``--calibration_eps NvTensorRtRtx --trt_rtx_backend abi``. The legacy backend remains
+available and uses TensorRT-RTX libraries supplied through ``PATH``.
+
+The ABI plugin requires an Ampere-or-newer RTX GPU and an NVIDIA driver with CUDA 13 support.
+Python 3.10 and Windows architectures other than x64 continue to use the legacy backend.
+
+The *onnxruntime-gpu* package (i.e. CUDA EP) has CUDA and cuDNN dependencies:
 
 - Install CUDA and cuDNN:
     - For the ONNX Runtime GPU package, you need to install the appropriate version of CUDA and cuDNN. Refer to the `CUDA Execution Provider requirements <https://onnxruntime.ai/docs/install/#cuda-and-cudnn/>`_ for compatible versions of CUDA and cuDNN.
 
-If you need to use any other EP for calibration, you can uninstall the existing *onnxruntime-gpu* package and install the corresponding package. For example, to use the DirectML EP, you can uninstall the existing *onnxruntime-gpu* package and install the *onnxruntime-directml* package:
+If you need to use an incompatible ORT package for calibration, uninstall the existing
+*onnxruntime-gpu* package before installing it. For example, to use the DirectML EP:
 
   .. code-block:: bash
 
@@ -95,9 +105,10 @@ Ensure the following steps are verified:
       - **Python Interpreter**: Open the command line and type python. The Python interpreter should start, displaying the Python version.
       - **Onnxruntime Package**: Ensure that exactly one of the following is installed:
             - *onnxruntime-directml* (DirectML EP)
-            - *onnxruntime-trt-rtx* (TensorRT-RTX EP)
             - *onnxruntime-gpu* (CUDA EP)
             - *onnxruntime* (CPU EP)
+        The *onnxruntime-ep-nv-tensorrt-rtx-cu13* plugin is installed alongside the selected
+        ONNX Runtime package; it does not replace *onnxruntime-gpu*.
       - **CUDA Toolkit**: For CUDA workflows, verify that the selected Toolkit is found first and that ``nvcc`` reports the expected major version:
 
             .. code-block:: bat

examples/windows/onnx_ptq/genai_llm/README.md

@@ -68,6 +68,7 @@ The table below lists key command-line arguments of the ONNX PTQ example script.
 | `--awqclip_alpha_min` | 0.5 (default) | Minimum AWQ weight-clipping threshold, user-defined |
 | `--awqclip_bsz_col` | 1024 (default) | Chunk size in columns during weight clipping, user-defined |
 | `--calibration_eps` | dml, cuda, cpu, NvTensorRtRtx (default: [cuda,cpu]) | List of execution-providers to use for session run during calibration |
+| `--trt_rtx_backend` | legacy (default), abi | TensorRT-RTX implementation used when calibration_eps includes NvTensorRtRtx. Legacy uses TensorRT-RTX libraries on PATH; ABI uses the standalone EP plugin. |
 | `--add_position_ids` | Default: position_ids input is disabled | Use this option to enable position_ids input in calibration data|
 | `--enable_mixed_quant` | Default: mixed-quant is disabled | Use this option to enable mixed precision quantization|
 | `--layers_8bit` | Default: None | Use this option to override default mixed-quant strategy|

examples/windows/onnx_ptq/genai_llm/quantize.py

@@ -390,6 +390,7 @@ def main(args):
         calibration_method=args.algo,
         calibration_data_reader=None if args.use_random_calib else calib_inputs,
         calibration_eps=args.calibration_eps,
+        trt_rtx_backend=args.trt_rtx_backend,
         use_zero_point=args.use_zero_point,
         block_size=args.block_size,
         input_shapes_profile=input_shapes_profile_data,
@@ -569,6 +570,12 @@ def main(args):
         default=["cuda", "cpu"],  # Default as a list
         help="Comma-separated list of calibration endpoints. Choose from 'cuda', 'cpu', 'dml', 'NvTensorRtRtx'.",
     )
+    parser.add_argument(
+        "--trt_rtx_backend",
+        choices=["legacy", "abi"],
+        default="legacy",
+        help="TensorRT-RTX backend used with NvTensorRtRtx calibration: legacy or abi.",
+    )
     parser.add_argument(
         "--trust_remote_code",
         help="Set trust_remote_code for Huggingface models and tokenizers",

modelopt/onnx/quantization/__main__.py

@@ -163,10 +163,21 @@ def get_parser() -> argparse.ArgumentParser:
         nargs="+",
         help=(
             "Priority order for the execution providers (EP) to calibrate the model. "
-            "Any subset of ['trt', 'cuda:x', dml:x, 'cpu'], where 'x' is the device id."
+            "Any subset of ['NvTensorRtRtx', 'trt', 'cuda:x', dml:x, 'cpu'], where 'x' is "
+            "the device id."
             "If a custom op is detected in the model, 'trt' will automatically be added to the EP list."
         ),
     )
+    argparser.add_argument(
+        "--trt_rtx_backend",
+        choices=["legacy", "abi"],
+        default="legacy",
+        help=(
+            "TensorRT-RTX implementation used with --calibration_eps NvTensorRtRtx. "
+            "The legacy backend uses TensorRT-RTX libraries on PATH; "
+            "the ABI backend uses the installed standalone EP plugin."
+        ),
+    )
     argparser.add_argument(
         "--override_shapes",
         type=str,
@@ -532,6 +543,7 @@ def main():
         calibration_cache_path=args.calibration_cache_path,
         calibration_shapes=args.calibration_shapes,
         calibration_eps=args.calibration_eps,
+        trt_rtx_backend=args.trt_rtx_backend,
         override_shapes=args.override_shapes,
         op_types_to_quantize=args.op_types_to_quantize,
         op_types_to_exclude=args.op_types_to_exclude,

modelopt/onnx/quantization/fp8.py

@@ -233,6 +233,7 @@ def quantize(
             calibration_eps,
             calibration_shapes,
             input_shapes_profile,
+            kwargs.get("trt_rtx_backend", "legacy"),
         )
         nodes_to_exclude.extend(matmul_nodes_to_exclude)  # type: ignore[union-attr]
         logger.debug(f"Excluding {len(matmul_nodes_to_exclude)} MatMul nodes due to GEMV pattern")
@@ -253,6 +254,7 @@ def quantize(
         custom_ops_to_quantize,
         kwargs.get("op_types_needing_output_quant"),
         input_shapes_profile,
+        kwargs.get("trt_rtx_backend", "legacy"),
     )
     logger.info(
         f"Quantizable op types in the model: {[t for t in op_types_to_quantize if t in op_types]}"

modelopt/onnx/quantization/graph_utils.py

@@ -1040,6 +1040,7 @@ def get_extended_model_outputs(
     calibration_data_reader: CalibrationDataReader,
     calibration_eps: list[str],
     input_shapes_profile: Sequence[dict[str, str]] | None = None,
+    trt_rtx_backend: str = "legacy",
 ) -> dict[str, np.ndarray]:
     """Run one inference step on an onnx model which has some intermediate tensor marked as model outputs.
 
@@ -1072,11 +1073,14 @@ def get_extended_model_outputs(
         save_onnx(extended_model, extended_onnx_path, save_as_external_data=True)
         intermediate_generated_files.append(extended_onnx_path)
         session = create_inference_session(
-            extended_onnx_path, calibration_eps, input_shapes_profile
+            extended_onnx_path, calibration_eps, input_shapes_profile, trt_rtx_backend
         )
     else:
         session = create_inference_session(
-            extended_model.SerializeToString(), calibration_eps, input_shapes_profile
+            extended_model.SerializeToString(),
+            calibration_eps,
+            input_shapes_profile,
+            trt_rtx_backend,
         )
 
     # Run extended model's inference.
@@ -1095,6 +1099,7 @@ def find_nodes_from_matmul_to_exclude(
     calibration_eps: list[str] = ["cpu", "cuda:0", "trt"],
     calibration_shapes: str | dict | None = None,
     input_shapes_profile: Sequence[dict[str, str]] | None = None,
+    trt_rtx_backend: str = "legacy",
 ) -> list[str]:
     """Find MatMul nodes that meet gemv or small-gemm conditions and should be excluded.
 
@@ -1147,6 +1152,7 @@ def find_nodes_from_matmul_to_exclude(
             calibration_data_reader,
             calibration_eps,
             input_shapes_profile,
+            trt_rtx_backend,
         )
 
     logger.debug(f"Matmul nodes to exclude: {nodes_to_exclude}")
@@ -1365,6 +1371,7 @@ def _exclude_matmuls_by_inference(
     calibration_data_reader: CalibrationDataReader,
     calibration_eps: list[str],
     input_shapes_profile: Sequence[dict[str, str]] | None = None,
+    trt_rtx_backend: str = "legacy",
 ) -> list[str]:
     """Use actual inference to find MatMuls with dimension 1 or small K/N."""
     # Add matmul outputs and second-input outputs to model outputs
@@ -1389,6 +1396,7 @@ def _exclude_matmuls_by_inference(
         calibration_data_reader,
         calibration_eps,
         input_shapes_profile,
+        trt_rtx_backend,
     )
 
     nodes_to_exclude = []
@@ -1432,6 +1440,7 @@ def find_nodes_from_mha_to_exclude(
     calibration_data_reader: CalibrationDataReader = None,
     calibration_eps: list[str] = ["cpu", "cuda:0", "trt"],
     input_shapes_profile: Sequence[dict[str, str]] | None = None,
+    trt_rtx_backend: str = "legacy",
 ) -> list[str]:
     """Find MatMul nodes in MHA pattern to exclude.
 
@@ -1493,6 +1502,7 @@ def find_nodes_from_mha_to_exclude(
             calibration_data_reader,
             calibration_eps,
             input_shapes_profile,
+            trt_rtx_backend,
         )
 
         # For each MHA block,

modelopt/onnx/quantization/int4.py

@@ -557,7 +557,10 @@ def _quantize_awq_clip(
 
         # Creating inference session and preparing inputs for calibration
         session = create_inference_session(
-            augmented_onnx_path, calibration_eps, input_shapes_profile
+            augmented_onnx_path,
+            calibration_eps,
+            input_shapes_profile,
+            kwargs.get("trt_rtx_backend", "legacy"),
         )
         inputs = []
         for inp_d in data_reader:
@@ -1114,7 +1117,10 @@ def _quantize_awq_lite(
 
         # Creating inference session and preparing inputs for calibration
         session = create_inference_session(
-            augmented_onnx_path, calibration_eps, input_shapes_profile
+            augmented_onnx_path,
+            calibration_eps,
+            input_shapes_profile,
+            kwargs.get("trt_rtx_backend", "legacy"),
         )
         inputs = []
         for inp_d in data_reader:

modelopt/onnx/quantization/int8.py

@@ -180,6 +180,7 @@ def quantize(
             calibration_eps,
             calibration_shapes,
             input_shapes_profile,
+            kwargs.get("trt_rtx_backend", "legacy"),
         )
         nodes_to_exclude.extend(matmul_nodes_to_exclude)  # type: ignore[union-attr]
         logger.debug(f"Excluding {len(matmul_nodes_to_exclude)} MatMul nodes due to GEMV pattern")
@@ -205,6 +206,7 @@ def quantize(
         custom_ops_to_quantize,
         kwargs.get("op_types_needing_output_quant"),
         input_shapes_profile,
+        kwargs.get("trt_rtx_backend", "legacy"),
     )
     logger.info(f"Quantizable op types: {[t for t in quantizable_op_types if t in op_types]}")
 

modelopt/onnx/quantization/ort_patching.py

@@ -292,6 +292,10 @@ def _create_inference_session_with_ep_config(calibrator, **kwargs):
     model_path = kwargs.get("model_path")
     logger.debug("Creating inference session with Execution Provider configuration")
 
+    trt_rtx_backend = kwargs.get("trt_rtx_backend", "legacy")
+    if trt_rtx_backend not in ("legacy", "abi"):
+        raise ValueError(f"trt_rtx_backend must be 'legacy' or 'abi', got {trt_rtx_backend!r}")
+
     sess_options = ort.SessionOptions()
     sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
     sess_options.add_session_config_entry("session.use_device_allocator_for_initializers", "1")
@@ -302,8 +306,7 @@ def _create_inference_session_with_ep_config(calibrator, **kwargs):
 
     # Note. This path can be an empty string, which denotes that the model has custom ops and TRT EP is needed.
     calibrator.trt_extra_plugin_lib_paths = kwargs.get("trt_extra_plugin_lib_paths")
-
-    if calibrator.trt_extra_plugin_lib_paths is not None:
+    if trt_rtx_backend != "abi" and calibrator.trt_extra_plugin_lib_paths is not None:
         logger.debug(f"TRT extra plugin paths: {calibrator.trt_extra_plugin_lib_paths}")
         if "TensorrtExecutionProvider" not in ort.get_available_providers():
             raise RuntimeError(
@@ -336,20 +339,12 @@ def _update_provider_config(provider, config):
                 providers[i], {"arena_extend_strategy": "kSameAsRequested"}
             )
 
-    if model_path is None:
-        # Create the inference session with EP configuration on augmented_model
-        calibrator.infer_session = ort.InferenceSession(
-            calibrator.augmented_model_path,
-            sess_options=sess_options,
-            providers=providers,
-        )
-    else:
-        # Create the inference session with EP configuration on provided model path
-        calibrator.infer_session = ort.InferenceSession(
-            model_path,
-            sess_options=sess_options,
-            providers=providers,
-        )
+    session_path = calibrator.augmented_model_path if model_path is None else model_path
+    calibrator.infer_session = ort.InferenceSession(
+        session_path,
+        sess_options=sess_options,
+        providers=providers,
+    )
 
     # Group qdq tensors will have the same scaling factor.
     calibrator.group_qdq_tensors = kwargs.get("group_qdq_tensors")
@@ -1568,6 +1563,8 @@ def _quantize_static(
             ExecutionProviders = list[string] :
                 Default is [("CUDAExecutionProvider", {"device_id": 0}), "CPUExecutionProvider",
                 "TensorrtExecutionProvider"]
+            TrtRtxBackend = string :
+                Selects the legacy or ABI TensorRT-RTX execution provider implementation.
     """
     logger.info("Starting static quantization")
     logger.debug(f"Quantization format: {quant_format}")
@@ -1607,6 +1604,7 @@ def _quantize_static(
         # ====================== Modification ======================
         ("TrtExtraPluginLibraryPaths", "trt_extra_plugin_lib_paths"),
         ("ExecutionProviders", "execution_providers"),
+        ("TrtRtxBackend", "trt_rtx_backend"),
         ("group_qdq_tensors", "group_qdq_tensors"),
         ("QDQDisableWeightAdjustForInt32Bias", "disable_int32_weight_adjustment"),
         # ==========================================================