From 6304c285d30ae4843229cf9a6939c227c2e60bb2 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 25 Feb 2026 09:10:18 +0900 Subject: [PATCH] [CUDA] FPxINT quantized matmul for Hopper (#3160) --- mlx/backend/cuda/CMakeLists.txt | 5 + mlx/backend/cuda/cutlass_utils.cuh | 2 +- mlx/backend/cuda/quantized/qmm.h | 24 +++ mlx/backend/cuda/quantized/qmm_sm90.cu | 244 +++++++++++++++++++++++ mlx/backend/cuda/quantized/qmv.cu | 16 +- mlx/backend/cuda/quantized/qmv.h | 5 +- mlx/backend/cuda/quantized/qqmm.cpp | 7 +- mlx/backend/cuda/quantized/quantized.cpp | 40 ++-- 8 files changed, 313 insertions(+), 30 deletions(-) create mode 100644 mlx/backend/cuda/quantized/qmm.h create mode 100644 mlx/backend/cuda/quantized/qmm_sm90.cu diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 013b24b2..7b4c47f8 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -56,6 +56,7 @@ 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/qmm_sm90.cu ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmv.cu ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp @@ -116,6 +117,10 @@ target_compile_options(mlx target_compile_options( mlx PRIVATE "$<$:--expt-relaxed-constexpr>") +# Required for generating optimized CUTLASS code. +target_compile_options( + mlx PRIVATE "$<$:-Xcompiler=-fno-strict-aliasing>") + # Suppress nvcc warnings on C++ headers. target_compile_options( mlx diff --git a/mlx/backend/cuda/cutlass_utils.cuh b/mlx/backend/cuda/cutlass_utils.cuh index c9bbf0d4..4770f01a 100644 --- a/mlx/backend/cuda/cutlass_utils.cuh +++ b/mlx/backend/cuda/cutlass_utils.cuh @@ -22,7 +22,7 @@ inline void check_cutlass_error(const char* name, cutlass::Status status) { } // The macro version that prints the command that failed. -#define CHECK_CUTLASS_ERROR(cmd) check_cutlass_error(#cmd, (cmd)) +#define CHECK_CUTLASS_ERROR(cmd) ::mlx::core::check_cutlass_error(#cmd, (cmd)) // Maps CPU types to CUTLASS types. template diff --git a/mlx/backend/cuda/quantized/qmm.h b/mlx/backend/cuda/quantized/qmm.h new file mode 100644 index 00000000..8c45140a --- /dev/null +++ b/mlx/backend/cuda/quantized/qmm.h @@ -0,0 +1,24 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include "mlx/backend/cuda/device.h" +#include "mlx/primitives.h" + +#include + +namespace mlx::core { + +void qmm_sm90( + 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, + Stream s); + +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qmm_sm90.cu b/mlx/backend/cuda/quantized/qmm_sm90.cu new file mode 100644 index 00000000..34b82734 --- /dev/null +++ b/mlx/backend/cuda/quantized/qmm_sm90.cu @@ -0,0 +1,244 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/cutlass_utils.cuh" +#include "mlx/backend/cuda/quantized/qmm.h" +#include "mlx/backend/cuda/quantized/quantized_utils.h" +#include "mlx/backend/gpu/copy.h" +#include "mlx/dtype_utils.h" + +#include +#include +#include +#include +#include +#include + +// We can't put kernel code in mlx::core due to name conflicts of "Shape". +namespace cutlass_gemm { + +template +void qmm_sm90( + const Element* A, + const Quant* B, + const Element* S, + const Element* Z, + Element* D, + int64_t m, + int64_t n, + int64_t k, + int64_t l, + GroupSize group_size, + F&& launch_kernel) { +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + using namespace cute; + + constexpr int kAlignmentA = 128 / sizeof_bits::value; + constexpr int kAlignmentB = 128 / sizeof_bits::value; + constexpr int kTileShapeK = + std::max(64, 128 * 8 / sizeof_bits::value); + static_assert(group_size % kTileShapeK == 0); + + using Arch = cutlass::arch::Sm90; + using Accumulator = float; + using TileShape = Shape<_128, _16, Int>; + using ClusterShape = Shape<_1, _1, _1>; + + using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + Arch, + cutlass::arch::OpClassTensorOp, + TileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + Accumulator, + Accumulator, + // ElementC: + void, + cutlass::layout::ColumnMajor, + kAlignmentA, + // ElementD: + Element, + cutlass::layout::ColumnMajor, + kAlignmentA, + cutlass::epilogue::TmaWarpSpecializedCooperative>::CollectiveOp; + + // Note that A/B are swapped and transposed to use TMA epilogue. + using Mainloop = typename cutlass::gemm::collective::CollectiveBuilder< + Arch, + cutlass::arch::OpClassTensorOp, + // ElementA: + cute::tuple, + cutlass::layout::RowMajor, + kAlignmentB, + // ElementB: + Element, + cutlass::layout::ColumnMajor, + kAlignmentA, + Accumulator, + TileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename Epilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel:: + GemmUniversal, Mainloop, Epilogue>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + auto dA = make_stride(k, Int<1>{}, m * k); + auto dB = make_stride(k, Int<1>{}, n * k); + auto dS = make_stride(Int<1>{}, n, n * k / group_size); + auto dD = make_stride(Int<1>{}, n, m * n); + + Gemm gemm; + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {int(n), int(m), int(k), int(l)}, + {B, dB, A, dA, S, dS, group_size, Z}, + {{1.f, 0.f}, D, dD, D, dD}}; + + CHECK_CUTLASS_ERROR(gemm.can_implement(args)); + CHECK_CUTLASS_ERROR(gemm.initialize(args, nullptr)); + + auto* kernel = &cutlass::device_kernel; + void* kernel_params[] = {const_cast(&gemm.params())}; + launch_kernel( + reinterpret_cast(kernel), + gemm.get_grid_shape(gemm.params()), + GemmKernel::get_block_shape(), + GemmKernel::SharedStorageSize, + kernel_params); +#else + throw std::runtime_error( + "[quantized_matmul] Hopper-only kernel is not available."); +#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) +} + +} // namespace cutlass_gemm + +namespace mlx::core { + +inline array transpose_last_2_dims( + const array& x, + cu::CommandEncoder& encoder, + const Stream& s) { + array transposed = swapaxes_in_eval(x, -1, -2); + array transposed_copy = contiguous_copy_gpu(transposed, s); + encoder.add_temporary(transposed_copy); + return transposed_copy; +} + +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, const char* tag, F&& f) { + 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 == 64) { + f(cute::Int<64>{}); + } else if (group_size == 128) { + f(cute::Int<128>{}); + } else { + throw std::invalid_argument( + fmt::format("{} Group size {} is not supported.", tag, group_size)); + } +} + +void qmm_sm90( + 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, + Stream s) { + if ((mode != QuantizationMode::Affine) || !biases_) { + throw std::runtime_error("qmm_sm90 NYI"); + } + + const char* tag = "[quantized_matmul]"; + int m = out.shape(-2); + 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 (!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)); + } + + // TODO: Support column-major x. + array x = ensure_row_contiguous(x_, encoder, s); + // FIXME: Copy happens for every call. + array scales = transpose_last_2_dims(scales_, encoder, s); + array biases = transpose_last_2_dims(*biases_, encoder, s); + + dispatch_element_types(out.dtype(), tag, [&]() { + dispatch_quant_types(bits, tag, [&]() { + dispatch_groups(group_size, tag, [&](auto group_size) { + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + encoder.set_input_array(biases); + encoder.set_output_array(out); + cutlass_gemm::qmm_sm90( + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(scales), + gpu_ptr(biases), + gpu_ptr(out), + m, + n, + k, + l, + group_size, + [&](auto* kernel, + dim3 num_blocks, + dim3 block_dims, + uint32_t smem_bytes, + void** args) { + encoder.add_kernel_node( + kernel, num_blocks, block_dims, smem_bytes, args); + }); + }); + }); + }); +} + +} // namespace mlx::core diff --git a/mlx/backend/cuda/quantized/qmv.cu b/mlx/backend/cuda/quantized/qmv.cu index 0cc00464..4a6c44c9 100644 --- a/mlx/backend/cuda/quantized/qmv.cu +++ b/mlx/backend/cuda/quantized/qmv.cu @@ -4,6 +4,7 @@ #include "mlx/backend/cuda/kernel_utils.cuh" #include "mlx/backend/cuda/quantized/qmv.h" #include "mlx/backend/cuda/quantized/quantized_utils.cuh" +#include "mlx/backend/cuda/quantized/quantized_utils.h" #include "mlx/dtype_utils.h" #include @@ -214,16 +215,23 @@ void dispatch_1_2_4(int n, F&& f) { } void fp_qmv( - const array& mat, - const array& scales, - const array& vec, + const array& x, + const array& w, + const array& scales_, array& out, int bits, int group_size, int M, int N, int K, - CommandEncoder& encoder) { + CommandEncoder& encoder, + Stream s) { + // 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); + array mat = ensure_row_contiguous_matrix(w, encoder, s); + array scales = ensure_row_contiguous_matrix(scales_, encoder, s); + encoder.set_input_array(mat); encoder.set_input_array(scales); encoder.set_input_array(vec); diff --git a/mlx/backend/cuda/quantized/qmv.h b/mlx/backend/cuda/quantized/qmv.h index d4f04195..9fc3719a 100644 --- a/mlx/backend/cuda/quantized/qmv.h +++ b/mlx/backend/cuda/quantized/qmv.h @@ -7,15 +7,16 @@ namespace mlx::core::cu { void fp_qmv( + const array& x, const array& w, const array& scales, - const array& vec, array& out, int bits, int group_size, int M, int N, int K, - CommandEncoder& encoder); + 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 665bdb45..7ff5f0a3 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/mlx/backend/cuda/quantized/qqmm.cpp @@ -106,16 +106,15 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { fp_quantize_dequantize( x, xhat, group_size_, bits_, global_scale, 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); + 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(w, scales, xhat, out, bits_, group_size_, M, N, K, encoder); + fp_qmv(xhat, w, scales, out, bits_, group_size_, M, N, K, encoder, s); return; } diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index 81ded942..d5e1212d 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/mlx/backend/cuda/quantized/quantized.cpp @@ -2,6 +2,7 @@ #include "mlx/backend/cuda/quantized/quantized.h" #include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/quantized/qmm.h" #include "mlx/backend/cuda/quantized/qmv.h" #include "mlx/backend/cuda/quantized/quantized_utils.h" #include "mlx/fast_primitives.h" @@ -14,34 +15,35 @@ namespace mlx::core { 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); + auto& encoder = cu::get_command_encoder(s); - out.set_data(cu::malloc_async(out.nbytes(), enc)); + out.set_data(cu::malloc_async(out.nbytes(), encoder)); - // 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); + const array& x = inputs[0]; + const array& w = inputs[1]; + const array& scales = inputs[2]; + std::optional biases; + if (inputs.size() > 3) { + biases = inputs[3]; } - bool non_batched = w.ndim() == 2 && x.flags().row_contiguous; + 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); + int vec_batch = non_batched ? x.size() / K : x.shape(-2); - 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); + 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_ && encoder.device().compute_capability_major() == 9) { + qmm_sm90(x, w, scales, biases, out, bits_, group_size_, mode_, encoder, s); + return; + } + + throw std::runtime_error("QMM NYI"); } void fast::Quantize::eval_gpu(