WaveCut's picture
Publish OrbitQuant 0.9.5 byte-pair decode and optional RMS activation fusion
b35ab1f verified
Raw
History Blame Contribute Delete
1.41 kB
#include <torch/library.h>
#include "registration.h"
#include "torch_binding.h"
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("rmsquant(Tensor! packed, Tensor! norms, Tensor x, Tensor weight, Tensor permutation, Tensor signs, Tensor boundaries, float rms_eps, float quant_eps) -> ()");
ops.impl("rmsquant", torch::kCUDA, &rmsquant);
ops.def("orbitquant_gemv(Tensor! out, Tensor x, Tensor w, Tensor xn, Tensor wn, Tensor ac, Tensor wc, Tensor bias, bool has_bias, float activation_scale, float weight_scale) -> ()");
ops.impl("orbitquant_gemv", torch::kCUDA, &orbitquant_gemv);
ops.def("yue2_rmsnorm(Tensor! out, Tensor x, Tensor weight, float eps) -> ()");
ops.impl("yue2_rmsnorm", torch::kCUDA, &yue2_rmsnorm);
ops.def("yue2_rope(Tensor! out, Tensor x, Tensor cos, Tensor sin, int heads) -> ()");
ops.impl("yue2_rope", torch::kCUDA, &yue2_rope);
ops.def("yue2_swiglu(Tensor! out, Tensor x) -> ()");
ops.impl("yue2_swiglu", torch::kCUDA, &yue2_swiglu);
ops.def("yue2_qkv_decode(Tensor! q, Tensor! keys, Tensor! values, Tensor x, Tensor qw, Tensor kw, Tensor cos, Tensor sin, Tensor positions, float eps) -> ()");
ops.impl("yue2_qkv_decode", torch::kCUDA, &yue2_qkv_decode);
ops.def("yue2_rmsnorm_exact(Tensor! out, Tensor x, Tensor weight, float eps) -> ()");
ops.impl("yue2_rmsnorm_exact", torch::kCUDA, &yue2_rmsnorm_exact);
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)