diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index b21310bc..3c3643a3 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -14,6 +14,19 @@ namespace mlx::core { namespace { +array ensure_row_contiguous( + const array& arr, + cpu::CommandEncoder& encoder, + Stream s) { + if (arr.flags().row_contiguous) { + return arr; + } else { + auto arr_cpy = contiguous_copy_cpu(arr, s); + encoder.add_temporary(arr_cpy); + return arr_cpy; + } +}; + const static float FP4_LUT[16] = { +0.0f, +0.5f, @@ -922,20 +935,9 @@ void QuantizedMatmul::eval_cpu(const std::vector& inputs, array& out) { auto& scales_pre = inputs[2]; auto& encoder = cpu::get_command_encoder(stream()); - auto ensure_row_contiguous = [s = stream(), &encoder](const array& arr) { - if (arr.flags().row_contiguous) { - return arr; - } else { - auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {}); - copy_cpu(arr, arr_cpy, CopyType::General, s); - encoder.add_temporary(arr_cpy); - return arr_cpy; - } - }; - - auto x = ensure_row_contiguous(x_pre); - auto w = ensure_row_contiguous(w_pre); - auto scales = ensure_row_contiguous(scales_pre); + auto x = ensure_row_contiguous(x_pre, encoder, stream()); + auto w = ensure_row_contiguous(w_pre, encoder, stream()); + auto scales = ensure_row_contiguous(scales_pre, encoder, stream()); out.set_data(allocator::malloc(out.nbytes())); @@ -944,7 +946,7 @@ void QuantizedMatmul::eval_cpu(const std::vector& inputs, array& out) { encoder.set_input_array(scales); encoder.set_output_array(out); if (mode_ == QuantizationMode::Affine) { - auto biases = ensure_row_contiguous(inputs[3]); + auto biases = ensure_row_contiguous(inputs[3], encoder, stream()); encoder.set_input_array(biases); encoder.dispatch([out = array::unsafe_weak_copy(out), x = array::unsafe_weak_copy(x), @@ -1052,6 +1054,105 @@ void GatherQMM::eval_cpu(const std::vector& inputs, array& out) { } } +uint8_t to_fp8_e8m0(float x) { + if (!std::isfinite(x)) { + return 0xFF; + } + if (x < 0.0f) { + return 0x00; + } + float le = std::log2(x); + int n = int(std::round(le)); + + n = n < -127 ? -127 : n; + n = n > 127 ? 127 : n; + return static_cast(n + 127); +} + +uint8_t to_fp4_e2m1(float x) { + if (std::isnan(x)) { + return 0x7; + } + + const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0; + x = std::abs(x); + + uint8_t bits; + if (x > 5.0f) { + bits = 0x7; + } else if (x >= 3.5f) { + bits = 0x6; + } else if (x > 2.5f) { + bits = 0x5; + } else if (x >= 1.75f) { + bits = 0x4; + } else if (x > 1.25f) { + bits = 0x3; + } else if (x >= 0.75f) { + bits = 0x2; + } else if (x > 0.25f) { + bits = 0x1; + } else { + bits = 0x0; + } + return bits | sign_bit; +} + +template +void fp_quantize_dequantize( + const array& w_arr, + array& out_arr, + int bits, + int group_size, + size_t w_size) { + auto w = w_arr.data(); + auto out = out_arr.data(); + + size_t n_groups = w_size / group_size; + + for (size_t i = 0; i < n_groups; ++i) { + size_t idx = i * group_size; + float scale = -std::numeric_limits::infinity(); + for (int j = 0; j < group_size; ++j) { + scale = std::max(scale, std::abs(w[idx + j])); + } + scale /= bits == 4 ? 6.0f : 448.0f; + if (group_size == 16) { + scale = dequantize_scale(detail::ToFP8()(scale)); + } else { + scale = dequantize_scale(to_fp8_e8m0(scale)); + } + + for (int j = 0; j < group_size; ++j) { + float w_el = scale == 0 ? 0.0f : w[idx + j] / scale; + float output; + if (bits == 8) { + output = detail::FromFP8()(detail::ToFP8()(w_el)); + } else { + output = FP4_LUT[to_fp4_e2m1(w_el)]; + } + out[idx + j] = static_cast(scale * output); + } + } +} + +void dispatch_quantize_dequantize( + const array& w, + array& out, + int bits, + int group_size) { + if (w.dtype() == float16) { + fp_quantize_dequantize(w, out, bits, group_size, w.size()); + } else if (w.dtype() == bfloat16) { + fp_quantize_dequantize(w, out, bits, group_size, w.size()); + } else if (w.dtype() == float32) { + fp_quantize_dequantize(w, out, bits, group_size, w.size()); + } else { + throw std::runtime_error( + "[quantize_dequantize] Only supports floating point inputs"); + } +} + template void quantize( const T* w, @@ -1136,15 +1237,8 @@ void dispatch_quantize( void fast::Quantize::eval_cpu( const std::vector& inputs, std::vector& outputs) { - auto ensure_row_contiguous = [s = stream()](const array& arr) { - if (arr.flags().row_contiguous) { - return std::make_pair(arr, false); - } else { - return std::make_pair(contiguous_copy_cpu(arr, s), true); - } - }; - - auto [w, copied] = ensure_row_contiguous(inputs[0]); + auto& encoder = cpu::get_command_encoder(stream()); + auto w = ensure_row_contiguous(inputs[0], encoder, stream()); auto& out = outputs[0]; out.set_data(allocator::malloc(out.nbytes())); @@ -1152,10 +1246,6 @@ void fast::Quantize::eval_cpu( auto& biases = outputs[2]; scales.set_data(allocator::malloc(scales.nbytes())); biases.set_data(allocator::malloc(biases.nbytes())); - auto& encoder = cpu::get_command_encoder(stream()); - if (copied) { - encoder.add_temporary(w); - } encoder.set_input_array(w); encoder.set_input_array(scales); encoder.set_input_array(biases); @@ -1238,6 +1328,43 @@ void fast::ConvertFP8::eval_cpu( } void QQMatmul::eval_cpu(const std::vector& inputs, array& out) { - throw std::runtime_error("QQMatmul not implemented on CPU."); + auto& encoder = cpu::get_command_encoder(stream()); + + bool w_quantized = (inputs[1].dtype() == uint32); + if (w_quantized && inputs[0].shape(-2) == 1) { + bool donate_x = inputs[0].is_donatable(); + auto x = ensure_row_contiguous(inputs[0], encoder, stream()); + auto w = ensure_row_contiguous(inputs[1], encoder, stream()); + auto scales = ensure_row_contiguous(inputs[2], encoder, stream()); + + out.set_data(allocator::malloc(out.nbytes())); + + // If x is a copy it should be donatable + donate_x |= x.is_donatable(); + auto xhat = donate_x + ? x + : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype()); + if (!donate_x) { + encoder.add_temporary(xhat); + } + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + encoder.set_output_array(out); + encoder.dispatch([out = array::unsafe_weak_copy(out), + x = array::unsafe_weak_copy(x), + xhat = array::unsafe_weak_copy(xhat), + w = array::unsafe_weak_copy(w), + scales = array::unsafe_weak_copy(scales), + group_size_ = group_size_, + bits_ = bits_]() mutable { + dispatch_quantize_dequantize(x, xhat, bits_, group_size_); + fp_qmm_dispatch(out, xhat, w, scales, group_size_, bits_, true); + }); + return; + } else { + throw std::runtime_error("[QQMatmul] NYI for the general case"); + } } + } // namespace mlx::core diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 2d00fd4c..601fbf74 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -56,7 +56,10 @@ 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 ${CMAKE_CURRENT_SOURCE_DIR}/quantized/convert_fp8.cu ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp) @@ -66,12 +69,12 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary) # fp4 is not available on < 12.8 if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.8.0) target_include_directories(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/) + target_sources(mlx + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/no_qqmm_impl.cpp) else() target_sources( - mlx - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu) + mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_impl.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp) endif() if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0) diff --git a/mlx/backend/cuda/primitives.cpp b/mlx/backend/cuda/primitives.cpp index 56e5969e..0caac8de 100644 --- a/mlx/backend/cuda/primitives.cpp +++ b/mlx/backend/cuda/primitives.cpp @@ -24,20 +24,12 @@ namespace mlx::core { throw std::runtime_error(#func " has no CUDA implementation."); \ } -#if CUDART_VERSION < 12080 -void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { - throw std::runtime_error( - "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher."); -} -#endif - NO_GPU(BlockMaskedMM) NO_GPU(FFT) NO_GPU(GatherQMM) NO_GPU(Hadamard) NO_GPU_MULTI(LUF) NO_GPU_MULTI(QRF) -NO_GPU(QuantizedMatmul) NO_GPU(SegmentedMM) NO_GPU_MULTI(SVD) NO_GPU(Inverse) diff --git a/mlx/backend/cuda/quantized/cuda_fp4.h b/mlx/backend/cuda/quantized/cuda_fp4.h index 10df4579..4500be6b 100644 --- a/mlx/backend/cuda/quantized/cuda_fp4.h +++ b/mlx/backend/cuda/quantized/cuda_fp4.h @@ -81,3 +81,20 @@ struct __nv_fp4_e2m1 { } uint8_t __x{0}; }; + +struct __nv_fp4x4_e2m1 { + __device__ operator float4() { + float4 out; + auto bits = __high & 0xf; + out.x = float(*(__nv_fp4_e2m1*)(&bits)); + bits = (__high >> 4) & 0xf; + out.y = float(*(__nv_fp4_e2m1*)(&bits)); + bits = (__low) & 0xf; + out.z = float(*(__nv_fp4_e2m1*)(&bits)); + bits = (__low >> 4) & 0xf; + out.w = float(*(__nv_fp4_e2m1*)(&bits)); + return out; + } + uint8_t __high{0}; + uint8_t __low{0}; +}; diff --git a/mlx/backend/cuda/quantized/fp_quantize.cu b/mlx/backend/cuda/quantized/fp_quantize.cu index 735ae5f2..3b4e96ef 100644 --- a/mlx/backend/cuda/quantized/fp_quantize.cu +++ b/mlx/backend/cuda/quantized/fp_quantize.cu @@ -11,8 +11,6 @@ #include #include -#include -#include namespace mlx::core { namespace cu { @@ -30,6 +28,69 @@ struct Dequantize { namespace cg = cooperative_groups; +template +__global__ void fp_quantize_dequantize(T* w, T* out, size_t size) { + using Tx2 = Vector2_t; + using Tx4 = Vector4_t; + uint32_t rbits = 0; // reserved bits for future use + auto block_size = cg::this_thread_block().dim_threads(); + auto block_idx = cg::this_thread_block().group_index(); + auto idx_in_block = cg::this_thread_block().thread_index(); + auto tidx = block_idx.x * block_size.x + idx_in_block.x; + auto tidy = block_idx.y * block_size.y + idx_in_block.y; + auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; + + size_t thread_idx = tidx + grid_dim_x * size_t(tidy); + size_t base_idx = thread_idx * group_size; + + if (base_idx >= size) { + return; + } + + auto w_tile = load_vector(w, thread_idx); + float scale = 0.0f; + + Tx2 amax_2x = Tx2{0.0f, 0.0f}; + +#pragma unroll + for (int i = 0; i < group_size; i += 2) { + auto pair = Tx2{w_tile[i], w_tile[i + 1]}; + abs_max_x2(amax_2x, amax_2x, pair); + } + + scale = static_cast( + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y)))); + + scale /= bits == 4 ? 6.0f : 448.0f; + // Convert to mx scale or nv scale + using ScaleType = + std::conditional_t; + auto s = ScaleType(scale); + scale = float(s); + AlignedVector w_hat; + +#pragma unroll + for (int i = 0; i < group_size / 4; i++) { + Tx4 w_Tx4 = *reinterpret_cast(&w_tile[i * 4]); + float4 dq; + if constexpr (bits == 8) { + uint32_t quantized_val = + scale_cvt_Tx4_to_fp8x4(w_Tx4, 1.0f / scale, rbits); + dq = dequant_fp8(quantized_val); + } else { + uint16_t quantized_val = + scale_cvt_Tx4_to_fp4x4(w_Tx4, 1.0f / scale, rbits); + dq = dequant_fp4(quantized_val); + } + w_hat[i * 4] = static_cast(dq.x * scale); + w_hat[i * 4 + 1] = static_cast(dq.y * scale); + w_hat[i * 4 + 2] = static_cast(dq.z * scale); + w_hat[i * 4 + 3] = static_cast(dq.w * scale); + } + store_vector(out, thread_idx, w_hat); +} + template __global__ void fp_quantize_rowwise(T* w, uint8_t* out, uint8_t* scales, size_t size) { @@ -280,6 +341,40 @@ get_columnwise_quantize_launch_args(size_t size, int group_size, int M, int K) { } // namespace cu +void fp_quantize_dequantize( + const array& w, + array& what, + int group_size, + int bits, + cu::CommandEncoder& enc, + const Stream& s) { + enc.set_input_array(w); + enc.set_output_array(what); + dispatch_float_types(w.dtype(), "fp_quantize_dequantize", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + auto kernel = cu::fp_quantize_dequantize; + if (bits == 8) { + kernel = cu::fp_quantize_dequantize; + } else if (group_size == 16) { + kernel = cu::fp_quantize_dequantize; + } + bool large = w.size() > UINT_MAX; + auto [num_blocks, block_dims] = + get_launch_args(w.size(), w.shape(), w.strides(), large, group_size); + + enc.add_kernel_node( + kernel, + num_blocks, + block_dims, + 0, + gpu_ptr(w), + gpu_ptr(what), + w.size()); + } + }); +} + void fp_quantize( const array& w, array& wq, diff --git a/mlx/backend/cuda/quantized/no_qqmm_impl.cpp b/mlx/backend/cuda/quantized/no_qqmm_impl.cpp new file mode 100644 index 00000000..5552e240 --- /dev/null +++ b/mlx/backend/cuda/quantized/no_qqmm_impl.cpp @@ -0,0 +1,26 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/quantized/qqmm_impl.h" + +namespace mlx::core { +void qqmm_impl( + cu::CommandEncoder&, + int, + int, + int, + bool, + int64_t, + bool, + int64_t, + array&, + const array&, + const array&, + const array&, + const array&, + Dtype, + QuantizationMode, + float) { + throw std::runtime_error( + "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher."); +} +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qmv.cu b/mlx/backend/cuda/quantized/qmv.cu new file mode 100644 index 00000000..71d9e7ad --- /dev/null +++ b/mlx/backend/cuda/quantized/qmv.cu @@ -0,0 +1,302 @@ +// Copyright © 2025 Apple Inc. + +#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/quantized_utils.cuh" +#include "mlx/dtype_utils.h" + +#include +#include + +namespace mlx::core::cu { + +namespace cg = cooperative_groups; + +static constexpr int rows_per_block = 8; + +template +__device__ void adjust_matrix_offsets( + const T*& x, + const uint32_t*& w, + const uint8_t*& scales, + T*& y, + int output_stride, + const int& x_batch_ndims, + const Shape x_shape, + const Strides x_strides, + const int& w_batch_ndims, + const Shape w_shape, + const Strides w_strides, + const Strides s_strides) { + uint32_t idx = cg::this_grid().block_index().z; + if (x_batch_ndims == 1) { + x += idx * x_strides[0]; + } else { + x += elem_to_loc(idx, x_shape.data(), x_strides.data(), x_batch_ndims); + } + if (w_batch_ndims == 1) { + w += idx * w_strides[0]; + scales += idx * s_strides[0]; + } else { + auto [w_idx, s_idx] = elem_to_loc( + idx, w_shape.data(), w_strides.data(), s_strides.data(), w_batch_ndims); + w += w_idx; + scales += s_idx; + } + y += idx * output_stride; +} + +template < + typename T, + int rows_per_block, + int n_per_thread, + int bits, + int group_size, + bool use_mx_scale> +__device__ void fp_qmv_impl( + const uint32_t* mat, + const uint8_t* scales_, + const T* vec, + T* out, + int rows, + int cols) { + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + constexpr int vals_per_item = bits == 8 ? 4 : 8; + constexpr int nv_per_thread = vals_per_item * n_per_thread; + auto g_idx = block.group_index(); + auto t_idx = block.thread_index(); + int row = g_idx.y * rows_per_block + t_idx.y; + + vec += g_idx.x * cols; + out += g_idx.x * rows; + + using ScaleType = + std::conditional_t; + auto scales = (ScaleType*)(scales_); + auto packed_cols = cols / vals_per_item; + + if (row < rows) { + constexpr int scales_per_step = std::max(nv_per_thread / group_size, 1); + constexpr int scale_step = (WARP_SIZE * nv_per_thread) / group_size; + constexpr int n_per_step = n_per_thread / scales_per_step; + // Offset scales to correct row + scales += row * (cols / group_size) + + (warp.thread_rank() * nv_per_thread) / group_size; + float sum = 0.0f; + for (int col = n_per_thread * warp.thread_rank(); col < packed_cols; + col += (WARP_SIZE * n_per_thread)) { + auto local_vec = + unsafe_load_vector(vec + vals_per_item * col, 0); + auto local_mat = + unsafe_load_vector(mat + row * packed_cols + col, 0); +#pragma unroll + for (int i = 0; i < scales_per_step; ++i) { + float2 local_sum = {0.0f, 0.0f}; +#pragma unroll + for (int j = 0; j < n_per_step; ++j) { + int k = n_per_step * i + j; + if constexpr (bits == 8) { + auto v = dequant_fp8(local_mat[k]); + local_sum.x += + v.x * static_cast(local_vec[vals_per_item * k]); + local_sum.x += + v.y * static_cast(local_vec[vals_per_item * k + 1]); + local_sum.y += + v.z * static_cast(local_vec[vals_per_item * k + 2]); + local_sum.y += + v.w * static_cast(local_vec[vals_per_item * k + 3]); + } else { + auto v = dequant_fp4(local_mat[k]); + local_sum.x += + v.x * static_cast(local_vec[vals_per_item * k]); + local_sum.y += + v.y * static_cast(local_vec[vals_per_item * k + 1]); + local_sum.x += + v.z * static_cast(local_vec[vals_per_item * k + 2]); + local_sum.y += + v.w * static_cast(local_vec[vals_per_item * k + 3]); + + v = dequant_fp4(local_mat[k] >> 16); + local_sum.x += + v.x * static_cast(local_vec[vals_per_item * k + 4]); + local_sum.y += + v.y * static_cast(local_vec[vals_per_item * k + 5]); + local_sum.x += + v.z * static_cast(local_vec[vals_per_item * k + 6]); + local_sum.y += + v.w * static_cast(local_vec[vals_per_item * k + 7]); + } + } + sum += (local_sum.x + local_sum.y) * float(scales[i]); + } + scales += scale_step; + } + + sum = cg::reduce(warp, sum, cg::plus{}); + if (warp.thread_rank() == 0) { + out[row] = static_cast(sum); + } + } +} + +template < + typename T, + int rows_per_block, + int n_per_thread, + int bits, + int group_size, + bool use_mx_scale> +__global__ void fp_qmv_single( + const uint32_t* mat, + const uint8_t* scales, + const T* vec, + T* out, + int rows, + int cols) { + fp_qmv_impl( + mat, scales, vec, out, rows, cols); +} + +template < + typename T, + int rows_per_block, + int n_per_thread, + int bits, + int group_size, + bool use_mx_scale> +__global__ void fp_qmv_batched( + const uint32_t* mat, + const uint8_t* scales, + const T* vec, + T* out, + int rows, + int cols, + int vec_batch_ndims, + const __grid_constant__ Shape vec_shape, + const __grid_constant__ Strides vec_strides, + int mat_batch_ndims, + const __grid_constant__ Shape mat_shape, + const __grid_constant__ Strides mat_strides, + const __grid_constant__ Strides scales_strides) { + adjust_matrix_offsets( + vec, + mat, + scales, + out, + rows * vec_shape[vec_batch_ndims], + vec_batch_ndims, + vec_shape, + vec_strides, + mat_batch_ndims, + mat_shape, + mat_strides, + scales_strides); + fp_qmv_impl( + mat, scales, vec, out, rows, cols); +} + +template +void dispatch_1_2_4(int n, F&& f) { + switch (n) { + case 1: + f(std::integral_constant{}); + break; + case 2: + f(std::integral_constant{}); + break; + case 4: + f(std::integral_constant{}); + break; + } +} + +void fp_qmv( + const array& mat, + const array& scales, + const array& vec, + array& out, + int bits, + int group_size, + int M, + int N, + int K, + CommandEncoder& encoder) { + encoder.set_input_array(mat); + encoder.set_input_array(scales); + encoder.set_input_array(vec); + encoder.set_output_array(out); + dispatch_float_types(out.dtype(), "qmv", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + dim3 block_dims{WARP_SIZE, rows_per_block}; + uint B = out.size() / (M * N); + uint 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); + int n = 1; + if (K % 32 == 0 && cu::is_aligned<4>(mat_ptr) && + ((bits == 4 && cu::is_aligned<8>(vec_ptr)) || + cu::is_aligned<4>(vec_ptr))) { + n = 4; + } else if ( + cu::is_aligned<2>(mat_ptr) && + ((bits == 4 && cu::is_aligned<4>(vec_ptr)) || + cu::is_aligned<2>(vec_ptr))) { + n = 2; + } + dispatch_1_2_4(n, [&](auto n) { + dispatch_bool(B > 1, [&](auto batched) { + if (!batched()) { + 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, + 0, + 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, + 0, + 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 diff --git a/mlx/backend/cuda/quantized/qmv.h b/mlx/backend/cuda/quantized/qmv.h new file mode 100644 index 00000000..d4f04195 --- /dev/null +++ b/mlx/backend/cuda/quantized/qmv.h @@ -0,0 +1,21 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include "mlx/backend/cuda/device.h" + +namespace mlx::core::cu { + +void fp_qmv( + const array& w, + const array& scales, + const array& vec, + array& out, + int bits, + int group_size, + int M, + int N, + int K, + CommandEncoder& encoder); + +} // namespace mlx::core::cu diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp index 122859ae..93724b2f 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/mlx/backend/cuda/quantized/qqmm.cpp @@ -1,10 +1,11 @@ // Copyright © 2025 Apple Inc. #include "mlx/backend/cuda/device.h" -#include "mlx/backend/cuda/quantized/cublas_qqmm.h" +#include "mlx/backend/cuda/quantized/qmv.h" +#include "mlx/backend/cuda/quantized/qqmm_impl.h" #include "mlx/backend/cuda/quantized/qqmm_utils.h" #include "mlx/backend/cuda/quantized/quantized.h" -#include "mlx/backend/gpu/copy.h" +#include "mlx/backend/cuda/quantized/quantized_utils.h" #include "mlx/primitives.h" #include @@ -13,16 +14,6 @@ namespace mlx::core { namespace { -inline array -ensure_contiguous(const array& x, cu::CommandEncoder& enc, const Stream& s) { - if (x.flags().row_contiguous || x.flags().col_contiguous) { - return x; - } - array x_copy = contiguous_copy_gpu(x, s); - enc.add_temporary(x_copy); - return x_copy; -} - array pad_and_swizzle_scales( const array& scale, cu::CommandEncoder& encoder, @@ -47,64 +38,54 @@ array pad_and_swizzle_scales( return scale_tiled; } -void qqmm_impl( - cu::CommandEncoder& encoder, - int M, - int N, - int K, - bool a_transposed, - int64_t lda, - bool b_transposed, - int64_t ldb, - array& out, - const array& a, - const array& b, - const array& a_scale, - const array& b_scale, - Dtype out_dtype, - QuantizationMode mode, - float alpha = 1.0f) { - // Invoke CublasQQMM - std::string qmode = quantization_mode_to_string(mode); - - // Currently only supports non-batched QQMM operations - // that covers all use cases for training, we will just collapse (batch, - // seq_len) into (tokens) - CublasQQMM qqmm( - encoder.device(), - a_transposed, - M, - K, - lda, - b_transposed, - K, - N, - ldb, - 1, // batch_count - 0, // a_batch_stride - 0, // b_batch_stride - out_dtype, - qmode); - - qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha); -} } // namespace void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { + assert( + (inputs.size() == 3 && inputs[1].dtype() == uint32) || + (inputs.size() == 2)); nvtx3::scoped_range r("QQMatmul::eval_gpu"); + auto& s = stream(); auto& encoder = cu::get_command_encoder(s); auto& device = encoder.device(); + + bool w_quantized = (inputs[1].dtype() == uint32); + if (w_quantized && inputs[0].shape(-2) == 1) { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + + bool donate_x = inputs[0].is_donatable(); + array x = ensure_row_contiguous(inputs[0], encoder, s); + // If x is a copy it should be donatable + donate_x |= x.is_donatable(); + auto xhat = donate_x + ? x + : array(cu::malloc_async(x.nbytes(), encoder), x.shape(), x.dtype()); + if (!donate_x) { + encoder.add_temporary(xhat); + } + fp_quantize_dequantize(x, xhat, group_size_, bits_, encoder, s); + + // Make sure the last two dims of w and s are contiguous + array w = ensure_row_contiguous_matrix(inputs[1], encoder, s); + array scales = ensure_row_contiguous_matrix(inputs[2], encoder, s); + + 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(w, scales, xhat, out, bits_, group_size_, M, N, K, encoder); + return; + } + std::cout << "RUNNING FULL?" << std::endl; + auto cc = device.compute_capability_major() * 100 + device.compute_capability_minor() * 10; if (cc < 1000) { throw std::runtime_error( "[QQMatmul::eval_gpu] QQMM is only supported on GPUs with compute capability 10.0 or higher."); } - assert( - (inputs.size() == 3 && inputs[1].dtype() == uint32) || - (inputs.size() == 2)); - auto quantize = [&](const array& input, cu::CommandEncoder& encoder, const Stream& s) -> std::pair { @@ -136,9 +117,8 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { return {x_q, scales_x}; }; auto [x_q, scale_x_pre] = quantize(inputs[0], encoder, s); - auto [w_q, scale_w_pre] = (inputs[1].dtype() != uint32) - ? quantize(inputs[1], encoder, s) - : std::make_pair(inputs[1], inputs[2]); + auto [w_q, scale_w_pre] = !w_quantized ? quantize(inputs[1], encoder, s) + : std::make_pair(inputs[1], inputs[2]); out.set_data(cu::malloc_async(out.nbytes(), encoder)); diff --git a/mlx/backend/cuda/quantized/qqmm_impl.cpp b/mlx/backend/cuda/quantized/qqmm_impl.cpp new file mode 100644 index 00000000..dd9407dc --- /dev/null +++ b/mlx/backend/cuda/quantized/qqmm_impl.cpp @@ -0,0 +1,50 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/quantized/qqmm_impl.h" +#include "mlx/backend/cuda/quantized/cublas_qqmm.h" + +namespace mlx::core { + +void qqmm_impl( + cu::CommandEncoder& encoder, + int M, + int N, + int K, + bool a_transposed, + int64_t lda, + bool b_transposed, + int64_t ldb, + array& out, + const array& a, + const array& b, + const array& a_scale, + const array& b_scale, + Dtype out_dtype, + QuantizationMode mode, + float alpha) { + // Invoke CublasQQMM + std::string qmode = quantization_mode_to_string(mode); + + // Currently only supports non-batched QQMM operations + // that covers all use cases for training, we will just collapse (batch, + // seq_len) into (tokens) + CublasQQMM qqmm( + encoder.device(), + a_transposed, + M, + K, + lda, + b_transposed, + K, + N, + ldb, + 1, // batch_count + 0, // a_batch_stride + 0, // b_batch_stride + out_dtype, + qmode); + + qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha); +} + +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qqmm_impl.h b/mlx/backend/cuda/quantized/qqmm_impl.h new file mode 100644 index 00000000..7288e2fd --- /dev/null +++ b/mlx/backend/cuda/quantized/qqmm_impl.h @@ -0,0 +1,26 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include "mlx/backend/cuda/device.h" +#include "mlx/primitives.h" + +namespace mlx::core { +void qqmm_impl( + cu::CommandEncoder& encoder, + int M, + int N, + int K, + bool a_transposed, + int64_t lda, + bool b_transposed, + int64_t ldb, + array& out, + const array& a, + const array& b, + const array& a_scale, + const array& b_scale, + Dtype out_dtype, + QuantizationMode mode, + float alpha = 1.0f); + +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index 4699ab78..30662d55 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/mlx/backend/cuda/quantized/quantized.cpp @@ -2,60 +2,48 @@ #include "mlx/backend/cuda/quantized/quantized.h" #include "mlx/backend/cuda/device.h" -#include "mlx/backend/gpu/copy.h" +#include "mlx/backend/cuda/quantized/qmv.h" +#include "mlx/backend/cuda/quantized/quantized_utils.h" #include "mlx/fast_primitives.h" +#include "mlx/primitives.h" #include namespace mlx::core { -namespace { +void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("QuantizedMatmul::eval_gpu"); + auto& s = stream(); + auto& d = cu::device(s.device); + auto& enc = d.get_command_encoder(s); -inline array ensure_row_contiguous( - const array& x, - cu::CommandEncoder& enc, - const Stream& s) { - if (!x.flags().row_contiguous) { - array x_copy = contiguous_copy_gpu(x, s); - enc.add_temporary(x_copy); - return x_copy; - } else { - return x; + out.set_data(cu::malloc_async(out.nbytes(), enc)); + + // Make sure the last two dims of x and w, s, b are contiguous. This should + // be relaxed for x. + array x = ensure_row_contiguous_matrix(inputs[0], enc, s); + array w = ensure_row_contiguous_matrix(inputs[1], enc, s); + array scales = ensure_row_contiguous_matrix(inputs[2], enc, s); + std::optional biases = std::nullopt; + if (inputs.size() == 4) { + biases = ensure_row_contiguous_matrix(inputs[3], enc, s); + } + + bool non_batched = w.ndim() == 2 && x.flags().row_contiguous; + int K = x.shape(-1); + int M = non_batched ? x.size() / K : x.shape(-2); + int N = out.shape(-1); + + if (M > 8 || !transpose_ || mode_ == QuantizationMode::Affine) { + throw std::runtime_error("QMM NYI"); + } + + if (transpose_) { + fp_qmv(w, scales, x, out, bits_, group_size_, M, N, K, enc); + return; } } -inline array -ensure_contiguous(const array& x, cu::CommandEncoder& enc, const Stream& s) { - if (x.flags().row_contiguous || x.flags().col_contiguous) { - return x; - } - array x_copy = contiguous_copy_gpu(x, s); - enc.add_temporary(x_copy); - return x_copy; -} - -inline array ensure_row_contiguous_matrix( - const array& x, - cu::CommandEncoder& enc, - const Stream& s) { - if (x.ndim() < 2) { - if (x.strides()[0] == 1) { - return x; - } - } else { - auto stride_0 = x.strides()[x.ndim() - 2]; - auto stride_1 = x.strides()[x.ndim() - 1]; - if (stride_0 == x.shape(-1) && stride_1 == 1) { - return x; - } - } - array x_copy = contiguous_copy_gpu(x, s); - enc.add_temporary(x_copy); - return x_copy; -} - -} // namespace - void fast::Quantize::eval_gpu( const std::vector& inputs, std::vector& outputs) { diff --git a/mlx/backend/cuda/quantized/quantized.h b/mlx/backend/cuda/quantized/quantized.h index 4f1980a9..744a12f5 100644 --- a/mlx/backend/cuda/quantized/quantized.h +++ b/mlx/backend/cuda/quantized/quantized.h @@ -42,4 +42,12 @@ void fp_dequantize( cu::CommandEncoder& enc, const Stream& s); +void fp_quantize_dequantize( + const array& w, + array& what, + int group_size, + int bits, + cu::CommandEncoder& enc, + const Stream& s); + } // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/quantized_utils.cuh b/mlx/backend/cuda/quantized/quantized_utils.cuh index e589c970..93d83292 100644 --- a/mlx/backend/cuda/quantized/quantized_utils.cuh +++ b/mlx/backend/cuda/quantized/quantized_utils.cuh @@ -1,9 +1,22 @@ // Copyright © 2025 Apple Inc. +#include +#include + namespace mlx::core { namespace cu { +inline __device__ float4 dequant_fp8(uint32_t bits) { + auto out = *(__nv_fp8x4_e4m3*)(&bits); + return out.operator float4(); +} + +inline __device__ float4 dequant_fp4(uint16_t bits) { + auto out = *(__nv_fp4x4_e2m1*)(&bits); + return out.operator float4(); +} + template inline constexpr __device__ short get_pack_factor() { return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits); diff --git a/mlx/backend/cuda/quantized/quantized_utils.h b/mlx/backend/cuda/quantized/quantized_utils.h new file mode 100644 index 00000000..66be7686 --- /dev/null +++ b/mlx/backend/cuda/quantized/quantized_utils.h @@ -0,0 +1,50 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/gpu/copy.h" + +namespace mlx::core { +inline array ensure_row_contiguous( + const array& x, + cu::CommandEncoder& enc, + const Stream& s) { + if (!x.flags().row_contiguous) { + array x_copy = contiguous_copy_gpu(x, s); + enc.add_temporary(x_copy); + return x_copy; + } else { + return x; + } +} + +inline array ensure_row_contiguous_matrix( + const array& x, + cu::CommandEncoder& enc, + const Stream& s) { + if (x.ndim() < 2) { + if (x.strides()[0] == 1) { + return x; + } + } else { + auto stride_0 = x.strides()[x.ndim() - 2]; + auto stride_1 = x.strides()[x.ndim() - 1]; + if (stride_0 == x.shape(-1) && stride_1 == 1) { + return x; + } + } + array x_copy = contiguous_copy_gpu(x, s); + enc.add_temporary(x_copy); + return x_copy; +} + +inline array +ensure_contiguous(const array& x, cu::CommandEncoder& enc, const Stream& s) { + if (x.flags().row_contiguous || x.flags().col_contiguous) { + return x; + } + array x_copy = contiguous_copy_gpu(x, s); + enc.add_temporary(x_copy); + return x_copy; +} + +} // namespace mlx::core diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 2bd52fa7..0591532f 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -1766,7 +1766,6 @@ template uint8_t q_scale = s.bits; scale = float(s); - // Write out the scales and biases size_t gindex = index / group_size; if (index % group_size == 0) { scales[gindex] = q_scale; @@ -1786,7 +1785,7 @@ template template [[kernel]] void fp_dequantize( const device uint8_t* w [[buffer(0)]], - const device T* scales [[buffer(1)]], + const device uint8_t* scales [[buffer(1)]], device T* out [[buffer(3)]], uint2 index [[thread_position_in_grid]], uint2 grid_dim [[threads_per_grid]]) { @@ -1814,3 +1813,32 @@ template out[i] = static_cast(scale * Dequantize{}(d)); } } + +template +[[kernel]] void fp_quantize_dequantize( + const device T* w [[buffer(0)]], + device T* out [[buffer(1)]], + uint2 tidx [[thread_position_in_grid]], + uint2 grid_dim [[threads_per_grid]]) { + constexpr bool use_mx_scale = group_size == 32; + size_t index = tidx.x + grid_dim.x * size_t(tidx.y); + + float scale; + float w_thread = w[index]; + if (use_mx_scale) { + scale = simd_max(abs(w_thread)); + } else { + float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0); + float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0); + scale = tidx.x < 16 ? w_max_l : w_max_r; + } + scale /= bits == 4 ? 6.0f : 448.0f; + + using ScaleType = metal::conditional_t; + auto s = ScaleType(scale); + scale = float(s); + + uint8_t output = Quantize{}(scale == 0 ? 0.0f : w_thread / scale); + + out[index] = static_cast(scale * Dequantize{}(output)); +} diff --git a/mlx/backend/metal/kernels/fp_quantized.metal b/mlx/backend/metal/kernels/fp_quantized.metal index 51ed8ca5..01a2df02 100644 --- a/mlx/backend/metal/kernels/fp_quantized.metal +++ b/mlx/backend/metal/kernels/fp_quantized.metal @@ -114,6 +114,12 @@ instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false, mode, group_size, bits) #define instantiate_quantize_dequantize(type, mode, group_size, bits) \ + instantiate_kernel( \ + #mode "_quantize_dequantize_" #type "_gs_" #group_size "_b_" #bits, \ + fp_quantize_dequantize, \ + type, \ + group_size, \ + bits) \ instantiate_kernel( \ #mode "_quantize_" #type "_gs_" #group_size "_b_" #bits, \ fp_quantize, \ diff --git a/mlx/backend/metal/primitives.cpp b/mlx/backend/metal/primitives.cpp index f3dfbb75..5f6376c5 100644 --- a/mlx/backend/metal/primitives.cpp +++ b/mlx/backend/metal/primitives.cpp @@ -239,8 +239,4 @@ void LUF::eval_gpu( throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI."); } -void QQMatmul::eval_gpu(const std::vector& inputs, array& output) { - throw std::runtime_error("[QQMatmul::eval_gpu] Metal QQMatmul NYI."); -} - } // namespace mlx::core diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index badc8883..cb67c74f 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1267,6 +1267,30 @@ void gather_qmm_rhs( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +void dispatch_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + // It is a qmv with a small inner dimension so route to qmv_quad kernel + if ((K == 128 || K == 64) && is_power_of_2(bits)) { + qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); + return; + } + + // Run of the mill qmv + qmv(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); +} + void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); @@ -1310,16 +1334,10 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { return; } - // It is a qmv with a small inner dimension so route to qmv_quad kernel - if (transpose_ && (K == 128 || K == 64) && is_power_of_2(bits_)) { - qmv_quad( - x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); - return; - } - // Run of the mill qmv if (transpose_) { - qmv(x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + dispatch_qmv( + x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); return; } @@ -1443,6 +1461,99 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { mode); } +void quantize_dequantize( + const array& in, + array& out, + std::string mode, + int group_size, + int bits, + metal::Device& d, + const Stream& s) { + auto& compute_encoder = d.get_command_encoder(s.index); + + auto w = ensure_row_contiguous(in, d, s); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_output_array(out, 1); + auto type_string = get_type_string(in.dtype()); + std::string kname; + concatenate( + kname, + mode + "_quantize_dequantize_", + type_string, + "_gs_", + group_size, + "_b_", + bits); + auto kernel = get_quantized_kernel_wrapped( + d, kname, "quantize_dequantize", mode, type_string, group_size, bits); + + compute_encoder.set_compute_pipeline_state(kernel); + + constexpr int uint8_per_uint32 = 4; + constexpr int simd_size = 32; + int packs_per_int = (bits == 3 || bits == 5) ? 8 : bits == 6 ? 4 : 8 / bits; + int per_thread = std::max(group_size / simd_size, 1); + size_t nthreads = w.size() / per_thread; + + NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); + if (thread_group_size > nthreads) { + thread_group_size = nthreads; + } + auto group_dims = MTL::Size(thread_group_size, 1, 1); + bool use_2d = nthreads > UINT_MAX; + auto grid_shape = w.shape(); + grid_shape.back() /= per_thread; + MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides()) + : MTL::Size(nthreads, 1, 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { + auto& s = stream(); + auto& d = metal::device(s.device); + + auto mode = quantization_mode_to_string(mode_); + bool w_quantized = (inputs[1].dtype() == uint32); + if (w_quantized && inputs[0].shape(-2) == 1) { + out.set_data(allocator::malloc(out.nbytes())); + + bool donate_x = inputs[0].is_donatable(); + array x = ensure_row_contiguous(inputs[0], d, s); + // If x is a copy it should be donatable + donate_x |= x.is_donatable(); + auto xhat = donate_x + ? x + : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype()); + quantize_dequantize(x, xhat, mode, group_size_, bits_, d, s); + + // Make sure the last two dims of w and s are contiguous + array w = ensure_row_contiguous_matrix(inputs[1], d, s); + array scales = ensure_row_contiguous_matrix(inputs[2], d, s); + + 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); + dispatch_qmv( + xhat, + w, + scales, + std::nullopt, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } else { + throw std::runtime_error("[QQMatmul] NYI for the general case"); + } +} + void fast::Quantize::eval_gpu( const std::vector& inputs, std::vector& outputs) { diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 3159f49e..076de699 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -4345,10 +4345,6 @@ array qqmm( const std::string& mode /* = "nvfp4" */, StreamOrDevice s /* = {} */) { auto stream = to_stream(s); - if (stream.device != Device::gpu || !cu::is_available()) { - throw std::invalid_argument( - "[qqmm] Only supported on GPU with the CUDA backend."); - } auto qmode = string_to_quantization_mode(mode, "qqmm"); // cuBLAS block scaled matmul only supports nvfp4 and mxfp8 if (qmode != QuantizationMode::Nvfp4 && qmode != QuantizationMode::Mxfp8) { diff --git a/python/tests/cuda_skip.py b/python/tests/cuda_skip.py index 15c9487c..47ed4a61 100644 --- a/python/tests/cuda_skip.py +++ b/python/tests/cuda_skip.py @@ -47,7 +47,6 @@ cuda_skip = { "TestQuantized.test_qmm_shapes", "TestQuantized.test_qmm_vjp", "TestQuantized.test_qmv", - "TestQuantized.test_fp_qmv", "TestQuantized.test_fp_qvm", "TestQuantized.test_qvm", "TestQuantized.test_qvm_splitk", diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 9e581970..6e4e9eee 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -160,6 +160,38 @@ class TestQuantized(mlx_tests.MLXTestCase): w_hat = mx.dequantize(w_q, scales, mode="nvfp4") self.assertTrue(mx.all(w_hat == 0)) + def test_qqmv(self): + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + tests = product( + [256, 512, 67], # M + [64, 256], # N + ) + modes = ["nvfp4", "mxfp8"] + for M, N in tests: + for mode in modes: + with self.subTest(shape=(M, N), mode=mode): + x_shape = (1, N) + w_shape = (M, N) + + x = mx.random.normal(shape=x_shape, key=k1) + x_hat = mx.dequantize( + *mx.quantize(x, mode=mode), mode=mode, dtype=mx.float32 + ) + + w = mx.random.normal(shape=w_shape, key=k2) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode, dtype=mx.float32) + y_q = mx.qqmm( + x, + w_q, + scales, + mode=mode, + ) + y_hat = x_hat @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_qmm(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) @@ -338,6 +370,28 @@ class TestQuantized(mlx_tests.MLXTestCase): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + # Test multiple of 16 but not 32 + M = 128 + N = 48 + mode = "nvfp4" + with self.subTest(shape=(B, M, N), mode=mode): + x_shape = (1, N) + w_shape = (M, N) + x = mx.random.normal(shape=x_shape, key=k1) + w = mx.random.normal(shape=w_shape, key=k2) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + y_q = mx.quantized_matmul( + x, + w_q, + scales, + transpose=True, + mode=mode, + ) + y_hat = x @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_qvm(self): key = mx.random.key(0) k1, k2 = mx.random.split(key)