Enable CUTLASS grouped GEMM for pretraining wgrad on GB200 and H100 (resubmit)

Summary:
- Enable CUTLASS grouped GEMM for llama4x pretraining wgrad on GB200 and H100
- Optimize performance of pretraining moe shapes on H100
- Support total_K in quantize_bench for wgrad
- The FBGEMM relocation issue has been released for short-term, so resubmit. Passed all tests in T238469849

Differential Revision: D83001505

Author: jiawenliu64

fbgemm_gpu/experimental/gen_ai/bench/quantize_bench.py

@@ -486,10 +486,16 @@ def print_kernels(kernels: Optional[List[str]]) -> List[QuantizeOpBase]:
     default=None,
     help="If set with grouped mode, repeat input shapes this many times. Comma separated list of groups to benchmark",
 )
+@click.option(
+    "--total-K",
+    default=None,
+    help="If set, adjusts the K values to sum to this number. "
+    "This can help simulate real grouped workloads in backward wgrad.",
+)
 @click.option(
     "--total-M",
     default=None,
-    help="If set, Adjusts the M values to sum to this number. "
+    help="If set, adjusts the M values to sum to this number. "
     "This can help simulate real grouped workloads.",
 )
 @click.option(
@@ -542,6 +548,7 @@ def invoke_main(
     pair_nk: bool,
     grouped: bool,
     groups: Optional[str],
+    total_k: Optional[str],
     total_m: Optional[str],
     no_cuda_graph: bool,
     use_rotating_buffer_bench: bool,
@@ -553,6 +560,14 @@ def invoke_main(
 ):
     if enable_amd_env_vars:
         set_amd_env_vars()
+
+    # Validate that total_m and total_k are mutually exclusive
+    if total_m is not None and total_k is not None:
+        raise ValueError(
+            "total_m and total_k cannot be specified at the same time. "
+            "Please provide only one of them."
+        )
+
     # If kernel filter is provided, parse it. Else, benchmark all kernels.
     all_kernels = kernels.strip().split(",") if kernels else None
     quantize_ops = collect_kernels_to_profile(all_kernels)
@@ -629,6 +644,17 @@ def invoke_main(
                 for g in groups_list
                 for b, _, n, k in MNK
             ]
+        elif total_k:
+            MNK = [
+                [
+                    [b] * g,
+                    [m] * g,
+                    [n] * g,
+                    generate_group_tensor(g, int(total_k)),
+                ]
+                for g in groups_list
+                for b, m, n, _ in MNK
+            ]
         else:
             MNK = [
                 [[b] * g, [m] * g, [n] * g, [k] * g]

fbgemm_gpu/experimental/gen_ai/bench/quantize_ops.py

@@ -2084,7 +2084,7 @@ def cuda(self) -> bool:
 @register_quantize_op
 class BF16GroupedGrad(QuantizeOpBase):
     """
-    BF16 grouped matmul with grad inputs backed by cutlass
+    BF16 grouped matmul with dgrad inputs in pretraining backed by cutlass
     """
 
     def preprocess(self, x, w):
@@ -2126,6 +2126,52 @@ def cuda(self) -> bool:
         return True
 
 
+@register_quantize_op
+class BF16GroupedWGrad(QuantizeOpBase):
+    """
+    BF16 grouped matmul with wgrad inputs in pretraining backed by cutlass
+    """
+
+    def preprocess(self, x, w):
+        # Get K values for each group
+        k_values = [xi.shape[1] for xi in x]  # K dimension for each group
+
+        # Convert k_values into sizes tensor
+        k_sizes = torch.tensor(k_values).to(dtype=torch.int64, device=x[0].device)
+
+        x = torch.concat(x, dim=1).contiguous()  # shape: (M, G*K)
+        w = torch.concat(w, dim=1).contiguous()  # shape: (N, G*K)
+
+        # Transpose the follows to simulate wgrad shapes
+        x = x.t().contiguous()  # shape: (G*K, M)
+        w = w.t().contiguous()  # shape: (G*K, N)
+
+        # Return processed tensors
+        return x, w, k_sizes
+
+    def quantize(self, x, w, k_sizes):
+        return x, w, k_sizes
+
+    def compute(self, x, w, k_sizes):
+        return torch.ops.fbgemm.bf16bf16bf16_grouped_wgrad(x, w, k_sizes)
+
+    def quantize_and_compute(self, x, w, k_sizes):
+        x, w, k_sizes = self.quantize(x, w, k_sizes)
+        return self.compute(x, w, k_sizes)
+
+    @property
+    def name(self) -> str:
+        return "bf16_grouped_wgrad"
+
+    @property
+    def hip(self) -> bool:
+        return False
+
+    @property
+    def cuda(self) -> bool:
+        return True
+
+
 @register_quantize_op
 class BF16GroupedStacked(QuantizeOpBase):
     """