From 74f8d830d742a0e0d79a5b169036f0db7737ac5d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 05:04:18 +0000 Subject: [PATCH 1/9] Initial plan From c04fb6dd94d8b9dfe7967bab332f2fa85bdc0020 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 05:21:43 +0000 Subject: [PATCH 2/9] Fix CUDA runtime hard dependency in pybind module by routing through ProviderInfo_CUDA interface - Remove CUDA source files (fpA_intB_gemm_adaptor.cu, fpA_intB_gemm_preprocessors_impl.cu) from onnxruntime_pybind11_state build, and remove CUDA::cudart linkage. This eliminates the NEEDED libcudart.so.13 entry that caused import failures on CPU-only machines. - Add PackWeightsForMixedGemm() virtual method to ProviderInfo_CUDA interface so the pybind module can call it via TryGetProviderInfo_CUDA() (dynamic load) instead of linking CUDA directly. - Implement PackWeightsForMixedGemm() in ProviderInfo_CUDA_Impl (cuda_provider_factory.cc) using existing GPU kernels. - Update onnxruntime_pybind_quant.cc to call through the provider interface; move PackFP4WeightsForMoE outside the USE_CUDA guard since it is pure CPU code. Fixes: import onnxruntime fails with libcudart.so.13 error on CPU-only Linux (issue #29500) --- cmake/onnxruntime_python.cmake | 21 +-- .../providers/cuda/cuda_provider_factory.cc | 67 ++++++++++ .../providers/cuda/cuda_provider_factory.h | 9 ++ .../python/onnxruntime_pybind_quant.cc | 123 +++++------------- 4 files changed, 113 insertions(+), 107 deletions(-) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 3f6976c3c8955..1f6564a257c70 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -231,22 +231,11 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE Python::NumPy ) -# Starting with Python 3.8 on Windows, PATH environment variable are no longer used to resolve DLL dependencies -# for extension modules or libraries loaded via ctypes. -# To avoid package import issues, we do not link pybind module against the CUDA runtime on Windows, instead of -# os.add_dll_directory() to deal with CUDA paths. -if (onnxruntime_USE_CUDA AND NOT WIN32) - target_sources(onnxruntime_pybind11_state PRIVATE - "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" - "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" - ) - include(cutlass) - target_include_directories(onnxruntime_pybind11_state PRIVATE ${cutlass_SOURCE_DIR}/include) - target_link_libraries(onnxruntime_pybind11_state PRIVATE CUDA::cudart) -endif() -if (onnxruntime_USE_CUDA AND WIN32) - target_compile_definitions(onnxruntime_pybind11_state PRIVATE ORT_NO_CUDA_IN_PYBIND) -endif() +# The CUDA quantization helpers (pack_weights_for_cuda_mixed_gemm) are called through the +# ProviderInfo_CUDA interface, which dynamically loads onnxruntime_providers_cuda at runtime. +# Do NOT compile CUDA source files directly into onnxruntime_pybind11_state or link +# CUDA::cudart from it: that would create a hard libcudart.so dependency that prevents +# importing the Python module on CPU-only machines. set(onnxruntime_pybind11_state_dependencies ${onnxruntime_EXTERNAL_DEPENDENCIES} diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index 66d3617c8da9a..acfca0be129b3 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -10,6 +10,7 @@ #include #include #include +#include #include @@ -20,6 +21,8 @@ #include "core/providers/cuda/cuda_stream_handle.h" #include "core/providers/cuda/gpu_data_transfer.h" #include "core/providers/cuda/math/unary_elementwise_ops_impl.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" #ifdef ENABLE_NVTX_PROFILE #include "nvtx_profile.h" @@ -198,6 +201,70 @@ struct ProviderInfo_CUDA_Impl final : ProviderInfo_CUDA { params.arena_cfg = default_memory_arena_cfg; return CUDAExecutionProvider::CreateCudaPinnedAllocator(params); } + + void PackWeightsForMixedGemm(const uint8_t* q_weights, int32_t N, int32_t K, + int32_t bits, int32_t force_arch, + int8_t* output) override { + size_t packed_weight_bytes = static_cast(N) * static_cast(K) / static_cast(8 / bits); + + struct CudaMemDeleter { + void operator()(void* p) const noexcept { + if (p) cudaFree(p); + } + }; + auto make_device_buf = [](size_t bytes) { + void* p = nullptr; + CUDA_CALL_THROW(cudaMalloc(&p, bytes)); + return std::unique_ptr(p); + }; + + auto d_input = make_device_buf(packed_weight_bytes); + auto d_transposed = make_device_buf(packed_weight_bytes); + auto d_preprocessed = make_device_buf(packed_weight_bytes); + auto d_permutation_map = make_device_buf(32 * sizeof(int32_t)); + + cudaStream_t stream = cudaStreamLegacy; + + CUDA_CALL_THROW(cudaMemcpyAsync(d_input.get(), q_weights, packed_weight_bytes, + cudaMemcpyHostToDevice, stream)); + + if (bits == 4) { + ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( + stream, d_transposed.get(), d_input.get(), N, K); + } else { + ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( + stream, + static_cast(d_transposed.get()), + static_cast(d_input.get()), + N, K); + } + + using ::onnxruntime::llm::kernels::weight_only::QuantType; + QuantType quant_type = (bits == 4) ? QuantType::W4_A16 : QuantType::W8_A16; + + int sm = force_arch; + if (sm < 0) { + int device_id = 0; + CUDA_CALL_THROW(cudaGetDevice(&device_id)); + cudaDeviceProp device_prop; + CUDA_CALL_THROW(cudaGetDeviceProperties(&device_prop, device_id)); + sm = device_prop.major * 10 + device_prop.minor; + } + sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); + + ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( + stream, sm, + static_cast(d_preprocessed.get()), + static_cast(d_transposed.get()), + static_cast(d_permutation_map.get()), + {static_cast(K), static_cast(N)}, + quant_type); + + CUDA_CALL_THROW(cudaGetLastError()); + CUDA_CALL_THROW(cudaMemcpyAsync(output, d_preprocessed.get(), packed_weight_bytes, + cudaMemcpyDeviceToHost, stream)); + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + } } g_info; struct CUDA_Provider : Provider { diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.h b/onnxruntime/core/providers/cuda/cuda_provider_factory.h index 1a4b19cb100d3..edb48dd1e2cc6 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.h +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.h @@ -55,6 +55,15 @@ struct ProviderInfo_CUDA { virtual std::shared_ptr CreateCudaAllocator(int16_t device_id, size_t gpu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, onnxruntime::CUDAExecutionProviderExternalAllocatorInfo& external_allocator_info, const OrtArenaCfg* default_memory_arena_cfg) = 0; virtual std::shared_ptr CreateCudaPinnedAllocator(int16_t device_id, size_t gpu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, const OrtArenaCfg* default_memory_arena_cfg) = 0; + // Preprocess quantized weights for CUDA mixed-precision GEMM (FpA_IntB format). + // q_weights: packed quantized weights from MatMulNBits in (N, K/(8/bits)) layout. + // output: caller-allocated buffer of size N * K / (8 / bits) bytes. + // bits: quantization bit-width (4 or 8). + // force_arch: SM version to use for permutation (-1 for auto-detect from current device). + virtual void PackWeightsForMixedGemm(const uint8_t* q_weights, int32_t N, int32_t K, + int32_t bits, int32_t force_arch, + int8_t* output) = 0; + // This function is the entry point to CUDA EP's UT cases. // All tests are only called from onnxruntime_provider_test. virtual void TestAll() { diff --git a/onnxruntime/python/onnxruntime_pybind_quant.cc b/onnxruntime/python/onnxruntime_pybind_quant.cc index 7220153b4fa17..f0a108729820e 100644 --- a/onnxruntime/python/onnxruntime_pybind_quant.cc +++ b/onnxruntime/python/onnxruntime_pybind_quant.cc @@ -9,10 +9,8 @@ #include "contrib_ops/cpu/quantization/dequantize_blockwise_bnb4.h" #include "core/util/thread_utils.h" -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -#include -#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" -#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" +#if defined(USE_CUDA) +#include "core/providers/cuda/cuda_provider_factory.h" #endif #include #include @@ -37,6 +35,13 @@ struct npy_format_descriptor { } // namespace detail } // namespace pybind11 +#if defined(USE_CUDA) +namespace onnxruntime { +// Forward declaration; defined in provider_bridge_ort.cc and linked into pybind11 module. +ProviderInfo_CUDA* TryGetProviderInfo_CUDA(); +} // namespace onnxruntime +#endif + namespace onnxruntime { namespace python { @@ -147,23 +152,7 @@ void QuantizeMatMulBnb4Blockwise( tp.get()); } -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -namespace cuda { -void ThrowIfCudaError(cudaError_t status, const char* expression) { - if (status != cudaSuccess) { - std::ostringstream oss; - oss << expression << " failed: " << cudaGetErrorString(status); - throw std::runtime_error(oss.str()); - } -} - -struct CudaDeleter { - void operator()(void* p) const { - if (p) cudaFree(p); - } -}; - -using CudaPtr = std::unique_ptr; +#if defined(USE_CUDA) // Preprocess quantized weights for CUDA mixed-precision GEMM kernels (FpA_IntB format). // @@ -189,8 +178,6 @@ py::array_t PackWeightsForMixedGemm( int32_t K, int32_t bits, int32_t force_arch = -1) { - py::buffer_info q_weights_buf = q_weights.request(); - if (bits != 4 && bits != 8) { throw std::invalid_argument("bits must be 4 or 8"); } @@ -200,79 +187,34 @@ py::array_t PackWeightsForMixedGemm( if (bits == 4 && K % 2 != 0) { throw std::invalid_argument("K must be even for 4-bit packed weights"); } + + py::buffer_info q_weights_buf = q_weights.request(); if (q_weights_buf.ndim != 2 || q_weights_buf.shape[0] != N || q_weights_buf.shape[1] != K / (8 / bits)) { throw std::invalid_argument("q_weights must have shape (N, K / (8 / bits))"); } - int n = static_cast(N); - int k = static_cast(K); - - size_t packed_weight_bytes = static_cast(n) * static_cast(k) / (8 / bits); - py::array_t processed_weights({static_cast(packed_weight_bytes)}); - py::buffer_info processed_weights_buf = processed_weights.request(); - - auto make_cuda_ptr = [](size_t bytes) -> CudaPtr { - void* p = nullptr; - ThrowIfCudaError(cudaMalloc(&p, bytes), "cudaMalloc"); - return CudaPtr(p); - }; - - auto packed_transposed_weight_space = make_cuda_ptr(packed_weight_bytes); - int8_t* packed_transposed_weight = reinterpret_cast(packed_transposed_weight_space.get()); - - auto fpA_intB_weight_buffer_ = make_cuda_ptr(packed_weight_bytes); - int8_t* preprocessed_weight = reinterpret_cast(fpA_intB_weight_buffer_.get()); - - const uint8_t* blob_data_cpu = static_cast(q_weights_buf.ptr); - - auto blob_data_gpu_buf = make_cuda_ptr(packed_weight_bytes); - uint8_t* blob_data_gpu = reinterpret_cast(blob_data_gpu_buf.get()); - - cudaStream_t stream = cudaStreamLegacy; - ThrowIfCudaError(cudaMemcpyAsync(blob_data_gpu, blob_data_cpu, packed_weight_bytes, cudaMemcpyHostToDevice, stream), - "cudaMemcpyAsync host-to-device"); - - if (bits == 4) { - ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( - stream, packed_transposed_weight, blob_data_gpu, n, k); - } else { - // 8 bits - ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( - stream, packed_transposed_weight, blob_data_gpu, n, k); - } - - using ::onnxruntime::llm::kernels::weight_only::QuantType; - QuantType quant_type = bits == 4 ? QuantType::W4_A16 : QuantType::W8_A16; - - int sm = force_arch; - if (sm < 0) { - int device_id = 0; - ThrowIfCudaError(cudaGetDevice(&device_id), "cudaGetDevice"); - cudaDeviceProp device_prop; - ThrowIfCudaError(cudaGetDeviceProperties(&device_prop, device_id), "cudaGetDeviceProperties"); - sm = device_prop.major * 10 + device_prop.minor; + auto* info = ::onnxruntime::TryGetProviderInfo_CUDA(); + if (info == nullptr) { + throw std::runtime_error( + "CUDA provider is not available. Ensure onnxruntime was built with CUDA support " + "and that a CUDA-capable device is present."); } - sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); - auto permutation_map_buffer = make_cuda_ptr(32 * sizeof(int32_t)); - - ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( - stream, - sm, - preprocessed_weight, - packed_transposed_weight, - reinterpret_cast(permutation_map_buffer.get()), - {static_cast(k), static_cast(n)}, - quant_type); + size_t packed_weight_bytes = static_cast(N) * static_cast(K) / + static_cast(8 / bits); + py::array_t processed_weights({static_cast(packed_weight_bytes)}); + py::buffer_info out_buf = processed_weights.request(); - ThrowIfCudaError(cudaGetLastError(), "preprocess CUDA kernel launch"); - ThrowIfCudaError(cudaMemcpyAsync(processed_weights_buf.ptr, preprocessed_weight, packed_weight_bytes, cudaMemcpyDeviceToHost, stream), - "cudaMemcpyAsync device-to-host"); - ThrowIfCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); + info->PackWeightsForMixedGemm( + static_cast(q_weights_buf.ptr), + N, K, bits, force_arch, + static_cast(out_buf.ptr)); return processed_weights; } +#endif // defined(USE_CUDA) + // Pack FP4 (MXFP4) weights for MoE GEMM kernels. // // Input: q_weights in [N, K/2] layout (FP4 packed 2 per byte along K dimension, row-major) @@ -280,6 +222,7 @@ py::array_t PackWeightsForMixedGemm( // // Unlike INT4 which requires architecture-specific row permutation and interleaving, // FP4 (SM90+ TMA path) only needs a simple transpose at the nibble level. +// This function is CPU-only and does not require CUDA to be present. py::array_t PackFP4WeightsForMoE( py::array_t q_weights, int32_t N, @@ -327,8 +270,6 @@ py::array_t PackFP4WeightsForMoE( return output; } -} // namespace cuda -#endif void CreateQuantPybindModule(py::module& m) { m.def("quantize_matmul_2bits", &QuantizeMatMulNBitsBlockwise); @@ -343,14 +284,14 @@ void CreateQuantPybindModule(py::module& m) { m.def("quantize_qdq_matmul_2bits", &QuantizeQDQMatMulNBitsBlockwise); m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) - m.def("pack_weights_for_cuda_mixed_gemm", &cuda::PackWeightsForMixedGemm, +#if defined(USE_CUDA) + m.def("pack_weights_for_cuda_mixed_gemm", &PackWeightsForMixedGemm, "Pack quantized weights for CUDA mixed-precision GEMM (FpA_IntB format)", py::arg("q_weights"), py::arg("N"), py::arg("K"), py::arg("bits"), py::arg("force_arch") = -1); - m.def("pack_fp4_weights_for_cuda_moe_gemm", &cuda::PackFP4WeightsForMoE, +#endif + m.def("pack_fp4_weights_for_cuda_moe_gemm", &PackFP4WeightsForMoE, "Pack FP4 (MXFP4) weights for CUDA MoE GEMM: transpose [N,K/2] to column-major [K,N/2]", py::arg("q_weights"), py::arg("N"), py::arg("K")); -#endif } } // namespace python From 91291d42e015449caa7be434fe525bf0b24a8ddd Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 05:22:42 +0000 Subject: [PATCH 3/9] Address code review feedback: fix type casts in PackWeightsForMixedGemm - Use static_cast(d_transposed.get()) explicitly for consistency with the 8-bit path in cuda_provider_factory.cc - Fix integer division order: compute (8/bits) as size_t division (static_cast(8) / static_cast(bits)) in both cuda_provider_factory.cc and onnxruntime_pybind_quant.cc --- onnxruntime/core/providers/cuda/cuda_provider_factory.cc | 5 +++-- onnxruntime/python/onnxruntime_pybind_quant.cc | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index acfca0be129b3..0c6b8a537703a 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -205,7 +205,8 @@ struct ProviderInfo_CUDA_Impl final : ProviderInfo_CUDA { void PackWeightsForMixedGemm(const uint8_t* q_weights, int32_t N, int32_t K, int32_t bits, int32_t force_arch, int8_t* output) override { - size_t packed_weight_bytes = static_cast(N) * static_cast(K) / static_cast(8 / bits); + size_t packed_weight_bytes = static_cast(N) * static_cast(K) / + (static_cast(8) / static_cast(bits)); struct CudaMemDeleter { void operator()(void* p) const noexcept { @@ -230,7 +231,7 @@ struct ProviderInfo_CUDA_Impl final : ProviderInfo_CUDA { if (bits == 4) { ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( - stream, d_transposed.get(), d_input.get(), N, K); + stream, static_cast(d_transposed.get()), d_input.get(), N, K); } else { ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( stream, diff --git a/onnxruntime/python/onnxruntime_pybind_quant.cc b/onnxruntime/python/onnxruntime_pybind_quant.cc index f0a108729820e..e9c41ed72dbeb 100644 --- a/onnxruntime/python/onnxruntime_pybind_quant.cc +++ b/onnxruntime/python/onnxruntime_pybind_quant.cc @@ -201,7 +201,7 @@ py::array_t PackWeightsForMixedGemm( } size_t packed_weight_bytes = static_cast(N) * static_cast(K) / - static_cast(8 / bits); + (static_cast(8) / static_cast(bits)); py::array_t processed_weights({static_cast(packed_weight_bytes)}); py::buffer_info out_buf = processed_weights.request(); From 5e9fa2558dbfece12f3da0b250ee50002a95678c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 7 Jul 2026 09:16:40 +0000 Subject: [PATCH 4/9] move cuda quant preprocess to a new dll --- cmake/onnxruntime_python.cmake | 73 +++++++- docs/contrib_ops/cuda/matmul_nbits.md | 6 +- .../providers/cuda/cuda_provider_factory.cc | 68 ------- .../providers/cuda/cuda_provider_factory.h | 9 - .../python/onnxruntime_pybind_cuda_quant.cc | 166 ++++++++++++++++++ .../python/onnxruntime_pybind_quant.cc | 78 -------- .../tools/quantization/cuda_quantizer.py | 24 ++- .../test_op_matmulnbits_prepacked_cuda.py | 9 +- setup.py | 2 + 9 files changed, 267 insertions(+), 168 deletions(-) create mode 100644 onnxruntime/python/onnxruntime_pybind_cuda_quant.cc diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 1f6564a257c70..9f4658ba84411 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -231,11 +231,11 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE Python::NumPy ) -# The CUDA quantization helpers (pack_weights_for_cuda_mixed_gemm) are called through the -# ProviderInfo_CUDA interface, which dynamically loads onnxruntime_providers_cuda at runtime. -# Do NOT compile CUDA source files directly into onnxruntime_pybind11_state or link -# CUDA::cudart from it: that would create a hard libcudart.so dependency that prevents -# importing the Python module on CPU-only machines. +# The CUDA quantization helpers (pack_weights_for_cuda_mixed_gemm) are built into a +# separate extension module (onnxruntime_cuda_quant_preprocess) that is imported on +# demand. Do NOT compile CUDA source files directly into onnxruntime_pybind11_state or +# link CUDA::cudart from it: that would create a hard libcudart.so dependency that +# prevents importing the Python module on CPU-only machines. set(onnxruntime_pybind11_state_dependencies ${onnxruntime_EXTERNAL_DEPENDENCIES} @@ -304,6 +304,69 @@ else() set_target_properties(onnxruntime_pybind11_state PROPERTIES SUFFIX ".so") endif() +# --------------------------------------------------------------------------- +# Standalone CUDA weight-preprocessing extension module. +# +# The CUDA weight-packing kernels (pack_weights_for_cuda_mixed_gemm) are compiled +# into their OWN Python extension module instead of onnxruntime_pybind11_state. +# This keeps the hard libcudart dependency out of the main pybind module so that +# `import onnxruntime` still works on CPU-only machines. The module is imported +# lazily by onnxruntime.python.tools.quantization.cuda_quantizer only when CUDA +# weight prepacking is requested. +# +# It does NOT go through the provider bridge / ProviderInfo_CUDA, so it works for +# both the legacy in-tree CUDA EP build and the CUDA-EP-as-plugin build. +# +# Not built on Windows: matching the previous behavior where CUDA runtime was not +# linked into Python extension modules (DLL search path constraints since +# Python 3.8), so pack_weights_for_cuda_mixed_gemm was unavailable there. +if (onnxruntime_USE_CUDA AND NOT WIN32) + onnxruntime_add_shared_library_module(onnxruntime_cuda_quant_preprocess + "${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_cuda_quant.cc" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" + ) + include(cutlass) + onnxruntime_add_include_to_target(onnxruntime_cuda_quant_preprocess Python::Module onnxruntime_common) + target_include_directories(onnxruntime_cuda_quant_preprocess PRIVATE + ${ONNXRUNTIME_ROOT} + ${pybind11_INCLUDE_DIRS} + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + ${cutlass_SOURCE_DIR}/include + ${cutlass_SOURCE_DIR}/tools/util/include + ) + target_compile_definitions(onnxruntime_cuda_quant_preprocess PRIVATE USE_CUDA) + target_link_libraries(onnxruntime_cuda_quant_preprocess PRIVATE + onnxruntime_common + Boost::mp11 + safeint_interface + ${ABSEIL_LIBS} + CUDA::cudart + ${pybind11_lib} + Python::NumPy + ) + if (NOT MSVC) + target_compile_options(onnxruntime_cuda_quant_preprocess PRIVATE "-fvisibility=hidden") + endif() + set_target_properties(onnxruntime_cuda_quant_preprocess PROPERTIES PREFIX "" SUFFIX ".so" FOLDER "ONNXRuntime") + if (APPLE) + set_target_properties(onnxruntime_cuda_quant_preprocess PROPERTIES + INSTALL_RPATH "@loader_path" + BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH_USE_LINK_PATH FALSE) + elseif (NOT CMAKE_SYSTEM_NAME MATCHES "AIX") + target_link_options(onnxruntime_cuda_quant_preprocess PRIVATE "LINKER:-rpath=\$ORIGIN") + endif() + # Place the module next to the main pybind module inside onnxruntime/capi. + add_custom_command( + TARGET onnxruntime_cuda_quant_preprocess POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/capi + COMMAND ${CMAKE_COMMAND} -E copy + $ + $/onnxruntime/capi/ + ) +endif() + # Generate version_info.py in Windows build. # Has to be done before onnxruntime_python_srcs is set. if (WIN32) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index a09a80a4732ed..3f8e1f1d8f616 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -87,9 +87,9 @@ step is **not** performed. The offline CUDA packer exposed through Python produces this layout: ```python -from onnxruntime.capi import _pybind_state as _pybind +from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant -prepacked_flat = _pybind.pack_weights_for_cuda_mixed_gemm( +prepacked_flat = _cuda_quant.pack_weights_for_cuda_mixed_gemm( q_weight.reshape(N, -1), N, K, bits, 80 ) prepacked_b = np.asarray(prepacked_flat, dtype=np.int8).view(np.uint8).reshape(q_weight.shape) @@ -309,7 +309,7 @@ present. `ComputeInternal` then: GEMV profiling helpers, e.g. `profile_qmoe_gemv.sh`). - CUDA prepacked-weight parity tests: [onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py](../../../onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py). - These use `_pybind_state.pack_weights_for_cuda_mixed_gemm(..., 80)` to produce + These use `onnxruntime_cuda_quant_preprocess.pack_weights_for_cuda_mixed_gemm(..., 80)` to produce `weight_prepacked=1` initializers and compare their outputs against runtime fpA_intB prepacking for int4/int8 and GEMV/GEMM-shaped `M` values. - Constructor failure tests for unsupported prepacked configurations live in diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index 0c6b8a537703a..66d3617c8da9a 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -10,7 +10,6 @@ #include #include #include -#include #include @@ -21,8 +20,6 @@ #include "core/providers/cuda/cuda_stream_handle.h" #include "core/providers/cuda/gpu_data_transfer.h" #include "core/providers/cuda/math/unary_elementwise_ops_impl.h" -#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" -#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" #ifdef ENABLE_NVTX_PROFILE #include "nvtx_profile.h" @@ -201,71 +198,6 @@ struct ProviderInfo_CUDA_Impl final : ProviderInfo_CUDA { params.arena_cfg = default_memory_arena_cfg; return CUDAExecutionProvider::CreateCudaPinnedAllocator(params); } - - void PackWeightsForMixedGemm(const uint8_t* q_weights, int32_t N, int32_t K, - int32_t bits, int32_t force_arch, - int8_t* output) override { - size_t packed_weight_bytes = static_cast(N) * static_cast(K) / - (static_cast(8) / static_cast(bits)); - - struct CudaMemDeleter { - void operator()(void* p) const noexcept { - if (p) cudaFree(p); - } - }; - auto make_device_buf = [](size_t bytes) { - void* p = nullptr; - CUDA_CALL_THROW(cudaMalloc(&p, bytes)); - return std::unique_ptr(p); - }; - - auto d_input = make_device_buf(packed_weight_bytes); - auto d_transposed = make_device_buf(packed_weight_bytes); - auto d_preprocessed = make_device_buf(packed_weight_bytes); - auto d_permutation_map = make_device_buf(32 * sizeof(int32_t)); - - cudaStream_t stream = cudaStreamLegacy; - - CUDA_CALL_THROW(cudaMemcpyAsync(d_input.get(), q_weights, packed_weight_bytes, - cudaMemcpyHostToDevice, stream)); - - if (bits == 4) { - ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( - stream, static_cast(d_transposed.get()), d_input.get(), N, K); - } else { - ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( - stream, - static_cast(d_transposed.get()), - static_cast(d_input.get()), - N, K); - } - - using ::onnxruntime::llm::kernels::weight_only::QuantType; - QuantType quant_type = (bits == 4) ? QuantType::W4_A16 : QuantType::W8_A16; - - int sm = force_arch; - if (sm < 0) { - int device_id = 0; - CUDA_CALL_THROW(cudaGetDevice(&device_id)); - cudaDeviceProp device_prop; - CUDA_CALL_THROW(cudaGetDeviceProperties(&device_prop, device_id)); - sm = device_prop.major * 10 + device_prop.minor; - } - sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); - - ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( - stream, sm, - static_cast(d_preprocessed.get()), - static_cast(d_transposed.get()), - static_cast(d_permutation_map.get()), - {static_cast(K), static_cast(N)}, - quant_type); - - CUDA_CALL_THROW(cudaGetLastError()); - CUDA_CALL_THROW(cudaMemcpyAsync(output, d_preprocessed.get(), packed_weight_bytes, - cudaMemcpyDeviceToHost, stream)); - CUDA_CALL_THROW(cudaStreamSynchronize(stream)); - } } g_info; struct CUDA_Provider : Provider { diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.h b/onnxruntime/core/providers/cuda/cuda_provider_factory.h index edb48dd1e2cc6..1a4b19cb100d3 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.h +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.h @@ -55,15 +55,6 @@ struct ProviderInfo_CUDA { virtual std::shared_ptr CreateCudaAllocator(int16_t device_id, size_t gpu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, onnxruntime::CUDAExecutionProviderExternalAllocatorInfo& external_allocator_info, const OrtArenaCfg* default_memory_arena_cfg) = 0; virtual std::shared_ptr CreateCudaPinnedAllocator(int16_t device_id, size_t gpu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, const OrtArenaCfg* default_memory_arena_cfg) = 0; - // Preprocess quantized weights for CUDA mixed-precision GEMM (FpA_IntB format). - // q_weights: packed quantized weights from MatMulNBits in (N, K/(8/bits)) layout. - // output: caller-allocated buffer of size N * K / (8 / bits) bytes. - // bits: quantization bit-width (4 or 8). - // force_arch: SM version to use for permutation (-1 for auto-detect from current device). - virtual void PackWeightsForMixedGemm(const uint8_t* q_weights, int32_t N, int32_t K, - int32_t bits, int32_t force_arch, - int8_t* output) = 0; - // This function is the entry point to CUDA EP's UT cases. // All tests are only called from onnxruntime_provider_test. virtual void TestAll() { diff --git a/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc b/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc new file mode 100644 index 0000000000000..3edadd786a8c9 --- /dev/null +++ b/onnxruntime/python/onnxruntime_pybind_cuda_quant.cc @@ -0,0 +1,166 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Standalone CUDA weight-preprocessing extension module. +// +// This module is intentionally kept SEPARATE from onnxruntime_pybind11_state so +// that `import onnxruntime` never triggers a load-time dependency on the CUDA +// runtime (libcudart). The CUDA weight packing kernels below link CUDA::cudart, +// so this module has a hard libcudart dependency -- but it is imported lazily by +// onnxruntime.python.tools.quantization.cuda_quantizer only when CUDA weight +// prepacking is actually requested. +// +// This approach works for both the legacy in-tree CUDA EP build and the +// CUDA-EP-as-plugin build (onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON), because it +// does not rely on the provider bridge / ProviderInfo_CUDA interface (which is +// not available in plugin builds). + +#include +#include + +#include + +#include +#include +#include + +#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" + +namespace py = pybind11; + +namespace { + +void ThrowIfCudaError(cudaError_t status, const char* expression) { + if (status != cudaSuccess) { + std::ostringstream oss; + oss << expression << " failed: " << cudaGetErrorString(status); + throw std::runtime_error(oss.str()); + } +} + +struct CudaDeleter { + void operator()(void* p) const { + if (p) cudaFree(p); + } +}; + +using CudaPtr = std::unique_ptr; + +// Preprocess quantized weights for CUDA mixed-precision GEMM kernels (FpA_IntB format). +// +// MatMulNBits/QMoE stores quantized weights in (N, K) layout: +// - N = number of output channels (columns in weight matrix W) +// - K = number of input features (rows in weight matrix W) +// - For 4-bit: shape is (N, K/2) bytes where each byte packs 2 elements +// - For 8-bit: shape is (N, K) bytes +// +// FpA_IntB GEMM kernels expect weights in (K, N) layout (transposed) for efficient +// memory access during matrix multiplication. This function: +// 1. Transposes from (N, K) to (K, N) layout +// 2. Converts unsigned quantized values to signed int8 with zero-point adjustment +// - 4-bit: uint4 -> int8 with zero_point=8 (range [0,15] -> [-8,7]) +// - 8-bit: uint8 -> int8 with zero_point=128 (range [0,255] -> [-128,127]) +// 3. Applies architecture-specific row permutation for optimized tensor core access +// +// Input: q_weights - Quantized weights from MatMulNBits in (N, K) layout +// Output: Preprocessed weights in (K, N) layout ready for fpA_intB GEMM kernels +py::array_t PackWeightsForMixedGemm( + py::array_t q_weights, + int32_t N, + int32_t K, + int32_t bits, + int32_t force_arch = -1) { + py::buffer_info q_weights_buf = q_weights.request(); + + if (bits != 4 && bits != 8) { + throw std::invalid_argument("bits must be 4 or 8"); + } + if (N <= 0 || K <= 0) { + throw std::invalid_argument("N and K must be positive"); + } + if (bits == 4 && K % 2 != 0) { + throw std::invalid_argument("K must be even for 4-bit packed weights"); + } + if (q_weights_buf.ndim != 2 || q_weights_buf.shape[0] != N || q_weights_buf.shape[1] != K / (8 / bits)) { + throw std::invalid_argument("q_weights must have shape (N, K / (8 / bits))"); + } + + int n = static_cast(N); + int k = static_cast(K); + + size_t packed_weight_bytes = static_cast(n) * static_cast(k) / (8 / bits); + py::array_t processed_weights({static_cast(packed_weight_bytes)}); + py::buffer_info processed_weights_buf = processed_weights.request(); + + auto make_cuda_ptr = [](size_t bytes) -> CudaPtr { + void* p = nullptr; + ThrowIfCudaError(cudaMalloc(&p, bytes), "cudaMalloc"); + return CudaPtr(p); + }; + + auto packed_transposed_weight_space = make_cuda_ptr(packed_weight_bytes); + int8_t* packed_transposed_weight = reinterpret_cast(packed_transposed_weight_space.get()); + + auto fpA_intB_weight_buffer_ = make_cuda_ptr(packed_weight_bytes); + int8_t* preprocessed_weight = reinterpret_cast(fpA_intB_weight_buffer_.get()); + + const uint8_t* blob_data_cpu = static_cast(q_weights_buf.ptr); + + auto blob_data_gpu_buf = make_cuda_ptr(packed_weight_bytes); + uint8_t* blob_data_gpu = reinterpret_cast(blob_data_gpu_buf.get()); + + cudaStream_t stream = cudaStreamLegacy; + ThrowIfCudaError(cudaMemcpyAsync(blob_data_gpu, blob_data_cpu, packed_weight_bytes, cudaMemcpyHostToDevice, stream), + "cudaMemcpyAsync host-to-device"); + + if (bits == 4) { + ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( + stream, packed_transposed_weight, blob_data_gpu, n, k); + } else { + // 8 bits + ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( + stream, packed_transposed_weight, blob_data_gpu, n, k); + } + + using ::onnxruntime::llm::kernels::weight_only::QuantType; + QuantType quant_type = bits == 4 ? QuantType::W4_A16 : QuantType::W8_A16; + + int sm = force_arch; + if (sm < 0) { + int device_id = 0; + ThrowIfCudaError(cudaGetDevice(&device_id), "cudaGetDevice"); + cudaDeviceProp device_prop; + ThrowIfCudaError(cudaGetDeviceProperties(&device_prop, device_id), "cudaGetDeviceProperties"); + sm = device_prop.major * 10 + device_prop.minor; + } + sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); + + auto permutation_map_buffer = make_cuda_ptr(32 * sizeof(int32_t)); + + ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( + stream, + sm, + preprocessed_weight, + packed_transposed_weight, + reinterpret_cast(permutation_map_buffer.get()), + {static_cast(k), static_cast(n)}, + quant_type); + + ThrowIfCudaError(cudaGetLastError(), "preprocess CUDA kernel launch"); + ThrowIfCudaError(cudaMemcpyAsync(processed_weights_buf.ptr, preprocessed_weight, packed_weight_bytes, + cudaMemcpyDeviceToHost, stream), + "cudaMemcpyAsync device-to-host"); + ThrowIfCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); + + return processed_weights; +} + +} // namespace + +PYBIND11_MODULE(onnxruntime_cuda_quant_preprocess, m) { + m.doc() = "CUDA weight-only quantization preprocessing helpers (loaded on demand)."; + m.def("pack_weights_for_cuda_mixed_gemm", &PackWeightsForMixedGemm, + "Pack quantized weights for CUDA mixed-precision GEMM (FpA_IntB format)", + py::arg("q_weights"), py::arg("N"), py::arg("K"), py::arg("bits"), py::arg("force_arch") = -1); +} diff --git a/onnxruntime/python/onnxruntime_pybind_quant.cc b/onnxruntime/python/onnxruntime_pybind_quant.cc index e9c41ed72dbeb..d0d326ebf12a8 100644 --- a/onnxruntime/python/onnxruntime_pybind_quant.cc +++ b/onnxruntime/python/onnxruntime_pybind_quant.cc @@ -9,9 +9,6 @@ #include "contrib_ops/cpu/quantization/dequantize_blockwise_bnb4.h" #include "core/util/thread_utils.h" -#if defined(USE_CUDA) -#include "core/providers/cuda/cuda_provider_factory.h" -#endif #include #include #include @@ -35,13 +32,6 @@ struct npy_format_descriptor { } // namespace detail } // namespace pybind11 -#if defined(USE_CUDA) -namespace onnxruntime { -// Forward declaration; defined in provider_bridge_ort.cc and linked into pybind11 module. -ProviderInfo_CUDA* TryGetProviderInfo_CUDA(); -} // namespace onnxruntime -#endif - namespace onnxruntime { namespace python { @@ -152,69 +142,6 @@ void QuantizeMatMulBnb4Blockwise( tp.get()); } -#if defined(USE_CUDA) - -// Preprocess quantized weights for CUDA mixed-precision GEMM kernels (FpA_IntB format). -// -// MatMulNBits/QMoE stores quantized weights in (N, K) layout: -// - N = number of output channels (columns in weight matrix W) -// - K = number of input features (rows in weight matrix W) -// - For 4-bit: shape is (N, K/2) bytes where each byte packs 2 elements -// - For 8-bit: shape is (N, K) bytes -// -// FpA_IntB GEMM kernels expect weights in (K, N) layout (transposed) for efficient -// memory access during matrix multiplication. This function: -// 1. Transposes from (N, K) to (K, N) layout -// 2. Converts unsigned quantized values to signed int8 with zero-point adjustment -// - 4-bit: uint4 -> int8 with zero_point=8 (range [0,15] -> [-8,7]) -// - 8-bit: uint8 -> int8 with zero_point=128 (range [0,255] -> [-128,127]) -// 3. Applies architecture-specific row permutation for optimized tensor core access -// -// Input: q_weights - Quantized weights from MatMulNBits in (N, K) layout -// Output: Preprocessed weights in (K, N) layout ready for fpA_intB GEMM kernels -py::array_t PackWeightsForMixedGemm( - py::array_t q_weights, - int32_t N, - int32_t K, - int32_t bits, - int32_t force_arch = -1) { - if (bits != 4 && bits != 8) { - throw std::invalid_argument("bits must be 4 or 8"); - } - if (N <= 0 || K <= 0) { - throw std::invalid_argument("N and K must be positive"); - } - if (bits == 4 && K % 2 != 0) { - throw std::invalid_argument("K must be even for 4-bit packed weights"); - } - - py::buffer_info q_weights_buf = q_weights.request(); - if (q_weights_buf.ndim != 2 || q_weights_buf.shape[0] != N || q_weights_buf.shape[1] != K / (8 / bits)) { - throw std::invalid_argument("q_weights must have shape (N, K / (8 / bits))"); - } - - auto* info = ::onnxruntime::TryGetProviderInfo_CUDA(); - if (info == nullptr) { - throw std::runtime_error( - "CUDA provider is not available. Ensure onnxruntime was built with CUDA support " - "and that a CUDA-capable device is present."); - } - - size_t packed_weight_bytes = static_cast(N) * static_cast(K) / - (static_cast(8) / static_cast(bits)); - py::array_t processed_weights({static_cast(packed_weight_bytes)}); - py::buffer_info out_buf = processed_weights.request(); - - info->PackWeightsForMixedGemm( - static_cast(q_weights_buf.ptr), - N, K, bits, force_arch, - static_cast(out_buf.ptr)); - - return processed_weights; -} - -#endif // defined(USE_CUDA) - // Pack FP4 (MXFP4) weights for MoE GEMM kernels. // // Input: q_weights in [N, K/2] layout (FP4 packed 2 per byte along K dimension, row-major) @@ -284,11 +211,6 @@ void CreateQuantPybindModule(py::module& m) { m.def("quantize_qdq_matmul_2bits", &QuantizeQDQMatMulNBitsBlockwise); m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); m.def("quantize_qdq_matmul_4bits", &QuantizeQDQMatMul4BitsBlockwise); -#if defined(USE_CUDA) - m.def("pack_weights_for_cuda_mixed_gemm", &PackWeightsForMixedGemm, - "Pack quantized weights for CUDA mixed-precision GEMM (FpA_IntB format)", - py::arg("q_weights"), py::arg("N"), py::arg("K"), py::arg("bits"), py::arg("force_arch") = -1); -#endif m.def("pack_fp4_weights_for_cuda_moe_gemm", &PackFP4WeightsForMoE, "Pack FP4 (MXFP4) weights for CUDA MoE GEMM: transpose [N,K/2] to column-major [K,N/2]", py::arg("q_weights"), py::arg("N"), py::arg("K")); diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index 6b03485280d62..d064883b0aa27 100644 --- a/onnxruntime/python/tools/quantization/cuda_quantizer.py +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -43,22 +43,40 @@ def _get_torch(): def _get_pack_weights_for_cuda_mixed_gemm(): - """Return the CUDA mixed-GEMM weight prepacker from the ORT pybind module.""" + """Return the CUDA mixed-GEMM weight prepacker from the standalone CUDA module. + + The prepacker lives in ``onnxruntime.capi.onnxruntime_cuda_quant_preprocess``, a + separate extension module that links the CUDA runtime. It is imported lazily here + (never at ``import onnxruntime`` time) so that CPU-only environments are unaffected. + """ try: - from onnxruntime.capi import _pybind_state as _pybind # noqa: PLC0415 + from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant # noqa: PLC0415 except ImportError as e: raise ImportError( "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." ) from e try: - return _pybind.pack_weights_for_cuda_mixed_gemm + return _cuda_quant.pack_weights_for_cuda_mixed_gemm except AttributeError as e: raise ImportError( "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." ) from e +def has_cuda_weight_prepacking() -> bool: + """Return True if the CUDA mixed-GEMM weight prepacker is importable. + + Callers use this to skip CUDA-prepack code paths (and tests) when running against a + CPU-only or non-CUDA onnxruntime build. + """ + try: + _get_pack_weights_for_cuda_mixed_gemm() + except ImportError: + return False + return True + + def _get_quantize_matmul_nbits(): """Return MatMulNBits blockwise quantizers from the ORT pybind module.""" try: diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 406eee21c059f..0ff97478072ae 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -17,6 +17,11 @@ import onnxruntime as ort from onnxruntime.capi import _pybind_state as _pybind +try: + from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant +except ImportError: + _cuda_quant = None + @contextmanager def set_env(name: str, value: str): @@ -32,7 +37,7 @@ def set_env(name: str, value: str): @unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") -@unittest.skipUnless(hasattr(_pybind, "pack_weights_for_cuda_mixed_gemm"), "fpA_intB weight packer is unavailable") +@unittest.skipUnless(_cuda_quant is not None, "fpA_intB weight packer is unavailable") class TestMatMulNBitsPrepackedCuda(unittest.TestCase): def _quantize_weight(self, weight: np.ndarray, bits: int, block_size: int): k, n = weight.shape @@ -118,7 +123,7 @@ def _check_prepacked_parity( bias = rng.normal(0.0, 1.0, size=(n,)).astype(np.float16) if has_bias else None q_weight, scales = self._quantize_weight(weight, bits, block_size) - prepacked_flat = _pybind.pack_weights_for_cuda_mixed_gemm(q_weight.reshape(n, -1), n, k, bits, force_arch) + prepacked_flat = _cuda_quant.pack_weights_for_cuda_mixed_gemm(q_weight.reshape(n, -1), n, k, bits, force_arch) prepacked_weight = np.asarray(prepacked_flat, dtype=np.int8).view(np.uint8).reshape(q_weight.shape) raw_model = self._make_model((m, k), q_weight, scales, bits, block_size, weight_prepacked=0, bias=bias) diff --git a/setup.py b/setup.py index 62ced38819f2c..ad5784896feca 100644 --- a/setup.py +++ b/setup.py @@ -375,6 +375,7 @@ def finalize_options(self): if platform.system() == "Linux" or platform.system() == "AIX": libs = [ "onnxruntime_pybind11_state.so", + "onnxruntime_cuda_quant_preprocess.so", "libdnnl.so.2", "libmklml_intel.so", "libmklml_gnu.so", @@ -422,6 +423,7 @@ def finalize_options(self): elif platform.system() == "Darwin": libs = [ "onnxruntime_pybind11_state.so", + "onnxruntime_cuda_quant_preprocess.so", "libdnnl.2.dylib", "mimalloc.so", "libonnxruntime*.dylib", From c948c7db22ec13953388af1392b2841c6d71320f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 7 Jul 2026 23:15:47 +0000 Subject: [PATCH 5/9] Fix CPU/Windows pybind build: isolate CUDA quant module, restore ORT_NO_CUDA_IN_PYBIND The new onnxruntime_pybind_cuda_quant.cc was picked up by the python/*.cc glob and compiled into onnxruntime_pybind11_state, breaking CPU-only builds (fatal error: cuda_runtime.h not found) and re-adding CUDA link deps on every platform. Remove it from the main pybind sources so it only builds inside the standalone onnxruntime_cuda_quant_preprocess module. Restore the ORT_NO_CUDA_IN_PYBIND compile definition on Windows CUDA builds: it is still consumed by onnxruntime_pybind_mlvalue.cc / onnxruntime_pybind_ortvalue.cc to avoid direct CUDA runtime calls (cudaMemcpy/cudaStreamSynchronize) in the pybind module. Dropping it caused LNK2001 unresolved cudart symbols on Windows. Add onnxruntime_cuda_quant_preprocess.so to dl_libs so manylinux wheels package it (the manylinux path builds 'data' from dl_libs only, not libs). --- cmake/onnxruntime_python.cmake | 17 +++++++++++++++++ setup.py | 4 ++++ 2 files changed, 21 insertions(+) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 9f4658ba84411..a8590d746a759 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -20,6 +20,13 @@ file(GLOB onnxruntime_pybind_srcs CONFIGURE_DEPENDS ${onnxruntime_pybind_srcs_pattern} ) +# onnxruntime_pybind_cuda_quant.cc is compiled into the standalone +# onnxruntime_cuda_quant_preprocess extension module (see below), not into +# onnxruntime_pybind11_state. It includes and links CUDA::cudart, +# so compiling it into the main pybind module would break CPU-only builds and +# re-introduce the hard libcudart dependency this design avoids. +list(REMOVE_ITEM onnxruntime_pybind_srcs ${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_cuda_quant.cc) + if(onnxruntime_ENABLE_TRAINING) list(REMOVE_ITEM onnxruntime_pybind_srcs ${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_module.cc) endif() @@ -237,6 +244,16 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE # link CUDA::cudart from it: that would create a hard libcudart.so dependency that # prevents importing the Python module on CPU-only machines. +# On Windows CUDA builds the main pybind module must still be built with +# ORT_NO_CUDA_IN_PYBIND so that onnxruntime_pybind_mlvalue.cc / onnxruntime_pybind_ortvalue.cc +# do not call CUDA runtime APIs (e.g. cudaMemcpy) directly. Starting with Python 3.8 on +# Windows, PATH is no longer used to resolve extension-module DLL dependencies, so we avoid +# linking the CUDA runtime into the pybind module and rely on os.add_dll_directory() / +# the provider bridge instead. +if (onnxruntime_USE_CUDA AND WIN32) + target_compile_definitions(onnxruntime_pybind11_state PRIVATE ORT_NO_CUDA_IN_PYBIND) +endif() + set(onnxruntime_pybind11_state_dependencies ${onnxruntime_EXTERNAL_DEPENDENCIES} ${pybind11_dep} diff --git a/setup.py b/setup.py index ad5784896feca..ab50af2d44308 100644 --- a/setup.py +++ b/setup.py @@ -389,6 +389,10 @@ def finalize_options(self): dl_libs.append(providers_cann) dl_libs.append(providers_qnn) dl_libs.append("libonnxruntime.so*") + # onnxruntime_cuda_quant_preprocess.so is a standalone CUDA extension module. It must be + # listed in dl_libs (not just libs) so that manylinux wheels include it: the manylinux + # packaging path builds "data" from dl_libs only (see the is_manylinux block below). + dl_libs.append("onnxruntime_cuda_quant_preprocess.so") # DNNL, TensorRT, OpenVINO, and QNN EPs are built as shared libs libs.extend(["libonnxruntime_providers_shared.so"]) libs.extend(["libonnxruntime_providers_dnnl.so"]) From 4f47b23bafa16590c5dd1e088a797db82f884542 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 01:02:01 +0000 Subject: [PATCH 6/9] Fix CI: define ORT_NO_CUDA_IN_PYBIND on all CUDA builds; drop braced scalar init Two CI regressions from removing the CUDA runtime link from the main pybind module: - Linux CUDA/TensorRT/Plugin test jobs failed at import with 'undefined symbol: cudaMemcpy'. The pybind module no longer links CUDA::cudart, but ORT_NO_CUDA_IN_PYBIND was only defined on Windows, so onnxruntime_pybind_mlvalue.cc still emitted direct cudaMemcpy calls on Linux. Define ORT_NO_CUDA_IN_PYBIND for all CUDA builds so host/device copies route through the CUDA provider bridge (ProviderInfo_CUDA) on every platform. - macOS arm64 CPU-only build jobs (cpu/coreml/xnnpack/webgpu) failed compiling PackFP4WeightsForMoE with clang -Werror,-Wbraced-scalar-init. This CPU-only function is now compiled unconditionally; replace the braced scalar py::array_t output({size}) with the array_t(ssize_t count) constructor output(size), which yields the identical 1-D array. --- cmake/onnxruntime_python.cmake | 16 +++++++++------- onnxruntime/python/onnxruntime_pybind_quant.cc | 2 +- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index a8590d746a759..8bfdf3b24e641 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -244,13 +244,15 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE # link CUDA::cudart from it: that would create a hard libcudart.so dependency that # prevents importing the Python module on CPU-only machines. -# On Windows CUDA builds the main pybind module must still be built with -# ORT_NO_CUDA_IN_PYBIND so that onnxruntime_pybind_mlvalue.cc / onnxruntime_pybind_ortvalue.cc -# do not call CUDA runtime APIs (e.g. cudaMemcpy) directly. Starting with Python 3.8 on -# Windows, PATH is no longer used to resolve extension-module DLL dependencies, so we avoid -# linking the CUDA runtime into the pybind module and rely on os.add_dll_directory() / -# the provider bridge instead. -if (onnxruntime_USE_CUDA AND WIN32) +# Because the main pybind module no longer links CUDA::cudart, it must be built with +# ORT_NO_CUDA_IN_PYBIND on all platforms so that onnxruntime_pybind_mlvalue.cc / +# onnxruntime_pybind_ortvalue.cc do not call CUDA runtime APIs (e.g. cudaMemcpy) directly +# (which would leave undefined symbols like cudaMemcpy in the module). Instead the module +# routes host/device copies through the CUDA provider bridge (ProviderInfo_CUDA). On Windows +# this also matches the pre-existing behavior where, starting with Python 3.8, PATH is no +# longer used to resolve extension-module DLL dependencies, so we rely on +# os.add_dll_directory() / the provider bridge rather than linking the CUDA runtime. +if (onnxruntime_USE_CUDA) target_compile_definitions(onnxruntime_pybind11_state PRIVATE ORT_NO_CUDA_IN_PYBIND) endif() diff --git a/onnxruntime/python/onnxruntime_pybind_quant.cc b/onnxruntime/python/onnxruntime_pybind_quant.cc index d0d326ebf12a8..b2b6ed77c6296 100644 --- a/onnxruntime/python/onnxruntime_pybind_quant.cc +++ b/onnxruntime/python/onnxruntime_pybind_quant.cc @@ -170,7 +170,7 @@ py::array_t PackFP4WeightsForMoE( int K_half = K / 2; int N_half = N / 2; size_t out_size = static_cast(K) * static_cast(N_half); - py::array_t output({static_cast(out_size)}); + py::array_t output(static_cast(out_size)); py::buffer_info out_buf = output.request(); uint8_t* dst = static_cast(out_buf.ptr); std::memset(dst, 0, out_size); From 8b32759c0f044202048cd38be4d9383e386ad663 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 02:17:57 +0000 Subject: [PATCH 7/9] Clean up --- cmake/onnxruntime_python.cmake | 12 -------- docs/cuda_plugin_ep/QUICK_START.md | 2 +- docs/cuda_plugin_ep/cuda_plugin_ep_design.md | 8 +++-- .../python/onnxruntime_pybind_mlvalue.cc | 29 ------------------- .../python/onnxruntime_pybind_ortvalue.cc | 9 ------ 5 files changed, 6 insertions(+), 54 deletions(-) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 8bfdf3b24e641..dd7d4ab6bdfa7 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -244,18 +244,6 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE # link CUDA::cudart from it: that would create a hard libcudart.so dependency that # prevents importing the Python module on CPU-only machines. -# Because the main pybind module no longer links CUDA::cudart, it must be built with -# ORT_NO_CUDA_IN_PYBIND on all platforms so that onnxruntime_pybind_mlvalue.cc / -# onnxruntime_pybind_ortvalue.cc do not call CUDA runtime APIs (e.g. cudaMemcpy) directly -# (which would leave undefined symbols like cudaMemcpy in the module). Instead the module -# routes host/device copies through the CUDA provider bridge (ProviderInfo_CUDA). On Windows -# this also matches the pre-existing behavior where, starting with Python 3.8, PATH is no -# longer used to resolve extension-module DLL dependencies, so we rely on -# os.add_dll_directory() / the provider bridge rather than linking the CUDA runtime. -if (onnxruntime_USE_CUDA) - target_compile_definitions(onnxruntime_pybind11_state PRIVATE ORT_NO_CUDA_IN_PYBIND) -endif() - set(onnxruntime_pybind11_state_dependencies ${onnxruntime_EXTERNAL_DEPENDENCIES} ${pybind11_dep} diff --git a/docs/cuda_plugin_ep/QUICK_START.md b/docs/cuda_plugin_ep/QUICK_START.md index ab7b4308f13e5..7b971d621083f 100644 --- a/docs/cuda_plugin_ep/QUICK_START.md +++ b/docs/cuda_plugin_ep/QUICK_START.md @@ -146,7 +146,7 @@ sess = ort.InferenceSession( **Python `OrtValue` host/device copies:** -`OrtValue.update_inplace()` and `OrtValue.numpy()` work with CUDA plugin tensors after the plugin has been registered. On Linux, the ONNX Runtime Python binding links the CUDA runtime and can fall back to direct `cudaMemcpy` if the legacy CUDA provider bridge is unavailable. On Windows, the Python binding is built with `ORT_NO_CUDA_IN_PYBIND`, so it cannot call CUDA runtime APIs directly; host/device copies must use the data-transfer implementation registered by the CUDA plugin library. If `OrtValue.update_inplace()` fails with a message about the CUDA provider interface or an unsupported GPU device, verify that the plugin library is registered before creating or updating CUDA `OrtValue` objects. +`OrtValue.update_inplace()` and `OrtValue.numpy()` work with CUDA plugin tensors after the plugin has been registered. The Python binding cannot call CUDA runtime APIs directly; host/device copies must use the data-transfer implementation registered by the CUDA plugin library. If `OrtValue.update_inplace()` fails with a message about the CUDA provider interface or an unsupported GPU device, verify that the plugin library is registered before creating or updating CUDA `OrtValue` objects. ### External GPU Allocator Options diff --git a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md index a06f1de95b99e..7dfa2bd2d8667 100644 --- a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md +++ b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md @@ -344,10 +344,12 @@ This is intentionally conservative and correct for the plugin EP's first sync in The Python `OrtValue` helpers (`update_inplace()` for host-to-device and `numpy()` for device-to-host) historically reached CUDA copies through the legacy provider bridge (`GetProviderInfo_CUDA()`). That bridge requires the provider shared library to export `GetProvider()`, which the CUDA plugin intentionally does not export. -The fallback path is platform-specific: +To keep working when the bridge is absent (as with the plugin EP), the pybind can reach CUDA copies two ways: the legacy provider bridge (`TryGetProviderInfo_CUDA()`) and a plugin-registered `OrtDataTransfer` copy function (`CreateDataTransferMemCpy()`, backed by the plugin EP's `IDataTransfer`). It tries whichever is available and throws if neither is, in which case a CUDA `OrtValue` copy cannot be performed. -- On non-Windows CUDA builds, `onnxruntime_pybind11_state` links `CUDA::cudart`. If `TryGetProviderInfo_CUDA()` fails, pybind can copy directly with `cudaMemcpy`; host-to-device copies synchronize the default stream, matching `ProviderInfo_CUDA::cudaMemcpy_HostToDevice()`. -- On Windows CUDA builds, pybind is compiled with `ORT_NO_CUDA_IN_PYBIND` and does not link CUDA runtime APIs. If `TryGetProviderInfo_CUDA()` fails, pybind must obtain an `OrtDataTransfer` copy function from the registered plugin EP. Without a registered plugin data-transfer implementation, CUDA `OrtValue.update_inplace()` cannot copy host data into the plugin-owned device tensor. +The two code paths differ only in which mechanism they try first, and this does not change the outcome (exactly one applies in a given build): + +- `OrtValue.update_inplace(numpy_array)` / `OrtValue.numpy()` (in `onnxruntime_pybind_ortvalue.cc`) try the provider bridge first, then fall back to the plugin `OrtDataTransfer`. +- `OrtValue.update_inplace(OrtValue)` (`UpdateOrtValueInplace` in `onnxruntime_pybind_mlvalue.cc`) tries the plugin `OrtDataTransfer` first, then falls back to the built-in CUDA provider copy functions. ### 5.2 Handle Access Path diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc index 10e55259f834e..5ea5f42958926 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc @@ -23,10 +23,6 @@ #include "core/framework/kernel_registry.h" #include "core/framework/provider_options_utils.h" -#if defined(USE_CUDA) && !defined(ORT_NO_CUDA_IN_PYBIND) -#include -#endif - #ifdef USE_DML using Microsoft::WRL::ComPtr; @@ -184,23 +180,6 @@ int32_t GetTensorProtoType(const OrtValue& ort_value) { } #ifdef USE_CUDA -namespace { - -#if !defined(ORT_NO_CUDA_IN_PYBIND) -void CudaRuntimeMemCpy(void* dst, const void* src, size_t num_bytes, cudaMemcpyKind kind) { - const auto copy_result = cudaMemcpy(dst, src, num_bytes, kind); - ORT_ENFORCE(copy_result == cudaSuccess, "cudaMemcpy failed: ", cudaGetErrorString(copy_result)); - - if (kind == cudaMemcpyHostToDevice) { - // Match ProviderInfo_CUDA::cudaMemcpy_HostToDevice: cudaMemcpy() uses the default - // stream, and pageable host-to-device copies can return before DMA to device is done. - const auto sync_result = cudaStreamSynchronize(0); - ORT_ENFORCE(sync_result == cudaSuccess, "cudaStreamSynchronize failed: ", cudaGetErrorString(sync_result)); - } -} -#endif - -} // namespace void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes) { if (TryGetProviderInfo_CUDA() != nullptr) { @@ -208,11 +187,7 @@ void CpuToCudaMemCpy(void* dst, const void* src, size_t num_bytes) { return; } -#if !defined(ORT_NO_CUDA_IN_PYBIND) - CudaRuntimeMemCpy(dst, src, num_bytes, cudaMemcpyHostToDevice); -#else ORT_THROW("CUDA provider interface is not available for host-to-device copy."); -#endif } void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { @@ -221,11 +196,7 @@ void CudaToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { return; } -#if !defined(ORT_NO_CUDA_IN_PYBIND) - CudaRuntimeMemCpy(dst, src, num_bytes, cudaMemcpyDeviceToHost); -#else ORT_THROW("CUDA provider interface is not available for device-to-host copy."); -#endif } const std::unordered_map* GetCudaToHostMemCpyFunction(const OrtDevice& device) { diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index cf7f86a0b9e41..7bf9325cf2208 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -205,7 +205,6 @@ void addOrtValueMethods(pybind11::module& m) { #ifdef USE_CUDA if (device.Vendor() == OrtDevice::VendorIds::NVIDIA) { MemCpyFunc cpu_to_device_copy_fn = CpuToCudaMemCpy; -#if defined(ORT_NO_CUDA_IN_PYBIND) if (TryGetProviderInfo_CUDA() != nullptr) { if (!IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); @@ -217,12 +216,6 @@ void addOrtValueMethods(pybind11::module& m) { "Unsupported GPU device: Cannot find the supported GPU device."); } } -#else - if (TryGetProviderInfo_CUDA() != nullptr && - !IsCudaDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { - throw std::runtime_error("The provided device id doesn't match any available GPUs on the machine."); - } -#endif onnxruntime::python::CopyDataToTensor( py_values, @@ -467,12 +460,10 @@ void addOrtValueMethods(pybind11::module& m) { switch (device.Vendor()) { #ifdef USE_CUDA case OrtDevice::VendorIds::NVIDIA: -#if defined(ORT_NO_CUDA_IN_PYBIND) if (TryGetProviderInfo_CUDA() == nullptr) { return GetPyObjFromTensor(*ml_value, nullptr, nullptr, /*zero_copy_non_owning=*/true); } -#endif return GetPyObjFromTensor(*ml_value, nullptr, GetCudaToHostMemCpyFunction(device), /*zero_copy_non_owning=*/true); #endif From 9897958570db7c63d3c5513e416dfecf3aa58f07 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 05:45:19 +0000 Subject: [PATCH 8/9] add a test --- .../test_op_matmulnbits_prepacked_cuda.py | 18 ++++++++++++++++++ tools/ci_build/build.py | 4 +++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 0ff97478072ae..6a894d37ff81b 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -8,6 +8,7 @@ from __future__ import annotations import os +import sys import unittest from contextlib import contextmanager @@ -23,6 +24,23 @@ _cuda_quant = None +@unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") +class TestCudaQuantPreprocessAvailability(unittest.TestCase): + """Guards against ``TestMatMulNBitsPrepackedCuda`` being silently skipped. + + The fpA_intB weight packer ships as a standalone extension module + (``onnxruntime_cuda_quant_preprocess``) in the ``capi/`` directory. + ``TestMatMulNBitsPrepackedCuda`` is skipped when that module fails to import, + so a broken/missing module would hide those tests without any failure. When + CUDA is available (and the module is built, i.e. non-Windows), the import + must succeed. + """ + + @unittest.skipIf(sys.platform.startswith("win"), "cuda quant preprocess module is not built on Windows") + def test_import_succeeds_when_library_exists(self): + self.assertIsNotNone(_cuda_quant, "onnxruntime_cuda_quant_preprocess exists but module-level import failed") + + @contextmanager def set_env(name: str, value: str): old_value = os.environ.get(name) diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 8317018d33c64..4061e4fafbe7d 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -1822,7 +1822,9 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): if not args.disable_contrib_ops: run_subprocess( - [sys.executable, "-m", "unittest", "discover", "-s", "quantization"], cwd=cwd, dll_path=dll_path + [sys.executable, "-m", "unittest", "discover", "-s", "quantization", "-v"], + cwd=cwd, + dll_path=dll_path, ) if args.enable_transformers_tool_test and (sys.version_info.major, sys.version_info.minor) < ( From 081ac225b5535cb4aec8aba50567e0c001c60b17 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 08:05:47 +0000 Subject: [PATCH 9/9] Add build option and torch reference --- cmake/CMakeLists.txt | 1 + cmake/onnxruntime_python.cmake | 10 +- .../tools/quantization/cuda_quantizer.py | 218 ++++++++++++++++-- .../test_op_matmulnbits_prepacked_cuda.py | 53 +++-- setup.py | 9 +- 5 files changed, 244 insertions(+), 47 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 09e307e124316..ad446214cfc8f 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -77,6 +77,7 @@ cmake_dependent_option(onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS "Build with CUD cmake_dependent_option(onnxruntime_USE_CUDA_NHWC_OPS "Build CUDA with NHWC op support" ON "onnxruntime_USE_CUDA" OFF) cmake_dependent_option(onnxruntime_BUILD_CUDA_EP_AS_PLUGIN "Build CUDA EP as a separate plugin shared library instead of the legacy in-tree provider" OFF "onnxruntime_USE_CUDA" OFF) +option(onnxruntime_BUILD_CUDA_QUANT_PREPROCESS "Build CUDA weight-packing module onnxruntime_cuda_quant_preprocess.so" ON) option(onnxruntime_CUDA_MINIMAL "Build CUDA without any operations apart from memcpy ops. Usefuel for a very minial TRT build" OFF) option(onnxruntime_ENABLE_CUDA_LINE_NUMBER_INFO "When building with CUDA support, generate device code line number information." OFF) option(onnxruntime_USE_OPENVINO "Build with OpenVINO support" OFF) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index dd7d4ab6bdfa7..14b01f1f25473 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -317,9 +317,11 @@ endif() # The CUDA weight-packing kernels (pack_weights_for_cuda_mixed_gemm) are compiled # into their OWN Python extension module instead of onnxruntime_pybind11_state. # This keeps the hard libcudart dependency out of the main pybind module so that -# `import onnxruntime` still works on CPU-only machines. The module is imported -# lazily by onnxruntime.python.tools.quantization.cuda_quantizer only when CUDA -# weight prepacking is requested. +# `import onnxruntime` still works on CPU-only machines. +# +# Production weight packing is done in PyTorch (cuda_quantizer.py); this module is +# retained only as a byte-parity oracle for that PyTorch packer. It is gated by +# onnxruntime_BUILD_CUDA_QUANT_PREPROCESS. # # It does NOT go through the provider bridge / ProviderInfo_CUDA, so it works for # both the legacy in-tree CUDA EP build and the CUDA-EP-as-plugin build. @@ -327,7 +329,7 @@ endif() # Not built on Windows: matching the previous behavior where CUDA runtime was not # linked into Python extension modules (DLL search path constraints since # Python 3.8), so pack_weights_for_cuda_mixed_gemm was unavailable there. -if (onnxruntime_USE_CUDA AND NOT WIN32) +if (onnxruntime_USE_CUDA AND NOT WIN32 AND onnxruntime_BUILD_CUDA_QUANT_PREPROCESS) onnxruntime_add_shared_library_module(onnxruntime_cuda_quant_preprocess "${ONNXRUNTIME_ROOT}/python/onnxruntime_pybind_cuda_quant.cc" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index d064883b0aa27..2ab532b3d0e58 100644 --- a/onnxruntime/python/tools/quantization/cuda_quantizer.py +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -7,9 +7,13 @@ """CUDA weight-only quantization helpers. This module contains small Python utilities for producing the weight layouts -consumed by CUDA weight-only kernels. The helpers deliberately wrap the same C++ -pybind entry points used by runtime prepacking so tests and model builders can -generate byte-identical quantized weights. +consumed by CUDA weight-only kernels. The blockwise quantizers wrap the same C++ +pybind entry points used by runtime prepacking, and the mixed-GEMM weight packer +is a PyTorch reimplementation of the runtime CUDA packing, so tests and model +builders can generate byte-identical quantized weights. The PyTorch packer runs +on CUDA when a device is available and falls back to CPU otherwise, which is the +only option on platforms where the standalone CUDA packer is not built (Windows). +A GPU-gated parity test validates it against that standalone CUDA packer. Two storage families are exposed: @@ -24,10 +28,14 @@ from __future__ import annotations +import functools +import logging from typing import TYPE_CHECKING import numpy as np +_logger = logging.getLogger(__name__) + if TYPE_CHECKING: import torch @@ -43,40 +51,209 @@ def _get_torch(): def _get_pack_weights_for_cuda_mixed_gemm(): - """Return the CUDA mixed-GEMM weight prepacker from the standalone CUDA module. + """Return the standalone CUDA mixed-GEMM weight packer (parity oracle). - The prepacker lives in ``onnxruntime.capi.onnxruntime_cuda_quant_preprocess``, a - separate extension module that links the CUDA runtime. It is imported lazily here - (never at ``import onnxruntime`` time) so that CPU-only environments are unaffected. + Production packing uses the PyTorch implementation (``_pack_weights_for_cuda_mixed_gemm``). + This standalone packer lives in ``onnxruntime.capi.onnxruntime_cuda_quant_preprocess``, a + separate extension module that links the CUDA runtime (built only on non-Windows CUDA + builds). It is imported lazily here (never at ``import onnxruntime`` time) and is used by + the parity test to validate the PyTorch packer byte-for-byte. """ try: from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant # noqa: PLC0415 except ImportError as e: raise ImportError( - "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." + "The standalone CUDA weight packer (onnxruntime_cuda_quant_preprocess) is unavailable; " + "it is built only on non-Windows onnxruntime-gpu CUDA builds." ) from e try: return _cuda_quant.pack_weights_for_cuda_mixed_gemm except AttributeError as e: - raise ImportError( - "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." - ) from e + raise ImportError("onnxruntime_cuda_quant_preprocess is missing pack_weights_for_cuda_mixed_gemm.") from e def has_cuda_weight_prepacking() -> bool: - """Return True if the CUDA mixed-GEMM weight prepacker is importable. + """Return True if mixed-GEMM weight prepacking is available. - Callers use this to skip CUDA-prepack code paths (and tests) when running against a - CPU-only or non-CUDA onnxruntime build. + Prepacking is implemented with PyTorch (CUDA when available, CPU otherwise), so it is + available whenever torch is importable. Callers use this to skip prepack code paths + (and tests) when torch is unavailable. """ try: - _get_pack_weights_for_cuda_mixed_gemm() + _get_torch() except ImportError: return False return True +@functools.lru_cache(maxsize=1) +def _warn_cpu_prepack_once() -> None: + _logger.warning( + "CUDA device is not available; packing mixed-GEMM weights on CPU with PyTorch. " + "This is correct but significantly slower for large Mixture-of-Experts models. " + "Pack on a CUDA-enabled machine for best performance." + ) + + +def _prepack_device(): + """Pick the torch device for mixed-GEMM weight packing (CUDA if available, else CPU).""" + torch = _get_torch() + if torch.cuda.is_available(): + return torch.device("cuda") + _warn_cpu_prepack_once() + return torch.device("cpu") + + +def _preprocess_weights_for_mixed_gemm_torch(tensor, bits: int, sm: int): + """PyTorch port of the runtime CUDA ``preprocess_weights_for_mixed_gemm``. + + ``tensor`` is a signed int8 weight in ``(K, N/pack)`` packed row-major layout on any + device. Returns the CUTLASS mixed-GEMM layout with the same shape/dtype/device. This + mirrors ``preprocess_weights_for_mixed_gemm_cuda`` (permute_B_rows -> subbyte_transpose + -> interleave_column_major -> add_bias_and_interleave) so its output is byte-identical + to the standalone CUDA packer, for both the SM80 (Ampere) and SM90 (Hopper) layouts. + """ + torch = _get_torch() + bits_a = 16 # fp16/bf16 activations + bits_b = 4 if bits == 4 else 8 + + if tensor.dim() == 2: + tensor = tensor.unsqueeze(0) + + permutation_map = { + "16_8": [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15], + "16_4": [ + 0, + 1, + 8, + 9, + 16, + 17, + 24, + 25, + 2, + 3, + 10, + 11, + 18, + 19, + 26, + 27, + 4, + 5, + 12, + 13, + 20, + 21, + 28, + 29, + 6, + 7, + 14, + 15, + 22, + 23, + 30, + 31, + ], + } + mma_shape_n = 8 + b_rows_per_mma = 8 * 16 // bits_b + + num_experts, num_rows, num_cols = tensor.shape[0], tensor.shape[1], tensor.shape[2] + if num_rows % b_rows_per_mma != 0 or num_cols % mma_shape_n != 0: + raise ValueError( + f"weight shape (rows={num_rows}, packed_cols={num_cols}) is incompatible with mixed-GEMM " + f"packing (rows must be a multiple of {b_rows_per_mma}, packed cols a multiple of {mma_shape_n})." + ) + + # permute_B_rows_for_mixed_gemm + if sm < 100: + pmap = permutation_map[f"{bits_a}_{bits_b}"] + row_idx = [(r // b_rows_per_mma) * b_rows_per_mma + pmap[r % b_rows_per_mma] for r in range(num_rows)] + tensor = tensor[:, row_idx, :] + + # subbyte_transpose + original_shape = tensor.shape + if bits_b == 4: + u = tensor.view(torch.uint8) + high = (u >> 4).permute(0, 2, 1).unsqueeze(2) + low = ((u << 4) >> 4).permute(0, 2, 1).unsqueeze(2) + merged = torch.cat([low, high], dim=2).reshape(u.shape[0], -1, u.shape[1]) + merged = merged[:, :, 0::2] + merged[:, :, 1::2] * 16 + tensor = merged.view(torch.int8).reshape(original_shape) + else: + tensor = tensor.permute(0, 2, 1).reshape(original_shape) + + # interleave_column_major_tensor + interleave = bits_a // bits_b + if interleave > 1 and sm < 90: + rows_per_tile = 128 * 8 // bits_a + elts_in_int32 = 32 // bits_b + if num_rows % elts_in_int32 != 0 or num_rows % rows_per_tile != 0: + raise ValueError(f"num_rows ({num_rows}) is incompatible with column-interleave tiling.") + tensor = tensor.reshape( + num_experts, -1, interleave, num_rows // rows_per_tile, rows_per_tile * 4 // elts_in_int32 + ) + tensor = tensor.permute(0, 1, 3, 2, 4).reshape(original_shape) + + # add_bias_and_interleave_quantized_tensor_inplace + if bits_b == 8: + t = tensor.to(torch.int64) # widen so the +128 rebias cannot overflow int8 + t += -256 * (t > 127).to(torch.int64) + 128 + t = t.reshape(-1, 4)[:, [0, 2, 1, 3]].reshape(original_shape) + tensor = t.to(torch.uint8).view(torch.int8) + else: + u = tensor.view(torch.uint8) + high = (u >> 4).unsqueeze(-1) + low = ((u << 4) >> 4).unsqueeze(-1) + merged = torch.cat([low, high], dim=-1).reshape(u.shape[0], u.shape[1], -1) + merged = merged.reshape(-1, 8)[:, [0, 2, 4, 6, 1, 3, 5, 7]].reshape(merged.shape) + merged = merged.to(torch.int16) + merged += -16 * (merged > 7).to(torch.int16) + 8 + merged = merged[:, :, 0::2] + merged[:, :, 1::2] * 16 + tensor = merged.to(torch.uint8).view(torch.int8) + + return tensor.squeeze(0).contiguous() + + +def _pack_weights_for_cuda_mixed_gemm(q_weights, n: int, k: int, bits: int, force_arch: int = 80) -> np.ndarray: + """PyTorch implementation of the CUDA ``pack_weights_for_cuda_mixed_gemm``. + + ``q_weights`` is ORT's unsigned MatMulNBits/QMoE storage ``(N, K/pack)`` (uint8). Returns + a flat ``int8`` numpy array with the CUTLASS mixed-GEMM layout, byte-identical to the + standalone CUDA packer. Runs on CUDA when available, otherwise on CPU. + """ + torch = _get_torch() + bits = int(bits) + force_arch = int(force_arch) + if bits not in (4, 8): + raise ValueError(f"bits must be 4 or 8, got {bits}.") + if force_arch not in (80, 90): + raise ValueError(f"force_arch must be 80 (SM80) or 90 (SM90), got {force_arch}.") + pack = 8 // bits + device = _prepack_device() + + q = torch.as_tensor(np.ascontiguousarray(q_weights)).view(torch.uint8).reshape(n, k // pack).to(device) + + # Front-end adaptor: transpose ORT (N, K) -> (K, N) and convert unsigned -> signed int8. + if bits == 4: + low = (q & 0x0F).to(torch.int16) + high = (q >> 4).to(torch.int16) + unpacked = torch.empty((n, k), dtype=torch.int16, device=device) + unpacked[:, 0::2] = low + unpacked[:, 1::2] = high + signed_t = (unpacked - 8).transpose(0, 1).contiguous() # (K, N), zero point 8 + packed_t = ((signed_t[:, 0::2] & 0x0F) | ((signed_t[:, 1::2] & 0x0F) << 4)).to(torch.uint8).view(torch.int8) + else: + signed_t = (q.to(torch.int16) - 128).transpose(0, 1).contiguous() # (K, N), zero point 128 + packed_t = signed_t.to(torch.uint8).view(torch.int8) + + out = _preprocess_weights_for_mixed_gemm_torch(packed_t.contiguous(), bits, force_arch) + return out.reshape(-1).cpu().numpy() + + def _get_quantize_matmul_nbits(): """Return MatMulNBits blockwise quantizers from the ORT pybind module.""" try: @@ -164,8 +341,7 @@ def qmoe_per_channel_quantize( When ``prepack`` is true, returned weights have shape ``[K, N/pack]``. Otherwise, returned weights keep raw per-channel storage ``[N, K/pack]``. - CUDA prepacking requires ``pack_weights_for_cuda_mixed_gemm`` from an - onnxruntime-gpu CUDA build. + Prepacking uses PyTorch (CUDA when available, CPU otherwise). """ torch = _get_torch() @@ -177,14 +353,12 @@ def qmoe_per_channel_quantize( if not prepack: return qweight, scales - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() - n, k = weights.shape pack = 8 // int(bits) if n % pack != 0: raise ValueError(f"N ({n}) must be divisible by {pack} for CUDA QMoE prepacked weights.") - packed = pack_weights_for_cuda_mixed_gemm(qweight.numpy(), n, k, int(bits), force_arch) + packed = _pack_weights_for_cuda_mixed_gemm(qweight.numpy(), n, k, int(bits), force_arch) packed = np.asarray(packed).view(np.uint8).reshape(k, n // pack) return torch.from_numpy(np.ascontiguousarray(packed)), scales @@ -348,7 +522,7 @@ def matmulnbits_prepacked_blockwise_quantize( unsigned_full_range=unsigned_full_range, ) - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() + pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm packed = pack_weights_for_cuda_mixed_gemm(qweight.reshape(n, -1).numpy(), n, k, bits, force_arch) packed = np.asarray(packed).view(np.uint8).reshape(qweight.shape) packed = torch.from_numpy(np.ascontiguousarray(packed)) @@ -393,7 +567,7 @@ def qmoe_prepacked_blockwise_quantize( unsigned_full_range=unsigned_full_range, ) - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() + pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm packed = pack_weights_for_cuda_mixed_gemm(qweight.reshape(n, -1).numpy(), n, k, bits, force_arch) packed = np.asarray(packed).view(np.uint8).reshape(k, n // pack) torch = _get_torch() diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 6a894d37ff81b..9d3a378f6db19 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -8,7 +8,6 @@ from __future__ import annotations import os -import sys import unittest from contextlib import contextmanager @@ -17,6 +16,7 @@ import onnxruntime as ort from onnxruntime.capi import _pybind_state as _pybind +from onnxruntime.quantization.cuda_quantizer import _pack_weights_for_cuda_mixed_gemm try: from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant @@ -24,23 +24,6 @@ _cuda_quant = None -@unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") -class TestCudaQuantPreprocessAvailability(unittest.TestCase): - """Guards against ``TestMatMulNBitsPrepackedCuda`` being silently skipped. - - The fpA_intB weight packer ships as a standalone extension module - (``onnxruntime_cuda_quant_preprocess``) in the ``capi/`` directory. - ``TestMatMulNBitsPrepackedCuda`` is skipped when that module fails to import, - so a broken/missing module would hide those tests without any failure. When - CUDA is available (and the module is built, i.e. non-Windows), the import - must succeed. - """ - - @unittest.skipIf(sys.platform.startswith("win"), "cuda quant preprocess module is not built on Windows") - def test_import_succeeds_when_library_exists(self): - self.assertIsNotNone(_cuda_quant, "onnxruntime_cuda_quant_preprocess exists but module-level import failed") - - @contextmanager def set_env(name: str, value: str): old_value = os.environ.get(name) @@ -194,5 +177,39 @@ def test_int8_sm90_prepacked_weight_matches_runtime_prepack(self): self._check_sm90_parity(bits=8, block_size=128, m=32) +@unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") +@unittest.skipUnless(_cuda_quant is not None, "standalone CUDA weight packer (parity oracle) is unavailable") +class TestCudaQuantizerTorchPackerParity(unittest.TestCase): + """Validate the PyTorch mixed-GEMM packer in cuda_quantizer.py against the CUDA oracle. + + ``cuda_quantizer._pack_weights_for_cuda_mixed_gemm`` (PyTorch, used in production, and the + only option on Windows where the standalone module is not built) must be byte-identical to + the standalone ``onnxruntime_cuda_quant_preprocess.pack_weights_for_cuda_mixed_gemm`` (the + CUDA code the runtime prepack uses). This test is the guard against silent drift; it only + runs where the oracle is built (non-Windows CUDA). + """ + + def _check(self, bits: int, force_arch: int, n: int, k: int): + pack = 8 // bits + rng = np.random.default_rng(20260708 + bits * 100 + force_arch + n + k) + q = rng.integers(0, 256, size=(n, k // pack), dtype=np.uint8) + oracle = np.asarray(_cuda_quant.pack_weights_for_cuda_mixed_gemm(q, n, k, bits, force_arch), dtype=np.int8) + torch_out = _pack_weights_for_cuda_mixed_gemm(q, n, k, bits, force_arch).astype(np.int8) + self.assertEqual(oracle.shape, torch_out.shape, f"shape mismatch bits={bits} arch={force_arch} N={n} K={k}") + np.testing.assert_array_equal( + torch_out, oracle, err_msg=f"byte mismatch bits={bits} arch={force_arch} N={n} K={k}" + ) + + def test_torch_packer_matches_cuda_oracle(self): + # Cover both weight bit-widths, both mixed-GEMM layouts (SM80/SM90), and a GPT-OSS-20B + # MoE shape (fused gate+up FC1 [5760, 2880] and down FC2 [2880, 2880]). + shapes = [(256, 256), (512, 256), (256, 512), (5760, 2880), (2880, 2880), (128, 128)] + for bits in (4, 8): + for force_arch in (80, 90): + for n, k in shapes: + with self.subTest(bits=bits, force_arch=force_arch, n=n, k=k): + self._check(bits, force_arch, n, k) + + if __name__ == "__main__": unittest.main() diff --git a/setup.py b/setup.py index ab50af2d44308..58aadc75b5010 100644 --- a/setup.py +++ b/setup.py @@ -389,9 +389,12 @@ def finalize_options(self): dl_libs.append(providers_cann) dl_libs.append(providers_qnn) dl_libs.append("libonnxruntime.so*") - # onnxruntime_cuda_quant_preprocess.so is a standalone CUDA extension module. It must be - # listed in dl_libs (not just libs) so that manylinux wheels include it: the manylinux - # packaging path builds "data" from dl_libs only (see the is_manylinux block below). + # onnxruntime_cuda_quant_preprocess.so is a standalone CUDA extension module used only as a + # byte-parity oracle for the PyTorch weight packer. It is built (and thus present here) only + # when the CMake option onnxruntime_BUILD_CUDA_QUANT_PREPROCESS is ON. The glob-based filters below + # drop missing files, so listing it here is a no-op when it was not built. It must be listed in + # dl_libs (not just libs) so that manylinux test wheels include it: the manylinux packaging path + # builds "data" from dl_libs only (see the is_manylinux block below). dl_libs.append("onnxruntime_cuda_quant_preprocess.so") # DNNL, TensorRT, OpenVINO, and QNN EPs are built as shared libs libs.extend(["libonnxruntime_providers_shared.so"])