Simplify CK FP8 Kernel Launch and enable FP16 Outputs. (#4233)

Summary:
X-link: facebookresearch/FBGEMM#1311

This diff does some template cleanup of FP8 rowwise AMD kernels, specifically moving as much specialization as possible from the individual kernels to the common header file. This makes it a lot easier to do auto-generation of kernels or add new features going forward and should be functionally the same as before.

I leverage this new cleaner infrastructure to also allow FP16 outputs in a pretty seamless way. This allows us to introduce f8f8f16_rowwise for the AMD backend, which is needed for some recommendation system use cases.

Pull Request resolved: #4233

Reviewed By: jianyuh, jiawenliu64

Differential Revision: D74770197

Pulled By: jwfromm

Author: jwfromm

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/fp8_rowwise_gemm.hip

@@ -496,7 +496,8 @@ RowwiseKernel rowwise_dispatch(int M, int N, int K) {
   return rowwise_heuristic_dispatch(M, N, K);
 }
 
-at::Tensor f8f8bf16_rowwise_wrapper(
+template <at::ScalarType OUTPUT_DTYPE>
+at::Tensor f8f8_rowwise_wrapper(
     at::Tensor XQ,
     at::Tensor WQ,
     at::Tensor x_scale,
@@ -513,6 +514,7 @@ at::Tensor f8f8bf16_rowwise_wrapper(
       (x_scale.dtype() == at::kFloat) && (w_scale.dtype() == at::kFloat),
       "Scales must be float32.");
   TORCH_CHECK(use_fast_accum, "AMD does not support disabling use_fast_accum.");
+  TORCH_CHECK(!bias.has_value(), "AMD does not support fused bias.");
 
   // Check inputs are in expected format.
   TORCH_CHECK(XQ.is_cuda() && XQ.is_contiguous());
@@ -530,7 +532,7 @@ at::Tensor f8f8bf16_rowwise_wrapper(
   // Handle case where an input dimension is zero.
   if (M == 0 || N == 0 || K == 0) {
     // Return a tensor of zeros to handle case where K is 0.
-    return at::zeros(out_sizes, XQ.options().dtype(at::kBFloat16));
+    return at::zeros(out_sizes, XQ.options().dtype(OUTPUT_DTYPE));
   }
 
   // Prepare output tensor if needed.
@@ -540,9 +542,9 @@ at::Tensor f8f8bf16_rowwise_wrapper(
     // Make sure the provided output has the proper shape and dtype.
     int Y_M = size_to_dim_(Y.dim() - 1, Y.sizes());
     TORCH_CHECK(Y_M == M && Y.sizes().vec().back() == N);
-    TORCH_CHECK(Y.dtype() == at::kBFloat16);
+    TORCH_CHECK(Y.dtype() == OUTPUT_DTYPE);
   } else {
-    Y = at::empty(out_sizes, XQ.options().dtype(at::kBFloat16));
+    Y = at::empty(out_sizes, XQ.options().dtype(OUTPUT_DTYPE));
   }
 
   RowwiseKernel rowwise_impl = rowwise_dispatch(M, N, K);
@@ -557,7 +559,19 @@ at::Tensor f8f8bf16_rowwise(
     std::optional<at::Tensor> bias,
     bool use_fast_accum) {
   // Invoke f8f8bf16 rowwise without preallocated output.
-  return f8f8bf16_rowwise_wrapper(
+  return f8f8_rowwise_wrapper<at::kBFloat16>(
+      XQ, WQ, x_scale, w_scale, bias, use_fast_accum);
+}
+
+at::Tensor f8f8f16_rowwise(
+    at::Tensor XQ,
+    at::Tensor WQ,
+    at::Tensor x_scale,
+    at::Tensor w_scale,
+    std::optional<at::Tensor> bias,
+    bool use_fast_accum) {
+  // Invoke f8f8bf16 rowwise without preallocated output.
+  return f8f8_rowwise_wrapper<at::kHalf>(
       XQ, WQ, x_scale, w_scale, bias, use_fast_accum);
 }
 
@@ -570,7 +584,7 @@ void f8f8bf16_rowwise_out(
     std::optional<at::Tensor> bias,
     bool use_fast_accum) {
   // Invoke f8f8bf16 rowwise with preallocated output.
-  f8f8bf16_rowwise_wrapper(
+  f8f8_rowwise_wrapper<at::kBFloat16>(
       XQ, WQ, x_scale, w_scale, bias, use_fast_accum, output);
 }
 

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x128x16x128_16x16_4x1_8x16x1_8x16x1_1x16x1x8_8x8x1_1x1_interwave_v2.hip

@@ -16,7 +16,7 @@ fp8_rowwise_128x128x16x128_16x16_4x1_8x16x1_8x16x1_1x16x1x8_8x8x1_1x1_interwave_
     at::Tensor w_scale,
     at::Tensor Y) {
   // A kernel that works well on small but not super tiny shapes.
-  using DeviceGemmInstance = DeviceGemmHelper<
+  return f8f8bf16_rowwise_wrapper<
       128,
       128,
       16,
@@ -32,7 +32,5 @@ fp8_rowwise_128x128x16x128_16x16_4x1_8x16x1_8x16x1_1x16x1x8_8x8x1_1x1_interwave_
       1,
       1,
       ck::BlockGemmPipelineScheduler::Interwave,
-      ck::BlockGemmPipelineVersion::v2>;
-  // Run kernel instance.
-  return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y);
+      ck::BlockGemmPipelineVersion::v2>(XQ, WQ, x_scale, w_scale, Y, 1);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x128x32x128_32x32_2x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v2.hip

@@ -15,55 +15,21 @@ fp8_rowwise_128x128x32x128_32x32_2x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor Y) {
-  // Check if this input needs to be padded.
-  int M = size_to_dim_(XQ.dim() - 1, XQ.sizes());
-  int N = WQ.size(0);
-  int K = WQ.size(1);
-  bool pad = (M % 128 != 0) || (N % 32 != 0) || (K % 128 != 0);
-
-  // This kernel seems optimal in the most purely compute bound tasks.
-  if (pad) {
-    using DeviceGemmInstance = DeviceGemmHelper<
-        128,
-        128,
-        32,
-        128,
-        32,
-        32,
-        2,
-        1,
-        S<8, 16, 1>,
-        S<8, 16, 1>,
-        S<1, 16, 1, 8>,
-        S<4, 4, 1>,
-        1,
-        1,
-        ck::BlockGemmPipelineScheduler::Intrawave,
-        ck::BlockGemmPipelineVersion::v2>;
-    // Run kernel instance.
-    return f8f8bf16_rowwise_impl<DeviceGemmInstance>(
-        XQ, WQ, x_scale, w_scale, Y);
-  } else {
-    using DeviceGemmInstance = DeviceGemmHelper<
-        128,
-        128,
-        32,
-        128,
-        32,
-        32,
-        2,
-        1,
-        S<8, 16, 1>,
-        S<8, 16, 1>,
-        S<1, 16, 1, 8>,
-        S<4, 4, 1>,
-        1,
-        1,
-        ck::BlockGemmPipelineScheduler::Intrawave,
-        ck::BlockGemmPipelineVersion::v2,
-        ck::tensor_operation::device::GemmSpecialization::Default>;
-    // Run kernel instance.
-    return f8f8bf16_rowwise_impl<DeviceGemmInstance>(
-        XQ, WQ, x_scale, w_scale, Y);
-  }
+  return f8f8bf16_rowwise_wrapper<
+      128,
+      128,
+      32,
+      128,
+      32,
+      32,
+      2,
+      1,
+      S<8, 16, 1>,
+      S<8, 16, 1>,
+      S<1, 16, 1, 8>,
+      S<4, 4, 1>,
+      1,
+      1,
+      ck::BlockGemmPipelineScheduler::Intrawave,
+      ck::BlockGemmPipelineVersion::v2>(XQ, WQ, x_scale, w_scale, Y, 1);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v2.hip

@@ -15,7 +15,7 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor Y) {
-  using DeviceGemmInstance = DeviceGemmHelper<
+  return f8f8bf16_rowwise_wrapper<
       128,
       16,
       32,
@@ -31,9 +31,5 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v
       1,
       1,
       ck::BlockGemmPipelineScheduler::Interwave,
-      ck::BlockGemmPipelineVersion::v2,
-      ck::tensor_operation::device::GemmSpecialization::Default>;
-  // Run kernel instance.
-  return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y);
+      ck::BlockGemmPipelineVersion::v2>(XQ, WQ, x_scale, w_scale, Y, 1);
 }
-

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v2_4_split_k.hip

@@ -15,7 +15,7 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor Y) {
-  using DeviceGemmInstance = DeviceGemmHelper<
+  return f8f8bf16_rowwise_wrapper<
       128,
       16,
       32,
@@ -31,8 +31,5 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v
       1,
       1,
       ck::BlockGemmPipelineScheduler::Interwave,
-      ck::BlockGemmPipelineVersion::v2,
-      ck::tensor_operation::device::GemmSpecialization::Default>;
-  // Run kernel instance.
-  return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y, 4);
+      ck::BlockGemmPipelineVersion::v2>(XQ, WQ, x_scale, w_scale, Y, 4);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v2_8_split_k.hip

@@ -15,7 +15,7 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor Y) {
-  using DeviceGemmInstance = DeviceGemmHelper<
+  return f8f8bf16_rowwise_wrapper<
       128,
       16,
       32,
@@ -31,8 +31,5 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_interwave_v
       1,
       1,
       ck::BlockGemmPipelineScheduler::Interwave,
-      ck::BlockGemmPipelineVersion::v2,
-      ck::tensor_operation::device::GemmSpecialization::Default>;
-  // Run kernel instance.
-  return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y, 8);
+      ck::BlockGemmPipelineVersion::v2>(XQ, WQ, x_scale, w_scale, Y, 8);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v1.hip

@@ -16,53 +16,21 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v
     at::Tensor w_scale,
     at::Tensor Y) {
   // The smallest kernel we have available. Works well for memory bound shapes.
-
-  // Check if this input needs to be padded.
-  int M = size_to_dim_(XQ.dim() - 1, XQ.sizes());
-  int N = WQ.size(0);
-  int K = WQ.size(1);
-  bool pad = (M % 16 != 0) || (N % 32 != 0) || (K % 128 != 0);
-  if (pad) {
-    using DeviceGemmInstance = DeviceGemmHelper<
-        128,
-        16,
-        32,
-        128,
-        16,
-        16,
-        1,
-        1,
-        S<8, 16, 1>,
-        S<8, 16, 1>,
-        S<1, 16, 1, 8>,
-        S<4, 4, 1>,
-        1,
-        1,
-        ck::BlockGemmPipelineScheduler::Intrawave,
-        ck::BlockGemmPipelineVersion::v1,
-        ck::tensor_operation::device::GemmSpecialization::MNKPadding>;
-    // Run kernel instance.
-    return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y);
-  } else{
-    using DeviceGemmInstance = DeviceGemmHelper<
-        128,
-        16,
-        32,
-        128,
-        16,
-        16,
-        1,
-        1,
-        S<8, 16, 1>,
-        S<8, 16, 1>,
-        S<1, 16, 1, 8>,
-        S<4, 4, 1>,
-        1,
-        1,
-        ck::BlockGemmPipelineScheduler::Intrawave,
-        ck::BlockGemmPipelineVersion::v1,
-        ck::tensor_operation::device::GemmSpecialization::Default>;
-    // Run kernel instance.
-    return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y);
-  }
+  return f8f8bf16_rowwise_wrapper<
+      128,
+      16,
+      32,
+      128,
+      16,
+      16,
+      1,
+      1,
+      S<8, 16, 1>,
+      S<8, 16, 1>,
+      S<1, 16, 1, 8>,
+      S<4, 4, 1>,
+      1,
+      1,
+      ck::BlockGemmPipelineScheduler::Intrawave,
+      ck::BlockGemmPipelineVersion::v1>(XQ, WQ, x_scale, w_scale, Y, 1);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v2.hip

@@ -16,52 +16,21 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v
     at::Tensor w_scale,
     at::Tensor Y) {
   // The smallest kernel we have available. Works well for memory bound shapes.
-
-  // Check if this input needs to be padded.
-  int M = size_to_dim_(XQ.dim() - 1, XQ.sizes());
-  int N = WQ.size(0);
-  int K = WQ.size(1);
-  bool pad = (M % 16 != 0) || (N % 32 != 0) || (K % 128 != 0);
-  if (pad) {
-    using DeviceGemmInstance = DeviceGemmHelper<
-        128,
-        16,
-        32,
-        128,
-        16,
-        16,
-        1,
-        1,
-        S<8, 16, 1>,
-        S<8, 16, 1>,
-        S<1, 16, 1, 8>,
-        S<4, 4, 1>,
-        1,
-        1,
-        ck::BlockGemmPipelineScheduler::Intrawave,
-        ck::BlockGemmPipelineVersion::v2>;
-    // Run kernel instance.
-    return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y);
-  } else{
-    using DeviceGemmInstance = DeviceGemmHelper<
-        128,
-        16,
-        32,
-        128,
-        16,
-        16,
-        1,
-        1,
-        S<8, 16, 1>,
-        S<8, 16, 1>,
-        S<1, 16, 1, 8>,
-        S<4, 4, 1>,
-        1,
-        1,
-        ck::BlockGemmPipelineScheduler::Intrawave,
-        ck::BlockGemmPipelineVersion::v2,
-        ck::tensor_operation::device::GemmSpecialization::Default>;
-    // Run kernel instance.
-    return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y);
-  }
+  return f8f8bf16_rowwise_wrapper<
+      128,
+      16,
+      32,
+      128,
+      16,
+      16,
+      1,
+      1,
+      S<8, 16, 1>,
+      S<8, 16, 1>,
+      S<1, 16, 1, 8>,
+      S<4, 4, 1>,
+      1,
+      1,
+      ck::BlockGemmPipelineScheduler::Intrawave,
+      ck::BlockGemmPipelineVersion::v2>(XQ, WQ, x_scale, w_scale, Y, 1);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v2_8_split_k.hip

@@ -15,7 +15,7 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor Y) {
-  using DeviceGemmInstance = DeviceGemmHelper<
+  return f8f8bf16_rowwise_wrapper<
       128,
       16,
       32,
@@ -31,8 +31,5 @@ fp8_rowwise_128x16x32x128_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4x4x1_1x1_intrawave_v
       1,
       1,
       ck::BlockGemmPipelineScheduler::Intrawave,
-      ck::BlockGemmPipelineVersion::v2,
-      ck::tensor_operation::device::GemmSpecialization::Default>;
-  // Run kernel instance.
-  return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y, 8);
+      ck::BlockGemmPipelineVersion::v2>(XQ, WQ, x_scale, w_scale, Y, 8);
 }

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise/kernels/fp8_rowwise_128x16x32x256_16x16_1x1_16x8x1_16x8x1_1x16x1x8_4x4x1_1x1_intrawave_v1.hip

@@ -15,7 +15,7 @@ fp8_rowwise_128x16x32x256_16x16_1x1_16x8x1_16x8x1_1x16x1x8_4x4x1_1x1_intrawave_v
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor Y) {
-  using DeviceGemmInstance = DeviceGemmHelper<
+  return f8f8bf16_rowwise_wrapper<
       128,
       16,
       32,
@@ -31,9 +31,5 @@ fp8_rowwise_128x16x32x256_16x16_1x1_16x8x1_16x8x1_1x16x1x8_4x4x1_1x1_intrawave_v
       1,
       1,
       ck::BlockGemmPipelineScheduler::Intrawave,
-      ck::BlockGemmPipelineVersion::v1,
-      ck::tensor_operation::device::GemmSpecialization::Default>;
-  // Run kernel instance.
-  return f8f8bf16_rowwise_impl<DeviceGemmInstance>(XQ, WQ, x_scale, w_scale, Y);
+      ck::BlockGemmPipelineVersion::v1>(XQ, WQ, x_scale, w_scale, Y, 1);
 }
-