File size: 2,619 Bytes
0154673 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
#include <torch/library.h>
#include "registration.h"
#include "torch_binding.h"
extern "C" void higgs_dequantize_2_256_ptr_cuda_portable(uint64_t x_ptr,
uint64_t grid_ptr,
uint64_t out_ptr,
int64_t out_dim);
extern "C" void higgs_quantize_2_256_ptr_f16_cuda_portable(
uint64_t x_ptr, uint64_t grid_ptr, uint64_t grid_norms_ptr,
uint64_t out_ptr, int64_t out_dim);
extern "C" void higgs_quantize_2_256_ptr_bf16_cuda_portable(
uint64_t x_ptr, uint64_t grid_ptr, uint64_t grid_norms_ptr,
uint64_t out_ptr, int64_t out_dim);
void higgs_dequantize_2_256(torch::Tensor x, torch::Tensor grid,
torch::Tensor out) {
int64_t out_dim = x.size(0);
higgs_dequantize_2_256_ptr_cuda_portable(
reinterpret_cast<uint64_t>(x.data_ptr()),
reinterpret_cast<uint64_t>(grid.data_ptr()),
reinterpret_cast<uint64_t>(out.data_ptr()), out_dim);
}
void higgs_quantize_2_256_f16(torch::Tensor x, torch::Tensor grid,
torch::Tensor grid_norms, torch::Tensor out) {
int64_t out_dim = x.size(0);
higgs_quantize_2_256_ptr_f16_cuda_portable(
reinterpret_cast<uint64_t>(x.data_ptr()),
reinterpret_cast<uint64_t>(grid.data_ptr()),
reinterpret_cast<uint64_t>(grid_norms.data_ptr()),
reinterpret_cast<uint64_t>(out.data_ptr()), out_dim);
}
void higgs_quantize_2_256_bf16(torch::Tensor x, torch::Tensor grid,
torch::Tensor grid_norms, torch::Tensor out) {
int64_t out_dim = x.size(0);
higgs_quantize_2_256_ptr_bf16_cuda_portable(
reinterpret_cast<uint64_t>(x.data_ptr()),
reinterpret_cast<uint64_t>(grid.data_ptr()),
reinterpret_cast<uint64_t>(grid_norms.data_ptr()),
reinterpret_cast<uint64_t>(out.data_ptr()), out_dim);
}
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"higgs_dequantize_2_256(Tensor x, Tensor grid, Tensor! out) -> ()");
ops.impl("higgs_dequantize_2_256", torch::kCUDA, &higgs_dequantize_2_256);
ops.def("higgs_quantize_2_256_f16(Tensor x, Tensor grid, Tensor "
"grid_norms, Tensor! out) -> ()");
ops.impl("higgs_quantize_2_256_f16", torch::kCUDA,
&higgs_quantize_2_256_f16);
ops.def("higgs_quantize_2_256_bf16(Tensor x, Tensor grid, Tensor "
"grid_norms, Tensor! out) -> ()");
ops.impl("higgs_quantize_2_256_bf16", torch::kCUDA,
&higgs_quantize_2_256_bf16);
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|