From 3c565437a5ac65fa7b6fe332bb57a5afad00d206 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 4 Mar 2026 08:59:31 +0900 Subject: [PATCH] [CUDA] Quantized GEMV (#3180) --- mlx/backend/cuda/CMakeLists.txt | 1 - mlx/backend/cuda/quantized/qmm/CMakeLists.txt | 2 + .../cuda/quantized/{qmv.cu => qmm/fp_qmv.cu} | 116 ++++---- mlx/backend/cuda/quantized/qmm/qmm.cpp | 107 +++++++ mlx/backend/cuda/quantized/qmm/qmm.h | 58 ++++ .../cuda/quantized/qmm/qmm_impl_sm90.cuh | 19 -- mlx/backend/cuda/quantized/qmm/qmv.cu | 280 ++++++++++++++++++ mlx/backend/cuda/quantized/qmv.h | 22 -- mlx/backend/cuda/quantized/qqmm.cpp | 10 +- mlx/backend/cuda/quantized/quantized.cpp | 83 ++++-- 10 files changed, 575 insertions(+), 123 deletions(-) rename mlx/backend/cuda/quantized/{qmv.cu => qmm/fp_qmv.cu} (79%) create mode 100644 mlx/backend/cuda/quantized/qmm/qmv.cu delete mode 100644 mlx/backend/cuda/quantized/qmv.h diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 223be22f..bc1bfa9f 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -56,7 +56,6 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/quantized/affine_quantize.cu ${CMAKE_CURRENT_SOURCE_DIR}/quantized/fp_quantize.cu - ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmv.cu ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu diff --git a/mlx/backend/cuda/quantized/qmm/CMakeLists.txt b/mlx/backend/cuda/quantized/qmm/CMakeLists.txt index 69266c65..74a1c8a6 100644 --- a/mlx/backend/cuda/quantized/qmm/CMakeLists.txt +++ b/mlx/backend/cuda/quantized/qmm/CMakeLists.txt @@ -1,6 +1,8 @@ target_sources( mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu + ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n16_m1.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n32_m1.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n64_m2.cu diff --git a/mlx/backend/cuda/quantized/qmv.cu b/mlx/backend/cuda/quantized/qmm/fp_qmv.cu similarity index 79% rename from mlx/backend/cuda/quantized/qmv.cu rename to mlx/backend/cuda/quantized/qmm/fp_qmv.cu index 8ddd93a9..33eb4adc 100644 --- a/mlx/backend/cuda/quantized/qmv.cu +++ b/mlx/backend/cuda/quantized/qmm/fp_qmv.cu @@ -2,7 +2,7 @@ #include "mlx/backend/cuda/device/utils.cuh" #include "mlx/backend/cuda/kernel_utils.cuh" -#include "mlx/backend/cuda/quantized/qmv.h" +#include "mlx/backend/cuda/quantized/qmm/qmm.h" #include "mlx/backend/cuda/quantized/quantized_utils.cuh" #include "mlx/backend/cuda/quantized/quantized_utils.h" #include "mlx/dtype_utils.h" @@ -10,12 +10,14 @@ #include #include -namespace mlx::core::cu { +namespace mlx::core { + +constexpr int rows_per_block = 8; + +namespace cu { namespace cg = cooperative_groups; -static constexpr int rows_per_block = 8; - template __device__ void adjust_matrix_offsets( const T*& x, @@ -199,6 +201,8 @@ __global__ void fp_qmv_batched( mat, scales, vec, out, rows, cols); } +} // namespace cu + template void dispatch_1_2_4(int n, F&& f) { switch (n) { @@ -221,11 +225,13 @@ void fp_qmv( array& out, int bits, int group_size, - int M, - int N, - int K, - CommandEncoder& encoder, + cu::CommandEncoder& encoder, Stream s) { + uint32_t M = x.shape(-2); + uint32_t N = out.shape(-1); + uint32_t K = x.shape(-1); + uint32_t B = out.size() / (M * N); + // Make sure the last two dims of x and w, s, b are contiguous. This should // be relaxed for x. array vec = ensure_row_contiguous_matrix(x, encoder, s); @@ -240,7 +246,6 @@ void fp_qmv( using T = cuda_type_t; if constexpr (!std::is_same_v) { dim3 block_dims{WARP_SIZE, rows_per_block}; - uint32_t B = out.size() / (M * N); uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block; const uint32_t* mat_ptr = gpu_ptr(mat); const T* vec_ptr = gpu_ptr(vec); @@ -256,55 +261,56 @@ void fp_qmv( n = 2; } dispatch_1_2_4(n, [&](auto n) { - dispatch_bool(B > 1, [&](auto batched) { - if (!batched.value) { - auto kernel = - fp_qmv_single; - if (bits == 8) { - kernel = fp_qmv_single; - } else if (group_size == 16) { - kernel = fp_qmv_single; - } - encoder.add_kernel_node( - kernel, - {static_cast(M), blocks_y}, - block_dims, - mat_ptr, - gpu_ptr(scales), - vec_ptr, - gpu_ptr(out), - N, - K); - } else { - auto kernel = - fp_qmv_batched; - if (bits == 8) { - kernel = fp_qmv_batched; - } else if (group_size == 16) { - kernel = fp_qmv_batched; - } - encoder.add_kernel_node( - kernel, - {static_cast(M), blocks_y, B}, - block_dims, - mat_ptr, - gpu_ptr(scales), - vec_ptr, - gpu_ptr(out), - N, - K, - vec.ndim() - 2, - const_param(vec.shape()), - const_param(vec.strides()), - mat.ndim() - 2, - const_param(mat.shape()), - const_param(mat.strides()), - const_param(scales.strides())); + if (B == 1) { + auto kernel = + cu::fp_qmv_single; + if (bits == 8) { + kernel = cu::fp_qmv_single; + } else if (group_size == 16) { + kernel = + cu::fp_qmv_single; } - }); + encoder.add_kernel_node( + kernel, + {uint32_t(x.size() / K), blocks_y}, + block_dims, + mat_ptr, + gpu_ptr(scales), + vec_ptr, + gpu_ptr(out), + N, + K); + } else { + auto kernel = + cu::fp_qmv_batched; + if (bits == 8) { + kernel = + cu::fp_qmv_batched; + } else if (group_size == 16) { + kernel = + cu::fp_qmv_batched; + } + encoder.add_kernel_node( + kernel, + {M, blocks_y, B}, + block_dims, + mat_ptr, + gpu_ptr(scales), + vec_ptr, + gpu_ptr(out), + N, + K, + vec.ndim() - 2, + const_param(vec.shape()), + const_param(vec.strides()), + mat.ndim() - 2, + const_param(mat.shape()), + const_param(mat.strides()), + const_param(scales.strides())); + } }); } }); } -} // namespace mlx::core::cu +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qmm/qmm.cpp b/mlx/backend/cuda/quantized/qmm/qmm.cpp index 613c0245..c26e2184 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm.cpp +++ b/mlx/backend/cuda/quantized/qmm/qmm.cpp @@ -21,6 +21,46 @@ void qmm_impl_sm90( Stream s); #endif // defined(MLX_CUDA_SM90A_ENABLED) +bool supports_qmm_sm90( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device) { + if (device.compute_capability_major() != 9) { + return false; + } + int k = x.shape(-1); + if (k % 64 != 0) { + return false; + } + if (!biases) { + return false; + } + if (!x.flags().row_contiguous || !w.flags().row_contiguous || + !scales.flags().row_contiguous || !biases->flags().row_contiguous) { + return false; + } + if (!transpose) { + return false; + } + if (bits % 2 != 0) { + return false; + } + if (group_size < k) { + return false; + } + if (mode != QuantizationMode::Affine) { + return false; + } + return true; +} + void qmm_sm90( const array& x, const array& w, @@ -57,4 +97,71 @@ void qmm_sm90( #endif // defined(MLX_CUDA_SM90A_ENABLED) } +bool supports_fp_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device) { + bool non_batched = w.ndim() == 2; + int k = x.shape(-1); + int n = out.shape(-1); + int vec_batch = non_batched ? x.size() / k : x.shape(-2); + if (vec_batch > 8) { + return false; + } + if (!transpose) { + return false; + } + if (mode == QuantizationMode::Affine) { + return false; + } + return true; +} + +bool supports_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device) { + int m = out.shape(-2); + int n = out.shape(-1); + int k = x.shape(-1); + int l = out.size() / (m * n); + if (l > 1) { + return false; + } + if (n % 8 != 0 || k % 8 != 0) { + return false; + } + if (!x.flags().row_contiguous || !w.flags().row_contiguous || + !scales.flags().row_contiguous) { + return false; + } + if (biases && !biases->flags().row_contiguous) { + return false; + } + if (!transpose) { + return false; + } + if (bits % 2 != 0) { + return false; + } + if (mode != QuantizationMode::Affine) { + return false; + } + return true; +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qmm/qmm.h b/mlx/backend/cuda/quantized/qmm/qmm.h index ad04ba11..efcd8eaf 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm.h +++ b/mlx/backend/cuda/quantized/qmm/qmm.h @@ -3,11 +3,24 @@ #pragma once #include "mlx/backend/cuda/device.h" +#include "mlx/primitives.h" #include namespace mlx::core { +bool supports_qmm_sm90( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device); + void qmm_sm90( const array& x, const array& w, @@ -19,4 +32,49 @@ void qmm_sm90( cu::CommandEncoder& encoder, Stream s); +bool supports_fp_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device); + +void fp_qmv( + const array& x, + const array& w, + const array& scales, + array& out, + int bits, + int group_size, + cu::CommandEncoder& encoder, + Stream s); + +bool supports_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device); + +void qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder); + } // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh index 17ad7ee9..b8dc63ef 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh +++ b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh @@ -186,25 +186,6 @@ void qmm_impl_sm90( int n = out.shape(-1); int k = x.shape(-1); int l = out.size() / (m * n); - if (k % 64 != 0) { - throw std::runtime_error(fmt::format("{} K must be multiples of 64.", tag)); - } - if (!x.flags().row_contiguous) { - throw std::runtime_error( - fmt::format("{} Activations must be row contiguous.", tag)); - } - if (!w.flags().row_contiguous) { - throw std::runtime_error( - fmt::format("{} Weights must be row contiguous.", tag)); - } - if (!scales_.flags().row_contiguous) { - throw std::runtime_error( - fmt::format("{} Scales must be row contiguous.", tag)); - } - if (!biases_.flags().row_contiguous) { - throw std::runtime_error( - fmt::format("{} Biases must be row contiguous.", tag)); - } // FIXME: Copy happens for every call. array scales = transpose_last_2_dims(scales_, encoder, s); diff --git a/mlx/backend/cuda/quantized/qmm/qmv.cu b/mlx/backend/cuda/quantized/qmm/qmv.cu new file mode 100644 index 00000000..c7109f5c --- /dev/null +++ b/mlx/backend/cuda/quantized/qmm/qmv.cu @@ -0,0 +1,280 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/cuda/quantized/qmm/qmm.h" +#include "mlx/dtype_utils.h" + +#include +#include +#include +#include + +namespace mlx::core { + +namespace cu { + +namespace cg = cooperative_groups; + +// Fused vectorized dequantize and multiply-add: +// w_dq = w * scale + bias +// out = fma(x, w_dq, out) +template +__device__ __forceinline__ void +dequant_fma(const T* x, const Q* w, T scale, T bias, float* out) { + // Read x/w into registers. + auto x_vec = *(reinterpret_cast*>(x)); + auto w_vec = *(reinterpret_cast*>(w)); + // Output is assumed to be registers. + auto* out_vec = reinterpret_cast*>(out); + + // Dequantize w. + cutlass::NumericArrayConverter converter_tq; + cutlass::Array w_dq = converter_tq(w_vec); + w_dq = w_dq * scale + bias; + + // Promote x/w to float. + static_assert(!cuda::std::is_same_v); + cutlass::NumericArrayConverter converter_ft; + cutlass::Array x_f = converter_ft(x_vec); + cutlass::Array w_f = converter_ft(w_dq); + + // Multiply and add. + *out_vec = cutlass::fma(x_f, w_f, *out_vec); +} + +// Specialized for float which does not need promotions. +template +__device__ __forceinline__ void +dequant_fma(const float* x, const Q* w, float scale, float bias, float* out) { + auto x_vec = *(reinterpret_cast*>(x)); + auto w_vec = *(reinterpret_cast*>(w)); + auto* out_vec = reinterpret_cast*>(out); + + cutlass::NumericArrayConverter converter; + cutlass::Array w_dq = converter(w_vec); +#pragma unroll + for (int i = 0; i < N; ++i) { + w_dq[i] = w_dq[i] * scale + bias; + } + + *out_vec = cutlass::fma(x_vec, w_dq, *out_vec); +} + +template < + int rows_per_block, + int elems_per_thread, + int group_size, + bool has_bias, + bool has_residue_k, + typename T, + typename Q> +__global__ void qmv_kernel( + const T* x, + const Q* w, + const T* scales, + const T* biases, + T* out, + int n, + int k) { + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + // The row that this warp handles. + int row = block.group_index().x * rows_per_block + warp.meta_group_rank(); + if (row >= n) { + return; + } + + // Advance pointers of x/out. + x += block.group_index().y * k; + out += block.group_index().y * n; + + // For sub-byte Q, pointer moves by 8bits for each advance, e.g. w += 1 would + // move past 2 elements for 4-bit Q. + constexpr int w_step = 8 / cuda::std::min(8, cute::sizeof_bits_v); + + // How many groups (and scales/biases) in a row. + int groups_per_row = k / group_size; + + // Advance w/scales/biases to current row. + w += static_cast(row) * k / w_step; + scales += static_cast(row) * groups_per_row; + if constexpr (has_bias) { + biases += static_cast(row) * groups_per_row; + } + + // Accumulations of current row. + float sums[elems_per_thread] = {}; + + auto dequant_fma_tile = [&](int idx) { + T scale = scales[idx / group_size]; + T bias{0}; + if constexpr (has_bias) { + bias = biases[idx / group_size]; + } + dequant_fma(x + idx, w + idx / w_step, scale, bias, sums); + }; + + // Loop over k dimension. + constexpr int elems_per_warp = WARP_SIZE * elems_per_thread; + for (int r = 0; r < k / elems_per_warp; ++r) { + int idx = warp.thread_rank() * elems_per_thread + r * elems_per_warp; + dequant_fma_tile(idx); + } + + // Handle remaining elements in k dimension. + if constexpr (has_residue_k) { + int rest = k % elems_per_warp; + int idx = warp.thread_rank() * elems_per_thread + k - rest; + if (idx < k) { + dequant_fma_tile(idx); + } + } + + // Result for current row. + float sum{0}; +#pragma unroll + for (int i = 0; i < elems_per_thread; ++i) { + sum += sums[i]; + } + sum = cg::reduce(warp, sum, cg::plus{}); + + // Write result for current warp, which maps to rows 1-to-1. + if (warp.thread_rank() == 0) { + out[row] = static_cast(sum); + } +} + +template +void qmv( + const T* x, + const Q* w, + const T* scales, + const T* biases, + T* out, + int m, + int n, + int k, + F&& launch_kernel) { + constexpr int rows_per_block = 8; + constexpr int elems_per_thread = 8; + + dim3 num_blocks{uint32_t(cuda::ceil_div(n, rows_per_block)), uint32_t(m)}; + dim3 block_dims{WARP_SIZE, rows_per_block}; + void* args[] = {&x, &w, &scales, &biases, &out, &n, &k}; + + dispatch_bool(k % (WARP_SIZE * elems_per_thread), [&](auto has_residue_k) { + auto* kernel = &qmv_kernel< + rows_per_block, + elems_per_thread, + group_size, + has_bias, + has_residue_k.value, + T, + Q>; + launch_kernel( + reinterpret_cast(kernel), num_blocks, block_dims, args); + }); +} + +} // namespace cu + +template +inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) { + if (dtype == float32) { + f.template operator()(); + } else if (dtype == float16) { + f.template operator()(); + } else if (dtype == bfloat16) { + f.template operator()(); + } else { + throw std::invalid_argument( + fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype))); + } +} + +template +inline void +dispatch_quant_types(int bits, QuantizationMode mode, const char* tag, F&& f) { + if (mode == QuantizationMode::Mxfp4) { + f.template operator()(); + } else if (mode == QuantizationMode::Mxfp8) { + f.template operator()(); + } else if (mode == QuantizationMode::Nvfp4) { + f.template operator()(); + } else { + if (bits == 2) { + f.template operator()(); + } else if (bits == 4) { + f.template operator()(); + } else if (bits == 8) { + f.template operator()(); + } else { + throw std::invalid_argument( + fmt::format("{} {}-bit quantization is not supported.", tag, bits)); + } + } +} + +template +inline void dispatch_groups(int group_size, const char* tag, F&& f) { + if (group_size == 16) { + f.template operator()<16>(); + } else if (group_size == 32) { + f.template operator()<32>(); + } else if (group_size == 64) { + f.template operator()<64>(); + } else if (group_size == 128) { + f.template operator()<128>(); + } else { + throw std::invalid_argument( + fmt::format("{} Group size {} is not supported.", tag, group_size)); + } +} + +void qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder) { + const char* tag = "[quantized_matmul]"; + int m = out.shape(-2); + int n = out.shape(-1); + int k = x.shape(-1); + + dispatch_element_types(out.dtype(), tag, [&]() { + dispatch_bool(biases.has_value(), [&](auto has_bias) { + dispatch_quant_types(bits, mode, tag, [&]() { + dispatch_groups(group_size, tag, [&]() { + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + if (biases) { + encoder.set_input_array(*biases); + } + encoder.set_output_array(out); + cu::qmv( + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(scales), + biases ? gpu_ptr(*biases) : nullptr, + gpu_ptr(out), + m, + n, + k, + [&](auto* kernel, dim3 num_blocks, dim3 block_dims, void** args) { + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args); + }); + }); + }); + }); + }); +} + +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qmv.h b/mlx/backend/cuda/quantized/qmv.h deleted file mode 100644 index 9fc3719a..00000000 --- a/mlx/backend/cuda/quantized/qmv.h +++ /dev/null @@ -1,22 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#pragma once - -#include "mlx/backend/cuda/device.h" - -namespace mlx::core::cu { - -void fp_qmv( - const array& x, - const array& w, - const array& scales, - array& out, - int bits, - int group_size, - int M, - int N, - int K, - CommandEncoder& encoder, - Stream s); - -} // namespace mlx::core::cu diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp index 7ff5f0a3..afce96e3 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/mlx/backend/cuda/quantized/qqmm.cpp @@ -1,7 +1,7 @@ // Copyright © 2025 Apple Inc. #include "mlx/backend/cuda/device.h" -#include "mlx/backend/cuda/quantized/qmv.h" +#include "mlx/backend/cuda/quantized/qmm/qmm.h" #include "mlx/backend/cuda/quantized/qqmm_impl.h" #include "mlx/backend/cuda/quantized/qqmm_utils.h" #include "mlx/backend/cuda/quantized/quantized.h" @@ -108,13 +108,7 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { const array& w = inputs[1]; const array& scales = inputs[2]; - - bool non_batched = w.ndim() == 2; - int K = x.shape(-1); - int M = non_batched ? x.size() / K : x.shape(-2); - int N = out.shape(-1); - - fp_qmv(xhat, w, scales, out, bits_, group_size_, M, N, K, encoder, s); + fp_qmv(xhat, w, scales, out, bits_, group_size_, encoder, s); return; } diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index e9f90e06..41279547 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/mlx/backend/cuda/quantized/quantized.cpp @@ -3,8 +3,8 @@ #include "mlx/backend/cuda/quantized/quantized.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/quantized/qmm/qmm.h" -#include "mlx/backend/cuda/quantized/qmv.h" #include "mlx/backend/cuda/quantized/quantized_utils.h" +#include "mlx/dtype_utils.h" #include "mlx/fast_primitives.h" #include "mlx/primitives.h" @@ -17,8 +17,6 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& encoder = cu::get_command_encoder(s); - out.set_data(cu::malloc_async(out.nbytes(), encoder)); - const array& x = inputs[0]; const array& w = inputs[1]; const array& scales = inputs[2]; @@ -27,25 +25,74 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { biases = inputs[3]; } - bool non_batched = w.ndim() == 2; - int K = x.shape(-1); - int N = out.shape(-1); - int vec_batch = non_batched ? x.size() / K : x.shape(-2); - - if (transpose_ && vec_batch <= 8 && mode_ != QuantizationMode::Affine) { - assert(!biases); - fp_qmv(x, w, scales, out, bits_, group_size_, vec_batch, N, K, encoder, s); - return; - } - - if (transpose_ && mode_ == QuantizationMode::Affine && - encoder.device().compute_capability_major() == 9) { - assert(biases); + auto call_qmm_sm90 = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); qmm_sm90(x, w, scales, *biases, out, bits_, group_size_, encoder, s); + }; + auto call_fp_qmv = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + fp_qmv(x, w, scales, out, bits_, group_size_, encoder, s); + }; + auto call_qmv = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + qmv(x, w, scales, biases, out, bits_, group_size_, mode_, encoder); + }; + + auto supports = [&](auto&& f) { + return f( + x, + w, + scales, + biases, + out, + transpose_, + bits_, + group_size_, + mode_, + encoder.device()); + }; + bool can_use_qmm_sm90 = supports(supports_qmm_sm90); + bool can_use_fp_qmv = supports(supports_fp_qmv); + bool can_use_qmv = supports(supports_qmv); + + int M = out.shape(-2); + int N = out.shape(-1); + int K = x.shape(-1); + int B = out.size() / (M * N); + bool prefer_qmv = M == 1 && B == 1 && N <= 16384 && K <= 16384; + + if (can_use_qmm_sm90) { + if (prefer_qmv) { + if (can_use_fp_qmv) { + call_fp_qmv(); + return; + } + if (can_use_qmv) { + call_qmv(); + return; + } + } + call_qmm_sm90(); return; } - throw std::runtime_error("QMM NYI"); + if (can_use_fp_qmv) { + call_fp_qmv(); + return; + } + if (can_use_qmv) { + call_qmv(); + return; + } + + throw std::runtime_error( + fmt::format( + "[quantized_matmul] No implementation for " + "activation: {}, bits: {}, group size: {}, mode: \"{}\".", + dtype_to_string(x.dtype()), + bits_, + group_size_, + quantization_mode_to_string(mode_))); } void fast::Quantize::eval_gpu(