Uploaded using `kernel-builder`.

Files changed (14) hide show

build/torch210-metal-aarch64-darwin/__init__.py ADDED Viewed

+from .functions import rmsnorm_forward, rmsnorm_backward
+from . import layers
+__all__ = ["layers", "rmsnorm_forward", "rmsnorm_backward"]

build/torch210-metal-aarch64-darwin/_mlx_rmsnorm_metal_b150ff3.abi3.so ADDED Viewed

+version https://git-lfs.github.com/spec/v1
+oid sha256:081e37329b9871e3547287660c1e3fa1d300f94cf92b1a86e9d4919645d7578d
+size 221840

build/torch210-metal-aarch64-darwin/_ops.py ADDED Viewed

+import torch
+from . import _mlx_rmsnorm_metal_b150ff3
+ops = torch.ops._mlx_rmsnorm_metal_b150ff3
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mlx_rmsnorm_metal_b150ff3::{op_name}"

build/torch210-metal-aarch64-darwin/functions.py ADDED Viewed

+import torch
+from ._ops import ops
+def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    original_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    weight = weight.view(-1)
+    output = torch.zeros_like(x)
+    ops.launch_forward_kernel(x, weight, output, epsilon)
+    output = output.view(original_shape)
+    return output
+def rmsnorm_backward(x: torch.Tensor, weight: torch.Tensor, grad_output: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    original_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    weight = weight.view(-1)
+    grad_output = grad_output.view(-1)
+    grad_input = torch.zeros_like(x)
+    grad_weight = torch.zeros_like(weight)
+    ops.launch_backward_kernel(x, weight, grad_output, grad_input, grad_weight, epsilon)
+    grad_input = grad_input.view(original_shape)
+    grad_weight = grad_weight.view(original_shape)
+    return grad_input, grad_weight

build/torch210-metal-aarch64-darwin/layers.py ADDED Viewed

+from .functions import rmsnorm_forward
+import torch
+class RMSNorm(torch.nn.Module):
+    weight: torch.Tensor
+    variance_epsilon: float
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return rmsnorm_forward(x, self.weight, self.variance_epsilon)

build/torch210-metal-aarch64-darwin/metadata.json ADDED Viewed

+{
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "metal"
+  }
+}

build/torch210-metal-aarch64-darwin/mlx_rmsnorm/__init__.py ADDED Viewed

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-metal-aarch64-darwin/__init__.py ADDED Viewed

+from .functions import rmsnorm_forward, rmsnorm_backward
+from . import layers
+__all__ = ["layers", "rmsnorm_forward", "rmsnorm_backward"]

build/torch211-metal-aarch64-darwin/_mlx_rmsnorm_metal_b150ff3.abi3.so ADDED Viewed

+version https://git-lfs.github.com/spec/v1
+oid sha256:e58acba77e25fca83fd4792ab7418780850d7c35354f74ef2cc1f0bcd16af13e
+size 221840

build/torch211-metal-aarch64-darwin/_ops.py ADDED Viewed

+import torch
+from . import _mlx_rmsnorm_metal_b150ff3
+ops = torch.ops._mlx_rmsnorm_metal_b150ff3
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mlx_rmsnorm_metal_b150ff3::{op_name}"

build/torch211-metal-aarch64-darwin/functions.py ADDED Viewed

+import torch
+from ._ops import ops
+def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    original_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    weight = weight.view(-1)
+    output = torch.zeros_like(x)
+    ops.launch_forward_kernel(x, weight, output, epsilon)
+    output = output.view(original_shape)
+    return output
+def rmsnorm_backward(x: torch.Tensor, weight: torch.Tensor, grad_output: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    original_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    weight = weight.view(-1)
+    grad_output = grad_output.view(-1)
+    grad_input = torch.zeros_like(x)
+    grad_weight = torch.zeros_like(weight)
+    ops.launch_backward_kernel(x, weight, grad_output, grad_input, grad_weight, epsilon)
+    grad_input = grad_input.view(original_shape)
+    grad_weight = grad_weight.view(original_shape)
+    return grad_input, grad_weight

build/torch211-metal-aarch64-darwin/layers.py ADDED Viewed

+from .functions import rmsnorm_forward
+import torch
+class RMSNorm(torch.nn.Module):
+    weight: torch.Tensor
+    variance_epsilon: float
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return rmsnorm_forward(x, self.weight, self.variance_epsilon)

build/torch211-metal-aarch64-darwin/metadata.json ADDED Viewed

+{
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "metal"
+  }
+}

build/torch211-metal-aarch64-darwin/mlx_rmsnorm/__init__.py ADDED Viewed

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))