Make `RMSNorm` layer pure, rename to `LlamaRMSNorm`

build/torch-universal/triton_layer_norm/layers.py

@@ -1,4 +1,24 @@
-from .layer_norm import RMSNorm
+import torch
+from torch import nn
 
+from .layer_norm import rms_norm_fn
 
-__all__ = ["RMSNorm"]
+
+class LlamaRMSNorm(nn.Module):
+    weight: torch.Tensor
+    variance_epsilon: float
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return rms_norm_fn(
+            hidden_states,
+            self.weight,
+            bias=None,
+            residual=None,
+            eps=self.variance_epsilon,
+            dropout_p=0.0,
+            prenorm=False,
+            residual_in_fp32=False,
+        )
+
+
+__all__ = ["LlamaRMSNorm"]

torch-ext/triton_layer_norm/layers.py

@@ -1,4 +1,24 @@
-from .layer_norm import RMSNorm
+import torch
+from torch import nn
 
+from .layer_norm import rms_norm_fn
 
-__all__ = ["RMSNorm"]
+
+class LlamaRMSNorm(nn.Module):
+    weight: torch.Tensor
+    variance_epsilon: float
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return rms_norm_fn(
+            hidden_states,
+            self.weight,
+            bias=None,
+            residual=None,
+            eps=self.variance_epsilon,
+            dropout_p=0.0,
+            prenorm=False,
+            residual_in_fp32=False,
+        )
+
+
+__all__ = ["LlamaRMSNorm"]