feat: add cuda build

build.toml

@@ -10,7 +10,7 @@ src = [
 
 [kernel.activation]
 backend = "rocm"
-rocm-archs = [ "gfx90a" ]
+rocm-archs = [ "gfx90a", "gfx942" ]
 src = [
   "activation/poly_norm.cu",
   "activation/rms_norm.cu",
@@ -21,3 +21,17 @@ src = [
   "activation/atomic_utils.h",
 ]
 depends = [ "torch" ]
+
+[kernel.activation_cuda]
+backend = "cuda"
+src = [
+  "activation/poly_norm.cu",
+  "activation/rms_norm.cu",
+  "activation/cuda_compat.h",
+  "activation/block_reduce.h",
+  "activation/dispatch_utils.h",
+  "activation/assert_utils.h",
+  "activation/atomic_utils.h",
+]
+depends = ["torch"]
+

build/torch26-cxx11-cu118-x86_64-linux/activation/__init__.py

@@ -0,0 +1,30 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/{torch26-cxx11-rocm62-x86_64-linux/activation/_activation_f3b99fb_dirty.abi3.so → torch26-cxx11-cu118-x86_64-linux/activation/_activation_605f22e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7bf9f4d85d15bc4869292e6a293ec53b7658cee61284457ea727c4be435062f7
-size 2656296
+oid sha256:b2fdb7378a1c907c3ff3ad0a5134a0a8ce4a464196404436470d7b4eb77ec305
+size 2957296

build/torch26-cxx11-cu118-x86_64-linux/activation/_ops.py

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

build/torch26-cxx11-cu118-x86_64-linux/activation/layers.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3, dtype=dtype) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

build/torch26-cxx11-cu118-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None

build/torch26-cxx11-cu118-x86_64-linux/activation/rms_norm.py

@@ -0,0 +1,34 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class RMSNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, eps):
+        output = torch.empty_like(input)
+        ops.rms_norm(output, input, weight, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+
+        ops.rms_norm_backward(input_grad, weight_grad, output_grad, input, weight, eps)
+
+        return input_grad, weight_grad, None

build/torch26-cxx11-cu124-x86_64-linux/activation/__init__.py

@@ -0,0 +1,30 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/{torch27-cxx11-rocm63-x86_64-linux/activation/_activation_f3b99fb_dirty.abi3.so → torch26-cxx11-cu124-x86_64-linux/activation/_activation_605f22e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ad289cf495aa7bcb7318535f2d76a6543bd44827369ec358ff7411e182ce089f
-size 2642736
+oid sha256:5baac6228e04fbb209cbc90a24702c14f4eb52d2698cea12a766d77412622096
+size 2981880

build/torch26-cxx11-cu124-x86_64-linux/activation/_ops.py

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

build/torch26-cxx11-cu124-x86_64-linux/activation/layers.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3, dtype=dtype) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

build/torch26-cxx11-cu124-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None

build/torch26-cxx11-cu124-x86_64-linux/activation/rms_norm.py

@@ -0,0 +1,34 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class RMSNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, eps):
+        output = torch.empty_like(input)
+        ops.rms_norm(output, input, weight, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+
+        ops.rms_norm_backward(input_grad, weight_grad, output_grad, input, weight, eps)
+
+        return input_grad, weight_grad, None

build/torch26-cxx11-cu126-x86_64-linux/activation/__init__.py

@@ -0,0 +1,30 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/torch26-cxx11-cu126-x86_64-linux/activation/_activation_605f22e_dirty.abi3.so

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

build/torch26-cxx11-cu126-x86_64-linux/activation/_ops.py

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

build/torch26-cxx11-cu126-x86_64-linux/activation/layers.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3, dtype=dtype) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

build/torch26-cxx11-cu126-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None

build/torch26-cxx11-cu126-x86_64-linux/activation/rms_norm.py

@@ -0,0 +1,34 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class RMSNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, eps):
+        output = torch.empty_like(input)
+        ops.rms_norm(output, input, weight, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+
+        ops.rms_norm_backward(input_grad, weight_grad, output_grad, input, weight, eps)
+
+        return input_grad, weight_grad, None

build/torch26-cxx11-rocm62-x86_64-linux/activation/_activation_605f22e_dirty.abi3.so

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

build/torch26-cxx11-rocm62-x86_64-linux/activation/_ops.py

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

build/torch26-cxx98-cu118-x86_64-linux/activation/__init__.py

@@ -0,0 +1,30 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/torch26-cxx98-cu118-x86_64-linux/activation/_activation_605f22e_dirty.abi3.so

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

build/torch26-cxx98-cu118-x86_64-linux/activation/_ops.py

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

build/torch26-cxx98-cu118-x86_64-linux/activation/layers.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3, dtype=dtype) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

build/torch26-cxx98-cu118-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None

build/torch26-cxx98-cu118-x86_64-linux/activation/rms_norm.py

@@ -0,0 +1,34 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class RMSNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, eps):
+        output = torch.empty_like(input)
+        ops.rms_norm(output, input, weight, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+
+        ops.rms_norm_backward(input_grad, weight_grad, output_grad, input, weight, eps)
+
+        return input_grad, weight_grad, None

build/torch26-cxx98-cu124-x86_64-linux/activation/__init__.py

@@ -0,0 +1,30 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/torch26-cxx98-cu124-x86_64-linux/activation/_activation_605f22e_dirty.abi3.so

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

build/torch26-cxx98-cu124-x86_64-linux/activation/_ops.py

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

build/torch26-cxx98-cu124-x86_64-linux/activation/layers.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3, dtype=dtype) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

build/torch26-cxx98-cu124-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None

build/torch26-cxx98-cu124-x86_64-linux/activation/rms_norm.py

@@ -0,0 +1,34 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class RMSNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, eps):
+        output = torch.empty_like(input)
+        ops.rms_norm(output, input, weight, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+
+        ops.rms_norm_backward(input_grad, weight_grad, output_grad, input, weight, eps)
+
+        return input_grad, weight_grad, None

build/torch26-cxx98-cu126-x86_64-linux/activation/__init__.py

@@ -0,0 +1,30 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/torch26-cxx98-cu126-x86_64-linux/activation/_activation_605f22e_dirty.abi3.so

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

build/torch26-cxx98-cu126-x86_64-linux/activation/_ops.py

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

build/torch26-cxx98-cu126-x86_64-linux/activation/layers.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3, dtype=dtype) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

build/torch26-cxx98-cu126-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None

build/torch26-cxx98-cu126-x86_64-linux/activation/rms_norm.py

@@ -0,0 +1,34 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class RMSNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, eps):
+        output = torch.empty_like(input)
+        ops.rms_norm(output, input, weight, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+
+        ops.rms_norm_backward(input_grad, weight_grad, output_grad, input, weight, eps)
+
+        return input_grad, weight_grad, None

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

@@ -0,0 +1,30 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/torch27-cxx11-cu118-x86_64-linux/activation/_activation_605f22e_dirty.abi3.so

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

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

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

build/torch27-cxx11-cu118-x86_64-linux/activation/layers.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3, dtype=dtype) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim, dtype=dtype))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

build/torch27-cxx11-cu118-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None

build/torch27-cxx11-cu118-x86_64-linux/activation/rms_norm.py

@@ -0,0 +1,34 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class RMSNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, eps):
+        output = torch.empty_like(input)
+        ops.rms_norm(output, input, weight, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+
+        ops.rms_norm_backward(input_grad, weight_grad, output_grad, input, weight, eps)
+
+        return input_grad, weight_grad, None

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

@@ -0,0 +1,30 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+from .rms_norm import RMSNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]