File size: 3,652 Bytes
e5e2eeb 44e9845 e5e2eeb 97825b8 e5e2eeb 97825b8 44e9845 f3b99fb e5e2eeb dc88599 e5e2eeb 79a877a dc88599 f3b99fb e5e2eeb 7e4334d e5e2eeb a1e5ca8 7e4334d a1e5ca8 0045757 906e125 6436ad6 0045757 6436ad6 0045757 6436ad6 0045757 6436ad6 0045757 44e9845 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | #include "torch_binding.h"
#include <torch/library.h>
#include "registration.h"
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// poly_norm
ops.def("poly_norm(Tensor! out, Tensor input, Tensor weight, Tensor bias, "
"float eps) -> ()");
ops.impl("poly_norm", torch::kCUDA, &poly_norm);
ops.def("poly_norm_backward(Tensor! input_grad, Tensor! weight_grad, Tensor! "
"bias_grad, Tensor output_grad, Tensor input, Tensor weight, float "
"eps) -> ()");
ops.impl("poly_norm_backward", torch::kCUDA, &poly_norm_backward);
// rms_norm
ops.def(
"rms_norm(Tensor input, Tensor weight, float eps) -> (Tensor, Tensor)");
ops.impl("rms_norm", torch::kCUDA, &rms_norm);
ops.def("rms_norm_backward(Tensor output_grad, Tensor input, Tensor weight, "
"Tensor inv_rms, float eps) -> (Tensor, Tensor)");
ops.impl("rms_norm_backward", torch::kCUDA, &rms_norm_backward);
// fused_mul_poly_norm
ops.def("fused_mul_poly_norm(Tensor! out, Tensor input, Tensor mul, Tensor "
"weight, Tensor bias, "
"float eps) -> ()");
ops.impl("fused_mul_poly_norm", torch::kCUDA, &fused_mul_poly_norm);
ops.def("fused_mul_poly_norm_backward(Tensor! input_grad, Tensor! mul_grad, "
"Tensor! weight_grad, Tensor! "
"bias_grad, Tensor output_grad, Tensor input, Tensor mul, Tensor "
"weight, Tensor "
"bias, float eps) -> ()");
ops.impl("fused_mul_poly_norm_backward", torch::kCUDA,
&fused_mul_poly_norm_backward);
// fused_add_rms_norm
ops.def("fused_add_rms_norm(Tensor input, Tensor residual, Tensor "
"weight, float eps) -> (Tensor, Tensor)");
ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);
ops.def(
"fused_add_rms_norm_backward(Tensor output_grad, Tensor add_output_grad,"
"Tensor input, Tensor weight, float eps, bool need_input_grad) -> "
"(Tensor, Tensor)");
ops.impl("fused_add_rms_norm_backward", torch::kCUDA,
&fused_add_rms_norm_backward);
// grouped_poly_norm (without scores, hidden_clamp < 0 = disabled)
ops.def(
"grouped_poly_norm_forward("
"Tensor input, Tensor mul, Tensor weight, "
"Tensor bias, Tensor offsets, "
"float eps, int expert_offset, float hidden_clamp) -> (Tensor, Tensor)");
ops.impl("grouped_poly_norm_forward", torch::kCUDA,
&grouped_poly_norm_forward);
ops.def("grouped_poly_norm_backward("
"Tensor grad_output, Tensor input, Tensor mul, Tensor weight, "
"Tensor bias, Tensor offsets, Tensor inv_rms, "
"float eps, int expert_offset, float hidden_clamp) -> (Tensor, "
"Tensor, Tensor, Tensor)");
ops.impl("grouped_poly_norm_backward", torch::kCUDA,
&grouped_poly_norm_backward);
// grouped_poly_norm (with scores)
ops.def(
"grouped_poly_norm_forward_scored("
"Tensor input, Tensor mul, Tensor weight, "
"Tensor bias, Tensor offsets, Tensor scores, "
"float eps, int expert_offset, float hidden_clamp) -> (Tensor, Tensor)");
ops.impl("grouped_poly_norm_forward_scored", torch::kCUDA,
&grouped_poly_norm_forward_scored);
ops.def("grouped_poly_norm_backward_scored("
"Tensor grad_output, Tensor input, Tensor mul, Tensor weight, "
"Tensor bias, Tensor offsets, Tensor inv_rms, Tensor scores, "
"float eps, int expert_offset, float hidden_clamp) -> (Tensor, "
"Tensor, Tensor, Tensor, Tensor)");
ops.impl("grouped_poly_norm_backward_scored", torch::kCUDA,
&grouped_poly_norm_backward_scored);
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|