Update build

build/torch-universal/triton_layer_norm/__init__.py

@@ -1,4 +1,4 @@
-"""Triton layer normalization kernels.
+"""Triton layer normalization kernels
 
 This kernel implements layers normalization using Triton. This kernel is from
 the `flash-attention <https://github.com/Dao-AILab/flash-attention>`_ project.

build/torch-universal/triton_layer_norm/_ops.py

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

build/torch-universal/triton_layer_norm/layers.py

@@ -5,10 +5,32 @@ from .layer_norm import rms_norm_fn
 
 
 class LlamaRMSNorm(nn.Module):
+    """
+    RMS Layer Norm for Llama models.
+
+    Triton-optimized RMS layer norm. The interface is compatible with `LLamaRMSNorm` in
+    `transformers`.
+
+    Attributes:
+        weight (`torch.Tensor`): The learnable scaling parameter.
+        variance_epsilon (`float`): The epsilon value for numerical stability.
+    """
     weight: torch.Tensor
     variance_epsilon: float
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Apply RMS normalization to the input hidden states.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                Input tensor of shape `(batch_size, sequence_length, hidden_size)` or any shape
+                where the last dimension is the feature dimension to be normalized.
+
+        Returns:
+            `torch.Tensor`:
+                The normalized tensor with the same shape as the input `hidden_states`.
+        """
         return rms_norm_fn(
             hidden_states,
             self.weight,