Remove e5m2 from f8f8bf16_rowwise (#4907)

Summary:
Pull Request resolved: #4907

X-link: facebookresearch/FBGEMM#1931

- e5m2 is not used, remove it to save code size and avoid more build issues
- possibly e5m2 was messed up to begin with, we hardcoded `using ElementB = cutlass::float_e4m3_t`

Reviewed By: jiawenliu64, jwfromm

Differential Revision: D82965586

fbshipit-source-id: 61e61bd316bc0664cdbca8aa761e44f1232d0ee2

Author: cthi

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

@@ -14,15 +14,18 @@
 #include "f8f8bf16_rowwise/f8f8bf16_rowwise_manifest.cuh"
 #include "fbgemm_gpu/quantize/tuning_cache.hpp"
 #include "fbgemm_gpu/quantize/utils.h"
+#include "fbgemm_gpu/quantize/utils_gpu.h"
 
 namespace fbgemm_gpu {
 
 #if CUDART_VERSION >= 12000
 
 // FP8 Rowwise Cutlass kernel dispatch.
 Kernel_f8f8bf16_rowwise
-get_kernel_via_heuristic(int arch, int M, int N, int K, bool use_fast_accum) {
+get_kernel_via_heuristic(int M, int N, int K, bool use_fast_accum) {
   // Use shape heuristics to dispatch to optimized kernel configuration.
+  const int arch = getDeviceArch();
+
   if (arch == 10) {
     if (M <= 128) {
       if (N <= 1024) {
@@ -115,7 +118,6 @@ get_kernel_via_heuristic(int arch, int M, int N, int K, bool use_fast_accum) {
 }
 
 Kernel_f8f8bf16_rowwise get_kernel_via_tuning(
-    int arch,
     int M,
     int N,
     int K,
@@ -134,6 +136,7 @@ Kernel_f8f8bf16_rowwise get_kernel_via_tuning(
   // Use (M, N, K) shape as the key.
   const std::string shape_key =
       std::to_string(M) + "_" + std::to_string(N) + "_" + std::to_string(K);
+  const int arch = getDeviceArch();
   const auto& kernels = get_f8f8bf16_rowwise_kernels(arch);
   auto kernel = cache.findBestKernelMaybeAutotune(
       shape_key,
@@ -158,44 +161,19 @@ at::Tensor dispatch_fp8_rowwise_kernel(
     bool use_fast_accum,
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
+  TORCH_CHECK(XQ.dtype() == at::kFloat8_e4m3fn);
+
   int M = size_to_dim_(XQ.dim() - 1, XQ.sizes());
   int N = size_to_dim_(WQ.dim() - 1, WQ.sizes());
   int K = XQ.size(-1);
 
-  static int arch = -1;
-  // Avoid expensive cudaGetDeviceProperties call.
-  if (arch < 0) {
-    cudaDeviceProp prop;
-    cudaGetDeviceProperties(&prop, 0);
-    if (prop.major >= 10) {
-      arch = 10;
-      int runtimeVersion;
-      C10_CUDA_CHECK(cudaRuntimeGetVersion(&runtimeVersion));
-      TORCH_CHECK(
-          runtimeVersion >= 12080,
-          "FP8 GEMM on sm100a or above requires cuda >= 12.8");
-    } else {
-      arch = 9;
-    }
-  }
-
   // Select kernel to run via heuristics or tuning.
   auto kernel = [&]() {
     if (std::getenv("FBGEMM_AUTOTUNE_ENABLE")) {
       return get_kernel_via_tuning(
-          arch,
-          M,
-          N,
-          K,
-          XQ,
-          WQ,
-          x_scale,
-          w_scale,
-          use_fast_accum,
-          bias,
-          output);
+          M, N, K, XQ, WQ, x_scale, w_scale, use_fast_accum, bias, output);
     } else {
-      return get_kernel_via_heuristic(arch, M, N, K, use_fast_accum);
+      return get_kernel_via_heuristic(M, N, K, use_fast_accum);
     }
   }();
   // Invoke kernel

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise/f8f8bf16_rowwise_common.cuh

@@ -36,7 +36,6 @@ template <
     bool PONG,
     bool COOP,
     bool FAST_ACCUM,
-    typename INPUT_DTYPE,
     typename BIAS_DTYPE>
 at::Tensor f8f8bf16_rowwise_impl(
     at::Tensor XQ, // FP8
@@ -76,7 +75,7 @@ at::Tensor f8f8bf16_rowwise_impl(
     Y = at::empty(out_sizes, XQ.options().dtype(at::kBFloat16));
   }
 
-  using ElementA = INPUT_DTYPE;
+  using ElementA = cutlass::float_e4m3_t;
   using LayoutA = cutlass::layout::RowMajor;
   constexpr int AlignmentInputA = 16 / sizeof(ElementA);
 
@@ -351,131 +350,61 @@ at::Tensor f8f8bf16_rowwise_wrapper(
   bool bf16_bias = bias.has_value() && bias.value().dtype() == at::kBFloat16;
 
   // Templatize based on input dtype.
-  bool use_e5m2 = XQ.dtype() == at::kFloat8_e5m2;
-
   if (bf16_bias) {
     if (use_fast_accum) {
-      if (use_e5m2) {
-        return f8f8bf16_rowwise_impl<
-            TB_M,
-            TB_N,
-            TB_K,
-            TBS_M,
-            TBS_N,
-            TBS_K,
-            ARCH,
-            PONG,
-            COOP,
-            true,
-            cutlass::float_e5m2_t,
-            cutlass::bfloat16_t>(XQ, WQ, x_scale, w_scale, bias, output);
-      } else {
-        return f8f8bf16_rowwise_impl<
-            TB_M,
-            TB_N,
-            TB_K,
-            TBS_M,
-            TBS_N,
-            TBS_K,
-            ARCH,
-            PONG,
-            COOP,
-            true,
-            cutlass::float_e4m3_t,
-            cutlass::bfloat16_t>(XQ, WQ, x_scale, w_scale, bias, output);
-      }
+      return f8f8bf16_rowwise_impl<
+          TB_M,
+          TB_N,
+          TB_K,
+          TBS_M,
+          TBS_N,
+          TBS_K,
+          ARCH,
+          PONG,
+          COOP,
+          true,
+          cutlass::bfloat16_t>(XQ, WQ, x_scale, w_scale, bias, output);
     } else {
-      if (use_e5m2) {
-        return f8f8bf16_rowwise_impl<
-            TB_M,
-            TB_N,
-            TB_K,
-            TBS_M,
-            TBS_N,
-            TBS_K,
-            ARCH,
-            PONG,
-            COOP,
-            false,
-            cutlass::float_e5m2_t,
-            cutlass::bfloat16_t>(XQ, WQ, x_scale, w_scale, bias, output);
-      } else {
-        return f8f8bf16_rowwise_impl<
-            TB_M,
-            TB_N,
-            TB_K,
-            TBS_M,
-            TBS_N,
-            TBS_K,
-            ARCH,
-            PONG,
-            COOP,
-            false,
-            cutlass::float_e4m3_t,
-            cutlass::bfloat16_t>(XQ, WQ, x_scale, w_scale, bias, output);
-      }
+      return f8f8bf16_rowwise_impl<
+          TB_M,
+          TB_N,
+          TB_K,
+          TBS_M,
+          TBS_N,
+          TBS_K,
+          ARCH,
+          PONG,
+          COOP,
+          false,
+          cutlass::bfloat16_t>(XQ, WQ, x_scale, w_scale, bias, output);
     }
   } else {
     if (use_fast_accum) {
-      if (use_e5m2) {
-        return f8f8bf16_rowwise_impl<
-            TB_M,
-            TB_N,
-            TB_K,
-            TBS_M,
-            TBS_N,
-            TBS_K,
-            ARCH,
-            PONG,
-            COOP,
-            true,
-            cutlass::float_e5m2_t,
-            float>(XQ, WQ, x_scale, w_scale, bias, output);
-      } else {
-        return f8f8bf16_rowwise_impl<
-            TB_M,
-            TB_N,
-            TB_K,
-            TBS_M,
-            TBS_N,
-            TBS_K,
-            ARCH,
-            PONG,
-            COOP,
-            true,
-            cutlass::float_e4m3_t,
-            float>(XQ, WQ, x_scale, w_scale, bias, output);
-      }
+      return f8f8bf16_rowwise_impl<
+          TB_M,
+          TB_N,
+          TB_K,
+          TBS_M,
+          TBS_N,
+          TBS_K,
+          ARCH,
+          PONG,
+          COOP,
+          true,
+          float>(XQ, WQ, x_scale, w_scale, bias, output);
     } else {
-      if (use_e5m2) {
-        return f8f8bf16_rowwise_impl<
-            TB_M,
-            TB_N,
-            TB_K,
-            TBS_M,
-            TBS_N,
-            TBS_K,
-            ARCH,
-            PONG,
-            COOP,
-            false,
-            cutlass::float_e5m2_t,
-            float>(XQ, WQ, x_scale, w_scale, bias, output);
-      } else {
-        return f8f8bf16_rowwise_impl<
-            TB_M,
-            TB_N,
-            TB_K,
-            TBS_M,
-            TBS_N,
-            TBS_K,
-            ARCH,
-            PONG,
-            COOP,
-            false,
-            cutlass::float_e4m3_t,
-            float>(XQ, WQ, x_scale, w_scale, bias, output);
-      }
+      return f8f8bf16_rowwise_impl<
+          TB_M,
+          TB_N,
+          TB_K,
+          TBS_M,
+          TBS_N,
+          TBS_K,
+          ARCH,
+          PONG,
+          COOP,
+          false,
+          float>(XQ, WQ, x_scale, w_scale, bias, output);
     }
   }
 }

fbgemm_gpu/experimental/gen_ai/test/quantize/quantize_test.py

@@ -349,9 +349,7 @@ def test_f8f8bf16(self, kernel: str, use_fast_accum: bool) -> None:
             ["rowwise", "blockwise"]
             + (["tensorwise_broadcast", "tensorwise"] if torch.version.cuda else [])
         ),
-        QType=(
-            st.sampled_from([fp8_e4m3, fp8_e5m2] if torch.version.cuda else [fp8_e4m3])
-        ),
+        QType=(st.sampled_from([fp8_e4m3])),
         Bias=st.sampled_from([True, False]),
         CudaGraph=st.sampled_from([True, False]),
         UseTriton=st.sampled_from([False] + ([True] if torch.version.cuda else [])),