Voice et bot modif

venv/lib/python3.12/site-packages/onnxruntime/ThirdPartyNotices.txt

@@ -5806,41 +5806,6 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 _____
 
-composable_kernel
-
-https://github.com/ROCmSoftwarePlatform/composable_kernel
-
-Copyright (c) 2018-    , Advanced Micro Devices, Inc. (Chao Liu, Jing Zhang)
-Copyright (c) 2019-    , Advanced Micro Devices, Inc. (Letao Qin, Qianfeng Zhang, Liang Huang, Shaojie Wang)
-Copyright (c) 2022-    , Advanced Micro Devices, Inc. (Anthony Chang, Chunyu Lai, Illia Silin, Adam Osewski, Poyen Chen, Jehandad Khan)
-Copyright (c) 2019-2021, Advanced Micro Devices, Inc. (Hanwen Chang)
-Copyright (c) 2019-2020, Advanced Micro Devices, Inc. (Tejash Shah)
-Copyright (c) 2020     , Advanced Micro Devices, Inc. (Xiaoyan Zhou)
-Copyright (c) 2021-2022, Advanced Micro Devices, Inc. (Jianfeng Yan)
-
-SPDX-License-Identifier: MIT
-Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-
-_____
-
 neural-speed
 
 https://github.com/intel/neural-speed

venv/lib/python3.12/site-packages/onnxruntime/__init__.py

@@ -8,7 +8,9 @@ For more information on ONNX Runtime, please see `aka.ms/onnxruntime <https://ak
 or the `Github project <https://github.com/microsoft/onnxruntime/>`_.
 """
 
-__version__ = "1.23.2"
+import contextlib
+
+__version__ = "1.26.0"
 __author__ = "Microsoft"
 
 # we need to do device version validation (for example to check Cuda version for an onnxruntime-training package).
@@ -32,6 +34,8 @@ try:
         OrtArenaCfg,  # noqa: F401
         OrtCompileApiFlags,  # noqa: F401
         OrtDeviceMemoryType,  # noqa: F401
+        OrtEpAssignedNode,  # noqa: F401
+        OrtEpAssignedSubgraph,  # noqa: F401
         OrtEpDevice,  # noqa: F401
         OrtExecutionProviderDevicePolicy,  # noqa: F401
         OrtExternalInitializerInfo,  # noqa: F401
@@ -79,6 +83,7 @@ from onnxruntime.capi.onnxruntime_inference_collection import (
     IOBinding,  # noqa: F401
     ModelCompiler,  # noqa: F401
     OrtDevice,  # noqa: F401
+    OrtDeviceVendorId,  # noqa: F401
     OrtValue,  # noqa: F401
     SparseTensor,  # noqa: F401
     copy_tensors,  # noqa: F401
@@ -133,14 +138,43 @@ def _get_package_root(package_name: str, directory_name: str | None = None):
     return None
 
 
+def _extract_cuda_major_version(version_str: str) -> str:
+    """Extract CUDA major version from version string (e.g., '12.1' -> '12').
+
+    Args:
+        version_str: CUDA version string to parse
+
+    Returns:
+        Major version as string, or "12" if parsing fails
+    """
+    return version_str.split(".")[0] if version_str else "12"
+
+
+def _get_cufft_version(cuda_major: str) -> str:
+    """Get cufft library version based on CUDA major version.
+
+    Args:
+        cuda_major: CUDA major version as string (e.g., "12", "13")
+
+    Returns:
+        cufft version as string
+    """
+    # cufft versions: CUDA 12.x -> 11, CUDA 13.x -> 12
+    return "12" if cuda_major == "13" else "11"
+
+
 def _get_nvidia_dll_paths(is_windows: bool, cuda: bool = True, cudnn: bool = True):
+    # Dynamically determine CUDA major version from build info
+    cuda_major_version = _extract_cuda_major_version(cuda_version)
+    cufft_version = _get_cufft_version(cuda_major_version)
+
     if is_windows:
         # Path is relative to site-packages directory.
         cuda_dll_paths = [
-            ("nvidia", "cublas", "bin", "cublasLt64_12.dll"),
-            ("nvidia", "cublas", "bin", "cublas64_12.dll"),
-            ("nvidia", "cufft", "bin", "cufft64_11.dll"),
-            ("nvidia", "cuda_runtime", "bin", "cudart64_12.dll"),
+            ("nvidia", "cublas", "bin", f"cublasLt64_{cuda_major_version}.dll"),
+            ("nvidia", "cublas", "bin", f"cublas64_{cuda_major_version}.dll"),
+            ("nvidia", "cufft", "bin", f"cufft64_{cufft_version}.dll"),
+            ("nvidia", "cuda_runtime", "bin", f"cudart64_{cuda_major_version}.dll"),
         ]
         cudnn_dll_paths = [
             ("nvidia", "cudnn", "bin", "cudnn_engines_runtime_compiled64_9.dll"),
@@ -154,12 +188,12 @@ def _get_nvidia_dll_paths(is_windows: bool, cuda: bool = True, cudnn: bool = Tru
     else:  # Linux
         # cublas64 depends on cublasLt64, so cublasLt64 should be loaded first.
         cuda_dll_paths = [
-            ("nvidia", "cublas", "lib", "libcublasLt.so.12"),
-            ("nvidia", "cublas", "lib", "libcublas.so.12"),
-            ("nvidia", "cuda_nvrtc", "lib", "libnvrtc.so.12"),
+            ("nvidia", "cublas", "lib", f"libcublasLt.so.{cuda_major_version}"),
+            ("nvidia", "cublas", "lib", f"libcublas.so.{cuda_major_version}"),
+            ("nvidia", "cuda_nvrtc", "lib", f"libnvrtc.so.{cuda_major_version}"),
             ("nvidia", "curand", "lib", "libcurand.so.10"),
-            ("nvidia", "cufft", "lib", "libcufft.so.11"),
-            ("nvidia", "cuda_runtime", "lib", "libcudart.so.12"),
+            ("nvidia", "cufft", "lib", f"libcufft.so.{cufft_version}"),
+            ("nvidia", "cuda_runtime", "lib", f"libcudart.so.{cuda_major_version}"),
         ]
 
         # Do not load cudnn sub DLLs (they will be dynamically loaded later) to be consistent with PyTorch in Linux.
@@ -201,15 +235,17 @@ def print_debug_info():
 
     if cuda_version:
         # Print version of installed packages that is related to CUDA or cuDNN DLLs.
+        cuda_major = _extract_cuda_major_version(cuda_version)
+
         packages = [
             "torch",
-            "nvidia-cuda-runtime-cu12",
-            "nvidia-cudnn-cu12",
-            "nvidia-cublas-cu12",
-            "nvidia-cufft-cu12",
-            "nvidia-curand-cu12",
-            "nvidia-cuda-nvrtc-cu12",
-            "nvidia-nvjitlink-cu12",
+            f"nvidia-cuda-runtime-cu{cuda_major}",
+            f"nvidia-cudnn-cu{cuda_major}",
+            f"nvidia-cublas-cu{cuda_major}",
+            f"nvidia-cufft-cu{cuda_major}",
+            f"nvidia-curand-cu{cuda_major}",
+            f"nvidia-cuda-nvrtc-cu{cuda_major}",
+            f"nvidia-nvjitlink-cu{cuda_major}",
         ]
         for package in packages:
             directory_name = "nvidia" if package.startswith("nvidia-") else None
@@ -220,9 +256,9 @@ def print_debug_info():
                 print(f"{package} not installed")
 
     if platform.system() == "Windows":
-        print(f"\nEnvironment variable:\nPATH={os.environ['PATH']}")
+        print(f"\nEnvironment variable:\nPATH={os.environ.get('PATH', '(unset)')}")
     elif platform.system() == "Linux":
-        print(f"\nEnvironment variable:\nLD_LIBRARY_PATH={os.environ['LD_LIBRARY_PATH']}")
+        print(f"\nEnvironment variable:\nLD_LIBRARY_PATH={os.environ.get('LD_LIBRARY_PATH', '(unset)')}")
 
     if importlib.util.find_spec("psutil"):
 
@@ -254,7 +290,7 @@ def print_debug_info():
 
 
 def preload_dlls(cuda: bool = True, cudnn: bool = True, msvc: bool = True, directory=None):
-    """Preload CUDA 12.x and cuDNN 9.x DLLs in Windows or Linux, and MSVC runtime DLLs in Windows.
+    """Preload CUDA 12.x+ and cuDNN 9.x DLLs in Windows or Linux, and MSVC runtime DLLs in Windows.
 
        When the installed PyTorch is compatible (using same major version of CUDA and cuDNN),
        there is no need to call this function if `import torch` is done before `import onnxruntime`.
@@ -289,30 +325,53 @@ def preload_dlls(cuda: bool = True, cudnn: bool = True, msvc: bool = True, direc
             print("Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.")
             print("It can be downloaded at https://aka.ms/vs/17/release/vc_redist.x64.exe.")
 
-    if not (cuda_version and cuda_version.startswith("12.")) and (cuda or cudnn):
-        print(
-            f"\033[33mWARNING: {package_name} is not built with CUDA 12.x support. "
-            "Please install a version that supports CUDA 12.x, or call preload_dlls with cuda=False and cudnn=False.\033[0m"
-        )
-        return
-
-    if not (cuda_version and cuda_version.startswith("12.") and (cuda or cudnn)):
+    # Check if CUDA version is supported (12.x or 13.x+)
+    ort_cuda_major = None
+    if cuda_version:
+        try:
+            ort_cuda_major = int(cuda_version.split(".")[0])
+            if ort_cuda_major < 12 and (cuda or cudnn):
+                print(
+                    f"\033[33mWARNING: {package_name} is built with CUDA {cuda_version}, which is not supported for preloading. "
+                    f"CUDA 12.x or newer is required. Call preload_dlls with cuda=False and cudnn=False.\033[0m"
+                )
+                return
+        except ValueError:
+            print(
+                f"\033[33mWARNING: Unable to parse CUDA version '{cuda_version}'. "
+                "Skipping DLL preloading. Call preload_dlls with cuda=False and cudnn=False.\033[0m"
+            )
+            return
+    elif cuda or cudnn:
+        # No CUDA version info available but CUDA/cuDNN preloading requested
         return
 
     is_cuda_cudnn_imported_by_torch = False
 
     if is_windows:
         torch_version = _get_package_version("torch")
-        is_torch_for_cuda_12 = torch_version and "+cu12" in torch_version
+        # Check if torch CUDA version matches onnxruntime CUDA version
+        torch_cuda_major = None
+        if torch_version and "+cu" in torch_version:
+            with contextlib.suppress(ValueError):
+                # Extract CUDA version from torch (e.g., "2.0.0+cu121" -> 12)
+                cu_part = torch_version.split("+cu")[1]
+                torch_cuda_major = int(cu_part[:2])  # First 2 digits are major version
+
+        is_torch_cuda_compatible = (
+            torch_cuda_major == ort_cuda_major if (torch_cuda_major and ort_cuda_major) else False
+        )
+
         if "torch" in sys.modules:
-            is_cuda_cudnn_imported_by_torch = is_torch_for_cuda_12
-            if (torch_version and "+cu" in torch_version) and not is_torch_for_cuda_12:
+            is_cuda_cudnn_imported_by_torch = is_torch_cuda_compatible
+            if torch_cuda_major and ort_cuda_major and torch_cuda_major != ort_cuda_major:
                 print(
-                    f"\033[33mWARNING: The installed PyTorch {torch_version} does not support CUDA 12.x. "
-                    f"Please install PyTorch for CUDA 12.x to be compatible with {package_name}.\033[0m"
+                    f"\033[33mWARNING: The installed PyTorch {torch_version} uses CUDA {torch_cuda_major}.x, "
+                    f"but {package_name} is built with CUDA {ort_cuda_major}.x. "
+                    f"Please install PyTorch for CUDA {ort_cuda_major}.x to be compatible.\033[0m"
                 )
 
-        if is_torch_for_cuda_12 and directory is None:
+        if is_torch_cuda_compatible and directory is None:
             torch_root = _get_package_root("torch", "torch")
             if torch_root:
                 directory = os.path.join(torch_root, "lib")

venv/lib/python3.12/site-packages/onnxruntime/backend/backend.py

@@ -17,6 +17,29 @@ from onnx.checker import check_model
 from onnxruntime import InferenceSession, SessionOptions, get_available_providers, get_device
 from onnxruntime.backend.backend_rep import OnnxRuntimeBackendRep
 
+# Allowlist of SessionOptions attributes that are safe to set via the backend API.
+# Dangerous attributes intentionally excluded:
+#   optimized_model_filepath  — triggers Model::Save(), overwrites arbitrary files
+#   profile_file_prefix       — writes profiling JSON to arbitrary path
+#   enable_profiling          — causes uncontrolled file writes to cwd
+_ALLOWED_SESSION_OPTIONS = frozenset(
+    {
+        "enable_cpu_mem_arena",
+        "enable_mem_pattern",
+        "enable_mem_reuse",
+        "execution_mode",
+        "execution_order",
+        "graph_optimization_level",
+        "inter_op_num_threads",
+        "intra_op_num_threads",
+        "log_severity_level",
+        "log_verbosity_level",
+        "logid",
+        "use_deterministic_compute",
+        "use_per_session_threads",
+    }
+)
+
 
 class OnnxRuntimeBackend(Backend):
     """
@@ -93,16 +116,18 @@ class OnnxRuntimeBackend(Backend):
     @classmethod
     def prepare(cls, model, device=None, **kwargs):
         """
-        Load the model and creates a :class:`onnxruntime.InferenceSession`
+        Load the model and creates an :class:`onnxruntime.backend.backend_rep.OnnxRuntimeBackendRep`
         ready to be used as a backend.
 
-        :param model: ModelProto (returned by `onnx.load`),
-            string for a filename or bytes for a serialized model
+        :param model: the model to prepare — accepts a file path (str), serialized
+            model (bytes), :class:`onnx.ModelProto`, :class:`onnxruntime.InferenceSession`,
+            or :class:`onnxruntime.backend.backend_rep.OnnxRuntimeBackendRep` (returned as-is)
         :param device: requested device for the computation,
             None means the default one which depends on
             the compilation settings
-        :param kwargs: see :class:`onnxruntime.SessionOptions`
-        :return: :class:`onnxruntime.InferenceSession`
+        :param kwargs: only a safe subset of :class:`onnxruntime.SessionOptions` attributes are
+            accepted; see ``_ALLOWED_SESSION_OPTIONS`` for the list
+        :return: :class:`onnxruntime.backend.backend_rep.OnnxRuntimeBackendRep`
         """
         if isinstance(model, OnnxRuntimeBackendRep):
             return model
@@ -111,8 +136,14 @@ class OnnxRuntimeBackend(Backend):
         elif isinstance(model, (str, bytes)):
             options = SessionOptions()
             for k, v in kwargs.items():
-                if hasattr(options, k):
+                if k in _ALLOWED_SESSION_OPTIONS:
                     setattr(options, k, v)
+                elif hasattr(options, k):
+                    raise RuntimeError(
+                        f"SessionOptions attribute '{k}' is not permitted via the backend API. "
+                        f"Allowed attributes: {', '.join(sorted(_ALLOWED_SESSION_OPTIONS))}"
+                    )
+                # else: silently ignore unknown keys
 
             excluded_providers = os.getenv("ORT_ONNX_BACKEND_EXCLUDE_PROVIDERS", default="").split(",")
             providers = [x for x in get_available_providers() if (x not in excluded_providers)]
@@ -148,13 +179,21 @@ class OnnxRuntimeBackend(Backend):
         """
         Compute the prediction.
 
-        :param model: :class:`onnxruntime.InferenceSession` returned
-            by function *prepare*
+        :param model: the model to run — accepts a file path (str), serialized
+            model (bytes), :class:`onnx.ModelProto`, :class:`onnxruntime.InferenceSession`,
+            or :class:`onnxruntime.backend.backend_rep.OnnxRuntimeBackendRep`
         :param inputs: inputs
         :param device: requested device for the computation,
             None means the default one which depends on
             the compilation settings
-        :param kwargs: see :class:`onnxruntime.RunOptions`
+        :param kwargs: ``run_model()`` forwards kwargs to both ``prepare()`` and ``rep.run()``.
+            ``prepare()`` validates and applies ``_ALLOWED_SESSION_OPTIONS`` only when creating
+            a new session from a model path or bytes; if ``model`` is already an
+            ``InferenceSession`` or ``OnnxRuntimeBackendRep``, session-option kwargs are
+            silently ignored. ``rep.run()`` always validates against ``_ALLOWED_RUN_OPTIONS``
+            and raises ``RuntimeError`` for known-but-blocked run attributes.
+            Logging-related kwargs (``log_severity_level``, ``log_verbosity_level``, ``logid``)
+            appear in both allowlists.
         :return: predictions
         """
         rep = cls.prepare(model, device, **kwargs)

venv/lib/python3.12/site-packages/onnxruntime/backend/backend_rep.py

@@ -10,11 +10,23 @@ from onnx.backend.base import BackendRep
 
 from onnxruntime import RunOptions
 
+# Allowlist of RunOptions attributes that are safe to set via the backend API.
+# 'terminate' excluded: setting it True would deny the current inference call.
+# 'training_mode' excluded: silently switches inference behavior in training builds.
+_ALLOWED_RUN_OPTIONS = frozenset(
+    {
+        "log_severity_level",
+        "log_verbosity_level",
+        "logid",
+        "only_execute_path_to_fetches",
+    }
+)
+
 
 class OnnxRuntimeBackendRep(BackendRep):
     """
-    Computes the prediction for a pipeline converted into
-    an :class:`onnxruntime.InferenceSession` node.
+    Wraps an :class:`onnxruntime.InferenceSession` to implement ONNX's
+    :class:`onnx.backend.base.BackendRep` interface for running predictions.
     """
 
     def __init__(self, session):
@@ -27,12 +39,24 @@ class OnnxRuntimeBackendRep(BackendRep):
         """
         Computes the prediction.
         See :meth:`onnxruntime.InferenceSession.run`.
+
+        :param inputs: a list of input arrays (one per model input) or a single
+            array when the model has exactly one input
+        :param kwargs: only a safe subset of :class:`onnxruntime.RunOptions` attributes are
+            accepted; see ``_ALLOWED_RUN_OPTIONS`` for the list
+        :return: list of output arrays
         """
 
         options = RunOptions()
         for k, v in kwargs.items():
-            if hasattr(options, k):
+            if k in _ALLOWED_RUN_OPTIONS:
                 setattr(options, k, v)
+            elif hasattr(options, k):
+                raise RuntimeError(
+                    f"RunOptions attribute '{k}' is not permitted via the backend API. "
+                    f"Allowed attributes: {', '.join(sorted(_ALLOWED_RUN_OPTIONS))}"
+                )
+            # else: silently ignore unknown keys
 
         if isinstance(inputs, list):
             inps = {}

venv/lib/python3.12/site-packages/onnxruntime/capi/build_and_package_info.py

@@ -1,2 +1,2 @@
 package_name = 'onnxruntime'
-__version__ = '1.23.2'
+__version__ = '1.26.0'

venv/lib/python3.12/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py

@@ -10,12 +10,14 @@ import os
 import typing
 import warnings
 from collections.abc import Callable, Sequence
+from enum import IntEnum
 from typing import Any
 
+import numpy as np
+
 from onnxruntime.capi import _pybind_state as C
 
 if typing.TYPE_CHECKING:
-    import numpy as np
     import numpy.typing as npt
 
     import onnxruntime
@@ -40,6 +42,30 @@ def get_ort_device_type(device_type: str) -> int:
         raise Exception("Unsupported device type: " + device_type)
 
 
+class OrtDeviceVendorId(IntEnum):
+    """Vendor IDs aligned with OrtDevice::VendorIds in ortdevice.h."""
+
+    NONE = 0x0000
+    AMD = 0x1002
+    NVIDIA = 0x10DE
+    ARM = 0x13B5
+    MICROSOFT = 0x1414
+    HUAWEI = 0x19E5
+    QUALCOMM = 0x5143
+    INTEL = 0x8086
+
+
+def get_vendor_id_for_device_type(device_type: str) -> OrtDeviceVendorId | None:
+    if device_type == "cuda":
+        return OrtDeviceVendorId.NVIDIA
+    elif device_type == "dml":
+        return OrtDeviceVendorId.MICROSOFT
+    elif device_type == "cann":
+        return OrtDeviceVendorId.HUAWEI
+    else:
+        return None
+
+
 class AdapterFormat:
     """
     This class is used to create adapter files from python structures
@@ -219,6 +245,15 @@ class Session:
         "Return registered execution providers' configurations."
         return self._provider_options
 
+    def get_provider_graph_assignment_info(self) -> Sequence[onnxruntime.OrtEpAssignedSubgraph]:
+        """
+        Get information about the subgraphs assigned to each execution provider and the nodes within.
+
+        Application must enable the recording of graph assignment information by setting the session configuration
+        for the key "session.record_ep_graph_assignment_info" to "1".
+        """
+        return self._sess.get_provider_graph_assignment_info()
+
     def set_providers(self, providers=None, provider_options=None) -> None:
         """
         Register the input list of execution providers. The underlying session is re-created.
@@ -397,6 +432,16 @@ class Session:
         """
         self._sess.run_with_iobinding(iobinding._iobinding, run_options)
 
+    def set_ep_dynamic_options(self, options: dict[str, str]):
+        """
+        Set dynamic options for execution providers.
+
+        :param options: Dictionary of key-value pairs where both keys and values are strings.
+                        These options will be passed to the execution providers to modify
+                        their runtime behavior.
+        """
+        self._sess.set_ep_dynamic_options(options)
+
     def get_tuning_results(self):
         return self._sess.get_tuning_results()
 
@@ -502,8 +547,42 @@ class InferenceSession(Session):
     def _create_inference_session(self, providers, provider_options, disabled_optimizers=None):
         available_providers = C.get_available_providers()
 
-        # Tensorrt can fall back to CUDA if it's explicitly assigned. All others fall back to CPU.
-        if "TensorrtExecutionProvider" in available_providers:
+        # Validate that TensorrtExecutionProvider and NvTensorRTRTXExecutionProvider are not both specified
+        if providers:
+            has_tensorrt = any(
+                provider == "TensorrtExecutionProvider"
+                or (isinstance(provider, tuple) and provider[0] == "TensorrtExecutionProvider")
+                for provider in providers
+            )
+            has_tensorrt_rtx = any(
+                provider == "NvTensorRTRTXExecutionProvider"
+                or (isinstance(provider, tuple) and provider[0] == "NvTensorRTRTXExecutionProvider")
+                for provider in providers
+            )
+            if has_tensorrt and has_tensorrt_rtx:
+                raise ValueError(
+                    "Cannot enable both 'TensorrtExecutionProvider' and 'NvTensorRTRTXExecutionProvider' "
+                    "in the same session."
+                )
+        # Tensorrt and TensorRT RTX can fall back to CUDA if it's explicitly assigned. All others fall back to CPU.
+        if "NvTensorRTRTXExecutionProvider" in available_providers:
+            if (
+                providers
+                and any(
+                    provider == "CUDAExecutionProvider"
+                    or (isinstance(provider, tuple) and provider[0] == "CUDAExecutionProvider")
+                    for provider in providers
+                )
+                and any(
+                    provider == "NvTensorRTRTXExecutionProvider"
+                    or (isinstance(provider, tuple) and provider[0] == "NvTensorRTRTXExecutionProvider")
+                    for provider in providers
+                )
+            ):
+                self._fallback_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+            else:
+                self._fallback_providers = ["CPUExecutionProvider"]
+        elif "TensorrtExecutionProvider" in available_providers:
             if (
                 providers
                 and any(
@@ -520,33 +599,6 @@ class InferenceSession(Session):
                 self._fallback_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
             else:
                 self._fallback_providers = ["CPUExecutionProvider"]
-        if "NvTensorRTRTXExecutionProvider" in available_providers:
-            if (
-                providers
-                and any(
-                    provider == "CUDAExecutionProvider"
-                    or (isinstance(provider, tuple) and provider[0] == "CUDAExecutionProvider")
-                    for provider in providers
-                )
-                and any(
-                    provider == "NvTensorRTRTXExecutionProvider"
-                    or (isinstance(provider, tuple) and provider[0] == "NvExecutionProvider")
-                    for provider in providers
-                )
-            ):
-                self._fallback_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
-            else:
-                self._fallback_providers = ["CPUExecutionProvider"]
-        # MIGraphX can fall back to ROCM if it's explicitly assigned. All others fall back to CPU.
-        elif "MIGraphXExecutionProvider" in available_providers:
-            if providers and any(
-                provider == "ROCMExecutionProvider"
-                or (isinstance(provider, tuple) and provider[0] == "ROCMExecutionProvider")
-                for provider in providers
-            ):
-                self._fallback_providers = ["ROCMExecutionProvider", "CPUExecutionProvider"]
-            else:
-                self._fallback_providers = ["CPUExecutionProvider"]
         else:
             self._fallback_providers = ["CPUExecutionProvider"]
 
@@ -986,7 +1038,9 @@ class OrtValue:
         return self._ortvalue
 
     @classmethod
-    def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0, vendor_id=-1) -> OrtValue:
+    def ortvalue_from_numpy(
+        cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0, vendor_id: int | OrtDeviceVendorId = -1
+    ) -> OrtValue:
         """
         Factory method to construct an OrtValue (which holds a Tensor) from a given Numpy object
         A copy of the data in the Numpy object is held by the OrtValue only if the device is NOT cpu
@@ -994,7 +1048,7 @@ class OrtValue:
         :param numpy_obj: The Numpy object to construct the OrtValue from
         :param device_type: e.g. cpu, cuda, cann, cpu by default
         :param device_id: device id, e.g. 0
-        :param vendor_id: The device's PCI vendor id. If provided, the device_type should be "gpu" or "npu".
+        :param vendor_id: The device's PCI vendor id as an int or OrtDeviceVendorId. If provided, the device_type should be "gpu" or "npu".
         """
         # Hold a reference to the numpy object (if device_type is 'cpu') as the OrtValue
         # is backed directly by the data buffer of the numpy object and so the numpy object
@@ -1023,7 +1077,12 @@ class OrtValue:
 
     @classmethod
     def ortvalue_from_shape_and_type(
-        cls, shape: Sequence[int], element_type, device_type: str = "cpu", device_id: int = 0, vendor_id: int = -1
+        cls,
+        shape: Sequence[int],
+        element_type,
+        device_type: str = "cpu",
+        device_id: int = 0,
+        vendor_id: int | OrtDeviceVendorId = -1,
     ) -> OrtValue:
         """
         Factory method to construct an OrtValue (which holds a Tensor) from given shape and element_type
@@ -1032,7 +1091,7 @@ class OrtValue:
         :param element_type: The data type of the elements. It can be either numpy type (like numpy.float32) or an integer for onnx type (like onnx.TensorProto.BFLOAT16).
         :param device_type: e.g. cpu, cuda, cann, cpu by default
         :param device_id: device id, e.g. 0
-        :param vendor_id: If provided the device type should be "gpu" or "npu".
+        :param vendor_id: The device's PCI vendor id as an int or OrtDeviceVendorId. If provided, the device type should be "gpu" or "npu".
         """
 
         device = OrtDevice.make(device_type, device_id, vendor_id)._get_c_device()
@@ -1141,15 +1200,126 @@ class OrtValue:
         """
         return self._ortvalue.numpy()
 
-    def update_inplace(self, np_arr) -> None:
+    def __array__(self, dtype=None, copy=None) -> np.ndarray:
         """
-        Update the OrtValue in place with a new Numpy array. The numpy contents
-        are copied over to the device memory backing the OrtValue. It can be used
-        to update the input valuess for an InferenceSession with CUDA graph
-        enabled or other scenarios where the OrtValue needs to be updated while
-        the memory address can not be changed.
+        Supports ``numpy.asarray(ortvalue)`` and ``numpy.array(ortvalue)`` via the
+        `numpy __array__ protocol <https://numpy.org/devdocs/user/basics.interoperability.html>`_.
+
+        Valid only for OrtValues holding Tensors on CPU.
+
+        :param dtype: Optional numpy dtype to cast the result to.
+        :param copy: Optional bool (numpy >= 2.0). If ``False``, a copy will
+            only be made if necessary. If ``True``, a copy is always forced.
+            If ``None`` (default), a copy will be made only if needed.
+        :return: A numpy array with the same data as the OrtValue.
         """
-        self._ortvalue.update_inplace(np_arr)
+        arr = self.numpy()
+
+        if copy is not None:
+            # numpy >= 2.0 added the copy kwarg to np.asarray;
+            # np.array has always accepted it but with weaker semantics pre-2.0.
+            arr = np.array(arr, dtype=dtype, copy=copy)
+        elif dtype is not None:
+            # np.asarray avoids a copy when the dtype already matches,
+            # preserving memory sharing with the underlying OrtValue.
+            arr = np.asarray(arr, dtype=dtype)
+
+        return arr
+
+    def __dlpack__(self, *, stream=None):
+        """
+        Returns a DLPack capsule representing the tensor (part of the
+        `DLPack protocol <https://dmlc.github.io/dlpack/latest/>`_).
+
+        This enables interoperability with other frameworks via
+        ``from_dlpack(ortvalue)`` (e.g. ``torch.from_dlpack``,
+        ``jax.dlpack.from_dlpack``, ``numpy.from_dlpack``).
+
+        The OrtValue must hold a contiguous tensor. No data is copied;
+        the consumer shares memory with this OrtValue, which must remain
+        alive while the capsule is in use.
+
+        :param stream: Optional stream on which the tensor data is accessible.
+            Currently unused; included for protocol compliance.
+        :return: A PyCapsule holding a DLManagedTensor.
+        """
+        return self._ortvalue.__dlpack__(stream=stream)
+
+    def __dlpack_device__(self) -> tuple[int, int]:
+        """
+        Returns ``(device_type, device_id)`` indicating where the tensor data
+        resides (part of the `DLPack protocol
+        <https://dmlc.github.io/dlpack/latest/>`_).
+
+        :return: Tuple of ``(device_type, device_id)`` as ints following DLPack
+            ``DLDeviceType`` enum values.
+        """
+        return self._ortvalue.__dlpack_device__()
+
+    @classmethod
+    def from_dlpack(cls, data, /) -> OrtValue:
+        """
+        Construct an OrtValue from an object that implements the DLPack protocol.
+
+        Accepts either:
+
+        * An object with ``__dlpack__`` / ``__dlpack_device__`` methods
+          (e.g. a PyTorch tensor, JAX array, or numpy array).
+        * A raw DLPack PyCapsule (legacy path).
+
+        Boolean tensors are automatically detected when the source object
+        exposes a ``dtype`` attribute (numpy, PyTorch, etc.) or is an
+        ``OrtValue``. For raw DLPack capsules where the original dtype cannot
+        be inspected, bool tensors encoded as uint8 by older DLPack versions
+        are not distinguishable from true uint8 tensors and will be imported
+        as uint8.
+
+        No data is copied; the new OrtValue shares memory with the source.
+
+        :param data: A tensor object supporting the DLPack protocol, or a raw
+            DLPack PyCapsule.
+        :return: An OrtValue wrapping the tensor data.
+        """
+        # Detect boolean dtype from the source object before consuming it,
+        # because DLPack encodes bool as uint8 and the capsule alone cannot
+        # distinguish between the two.
+        is_bool = False
+        if isinstance(data, OrtValue):
+            is_bool = data.data_type() == "tensor(bool)"
+        elif hasattr(data, "dtype"):
+            dtype_obj = data.dtype
+            # Use .name when available (numpy, cupy, tensorflow all expose it).
+            # Fall back to str() for frameworks that don't (e.g. PyTorch).
+            dtype_name = getattr(dtype_obj, "name", str(dtype_obj))
+            is_bool = dtype_name in ("bool", "bool_", "torch.bool")
+
+        # If the input supports the __dlpack__ protocol, call it to get the capsule.
+        if hasattr(data, "__dlpack__"):
+            capsule = data.__dlpack__()
+        else:
+            capsule = data
+
+        return cls(C.OrtValue.from_dlpack(capsule, is_bool))
+
+    def update_inplace(self, data) -> None:
+        """
+        Update the OrtValue in place. The source data is copied over to the device
+        memory backing the OrtValue. It can be used to update the input values for
+        an InferenceSession with CUDA graph enabled or other scenarios where the
+        OrtValue needs to be updated while the memory address can not be changed.
+
+        :param data: The source data, which can be a Numpy array or another OrtValue.
+            When an OrtValue is provided, data can be copied between devices (e.g.,
+            GPU to GPU) without going through the CPU.
+        """
+        if isinstance(data, OrtValue):
+            self._ortvalue.update_inplace(data._ortvalue)
+            return
+
+        if not isinstance(data, np.ndarray):
+            raise TypeError("data must be a numpy.ndarray or an OrtValue.")
+
+        self._ortvalue.update_inplace(data)
 
 
 def copy_tensors(src: Sequence[OrtValue], dst: Sequence[OrtValue], stream=None) -> None:
@@ -1185,9 +1355,24 @@ class OrtDevice:
         return self._ort_device
 
     @staticmethod
-    def make(ort_device_name, device_id, vendor_id=-1):
+    def make(ort_device_name, device_id, vendor_id: int | OrtDeviceVendorId = -1):
         if vendor_id < 0:
-            # backwards compatibility with predefined OrtDevice names
+            # Preserve the historical convenience aliases ("cuda", "dml", "cann")
+            # while making them work with plugin EP shared allocators. Those
+            # allocators are keyed by vendor-specific OrtDevice values even when the
+            # Python package itself was built without the corresponding built-in EP.
+            alias_vendor_id = get_vendor_id_for_device_type(ort_device_name)
+            if alias_vendor_id is not None:
+                return OrtDevice(
+                    C.OrtDevice(
+                        get_ort_device_type(ort_device_name),
+                        C.OrtDevice.default_memory(),
+                        int(alias_vendor_id),
+                        device_id,
+                    )
+                )
+
+            # backwards compatibility with generic predefined OrtDevice names
             return OrtDevice(
                 C.OrtDevice(
                     get_ort_device_type(ort_device_name),
@@ -1202,7 +1387,7 @@ class OrtDevice:
                 C.OrtDevice(
                     get_ort_device_type(ort_device_name),
                     C.OrtDevice.default_memory(),
-                    vendor_id,
+                    int(vendor_id),
                     device_id,
                 )
             )

venv/lib/python3.12/site-packages/onnxruntime/capi/onnxruntime_validation.py

@@ -23,9 +23,9 @@ def check_distro_info():
         __my_distro__ = __my_system__
         __my_distro_ver__ = platform.release().lower()
 
-        if __my_distro_ver__ not in ["10", "11"]:
+        if __my_distro_ver__ not in ["10", "11", "2016server", "2019server", "2022server", "2025server"]:
             warnings.warn(
-                f"Unsupported Windows version ({__my_distro_ver__}). ONNX Runtime supports Windows 10 and above, only."
+                f"Unsupported Windows version ({__my_distro_ver__}). ONNX Runtime supports Windows 10 and above, or Windows Server 2016 and above."
             )
     elif __my_system__ == "linux":
         """Although the 'platform' python module for getting Distro information works well on standard OS images

venv/lib/python3.12/site-packages/onnxruntime/quantization/calibrate.py

@@ -353,6 +353,14 @@ class MinMaxCalibrater(CalibraterBase):
                     return opset_import.version
             raise RuntimeError(f"Model does not contain a version for '{op_type}'.")
 
+        def insert_nodes(tensor_name, new_nodes):
+            index = next(
+                (i for i, x in enumerate(self.model.graph.node) if tensor_name in x.input), len(self.model.graph.node)
+            )
+            for node in new_nodes:
+                self.model.graph.node.insert(index, node)
+                index += 1
+
         def add_reduce_min_max(tensor_name, reduce_op_name):
             # When doing ReduceMax/ReduceMin, ORT can't reduce on dim with value of 0 if 'keepdims' is false.
             # To make the code simple, we always let keepdims to be 1.
@@ -396,7 +404,7 @@ class MinMaxCalibrater(CalibraterBase):
                     reduce_node.input.append(reduce_axes_name)
                     self.model.graph.initializer.append(reduce_axes)
 
-            self.model.graph.node.extend([reduce_node, reshape_node])
+            insert_nodes(tensor_name, [reduce_node, reshape_node])
             self.model.graph.output.append(helper.make_tensor_value_info(reduce_output, onnx_type, [None]))
 
         for tensor in tensors:
@@ -417,7 +425,14 @@ class MinMaxCalibrater(CalibraterBase):
             inputs = data_reader.get_next()
             if not inputs:
                 break
-            self.intermediate_outputs.append(self.infer_session.run(None, inputs))
+            self.intermediate_outputs.append(
+                [
+                    value if sess_o.name not in self.model_original_outputs else None
+                    for sess_o, value in zip(
+                        self.infer_session.get_outputs(), self.infer_session.run(None, inputs), strict=False
+                    )
+                ]
+            )
             if (
                 self.max_intermediate_outputs is not None
                 and len(self.intermediate_outputs) == self.max_intermediate_outputs

venv/lib/python3.12/site-packages/onnxruntime/quantization/execution_providers/qnn/preprocess.py

@@ -6,15 +6,15 @@
 from __future__ import annotations
 
 import logging
+import tempfile
 from pathlib import Path
 
 import onnx
 
-from ....tools.onnx_model_utils import fix_output_shapes, make_input_shape_fixed
+from ....tools.onnx_model_utils import fix_output_shapes, make_input_shape_fixed, optimize_model
 from ....tools.remove_initializer_from_input import remove_initializer_from_input
 from ...fusions import FusionGelu, FusionLayerNormalization
 from ...onnx_model import ONNXModel
-from ...quant_utils import save_and_reload_model_with_shape_infer
 from .fusion_lpnorm import FusionLpNormalization
 from .fusion_spacetodepth import FusionSpaceToDepth
 
@@ -93,7 +93,7 @@ def qnn_preprocess_model(
     """
     modified = False
     model = model_input if isinstance(model_input, onnx.ModelProto) else onnx.load_model(model_input)
-    model = save_and_reload_model_with_shape_infer(model)
+    model = save_and_reload_optimize_model(model, shape_infer=True)
     onnx_model = ONNXModel(model)
 
     # Optionally, fix the dynamic input shapes.
@@ -178,6 +178,24 @@ def qnn_preprocess_model(
     return modified
 
 
+def save_and_reload_optimize_model(model: onnx.ModelProto, shape_infer: bool) -> onnx.ModelProto:
+    with tempfile.TemporaryDirectory(prefix="ort.qnn_preproc.") as qnn_preproc_tmp_dir:
+        model_in_path = Path(qnn_preproc_tmp_dir).joinpath("qnn_proc_input.onnx")
+        onnx.save_model(model, model_in_path, save_as_external_data=True)
+        if shape_infer:
+            model_infer_path = Path(qnn_preproc_tmp_dir).joinpath("qnn_proc_infer.onnx")
+            onnx.shape_inference.infer_shapes_path(str(model_in_path), str(model_infer_path))
+            model_in_path = model_infer_path
+        model_out_path = Path(qnn_preproc_tmp_dir).joinpath("qnn_proc_output.onnx")
+        optimize_model(model_in_path, model_out_path)
+        ret_model = onnx.load_model(model_out_path)
+        ret_metaprops = {"onnx.infer": "onnxruntime.tools.qnn.preprocess"}
+        if ret_model.metadata_props:
+            ret_metaprops.update(ret_model.metadata_props)
+        onnx.helper.set_model_props(ret_model, ret_metaprops)
+        return ret_model
+
+
 class InputOutputNameMap:
     def __init__(
         self,

venv/lib/python3.12/site-packages/onnxruntime/quantization/execution_providers/qnn/quant_config.py

@@ -331,23 +331,6 @@ class QnnCompatibilityOverrides:
 
         if not self.per_channel:
             self._make_static_inputs_use_default_weight_type(node)
-            return
-
-        has_weight_no_overrides = node.input[1] in self.initializers and node.input[1] not in self.overrides
-        has_bias_no_overrides = (
-            len(node.input) > 2
-            and node.input[2]
-            and node.input[2] in self.initializers
-            and node.input[2] not in self.overrides
-        )
-
-        if has_weight_no_overrides or has_bias_no_overrides:
-            # TODO: Make bias input not per-channel. QNN needs it to be per-tensor, but quantizer
-            # tries to makes it per-channel if the weight is also per-channel.
-            raise ValueError(
-                "get_qnn_qdq_config() does not currently support the global per_channel option with LayerNormalization."
-                " Please try using custom overrides that make bias per-tensor quantized."
-            )
 
     def _process_sigmoid(self, node: onnx.NodeProto):
         """

venv/lib/python3.12/site-packages/onnxruntime/quantization/fusions/fusion_layernorm.py

@@ -33,6 +33,16 @@ class FusionLayerNormalization(Fusion):
                                      |                                                 |
                                      +-------------------------------------------------+
 
+         Or, using Mul instead of Pow:
+
+              +----------------------+
+              |                      |
+              |                      v
+          [Root] --> ReduceMean -->  Sub  --> Mul --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add
+                     (axis=2 or -1)  |     (in0=in1)   (axis=2 or -1)  (E-6 or E-12 or 0) ^
+                                     |                                                 |
+                                     +-------------------------------------------------+
+
          It also handles cases of duplicated sub nodes exported from older version of PyTorch:
 
               +----------------------+
@@ -40,7 +50,7 @@ class FusionLayerNormalization(Fusion):
               |           +-------> Sub-----------------------------------------------+
               |           |                                                           |
               |           |                                                           v
-          [Root] --> ReduceMean -->  Sub  --> Pow --> ReduceMean --> Add --> Sqrt --> Div  --> Mul --> Add
+          [Root] --> ReduceMean -->  Sub  --> (Pow or Mul) --> ReduceMean --> Add --> Sqrt --> Div  --> Mul --> Add
               |                      ^
               |                      |
               +----------------------+
@@ -70,10 +80,9 @@ class FusionLayerNormalization(Fusion):
             div_node,
             [
                 (["Sqrt", "Add", "ReduceMean", "Pow", "Sub"], [1, 0, 0, 0, 0]),
-                (
-                    ["Sqrt", "Add", "ReduceMean", "Pow", "Cast", "Sub"],
-                    [1, 0, 0, 0, 0, 0],
-                ),
+                (["Sqrt", "Add", "ReduceMean", "Pow", "Cast", "Sub"], [1, 0, 0, 0, 0, 0]),
+                (["Sqrt", "Add", "ReduceMean", "Mul", "Sub"], [1, 0, 0, 0, 0]),
+                (["Sqrt", "Add", "ReduceMean", "Mul", "Cast", "Sub"], [1, 0, 0, 0, 0, 0]),
             ],
             output_name_to_node,
         )
@@ -90,8 +99,10 @@ class FusionLayerNormalization(Fusion):
             # Skip fusion since epsilon value is not expected.
             return
 
-        pow_node = parent_nodes[3]
-        if self.find_constant_input(pow_node, 2.0) != 1:
+        pow_or_mul_node = parent_nodes[3]
+        if pow_or_mul_node.op_type == "Pow" and self.find_constant_input(pow_or_mul_node, 2.0) != 1:
+            return
+        elif pow_or_mul_node.op_type == "Mul" and pow_or_mul_node.input[0] != pow_or_mul_node.input[1]:
             return
 
         mul_node = input_name_to_nodes[div_node.output[0]][0]

venv/lib/python3.12/site-packages/onnxruntime/quantization/matmul_nbits_quantizer.py

@@ -11,12 +11,19 @@ import copy
 import logging
 import os
 
+import ml_dtypes
 import numpy as np
 import numpy.typing as npt
 import onnx
+import onnx_ir as ir
 from onnx.onnx_pb import GraphProto, ModelProto, NodeProto, TensorProto
 
-from onnxruntime.capi._pybind_state import quantize_matmul_4bits, quantize_matmul_8bits, quantize_qdq_matmul_4bits
+from onnxruntime.capi._pybind_state import (
+    quantize_matmul_2bits,
+    quantize_matmul_4bits,
+    quantize_matmul_8bits,
+    quantize_qdq_matmul_4bits,
+)
 
 from .calibrate import CalibrationDataReader
 from .neural_compressor import gptq_quantize, rtn_quantize
@@ -816,9 +823,13 @@ class DefaultWeightOnlyQuantizer:
 
             # block wise quantization, each block comes from a single column
             packed = np.zeros((cols, k_blocks, blob_size), dtype="uint8")
-            zero_point = np.zeros(cols * ((k_blocks + kpack - 1) // kpack), dtype="uint8")
-            scales = np.zeros((cols * k_blocks), dtype=fp32weight.dtype)
-            if qbits == 8:
+            zero_point = np.zeros((cols, ((k_blocks + kpack - 1) // kpack)), dtype="uint8")
+            scales = np.zeros((cols, k_blocks), dtype=fp32weight.dtype)
+            if qbits == 2:
+                quantize_matmul_2bits(
+                    packed, fp32weight, scales, zero_point, block_size, cols, rows, self.config.is_symmetric
+                )
+            elif qbits == 8:
                 quantize_matmul_8bits(
                     packed, fp32weight, scales, zero_point, block_size, cols, rows, self.config.is_symmetric
                 )
@@ -857,21 +868,27 @@ class DefaultWeightOnlyQuantizer:
             logger.info("MatMul doesn't have const weight. Skip to quantize")
             return [node]  # only care about constant weight
 
-        b_ndarray = onnx.numpy_helper.to_array(b_tensor)
+        b_ndarray = ir.from_proto(b_tensor).numpy()
         if len(b_ndarray.shape) != 2:
             logger.info("MatMul weight is not 2D. Skip to quantize")
             return [node]  # can only process 2-D matrix
 
+        bfloat16 = b_ndarray.dtype == "bfloat16"
+        if bfloat16:
+            b_ndarray = b_ndarray.astype(np.float32)
+
         packed, scales, zero_points = self.qbits_block_quant(b_ndarray)
+        if bfloat16:
+            scales = scales.astype(ml_dtypes.bfloat16)
 
         if self.config.quant_format == QuantFormat.QOperator:
-            b_quant = onnx.numpy_helper.from_array(packed, b_tensor.name + f"_Q{bits}")
-            scales_tensor = onnx.numpy_helper.from_array(scales, b_tensor.name + "_scales")
+            b_quant = ir.serde.serialize_tensor(ir.Tensor(packed, name=b_tensor.name + f"_Q{bits}"))
+            scales_tensor = ir.serde.serialize_tensor(ir.Tensor(scales, name=b_tensor.name + "_scales"))
         else:
             b_quant = onnx.helper.make_tensor(
                 b_tensor.name + f"_DQ_Q{bits}", qtype, b_ndarray.shape, packed.tobytes(), True
             )
-            scales_tensor = onnx.numpy_helper.from_array(scales, b_tensor.name + "_DQ_scales")
+            scales_tensor = ir.serde.serialize_tensor(ir.Tensor(scales, name=b_tensor.name + "_DQ_scales"))
 
         # if QDQ, CW and SYM enabled, optimize for Intel NPU, tranpose the weight to NHWC format will increase performance
         qdq_opt_for_intel_npu_enabled = (
@@ -886,7 +903,7 @@ class DefaultWeightOnlyQuantizer:
             b_quant = onnx.helper.make_tensor(
                 b_tensor.name + f"_DQ_Q{bits}", qtype, [cols, rows], packed.tobytes(), True
             )
-            scales_tensor = onnx.numpy_helper.from_array(scales, b_tensor.name + "_DQ_scales")
+            scales_tensor = ir.serde.serialize_tensor(ir.Tensor(scales, name=b_tensor.name + "_DQ_scales"))
 
         for input in b_graph.input:
             if input.name == input_b:
@@ -1206,7 +1223,7 @@ class MatMulNBitsQuantizer:
     MatMul              MatMulNBits                DeQuantizeLinear -> MatMul
     Gather              GatherBlockQuantized       Gather, Gather, Gather (optional) -> DequantizeLinear
 
-    Perform 4/8 bits quantization of constant weights for target nodes.
+    Perform 2/4/8 bits quantization of constant weights for target nodes.
     If algo_config.quant_format is QOperator:
       - nodes are replaced by the corresponding QOperator nodes.
       - quantized weights are stored in the contrib ops.
@@ -1224,6 +1241,7 @@ class MatMulNBitsQuantizer:
     def __init__(
         self,
         model: ModelProto | str,
+        bits: int = 4,  # default to 4bit
         block_size: int = 128,
         is_symmetric: bool = False,
         accuracy_level: int | None = None,
@@ -1239,6 +1257,7 @@ class MatMulNBitsQuantizer:
             nodes_to_exclude = []
         self.model = ONNXModel(onnx.load(model)) if isinstance(model, str) else ONNXModel(model)
         self.model_path = model if isinstance(model, str) else None
+        self.bits = bits
         self.block_size = block_size
         self.is_symmetric = is_symmetric
         self.accuracy_level = accuracy_level
@@ -1254,13 +1273,13 @@ class MatMulNBitsQuantizer:
                 quant_format=quant_format,
                 op_types_to_quantize=op_types_to_quantize,
                 quant_axes=quant_axes,
-                bits=4,  # default to 4 bits
+                bits=bits,
                 channel_wised_quantize=channel_wised_quantize,
             )
 
         self.algo_config = algo_config
         if hasattr(self.algo_config, "bits"):
-            assert self.algo_config.bits in [4, 8], "Only support 4 or 8 bits quantization"
+            assert self.algo_config.bits in [2, 4, 8], "Only support 2, 4 or 8 bits quantization"
 
         if algo_config.algorithm == "HQQ":
             self.node_quantizer = HQQWeightOnlyQuantizer(self.algo_config)
@@ -1609,6 +1628,7 @@ if __name__ == "__main__":
 
     quant = MatMulNBitsQuantizer(
         model=model,
+        bits=args.bits,
         accuracy_level=args.accuracy_level,
         nodes_to_exclude=args.nodes_to_exclude,
         nodes_to_include=args.nodes_to_include,

venv/lib/python3.12/site-packages/onnxruntime/quantization/neural_compressor/onnx_model.py

@@ -21,12 +21,12 @@
 import copy
 import logging
 import os
-import sys
 from collections import deque
 from pathlib import Path
 
 import onnx
 import onnx.external_data_helper
+import onnx_ir as ir
 
 from .util import MAXIMUM_PROTOBUF, find_by_name
 
@@ -73,26 +73,11 @@ class ONNXModel:
 
     def check_is_large_model(self):
         """Check model > 2GB."""
-        init_size = 0
-        for init in self._model.graph.initializer:
-            # if initializer has external data location, return True
-            if init.HasField("data_location") and init.data_location == onnx.TensorProto.EXTERNAL:
-                self._is_large_model = True
-                return
-            # if raise error of initializer size > 2GB, return True
-            try:
-                init_bytes = init.SerializeToString()
-                init_size += sys.getsizeof(init_bytes)
-            except Exception as e:
-                if "exceeds maximum protobuf size of 2GB" in str(e):
-                    self._is_large_model = True
-                    return
-                else:  # pragma: no cover
-                    raise e
-            if init_size > MAXIMUM_PROTOBUF:
-                self._is_large_model = True
-                return
-        self._is_large_model = False
+        ir_graph = ir.from_proto(self._model.graph)
+        initializer_size = sum(
+            v.const_value.nbytes for v in ir_graph.initializers.values() if v.const_value is not None
+        )
+        self._is_large_model = initializer_size > MAXIMUM_PROTOBUF
 
     @property
     def is_large_model(self):

venv/lib/python3.12/site-packages/onnxruntime/quantization/qdq_quantizer.py

@@ -12,7 +12,6 @@ from typing import Any
 
 import numpy as np
 import onnx
-import onnx.numpy_helper
 from onnx import TensorProto
 from onnx import onnx_pb as onnx_proto
 

venv/lib/python3.12/site-packages/onnxruntime/quantization/quant_utils.py

@@ -14,6 +14,7 @@ from pathlib import Path
 
 import numpy
 import onnx
+from ml_dtypes import float8_e4m3fn, int4, uint4
 from onnx import ModelProto, TensorProto, external_data_helper
 from onnx import onnx_pb as onnx_proto
 from onnx.helper import make_graph, make_model, make_node, make_tensor_value_info
@@ -21,19 +22,6 @@ from onnx.reference import ReferenceEvaluator
 
 from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions
 
-try:
-    from onnx.reference.custom_element_types import float8e4m3fn
-except ImportError:
-    float8e4m3fn = None
-
-# INT4 np.dtypes added in ONNX 1.16. These map to np.int8/np.uint8 because numpy
-# does not support sub-byte types.
-try:
-    from onnx.reference.custom_element_types import int4, uint4
-except ImportError:
-    int4 = None
-    uint4 = None
-
 try:
     from onnx.reference.op_run import to_array_extended
 except ImportError:
@@ -149,9 +137,9 @@ ONNX_TYPE_TO_NP_TYPE = {
     onnx_proto.TensorProto.UINT8: numpy.dtype("uint8"),
     onnx_proto.TensorProto.INT16: numpy.dtype("int16"),
     onnx_proto.TensorProto.UINT16: numpy.dtype("uint16"),
-    onnx_proto.TensorProto.FLOAT8E4M3FN: float8e4m3fn,
-    onnx_proto.TensorProto.INT4: int4,  # base_dtype is np.int8
-    onnx_proto.TensorProto.UINT4: uint4,  # base_dtype is np.uint8
+    onnx_proto.TensorProto.FLOAT8E4M3FN: float8_e4m3fn,
+    onnx_proto.TensorProto.INT4: int4,
+    onnx_proto.TensorProto.UINT4: uint4,
 }
 
 ONNX_INT_TYPE_RANGE = {
@@ -164,9 +152,7 @@ ONNX_INT_TYPE_RANGE = {
 }
 
 ONNX_INT_TYPE_SYMMETRIC_RANGE = {
-    onnx_proto.TensorProto.UINT8: (numpy.array(0, dtype=numpy.uint8), numpy.array(254, dtype=numpy.uint8)),
     onnx_proto.TensorProto.INT8: (numpy.array(-127, dtype=numpy.int8), numpy.array(127, dtype=numpy.int8)),
-    onnx_proto.TensorProto.UINT16: (numpy.array(0, dtype=numpy.uint16), numpy.array(65534, dtype=numpy.uint16)),
     onnx_proto.TensorProto.INT16: (numpy.array(-32767, dtype=numpy.int16), numpy.array(32767, dtype=numpy.int16)),
 }
 
@@ -175,7 +161,7 @@ ONNX_INT_TYPE_REDUCED_RANGE = {
     onnx_proto.TensorProto.INT8: (numpy.array(-64, dtype=numpy.int8), numpy.array(64, dtype=numpy.int8)),
     onnx_proto.TensorProto.UINT16: (numpy.array(0, dtype=numpy.uint16), numpy.array(32767, dtype=numpy.uint16)),
     onnx_proto.TensorProto.INT16: (numpy.array(-16384, dtype=numpy.int16), numpy.array(16384, dtype=numpy.int16)),
-    onnx_proto.TensorProto.UINT4: (numpy.array(0, dtype=int4), numpy.array(7, dtype=int4)),
+    onnx_proto.TensorProto.UINT4: (numpy.array(0, dtype=uint4), numpy.array(7, dtype=uint4)),
     onnx_proto.TensorProto.INT4: (numpy.array(-4, dtype=int4), numpy.array(3, dtype=int4)),
 }
 
@@ -324,11 +310,10 @@ def compute_scale_zp_float8(element_type, std):
     zp_dtype = None
     if element_type not in FLOAT8_DISTRIBUTIONS:
         if element_type == TensorProto.FLOAT8E4M3FN:
-            from onnx.numpy_helper import float8e4m3_to_float32  # noqa: PLC0415
-            from onnx.reference.custom_element_types import float8e4m3fn  # noqa: PLC0415
+            from ml_dtypes import float8_e4m3fn  # noqa: PLC0415
 
-            zp_dtype = float8e4m3fn
-            all_values = [float8e4m3_to_float32(i) for i in range(256)]
+            zp_dtype = float8_e4m3fn
+            all_values = [float(i) for i in range(256)]
             values = numpy.array(
                 [f for f in all_values if not numpy.isnan(f) and not numpy.isinf(f)], dtype=numpy.float32
             )
@@ -336,9 +321,9 @@ def compute_scale_zp_float8(element_type, std):
             raise ValueError(f"Quantization to element_type={element_type} not implemented.")
         FLOAT8_DISTRIBUTIONS[element_type] = values
     elif element_type == TensorProto.FLOAT8E4M3FN:
-        from onnx.reference.custom_element_types import float8e4m3fn  # noqa: PLC0415
+        from ml_dtypes import float8_e4m3fn  # noqa: PLC0415
 
-        zp_dtype = float8e4m3fn
+        zp_dtype = float8_e4m3fn
 
     if zp_dtype is None:
         raise TypeError(f"Unexpected element_type {element_type}.")
@@ -449,7 +434,7 @@ def quantize_data(
     )
     if qType == TensorProto.FLOAT8E4M3FN:
         quantized_data = quantize_nparray(qType, data, scale, zero_point)
-        if any((quantized_data.astype(numpy.uint8).ravel() & 127) == 127):
+        if any((quantized_data.view(numpy.uint8).ravel() & 127) == 127):
             np_data = numpy.asarray(data)
             raise RuntimeError(
                 f"One of the quantized value is NaN data in [{np_data.min()}, {np_data.max()}], "
@@ -533,7 +518,7 @@ def quantize_onnx_initializer(
                     f"\nraw={str(q_weight_initializer)[:200]}."
                 )
     elif quant_type in (onnx.TensorProto.INT4, onnx.TensorProto.UINT4):
-        if q_weight_data.dtype not in (numpy.int8, numpy.uint8):
+        if q_weight_data.dtype not in (int4, uint4):
             raise RuntimeError(f"Quantized weights for {q_weight_name} must be 8-bit before packing as 4-bit values.")
 
         # We do not use onnx.helper.pack_float32_to_4bit() due to performance.

venv/lib/python3.12/site-packages/onnxruntime/quantization/registry.py

@@ -87,6 +87,7 @@ QDQRegistry = {
     "LayerNormalization": QDQNormalization,
     "BatchNormalization": QDQNormalization,
     "TopK": QDQDirect8BitOp,
+    "CumSum": QDQOperatorBase,
 }
 
 

venv/lib/python3.12/site-packages/onnxruntime/quantization/shape_inference.py

@@ -74,34 +74,29 @@ def quant_pre_process(
 
     with tempfile.TemporaryDirectory(prefix="pre.quant.") as quant_tmp_dir:
         temp_path = Path(quant_tmp_dir)
-        model = None
+        model = input_model if isinstance(input_model, onnx.ModelProto) else onnx.load(input_model)
+
+        # Since Upsample is deprecated after opset v10, and the model's opset will
+        # be upgraded to at least v11 during quantization, we need to replace Upsample
+        # with Resize first to avoid generating an invalid model.
+        ai_onnx_domain = [opset for opset in model.opset_import if not opset.domain or opset.domain == "ai.onnx"]
+        if len(ai_onnx_domain) == 1:
+            opset_version = ai_onnx_domain[0].version
+            if opset_version <= 10:
+                ReplaceUpsampleWithResize(ONNXModel(model), opset_version).apply()
+                model = onnx.version_converter.convert_version(model, 11)
+                model = save_and_reload_model_with_shape_infer(model)
 
         if not skip_symbolic_shape:
             logger.info("Performing symbolic shape inference...")
-            loaded_model = input_model if isinstance(input_model, onnx.ModelProto) else onnx.load(input_model)
             model = SymbolicShapeInference.infer_shapes(
-                loaded_model,
+                model,
                 int_max,
                 auto_merge,
                 guess_output_rank,
                 verbose,
             )
 
-        # Since Upsample is deprecated after opset v10, and the model's opset will
-        # be upgraded to at least v11 during quantization, we need to replace Upsample
-        # with Resize first to avoid generating an invalid model.
-        if model:
-            ai_onnx_domain = [opset for opset in model.opset_import if not opset.domain or opset.domain == "ai.onnx"]
-            if len(ai_onnx_domain) == 1:
-                opset_version = ai_onnx_domain[0].version
-                if opset_version < 10:
-                    ReplaceUpsampleWithResize(ONNXModel(model), opset_version).apply()
-                    model.opset_import.remove(ai_onnx_domain[0])
-                    opset_version = 11
-                    model.opset_import.extend([onnx.helper.make_opsetid("", opset_version)])
-                    model = onnx.version_converter.convert_version(model, opset_version)
-                    model = save_and_reload_model_with_shape_infer(model)
-
         if not skip_optimization:
             # Use ORT optimizers (native code) to optimize model
             if not skip_symbolic_shape:

venv/lib/python3.12/site-packages/onnxruntime/quantization/static_quantize_runner.py

@@ -20,7 +20,7 @@ class OnnxModelCalibrationDataReader(CalibrationDataReader):
         name2tensors = []
         for data_dir in data_dirs:
             name2tensor = {}
-            data_paths = [os.path.join(data_dir, a) for a in sorted(os.listdir(data_dir))]
+            data_paths = [os.path.join(data_dir, f"input_{input_idx}.pb") for input_idx in range(len(model_inputs))]
             data_ndarrays = [self.read_onnx_pb_data(data_path) for data_path in data_paths]
             for model_input, data_ndarray in zip(model_inputs, data_ndarrays, strict=False):
                 name2tensor[model_input.name] = data_ndarray