Jelajahi Sumber

fix aqlm compilation

AlpinDale 7 bulan lalu
induk
melakukan
a6a627d745
2 mengubah file dengan 5 tambahan dan 0 penghapusan
  1. 1 0
      kernels/quantization/quant_ops.cpp
  2. 4 0
      kernels/quantization/quant_ops.h

+ 1 - 0
kernels/quantization/quant_ops.cpp

@@ -8,6 +8,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 #ifndef USE_ROCM
   // AQLM
   quant_ops.def("aqlm_gemm", &aqlm_gemm, "Quantized GEMM for AQLM");
+  quant_ops.def("aqlm_dequant", &aqlm_dequant, "Dequantization for AQLM");
   // AWQ
   quant_ops.def("awq_gemm", &awq_gemm, "Quantized GEMM for AWQ");
   quant_ops.def("awq_dequantize", &awq_dequantize, "Dequantization for AWQ");

+ 4 - 0
kernels/quantization/quant_ops.h

@@ -13,6 +13,10 @@ torch::Tensor aqlm_gemm(
   const std::optional<torch::Tensor>& bias
 );
 
+torch::Tensor aqlm_dequant(const torch::Tensor& codes,
+                           const torch::Tensor& codebooks,
+                           const torch::Tensor& codebook_partition_sizes);
+
 // AWQ
 torch::Tensor awq_gemm(
   torch::Tensor _in_feats,