[mxfp8 moe training] add compile support

Author: danielvegamyhre

test/prototype/moe_training/test_training.py

@@ -34,10 +34,7 @@
 
 @pytest.mark.parametrize(
     "target_fqns",
-    [
-        ["experts"],
-        ["does.not.exist"],
-    ],
+    [["experts"], ["experts,shared_expert"], ["invalid.fqns"]],
 )
 @pytest.mark.parametrize("compile", [False, True])
 @pytest.mark.parametrize(

torchao/prototype/moe_training/kernels/mxfp8_blocked_scales.py

@@ -4,6 +4,7 @@
 import triton
 import triton.language as tl
 from torch import Tensor
+from torch.library import triton_op, wrap_triton
 
 from torchao.prototype.mx_formats.utils import to_blocked
 from torchao.utils import ceil_div
@@ -192,6 +193,7 @@ def compute_blocked_scale_offsets_for_K_groups(
     return group_sizes, starting_col_after_padding
 
 
+@triton_op("torchao::triton_mx_block_rearrange_2d_M_groups", mutates_args={})
 def triton_mx_block_rearrange_2d_M_groups(
     scales_tensor: torch.Tensor,
     input_group_end_offsets: torch.Tensor,
@@ -238,7 +240,7 @@ def triton_mx_block_rearrange_2d_M_groups(
         num_groups,
         num_col_blocks,
     )
-    triton_scale_swizzle_M_groups[grid](
+    wrap_triton(triton_scale_swizzle_M_groups)[grid](
         # Input scales
         scales_tensor.view(torch.uint8),
         scales_tensor.stride(0),
@@ -336,6 +338,7 @@ def triton_scale_swizzle_M_groups(
         current_start_row += BLOCK_ROWS
 
 
+@triton_op("torchao::triton_mx_block_rearrange_per_group_3d", mutates_args={})
 def triton_mx_block_rearrange_per_group_3d(scale_tensor: torch.Tensor) -> torch.Tensor:
     """
     Rearranges an E8M0 tensor scale to block-scaled swizzle format.
@@ -379,7 +382,7 @@ def triton_mx_block_rearrange_per_group_3d(scale_tensor: torch.Tensor) -> torch.
         num_col_blocks,
     )
 
-    triton_scale_swizzle_per_group_3d[grid](
+    wrap_triton(triton_scale_swizzle_per_group_3d)[grid](
         scale_tensor.view(torch.uint8),
         input_stride_dim0,
         input_stride_dim1,
@@ -454,6 +457,7 @@ def triton_scale_swizzle_per_group_3d(
     )
 
 
+@triton_op("torchao::triton_mx_block_rearrange_2d_K_groups", mutates_args={})
 def triton_mx_block_rearrange_2d_K_groups(
     scales_tensor: torch.Tensor,
     input_group_end_offsets: torch.Tensor,
@@ -497,7 +501,7 @@ def triton_mx_block_rearrange_2d_K_groups(
         num_groups,
         num_row_blocks,
     )
-    triton_scale_swizzle_2d_K_groups[grid](
+    wrap_triton(triton_scale_swizzle_2d_K_groups)[grid](
         # Input scales
         scales_tensor.view(torch.uint8),
         scales_tensor.stride(0),

torchao/prototype/moe_training/scaled_grouped_mm.py

@@ -66,7 +66,7 @@ def _scaled_grouped_mm(
             out_dtype,
         )
     elif scaling_type == MoEScalingType.MXFP8:
-        logger.info("Using mxfp8 for _scaled_grouped_mm")
+        # logger.info("Using mxfp8 for _scaled_grouped_mm")
         block_size = 32  # TODO: should we make this configurable? plumb it through in a config somehow?
         return _MXFP8GroupedMM.apply(
             A,

torchao/prototype/moe_training/tensor.py

@@ -27,7 +27,7 @@
     torch.ops.aten.copy_.default,
     torch.ops.aten.view.default,
     torch.ops.aten.as_strided.default,
-    torch.ops.aten._to_copy.default,
+    torch.ops.aten._to_copy.default,  # for *.to(dtype)
     torch.ops.aten._pin_memory.default,
     torch.ops.aten.split.Tensor,
     torch.ops.aten.clone.default,
@@ -94,11 +94,11 @@ def __torch_function__(cls, func, types, args, kwargs={}):
                 "B should be a ScaledGroupedMMTensor"
             )
             scaling_type = B.scaling_type
-            A_is_2d = A.dim() == 2
-            B_is_3d = B.dim() == 3
+            A_is_2d = A.ndim == 2
+            B_is_2d_or_3d = B.ndim == 2 or B.ndim == 3
             has_offs = kwargs.get(cls.offs_arg_name) is not None
             other_args = args[2:]
-            if A_is_2d and B_is_3d and has_offs:
+            if A_is_2d and B_is_2d_or_3d and has_offs:
                 return _scaled_grouped_mm(
                     A,
                     B,
@@ -125,28 +125,41 @@ def unwrap(t):
                 assert t.scaling_type == scaling_type
             return t._data
 
-        args, kwargs = pytree.tree_map_only(
+        args_unwrapped, kwargs_unwrapped = pytree.tree_map_only(
             ScaledGroupedMMTensor, unwrap, (args, kwargs or {})
         )
-        assert scaling_type is not None
+        assert scaling_type is not None, (
+            f"__torch_dispatch__ called on {func.__name__} without any ScaledGroupedMMTensor arguments"
+        )
 
-        # detach is special case
+        # # detach is special case
         if func == torch.ops.aten.detach.default:
             return ScaledGroupedMMTensor(args[0], scaling_type)
 
         # perform op
-        out = func(*args, **kwargs)
+        with torch._C.DisableTorchFunctionSubclass():
+            out = func(*args_unwrapped, **kwargs_unwrapped)
 
         # return regular tensors for ops that don't preserve subclass
         if func not in _ops_to_preserve_subclass:
             return out
 
-        # wrap outputs back into ScaledGroupedMMTensor for ops that do preserve subclass
-        return pytree.tree_map_only(
-            torch.Tensor,
-            lambda x: ScaledGroupedMMTensor(x, scaling_type),
-            out,
+        # For ops that do preserve subclass, only preserve it if the primary input
+        # was a ScaledGroupedMMTensor. Most ops in _ops_to_preserve_subclass are
+        # unary operations where the first argument determines the output type.
+        unary_op_on_subclass = len(args) > 0 and isinstance(
+            args[0], ScaledGroupedMMTensor
         )
+        if not unary_op_on_subclass:
+            return out
+
+        # Only wrap tensor outputs, preserving any existing subclasses
+        def selective_wrap(x):
+            if isinstance(x, torch.Tensor) and type(x) == torch.Tensor:
+                return ScaledGroupedMMTensor(x, scaling_type)
+            return x
+
+        return pytree.tree_map(selective_wrap, out)
 
     def __repr__(self):
         return f"ScaledGroupedMMTensor(data={self._data}, scaling_type={self.scaling_type})"