File size: 1,129 Bytes
3ca8249 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | #include <torch/library.h>
#include "registration.h"
#include "torch_binding.h"
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("quantize_act(Tensor! Aq, Tensor! scale, Tensor a) -> ()");
ops.def("bitnet_gemv_fused(Tensor! out, Tensor a_bf16, Tensor w_packed, Tensor scale_wt) -> ()");
ops.def("bitnet_gemm(Tensor! out, Tensor a_int8, Tensor w_packed, Tensor scale_act, Tensor scale_wt, Tensor? scratch) -> ()");
ops.def("ternarize_pack(Tensor! w_packed, Tensor! gamma, Tensor W) -> ()");
ops.def("dequant_w(Tensor! W_hat, Tensor w_packed, Tensor gamma) -> ()");
ops.def("dequant_x(Tensor! X_hat, Tensor x_q, Tensor scale) -> ()");
#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
ops.impl("quantize_act", torch::kCUDA, &quantize_act);
ops.impl("bitnet_gemv_fused", torch::kCUDA, &bitnet_gemv_fused);
ops.impl("bitnet_gemm", torch::kCUDA, &bitnet_gemm);
ops.impl("ternarize_pack", torch::kCUDA, &ternarize_pack);
ops.impl("dequant_w", torch::kCUDA, &dequant_w);
ops.impl("dequant_x", torch::kCUDA, &dequant_x);
#endif
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|