| #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) |
|
|
