L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[CUDA] Quantized GEMV (#3180)

Browse Source
This commit is contained in:
Cheng
2026-03-04 08:59:31 +09:00
committed by GitHub
parent 9eef9f1774
commit 3c565437a5
10 changed files with 575 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/cuda/CMakeLists.txt
-1
View File
@@ -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
mlx/backend/cuda/quantized/qmm/CMakeLists.txt
+2
View File
@@ -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
mlx/backend/cuda/quantized/qmv.cu → mlx/backend/cuda/quantized/qmm/fp_qmv.cu
+61 -55
View File
@@ -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 <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
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 <typename T>
__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 <typename F>
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<MLX_GET_TYPE(type_tag)>;
if constexpr (!std::is_same_v<T, double>) {
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<uint32_t>(mat);
const T* vec_ptr = gpu_ptr<T>(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<T, rows_per_block, n.value, 4, 32, true>;
if (bits == 8) {
kernel = fp_qmv_single<T, rows_per_block, n.value, 8, 32, true>;
} else if (group_size == 16) {
kernel = fp_qmv_single<T, rows_per_block, n.value, 4, 16, false>;
}
encoder.add_kernel_node(
kernel,
{static_cast<uint32_t>(M), blocks_y},
block_dims,
mat_ptr,
gpu_ptr<uint8_t>(scales),
vec_ptr,
gpu_ptr<T>(out),
N,
K);
} else {
auto kernel =
fp_qmv_batched<T, rows_per_block, n.value, 4, 32, true>;
if (bits == 8) {
kernel = fp_qmv_batched<T, rows_per_block, n.value, 8, 32, true>;
} else if (group_size == 16) {
kernel = fp_qmv_batched<T, rows_per_block, n.value, 4, 16, false>;
}
encoder.add_kernel_node(
kernel,
{static_cast<uint32_t>(M), blocks_y, B},
block_dims,
mat_ptr,
gpu_ptr<uint8_t>(scales),
vec_ptr,
gpu_ptr<T>(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<T, rows_per_block, n.value, 4, 32, true>;
if (bits == 8) {
kernel = cu::fp_qmv_single<T, rows_per_block, n.value, 8, 32, true>;
} else if (group_size == 16) {
kernel =
cu::fp_qmv_single<T, rows_per_block, n.value, 4, 16, false>;
}
});
encoder.add_kernel_node(
kernel,
{uint32_t(x.size() / K), blocks_y},
block_dims,
mat_ptr,
gpu_ptr<uint8_t>(scales),
vec_ptr,
gpu_ptr<T>(out),
N,
K);
} else {
auto kernel =
cu::fp_qmv_batched<T, rows_per_block, n.value, 4, 32, true>;
if (bits == 8) {
kernel =
cu::fp_qmv_batched<T, rows_per_block, n.value, 8, 32, true>;
} else if (group_size == 16) {
kernel =
cu::fp_qmv_batched<T, rows_per_block, n.value, 4, 16, false>;
}
encoder.add_kernel_node(
kernel,
{M, blocks_y, B},
block_dims,
mat_ptr,
gpu_ptr<uint8_t>(scales),
vec_ptr,
gpu_ptr<T>(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
mlx/backend/cuda/quantized/qmm/qmm.cpp
+107
View File
@@ -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<array>& 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<array>& 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<array>& 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
mlx/backend/cuda/quantized/qmm/qmm.h
+58
View File
@@ -3,11 +3,24 @@
#pragma once
#include "mlx/backend/cuda/device.h"
#include "mlx/primitives.h"
#include <optional>
namespace mlx::core {
bool supports_qmm_sm90(
const array& x,
const array& w,
const array& scales,
const std::optional<array>& 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<array>& 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<array>& 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<array>& biases,
array& out,
int bits,
int group_size,
QuantizationMode mode,
cu::CommandEncoder& encoder);
} // namespace mlx::core
mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh
-19
View File
@@ -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);
mlx/backend/cuda/quantized/qmm/qmv.cu
+280
View File
@@ -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 <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <cute/numeric/numeric_types.hpp>
#include <cutlass/numeric_conversion.h>
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 <int N, typename T, typename Q>
__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<const cutlass::AlignedArray<T, N>*>(x));
auto w_vec = *(reinterpret_cast<const cutlass::AlignedArray<Q, N>*>(w));
// Output is assumed to be registers.
auto* out_vec = reinterpret_cast<cutlass::Array<float, N>*>(out);
// Dequantize w.
cutlass::NumericArrayConverter<T, Q, N> converter_tq;
cutlass::Array<T, N> w_dq = converter_tq(w_vec);
w_dq = w_dq * scale + bias;
// Promote x/w to float.
static_assert(!cuda::std::is_same_v<T, float>);
cutlass::NumericArrayConverter<float, T, N> converter_ft;
cutlass::Array<float, N> x_f = converter_ft(x_vec);
cutlass::Array<float, N> 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 <int N, typename Q>
__device__ __forceinline__ void
dequant_fma(const float* x, const Q* w, float scale, float bias, float* out) {
auto x_vec = *(reinterpret_cast<const cutlass::AlignedArray<float, N>*>(x));
auto w_vec = *(reinterpret_cast<const cutlass::AlignedArray<Q, N>*>(w));
auto* out_vec = reinterpret_cast<cutlass::Array<float, N>*>(out);
cutlass::NumericArrayConverter<float, Q, N> converter;
cutlass::Array<float, N> 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<WARP_SIZE>(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<Q>);
// 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<int64_t>(row) * k / w_step;
scales += static_cast<int64_t>(row) * groups_per_row;
if constexpr (has_bias) {
biases += static_cast<int64_t>(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<elems_per_thread>(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<float>{});
// Write result for current warp, which maps to rows 1-to-1.
if (warp.thread_rank() == 0) {
out[row] = static_cast<T>(sum);
}
}
template <int group_size, bool has_bias, typename T, typename Q, typename F>
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<void*>(kernel), num_blocks, block_dims, args);
});
}
} // namespace cu
template <typename F>
inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
if (dtype == float32) {
f.template operator()<float>();
} else if (dtype == float16) {
f.template operator()<cutlass::half_t>();
} else if (dtype == bfloat16) {
f.template operator()<cutlass::bfloat16_t>();
} else {
throw std::invalid_argument(
fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype)));
}
}
template <typename F>
inline void
dispatch_quant_types(int bits, QuantizationMode mode, const char* tag, F&& f) {
if (mode == QuantizationMode::Mxfp4) {
f.template operator()<cutlass::float_e2m1_t>();
} else if (mode == QuantizationMode::Mxfp8) {
f.template operator()<cutlass::float_e4m3_t>();
} else if (mode == QuantizationMode::Nvfp4) {
f.template operator()<cutlass::float_e2m1_t>();
} else {
if (bits == 2) {
f.template operator()<cutlass::uint2b_t>();
} else if (bits == 4) {
f.template operator()<cutlass::uint4b_t>();
} else if (bits == 8) {
f.template operator()<uint8_t>();
} else {
throw std::invalid_argument(
fmt::format("{} {}-bit quantization is not supported.", tag, bits));
}
}
}
template <typename F>
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<array>& 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, [&]<typename T>() {
dispatch_bool(biases.has_value(), [&](auto has_bias) {
dispatch_quant_types(bits, mode, tag, [&]<typename Q>() {
dispatch_groups(group_size, tag, [&]<int group_size>() {
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<group_size, has_bias.value>(
gpu_ptr<T>(x),
gpu_ptr<Q>(w),
gpu_ptr<T>(scales),
biases ? gpu_ptr<T>(*biases) : nullptr,
gpu_ptr<T>(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
mlx/backend/cuda/quantized/qmv.h
-22
View File
@@ -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
mlx/backend/cuda/quantized/qqmm.cpp
+2 -8
View File
@@ -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<array>& 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;
}
mlx/backend/cuda/quantized/quantized.cpp
+65 -18
View File
@@ -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<array>& 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<array>& 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(