Optimize wgrad CUTLASS grouped gemm (#4891)

Summary:
Pull Request resolved: #4891

X-link: facebookresearch/FBGEMM#1916

- Make wgrad CUTLASS grouped gemm return float32 output when wgrad is provided, respecting e2e
- Optimize general heuristic
- Make tests cover wgrad accum with float32 output

Reviewed By: q10

Differential Revision: D82700455

Author: jiawenliu64

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad.cu

@@ -10,25 +10,18 @@
 
 namespace fbgemm_gpu {
 
-at::Tensor bf16bf16bf16_grouped_wgrad_256_128_128_1_2_1_9_f(
+at::Tensor bf16bf16bf16_grouped_wgrad_128_128_128_1_1_1_9_t(
     at::Tensor X, // BF16
     at::Tensor W, // BF16
     at::Tensor M_sizes,
     at::Tensor output,
     bool output_accum) {
   if (output_accum) {
-    return bf16bf16bf16_grouped_wgrad_impl<256, 128, 128, 1, 2, 1, true, false>(
+    return bf16bf16bf16_grouped_wgrad_impl<128, 128, 128, 1, 1, 1, true, true>(
         X, W, M_sizes, output);
   } else {
-    return bf16bf16bf16_grouped_wgrad_impl<
-        256,
-        128,
-        128,
-        1,
-        2,
-        1,
-        false,
-        false>(X, W, M_sizes, output);
+    return bf16bf16bf16_grouped_wgrad_impl<128, 128, 128, 1, 1, 1, false, true>(
+        X, W, M_sizes, output);
   }
 }
 

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_256_128_128_1_2_1_9_f.cu renamed to fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_128_128_128_1_1_1_9_t.cu

@@ -10,22 +10,22 @@
 
 namespace fbgemm_gpu {
 
-at::Tensor bf16bf16bf16_grouped_wgrad_128_256_128_1_1_1_9_f(
+at::Tensor bf16bf16bf16_grouped_wgrad_128_128_128_1_4_1_9_f(
     at::Tensor X, // BF16
     at::Tensor W, // BF16
     at::Tensor M_sizes,
     at::Tensor output,
     bool output_accum) {
   if (output_accum) {
-    return bf16bf16bf16_grouped_wgrad_impl<128, 256, 128, 1, 1, 1, true, false>(
+    return bf16bf16bf16_grouped_wgrad_impl<128, 128, 128, 1, 4, 1, true, false>(
         X, W, M_sizes, output);
   } else {
     return bf16bf16bf16_grouped_wgrad_impl<
         128,
-        256,
+        128,
         128,
         1,
-        1,
+        4,
         1,
         false,
         false>(X, W, M_sizes, output);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_128_256_128_1_1_1_9_f.cu renamed to fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_128_128_128_1_4_1_9_f.cu

@@ -10,23 +10,23 @@
 
 namespace fbgemm_gpu {
 
-at::Tensor bf16bf16bf16_grouped_wgrad_128_256_128_1_2_1_9_f(
+at::Tensor bf16bf16bf16_grouped_wgrad_128_128_128_2_1_1_9_f(
     at::Tensor X, // BF16
     at::Tensor W, // BF16
     at::Tensor M_sizes,
     at::Tensor output,
     bool output_accum) {
   if (output_accum) {
-    return bf16bf16bf16_grouped_wgrad_impl<128, 256, 128, 1, 2, 1, true, false>(
+    return bf16bf16bf16_grouped_wgrad_impl<128, 128, 128, 2, 1, 1, true, false>(
         X, W, M_sizes, output);
   } else {
     return bf16bf16bf16_grouped_wgrad_impl<
         128,
-        256,
         128,
-        1,
+        128,
         2,
         1,
+        1,
         false,
         false>(X, W, M_sizes, output);
   }

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_128_256_128_1_2_1_9_f.cu renamed to fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_128_128_128_2_1_1_9_f.cu

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "bf16bf16bf16_grouped_wgrad_common.cuh"
+
+namespace fbgemm_gpu {
+
+at::Tensor bf16bf16bf16_grouped_wgrad_128_128_128_4_1_1_9_t(
+    at::Tensor X, // BF16
+    at::Tensor W, // BF16
+    at::Tensor M_sizes,
+    at::Tensor output,
+    bool output_accum) {
+  if (output_accum) {
+    return bf16bf16bf16_grouped_wgrad_impl<128, 128, 128, 4, 1, 1, true, true>(
+        X, W, M_sizes, output);
+  } else {
+    return bf16bf16bf16_grouped_wgrad_impl<128, 128, 128, 4, 1, 1, false, true>(
+        X, W, M_sizes, output);
+  }
+}
+
+} // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_128_128_128_4_1_1_9_t.cu

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "bf16bf16bf16_grouped_wgrad_common.cuh"
+
+namespace fbgemm_gpu {
+
+at::Tensor bf16bf16bf16_grouped_wgrad_128_32_128_1_2_1_9_f(
+    at::Tensor X, // BF16
+    at::Tensor W, // BF16
+    at::Tensor M_sizes,
+    at::Tensor output,
+    bool output_accum) {
+  if (output_accum) {
+    return bf16bf16bf16_grouped_wgrad_impl<128, 32, 128, 1, 2, 1, true, false>(
+        X, W, M_sizes, output);
+  } else {
+    return bf16bf16bf16_grouped_wgrad_impl<128, 32, 128, 1, 2, 1, false, false>(
+        X, W, M_sizes, output);
+  }
+}
+
+} // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_128_32_128_1_2_1_9_f.cu

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "bf16bf16bf16_grouped_wgrad_common.cuh"
+
+namespace fbgemm_gpu {
+
+at::Tensor bf16bf16bf16_grouped_wgrad_128_64_128_1_1_1_9_f(
+    at::Tensor X, // BF16
+    at::Tensor W, // BF16
+    at::Tensor M_sizes,
+    at::Tensor output,
+    bool output_accum) {
+  if (output_accum) {
+    return bf16bf16bf16_grouped_wgrad_impl<128, 64, 128, 1, 1, 1, true, false>(
+        X, W, M_sizes, output);
+  } else {
+    return bf16bf16bf16_grouped_wgrad_impl<128, 64, 128, 1, 1, 1, false, false>(
+        X, W, M_sizes, output);
+  }
+}
+
+} // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_128_64_128_1_1_1_9_f.cu

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "bf16bf16bf16_grouped_wgrad_common.cuh"
+
+namespace fbgemm_gpu {
+
+at::Tensor bf16bf16bf16_grouped_wgrad_256_64_128_1_1_1_9_f(
+    at::Tensor X, // BF16
+    at::Tensor W, // BF16
+    at::Tensor M_sizes,
+    at::Tensor output,
+    bool output_accum) {
+  if (output_accum) {
+    return bf16bf16bf16_grouped_wgrad_impl<256, 64, 128, 1, 1, 1, true, false>(
+        X, W, M_sizes, output);
+  } else {
+    return bf16bf16bf16_grouped_wgrad_impl<256, 64, 128, 1, 1, 1, false, false>(
+        X, W, M_sizes, output);
+  }
+}
+
+} // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_256_64_128_1_1_1_9_f.cu

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "bf16bf16bf16_grouped_wgrad_common.cuh"
+
+namespace fbgemm_gpu {
+
+at::Tensor bf16bf16bf16_grouped_wgrad_256_64_128_1_2_1_9_f(
+    at::Tensor X, // BF16
+    at::Tensor W, // BF16
+    at::Tensor M_sizes,
+    at::Tensor output,
+    bool output_accum) {
+  if (output_accum) {
+    return bf16bf16bf16_grouped_wgrad_impl<256, 64, 128, 1, 2, 1, true, false>(
+        X, W, M_sizes, output);
+  } else {
+    return bf16bf16bf16_grouped_wgrad_impl<256, 64, 128, 1, 2, 1, false, false>(
+        X, W, M_sizes, output);
+  }
+}
+
+} // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_256_64_128_1_2_1_9_f.cu

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "bf16bf16bf16_grouped_wgrad_common.cuh"
+
+namespace fbgemm_gpu {
+
+at::Tensor bf16bf16bf16_grouped_wgrad_256_64_128_1_4_1_9_f(
+    at::Tensor X, // BF16
+    at::Tensor W, // BF16
+    at::Tensor M_sizes,
+    at::Tensor output,
+    bool output_accum) {
+  if (output_accum) {
+    return bf16bf16bf16_grouped_wgrad_impl<256, 64, 128, 1, 4, 1, true, false>(
+        X, W, M_sizes, output);
+  } else {
+    return bf16bf16bf16_grouped_wgrad_impl<256, 64, 128, 1, 4, 1, false, false>(
+        X, W, M_sizes, output);
+  }
+}
+
+} // namespace fbgemm_gpu