some updates

torch-ext/rmsnorm_kernel/__init__.py

@@ -1,21 +1,13 @@
 import torch
-import torch.nn as nn
 
 from ._ops import ops
 
+from . import layers
 
-class LlamaRMSNorm(nn.Module):
-    weight: torch.Tensor
-    variance_epsilon: float
+def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    return ops.rmsnorm_forward(x, weight)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return ops.rmsnorm_forward(
-            hidden_states,
-            self.weight,
-            bias=None,
-            residual=None,
-            eps=self.variance_epsilon,
-            dropout_p=0.0,
-            prenorm=False,
-            residual_in_fp32=False,
-        )
+__all__ = [
+    "layers",
+    "rmsnorm_forward",
+]

torch-ext/rmsnorm_kernel/layers.py

@@ -0,0 +1,15 @@
+import torch
+import torch.nn as nn
+
+from ._ops import ops
+
+
+class LlamaRMSNorm(nn.Module):
+    weight: torch.Tensor
+    variance_epsilon: float
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return ops.rmsnorm_forward(
+            hidden_states,
+            self.weight,
+        )