| #include <torch/library.h> |
|
|
| #include "registration.h" |
| #include "torch_binding.h" |
|
|
| TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| ops.def( |
| "qgemm_raw_simple(Tensor input, Tensor weight, Tensor scales, " |
| "Tensor table, Tensor table2, Tensor(a!) workspace, " |
| "int num_bits, int group_size, int template_id, int num_sms) -> Tensor"); |
| ops.def( |
| "qgemm_raw_simple_hadamard(Tensor input, Tensor weight, Tensor scales, " |
| "Tensor table, Tensor table2, Tensor(a!) workspace, " |
| "int num_bits, int group_size, int hadamard_size, " |
| "int template_id, int num_sms) -> Tensor"); |
| ops.def("hadamard_transform(Tensor(a!) input, bool inplace) -> Tensor"); |
| #if defined(CUDA_KERNEL) || defined(ROCM_KERNEL) |
| ops.impl("qgemm_raw_simple", c10::kCUDA, &qgemm_raw_simple); |
| ops.impl("qgemm_raw_simple_hadamard", c10::kCUDA, &qgemm_raw_simple_hadamard); |
| ops.impl("hadamard_transform", c10::kCUDA, &hadamard_transform); |
| #endif |
| } |
|
|
| REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |
|
|
