Awni Hannun
86 lines
2.7 KiB
C++
86 lines
2.7 KiB
C++
// Copyright © 2025 Apple Inc.
|
|
|
|
#include "mlx/backend/cuda/quantized/quantized.h"
|
|
#include "mlx/backend/cuda/device.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 <nvtx3/nvtx3.hpp>
|
|
|
|
namespace mlx::core {
|
|
|
|
void QuantizedMatmul::eval_gpu(const std::vector<array>& 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);
|
|
|
|
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<array> 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;
|
|
}
|
|
}
|
|
|
|
void fast::Quantize::eval_gpu(
|
|
const std::vector<array>& inputs,
|
|
std::vector<array>& outputs) {
|
|
nvtx3::scoped_range r("Quantize::eval_gpu");
|
|
auto& s = stream();
|
|
auto& d = cu::device(s.device);
|
|
auto& enc = d.get_command_encoder(s);
|
|
|
|
if (dequantize_) {
|
|
auto wq = ensure_row_contiguous(inputs[0], enc, s);
|
|
auto scales = ensure_row_contiguous(inputs[1], enc, s);
|
|
auto& w = outputs[0];
|
|
|
|
w.set_data(cu::malloc_async(w.nbytes(), enc));
|
|
|
|
if (mode_ == QuantizationMode::Affine) {
|
|
auto biases = ensure_row_contiguous(inputs[2], enc, s);
|
|
affine_dequantize(wq, scales, biases, w, group_size_, bits_, enc, s);
|
|
} else {
|
|
fp_dequantize(wq, scales, w, group_size_, bits_, enc, s);
|
|
}
|
|
} else {
|
|
auto w = ensure_contiguous(inputs[0], enc, s);
|
|
auto& wq = outputs[0];
|
|
auto& scales = outputs[1];
|
|
|
|
wq.set_data(cu::malloc_async(wq.nbytes(), enc));
|
|
scales.set_data(cu::malloc_async(scales.nbytes(), enc));
|
|
if (mode_ == QuantizationMode::Affine) {
|
|
auto& biases = outputs[2];
|
|
biases.set_data(cu::malloc_async(biases.nbytes(), enc));
|
|
affine_quantize(w, wq, scales, biases, group_size_, bits_, enc, s);
|
|
} else {
|
|
fp_quantize(w, wq, scales, group_size_, bits_, enc, s);
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace mlx::core
|