Fix sm75, sm86 arch error of weight-only int8 gemm kernel (#59778)

wwbitejotunn · web-flow · commit 7fde5615d9ef · 2023-12-08T10:05:01.000+08:00
* fix sm75

* fix sm86
diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc
@@ -1168,10 +1168,10 @@ void WeightOnlyLinearGradInferMeta(const MetaTensor& x,
                                    const int32_t arch,
                                    MetaTensor* x_grad) {
   PADDLE_ENFORCE_EQ(
-      arch,
-      80,
+      ((arch == 80) || (arch == 86)),
+      true,
       phi::errors::InvalidArgument(
-          "Currently weightonly linear grad only support arch = 80. "));
+          "Currently weightonly linear grad only support arch = 80 or 86. "));
   x_grad->set_dims(x.dims());
   x_grad->set_dtype(x.dtype());
 }
diff --git a/paddle/phi/kernels/fusion/cutlass/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h b/paddle/phi/kernels/fusion/cutlass/cutlass_extensions/gemm/warp/mma_tensorop_dequantizer.h
@@ -173,13 +173,23 @@ class MmaTensorOpDequantizer<
     for (int mma_n_iter = 0; mma_n_iter < MmaOperator::MmaIterations::kColumn;
          ++mma_n_iter) {
       static_assert(ExpandedMmaOperandB::kElements % 2 == 0, "");
-
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+      __nv_bfloat162 scalex2;
+      scalex2.x = scale_ptr[mma_n_iter];
+      scalex2.y = scale_ptr[mma_n_iter];
+#else
       __nv_bfloat162 scalex2 = __bfloat162bfloat162(scale_ptr[mma_n_iter]);
+#endif
       __nv_bfloat162* operand_bf16x2_ptr =
           reinterpret_cast<__nv_bfloat162*>(&operand_frag_ptr[mma_n_iter]);
       CUTLASS_PRAGMA_UNROLL
       for (int ii = 0; ii < ExpandedMmaOperandB::kElements / 2; ++ii) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+        operand_bf16x2_ptr[ii].x = operand_bf16x2_ptr[ii].x * scalex2.x;
+        operand_bf16x2_ptr[ii].y = operand_bf16x2_ptr[ii].y * scalex2.y;
+#else
         operand_bf16x2_ptr[ii] = __hmul2(operand_bf16x2_ptr[ii], scalex2);
+#endif
       }
     }
   }
diff --git a/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/cutlass_heuristic.h b/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/cutlass_heuristic.h
@@ -127,6 +127,7 @@ static std::vector<CutlassTileConfig> get_candidate_tiles(
       CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64};
   std::vector<CutlassTileConfig> quant_B_configs;
   switch (sm) {
+    case 86:
     case 80:
       quant_B_configs = quant_B_configs_sm80;
       break;
@@ -154,7 +155,7 @@ static std::vector<CutlassGemmConfig> get_candidate_configs(
       is_weight_only, is_weight_only_encoder, simt_configs_only, sm);
 
   std::vector<CutlassGemmConfig> candidate_configs;
-  const int min_stages = 5;
+  const int min_stages = 2;
   const int max_stages = sm >= 80 ? 5 : 2;
 
   for (const auto& tile_config : tiles) {
diff --git a/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.cu b/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.cu
@@ -173,7 +173,7 @@ void dispatch_gemm_to_cutlass(const T* A,
   // fpA_intB. We also only instantiate configs here where threadblockShapeM ==
   // warpShapeM since those usually perform the best for mixed type gemms.
   switch (gemm_config.tile_config) {
-#ifndef USE_FPAINTB_GEMM_WITH_SM70
+#if defined(USE_FPAINTB_GEMM_WITH_SM80) || defined(USE_FPAINTB_GEMM_WITH_SM86)
     case CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64:
       dispatch_gemm_config<T,
                            WeightType,
@@ -238,7 +238,7 @@ void dispatch_gemm_to_cutlass(const T* A,
           stream,
           occupancy);
       break;
-#ifndef USE_FPAINTB_GEMM_WITH_SM70
+#if defined(USE_FPAINTB_GEMM_WITH_SM80) || defined(USE_FPAINTB_GEMM_WITH_SM86)
     case CutlassTileConfig::CtaShape128x128x64_WarpShape64x64x64:
       dispatch_gemm_config<T,
                            WeightType,
@@ -300,6 +300,84 @@ void dispatch_gemm_to_cutlass(const T* A,
   }
 }
 
+template <typename T, typename WeightType, typename arch, typename EpilogueTag>
+void dispatch_gemm_to_cutlass_sm7x(const T* A,
+                                   const WeightType* B,
+                                   const T* weight_scales,
+                                   const T* biases,
+                                   T* C,
+                                   int m,
+                                   int n,
+                                   int k,
+                                   char* workspace,
+                                   size_t workspace_bytes,
+                                   CutlassGemmConfig gemm_config,
+                                   cudaStream_t stream,
+                                   int* occupancy) {
+  // VLOG(3)<<__PRETTY_FUNCTION__;
+  // Note that SIMT configs are omitted here since they are not supported for
+  // fpA_intB. We also only instantiate configs here where threadblockShapeM ==
+  // warpShapeM since those usually perform the best for mixed type gemms.
+  switch (gemm_config.tile_config) {
+    case CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64:
+      dispatch_gemm_config<T,
+                           WeightType,
+                           arch,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<32, 128, 64>,
+                           cutlass::gemm::GemmShape<32, 32, 64>>(
+          A,
+          B,
+          weight_scales,
+          biases,
+          C,
+          m,
+          n,
+          k,
+          gemm_config,
+          workspace,
+          workspace_bytes,
+          stream,
+          occupancy);
+      break;
+    case CutlassTileConfig::CtaShape64x128x64_WarpShape64x64x64:
+      dispatch_gemm_config<T,
+                           WeightType,
+                           arch,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<64, 128, 64>,
+                           cutlass::gemm::GemmShape<64, 64, 64>>(
+          A,
+          B,
+          weight_scales,
+          biases,
+          C,
+          m,
+          n,
+          k,
+          gemm_config,
+          workspace,
+          workspace_bytes,
+          stream,
+          occupancy);
+      break;
+    case CutlassTileConfig::Undefined:
+      throw std::runtime_error(
+          "[fpA_intB][dispatch_gemm_to_cutlass] gemm config undefined.");
+      break;
+    case CutlassTileConfig::ChooseWithHeuristic:
+      throw std::runtime_error(
+          "[fpA_intB][dispatch_gemm_to_cutlass] gemm config should have "
+          "already been set by heuristic.");
+      break;
+    default:
+      throw std::runtime_error(
+          "[fpA_intB][dispatch_gemm_to_cutlass] Config is invalid for mixed "
+          "type GEMM.");
+      break;
+  }
+}
+
 template <typename T, typename WeightType>
 CutlassFpAIntBGemmRunner<T, WeightType>::CutlassFpAIntBGemmRunner() {
   // VLOG(3)<<__PRETTY_FUNCTION__;
@@ -334,41 +412,45 @@ void CutlassFpAIntBGemmRunner<T, WeightType>::dispatch_to_arch<EpilogueTag>(
   // VLOG(3)<<__PRETTY_FUNCTION__;
   if (sm_ >= 70 && sm_ < 75) {
 #if defined(USE_FPAINTB_GEMM_WITH_SM70)
-    dispatch_gemm_to_cutlass<T, WeightType, cutlass::arch::Sm70, EpilogueTag>(
-        A,
-        B,
-        weight_scales,
-        biases,
-        C,
-        m,
-        n,
-        k,
-        workspace_ptr,
-        workspace_bytes,
-        gemm_config,
-        stream,
-        occupancy);
+    dispatch_gemm_to_cutlass_sm7x<T,
+                                  WeightType,
+                                  cutlass::arch::Sm70,
+                                  EpilogueTag>(A,
+                                               B,
+                                               weight_scales,
+                                               biases,
+                                               C,
+                                               m,
+                                               n,
+                                               k,
+                                               workspace_ptr,
+                                               workspace_bytes,
+                                               gemm_config,
+                                               stream,
+                                               occupancy);
 #else
     throw std::runtime_error(
         "[CutlassFpAIntBGemmRunner][GEMM Dispatch] Arch unsupported for "
         "CUTLASS mixed type GEMM");
 #endif
   } else if (sm_ >= 75 && sm_ < 80) {
 #if defined(USE_FPAINTB_GEMM_WITH_SM75)
-    dispatch_gemm_to_cutlass<T, WeightType, cutlass::arch::Sm75, EpilogueTag>(
-        A,
-        B,
-        weight_scales,
-        biases,
-        C,
-        m,
-        n,
-        k,
-        workspace_ptr,
-        workspace_bytes,
-        gemm_config,
-        stream,
-        occupancy);
+    dispatch_gemm_to_cutlass_sm7x<T,
+                                  WeightType,
+                                  cutlass::arch::Sm75,
+                                  EpilogueTag>(A,
+                                               B,
+                                               weight_scales,
+                                               biases,
+                                               C,
+                                               m,
+                                               n,
+                                               k,
+                                               workspace_ptr,
+                                               workspace_bytes,
+                                               gemm_config,
+                                               stream,
+                                               occupancy);
 #else
     throw std::runtime_error(
         "[CutlassFpAIntBGemmRunner][GEMM Dispatch] Arch unsupported for "
diff --git a/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/fpA_intB_gemm/generic_mixed_gemm_kernelLauncher.py b/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/fpA_intB_gemm/generic_mixed_gemm_kernelLauncher.py
@@ -104,7 +104,7 @@
     "cutlass::gemm::GemmShape<32, 32, 64>",
     "cutlass::gemm::GemmShape<64, 64, 64>",
 ]
-StagesList = {70: [2], 75: [2, 3, 4, 5], 80: [2, 3, 4, 5]}
+StagesList = {70: [2], 75: [2], 80: [2, 3, 4, 5]}
 
 ElementTypes = {"fp16": "half", "bf16": "__nv_bfloat16"}
 Archs = {
@@ -179,7 +179,7 @@ def generate_source_cu(
     all_code = CommonHead
     ThreadblockShapes_arch = ThreadblockShapes
     WarpShapes_arch = WarpShapes
-    if arch == 70:
+    if arch < 80:
         ThreadblockShapes_arch = ThreadblockShapes_sm70
         WarpShapes_arch = WarpShapes_sm70
     for WeightType in WeightTypes:
diff --git a/paddle/phi/kernels/gpu/weight_only_linear_grad_kernel.cu b/paddle/phi/kernels/gpu/weight_only_linear_grad_kernel.cu
@@ -36,10 +36,10 @@ void WeightOnlyLinearGradKernel(const Context& dev_ctx,
                                 DenseTensor* x_grad) {
 #if defined(PADDLE_WITH_CUTLASS)
   PADDLE_ENFORCE_EQ(
-      arch,
-      80,
+      ((arch == 80) || (arch == 86)),
+      true,
       phi::errors::InvalidArgument(
-          "Currently weightonly linear grad only support arch = 80. "));
+          "Currently weightonly linear grad only support arch = 80 or 86. "));
 
   int n = weight_scale.dims()[0];
   int k = weight.dims()[1];
diff --git a/paddle/phi/kernels/gpu/weight_only_linear_kernel.cu b/paddle/phi/kernels/gpu/weight_only_linear_kernel.cu
@@ -34,7 +34,7 @@ void WeightOnlyLinearKernel(const Context& dev_ctx,
                             DenseTensor* out) {
 #if defined(PADDLE_WITH_CUTLASS)
   PADDLE_ENFORCE_EQ(
-      ((arch == 80) || (arch == 70)),
+      ((arch == 80) || (arch == 70) || (arch == 75) || (arch == 86)),
       true,
       phi::errors::InvalidArgument("Currently, arch only support 70, 80."));
 #else
diff --git a/python/paddle/nn/quant/quantized_linear.py b/python/paddle/nn/quant/quantized_linear.py
@@ -172,8 +172,8 @@ def weight_only_linear(
         arch = _get_arch_info()
 
     assert (
-        arch == 70 or arch == 80
-    ), "Currently weight_quantize only support SM70/80. "
+        arch == 70 or arch == 80 or arch == 86 or arch == 75
+    ), f"Currently weight_quantize only support SM70/75/80/86. but got {arch} "
 
     if in_dynamic_mode():
         out = _C_ops.weight_only_linear(
diff --git a/test/quantization/test_weight_only_linear.py b/test/quantization/test_weight_only_linear.py
@@ -57,9 +57,7 @@ def convert_uint16_to_float(in_list):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or get_cuda_version() < 11020
-    or paddle.device.cuda.get_device_capability()[0] < 8,
+    not core.is_compiled_with_cuda() or get_cuda_version() < 11020,
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearTestCase(unittest.TestCase):
@@ -216,9 +214,7 @@ def test_weight_only_linear(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or get_cuda_version() < 11020
-    or paddle.device.cuda.get_device_capability()[0] < 8,
+    not core.is_compiled_with_cuda() or get_cuda_version() < 11020,
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearTestCase1(WeightOnlyLinearTestCase):
@@ -229,9 +225,7 @@ def config(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or get_cuda_version() < 11020
-    or paddle.device.cuda.get_device_capability()[0] < 8,
+    not core.is_compiled_with_cuda() or get_cuda_version() < 11020,
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearTestCase2(WeightOnlyLinearTestCase):
@@ -270,9 +264,7 @@ def config(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or get_cuda_version() < 11020
-    or paddle.device.cuda.get_device_capability()[0] < 8,
+    not core.is_compiled_with_cuda() or get_cuda_version() < 11020,
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearTestCase5(WeightOnlyLinearTestCase):
@@ -298,9 +290,7 @@ def config(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or get_cuda_version() < 11020
-    or paddle.device.cuda.get_device_capability()[0] < 8,
+    not core.is_compiled_with_cuda() or get_cuda_version() < 11020,
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearTestCase7(WeightOnlyLinearTestCase):
@@ -313,9 +303,7 @@ def config(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or get_cuda_version() < 11020
-    or paddle.device.cuda.get_device_capability()[0] < 8,
+    not core.is_compiled_with_cuda() or get_cuda_version() < 11020,
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearTestCase8(WeightOnlyLinearTestCase):
@@ -360,9 +348,7 @@ def config(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda()
-    or get_cuda_version() < 11020
-    or paddle.device.cuda.get_device_capability()[0] < 8,
+    not core.is_compiled_with_cuda() or get_cuda_version() < 11020,
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearTestCaseStatic(WeightOnlyLinearTestCase):
