refactor RMSNorm by kernels-builder

build/torch-universal/RMSNorm/{layers.py → RMSNorm.py}

@@ -1,17 +1,29 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
 import torch
 import torch_npu
+from ._ops import add_op_namespace_prefix
 
 
-class RMSNorm(torch.nn.Module):
-    def forward(self, x):
+@torch.library.custom_op(add_op_namespace_prefix("RMSNorm"), mutates_args=())
+class _RMSNorm(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass through the RMSNorm layer.
-
         Args:
             x (torch.Tensor): The input tensor.
-
         Returns:
             torch.Tensor: The output tensor after applying RMSNorm.
-
         """
         return torch_npu.npu_rms_norm(x, self.weight, epsilon=self.variance_epsilon)[0]

build/torch-universal/RMSNorm/__init__.py

@@ -1,8 +1,30 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
 import torch
-import torch_npu
-from . import layers
+from ._ops import ops
+from .RMSNorm import _RMSNorm
+
+
+def RMSNorm(x: torch.Tensor) -> torch.Tensor:
+    """
+    Apply RMSNorm to the input tensor.
+    Args:
+        x (torch.Tensor): The input tensor.
+    Returns:
+        torch.Tensor: The output tensor after applying RMSNorm.
+    """
+    return ops.RMSNorm(x)
 
 
-__all__ = [
-    "layers",
-]
+__all__ = ["RMSNorm"]

build/torch-universal/RMSNorm/_ops.py

@@ -0,0 +1,8 @@
+import torch
+ops = torch.ops._RMSNorm_20251112073655
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_RMSNorm_20251112073655::{op_name}"