diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index a65396d64ff..3eac516d287 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -15,8 +15,12 @@ """Module to load C++ / CUDA extensions.""" +import os from pathlib import Path +import torch +from packaging.version import Version + from modelopt.torch.utils import load_cpp_extension __all__ = ["get_cuda_ext", "get_cuda_ext_fp8", "get_cuda_ext_mx", "precompile"] @@ -25,6 +29,18 @@ kernels_gemm = path.parent / "kernels" / "quantization" / "gemm" +def _get_mx_cuda_cflags() -> list[str]: + flags = ["--use_fast_math"] + if ( + os.name == "nt" + and torch.version.cuda is not None + and Version(torch.version.cuda) >= Version("13") + ): + # CUDA 13 CCCL requires MSVC's standard-conforming preprocessor. + flags.extend(["-Xcompiler", "/Zc:preprocessor"]) + return flags + + def get_cuda_ext(raise_if_failed: bool = False): """Returns the cuda extension for tensor_quant.""" if not hasattr(get_cuda_ext, "extension"): @@ -66,7 +82,7 @@ def get_cuda_ext_mx(raise_if_failed: bool = False): "CUDA extension for MX quantization could not be built and loaded, MX simulated" " quantization will not be available." ), - extra_cuda_cflags=["--use_fast_math"], + extra_cuda_cflags=_get_mx_cuda_cflags(), raise_if_failed=raise_if_failed, ) return get_cuda_ext_mx.extension # type:ignore[attr-defined] diff --git a/pyproject.toml b/pyproject.toml index f23ce722e5c..be6ee364e3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dependencies = [ onnx = [ "cppimport", "cupy-cuda12x; platform_machine != 'aarch64' and platform_machine != 'ARM64' and platform_system != 'Darwin'", + "cupy-cuda13x>=14.2.0; python_version > '3.10' and platform_system == 'Windows' and platform_machine == 'ARM64'", "lief", "ml_dtypes", "onnx-graphsurgeon>=0.6.1", @@ -65,7 +66,7 @@ onnx = [ # ORT for Windows x64 "onnxruntime-gpu==1.22.0; platform_system == 'Windows' and platform_machine == 'AMD64'", # ORT host and standalone TensorRT-RTX ABI EP for native Windows ARM64. - "onnxruntime~=1.24.2; python_version > '3.10' and platform_system == 'Windows' and platform_machine == 'ARM64'", + "onnxruntime~=1.26.0; python_version > '3.10' and platform_system == 'Windows' and platform_machine == 'ARM64'", "onnxruntime-ep-nv-tensorrt-rtx-cu13==0.4.0; python_version > '3.10' and platform_system == 'Windows' and platform_machine == 'ARM64'", # ORT with Python <= 3.10 on supported platforms "onnxruntime~=1.22.0; python_version <= '3.10' and (platform_machine == 'aarch64' or platform_system == 'Darwin')",