[mxfp8 moe training] register sharding rules for cutedsl 2d quant kernel (#4178)

Author: danielvegamyhre

torchao/prototype/moe_training/kernels/mxfp8/quant.py

@@ -8,6 +8,8 @@
 
 import torch
 from torch import Tensor
+from torch.distributed.tensor import Replicate, Shard
+from torch.distributed.tensor.experimental import register_sharding
 from torch.utils._triton import has_triton
 
 from torchao.prototype.moe_training.kernels.mxfp8.cute_utils import (
@@ -975,6 +977,20 @@ def _fake_mxfp8_quantize_2d_cutedsl_custom_op(
     return q_data, scales
 
 
+if _mxfp8_cutedsl_kernels_available:
+
+    @register_sharding(torch.ops.torchao.mxfp8_quantize_2d_cutedsl.default)
+    def custom_sharding_for_cutedsl_mxfp8_dim0_kernel(
+        x, block_size=32, scaling_mode: str = "rceil", stage_count: int = 2
+    ):
+        # order is: ([outputs, ...], [inputs, ...])
+        replicate = ([Replicate(), Replicate()], [Replicate(), None, None, None])
+        shard_dim0 = ([Shard(0), Shard(0)], [Shard(0), None, None, None])
+        shard_dim1 = ([Shard(1), Shard(1)], [Shard(1), None, None, None])
+        acceptable_shardings = [replicate, shard_dim0, shard_dim1]
+        return acceptable_shardings
+
+
 if _mxfp8_cuda_kernels_available:
     # CUDA kernel for per group blocked layout transform with groups along M
     def mx_block_rearrange_2d_M_groups_cuda(