feat: refine rms norm ops

csrc/pytorch/bindings_pytorch.cpp

@@ -58,10 +58,6 @@ TORCH_LIBRARY(primus_turbo_cpp_extension, m) {
     m.def("shuffle_scale(Tensor scale, int[] layout) -> Tensor");
     m.def("shuffle_weight(Tensor weight, int[] layout) -> Tensor");
 
-    // ********* RMSNorm *********
-    m.def("rmsnorm_fwd(Tensor input, Tensor gamma, float eps) -> Tensor");
-    m.def("rmsnorm_bwd(Tensor input, Tensor gamma, Tensor grad_out, float eps) -> Tensor[]");
-
     // ********* Grouped Gemm *********
     m.def("ck_grouped_gemm(Tensor a, Tensor b, Tensor group_lens, Tensor group_offs, bool transA, "
           "bool transB, int? num_cu=None) -> Tensor");
@@ -105,10 +101,6 @@ TORCH_LIBRARY_IMPL(primus_turbo_cpp_extension, CUDA, m) {
     m.impl("shuffle_scale", shuffle_scale_impl);
     m.impl("shuffle_weight", shuffle_weight_impl);
 
-    // ********* RMSNorm *********
-    m.impl("rmsnorm_fwd", rmsnorm_fwd);
-    m.impl("rmsnorm_bwd", rmsnorm_bwd);
-
     // ********* Grouped Gemm *********
     m.impl("ck_grouped_gemm", ck_grouped_gemm);
     m.impl("ck_grouped_gemm_variable_k", ck_grouped_gemm_variable_k);
@@ -143,10 +135,6 @@ TORCH_LIBRARY_IMPL(primus_turbo_cpp_extension, Meta, m) {
     m.impl("shuffle_scale", shuffle_scale_impl_meta);
     m.impl("shuffle_weight", shuffle_weight_impl_meta);
 
-    // ********* RMSNorm *********
-    m.impl("rmsnorm_fwd", rmsnorm_fwd_meta);
-    m.impl("rmsnorm_bwd", rmsnorm_bwd_meta);
-
     // ********* Grouped Gemm *********
     m.impl("ck_grouped_gemm", ck_grouped_gemm_meta);
     m.impl("ck_grouped_gemm_variable_k", ck_grouped_gemm_variable_k_meta);

csrc/pytorch/extensions.h

@@ -153,20 +153,6 @@ at::Tensor turbo_gemm_fp8_meta(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B
                                at::Tensor scaleB_inv, const at::ScalarType out_dtype, bool transA,
                                bool transB, bool transC, const std::string &granularity);
 
-//==================================================================
-//  Normalization
-//==================================================================
-
-at::Tensor rmsnorm_fwd(const at::Tensor &input, const at::Tensor &gamma, const double eps);
-
-at::Tensor rmsnorm_fwd_meta(const at::Tensor &input, const at::Tensor &gamma, const double eps);
-
-std::vector<at::Tensor> rmsnorm_bwd(const at::Tensor &input, const at::Tensor &gamma,
-                                    const at::Tensor &grad_output, const double eps);
-
-std::vector<at::Tensor> rmsnorm_bwd_meta(const at::Tensor &input, const at::Tensor &gamma,
-                                         const at::Tensor &grad_output, const double eps);
-
 //==================================================================
 //  Grouped GEMM
 //==================================================================