Build uploaded using `kernels`.

.gitattributes

@@ -64,3 +64,9 @@ build/torch28-cxx11-cpu-x86_64-linux/_rmsnorm_a8702c9.abi3.so filter=lfs diff=lf
 build/torch28-cxx11-xpu20251-x86_64-linux/_rmsnorm_a8702c9.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-cpu-x86_64-linux/_rmsnorm_a8702c9.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-cxx11-xpu20252-x86_64-linux/_rmsnorm_a8702c9.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cpu-x86_64-linux/_rmsnorm_235cde1.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-xpu20253-x86_64-linux/_rmsnorm_235cde1.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch28-cxx11-cpu-x86_64-linux/_rmsnorm_235cde1.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch28-cxx11-xpu20251-x86_64-linux/_rmsnorm_235cde1.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cpu-x86_64-linux/_rmsnorm_235cde1.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-xpu20252-x86_64-linux/_rmsnorm_235cde1.abi3.so filter=lfs diff=lfs merge=lfs -text

build/torch210-cxx11-cpu-x86_64-linux/__init__.py

@@ -4,11 +4,24 @@ from ._ops import ops
 
 
 def apply_rms_norm(input, weight, eps):
+    # ops.apply_rms_norm returns [output, rstd]
     return ops.apply_rms_norm(
             input,
             weight,
             eps,
+    )[0]
+
+def apply_rms_norm_backward(grad_output, input, weight, output, rstd, eps, input_requires_grad=True, weight_requires_grad=True):
+    return ops.apply_rms_norm_backward(
+            grad_output,
+            input,
+            weight,
+            output,
+            rstd,
+            eps,
+            input_requires_grad,
+            weight_requires_grad
     )
 
-__all__ = ["layers", "apply_rms_norm"]
+__all__ = ["layers", "apply_rms_norm_forward", "apply_rms_norm_backward"]
 

build/torch210-cxx11-cpu-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _rmsnorm_a8702c9
-ops = torch.ops._rmsnorm_a8702c9
+from . import _rmsnorm_235cde1
+ops = torch.ops._rmsnorm_235cde1
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rmsnorm_a8702c9::{op_name}"
+    return f"_rmsnorm_235cde1::{op_name}"

build/torch210-cxx11-cpu-x86_64-linux/{_rmsnorm_a8702c9.abi3.so → _rmsnorm_235cde1.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1662b01d67212b84f7395898dc325e2ce58c2c988d045f0c1ac61e204ca6e3f2
-size 326352
+oid sha256:08b2b07a7190f17e6463890cec75f84f96b10235e6daafc6adc3d4807c868607
+size 158864

build/torch210-cxx11-cpu-x86_64-linux/layers.py

@@ -1,6 +1,29 @@
 import torch
 from ._ops import ops
 
+class RMSNormFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, hidden_states, weight, variance_epsilon):
+        ctx.variance_epsilon = variance_epsilon
+        output, rstd = ops.apply_rms_norm(hidden_states, weight, variance_epsilon)
+        ctx.save_for_backward(hidden_states, weight, output, rstd)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        hidden_states, weight, output, rstd = ctx.saved_tensors
+        grads = ops.apply_rms_norm_backward(
+            grad_output,
+            hidden_states,
+            weight,
+            output,
+            rstd,
+            ctx.variance_epsilon,
+            ctx.needs_input_grad[0],
+            ctx.needs_input_grad[1]
+        )
+        return grads[0], grads[1], None
+
 class RMSNorm(torch.nn.Module):
     """
     RMSNorm module that uses the optimized LigerRMSNormFunction.
@@ -27,7 +50,7 @@ class RMSNorm(torch.nn.Module):
         Returns:
             torch.Tensor: Normalized tensor of the same shape as input
         """
-        return ops.apply_rms_norm(
+        return RMSNormFunction.apply(
             hidden_states,
             self.weight, 
             self.variance_epsilon,

build/torch210-cxx11-xpu20253-x86_64-linux/__init__.py

@@ -4,11 +4,24 @@ from ._ops import ops
 
 
 def apply_rms_norm(input, weight, eps):
+    # ops.apply_rms_norm returns [output, rstd]
     return ops.apply_rms_norm(
             input,
             weight,
             eps,
+    )[0]
+
+def apply_rms_norm_backward(grad_output, input, weight, output, rstd, eps, input_requires_grad=True, weight_requires_grad=True):
+    return ops.apply_rms_norm_backward(
+            grad_output,
+            input,
+            weight,
+            output,
+            rstd,
+            eps,
+            input_requires_grad,
+            weight_requires_grad
     )
 
-__all__ = ["layers", "apply_rms_norm"]
+__all__ = ["layers", "apply_rms_norm_forward", "apply_rms_norm_backward"]
 

build/torch210-cxx11-xpu20253-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _rmsnorm_a8702c9
-ops = torch.ops._rmsnorm_a8702c9
+from . import _rmsnorm_235cde1
+ops = torch.ops._rmsnorm_235cde1
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rmsnorm_a8702c9::{op_name}"
+    return f"_rmsnorm_235cde1::{op_name}"

build/torch210-cxx11-xpu20253-x86_64-linux/{_rmsnorm_a8702c9.abi3.so → _rmsnorm_235cde1.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e1083ef90913f7ae1cabbf9e1c184ddfe6abe9b31fc9de7d29f004fa807c3a9c
-size 104941392
+oid sha256:7a3adbeed741dcfb0ae701761009dba190d6432383906831741cedc133a519d7
+size 104793352

build/torch210-cxx11-xpu20253-x86_64-linux/layers.py

@@ -1,6 +1,29 @@
 import torch
 from ._ops import ops
 
+class RMSNormFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, hidden_states, weight, variance_epsilon):
+        ctx.variance_epsilon = variance_epsilon
+        output, rstd = ops.apply_rms_norm(hidden_states, weight, variance_epsilon)
+        ctx.save_for_backward(hidden_states, weight, output, rstd)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        hidden_states, weight, output, rstd = ctx.saved_tensors
+        grads = ops.apply_rms_norm_backward(
+            grad_output,
+            hidden_states,
+            weight,
+            output,
+            rstd,
+            ctx.variance_epsilon,
+            ctx.needs_input_grad[0],
+            ctx.needs_input_grad[1]
+        )
+        return grads[0], grads[1], None
+
 class RMSNorm(torch.nn.Module):
     """
     RMSNorm module that uses the optimized LigerRMSNormFunction.
@@ -27,7 +50,7 @@ class RMSNorm(torch.nn.Module):
         Returns:
             torch.Tensor: Normalized tensor of the same shape as input
         """
-        return ops.apply_rms_norm(
+        return RMSNormFunction.apply(
             hidden_states,
             self.weight, 
             self.variance_epsilon,

build/torch28-cxx11-cpu-x86_64-linux/__init__.py

@@ -4,11 +4,24 @@ from ._ops import ops
 
 
 def apply_rms_norm(input, weight, eps):
+    # ops.apply_rms_norm returns [output, rstd]
     return ops.apply_rms_norm(
             input,
             weight,
             eps,
+    )[0]
+
+def apply_rms_norm_backward(grad_output, input, weight, output, rstd, eps, input_requires_grad=True, weight_requires_grad=True):
+    return ops.apply_rms_norm_backward(
+            grad_output,
+            input,
+            weight,
+            output,
+            rstd,
+            eps,
+            input_requires_grad,
+            weight_requires_grad
     )
 
-__all__ = ["layers", "apply_rms_norm"]
+__all__ = ["layers", "apply_rms_norm_forward", "apply_rms_norm_backward"]
 

build/torch28-cxx11-cpu-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _rmsnorm_a8702c9
-ops = torch.ops._rmsnorm_a8702c9
+from . import _rmsnorm_235cde1
+ops = torch.ops._rmsnorm_235cde1
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rmsnorm_a8702c9::{op_name}"
+    return f"_rmsnorm_235cde1::{op_name}"

build/{torch29-cxx11-cpu-x86_64-linux/_rmsnorm_a8702c9.abi3.so → torch28-cxx11-cpu-x86_64-linux/_rmsnorm_235cde1.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c828a99dff803d9cdf8c0cca8353970b662bfee79cce6e30c30b4c852f797997
-size 324592
+oid sha256:16c92de9cefabeeadc60ffff87189a1e66ecb9ea19b343570ac55e9d9c7d98fe
+size 156648

build/torch28-cxx11-cpu-x86_64-linux/layers.py

@@ -1,6 +1,29 @@
 import torch
 from ._ops import ops
 
+class RMSNormFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, hidden_states, weight, variance_epsilon):
+        ctx.variance_epsilon = variance_epsilon
+        output, rstd = ops.apply_rms_norm(hidden_states, weight, variance_epsilon)
+        ctx.save_for_backward(hidden_states, weight, output, rstd)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        hidden_states, weight, output, rstd = ctx.saved_tensors
+        grads = ops.apply_rms_norm_backward(
+            grad_output,
+            hidden_states,
+            weight,
+            output,
+            rstd,
+            ctx.variance_epsilon,
+            ctx.needs_input_grad[0],
+            ctx.needs_input_grad[1]
+        )
+        return grads[0], grads[1], None
+
 class RMSNorm(torch.nn.Module):
     """
     RMSNorm module that uses the optimized LigerRMSNormFunction.
@@ -27,7 +50,7 @@ class RMSNorm(torch.nn.Module):
         Returns:
             torch.Tensor: Normalized tensor of the same shape as input
         """
-        return ops.apply_rms_norm(
+        return RMSNormFunction.apply(
             hidden_states,
             self.weight, 
             self.variance_epsilon,

build/torch28-cxx11-xpu20251-x86_64-linux/__init__.py

@@ -4,11 +4,24 @@ from ._ops import ops
 
 
 def apply_rms_norm(input, weight, eps):
+    # ops.apply_rms_norm returns [output, rstd]
     return ops.apply_rms_norm(
             input,
             weight,
             eps,
+    )[0]
+
+def apply_rms_norm_backward(grad_output, input, weight, output, rstd, eps, input_requires_grad=True, weight_requires_grad=True):
+    return ops.apply_rms_norm_backward(
+            grad_output,
+            input,
+            weight,
+            output,
+            rstd,
+            eps,
+            input_requires_grad,
+            weight_requires_grad
     )
 
-__all__ = ["layers", "apply_rms_norm"]
+__all__ = ["layers", "apply_rms_norm_forward", "apply_rms_norm_backward"]
 

build/torch28-cxx11-xpu20251-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _rmsnorm_a8702c9
-ops = torch.ops._rmsnorm_a8702c9
+from . import _rmsnorm_235cde1
+ops = torch.ops._rmsnorm_235cde1
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rmsnorm_a8702c9::{op_name}"
+    return f"_rmsnorm_235cde1::{op_name}"

build/torch28-cxx11-xpu20251-x86_64-linux/{_rmsnorm_a8702c9.abi3.so → _rmsnorm_235cde1.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ad7611e794711a0b4016d2c6ec928a0b69ec852353a68cf6b01b895e15021204
-size 103861336
+oid sha256:77c4b43d63dc74b210633da81630023a6d6e359a7a1115bff55da9f4436053d9
+size 103700632

build/torch28-cxx11-xpu20251-x86_64-linux/layers.py

@@ -1,6 +1,29 @@
 import torch
 from ._ops import ops
 
+class RMSNormFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, hidden_states, weight, variance_epsilon):
+        ctx.variance_epsilon = variance_epsilon
+        output, rstd = ops.apply_rms_norm(hidden_states, weight, variance_epsilon)
+        ctx.save_for_backward(hidden_states, weight, output, rstd)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        hidden_states, weight, output, rstd = ctx.saved_tensors
+        grads = ops.apply_rms_norm_backward(
+            grad_output,
+            hidden_states,
+            weight,
+            output,
+            rstd,
+            ctx.variance_epsilon,
+            ctx.needs_input_grad[0],
+            ctx.needs_input_grad[1]
+        )
+        return grads[0], grads[1], None
+
 class RMSNorm(torch.nn.Module):
     """
     RMSNorm module that uses the optimized LigerRMSNormFunction.
@@ -27,7 +50,7 @@ class RMSNorm(torch.nn.Module):
         Returns:
             torch.Tensor: Normalized tensor of the same shape as input
         """
-        return ops.apply_rms_norm(
+        return RMSNormFunction.apply(
             hidden_states,
             self.weight, 
             self.variance_epsilon,

build/torch29-cxx11-cpu-x86_64-linux/__init__.py

@@ -4,11 +4,24 @@ from ._ops import ops
 
 
 def apply_rms_norm(input, weight, eps):
+    # ops.apply_rms_norm returns [output, rstd]
     return ops.apply_rms_norm(
             input,
             weight,
             eps,
+    )[0]
+
+def apply_rms_norm_backward(grad_output, input, weight, output, rstd, eps, input_requires_grad=True, weight_requires_grad=True):
+    return ops.apply_rms_norm_backward(
+            grad_output,
+            input,
+            weight,
+            output,
+            rstd,
+            eps,
+            input_requires_grad,
+            weight_requires_grad
     )
 
-__all__ = ["layers", "apply_rms_norm"]
+__all__ = ["layers", "apply_rms_norm_forward", "apply_rms_norm_backward"]
 

build/torch29-cxx11-cpu-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _rmsnorm_a8702c9
-ops = torch.ops._rmsnorm_a8702c9
+from . import _rmsnorm_235cde1
+ops = torch.ops._rmsnorm_235cde1
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rmsnorm_a8702c9::{op_name}"
+    return f"_rmsnorm_235cde1::{op_name}"

build/{torch28-cxx11-cpu-x86_64-linux/_rmsnorm_a8702c9.abi3.so → torch29-cxx11-cpu-x86_64-linux/_rmsnorm_235cde1.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:549a788f18fc51a309681f459fcf55f80eeebbdcb8cf082d34068a7d53822268
-size 324616
+oid sha256:e9ea3ebe5949d376bb44127a66dc2946e72620ff7035e2f34e81a652b0c69ded
+size 156608

build/torch29-cxx11-cpu-x86_64-linux/layers.py

@@ -1,6 +1,29 @@
 import torch
 from ._ops import ops
 
+class RMSNormFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, hidden_states, weight, variance_epsilon):
+        ctx.variance_epsilon = variance_epsilon
+        output, rstd = ops.apply_rms_norm(hidden_states, weight, variance_epsilon)
+        ctx.save_for_backward(hidden_states, weight, output, rstd)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        hidden_states, weight, output, rstd = ctx.saved_tensors
+        grads = ops.apply_rms_norm_backward(
+            grad_output,
+            hidden_states,
+            weight,
+            output,
+            rstd,
+            ctx.variance_epsilon,
+            ctx.needs_input_grad[0],
+            ctx.needs_input_grad[1]
+        )
+        return grads[0], grads[1], None
+
 class RMSNorm(torch.nn.Module):
     """
     RMSNorm module that uses the optimized LigerRMSNormFunction.
@@ -27,7 +50,7 @@ class RMSNorm(torch.nn.Module):
         Returns:
             torch.Tensor: Normalized tensor of the same shape as input
         """
-        return ops.apply_rms_norm(
+        return RMSNormFunction.apply(
             hidden_states,
             self.weight, 
             self.variance_epsilon,

build/torch29-cxx11-xpu20252-x86_64-linux/__init__.py

@@ -4,11 +4,24 @@ from ._ops import ops
 
 
 def apply_rms_norm(input, weight, eps):
+    # ops.apply_rms_norm returns [output, rstd]
     return ops.apply_rms_norm(
             input,
             weight,
             eps,
+    )[0]
+
+def apply_rms_norm_backward(grad_output, input, weight, output, rstd, eps, input_requires_grad=True, weight_requires_grad=True):
+    return ops.apply_rms_norm_backward(
+            grad_output,
+            input,
+            weight,
+            output,
+            rstd,
+            eps,
+            input_requires_grad,
+            weight_requires_grad
     )
 
-__all__ = ["layers", "apply_rms_norm"]
+__all__ = ["layers", "apply_rms_norm_forward", "apply_rms_norm_backward"]
 

build/torch29-cxx11-xpu20252-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _rmsnorm_a8702c9
-ops = torch.ops._rmsnorm_a8702c9
+from . import _rmsnorm_235cde1
+ops = torch.ops._rmsnorm_235cde1
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_rmsnorm_a8702c9::{op_name}"
+    return f"_rmsnorm_235cde1::{op_name}"

build/torch29-cxx11-xpu20252-x86_64-linux/_rmsnorm_235cde1.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9cf81702c9b0a75c84f961b1b7555db936dbdd7b889c1a41dd048c1100a0aa2
+size 102179528

build/torch29-cxx11-xpu20252-x86_64-linux/_rmsnorm_a8702c9.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e00f18c2a5222d5782f37d052526d32ba71de2c28a271de5c85a5a8fb0efe6fa
-size 102340240

build/torch29-cxx11-xpu20252-x86_64-linux/layers.py

@@ -1,6 +1,29 @@
 import torch
 from ._ops import ops
 
+class RMSNormFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, hidden_states, weight, variance_epsilon):
+        ctx.variance_epsilon = variance_epsilon
+        output, rstd = ops.apply_rms_norm(hidden_states, weight, variance_epsilon)
+        ctx.save_for_backward(hidden_states, weight, output, rstd)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        hidden_states, weight, output, rstd = ctx.saved_tensors
+        grads = ops.apply_rms_norm_backward(
+            grad_output,
+            hidden_states,
+            weight,
+            output,
+            rstd,
+            ctx.variance_epsilon,
+            ctx.needs_input_grad[0],
+            ctx.needs_input_grad[1]
+        )
+        return grads[0], grads[1], None
+
 class RMSNorm(torch.nn.Module):
     """
     RMSNorm module that uses the optimized LigerRMSNormFunction.
@@ -27,7 +50,7 @@ class RMSNorm(torch.nn.Module):
         Returns:
             torch.Tensor: Normalized tensor of the same shape as input
         """
-        return ops.apply_rms_norm(
+        return RMSNormFunction.apply(
             hidden_states,
             self.weight, 
             self.variance_epsilon,