#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);
  ops.def("fused_gemv(Tensor! out, Tensor x, Tensor gamma, bool has_gamma, Tensor permutation, Tensor signs, Tensor boundaries, float rms_eps, float quant_eps, float inv_sqrt_block, Tensor w, Tensor wn, Tensor ac, Tensor wc, Tensor bias, bool has_bias, float activation_scale, float weight_scale, int virtual_threads) -> ()");
  ops.impl("fused_gemv", torch::kCUDA, &fused_gemv);
  ops.def("fused_gemv2(Tensor! out, Tensor! x_new, bool write_x_new, Tensor x, Tensor h, bool has_h, Tensor gamma, bool has_gamma, Tensor permutation, Tensor signs, Tensor boundaries, float rms_eps, float quant_eps, float inv_sqrt_block, Tensor w, Tensor wn, Tensor ac, Tensor wc, Tensor bias, bool has_bias, float activation_scale, float weight_scale, int virtual_threads) -> ()");
  ops.impl("fused_gemv2", torch::kCUDA, &fused_gemv2);
  ops.def("rmsquant_add(Tensor! packed, Tensor! norms, Tensor! x_new, Tensor x, Tensor h, Tensor weight, Tensor permutation, Tensor signs, Tensor boundaries, float rms_eps, float quant_eps) -> ()");
  ops.impl("rmsquant_add", torch::kCUDA, &rmsquant_add);
  ops.def("quantize_rows_int8(Tensor! out, Tensor! scales, Tensor x) -> ()");
  ops.impl("quantize_rows_int8", torch::kCUDA, &quantize_rows_int8);
  ops.def("gemv_w8a8(Tensor! out, Tensor x, Tensor xs, Tensor w, Tensor ws) -> ()");
  ops.impl("gemv_w8a8", torch::kCUDA, &gemv_w8a8);
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)