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 3f6976c3c8955..14b01f1f25473 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() @@ -231,22 +238,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 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} @@ -315,6 +311,71 @@ 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. +# +# 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. +# +# 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 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" + "${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/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_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_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 diff --git a/onnxruntime/python/onnxruntime_pybind_quant.cc b/onnxruntime/python/onnxruntime_pybind_quant.cc index 7220153b4fa17..b2b6ed77c6296 100644 --- a/onnxruntime/python/onnxruntime_pybind_quant.cc +++ b/onnxruntime/python/onnxruntime_pybind_quant.cc @@ -9,11 +9,6 @@ #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" -#endif #include #include #include @@ -147,132 +142,6 @@ 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; - -// 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; -} - // 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 +149,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, @@ -300,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); @@ -327,8 +197,6 @@ py::array_t PackFP4WeightsForMoE( return output; } -} // namespace cuda -#endif void CreateQuantPybindModule(py::module& m) { m.def("quantize_matmul_2bits", &QuantizeMatMulNBitsBlockwise); @@ -343,14 +211,9 @@ 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, - "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, + 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 diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index 6b03485280d62..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,20 +51,207 @@ 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 standalone CUDA mixed-GEMM weight packer (parity oracle). + + 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 _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." + "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 _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 + 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 mixed-GEMM weight prepacking is available. + + 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_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(): @@ -146,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() @@ -159,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 @@ -330,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)) @@ -375,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 406eee21c059f..9d3a378f6db19 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -16,6 +16,12 @@ 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 +except ImportError: + _cuda_quant = None @contextmanager @@ -32,7 +38,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 +124,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) @@ -171,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 62ced38819f2c..58aadc75b5010 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", @@ -388,6 +389,13 @@ 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 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"]) libs.extend(["libonnxruntime_providers_dnnl.so"]) @@ -422,6 +430,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", 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) < (