WaveCut's picture
Runtime v11: lookahead decode, ABC head crop, resident compiled VAE, KV handoff, INT8/W4A4 output heads; SM89 validated, SM120 built
9a2611f verified
Raw
History Blame Contribute Delete
2.77 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);
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)