[Optimization] Implement fused add rmsnorm (#1667)

csrc/layernorm.cpp

@@ -6,9 +6,19 @@ void rms_norm(
   torch::Tensor& weight,
   float epsilon);
 
+void fused_add_rms_norm(
+  torch::Tensor& input,
+  torch::Tensor& residual,
+  torch::Tensor& weight,
+  float epsilon);
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def(
     "rms_norm",
     &rms_norm,
     "Apply Root Mean Square (RMS) Normalization to the input tensor.");
+  m.def(
+    "fused_add_rms_norm",
+    &fused_add_rms_norm,
+    "In-place fused Add and RMS Normalization");
 }