First builds

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text

build/torch27-cxx11-cu118-x86_64-linux/rmsnorm_kernel/__init__.py

@@ -0,0 +1,21 @@
+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,
+            bias=None,
+            residual=None,
+            eps=self.variance_epsilon,
+            dropout_p=0.0,
+            prenorm=False,
+            residual_in_fp32=False,
+        )

build/torch27-cxx11-cu118-x86_64-linux/rmsnorm_kernel/_ops.py

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

build/torch27-cxx11-cu118-x86_64-linux/rmsnorm_kernel/_rmsnorm_kernel_538355f_dirty.abi3.so

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

build/torch27-cxx11-cu126-x86_64-linux/rmsnorm_kernel/__init__.py

@@ -0,0 +1,21 @@
+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,
+            bias=None,
+            residual=None,
+            eps=self.variance_epsilon,
+            dropout_p=0.0,
+            prenorm=False,
+            residual_in_fp32=False,
+        )

build/torch27-cxx11-cu126-x86_64-linux/rmsnorm_kernel/_ops.py

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

build/torch27-cxx11-cu126-x86_64-linux/rmsnorm_kernel/_rmsnorm_kernel_538355f_dirty.abi3.so

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

build/torch27-cxx11-cu128-x86_64-linux/rmsnorm_kernel/__init__.py

@@ -0,0 +1,21 @@
+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,
+            bias=None,
+            residual=None,
+            eps=self.variance_epsilon,
+            dropout_p=0.0,
+            prenorm=False,
+            residual_in_fp32=False,
+        )

build/torch27-cxx11-cu128-x86_64-linux/rmsnorm_kernel/_ops.py

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

build/torch27-cxx11-cu128-x86_64-linux/rmsnorm_kernel/_rmsnorm_kernel_538355f_dirty.abi3.so

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

build/torch28-cxx11-cu126-x86_64-linux/rmsnorm_kernel/__init__.py

@@ -0,0 +1,21 @@
+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,
+            bias=None,
+            residual=None,
+            eps=self.variance_epsilon,
+            dropout_p=0.0,
+            prenorm=False,
+            residual_in_fp32=False,
+        )

build/torch28-cxx11-cu126-x86_64-linux/rmsnorm_kernel/_ops.py

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

build/torch28-cxx11-cu126-x86_64-linux/rmsnorm_kernel/_rmsnorm_kernel_538355f_dirty.abi3.so

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

build/torch28-cxx11-cu128-x86_64-linux/rmsnorm_kernel/__init__.py

@@ -0,0 +1,21 @@
+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,
+            bias=None,
+            residual=None,
+            eps=self.variance_epsilon,
+            dropout_p=0.0,
+            prenorm=False,
+            residual_in_fp32=False,
+        )

build/torch28-cxx11-cu128-x86_64-linux/rmsnorm_kernel/_ops.py

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

build/torch28-cxx11-cu128-x86_64-linux/rmsnorm_kernel/_rmsnorm_kernel_538355f_dirty.abi3.so

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

build/torch28-cxx11-cu129-x86_64-linux/rmsnorm_kernel/__init__.py

@@ -0,0 +1,21 @@
+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,
+            bias=None,
+            residual=None,
+            eps=self.variance_epsilon,
+            dropout_p=0.0,
+            prenorm=False,
+            residual_in_fp32=False,
+        )

build/torch28-cxx11-cu129-x86_64-linux/rmsnorm_kernel/_ops.py

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

build/torch28-cxx11-cu129-x86_64-linux/rmsnorm_kernel/_rmsnorm_kernel_538355f_dirty.abi3.so

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

rmsnorm_kernel/rmsnorm.cu

@@ -65,7 +65,7 @@ struct RmsnormFunctor {
     }
 };
 
-torch::Tensor rmsnorm_forward(torch::Tensor x, torch::Tensor gamma) {
+torch::Tensor rmsnorm_forward(torch::Tensor &x, torch::Tensor &gamma) {
     int B = x.size(0), S = x.size(1), H = x.size(2);
     int rows = B * S;
 
@@ -160,4 +160,4 @@ torch::Tensor rmsnorm_forward(torch::Tensor x, torch::Tensor gamma) {
 //     std::cout << (max_diff < 1e-5 ? "PASSED: Results match!" : "FAILED: Results don't match!") << std::endl;
 
 //     return 0;
-// }
+// }

torch-ext/torch_binding.h

@@ -2,4 +2,4 @@
 
 #include <torch/torch.h>
 
-torch::Tensor rmsnorm_forward(torch::Tensor const &input, torch::Tensor const &gamma);
+torch::Tensor rmsnorm_forward(torch::Tensor &input, torch::Tensor &gamma);