feat: add assert is_contiguous

activation/assert_utils.h

@@ -8,6 +8,18 @@ inline void AssertTensorNotNull(const torch::Tensor &tensor,
   TORCH_INTERNAL_ASSERT(tensor.defined(), name + " tensor should not be null.");
 }
 
+inline void AssertTensorContiguous(const torch::Tensor &tensor,
+                                   const std::string &name,
+                                   bool nullable = false) {
+  if (nullable && !tensor.defined()) {
+    return;
+  }
+
+  AssertTensorNotNull(tensor, name);
+  TORCH_INTERNAL_ASSERT(tensor.is_contiguous(),
+                        name + " tensor should be contiguous.");
+}
+
 inline void AssertTensorShapeEqual(const torch::Tensor &tensor_a,
                                    const torch::Tensor &tensor_b,
                                    const std::string &name_a,

activation/fused_add_rms_norm.cu

@@ -307,6 +307,12 @@ void fused_add_rms_norm(torch::Tensor &out,            // [..., d]
   AssertTensorNotNull(weight, "weight");
   // TODO shape check
 
+  AssertTensorContiguous(out, "out");
+  AssertTensorContiguous(add_out, "add_out");
+  AssertTensorContiguous(input, "input");
+  AssertTensorContiguous(residual, "residual");
+  AssertTensorContiguous(weight, "weight");
+
   int d = input.size(-1);
   int64_t num_tokens = input.numel() / input.size(-1);
   dim3 grid(num_tokens);
@@ -346,6 +352,14 @@ void fused_add_rms_norm_backward(
   AssertTensorShapeEqual(input, output_grad, "input", "output_grad");
   AssertTensorShapeEqual(input, output_grad, "input", "add_output_grad");
   AssertTensorNotNull(weight, "weight");
+
+  constexpr bool ALLOW_NULL = true;
+  AssertTensorContiguous(input_grad, "input_grad", ALLOW_NULL);
+  AssertTensorContiguous(weight_grad, "weight_grad", ALLOW_NULL);
+  AssertTensorContiguous(output_grad, "output_grad");
+  AssertTensorContiguous(add_output_grad, "add_output_grad");
+  AssertTensorContiguous(input, "input");
+  AssertTensorContiguous(weight, "weight");
   // TODO shape check
   // weight_grad, input_grad can be nullable
 

activation/fused_mul_poly_norm.cu

@@ -556,6 +556,12 @@ void fused_mul_poly_norm(torch::Tensor &out,          // [..., d]
   AssertTensorShapeEqual(input, mul, "input", "mul");
   AssertTensorNotNull(weight, "weight");
   AssertTensorNotNull(bias, "bias");
+
+  AssertTensorContiguous(out, "out");
+  AssertTensorContiguous(input, "input");
+  AssertTensorContiguous(mul, "mul");
+  AssertTensorContiguous(weight, "weight");
+  AssertTensorContiguous(bias, "bias");
   // TODO shape check
 
   int d = input.size(-1);
@@ -602,6 +608,17 @@ void fused_mul_poly_norm_backward(torch::Tensor &input_grad,        // [..., d]
   AssertTensorShapeEqual(input, mul_grad, "input", "mul_grad");
   AssertTensorShapeEqual(input, mul, "input", "mul");
   AssertTensorNotNull(weight, "weight");
+
+  constexpr bool ALLOW_NULL = true;
+  AssertTensorContiguous(input_grad, "input_grad", ALLOW_NULL);
+  AssertTensorContiguous(mul_grad, "mul_grad", ALLOW_NULL);
+  AssertTensorContiguous(weight_grad, "weight_grad", ALLOW_NULL);
+  AssertTensorContiguous(bias_grad, "bias_grad", ALLOW_NULL);
+  AssertTensorContiguous(output_grad, "output_grad");
+  AssertTensorContiguous(input, "input");
+  AssertTensorContiguous(mul, "mul");
+  AssertTensorContiguous(weight, "weight");
+  AssertTensorContiguous(bias, "bias");
   // TODO shape check
   // weight_grad, bias_grad, mul_grad and input_grad can be nullable
 

activation/poly_norm.cu

@@ -508,6 +508,11 @@ void poly_norm(torch::Tensor &out,          // [..., d]
   AssertTensorNotNull(bias, "bias");
   // TODO shape check
 
+  AssertTensorContiguous(out, "out");
+  AssertTensorContiguous(input, "input");
+  AssertTensorContiguous(weight, "weight");
+  AssertTensorContiguous(bias, "bias");
+
   int d = input.size(-1);
   int64_t num_tokens = input.numel() / d;
   dim3 grid(num_tokens);
@@ -548,6 +553,14 @@ void poly_norm_backward(torch::Tensor &input_grad,        // [..., d]
   // TODO shape check
   // weight_grad, bias_grad and input_grad can be nullable
 
+  constexpr bool ALLOW_NULL = true;
+  AssertTensorContiguous(input_grad, "input_grad", ALLOW_NULL);
+  AssertTensorContiguous(weight_grad, "weight_grad", ALLOW_NULL);
+  AssertTensorContiguous(bias_grad, "bias_grad", ALLOW_NULL);
+  AssertTensorContiguous(output_grad, "output_grad");
+  AssertTensorContiguous(input, "input");
+  AssertTensorContiguous(weight, "weight");
+
   int d = input.size(-1);
   int64_t num_tokens = input.numel() / d;
   dim3 grid(num_tokens);

activation/rms_norm.cu

@@ -276,6 +276,8 @@ torch::Tensor rms_norm(const torch::Tensor &input,  // [..., d]
                        double eps) {
   AssertTensorNotNull(weight, "weight");
   // TODO shape check
+  AssertTensorContiguous(input, "input");
+  AssertTensorContiguous(weight, "weight");
 
   torch::Tensor out = torch::empty_like(input);
   int d = input.size(-1);
@@ -314,6 +316,10 @@ rms_norm_backward(const torch::Tensor &output_grad, // [..., d]
   torch::Tensor input_grad = torch::empty_like(input);
   torch::Tensor weight_grad = torch::empty_like(weight);
 
+  AssertTensorContiguous(output_grad, "output_grad");
+  AssertTensorContiguous(input, "input");
+  AssertTensorContiguous(weight, "weight");
+
   AssertTensorShapeEqual(input, input_grad, "input", "input_grad");
   AssertTensorShapeEqual(input, output_grad, "input", "output_grad");
   AssertTensorNotNull(weight, "weight");