[mxfp8 moe] add support for fbgemm 2d-3d mx8mx8bf16 grouped gemm

stack-info: PR: #2848, branch: danielvegamyhre/stack/55

Author: danielvegamyhre

test/prototype/moe_training/test_scaled_grouped_mm.py

@@ -230,25 +230,27 @@ def compute_reference_forward(
 @pytest.mark.parametrize("num_experts", (1, 8, 16))
 def test_emulate_mxfp8_grouped_gemm_2d_3d(M, K, N, num_experts):
     x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
-    w_t = torch.randn(num_experts, K, N, dtype=torch.bfloat16, device="cuda")
+    w = torch.randn(num_experts, N, K, dtype=torch.bfloat16, device="cuda")
     offs = generate_jagged_offs(num_experts, M)
-    x_ref, w_t_ref, offs_ref = x.clone(), w_t.clone(), offs.clone()
+    x_ref, w_ref, offs_ref = x.clone(), w.clone(), offs.clone()
 
     # Quantize inputs to mxpf8 for emulated mxfp8 scaled grouped mm
     block_size = 32
-    x_scale, x_mx = to_mx(x, elem_dtype=torch.float8_e4m3fn, block_size=block_size)
+    x_scale, x_fp8 = to_mx(x, elem_dtype=torch.float8_e4m3fn, block_size=block_size)
 
     # To cast B_t per-expert to mxfp8 across dim1, we transpose the experts, cast along dim -1, then untranspose.
-    w_scale, w_mx = to_mx(
-        w_t.transpose(-2, -1).contiguous(),
+    w_scale, w_fp8 = to_mx(
+        w,
         elem_dtype=torch.float8_e4m3fn,
         block_size=block_size,
     )
-    w_t_scale, w_t_mx = w_scale.transpose(-2, -1), w_mx.transpose(-2, -1)
+    w_t_scale, w_t_fp8 = w_scale.transpose(-2, -1), w_fp8.transpose(-2, -1)
 
-    ref_out = torch._grouped_mm(x_ref, w_t_ref, offs=offs_ref, out_dtype=torch.bfloat16)
+    ref_out = torch._grouped_mm(
+        x_ref, w_ref.transpose(-2, -1), offs=offs_ref, out_dtype=torch.bfloat16
+    )
     out = _emulated_mxfp8_scaled_grouped_mm_2d_3d(
-        x_mx, x_scale, w_t_mx, w_t_scale, offs=offs, out_dtype=torch.bfloat16
+        x_fp8, x_scale, w_t_fp8, w_t_scale, offs=offs, out_dtype=torch.bfloat16
     )
 
     sqnr = compute_error(ref_out, out)
@@ -305,19 +307,26 @@ def test_emulate_mxfp8_grouped_gemm_2d_2d(M, N, num_experts):
 
 
 @skip_if_rocm("ROCm not supported")
-@pytest.mark.parametrize("M,K,N", [(1024, 1024, 1024), (1024, 2048, 4096)])
-@pytest.mark.parametrize("num_experts", (1, 8, 16))
+@pytest.mark.parametrize("M,K,N", [(256, 512, 512)])
+@pytest.mark.parametrize("num_experts", (2,))
 def test_mxfp8_grouped_gemm_with_dq_fwd_bwd(M, K, N, num_experts):
     from torchao.prototype.moe_training.scaled_grouped_mm import (
         _MXFP8GroupedMM,
     )
 
     block_size = 32
     x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True)
-    w_t = torch.randn(
-        num_experts, K, N, dtype=torch.bfloat16, device="cuda", requires_grad=True
+    w = torch.randn(
+        num_experts,
+        N,
+        K,
+        dtype=torch.bfloat16,
+        device="cuda",
     )
-    offs = generate_jagged_offs(num_experts, M, multiple_of=block_size)
+    w_t = w.transpose(-2, -1).requires_grad_(True)
+    #    offs = generate_jagged_offs(num_experts, M, multiple_of=block_size)
+    group_size = M // num_experts
+    offs = torch.arange(group_size, M + 1, group_size, device="cuda", dtype=torch.int32)
     x_ref, w_t_ref, offs_ref = (
         x.clone().detach().requires_grad_(True),
         w_t.clone().detach().requires_grad_(True),

test/prototype/moe_training/test_training.py

@@ -136,7 +136,7 @@ def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
         ["does.not.exist"],
     ],
 )
-@pytest.mark.parametrize("compile", [False, True])
+@pytest.mark.parametrize("compile", [False])
 def test_moe_mxfp8_training(target_fqns: list[str], compile: bool):
     block_size = 32
 

torchao/prototype/moe_training/kernels/__init__.py

@@ -7,3 +7,6 @@
 from torchao.prototype.moe_training.kernels.jagged_float8_scales import (
     triton_fp8_per_group_rowwise_scales as triton_fp8_per_group_rowwise_scales,
 )
+from torchao.prototype.moe_training.kernels.mxfp8 import (
+    fbgemm_mxfp8_grouped_mm_2d_3d as fbgemm_mxfp8_grouped_mm_2d_3d,
+)

torchao/prototype/moe_training/kernels/mxfp8.py

@@ -0,0 +1,129 @@
+import logging
+
+import torch
+
+from torchao.prototype.mx_formats.utils import (
+    to_blocked_per_group_2d,
+    to_blocked_per_group_3d,
+)
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+try:
+    import fbgemm_gpu.experimental.gen_ai  # noqa: F401
+except Exception as e:
+    logging.warning(
+        f"fbgemm_gpu_genai package is required for this feature but import failed with exception: {e}"
+        "Please install nightly builds of pytorch and fbgemm_gpu_genai build using this command and try again: "
+        "pip3 install --force-reinstall --pre torch fbgemm-gpu-genai --index-url https://download.pytorch.org/whl/nightly/cu129"
+        "If errors persist, please file a bug report."
+    )
+
+
+@torch.library.custom_op("torchao::fbgemm_mxfp8_grouped_mm_2d_3d", mutates_args={})
+def fbgemm_mxfp8_grouped_mm_2d_3d(
+    A_mx: torch.Tensor,
+    A_scale: torch.Tensor,
+    B_t_mx_dim1: torch.Tensor,
+    B_t_scales_dim1: torch.Tensor,
+    offs: torch.Tensor,
+    block_size: int = 32,
+    out_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    assert A_mx.ndim == 2, "A_mx tensor must be 2D"
+    assert B_t_mx_dim1.ndim == 3, "B_t_mx_dim1 tensor must be 3D"
+    assert block_size == 32, "Only block_size=32 is supported"
+    assert out_dtype == torch.bfloat16, "Only out_dtype=bfloat16 is supported"
+
+    # Convert scales for each group to blocked format.
+    Mg, K = A_mx.shape
+    A_scale_blocked, starting_row_after_padding = to_blocked_per_group_2d(
+        A_scale, offs, Mg, K
+    )
+    B_t_scales_dim1_blocked = to_blocked_per_group_3d(B_t_scales_dim1)
+
+    # From this, we compute `group_sizes` and `starting_row_after_padding`:
+    # group_sizes = [32, 32, 64]
+    # starting_row_after_padding = [0, 32, 64, 128]
+    group_sizes = torch.diff(starting_row_after_padding).to(torch.int64)
+
+    # TODO: remove debug logging once prototype is more mature.
+    logger.debug("A_mx.shape", A_mx.shape, "stride", A_mx.stride(), "dtype", A_mx.dtype)
+    logger.debug(
+        "B_t_mx_dim1.shape",
+        B_t_mx_dim1.shape,
+        "stride",
+        B_t_mx_dim1.stride(),
+        "dtype",
+        B_t_mx_dim1.dtype,
+    )
+    logger.debug(
+        "A_scales_blocked.shape",
+        A_scale_blocked.shape,
+        "stride",
+        A_scale_blocked.stride(),
+        "dtype",
+        A_scale_blocked.dtype,
+    )
+    logger.debug(
+        "B_t_scales_dim (non-blocked)",
+        B_t_scales_dim1.shape,
+        "stride",
+        B_t_scales_dim1.stride(),
+        "dtype",
+        B_t_scales_dim1.dtype,
+    )
+    logger.debug(
+        "B_t_scales_dim1_blocked.shape",
+        B_t_scales_dim1_blocked.shape,
+        "stride",
+        B_t_scales_dim1_blocked.stride(),
+        "dtype",
+        B_t_scales_dim1_blocked.dtype,
+    )
+    logger.debug(
+        "group_sizes",
+        group_sizes,
+        "group_sizes.stride",
+        group_sizes.stride(),
+        "dtype",
+        group_sizes.dtype,
+    )
+    logger.debug(
+        "starting_row_after_padding",
+        starting_row_after_padding,
+        "stride",
+        starting_row_after_padding.stride(),
+        "dtype",
+        starting_row_after_padding.dtype,
+    )
+
+    out = torch.ops.fbgemm.mx8mx8bf16_grouped_stacked(
+        A_mx,
+        B_t_mx_dim1,
+        A_scale_blocked,
+        B_t_scales_dim1_blocked,
+        group_sizes,
+        starting_row_after_padding=starting_row_after_padding,
+    )
+    return out
+
+
+@fbgemm_mxfp8_grouped_mm_2d_3d.register_fake
+def _fbgemm_mxfp8_grouped_mm_2d_3d_fake(
+    A_mx: torch.Tensor,
+    B_t_mx_dim1: torch.Tensor,
+    A_scale: torch.Tensor,
+    B_t_scales_dim1: torch.Tensor,
+    offs: torch.Tensor,
+) -> torch.Tensor:
+    assert A_mx.ndim == 2, "A_mx tensor must be 2D"
+    assert B_t_mx_dim1.ndim == 3, "B_t_mx_dim1 tensor must be 3D"
+    mg, k = A_mx.shape
+    e, k, n = B_t_mx_dim1.shape
+    n_groups = offs.numel()
+    assert n_groups == e, (
+        "Size of `offs` (number of groups) must match first dim of `B_t_mx_dim1`"
+    )
+    output = torch.empty((mg, n), dtype=torch.bfloat16, device=A_mx.device)
+    return output

torchao/prototype/moe_training/scaled_grouped_mm.py

@@ -13,6 +13,7 @@
 from torchao.float8.float8_utils import tensor_to_scale, to_fp8_saturated
 from torchao.prototype.moe_training.conversion_utils import MoEScalingType
 from torchao.prototype.moe_training.kernels import (
+    fbgemm_mxfp8_grouped_mm_2d_3d,
     triton_fp8_per_group_colwise_scales,
     triton_fp8_per_group_rowwise_scales,
     triton_fp8_rowwise_3d_transpose_rhs,
@@ -277,52 +278,46 @@ def forward(
         offs: Optional[torch.Tensor] = None,
         block_size: int = 32,
         out_dtype: Optional[torch.dtype] = torch.bfloat16,
-        emulated: bool = True,
+        emulated: bool = False,
     ) -> torch.Tensor:
         # torchao _scaled_grouped_mm only supports A=2D and B=3D.
         assert A.ndim == 2, "A must be 2D"
         assert B_t.ndim == 3, "B must be 3D"
         assert block_size == 32, "Only block_size=32 is supported"
-        assert emulated, "Only emulated mxfp8 grouped gemm is supported"
+
+        # Store what we need for backward.
+        ctx.save_for_backward(A, B_t, offs)
+        ctx.block_size = block_size
+        ctx.out_dtype = out_dtype
+        ctx.emulated = emulated
 
         # Cast to mxpf8 across dim -1.
         # A_mx shape: (M, K)
         # A_scale shape: (M, K//block_size)
         A_scale, A_mx = to_mx(A, elem_dtype=torch.float8_e4m3fn, block_size=block_size)
 
-        # Cast B_t per-expert to mxfp8 across dim1.
-        # B_t_mx shape: (E, K, N)
-        # B_t_scale shape: (E, K//block_size, N)
-
-        # To cast B_t per-expert to mxfp8 across dim1, we transpose the experts, cast along dim -1, then untranspose.
+        # Cast B_t per-expert to mxfp8 across K dim.
         # B_mx shape: (E, N, K)
         # B_scale shape: (E, N, K//block_size)
-        B_scales_dim2, B_mx_dim2 = to_mx(
-            B_t.transpose(-2, -1),  # (E,K,N) -> (E,N,K)
+        B_scales, B_mx = to_mx(
+            B_t.transpose(-2, -1).contiguous(),
             elem_dtype=torch.float8_e4m3fn,
             block_size=block_size,
         )
 
-        # B_t_mx shape: (E, K, N)
-        # B_t_scale shape: (E, K//block_size, N)
-        B_t_scales_dim1 = B_scales_dim2.transpose(
-            -2, -1
-        )  # (E,N,K//block_size) -> (E,K//block_size,N)
-        B_t_mx_dim1 = B_mx_dim2.transpose(-2, -1)  # (E,N,K) -> (E,K,N)
-
-        # Store what we need for backward.
-        ctx.save_for_backward(A, B_t, offs)
-        ctx.block_size = block_size
-        ctx.out_dtype = out_dtype
-
         # Perform scaled grouped GEMM and return result.
         # output = input @ weight.T
         # output shape: (M, N)
-        out = _emulated_mxfp8_scaled_grouped_mm_2d_3d(
+        mxfp8_2d_3d_grouped_mm = (
+            _emulated_mxfp8_scaled_grouped_mm_2d_3d
+            if emulated
+            else fbgemm_mxfp8_grouped_mm_2d_3d
+        )
+        out = mxfp8_2d_3d_grouped_mm(
             A_mx,
             A_scale,
-            B_t_mx_dim1,
-            B_t_scales_dim1,
+            B_mx,
+            B_scales,
             offs=offs,
             block_size=block_size,
             out_dtype=out_dtype,
@@ -334,6 +329,7 @@ def backward(ctx, grad_out: torch.Tensor):
         A, B_t, offs = ctx.saved_tensors
         block_size = ctx.block_size
         out_dtype = ctx.out_dtype
+        emulated = ctx.emulated
 
         # grad_out_mx shape: (M, N)
         # grad_out_scale shape: (M, N//block_size)
@@ -343,23 +339,24 @@ def backward(ctx, grad_out: torch.Tensor):
 
         # B_mx shape: (E, K, N)
         # B_scale shape: (E, K, N//block_size)
-        B_t_scale_dim2, B_t_mx_dim2 = to_mx(
+        B_scales, B_mx = to_mx(
             B_t.contiguous(),
             elem_dtype=torch.float8_e4m3fn,
             block_size=block_size,
         )
-        B_scale_dim1 = B_t_scale_dim2.transpose(
-            -2, -1
-        )  # (E,K,N//block_size) -> (E,N//block_size,K)
-        B_mx_dim1 = B_t_mx_dim2.transpose(-2, -1)  # (E,K,N) -> (E,N,K)
 
         # Compute grad_A.
         # grad_A = scaled grouped mm of (M,N) @ (B,N,K) = (M,K)
-        grad_A = _emulated_mxfp8_scaled_grouped_mm_2d_3d(
+        mxfp8_2d_3d_grouped_mm = (
+            _emulated_mxfp8_scaled_grouped_mm_2d_3d
+            if emulated
+            else fbgemm_mxfp8_grouped_mm_2d_3d
+        )
+        grad_A = mxfp8_2d_3d_grouped_mm(
             grad_out_mx,
             grad_out_scale,
-            B_mx_dim1,
-            B_scale_dim1,
+            B_mx,
+            B_scales,
             offs=offs,
             out_dtype=out_dtype,
         )
@@ -385,7 +382,6 @@ def backward(ctx, grad_out: torch.Tensor):
         # Compute grad_B = grad_output_t @ A
         # grad_B_t = scaled grouped mm of (N,M) @ (M,K) = (E,N,K)
         # grad_B = grad_B_t.transpose(-2, -1) = (E,K,N)
-
         grad_B = _emulated_mxfp8_scaled_grouped_mm_2d_2d(
             grad_out_t_mx,
             grad_out_t_scales,
@@ -402,12 +398,24 @@ def backward(ctx, grad_out: torch.Tensor):
 def _emulated_mxfp8_scaled_grouped_mm_2d_3d(
     A_mx: torch.Tensor,
     A_scale: torch.Tensor,
-    B_t_mx: torch.Tensor,
-    B_t_scale: torch.Tensor,
+    B_mx: torch.Tensor,
+    B_scale: torch.Tensor,
     offs: Optional[torch.Tensor] = None,
     out_dtype: Optional[torch.dtype] = torch.bfloat16,
     block_size: int = 32,
 ) -> torch.Tensor:
+    assert A_mx.ndim == 2, "A must be 2D"
+    assert B_mx.ndim == 3, "B must be 3D"
+    assert A_scale.shape[0] == A_mx.shape[0], "A_scale must have same M dim as A_mx"
+    assert A_scale.shape[1] == A_mx.shape[1] // block_size, (
+        "A_scale dim1 should be size K//block_size"
+    )
+    assert B_scale.shape[0] == B_mx.shape[0], "B_scale must have same E dim as B_mx"
+    assert B_scale.shape[1] == B_mx.shape[1], "B_scale must have same N dim as B_mx"
+    assert B_scale.shape[2] == B_mx.shape[2] // block_size, (
+        "B_scale dim2 should be size K//block_size"
+    )
+
     # Dequantize input
     # A_mx shape: (M, K)
     # A_scale shape: (M, K//block_size)
@@ -427,14 +435,10 @@ def _emulated_mxfp8_scaled_grouped_mm_2d_3d(
     A = A.reshape(A_orig_shape)
 
     # Dequantize weights
-    # B_t_mx shape: (E, K, N)
-    # B_t_scale shape: (E, K//block_size, N)
-    E, K, N = B_t_mx.shape
-
     # Tranpose to get block_size on rightmost dim
     # B_mx shape: (E, N, K)
     # B_scale shape: (E, N, K//block_size)
-    B_mx, B_scale = B_t_mx.transpose(-2, -1), B_t_scale.transpose(-2, -1)
+    E, N, K = B_mx.shape
 
     # Reshape to be able to do per-scaling group multiplication
     # B_mx shape: (E, N, K//block_size, block_size)