Update

[ghstack-poisoned]

Author: alexsamardzic

test/prototype/moe_training/test_kernels.py

@@ -449,7 +449,10 @@ def test_cuda_mx_3d_cutedsl_numerics(E, N, K, input_dtype, scaling_mode, variant
         s_rows, s_cols = x_t.shape[-1], x_t.shape[-2] // block_size
         s_logical = (
             torch.stack(
-                [from_blocked(s[e], s_rows, s_cols).view(torch.uint8) for e in range(E)],
+                [
+                    from_blocked(s[e], s_rows, s_cols).view(torch.uint8)
+                    for e in range(E)
+                ],
                 dim=0,
             )
             .view(torch.float8_e8m0fnu)

torchao/prototype/moe_training/mxfp8_grouped_mm.py

@@ -17,6 +17,7 @@
     mxfp8_quantize_2d_1x32_cutedsl,
     mxfp8_quantize_cuda_3d,
     triton_mx_block_rearrange_2d_K_groups,
+    triton_mx_block_rearrange_per_group_3d,
 )
 from torchao.prototype.moe_training.utils import (
     conditional_nostrict_trace,