File size: 922 Bytes
1dc29e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
#include <torch/library.h>
#include "registration.h"
#include "torch_binding.h"
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("w8_a16_gemm(Tensor input, Tensor weight, Tensor scale) -> Tensor");
ops.impl("w8_a16_gemm", torch::kCUDA, &w8_a16_gemm_forward_cuda);
ops.def("w8_a16_gemm_(Tensor input, Tensor weight, Tensor scale, Tensor! output,"
"int m, int n, int k) -> Tensor");
ops.impl("w8_a16_gemm_", torch::kCUDA, &w8_a16_gemm_forward_cuda_);
ops.def("preprocess_weights(Tensor origin_weight, bool is_int4) -> Tensor");
ops.impl("preprocess_weights", torch::kCUDA, &preprocess_weights_cuda);
ops.def("quant_weights(Tensor origin_weight, ScalarType quant_type,"
"bool return_unprocessed_quantized_tensor) -> Tensor[]");
ops.impl("quant_weights", torch::kCUDA, &symmetric_quantize_last_axis_of_tensor);
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|