|
| 1 | +// Copyright (c) Microsoft Corporation. All rights reserved. |
| 2 | +// Licensed under the MIT License. |
| 3 | + |
| 4 | +// Standalone CUDA weight-preprocessing extension module. |
| 5 | +// |
| 6 | +// This module is intentionally kept SEPARATE from onnxruntime_pybind11_state so |
| 7 | +// that `import onnxruntime` never triggers a load-time dependency on the CUDA |
| 8 | +// runtime (libcudart). The CUDA weight packing kernels below link CUDA::cudart, |
| 9 | +// so this module has a hard libcudart dependency -- but it is imported lazily by |
| 10 | +// onnxruntime.python.tools.quantization.cuda_quantizer only when CUDA weight |
| 11 | +// prepacking is actually requested. |
| 12 | +// |
| 13 | +// This approach works for both the legacy in-tree CUDA EP build and the |
| 14 | +// CUDA-EP-as-plugin build (onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON), because it |
| 15 | +// does not rely on the provider bridge / ProviderInfo_CUDA interface (which is |
| 16 | +// not available in plugin builds). |
| 17 | + |
| 18 | +#include <pybind11/pybind11.h> |
| 19 | +#include <pybind11/numpy.h> |
| 20 | + |
| 21 | +#include <cuda_runtime.h> |
| 22 | + |
| 23 | +#include <memory> |
| 24 | +#include <sstream> |
| 25 | +#include <stdexcept> |
| 26 | + |
| 27 | +#include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" |
| 28 | +#include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" |
| 29 | + |
| 30 | +namespace py = pybind11; |
| 31 | + |
| 32 | +namespace { |
| 33 | + |
| 34 | +void ThrowIfCudaError(cudaError_t status, const char* expression) { |
| 35 | + if (status != cudaSuccess) { |
| 36 | + std::ostringstream oss; |
| 37 | + oss << expression << " failed: " << cudaGetErrorString(status); |
| 38 | + throw std::runtime_error(oss.str()); |
| 39 | + } |
| 40 | +} |
| 41 | + |
| 42 | +struct CudaDeleter { |
| 43 | + void operator()(void* p) const { |
| 44 | + if (p) cudaFree(p); |
| 45 | + } |
| 46 | +}; |
| 47 | + |
| 48 | +using CudaPtr = std::unique_ptr<void, CudaDeleter>; |
| 49 | + |
| 50 | +// Preprocess quantized weights for CUDA mixed-precision GEMM kernels (FpA_IntB format). |
| 51 | +// |
| 52 | +// MatMulNBits/QMoE stores quantized weights in (N, K) layout: |
| 53 | +// - N = number of output channels (columns in weight matrix W) |
| 54 | +// - K = number of input features (rows in weight matrix W) |
| 55 | +// - For 4-bit: shape is (N, K/2) bytes where each byte packs 2 elements |
| 56 | +// - For 8-bit: shape is (N, K) bytes |
| 57 | +// |
| 58 | +// FpA_IntB GEMM kernels expect weights in (K, N) layout (transposed) for efficient |
| 59 | +// memory access during matrix multiplication. This function: |
| 60 | +// 1. Transposes from (N, K) to (K, N) layout |
| 61 | +// 2. Converts unsigned quantized values to signed int8 with zero-point adjustment |
| 62 | +// - 4-bit: uint4 -> int8 with zero_point=8 (range [0,15] -> [-8,7]) |
| 63 | +// - 8-bit: uint8 -> int8 with zero_point=128 (range [0,255] -> [-128,127]) |
| 64 | +// 3. Applies architecture-specific row permutation for optimized tensor core access |
| 65 | +// |
| 66 | +// Input: q_weights - Quantized weights from MatMulNBits in (N, K) layout |
| 67 | +// Output: Preprocessed weights in (K, N) layout ready for fpA_intB GEMM kernels |
| 68 | +py::array_t<int8_t> PackWeightsForMixedGemm( |
| 69 | + py::array_t<uint8_t> q_weights, |
| 70 | + int32_t N, |
| 71 | + int32_t K, |
| 72 | + int32_t bits, |
| 73 | + int32_t force_arch = -1) { |
| 74 | + py::buffer_info q_weights_buf = q_weights.request(); |
| 75 | + |
| 76 | + if (bits != 4 && bits != 8) { |
| 77 | + throw std::invalid_argument("bits must be 4 or 8"); |
| 78 | + } |
| 79 | + if (N <= 0 || K <= 0) { |
| 80 | + throw std::invalid_argument("N and K must be positive"); |
| 81 | + } |
| 82 | + if (bits == 4 && K % 2 != 0) { |
| 83 | + throw std::invalid_argument("K must be even for 4-bit packed weights"); |
| 84 | + } |
| 85 | + if (q_weights_buf.ndim != 2 || q_weights_buf.shape[0] != N || q_weights_buf.shape[1] != K / (8 / bits)) { |
| 86 | + throw std::invalid_argument("q_weights must have shape (N, K / (8 / bits))"); |
| 87 | + } |
| 88 | + |
| 89 | + int n = static_cast<int>(N); |
| 90 | + int k = static_cast<int>(K); |
| 91 | + |
| 92 | + size_t packed_weight_bytes = static_cast<size_t>(n) * static_cast<size_t>(k) / (8 / bits); |
| 93 | + py:: array_t< int8_t> processed_weights({ static_cast<pybind11:: ssize_t>(packed_weight_bytes)}); |
| 94 | + py::buffer_info processed_weights_buf = processed_weights.request(); |
| 95 | + |
| 96 | + auto make_cuda_ptr = [](size_t bytes) -> CudaPtr { |
| 97 | + void* p = nullptr; |
| 98 | + ThrowIfCudaError(cudaMalloc(&p, bytes), "cudaMalloc"); |
| 99 | + return CudaPtr(p); |
| 100 | + }; |
| 101 | + |
| 102 | + auto packed_transposed_weight_space = make_cuda_ptr(packed_weight_bytes); |
| 103 | + int8_t* packed_transposed_weight = reinterpret_cast<int8_t*>(packed_transposed_weight_space.get()); |
| 104 | + |
| 105 | + auto fpA_intB_weight_buffer_ = make_cuda_ptr(packed_weight_bytes); |
| 106 | + int8_t* preprocessed_weight = reinterpret_cast<int8_t*>(fpA_intB_weight_buffer_.get()); |
| 107 | + |
| 108 | + const uint8_t* blob_data_cpu = static_cast<const uint8_t*>(q_weights_buf.ptr); |
| 109 | + |
| 110 | + auto blob_data_gpu_buf = make_cuda_ptr(packed_weight_bytes); |
| 111 | + uint8_t* blob_data_gpu = reinterpret_cast<uint8_t*>(blob_data_gpu_buf.get()); |
| 112 | + |
| 113 | + cudaStream_t stream = cudaStreamLegacy; |
| 114 | + ThrowIfCudaError(cudaMemcpyAsync(blob_data_gpu, blob_data_cpu, packed_weight_bytes, cudaMemcpyHostToDevice, stream), |
| 115 | + "cudaMemcpyAsync host-to-device"); |
| 116 | + |
| 117 | + if (bits == 4) { |
| 118 | + ::onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( |
| 119 | + stream, packed_transposed_weight, blob_data_gpu, n, k); |
| 120 | + } else { |
| 121 | + // 8 bits |
| 122 | + ::onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( |
| 123 | + stream, packed_transposed_weight, blob_data_gpu, n, k); |
| 124 | + } |
| 125 | + |
| 126 | + using ::onnxruntime::llm::kernels::weight_only::QuantType; |
| 127 | + QuantType quant_type = bits == 4 ? QuantType::W4_A16 : QuantType::W8_A16; |
| 128 | + |
| 129 | + int sm = force_arch; |
| 130 | + if (sm < 0) { |
| 131 | + int device_id = 0; |
| 132 | + ThrowIfCudaError(cudaGetDevice(&device_id), "cudaGetDevice"); |
| 133 | + cudaDeviceProp device_prop; |
| 134 | + ThrowIfCudaError(cudaGetDeviceProperties(&device_prop, device_id), "cudaGetDeviceProperties"); |
| 135 | + sm = device_prop.major * 10 + device_prop.minor; |
| 136 | + } |
| 137 | + sm = ::onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(sm); |
| 138 | + |
| 139 | + auto permutation_map_buffer = make_cuda_ptr(32 * sizeof(int32_t)); |
| 140 | + |
| 141 | + ::onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( |
| 142 | + stream, |
| 143 | + sm, |
| 144 | + preprocessed_weight, |
| 145 | + packed_transposed_weight, |
| 146 | + reinterpret_cast<int32_t*>(permutation_map_buffer.get()), |
| 147 | + {static_cast<size_t>(k), static_cast<size_t>(n)}, |
| 148 | + quant_type); |
| 149 | + |
| 150 | + ThrowIfCudaError(cudaGetLastError(), "preprocess CUDA kernel launch"); |
| 151 | + ThrowIfCudaError(cudaMemcpyAsync(processed_weights_buf.ptr, preprocessed_weight, packed_weight_bytes, |
| 152 | + cudaMemcpyDeviceToHost, stream), |
| 153 | + "cudaMemcpyAsync device-to-host"); |
| 154 | + ThrowIfCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); |
| 155 | + |
| 156 | + return processed_weights; |
| 157 | +} |
| 158 | + |
| 159 | +} // namespace |
| 160 | + |
| 161 | +PYBIND11_MODULE(onnxruntime_cuda_quant_preprocess, m) { |
| 162 | + m.doc() = "CUDA weight-only quantization preprocessing helpers (loaded on demand)."; |
| 163 | + m.def("pack_weights_for_cuda_mixed_gemm", &PackWeightsForMixedGemm, |
| 164 | + "Pack quantized weights for CUDA mixed-precision GEMM (FpA_IntB format)", |
| 165 | + py::arg("q_weights"), py::arg("N"), py::arg("K"), py::arg("bits"), py::arg("force_arch") = -1); |
| 166 | +} |
0 commit comments