[mxfp8 moe training] use dim1 cast cuda kernel for 3d weights by reshaping to 2d

Author: danielvegamyhre

torchao/prototype/moe_training/scaled_grouped_mm.py

@@ -358,13 +358,38 @@ def backward(ctx, grad_out: torch.Tensor):
 
         # B_data shape: (E, K, N)
         # B_scale shape: (E, K, N//block_size)
-        B_scales, B_data = to_mx(
+        B_scales_ref, B_data_ref = to_mx(
             # TODO: can we support non-contiguous input tensor in to_mx to eliminate this inefficiency?
             B_t.contiguous(),
             elem_dtype=torch.float8_e4m3fn,
             block_size=block_size,
         )
 
+        # Experiment with cuda kernel
+        B = B_t.transpose(-2, -1)
+        E, N, K = B.shape
+        B_reshaped = B.reshape(E * N, K)
+        B_t_mx = _to_mxfp8_dim1_kernel_wrapper(
+            B_reshaped,
+            block_size,
+            elem_dtype=torch.float8_e4m3fn,
+            hp_dtype=B_reshaped.dtype,
+            gemm_kernel_choice=MXGemmKernelChoice.CUTLASS,  # Not used
+            cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA,
+            scale_calculation_mode=ScaleCalculationMode.FLOOR,
+        )
+        B_data = B_t_mx.qdata.t()  # (K, E*N) -> (E*N, K)
+        B_data = B_data.reshape(E, N, K)  # (E*N, K) -> (E, N, K)
+        B_scales = B_t_mx._scale_e8m0.view(torch.uint8)  # (K, E*N//block_size)
+        B_scales = B_scales.reshape(
+            K, E, N // block_size
+        )  # (K, E*N//block_size) -> (K, E, N//block_size)
+        B_scales = B_scales.permute(
+            1, 0, 2
+        )  # (K, E, N//block_size) -> (E, K, N//block_size)
+        B_scales = B_scales.view(torch.float8_e8m0fnu)
+        B_data = B_data.transpose(-2, -1).contiguous().transpose(-2, -1)
+
         # Convert scales to blocked format for 2d-3d grouped mm
         grad_out_scales_blocked = triton_mx_block_rearrange_2d_M_groups(
             grad_out_scale,
@@ -376,7 +401,7 @@ def backward(ctx, grad_out: torch.Tensor):
         # grad_A = scaled grouped mm of (M,N) @ (B,N,K) = (M,K)
         grad_A = torch._scaled_grouped_mm(
             grad_out_data,
-            B_data.transpose(-2, -1),
+            B_data,
             grad_out_scales_blocked,
             B_scales_blocked,
             offs=offs,
@@ -606,3 +631,7 @@ def _emulated_mxfp8_scaled_grouped_mm_2d_2d(
     # Perform bf16 grouped GEMM using dequantized A and B.
     out = torch._grouped_mm(A, B, offs=offs, out_dtype=out_dtype)
     return out
+
+
+def round_up(x, y):
+    return ((x + y - 1) // y) * y