Improve general shape performance for fprop cutlass grouped gemm (pytorch#4988)

Summary:
Pull Request resolved: pytorch#4988

X-link: https://github.com/facebookresearch/FBGEMM/pull/2003

Improve general shape performance for fprop cutlass grouped gemm

Reviewed By: q10, cthi

Differential Revision: D84194717

Author: jiawenliu64

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

@@ -10,17 +10,17 @@
 
 namespace fbgemm_gpu {
 
-at::Tensor bf16bf16bf16_grouped_128_256_128_4_1_1_9_t(
+at::Tensor bf16bf16bf16_grouped_128_128_128_2_4_1_9_t(
     at::Tensor X, // BF16
     at::Tensor W, // BF16
     at::Tensor output,
     std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes) {
-  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 256, 128, 4, 1, 1, true>(
+  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 128, 128, 2, 4, 1, true>(
       X, W, output, zero_start_index_M, M_sizes);
 }
 
-at::Tensor bf16bf16bf16_grouped_128_256_128_4_1_1_9_t(
+at::Tensor bf16bf16bf16_grouped_128_128_128_2_4_1_9_t(
     at::TensorList X, // BF16
     at::TensorList W, // BF16
     at::Tensor output,
@@ -29,11 +29,11 @@ at::Tensor bf16bf16bf16_grouped_128_256_128_4_1_1_9_t(
   return bf16bf16bf16_grouped_impl<
       at::TensorList,
       128,
-      256,
       128,
+      128,
+      2,
       4,
       1,
-      1,
       true>(X, W, output, zero_start_index_M, M_sizes);
 }
 

…f16bf16_grouped_128_256_128_4_1_1_9_t.cu …f16bf16_grouped_128_128_128_2_4_1_9_t.cufbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_256_128_4_1_1_9_t.cu renamed to fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_128_128_2_4_1_9_t.cu fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_256_128_4_1_1_9_t.cu renamed to fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_128_128_2_4_1_9_t.cu

@@ -10,17 +10,17 @@
 
 namespace fbgemm_gpu {
 
-at::Tensor bf16bf16bf16_grouped_128_128_128_4_1_1_9_t(
+at::Tensor bf16bf16bf16_grouped_128_128_128_4_2_1_9_t(
     at::Tensor X, // BF16
     at::Tensor W, // BF16
     at::Tensor output,
     std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes) {
-  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 128, 128, 4, 1, 1, true>(
+  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 128, 128, 4, 2, 1, true>(
       X, W, output, zero_start_index_M, M_sizes);
 }
 
-at::Tensor bf16bf16bf16_grouped_128_128_128_4_1_1_9_t(
+at::Tensor bf16bf16bf16_grouped_128_128_128_4_2_1_9_t(
     at::TensorList X, // BF16
     at::TensorList W, // BF16
     at::Tensor output,
@@ -32,7 +32,7 @@ at::Tensor bf16bf16bf16_grouped_128_128_128_4_1_1_9_t(
       128,
       128,
       4,
-      1,
+      2,
       1,
       true>(X, W, output, zero_start_index_M, M_sizes);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_128_128_4_1_1_9_f.cu

@@ -10,29 +10,29 @@
 
 namespace fbgemm_gpu {
 
-at::Tensor bf16bf16bf16_grouped_256_128_128_2_1_1_9_f(
+at::Tensor bf16bf16bf16_grouped_128_256_128_1_2_1_9_f(
     at::Tensor X, // BF16
     at::Tensor W, // BF16
     at::Tensor output,
     std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes) {
-  return bf16bf16bf16_grouped_impl<at::Tensor, 256, 128, 128, 2, 1, 1, false>(
+  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 256, 128, 1, 2, 1, false>(
       X, W, output, zero_start_index_M, M_sizes);
 }
 
-at::Tensor bf16bf16bf16_grouped_256_128_128_2_1_1_9_f(
+at::Tensor bf16bf16bf16_grouped_128_256_128_1_2_1_9_f(
     at::TensorList X, // BF16
     at::TensorList W, // BF16
     at::Tensor output,
     std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes) {
   return bf16bf16bf16_grouped_impl<
       at::TensorList,
-      256,
       128,
+      256,
       128,
-      2,
       1,
+      2,
       1,
       false>(X, W, output, zero_start_index_M, M_sizes);
 }

…f16bf16_grouped_128_128_128_4_1_1_9_t.cu …f16bf16_grouped_128_128_128_4_2_1_9_t.cufbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_128_128_4_1_1_9_t.cu renamed to fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_128_128_4_2_1_9_t.cu fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_128_128_4_1_1_9_t.cu renamed to fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_128_128_4_2_1_9_t.cu

@@ -10,17 +10,17 @@
 
 namespace fbgemm_gpu {
 
-at::Tensor bf16bf16bf16_grouped_128_16_128_2_1_1_9_f(
+at::Tensor bf16bf16bf16_grouped_128_32_128_1_2_1_9_f(
     at::Tensor X, // BF16
     at::Tensor W, // BF16
     at::Tensor output,
     std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes) {
-  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 16, 128, 2, 1, 1, false>(
+  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 32, 128, 1, 2, 1, false>(
       X, W, output, zero_start_index_M, M_sizes);
 }
 
-at::Tensor bf16bf16bf16_grouped_128_16_128_2_1_1_9_f(
+at::Tensor bf16bf16bf16_grouped_128_32_128_1_2_1_9_f(
     at::TensorList X, // BF16
     at::TensorList W, // BF16
     at::Tensor output,
@@ -29,10 +29,10 @@ at::Tensor bf16bf16bf16_grouped_128_16_128_2_1_1_9_f(
   return bf16bf16bf16_grouped_impl<
       at::TensorList,
       128,
-      16,
+      32,
       128,
-      2,
       1,
+      2,
       1,
       false>(X, W, output, zero_start_index_M, M_sizes);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16bf16bf16_grouped/bf16bf16bf16_grouped_128_256_128_1_1_1_9_t.cu

@@ -10,17 +10,17 @@
 
 namespace fbgemm_gpu {
 
-at::Tensor bf16bf16bf16_grouped_128_16_128_4_1_1_9_f(
+at::Tensor bf16bf16bf16_grouped_128_32_128_1_4_1_9_f(
     at::Tensor X, // BF16
     at::Tensor W, // BF16
     at::Tensor output,
     std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes) {
-  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 16, 128, 4, 1, 1, false>(
+  return bf16bf16bf16_grouped_impl<at::Tensor, 128, 32, 128, 1, 4, 1, false>(
       X, W, output, zero_start_index_M, M_sizes);
 }
 
-at::Tensor bf16bf16bf16_grouped_128_16_128_4_1_1_9_f(
+at::Tensor bf16bf16bf16_grouped_128_32_128_1_4_1_9_f(
     at::TensorList X, // BF16
     at::TensorList W, // BF16
     at::Tensor output,
@@ -29,10 +29,10 @@ at::Tensor bf16bf16bf16_grouped_128_16_128_4_1_1_9_f(
   return bf16bf16bf16_grouped_impl<
       at::TensorList,
       128,
-      16,
+      32,
       128,
-      4,
       1,
+      4,
       1,
       false>(X, W, output, zero_start_index_M, M_sizes);
 }