Add built binary [skip-build]

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

@@ -1,6 +1,6 @@
 import torch
 
-from . import layers
+from . import layers, parallel_style
 from ._ops import ops
 from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
 from .rms_norm import FusedAddRMSNormFunction, RMSNormFunction
@@ -48,5 +48,6 @@ __all__ = [
     "rms_norm",
     "fused_add_rms_norm",
     "layers",
+    "parallel_style",
     "ops",
 ]

build/torch27-cxx11-cu118-x86_64-linux/activation/{_activation_20250907180255.abi3.so → _activation_53ed492_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d21a85bf21aa74f1281541e658acfd4f4326d902efe3578b059eccf054443284
-size 8089696
+oid sha256:80267a0391fa4cb22aa3eb04b05d8214c2bfaed968b714185bc20214596072e3
+size 8618232

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

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

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

@@ -0,0 +1,199 @@
+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.utils import (generate_redistribute_costs,
+                                                 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
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@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(), ),
+            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
+    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(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[
+                    output_grad_tgt if need_input_grad else None, weight_tgt
+                ],
+                input_specs=[
+                    output_grad_tgt, add_output_grad_tgt, add_output_tgt,
+                    weight_tgt
+                ],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

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

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from torch.distributed.tensor import (DeviceMesh, DTensor, Replicate, Shard,
+                                      distribute_module, distribute_tensor)
+from torch.distributed.tensor.parallel import SequenceParallel
+from torch.distributed.tensor.placement_types import Placement
+
+
+class ResidualSequenceParallel(SequenceParallel):
+    """ Consider the case where we have a residual connection across a sequence parallel layer."""
+
+    @staticmethod
+    def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh):
+        input_tensor = inputs[0]
+        residual_tensor = inputs[1]
+
+        assert isinstance(input_tensor,
+                          DTensor) == isinstance(residual_tensor, DTensor)
+        assert isinstance(input_tensor,
+                          torch.Tensor) == isinstance(residual_tensor,
+                                                      torch.Tensor)
+
+        if isinstance(input_tensor, DTensor):
+            # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it
+            if input_tensor.placements != sequence_sharding:
+                input_tensor = input_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            if residual_tensor.placements != sequence_sharding:
+                residual_tensor = residual_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            return input_tensor, residual_tensor
+
+        elif isinstance(input_tensor, torch.Tensor):
+            # assume the input passed in already sharded on the sequence dim and create the DTensor
+            return DTensor.from_local(input_tensor,
+                                      device_mesh,
+                                      sequence_sharding,
+                                      run_check=False), DTensor.from_local(
+                                          residual_tensor,
+                                          device_mesh,
+                                          sequence_sharding,
+                                          run_check=False)
+        else:
+            raise ValueError(
+                f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}"
+            )

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

@@ -1,4 +1,7 @@
+from collections.abc import Sequence
+
 import torch
+from packaging import version
 
 from ._ops import ops
 
@@ -8,9 +11,7 @@ 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
+        return ops.rms_norm(input, weight, eps)
 
     @staticmethod
     # inputs is a Tuple of all of the inputs passed to forward.
@@ -26,13 +27,8 @@ class RMSNormFunction(torch.autograd.Function):
         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)
+        input_grad, weight_grad = ops.rms_norm_backward(
+            output_grad, input, weight, eps)
 
         return input_grad, weight_grad, None
 
@@ -42,10 +38,8 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
     # Note that forward, setup_context, and backward are @staticmethods
     @staticmethod
     def forward(input, residual, weight, eps):
-        output = torch.empty_like(input)
-        add_output = torch.empty_like(input)
-        ops.fused_add_rms_norm(output, add_output, input, residual, weight,
-                               eps)
+        output, add_output = ops.fused_add_rms_norm(input, residual, weight,
+                                                    eps)
         return output, add_output
 
     @staticmethod
@@ -65,14 +59,47 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
         need_in = ctx.needs_input_grad[0]
         need_res = ctx.needs_input_grad[1]
 
-        grad = torch.empty_like(output_grad) if need_in or need_res else None
+        # TODO(ai-system): kernels currently do not support no input gradients
+        assert need_in or need_res, "Not implemented for no input gradients yet"
 
-        weight_grad = torch.empty_like(
-            weight) if ctx.needs_input_grad[2] else None
-
-        ops.fused_add_rms_norm_backward(grad, weight_grad, output_grad, add_output_grad, add_output,
-                              weight, eps)
+        grad, weight_grad = ops.fused_add_rms_norm_backward(
+            output_grad,
+            add_output_grad,
+            add_output,
+            weight,
+            eps,
+            need_input_grad=need_in or need_res)
         input_grad = grad if need_in else None
         residual_grad = grad if need_res else None
 
         return input_grad, residual_grad, weight_grad, None
+
+
+@torch.library.register_fake(ops.rms_norm.default)
+def rms_norm_abstract(x, weight, eps):
+    return torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.rms_norm_backward.default)
+def rms_norm_backward_abstract(output_grad, x, weight, eps):
+    return torch.empty_like(x), torch.empty_like(weight)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm.default)
+def fused_add_rms_norm_abstract(x, residual, weight, eps):
+    return torch.empty_like(x), torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm_backward.default)
+def fused_add_rms_norm_backward_abstract(output_grad, add_output_grad,
+                                         add_output, weight, eps,
+                                         need_input_grad: bool):
+    return torch.empty_like(
+        output_grad) if need_input_grad else None, torch.empty_like(weight)
+
+
+if version.parse(torch.__version__) >= version.parse("2.8"):
+    from .fused_add_rms_norm_meta import register_fused_add_rms_norm_meta
+    from .rms_norm_meta import register_rms_norm_meta
+    register_fused_add_rms_norm_meta()
+    register_rms_norm_meta()

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

@@ -0,0 +1,164 @@
+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.utils import (generate_redistribute_costs,
+                                                 register_op_strategy)
+from torch.distributed.tensor.placement_types import (Placement, Replicate,
+                                                      Shard)
+
+from ._ops import ops
+
+
+def register_rms_norm_meta():
+    """Dummy function to register the meta functions.
+    Registration happens at import time by the decorators below.
+    """
+    pass
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@register_op_strategy(ops.rms_norm.default, schema_info=RuntimeSchemaInfo(1))
+def rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 3
+    (
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    assert len(input_strategy.strategies) == len(weight_strategy.strategies)
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for input, weight in zip(input_strategy.strategies,
+                             weight_strategy.strategies):
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(input_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))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            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_specs=[input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy
+
+
+@register_op_strategy(ops.rms_norm_backward.default,
+                      schema_info=RuntimeSchemaInfo(1))
+def rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 4
+    (
+        output_grad_strategy,
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(output_grad_strategy, OpStrategy)
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    lengths = {
+        "output_grad": len(output_grad_strategy.strategies),
+        "input": len(input_strategy.strategies),
+        "weight": len(weight_strategy.strategies),
+    }
+
+    assert len(set(
+        lengths.values())) == 1, f"Strategies length mismatch {lengths}"
+
+    zipped = zip(
+        output_grad_strategy.strategies,
+        input_strategy.strategies,
+        weight_strategy.strategies,
+    )
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for output_grad, input, weight in zipped:
+        output_grad_src = output_grad.output_spec
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(output_grad_src, DTensorSpec)
+        assert isinstance(input_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))
+
+        # Input must have the same sharding as output grad.
+        input_tgt = output_grad_tgt
+        redistribute_costs.append(
+            generate_redistribute_costs(input_strategy, input_tgt))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[input_tgt, weight_tgt],
+                input_specs=[output_grad_tgt, input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

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

@@ -1,6 +1,6 @@
 import torch
 
-from . import layers
+from . import layers, parallel_style
 from ._ops import ops
 from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
 from .rms_norm import FusedAddRMSNormFunction, RMSNormFunction
@@ -48,5 +48,6 @@ __all__ = [
     "rms_norm",
     "fused_add_rms_norm",
     "layers",
+    "parallel_style",
     "ops",
 ]

build/{torch27-cxx11-cu118-x86_64-linux/activation/_activation_e5e2eeb_dirty.abi3.so → torch27-cxx11-cu126-x86_64-linux/activation/_activation_53ed492_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ec9ea7edc8b27f7983e20d615ab470cef6b82975afc214becfddfd05a867a839
-size 8600336
+oid sha256:ef6e4eb51daac20f0d7ed9825052ecca9d8451825784c87d58fa69092c145f35
+size 8793008

build/torch27-cxx11-cu126-x86_64-linux/activation/_activation_e5e2eeb_dirty.abi3.so

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

build/torch27-cxx11-cu126-x86_64-linux/activation/_activation_f517c97_dirty.abi3.so

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

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

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

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

@@ -0,0 +1,199 @@
+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.utils import (generate_redistribute_costs,
+                                                 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
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@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(), ),
+            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
+    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(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[
+                    output_grad_tgt if need_input_grad else None, weight_tgt
+                ],
+                input_specs=[
+                    output_grad_tgt, add_output_grad_tgt, add_output_tgt,
+                    weight_tgt
+                ],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

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

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from torch.distributed.tensor import (DeviceMesh, DTensor, Replicate, Shard,
+                                      distribute_module, distribute_tensor)
+from torch.distributed.tensor.parallel import SequenceParallel
+from torch.distributed.tensor.placement_types import Placement
+
+
+class ResidualSequenceParallel(SequenceParallel):
+    """ Consider the case where we have a residual connection across a sequence parallel layer."""
+
+    @staticmethod
+    def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh):
+        input_tensor = inputs[0]
+        residual_tensor = inputs[1]
+
+        assert isinstance(input_tensor,
+                          DTensor) == isinstance(residual_tensor, DTensor)
+        assert isinstance(input_tensor,
+                          torch.Tensor) == isinstance(residual_tensor,
+                                                      torch.Tensor)
+
+        if isinstance(input_tensor, DTensor):
+            # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it
+            if input_tensor.placements != sequence_sharding:
+                input_tensor = input_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            if residual_tensor.placements != sequence_sharding:
+                residual_tensor = residual_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            return input_tensor, residual_tensor
+
+        elif isinstance(input_tensor, torch.Tensor):
+            # assume the input passed in already sharded on the sequence dim and create the DTensor
+            return DTensor.from_local(input_tensor,
+                                      device_mesh,
+                                      sequence_sharding,
+                                      run_check=False), DTensor.from_local(
+                                          residual_tensor,
+                                          device_mesh,
+                                          sequence_sharding,
+                                          run_check=False)
+        else:
+            raise ValueError(
+                f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}"
+            )

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

@@ -1,4 +1,7 @@
+from collections.abc import Sequence
+
 import torch
+from packaging import version
 
 from ._ops import ops
 
@@ -8,9 +11,7 @@ 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
+        return ops.rms_norm(input, weight, eps)
 
     @staticmethod
     # inputs is a Tuple of all of the inputs passed to forward.
@@ -26,13 +27,8 @@ class RMSNormFunction(torch.autograd.Function):
         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)
+        input_grad, weight_grad = ops.rms_norm_backward(
+            output_grad, input, weight, eps)
 
         return input_grad, weight_grad, None
 
@@ -42,10 +38,8 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
     # Note that forward, setup_context, and backward are @staticmethods
     @staticmethod
     def forward(input, residual, weight, eps):
-        output = torch.empty_like(input)
-        add_output = torch.empty_like(input)
-        ops.fused_add_rms_norm(output, add_output, input, residual, weight,
-                               eps)
+        output, add_output = ops.fused_add_rms_norm(input, residual, weight,
+                                                    eps)
         return output, add_output
 
     @staticmethod
@@ -65,14 +59,47 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
         need_in = ctx.needs_input_grad[0]
         need_res = ctx.needs_input_grad[1]
 
-        grad = torch.empty_like(output_grad) if need_in or need_res else None
+        # TODO(ai-system): kernels currently do not support no input gradients
+        assert need_in or need_res, "Not implemented for no input gradients yet"
 
-        weight_grad = torch.empty_like(
-            weight) if ctx.needs_input_grad[2] else None
-
-        ops.fused_add_rms_norm_backward(grad, weight_grad, output_grad, add_output_grad, add_output,
-                              weight, eps)
+        grad, weight_grad = ops.fused_add_rms_norm_backward(
+            output_grad,
+            add_output_grad,
+            add_output,
+            weight,
+            eps,
+            need_input_grad=need_in or need_res)
         input_grad = grad if need_in else None
         residual_grad = grad if need_res else None
 
         return input_grad, residual_grad, weight_grad, None
+
+
+@torch.library.register_fake(ops.rms_norm.default)
+def rms_norm_abstract(x, weight, eps):
+    return torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.rms_norm_backward.default)
+def rms_norm_backward_abstract(output_grad, x, weight, eps):
+    return torch.empty_like(x), torch.empty_like(weight)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm.default)
+def fused_add_rms_norm_abstract(x, residual, weight, eps):
+    return torch.empty_like(x), torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm_backward.default)
+def fused_add_rms_norm_backward_abstract(output_grad, add_output_grad,
+                                         add_output, weight, eps,
+                                         need_input_grad: bool):
+    return torch.empty_like(
+        output_grad) if need_input_grad else None, torch.empty_like(weight)
+
+
+if version.parse(torch.__version__) >= version.parse("2.8"):
+    from .fused_add_rms_norm_meta import register_fused_add_rms_norm_meta
+    from .rms_norm_meta import register_rms_norm_meta
+    register_fused_add_rms_norm_meta()
+    register_rms_norm_meta()

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

@@ -0,0 +1,164 @@
+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.utils import (generate_redistribute_costs,
+                                                 register_op_strategy)
+from torch.distributed.tensor.placement_types import (Placement, Replicate,
+                                                      Shard)
+
+from ._ops import ops
+
+
+def register_rms_norm_meta():
+    """Dummy function to register the meta functions.
+    Registration happens at import time by the decorators below.
+    """
+    pass
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@register_op_strategy(ops.rms_norm.default, schema_info=RuntimeSchemaInfo(1))
+def rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 3
+    (
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    assert len(input_strategy.strategies) == len(weight_strategy.strategies)
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for input, weight in zip(input_strategy.strategies,
+                             weight_strategy.strategies):
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(input_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))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            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_specs=[input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy
+
+
+@register_op_strategy(ops.rms_norm_backward.default,
+                      schema_info=RuntimeSchemaInfo(1))
+def rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 4
+    (
+        output_grad_strategy,
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(output_grad_strategy, OpStrategy)
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    lengths = {
+        "output_grad": len(output_grad_strategy.strategies),
+        "input": len(input_strategy.strategies),
+        "weight": len(weight_strategy.strategies),
+    }
+
+    assert len(set(
+        lengths.values())) == 1, f"Strategies length mismatch {lengths}"
+
+    zipped = zip(
+        output_grad_strategy.strategies,
+        input_strategy.strategies,
+        weight_strategy.strategies,
+    )
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for output_grad, input, weight in zipped:
+        output_grad_src = output_grad.output_spec
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(output_grad_src, DTensorSpec)
+        assert isinstance(input_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))
+
+        # Input must have the same sharding as output grad.
+        input_tgt = output_grad_tgt
+        redistribute_costs.append(
+            generate_redistribute_costs(input_strategy, input_tgt))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[input_tgt, weight_tgt],
+                input_specs=[output_grad_tgt, input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

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

@@ -1,6 +1,6 @@
 import torch
 
-from . import layers
+from . import layers, parallel_style
 from ._ops import ops
 from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
 from .rms_norm import FusedAddRMSNormFunction, RMSNormFunction
@@ -48,5 +48,6 @@ __all__ = [
     "rms_norm",
     "fused_add_rms_norm",
     "layers",
+    "parallel_style",
     "ops",
 ]

build/torch27-cxx11-cu128-x86_64-linux/activation/_activation_20250907180255.abi3.so

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

build/torch27-cxx11-cu128-x86_64-linux/activation/_activation_53ed492_dirty.abi3.so

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

build/torch27-cxx11-cu128-x86_64-linux/activation/_activation_e5e2eeb_dirty.abi3.so

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

build/torch27-cxx11-cu128-x86_64-linux/activation/_activation_f517c97_dirty.abi3.so

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

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

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

build/torch27-cxx11-cu128-x86_64-linux/activation/fused_add_rms_norm_meta.py

@@ -0,0 +1,199 @@
+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.utils import (generate_redistribute_costs,
+                                                 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
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@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(), ),
+            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
+    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(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[
+                    output_grad_tgt if need_input_grad else None, weight_tgt
+                ],
+                input_specs=[
+                    output_grad_tgt, add_output_grad_tgt, add_output_tgt,
+                    weight_tgt
+                ],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

build/torch27-cxx11-cu128-x86_64-linux/activation/parallel_style.py

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from torch.distributed.tensor import (DeviceMesh, DTensor, Replicate, Shard,
+                                      distribute_module, distribute_tensor)
+from torch.distributed.tensor.parallel import SequenceParallel
+from torch.distributed.tensor.placement_types import Placement
+
+
+class ResidualSequenceParallel(SequenceParallel):
+    """ Consider the case where we have a residual connection across a sequence parallel layer."""
+
+    @staticmethod
+    def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh):
+        input_tensor = inputs[0]
+        residual_tensor = inputs[1]
+
+        assert isinstance(input_tensor,
+                          DTensor) == isinstance(residual_tensor, DTensor)
+        assert isinstance(input_tensor,
+                          torch.Tensor) == isinstance(residual_tensor,
+                                                      torch.Tensor)
+
+        if isinstance(input_tensor, DTensor):
+            # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it
+            if input_tensor.placements != sequence_sharding:
+                input_tensor = input_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            if residual_tensor.placements != sequence_sharding:
+                residual_tensor = residual_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            return input_tensor, residual_tensor
+
+        elif isinstance(input_tensor, torch.Tensor):
+            # assume the input passed in already sharded on the sequence dim and create the DTensor
+            return DTensor.from_local(input_tensor,
+                                      device_mesh,
+                                      sequence_sharding,
+                                      run_check=False), DTensor.from_local(
+                                          residual_tensor,
+                                          device_mesh,
+                                          sequence_sharding,
+                                          run_check=False)
+        else:
+            raise ValueError(
+                f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}"
+            )

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

@@ -1,4 +1,7 @@
+from collections.abc import Sequence
+
 import torch
+from packaging import version
 
 from ._ops import ops
 
@@ -8,9 +11,7 @@ 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
+        return ops.rms_norm(input, weight, eps)
 
     @staticmethod
     # inputs is a Tuple of all of the inputs passed to forward.
@@ -26,13 +27,8 @@ class RMSNormFunction(torch.autograd.Function):
         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)
+        input_grad, weight_grad = ops.rms_norm_backward(
+            output_grad, input, weight, eps)
 
         return input_grad, weight_grad, None
 
@@ -42,10 +38,8 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
     # Note that forward, setup_context, and backward are @staticmethods
     @staticmethod
     def forward(input, residual, weight, eps):
-        output = torch.empty_like(input)
-        add_output = torch.empty_like(input)
-        ops.fused_add_rms_norm(output, add_output, input, residual, weight,
-                               eps)
+        output, add_output = ops.fused_add_rms_norm(input, residual, weight,
+                                                    eps)
         return output, add_output
 
     @staticmethod
@@ -65,14 +59,47 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
         need_in = ctx.needs_input_grad[0]
         need_res = ctx.needs_input_grad[1]
 
-        grad = torch.empty_like(output_grad) if need_in or need_res else None
+        # TODO(ai-system): kernels currently do not support no input gradients
+        assert need_in or need_res, "Not implemented for no input gradients yet"
 
-        weight_grad = torch.empty_like(
-            weight) if ctx.needs_input_grad[2] else None
-
-        ops.fused_add_rms_norm_backward(grad, weight_grad, output_grad, add_output_grad, add_output,
-                              weight, eps)
+        grad, weight_grad = ops.fused_add_rms_norm_backward(
+            output_grad,
+            add_output_grad,
+            add_output,
+            weight,
+            eps,
+            need_input_grad=need_in or need_res)
         input_grad = grad if need_in else None
         residual_grad = grad if need_res else None
 
         return input_grad, residual_grad, weight_grad, None
+
+
+@torch.library.register_fake(ops.rms_norm.default)
+def rms_norm_abstract(x, weight, eps):
+    return torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.rms_norm_backward.default)
+def rms_norm_backward_abstract(output_grad, x, weight, eps):
+    return torch.empty_like(x), torch.empty_like(weight)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm.default)
+def fused_add_rms_norm_abstract(x, residual, weight, eps):
+    return torch.empty_like(x), torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm_backward.default)
+def fused_add_rms_norm_backward_abstract(output_grad, add_output_grad,
+                                         add_output, weight, eps,
+                                         need_input_grad: bool):
+    return torch.empty_like(
+        output_grad) if need_input_grad else None, torch.empty_like(weight)
+
+
+if version.parse(torch.__version__) >= version.parse("2.8"):
+    from .fused_add_rms_norm_meta import register_fused_add_rms_norm_meta
+    from .rms_norm_meta import register_rms_norm_meta
+    register_fused_add_rms_norm_meta()
+    register_rms_norm_meta()

build/torch27-cxx11-cu128-x86_64-linux/activation/rms_norm_meta.py

@@ -0,0 +1,164 @@
+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.utils import (generate_redistribute_costs,
+                                                 register_op_strategy)
+from torch.distributed.tensor.placement_types import (Placement, Replicate,
+                                                      Shard)
+
+from ._ops import ops
+
+
+def register_rms_norm_meta():
+    """Dummy function to register the meta functions.
+    Registration happens at import time by the decorators below.
+    """
+    pass
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@register_op_strategy(ops.rms_norm.default, schema_info=RuntimeSchemaInfo(1))
+def rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 3
+    (
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    assert len(input_strategy.strategies) == len(weight_strategy.strategies)
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for input, weight in zip(input_strategy.strategies,
+                             weight_strategy.strategies):
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(input_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))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            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_specs=[input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy
+
+
+@register_op_strategy(ops.rms_norm_backward.default,
+                      schema_info=RuntimeSchemaInfo(1))
+def rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 4
+    (
+        output_grad_strategy,
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(output_grad_strategy, OpStrategy)
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    lengths = {
+        "output_grad": len(output_grad_strategy.strategies),
+        "input": len(input_strategy.strategies),
+        "weight": len(weight_strategy.strategies),
+    }
+
+    assert len(set(
+        lengths.values())) == 1, f"Strategies length mismatch {lengths}"
+
+    zipped = zip(
+        output_grad_strategy.strategies,
+        input_strategy.strategies,
+        weight_strategy.strategies,
+    )
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for output_grad, input, weight in zipped:
+        output_grad_src = output_grad.output_spec
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(output_grad_src, DTensorSpec)
+        assert isinstance(input_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))
+
+        # Input must have the same sharding as output grad.
+        input_tgt = output_grad_tgt
+        redistribute_costs.append(
+            generate_redistribute_costs(input_strategy, input_tgt))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[input_tgt, weight_tgt],
+                input_specs=[output_grad_tgt, input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

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

@@ -1,6 +1,6 @@
 import torch
 
-from . import layers
+from . import layers, parallel_style
 from ._ops import ops
 from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
 from .rms_norm import FusedAddRMSNormFunction, RMSNormFunction
@@ -48,5 +48,6 @@ __all__ = [
     "rms_norm",
     "fused_add_rms_norm",
     "layers",
+    "parallel_style",
     "ops",
 ]

build/torch27-cxx11-rocm63-x86_64-linux/activation/_activation_20250907180255.abi3.so

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

build/{torch27-cxx11-cu118-x86_64-linux/activation/_activation_f517c97_dirty.abi3.so → torch27-cxx11-rocm63-x86_64-linux/activation/_activation_53ed492_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bd84c828d4c15e96d65d6c8f0eb7a945ee8167d92e978b2ebce03eeaf41e7fce
-size 4405112
+oid sha256:d973bad96565705f9e27514a9dbfb37343d0220da4a3ae7156b1cf6a27813643
+size 2773952

build/torch27-cxx11-rocm63-x86_64-linux/activation/_activation_e5e2eeb_dirty.abi3.so

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

build/torch27-cxx11-rocm63-x86_64-linux/activation/_activation_f517c97_dirty.abi3.so

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

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

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

build/torch27-cxx11-rocm63-x86_64-linux/activation/fused_add_rms_norm_meta.py

@@ -0,0 +1,199 @@
+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.utils import (generate_redistribute_costs,
+                                                 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
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@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(), ),
+            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
+    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(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[
+                    output_grad_tgt if need_input_grad else None, weight_tgt
+                ],
+                input_specs=[
+                    output_grad_tgt, add_output_grad_tgt, add_output_tgt,
+                    weight_tgt
+                ],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

build/torch27-cxx11-rocm63-x86_64-linux/activation/parallel_style.py

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from torch.distributed.tensor import (DeviceMesh, DTensor, Replicate, Shard,
+                                      distribute_module, distribute_tensor)
+from torch.distributed.tensor.parallel import SequenceParallel
+from torch.distributed.tensor.placement_types import Placement
+
+
+class ResidualSequenceParallel(SequenceParallel):
+    """ Consider the case where we have a residual connection across a sequence parallel layer."""
+
+    @staticmethod
+    def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh):
+        input_tensor = inputs[0]
+        residual_tensor = inputs[1]
+
+        assert isinstance(input_tensor,
+                          DTensor) == isinstance(residual_tensor, DTensor)
+        assert isinstance(input_tensor,
+                          torch.Tensor) == isinstance(residual_tensor,
+                                                      torch.Tensor)
+
+        if isinstance(input_tensor, DTensor):
+            # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it
+            if input_tensor.placements != sequence_sharding:
+                input_tensor = input_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            if residual_tensor.placements != sequence_sharding:
+                residual_tensor = residual_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            return input_tensor, residual_tensor
+
+        elif isinstance(input_tensor, torch.Tensor):
+            # assume the input passed in already sharded on the sequence dim and create the DTensor
+            return DTensor.from_local(input_tensor,
+                                      device_mesh,
+                                      sequence_sharding,
+                                      run_check=False), DTensor.from_local(
+                                          residual_tensor,
+                                          device_mesh,
+                                          sequence_sharding,
+                                          run_check=False)
+        else:
+            raise ValueError(
+                f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}"
+            )

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

@@ -1,4 +1,7 @@
+from collections.abc import Sequence
+
 import torch
+from packaging import version
 
 from ._ops import ops
 
@@ -8,9 +11,7 @@ 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
+        return ops.rms_norm(input, weight, eps)
 
     @staticmethod
     # inputs is a Tuple of all of the inputs passed to forward.
@@ -26,13 +27,8 @@ class RMSNormFunction(torch.autograd.Function):
         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)
+        input_grad, weight_grad = ops.rms_norm_backward(
+            output_grad, input, weight, eps)
 
         return input_grad, weight_grad, None
 
@@ -42,10 +38,8 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
     # Note that forward, setup_context, and backward are @staticmethods
     @staticmethod
     def forward(input, residual, weight, eps):
-        output = torch.empty_like(input)
-        add_output = torch.empty_like(input)
-        ops.fused_add_rms_norm(output, add_output, input, residual, weight,
-                               eps)
+        output, add_output = ops.fused_add_rms_norm(input, residual, weight,
+                                                    eps)
         return output, add_output
 
     @staticmethod
@@ -65,14 +59,47 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
         need_in = ctx.needs_input_grad[0]
         need_res = ctx.needs_input_grad[1]
 
-        grad = torch.empty_like(output_grad) if need_in or need_res else None
+        # TODO(ai-system): kernels currently do not support no input gradients
+        assert need_in or need_res, "Not implemented for no input gradients yet"
 
-        weight_grad = torch.empty_like(
-            weight) if ctx.needs_input_grad[2] else None
-
-        ops.fused_add_rms_norm_backward(grad, weight_grad, output_grad, add_output_grad, add_output,
-                              weight, eps)
+        grad, weight_grad = ops.fused_add_rms_norm_backward(
+            output_grad,
+            add_output_grad,
+            add_output,
+            weight,
+            eps,
+            need_input_grad=need_in or need_res)
         input_grad = grad if need_in else None
         residual_grad = grad if need_res else None
 
         return input_grad, residual_grad, weight_grad, None
+
+
+@torch.library.register_fake(ops.rms_norm.default)
+def rms_norm_abstract(x, weight, eps):
+    return torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.rms_norm_backward.default)
+def rms_norm_backward_abstract(output_grad, x, weight, eps):
+    return torch.empty_like(x), torch.empty_like(weight)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm.default)
+def fused_add_rms_norm_abstract(x, residual, weight, eps):
+    return torch.empty_like(x), torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm_backward.default)
+def fused_add_rms_norm_backward_abstract(output_grad, add_output_grad,
+                                         add_output, weight, eps,
+                                         need_input_grad: bool):
+    return torch.empty_like(
+        output_grad) if need_input_grad else None, torch.empty_like(weight)
+
+
+if version.parse(torch.__version__) >= version.parse("2.8"):
+    from .fused_add_rms_norm_meta import register_fused_add_rms_norm_meta
+    from .rms_norm_meta import register_rms_norm_meta
+    register_fused_add_rms_norm_meta()
+    register_rms_norm_meta()

build/torch27-cxx11-rocm63-x86_64-linux/activation/rms_norm_meta.py

@@ -0,0 +1,164 @@
+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.utils import (generate_redistribute_costs,
+                                                 register_op_strategy)
+from torch.distributed.tensor.placement_types import (Placement, Replicate,
+                                                      Shard)
+
+from ._ops import ops
+
+
+def register_rms_norm_meta():
+    """Dummy function to register the meta functions.
+    Registration happens at import time by the decorators below.
+    """
+    pass
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@register_op_strategy(ops.rms_norm.default, schema_info=RuntimeSchemaInfo(1))
+def rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 3
+    (
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    assert len(input_strategy.strategies) == len(weight_strategy.strategies)
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for input, weight in zip(input_strategy.strategies,
+                             weight_strategy.strategies):
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(input_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))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            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_specs=[input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy
+
+
+@register_op_strategy(ops.rms_norm_backward.default,
+                      schema_info=RuntimeSchemaInfo(1))
+def rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 4
+    (
+        output_grad_strategy,
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(output_grad_strategy, OpStrategy)
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    lengths = {
+        "output_grad": len(output_grad_strategy.strategies),
+        "input": len(input_strategy.strategies),
+        "weight": len(weight_strategy.strategies),
+    }
+
+    assert len(set(
+        lengths.values())) == 1, f"Strategies length mismatch {lengths}"
+
+    zipped = zip(
+        output_grad_strategy.strategies,
+        input_strategy.strategies,
+        weight_strategy.strategies,
+    )
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for output_grad, input, weight in zipped:
+        output_grad_src = output_grad.output_spec
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(output_grad_src, DTensorSpec)
+        assert isinstance(input_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))
+
+        # Input must have the same sharding as output grad.
+        input_tgt = output_grad_tgt
+        redistribute_costs.append(
+            generate_redistribute_costs(input_strategy, input_tgt))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[input_tgt, weight_tgt],
+                input_specs=[output_grad_tgt, input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

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

@@ -1,6 +1,6 @@
 import torch
 
-from . import layers
+from . import layers, parallel_style
 from ._ops import ops
 from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
 from .rms_norm import FusedAddRMSNormFunction, RMSNormFunction
@@ -48,5 +48,6 @@ __all__ = [
     "rms_norm",
     "fused_add_rms_norm",
     "layers",
+    "parallel_style",
     "ops",
 ]

build/torch28-cxx11-cu126-x86_64-linux/activation/_activation_20250907180255.abi3.so

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

build/{torch27-cxx11-cu126-x86_64-linux/activation/_activation_20250907180255.abi3.so → torch28-cxx11-cu126-x86_64-linux/activation/_activation_53ed492_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:74d4955271509451b946495da75f69a0f978e7258b8303fe3c077e585c0d3e6a
-size 8272456
+oid sha256:c301db3d37625ebf0cecf016948ec18fbeddb497acca8c870d2d8eff0a1d1203
+size 8735952

build/torch28-cxx11-cu126-x86_64-linux/activation/_activation_e5e2eeb_dirty.abi3.so

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

build/torch28-cxx11-cu126-x86_64-linux/activation/_activation_f517c97_dirty.abi3.so

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

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

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

build/torch28-cxx11-cu126-x86_64-linux/activation/fused_add_rms_norm_meta.py

@@ -0,0 +1,199 @@
+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.utils import (generate_redistribute_costs,
+                                                 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
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@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(), ),
+            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
+    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(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[
+                    output_grad_tgt if need_input_grad else None, weight_tgt
+                ],
+                input_specs=[
+                    output_grad_tgt, add_output_grad_tgt, add_output_tgt,
+                    weight_tgt
+                ],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

build/torch28-cxx11-cu126-x86_64-linux/activation/parallel_style.py

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from torch.distributed.tensor import (DeviceMesh, DTensor, Replicate, Shard,
+                                      distribute_module, distribute_tensor)
+from torch.distributed.tensor.parallel import SequenceParallel
+from torch.distributed.tensor.placement_types import Placement
+
+
+class ResidualSequenceParallel(SequenceParallel):
+    """ Consider the case where we have a residual connection across a sequence parallel layer."""
+
+    @staticmethod
+    def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh):
+        input_tensor = inputs[0]
+        residual_tensor = inputs[1]
+
+        assert isinstance(input_tensor,
+                          DTensor) == isinstance(residual_tensor, DTensor)
+        assert isinstance(input_tensor,
+                          torch.Tensor) == isinstance(residual_tensor,
+                                                      torch.Tensor)
+
+        if isinstance(input_tensor, DTensor):
+            # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it
+            if input_tensor.placements != sequence_sharding:
+                input_tensor = input_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            if residual_tensor.placements != sequence_sharding:
+                residual_tensor = residual_tensor.redistribute(
+                    placements=sequence_sharding, async_op=True)
+            return input_tensor, residual_tensor
+
+        elif isinstance(input_tensor, torch.Tensor):
+            # assume the input passed in already sharded on the sequence dim and create the DTensor
+            return DTensor.from_local(input_tensor,
+                                      device_mesh,
+                                      sequence_sharding,
+                                      run_check=False), DTensor.from_local(
+                                          residual_tensor,
+                                          device_mesh,
+                                          sequence_sharding,
+                                          run_check=False)
+        else:
+            raise ValueError(
+                f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}"
+            )

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

@@ -1,4 +1,7 @@
+from collections.abc import Sequence
+
 import torch
+from packaging import version
 
 from ._ops import ops
 
@@ -8,9 +11,7 @@ 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
+        return ops.rms_norm(input, weight, eps)
 
     @staticmethod
     # inputs is a Tuple of all of the inputs passed to forward.
@@ -26,13 +27,8 @@ class RMSNormFunction(torch.autograd.Function):
         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)
+        input_grad, weight_grad = ops.rms_norm_backward(
+            output_grad, input, weight, eps)
 
         return input_grad, weight_grad, None
 
@@ -42,10 +38,8 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
     # Note that forward, setup_context, and backward are @staticmethods
     @staticmethod
     def forward(input, residual, weight, eps):
-        output = torch.empty_like(input)
-        add_output = torch.empty_like(input)
-        ops.fused_add_rms_norm(output, add_output, input, residual, weight,
-                               eps)
+        output, add_output = ops.fused_add_rms_norm(input, residual, weight,
+                                                    eps)
         return output, add_output
 
     @staticmethod
@@ -65,14 +59,47 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
         need_in = ctx.needs_input_grad[0]
         need_res = ctx.needs_input_grad[1]
 
-        grad = torch.empty_like(output_grad) if need_in or need_res else None
+        # TODO(ai-system): kernels currently do not support no input gradients
+        assert need_in or need_res, "Not implemented for no input gradients yet"
 
-        weight_grad = torch.empty_like(
-            weight) if ctx.needs_input_grad[2] else None
-
-        ops.fused_add_rms_norm_backward(grad, weight_grad, output_grad, add_output_grad, add_output,
-                              weight, eps)
+        grad, weight_grad = ops.fused_add_rms_norm_backward(
+            output_grad,
+            add_output_grad,
+            add_output,
+            weight,
+            eps,
+            need_input_grad=need_in or need_res)
         input_grad = grad if need_in else None
         residual_grad = grad if need_res else None
 
         return input_grad, residual_grad, weight_grad, None
+
+
+@torch.library.register_fake(ops.rms_norm.default)
+def rms_norm_abstract(x, weight, eps):
+    return torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.rms_norm_backward.default)
+def rms_norm_backward_abstract(output_grad, x, weight, eps):
+    return torch.empty_like(x), torch.empty_like(weight)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm.default)
+def fused_add_rms_norm_abstract(x, residual, weight, eps):
+    return torch.empty_like(x), torch.empty_like(x)
+
+
+@torch.library.register_fake(ops.fused_add_rms_norm_backward.default)
+def fused_add_rms_norm_backward_abstract(output_grad, add_output_grad,
+                                         add_output, weight, eps,
+                                         need_input_grad: bool):
+    return torch.empty_like(
+        output_grad) if need_input_grad else None, torch.empty_like(weight)
+
+
+if version.parse(torch.__version__) >= version.parse("2.8"):
+    from .fused_add_rms_norm_meta import register_fused_add_rms_norm_meta
+    from .rms_norm_meta import register_rms_norm_meta
+    register_fused_add_rms_norm_meta()
+    register_rms_norm_meta()

build/torch28-cxx11-cu126-x86_64-linux/activation/rms_norm_meta.py

@@ -0,0 +1,164 @@
+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.utils import (generate_redistribute_costs,
+                                                 register_op_strategy)
+from torch.distributed.tensor.placement_types import (Placement, Replicate,
+                                                      Shard)
+
+from ._ops import ops
+
+
+def register_rms_norm_meta():
+    """Dummy function to register the meta functions.
+    Registration happens at import time by the decorators below.
+    """
+    pass
+
+
+def _replicate_dims_start_at(placements: Sequence[Placement],
+                             start_dim: int = 0) -> tuple[Placement, ...]:
+    new_placements: list[Placement] = []
+    for p in placements:
+        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
+            new_placements.append(Replicate())  # make it replicate
+        else:
+            new_placements.append(p)  # keep the placement
+    return tuple(new_placements)
+
+
+@register_op_strategy(ops.rms_norm.default, schema_info=RuntimeSchemaInfo(1))
+def rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 3
+    (
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    assert len(input_strategy.strategies) == len(weight_strategy.strategies)
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for input, weight in zip(input_strategy.strategies,
+                             weight_strategy.strategies):
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(input_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))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            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_specs=[input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy
+
+
+@register_op_strategy(ops.rms_norm_backward.default,
+                      schema_info=RuntimeSchemaInfo(1))
+def rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
+    mesh = op_schema.get_mesh_from_args()
+
+    assert len(op_schema.args_schema) == 4
+    (
+        output_grad_strategy,
+        input_strategy,
+        weight_strategy,
+        _,  # eps
+    ) = op_schema.args_schema
+
+    assert isinstance(output_grad_strategy, OpStrategy)
+    assert isinstance(input_strategy, OpStrategy)
+    assert isinstance(weight_strategy, OpStrategy)
+
+    lengths = {
+        "output_grad": len(output_grad_strategy.strategies),
+        "input": len(input_strategy.strategies),
+        "weight": len(weight_strategy.strategies),
+    }
+
+    assert len(set(
+        lengths.values())) == 1, f"Strategies length mismatch {lengths}"
+
+    zipped = zip(
+        output_grad_strategy.strategies,
+        input_strategy.strategies,
+        weight_strategy.strategies,
+    )
+
+    last_dim = input_strategy.ndim - 1
+    strategy = OpStrategy([])
+    for output_grad, input, weight in zipped:
+        output_grad_src = output_grad.output_spec
+        input_src = input.output_spec
+        weight_src = weight.output_spec
+
+        assert isinstance(output_grad_src, DTensorSpec)
+        assert isinstance(input_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))
+
+        # Input must have the same sharding as output grad.
+        input_tgt = output_grad_tgt
+        redistribute_costs.append(
+            generate_redistribute_costs(input_strategy, input_tgt))
+
+        # Weight cannot be sharded, so always replicate it.
+        weight_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=(Replicate(), ),
+            tensor_meta=weight_src.tensor_meta,
+        )
+        redistribute_costs.append(
+            generate_redistribute_costs(weight_strategy, weight_tgt))
+
+        strategy.strategies.append(
+            OpSpec(
+                output_specs=[input_tgt, weight_tgt],
+                input_specs=[output_grad_tgt, input_tgt, weight_tgt],
+                redistribute_cost=redistribute_costs,
+            ))
+    return strategy

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

@@ -1,6 +1,6 @@
 import torch
 
-from . import layers
+from . import layers, parallel_style
 from ._ops import ops
 from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
 from .rms_norm import FusedAddRMSNormFunction, RMSNormFunction
@@ -48,5 +48,6 @@ __all__ = [
     "rms_norm",
     "fused_add_rms_norm",
     "layers",
+    "parallel_style",
     "ops",
 ]

build/torch28-cxx11-cu128-x86_64-linux/activation/_activation_20250907180255.abi3.so

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

build/torch28-cxx11-cu128-x86_64-linux/activation/_activation_53ed492_dirty.abi3.so

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

build/torch28-cxx11-cu128-x86_64-linux/activation/_activation_e5e2eeb_dirty.abi3.so

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