feat: make rms_norm as out-place

activation/rms_norm.cu

@@ -271,14 +271,13 @@ rms_norm_backward_kernel(scalar_t *__restrict__ input_grad,        // [..., d]
                                      weight.data_ptr<scalar_t>(), eps, d);     \
   });
 
-void rms_norm(torch::Tensor &out,          // [..., d]
-              const torch::Tensor &input,  // [..., d]
-              const torch::Tensor &weight, // [d]
-              double eps) {
-  AssertTensorShapeEqual(input, out, "input", "out");
+torch::Tensor rms_norm(const torch::Tensor &input,  // [..., d]
+                       const torch::Tensor &weight, // [d]
+                       double eps) {
   AssertTensorNotNull(weight, "weight");
   // TODO shape check
 
+  torch::Tensor out = torch::empty_like(input);
   int d = input.size(-1);
   int64_t num_tokens = input.numel() / input.size(-1);
   dim3 grid(num_tokens);
@@ -292,6 +291,8 @@ void rms_norm(torch::Tensor &out,          // [..., d]
   } else {
     LAUNCH_RMS_NORM(0);
   }
+
+  return out;
 }
 
 #define LAUNCH_RMS_NORM_BWD(width)                                             \
@@ -305,12 +306,14 @@ void rms_norm(torch::Tensor &out,          // [..., d]
                                          weight.data_ptr<scalar_t>(), eps, d); \
       });
 
-void rms_norm_backward(torch::Tensor &input_grad,        // [..., d]
-                       torch::Tensor &weight_grad,       // [d]
-                       const torch::Tensor &output_grad, // [..., d]
-                       const torch::Tensor &input,       // [..., d]
-                       const torch::Tensor &weight,      // [d]
-                       double eps) {
+std::tuple<torch::Tensor, torch::Tensor>
+rms_norm_backward(const torch::Tensor &output_grad, // [..., d]
+                  const torch::Tensor &input,       // [..., d]
+                  const torch::Tensor &weight,      // [d]
+                  double eps) {
+  torch::Tensor input_grad = torch::empty_like(input);
+  torch::Tensor weight_grad = torch::empty_like(weight);
+
   AssertTensorShapeEqual(input, input_grad, "input", "input_grad");
   AssertTensorShapeEqual(input, output_grad, "input", "output_grad");
   AssertTensorNotNull(weight, "weight");
@@ -340,4 +343,6 @@ void rms_norm_backward(torch::Tensor &input_grad,        // [..., d]
     at::sum_out(acc, temp_weight_grad, {0});
     weight_grad.copy_(acc);
   }
+
+  return {input_grad, weight_grad};
 }

tests/test_rms_norm.py

@@ -51,8 +51,7 @@ def test_rms_norm(
     layer = activation.layers.RMSNorm(d, eps=eps, dtype=dtype)
     layer.weight = torch.nn.Parameter(weight)
 
-    out = torch.empty(x.shape, dtype=x.dtype, device=x.device)
-    opcheck(op, (out, x, weight, eps))
+    opcheck(op, (x, weight, eps))
 
     out = fn(x, weight, eps)
     mod_out = layer(x)

torch-ext/activation/rms_norm.py

@@ -8,9 +8,7 @@ class RMSNormFunction(torch.autograd.Function):
     # Note that forward, setup_context, and backward are @staticmethods
     @staticmethod
     def forward(input, weight, eps):
-        output = torch.empty_like(input)
-        ops.rms_norm(output, input, weight, eps)
-        return output
+        return ops.rms_norm(input, weight, eps)
 
     @staticmethod
     # inputs is a Tuple of all of the inputs passed to forward.
@@ -26,13 +24,8 @@ class RMSNormFunction(torch.autograd.Function):
         input, weight = ctx.saved_tensors
         eps = ctx.eps
 
-        input_grad = torch.empty_like(
-            input) if ctx.needs_input_grad[0] else None
-        weight_grad = torch.empty_like(
-            weight) if ctx.needs_input_grad[1] else None
-
-        ops.rms_norm_backward(input_grad, weight_grad, output_grad, input,
-                              weight, eps)
+        input_grad, weight_grad = ops.rms_norm_backward(
+            output_grad, input, weight, eps)
 
         return input_grad, weight_grad, None
 

torch-ext/torch_binding.cpp

@@ -16,12 +16,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("poly_norm_backward", torch::kCUDA, &poly_norm_backward);
 
   // rms_norm
-  ops.def(
-      "rms_norm(Tensor! out, Tensor input, Tensor weight, float eps) -> ()");
+  ops.def("rms_norm(Tensor input, Tensor weight, float eps) -> Tensor");
   ops.impl("rms_norm", torch::kCUDA, &rms_norm);
 
-  ops.def("rms_norm_backward(Tensor! input_grad, Tensor! weight_grad, Tensor "
-          "output_grad, Tensor input, Tensor weight, float eps) -> ()");
+  ops.def("rms_norm_backward(Tensor output_grad, Tensor input, Tensor weight, "
+          "float eps) -> (Tensor, Tensor)");
   ops.impl("rms_norm_backward", torch::kCUDA, &rms_norm_backward);
 
   // fused_mul_poly_norm

torch-ext/torch_binding.h

@@ -11,12 +11,11 @@ void poly_norm_backward(torch::Tensor &input_grad, torch::Tensor &weight_grad,
                         const torch::Tensor &input, const torch::Tensor &weight,
                         double eps);
 
-void rms_norm(torch::Tensor &out, const torch::Tensor &input,
-              const torch::Tensor &weights, double eps);
-void rms_norm_backward(torch::Tensor &input_grad, torch::Tensor &weight_grad,
-                       const torch::Tensor &output_grad,
-                       const torch::Tensor &input, const torch::Tensor &weight,
+torch::Tensor rms_norm(const torch::Tensor &input, const torch::Tensor &weights,
                        double eps);
+std::tuple<torch::Tensor, torch::Tensor>
+rms_norm_backward(const torch::Tensor &output_grad, const torch::Tensor &input,
+                  const torch::Tensor &weight, double eps);
 
 void fused_mul_poly_norm(torch::Tensor &out, const torch::Tensor &input,
                          const torch::Tensor &mul, const torch::Tensor &weights,