Fix tuning cache for f8f8bf16_rowwise_grouped on SM100 (#4843)

Summary:
Pull Request resolved: #4843

X-link: facebookresearch/FBGEMM#1871

It would run SM90 kernels before by accident.

Reviewed By: q10

Differential Revision: D82022651

fbshipit-source-id: ee739499faf61f73e5e9fbdb9d244cacb50c92e0

Author: cthi

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

@@ -202,7 +202,16 @@ Kernel_f8f8bf16_rowwise_grouped<InputType> get_kernel_via_tuning(
   const std::string shape_key = std::to_string(total_M) + "_" +
       std::to_string(max_N) + "_" + std::to_string(max_K) + "_" +
       std::to_string(G);
-  const auto& kernels = get_f8f8bf16_rowwise_grouped_kernels<InputType>();
+
+  const auto& kernels = []() {
+    const int arch = getDeviceArch();
+    if (arch == 9) {
+      return get_f8f8bf16_rowwise_grouped_kernels<InputType>();
+    } else {
+      return get_f8f8bf16_rowwise_grouped_kernels_sm100<InputType>();
+    }
+  }();
+
   auto kernel = cache.findBestKernelMaybeAutotune(
       shape_key,
       kernels,

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_grouped_sm100/f8f8bf16_rowwise_grouped_manifest.cuh

@@ -154,4 +154,31 @@ at::Tensor f8f8bf16_rowwise_grouped_256_256_128_2_1_1_10_f(
     std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes);
 
+template <typename InputType>
+const std::
+    unordered_map<std::string, Kernel_f8f8bf16_rowwise_grouped<InputType>>&
+    get_f8f8bf16_rowwise_grouped_kernels_sm100() {
+  static const std::
+      unordered_map<std::string, Kernel_f8f8bf16_rowwise_grouped<InputType>>
+          kernels = {
+              {"f8f8bf16_rowwise_grouped_128_32_128_2_1_1_10_f",
+               f8f8bf16_rowwise_grouped_128_32_128_2_1_1_10_f},
+              {"f8f8bf16_rowwise_grouped_128_64_128_2_1_1_10_f",
+               f8f8bf16_rowwise_grouped_128_64_128_2_1_1_10_f},
+              {"f8f8bf16_rowwise_grouped_128_128_128_2_1_1_10_f",
+               f8f8bf16_rowwise_grouped_128_128_128_2_1_1_10_f},
+              {"f8f8bf16_rowwise_grouped_128_256_128_2_1_1_10_f",
+               f8f8bf16_rowwise_grouped_128_256_128_2_1_1_10_f},
+              {"f8f8bf16_rowwise_grouped_256_32_128_2_1_1_10_f",
+               f8f8bf16_rowwise_grouped_256_32_128_2_1_1_10_f},
+              {"f8f8bf16_rowwise_grouped_256_64_128_2_1_1_10_f",
+               f8f8bf16_rowwise_grouped_256_64_128_2_1_1_10_f},
+              {"f8f8bf16_rowwise_grouped_256_128_128_2_1_1_10_f",
+               f8f8bf16_rowwise_grouped_256_128_128_2_1_1_10_f},
+              {"f8f8bf16_rowwise_grouped_256_256_128_2_1_1_10_f",
+               f8f8bf16_rowwise_grouped_256_256_128_2_1_1_10_f},
+          };
+  return kernels;
+}
+
 } // namespace fbgemm_gpu