File size: 689 Bytes
ce8679e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | #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, 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) -> ()");
#if defined(CPU_KERNEL)
ops.impl("quantize_act", torch::kCPU, &quantize_act);
ops.impl("bitnet_gemv_fused", torch::kCPU, &bitnet_gemv_fused);
ops.impl("bitnet_gemm", torch::kCPU, &bitnet_gemm);
#endif
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|