Remove e5m2 from f8f8bf16_rowwise_batched (pytorch#4908)

Summary:
Pull Request resolved: pytorch#4908

X-link: facebookresearch/FBGEMM#1934

Fully remove this from f8f8bf16_rowwise_batched, we have already removed most of the kernel instances, this just cleans up the remaining one.

Also moved some `TORCH_CHECK` into the main non-templated API level, as that would reduce some binary size (although likely trivial).

Reviewed By: jiawenliu64

Differential Revision: D82976652

fbshipit-source-id: 9fe1058a6993942d097d27b272368d172f5ca55f

Author: cthi

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

@@ -32,23 +32,27 @@ at::Tensor dispatch_fp8_rowwise_batched_kernel(
     at::Tensor w_scale, // FP32
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
-  const int arch = getDeviceArch();
-
   TORCH_CHECK(
       (XQ.dim() == 3 && WQ.dim() == 3),
       "FP8 rowwise batched GEMM only supports 3D inputs");
   int M, N;
   M = XQ.size(1);
   N = WQ.size(1);
 
-  const bool use_e5m2 = XQ.dtype() == at::kFloat8_e5m2;
-  if (use_e5m2) {
+  TORCH_CHECK(XQ.is_cuda() && XQ.is_contiguous());
+  TORCH_CHECK(WQ.is_cuda() && WQ.is_contiguous());
+  TORCH_CHECK(XQ.dtype() == at::kFloat8_e4m3fn, "XQ must be FP8 e4m3fn");
+  TORCH_CHECK(WQ.dtype() == at::kFloat8_e4m3fn, "WQ must be FP8 e4m3fn");
+  TORCH_CHECK(
+      x_scale.dtype() == at::kFloat && w_scale.dtype() == at::kFloat,
+      "Scale tensors must be float32.");
+  if (bias.has_value()) {
     TORCH_CHECK(
-        arch == 9, "f8f8bf16_rowwise_batched only supports FP8 e5m2 on SM90");
-    return f8f8bf16_rowwise_batched_64_128_128_2_1_1_9_f_e5m2(
-        XQ, WQ, x_scale, w_scale, bias, output);
+        bias.value().dtype() == at::kFloat,
+        "Bias type must be float32 if provided.");
   }
 
+  const int arch = getDeviceArch();
   if (arch == 10) {
     if ((M * N <= 4096 * 4096) || (N % 256 > 0 && M % 256 == 0) ||
         (M % 256 > 0 && N % 256 > 0) || M >= 1024 && N >= 1024) {

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched/f8f8bf16_rowwise_batched_128_128_128_1_2_1_10_t.cu

@@ -18,16 +18,8 @@ at::Tensor f8f8bf16_rowwise_batched_128_128_128_1_2_1_10_t(
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
   // Dispatch this kernel to the correct underlying implementation.
-  return f8f8bf16_rowwise_batched_wrapper<
-      128,
-      128,
-      128,
-      1,
-      2,
-      1,
-      10,
-      true,
-      cutlass::float_e4m3_t>(XQ, WQ, x_scale, w_scale, bias, output);
+  return f8f8bf16_rowwise_batched_impl<128, 128, 128, 1, 2, 1, 10, true>(
+      XQ, WQ, x_scale, w_scale, bias, output);
 }
 
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched/f8f8bf16_rowwise_batched_128_128_128_1_2_1_9_t.cu

@@ -18,16 +18,8 @@ at::Tensor f8f8bf16_rowwise_batched_128_128_128_1_2_1_9_t(
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
   // Dispatch this kernel to the correct underlying implementation.
-  return f8f8bf16_rowwise_batched_wrapper<
-      128,
-      128,
-      128,
-      1,
-      2,
-      1,
-      9,
-      true,
-      cutlass::float_e4m3_t>(XQ, WQ, x_scale, w_scale, bias, output);
+  return f8f8bf16_rowwise_batched_impl<128, 128, 128, 1, 2, 1, 9, true>(
+      XQ, WQ, x_scale, w_scale, bias, output);
 }
 
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched/f8f8bf16_rowwise_batched_128_128_128_2_1_1_10_t.cu

@@ -18,16 +18,8 @@ at::Tensor f8f8bf16_rowwise_batched_128_128_128_2_1_1_10_t(
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
   // Dispatch this kernel to the correct underlying implementation.
-  return f8f8bf16_rowwise_batched_wrapper<
-      128,
-      128,
-      128,
-      2,
-      1,
-      1,
-      10,
-      true,
-      cutlass::float_e4m3_t>(XQ, WQ, x_scale, w_scale, bias, output);
+  return f8f8bf16_rowwise_batched_impl<128, 128, 128, 2, 1, 1, 10, true>(
+      XQ, WQ, x_scale, w_scale, bias, output);
 }
 
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched/f8f8bf16_rowwise_batched_128_128_128_2_1_1_9_t.cu

@@ -18,16 +18,8 @@ at::Tensor f8f8bf16_rowwise_batched_128_128_128_2_1_1_9_t(
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
   // Dispatch this kernel to the correct underlying implementation.
-  return f8f8bf16_rowwise_batched_wrapper<
-      128,
-      128,
-      128,
-      2,
-      1,
-      1,
-      9,
-      true,
-      cutlass::float_e4m3_t>(XQ, WQ, x_scale, w_scale, bias, output);
+  return f8f8bf16_rowwise_batched_impl<128, 128, 128, 2, 1, 1, 9, true>(
+      XQ, WQ, x_scale, w_scale, bias, output);
 }
 
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched/f8f8bf16_rowwise_batched_64_128_128_1_2_1_10_f.cu

@@ -18,16 +18,8 @@ at::Tensor f8f8bf16_rowwise_batched_64_128_128_1_2_1_10_f(
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
   // Dispatch this kernel to the correct underlying implementation.
-  return f8f8bf16_rowwise_batched_wrapper<
-      64,
-      128,
-      128,
-      1,
-      2,
-      1,
-      10,
-      false,
-      cutlass::float_e4m3_t>(XQ, WQ, x_scale, w_scale, bias, output);
+  return f8f8bf16_rowwise_batched_impl<64, 128, 128, 1, 2, 1, 10, false>(
+      XQ, WQ, x_scale, w_scale, bias, output);
 }
 
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched/f8f8bf16_rowwise_batched_64_128_128_1_2_1_9_f.cu

@@ -18,16 +18,8 @@ at::Tensor f8f8bf16_rowwise_batched_64_128_128_1_2_1_9_f(
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
   // Dispatch this kernel to the correct underlying implementation.
-  return f8f8bf16_rowwise_batched_wrapper<
-      64,
-      128,
-      128,
-      1,
-      2,
-      1,
-      9,
-      false,
-      cutlass::float_e4m3_t>(XQ, WQ, x_scale, w_scale, bias, output);
+  return f8f8bf16_rowwise_batched_impl<64, 128, 128, 1, 2, 1, 9, false>(
+      XQ, WQ, x_scale, w_scale, bias, output);
 }
 
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched/f8f8bf16_rowwise_batched_64_128_128_2_1_1_10_f.cu

@@ -18,16 +18,8 @@ at::Tensor f8f8bf16_rowwise_batched_64_128_128_2_1_1_10_f(
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
   // Dispatch this kernel to the correct underlying implementation.
-  return f8f8bf16_rowwise_batched_wrapper<
-      64,
-      128,
-      128,
-      2,
-      1,
-      1,
-      10,
-      false,
-      cutlass::float_e4m3_t>(XQ, WQ, x_scale, w_scale, bias, output);
+  return f8f8bf16_rowwise_batched_impl<64, 128, 128, 2, 1, 1, 10, false>(
+      XQ, WQ, x_scale, w_scale, bias, output);
 }
 
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_batched/f8f8bf16_rowwise_batched_64_128_128_2_1_1_9_f.cu

@@ -18,16 +18,8 @@ at::Tensor f8f8bf16_rowwise_batched_64_128_128_2_1_1_9_f(
     std::optional<at::Tensor> bias = std::nullopt,
     std::optional<at::Tensor> output = std::nullopt) {
   // Dispatch this kernel to the correct underlying implementation.
-  return f8f8bf16_rowwise_batched_wrapper<
-      64,
-      128,
-      128,
-      2,
-      1,
-      1,
-      9,
-      false,
-      cutlass::float_e4m3_t>(XQ, WQ, x_scale, w_scale, bias, output);
+  return f8f8bf16_rowwise_batched_impl<64, 128, 128, 2, 1, 1, 9, false>(
+      XQ, WQ, x_scale, w_scale, bias, output);
 }
 
 } // namespace fbgemm_gpu