Add GPTQ support (#916)

2023-12-15 19:04:22 +08:00
parent c06170cc8e
commit 0fbfc4b81b
35 changed files with 1782 additions and 82 deletions
--- a/csrc/pybind.cpp
+++ b/csrc/pybind.cpp
@ -52,8 +52,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  // Quantization ops
  ops.def("awq_gemm", &awq_gemm, "Quantized GEMM for AWQ");
 #endif
-
-
+  ops.def("gptq_gemm", &gptq_gemm, "Quantized GEMM for GPTQ");
+  ops.def("gptq_shuffle", &gptq_shuffle, "Post processing for GPTQ");
  ops.def("squeezellm_gemm", &squeezellm_gemm, "Quantized GEMM for SqueezeLLM");

  // Cache ops