Add built binary [skip-build]

build/torch210-cxx11-cu126-x86_64-linux/__init__.py

@@ -0,0 +1,53 @@
+import torch
+
+from . import layers, parallel_style
+from ._ops import ops
+from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
+from .rms_norm import FusedAddRMSNormFunction, 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 fused_mul_poly_norm(
+    x: torch.Tensor,
+    mul: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedMulPolyNormFunction.apply(x, mul, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+def fused_add_rms_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedAddRMSNormFunction.apply(x, residual, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "fused_mul_poly_norm",
+    "rms_norm",
+    "fused_add_rms_norm",
+    "layers",
+    "parallel_style",
+    "ops",
+]

build/{torch28-cxx11-cu128-x86_64-linux/activation/_activation_496308d_dirty.abi3.so → torch210-cxx11-cu126-x86_64-linux/_activation_18b7543_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3e914bbcc41277358db1417cd12162ec7d4d0d63f597c137301f1c57df2a896a
-size 10122960
+oid sha256:39a7e25002120a73ea83ac813276c0518086fae2236f528dadf96bac4876a270
+size 10775296

build/{torch28-cxx11-cu128-x86_64-linux/activation → torch210-cxx11-cu126-x86_64-linux}/_ops.py

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

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

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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/{torch28-cxx11-cu128-x86_64-linux/activation → torch210-cxx11-cu126-x86_64-linux}/fused_add_rms_norm_meta.py

@@ -7,8 +7,13 @@ from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
 from torch.distributed.tensor._ops._math_ops import (
     _infer_reduce_dims_map, _replicate_dims_start_at,
     map_placements_after_reduction)
-from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
-                                                 register_op_strategy)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
                                                       Shard)
 

build/torch210-cxx11-cu126-x86_64-linux/metadata.json

@@ -0,0 +1,3 @@
+{
+  "python-depends": []
+}

build/{torch28-cxx11-cu129-x86_64-linux/activation → torch210-cxx11-cu126-x86_64-linux}/rms_norm_meta.py

@@ -7,8 +7,13 @@ from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
 from torch.distributed.tensor._ops._math_ops import (
     _infer_reduce_dims_map, _replicate_dims_start_at,
     map_placements_after_reduction)
-from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
-                                                 register_op_strategy)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
                                                       Shard)
 

build/torch210-cxx11-cu128-x86_64-linux/__init__.py

@@ -0,0 +1,53 @@
+import torch
+
+from . import layers, parallel_style
+from ._ops import ops
+from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
+from .rms_norm import FusedAddRMSNormFunction, 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 fused_mul_poly_norm(
+    x: torch.Tensor,
+    mul: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedMulPolyNormFunction.apply(x, mul, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+def fused_add_rms_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedAddRMSNormFunction.apply(x, residual, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "fused_mul_poly_norm",
+    "rms_norm",
+    "fused_add_rms_norm",
+    "layers",
+    "parallel_style",
+    "ops",
+]

build/{torch28-cxx11-cu129-x86_64-linux/activation/_activation_496308d_dirty.abi3.so → torch210-cxx11-cu128-x86_64-linux/_activation_18b7543_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5de606ea13ce268356eeae89b662f2b62a09acd94fe1bd66abea832bdd6d7b9e
-size 10147000
+oid sha256:078853c2db399a227822ea0c8e70c2e13bad41bfa370657dd19aa2efb3b503e9
+size 15815392

build/{torch28-cxx11-cu126-x86_64-linux/activation → torch210-cxx11-cu128-x86_64-linux}/_ops.py

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

build/torch210-cxx11-cu128-x86_64-linux/activation/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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/{torch28-cxx11-cu129-x86_64-linux/activation → torch210-cxx11-cu128-x86_64-linux}/fused_add_rms_norm_meta.py

@@ -7,8 +7,13 @@ from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
 from torch.distributed.tensor._ops._math_ops import (
     _infer_reduce_dims_map, _replicate_dims_start_at,
     map_placements_after_reduction)
-from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
-                                                 register_op_strategy)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
                                                       Shard)
 

build/torch210-cxx11-cu128-x86_64-linux/metadata.json

@@ -0,0 +1,3 @@
+{
+  "python-depends": []
+}

build/{torch28-cxx11-cu128-x86_64-linux/activation → torch210-cxx11-cu128-x86_64-linux}/rms_norm_meta.py

@@ -7,8 +7,13 @@ from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
 from torch.distributed.tensor._ops._math_ops import (
     _infer_reduce_dims_map, _replicate_dims_start_at,
     map_placements_after_reduction)
-from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
-                                                 register_op_strategy)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
                                                       Shard)
 

build/torch210-cxx11-cu130-x86_64-linux/__init__.py

@@ -0,0 +1,53 @@
+import torch
+
+from . import layers, parallel_style
+from ._ops import ops
+from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
+from .rms_norm import FusedAddRMSNormFunction, 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 fused_mul_poly_norm(
+    x: torch.Tensor,
+    mul: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedMulPolyNormFunction.apply(x, mul, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+def fused_add_rms_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedAddRMSNormFunction.apply(x, residual, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "fused_mul_poly_norm",
+    "rms_norm",
+    "fused_add_rms_norm",
+    "layers",
+    "parallel_style",
+    "ops",
+]

build/torch210-cxx11-cu130-x86_64-linux/_activation_18b7543_dirty.abi3.so

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

build/{torch28-cxx11-cu129-x86_64-linux/activation → torch210-cxx11-cu130-x86_64-linux}/_ops.py

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

build/torch210-cxx11-cu130-x86_64-linux/activation/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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/{torch28-cxx11-cu126-x86_64-linux/activation → torch210-cxx11-cu130-x86_64-linux}/fused_add_rms_norm_meta.py

@@ -7,8 +7,13 @@ from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
 from torch.distributed.tensor._ops._math_ops import (
     _infer_reduce_dims_map, _replicate_dims_start_at,
     map_placements_after_reduction)
-from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
-                                                 register_op_strategy)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
                                                       Shard)
 

build/torch210-cxx11-cu130-x86_64-linux/metadata.json

@@ -0,0 +1,3 @@
+{
+  "python-depends": []
+}

build/{torch28-cxx11-rocm63-x86_64-linux/activation → torch210-cxx11-cu130-x86_64-linux}/rms_norm_meta.py

@@ -7,8 +7,13 @@ from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
 from torch.distributed.tensor._ops._math_ops import (
     _infer_reduce_dims_map, _replicate_dims_start_at,
     map_placements_after_reduction)
-from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
-                                                 register_op_strategy)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
                                                       Shard)
 

build/torch210-cxx11-rocm70-x86_64-linux/__init__.py

@@ -0,0 +1,53 @@
+import torch
+
+from . import layers, parallel_style
+from ._ops import ops
+from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
+from .rms_norm import FusedAddRMSNormFunction, 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 fused_mul_poly_norm(
+    x: torch.Tensor,
+    mul: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedMulPolyNormFunction.apply(x, mul, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+def fused_add_rms_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedAddRMSNormFunction.apply(x, residual, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "fused_mul_poly_norm",
+    "rms_norm",
+    "fused_add_rms_norm",
+    "layers",
+    "parallel_style",
+    "ops",
+]

build/{torch28-cxx11-cu126-x86_64-linux/activation/_activation_496308d_dirty.abi3.so → torch210-cxx11-rocm70-x86_64-linux/_activation_18b7543_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6981f362de22807e15574e81e19b4816915a0e7a2863818ee6729641e7d28568
-size 6742080
+oid sha256:45ff2b71abb33d840d92116980e519786ed06f1e337d681d0e3301dba241ff63
+size 2919488

build/{torch28-cxx11-rocm63-x86_64-linux/activation → torch210-cxx11-rocm70-x86_64-linux}/_ops.py

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

build/torch210-cxx11-rocm70-x86_64-linux/activation/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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/{torch28-cxx11-rocm63-x86_64-linux/activation → torch210-cxx11-rocm70-x86_64-linux}/fused_add_rms_norm_meta.py

@@ -7,8 +7,13 @@ from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
 from torch.distributed.tensor._ops._math_ops import (
     _infer_reduce_dims_map, _replicate_dims_start_at,
     map_placements_after_reduction)
-from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
-                                                 register_op_strategy)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
                                                       Shard)
 

build/torch210-cxx11-rocm70-x86_64-linux/metadata.json

@@ -0,0 +1,3 @@
+{
+  "python-depends": []
+}

build/{torch28-cxx11-cu126-x86_64-linux/activation → torch210-cxx11-rocm70-x86_64-linux}/rms_norm_meta.py

@@ -7,8 +7,13 @@ from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
 from torch.distributed.tensor._ops._math_ops import (
     _infer_reduce_dims_map, _replicate_dims_start_at,
     map_placements_after_reduction)
-from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
-                                                 register_op_strategy)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
                                                       Shard)
 

build/torch210-cxx11-rocm71-x86_64-linux/__init__.py

@@ -0,0 +1,53 @@
+import torch
+
+from . import layers, parallel_style
+from ._ops import ops
+from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
+from .rms_norm import FusedAddRMSNormFunction, 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 fused_mul_poly_norm(
+    x: torch.Tensor,
+    mul: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedMulPolyNormFunction.apply(x, mul, weight, bias, eps)
+
+
+def rms_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return RMSNormFunction.apply(x, weight, eps)
+
+
+def fused_add_rms_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return FusedAddRMSNormFunction.apply(x, residual, weight, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "fused_mul_poly_norm",
+    "rms_norm",
+    "fused_add_rms_norm",
+    "layers",
+    "parallel_style",
+    "ops",
+]

build/{torch28-cxx11-rocm64-x86_64-linux/activation/_activation_496308d_dirty.abi3.so → torch210-cxx11-rocm71-x86_64-linux/_activation_18b7543_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2cfc63ad98abf44fcb2a491d41fe99ea98be7a5c8f59dbacaedbaa32b3700313
-size 2794144
+oid sha256:af4db38e8d5ad56226f5a95a86c2b5fc726bd9d576d07df2f07d3f03c1b6b35b
+size 2911200

build/torch210-cxx11-rocm71-x86_64-linux/_ops.py

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

build/torch210-cxx11-rocm71-x86_64-linux/activation/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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/torch210-cxx11-rocm71-x86_64-linux/fused_add_rms_norm_meta.py

@@ -0,0 +1,217 @@
+from collections.abc import Sequence
+
+import torch
+from torch.distributed.tensor._dtensor_spec import DTensorSpec
+from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
+                                                 RuntimeSchemaInfo)
+from torch.distributed.tensor._ops._math_ops import (
+    _infer_reduce_dims_map, _replicate_dims_start_at,
+    map_placements_after_reduction)
+from torch.distributed.tensor._ops.utils import generate_redistribute_costs
+
+try:
+    from torch.distributed.tensor._ops.utils import register_op_strategy
+except ImportError:
+    # PyTorch 2.10+ moved register_op_strategy to a separate module
+    from torch.distributed.tensor._ops.registration import register_op_strategy
+from torch.distributed.tensor.placement_types import (Placement, Replicate,
+                                                      Shard)
+
+from ._ops import ops
+
+
+def register_fused_add_rms_norm_meta():
+    """Dummy function to register the meta functions.
+    Registration happens at import time by the decorators below.
+    """
+    pass
+
+
+@register_op_strategy(ops.fused_add_rms_norm.default,
+                      schema_info=RuntimeSchemaInfo(1))
+def fused_add_rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 4
+    (
+        input_strategy,
+        residual_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(residual_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    lengths = {
+        "input": len(input_strategy.strategies),
+        "residual": len(residual_strategy.strategies),
+        "weight": len(weight_strategy.strategies),
+    }
+    assert len(set(
+        lengths.values())) == 1, f"Strategy length mismatch: {lengths}"
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for input, residual, weight in zip(input_strategy.strategies,
+                                       residual_strategy.strategies,
+                                       weight_strategy.strategies):
+
+        input_src = input.output_spec
+        residual_src = residual.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(input_src, DTensorSpec)
+        assert isinstance(residual_src, DTensorSpec)
+        assert isinstance(weight_src, DTensorSpec)
+
+        redistribute_costs = []
+
+        # Input can be sharded in any dim except the last dim.
+        input_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=_replicate_dims_start_at(input_src.placements,
+                                                last_dim),
+            tensor_meta=input_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(input_strategy, input_tgt))
+
+        # Residual add must have the same sharding as input.
+        residual_tgt = input_tgt
+        redistribute_costs.append(
+            generate_redistribute_costs(residual_strategy, residual_tgt))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=_replicate_dims_start_at(weight_src.placements),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[input_tgt, input_tgt],
+                input_specs=[input_tgt, residual_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy
+
+
+@register_op_strategy(ops.fused_add_rms_norm_backward.default,
+                      schema_info=RuntimeSchemaInfo(2))
+def fused_add_rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 6
+    (
+        output_grad_strategy,
+        add_output_grad_strategy,
+        add_output_strategy,
+        weight_strategy,
+        _,  # eps
+        need_input_grad,  # need_input_grad
+    ) = op_schema.args_schema
+
+    assert isinstance(output_grad_strategy, OpStrategy)
+    assert isinstance(add_output_grad_strategy, OpStrategy)
+    assert isinstance(add_output_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    lengths = {
+        "output_grad": len(output_grad_strategy.strategies),
+        "add_output_grad": len(add_output_grad_strategy.strategies),
+        "add_output": len(add_output_strategy.strategies),
+        "weight": len(weight_strategy.strategies),
+    }
+    assert len(set(
+        lengths.values())) == 1, f"Strategy length mismatch: {lengths}"
+
+    zipped = zip(
+        output_grad_strategy.strategies,
+        add_output_grad_strategy.strategies,
+        add_output_strategy.strategies,
+        weight_strategy.strategies,
+    )
+
+    last_dim = output_grad_strategy.ndim - 1
+    outer_dims = list(range(last_dim))
+
+    strategy = OpStrategy([])
+    for output_grad, add_output_grad, add_output, weight in zipped:
+        output_grad_src = output_grad.output_spec
+        add_output_grad_src = add_output_grad.output_spec
+        add_output_src = add_output.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(output_grad_src, DTensorSpec)
+        assert isinstance(add_output_grad_src, DTensorSpec)
+        assert isinstance(add_output_src, DTensorSpec)
+        assert isinstance(weight_src, DTensorSpec)
+
+        redistribute_costs = []
+
+        # output grad can be sharded in any dim except the last dim.
+        output_grad_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=_replicate_dims_start_at(output_grad_src.placements,
+                                                last_dim),
+            tensor_meta=output_grad_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(output_grad_strategy, output_grad_tgt))
+
+        # add_output_grad must have the same sharding as output_grad.
+        add_output_grad_tgt = output_grad_tgt
+        redistribute_costs.append(
+            generate_redistribute_costs(add_output_grad_strategy,
+                                        add_output_grad_tgt))
+
+        # add_output must have the same sharding as output_grad.
+        add_output_tgt = output_grad_tgt
+        redistribute_costs.append(
+            generate_redistribute_costs(add_output_strategy, add_output_tgt))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=_replicate_dims_start_at(weight_src.placements),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        # from torch/distributed/tensor/_ops/_math_ops.py::layer_norm_bwd_strategy()
+
+        # Weight cannot be sharded, so always replicate it.
+        # TODO: now d_weight spec follows input spec w/ a reduction.
+        # we may need to change to a pointwise rule over grad_out and
+        # input, then apply a reduction.
+        inp_placements = _replicate_dims_start_at(output_grad_src.placements,
+                                                  last_dim)
+        reduce_dims_map = _infer_reduce_dims_map(outer_dims,
+                                                 output_grad_src.ndim, False)
+        out_placements = map_placements_after_reduction(
+            inp_placements, outer_dims, reduce_dims_map, "sum")
+        weight_grad_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=out_placements,
+            tensor_meta=weight_src.tensor_meta,
+        )
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[
+                    output_grad_tgt if need_input_grad else None,
+                    weight_grad_tgt
+                ],
+                input_specs=[
+                    output_grad_tgt, add_output_grad_tgt, add_output_tgt,
+                    weight_tgt
+                ],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy