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