init (#6642)

Co-authored-by: gongweibao <gognweibao@baidu.com>

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/cutlass_gemm_caller.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
 #pragma once
 
 // clang-format will break include orders
@@ -21,15 +22,16 @@
 
 namespace fastdeploy::c3x {
 
-static inline cute::Shape<int, int, int, int>
-get_problem_shape(paddle::Tensor const &a, paddle::Tensor const &b) {
+static inline cute::Shape<int, int, int, int> get_problem_shape(
+    paddle::Tensor const &a, paddle::Tensor const &b) {
   int32_t m = a.dims()[0], n = b.dims()[0], k = a.dims()[1];
   return {m, n, k, 1};
 }
 
 template <typename GemmKernel>
 void cutlass_gemm_caller(
-    phi::Place device, cute::Shape<int, int, int, int> prob_shape,
+    phi::Place device,
+    cute::Shape<int, int, int, int> prob_shape,
     typename GemmKernel::MainloopArguments mainloop_args,
     typename GemmKernel::EpilogueArguments epilogue_args,
     typename GemmKernel::TileSchedulerArguments scheduler = {}) {
@@ -57,7 +59,8 @@ void cutlass_gemm_caller(
 }
 
 template <typename Gemm, typename... EpilogueArgs>
-void cutlass_gemm_caller(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_gemm_caller(paddle::Tensor &out,
+                         paddle::Tensor const &a,
                          paddle::Tensor const &b,
                          EpilogueArgs &&...epilogue_params) {
   using ElementAB = typename Gemm::ElementAB;
@@ -86,17 +89,20 @@ void cutlass_gemm_caller(paddle::Tensor &out, paddle::Tensor const &a,
 
   auto a_ptr = static_cast<ElementAB *>(const_cast<void *>(a.data()));
   auto b_ptr = static_cast<ElementAB *>(const_cast<void *>(b.data()));
-  typename GemmKernel::MainloopArguments mainloop_args{a_ptr, a_stride, b_ptr,
-                                                       b_stride};
+  typename GemmKernel::MainloopArguments mainloop_args{
+      a_ptr, a_stride, b_ptr, b_stride};
 
   auto c_ptr = static_cast<ElementD *>(const_cast<void *>(out.data()));
   typename GemmKernel::EpilogueArguments epilogue_args{
       Gemm::Epilogue::prepare_args(
           std::forward<EpilogueArgs>(epilogue_params)...),
-      c_ptr, c_stride, c_ptr, d_stride};
+      c_ptr,
+      c_stride,
+      c_ptr,
+      d_stride};
 
-  cutlass_gemm_caller<GemmKernel>(a.place(), prob_shape, mainloop_args,
-                                  epilogue_args);
+  cutlass_gemm_caller<GemmKernel>(
+      a.place(), prob_shape, mainloop_args, epilogue_args);
 }
 
-} // namespace fastdeploy::c3x
+}  // namespace fastdeploy::c3x

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm.cuh
 
 #pragma once
 
@@ -31,16 +32,19 @@ using namespace cute;
 
 namespace fastdeploy {
 
-template <typename ElementAB_, typename ElementD_,
-          template <typename, typename, typename> typename Epilogue_,
-          typename TileShape, typename ClusterShape, typename KernelSchedule,
+template <typename ElementAB_,
+          typename ElementD_,
+          template <typename, typename, typename>
+          typename Epilogue_,
+          typename TileShape,
+          typename ClusterShape,
+          typename KernelSchedule,
           typename EpilogueSchedule>
 struct cutlass_3x_gemm {
   using ElementAB = ElementAB_;
   using ElementD = ElementD_;
-  using ElementAcc =
-      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
-                                float>::type;
+  using ElementAcc = typename std::
+      conditional<std::is_same_v<ElementAB, int8_t>, int32_t, float>::type;
 
   using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
 
@@ -57,10 +61,21 @@ struct cutlass_3x_gemm {
 
   using CollectiveEpilogue =
       typename cutlass::epilogue::collective::CollectiveBuilder<
-          cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape,
-          ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
-          ElementAcc, float, ElementC, StrideC, AlignmentCD, ElementD, StrideD,
-          AlignmentCD, EpilogueSchedule, EVTCompute>::CollectiveOp;
+          cutlass::arch::Sm90,
+          cutlass::arch::OpClassTensorOp,
+          TileShape,
+          ClusterShape,
+          cutlass::epilogue::collective::EpilogueTileAuto,
+          ElementAcc,
+          float,
+          ElementC,
+          StrideC,
+          AlignmentCD,
+          ElementD,
+          StrideD,
+          AlignmentCD,
+          EpilogueSchedule,
+          EVTCompute>::CollectiveOp;
 
   static constexpr size_t CEStorageSize =
       sizeof(typename CollectiveEpilogue::SharedStorage);
@@ -78,16 +93,22 @@ struct cutlass_3x_gemm {
           KernelSchedule>::CollectiveOp;
   // clang-format on
 
-  using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
-      cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
-      cutlass::gemm::PersistentScheduler>>;
+  using KernelType = enable_sm90_or_later<
+      cutlass::gemm::kernel::GemmUniversal<cute::Shape<int, int, int, int>,
+                                           CollectiveMainloop,
+                                           CollectiveEpilogue,
+                                           cutlass::gemm::PersistentScheduler>>;
 
   struct GemmKernel : public KernelType {};
 };
 
-template <typename ElementAB_, typename ElementD_,
-          template <typename, typename, typename> typename Epilogue_,
-          typename TileShape, typename ClusterShape, typename KernelSchedule,
+template <typename ElementAB_,
+          typename ElementD_,
+          template <typename, typename, typename>
+          typename Epilogue_,
+          typename TileShape,
+          typename ClusterShape,
+          typename KernelSchedule,
           typename EpilogueSchedule>
 struct cutlass_3x_gemm_sm100 {
   using ElementAB = ElementAB_;
@@ -108,9 +129,8 @@ struct cutlass_3x_gemm_sm100 {
   using LayoutD = cutlass::layout::RowMajor;
   static constexpr int AlignmentD = AlignmentC;
 
-  using ElementAcc =
-      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
-                                float>::type;
+  using ElementAcc = typename std::
+      conditional<std::is_same_v<ElementAB, int8_t>, int32_t, float>::type;
   using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
 
   // MMA type
@@ -127,23 +147,44 @@ struct cutlass_3x_gemm_sm100 {
 
   using CollectiveEpilogue =
       typename cutlass::epilogue::collective::CollectiveBuilder<
-          cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, TileShape,
-          ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
-          ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC,
-          ElementD, LayoutD, AlignmentD, EpilogueSchedule,
+          cutlass::arch::Sm100,
+          cutlass::arch::OpClassTensorOp,
+          TileShape,
+          ClusterShape,
+          cutlass::epilogue::collective::EpilogueTileAuto,
+          ElementAccumulator,
+          ElementCompute,
+          ElementC,
+          LayoutC,
+          AlignmentC,
+          ElementD,
+          LayoutD,
+          AlignmentD,
+          EpilogueSchedule,
           EVTCompute>::CollectiveOp;
 
   using CollectiveMainloop =
       typename cutlass::gemm::collective::CollectiveBuilder<
-          cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, ElementAB,
-          LayoutA, AlignmentA, ElementAB, LayoutB, AlignmentB,
-          ElementAccumulator, TileShape, ClusterShape,
+          cutlass::arch::Sm100,
+          cutlass::arch::OpClassTensorOp,
+          ElementAB,
+          LayoutA,
+          AlignmentA,
+          ElementAB,
+          LayoutB,
+          AlignmentB,
+          ElementAccumulator,
+          TileShape,
+          ClusterShape,
           cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
               sizeof(typename CollectiveEpilogue::SharedStorage))>,
           KernelSchedule>::CollectiveOp;
 
-  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
-      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>,
+                                           CollectiveMainloop,
+                                           CollectiveEpilogue,
+                                           void>;
 };
 
-} // namespace fastdeploy
+}  // namespace fastdeploy

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_azp_sm90_int8.cu

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_azp_sm90_int8.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_azp_sm90_int8.cu
 
 // clang-format will break include orders
 // clang-format off
@@ -10,18 +11,22 @@
 namespace fastdeploy {
 
 void cutlass_scaled_mm_azp_sm90_int8(
-    paddle::Tensor &out, paddle::Tensor const &a, paddle::Tensor const &b,
-    paddle::Tensor const &a_scales, paddle::Tensor const &b_scales,
-    paddle::Tensor const &azp_adj, paddle::optional<paddle::Tensor> const &azp,
+    paddle::Tensor &out,
+    paddle::Tensor const &a,
+    paddle::Tensor const &b,
+    paddle::Tensor const &a_scales,
+    paddle::Tensor const &b_scales,
+    paddle::Tensor const &azp_adj,
+    paddle::optional<paddle::Tensor> const &azp,
     paddle::optional<paddle::Tensor> const &bias) {
   if (azp) {
     return cutlass_scaled_mm_sm90_int8_epilogue<
-        c3x::ScaledEpilogueBiasAzpToken>(out, a, b, a_scales, b_scales, azp_adj,
-                                         *azp, bias);
+        c3x::ScaledEpilogueBiasAzpToken>(
+        out, a, b, a_scales, b_scales, azp_adj, *azp, bias);
   } else {
     return cutlass_scaled_mm_sm90_int8_epilogue<c3x::ScaledEpilogueBiasAzp>(
         out, a, b, a_scales, b_scales, azp_adj, bias);
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_helper.hpp

@@ -1,34 +1,38 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_helper.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_helper.hpp
 
 #include "helper.h"
 
 template <typename Fp8Func, typename Int8Func>
-void dispatch_scaled_mm(paddle::Tensor &c, paddle::Tensor const &a,
-                        paddle::Tensor const &b, paddle::Tensor const &a_scales,
+void dispatch_scaled_mm(paddle::Tensor &c,
+                        paddle::Tensor const &a,
+                        paddle::Tensor const &b,
+                        paddle::Tensor const &a_scales,
                         paddle::Tensor const &b_scales,
                         paddle::optional<paddle::Tensor> const &bias,
-                        Fp8Func fp8_func, Int8Func int8_func) {
-    PD_CHECK(a_scales.dtype() == paddle::DataType::FLOAT32);
-    PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
+                        Fp8Func fp8_func,
+                        Int8Func int8_func) {
+  PD_CHECK(a_scales.dtype() == paddle::DataType::FLOAT32);
+  PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
 
-    int M = a.dims()[0], N = b.dims()[0], K = a.dims()[1];
+  int M = a.dims()[0], N = b.dims()[0], K = a.dims()[1];
 
-    if ((a_scales.numel() == 1 || a_scales.numel() == a.dims()[0]) &&
-        (b_scales.numel() == 1 || b_scales.numel() == b.dims()[0])) {
-      // Standard per-tensor/per-token/per-channel scaling
-      PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
-      if (a.dtype() == phi::DataType::FLOAT8_E4M3FN) {
-        fp8_func(c, a, b, a_scales, b_scales, bias);
-      } else {
-        PD_CHECK(a.dtype() == paddle::DataType::INT8);
-        if constexpr (!std::is_same_v<Int8Func, std::nullptr_t>) {
-          int8_func(c, a, b, a_scales, b_scales, bias);
-        } else {
-          PD_CHECK(false, "Int8 not supported for this architecture");
-        }
-      }
+  if ((a_scales.numel() == 1 || a_scales.numel() == a.dims()[0]) &&
+      (b_scales.numel() == 1 || b_scales.numel() == b.dims()[0])) {
+    // Standard per-tensor/per-token/per-channel scaling
+    PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
+    if (a.dtype() == phi::DataType::FLOAT8_E4M3FN) {
+      fp8_func(c, a, b, a_scales, b_scales, bias);
     } else {
-      PADDLE_THROW(phi::errors::Unimplemented(
-          "No kernel for this combination of input dtypes is implemented."));
+      PD_CHECK(a.dtype() == paddle::DataType::INT8);
+      if constexpr (!std::is_same_v<Int8Func, std::nullptr_t>) {
+        int8_func(c, a, b, a_scales, b_scales, bias);
+      } else {
+        PD_CHECK(false, "Int8 not supported for this architecture");
+      }
     }
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "No kernel for this combination of input dtypes is implemented."));
+  }
 }

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_kernels.hpp

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_kernels.hpp
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_kernels.hpp
 
 #pragma once
 
@@ -6,30 +7,35 @@
 
 namespace fastdeploy {
 
-void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out,
+                                paddle::Tensor const &a,
                                 paddle::Tensor const &b,
                                 paddle::Tensor const &a_scales,
                                 paddle::Tensor const &b_scales,
                                 paddle::optional<paddle::Tensor> const &bias);
 
-void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out,
+                                 paddle::Tensor const &a,
                                  paddle::Tensor const &b,
                                  paddle::Tensor const &a_scales,
                                  paddle::Tensor const &b_scales,
                                  paddle::optional<paddle::Tensor> const &bias);
 
-void cutlass_scaled_mm_azp_sm90_int8(paddle::Tensor& out, paddle::Tensor const& a,
-                                     paddle::Tensor const& b,
-                                     paddle::Tensor const& a_scales,
-                                     paddle::Tensor const& b_scales,
-                                     paddle::Tensor const& azp_adj,
-                                     paddle::optional<paddle::Tensor> const& azp,
-                                     paddle::optional<paddle::Tensor> const& bias);
+void cutlass_scaled_mm_azp_sm90_int8(
+    paddle::Tensor &out,
+    paddle::Tensor const &a,
+    paddle::Tensor const &b,
+    paddle::Tensor const &a_scales,
+    paddle::Tensor const &b_scales,
+    paddle::Tensor const &azp_adj,
+    paddle::optional<paddle::Tensor> const &azp,
+    paddle::optional<paddle::Tensor> const &bias);
 
-void cutlass_scaled_mm_sm100_fp8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm100_fp8(paddle::Tensor &out,
+                                 paddle::Tensor const &a,
                                  paddle::Tensor const &b,
                                  paddle::Tensor const &a_scales,
                                  paddle::Tensor const &b_scales,
                                  paddle::optional<paddle::Tensor> const &bias);
 
-} // namespace fastdeploy
+}  // namespace fastdeploy

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8.cu

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8.cu
 
 // clang-format will break include orders
 // clang-format off
@@ -9,7 +10,8 @@
 
 namespace fastdeploy {
 
-void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out,
+                                paddle::Tensor const &a,
                                 paddle::Tensor const &b,
                                 paddle::Tensor const &a_scales,
                                 paddle::Tensor const &b_scales,
@@ -17,7 +19,8 @@ void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out, paddle::Tensor const &a,
   PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-             "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm90_fp8_epilogue<c3x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -25,4 +28,4 @@ void cutlass_scaled_mm_sm90_fp8(paddle::Tensor &out, paddle::Tensor const &a,
         out, a, b, a_scales, b_scales);
   }
 }
-} // namespace fastdeploy
+}  // namespace fastdeploy

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8_dispatch.cuh
 
 #pragma once
 
@@ -17,8 +18,10 @@ namespace fastdeploy {
 
 using c3x::cutlass_gemm_caller;
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_fp8_config_default {
   // M in (128, inf)
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
@@ -27,13 +30,19 @@ struct sm90_fp8_config_default {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_128, _128, _128>;
   using ClusterShape = Shape<_2, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_fp8_config_M128 {
   // M in (64, 128]
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
@@ -42,13 +51,19 @@ struct sm90_fp8_config_M128 {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _128, _128>;
   using ClusterShape = Shape<_2, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_fp8_config_M64 {
   // M in [1, 64]
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
@@ -58,13 +73,19 @@ struct sm90_fp8_config_M64 {
   using TileShape = Shape<_64, _64, _128>;
   using ClusterShape = Shape<_1, _8, _1>;
 
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm90_fp8_dispatch(paddle::Tensor &out,
                                            paddle::Tensor const &a,
@@ -75,8 +96,8 @@ inline void cutlass_gemm_sm90_fp8_dispatch(paddle::Tensor &out,
   PD_CHECK(b.dtype() == phi::DataType::FLOAT8_E4M3FN);
 
   using Cutlass3xGemmDefault =
-      typename sm90_fp8_config_default<InType, OutType,
-                                       Epilogue>::Cutlass3xGemm;
+      typename sm90_fp8_config_default<InType, OutType, Epilogue>::
+          Cutlass3xGemm;
   using Cutlass3xGemmM64 =
       typename sm90_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
   using Cutlass3xGemmM128 =
@@ -84,7 +105,7 @@ inline void cutlass_gemm_sm90_fp8_dispatch(paddle::Tensor &out,
 
   uint32_t const m = a.dims()[0];
   uint32_t const mp2 =
-      std::max(static_cast<uint32_t>(64), next_pow_2(m)); // next power of 2
+      std::max(static_cast<uint32_t>(64), next_pow_2(m));  // next power of 2
 
   if (mp2 <= 64) {
     // m in [1, 64]
@@ -112,14 +133,16 @@ void cutlass_scaled_mm_sm90_fp8_epilogue(paddle::Tensor &out,
 
   if (out.dtype() == paddle::DataType::BFLOAT16) {
     return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
-                                          cutlass::bfloat16_t, Epilogue>(
+                                          cutlass::bfloat16_t,
+                                          Epilogue>(
         out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
   } else {
     PD_CHECK(out.dtype() == paddle::DataType::FLOAT16);
     return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
-                                          cutlass::half_t, Epilogue>(
+                                          cutlass::half_t,
+                                          Epilogue>(
         out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8.cu

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8.cu
 
 // clang-format will break include orders
 // clang-format off
@@ -9,7 +10,8 @@
 
 namespace fastdeploy {
 
-void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out, paddle::Tensor const &a,
+void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out,
+                                 paddle::Tensor const &a,
                                  paddle::Tensor const &b,
                                  paddle::Tensor const &a_scales,
                                  paddle::Tensor const &b_scales,
@@ -17,7 +19,8 @@ void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out, paddle::Tensor const &a,
   PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-             "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm90_int8_epilogue<c3x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -26,4 +29,4 @@ void cutlass_scaled_mm_sm90_int8(paddle::Tensor &out, paddle::Tensor const &a,
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy

custom_ops/gpu_ops/cutlass_kernels/w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8_dispatch.cuh
 
 #pragma once
 
@@ -17,8 +18,10 @@ namespace fastdeploy {
 
 using c3x::cutlass_gemm_caller;
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_default {
   // For M > 128 and any N
   static_assert(std::is_same<InType, int8_t>());
@@ -27,13 +30,19 @@ struct sm90_int8_config_default {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_128, _128, _128>;
   using ClusterShape = Shape<_2, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_M128 {
   // For M in (64, 128] and any N
   static_assert(std::is_same<InType, int8_t>());
@@ -42,13 +51,19 @@ struct sm90_int8_config_M128 {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _128, _128>;
   using ClusterShape = Shape<_2, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_M64 {
   // For M in (32, 64] and any N
   static_assert(std::is_same<InType, int8_t>());
@@ -56,13 +71,19 @@ struct sm90_int8_config_M64 {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _64, _256>;
   using ClusterShape = Shape<_1, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_M32_NBig {
   // For M in [1, 32] and N >= 8192
   static_assert(std::is_same<InType, int8_t>());
@@ -70,13 +91,19 @@ struct sm90_int8_config_M32_NBig {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _128, _256>;
   using ClusterShape = Shape<_1, _4, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue>
 struct sm90_int8_config_M32_NSmall {
   // For M in [1, 32] and N < 8192
   static_assert(std::is_same<InType, int8_t>());
@@ -84,13 +111,19 @@ struct sm90_int8_config_M32_NSmall {
   using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
   using TileShape = Shape<_64, _64, _256>;
   using ClusterShape = Shape<_1, _8, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
-                      KernelSchedule, EpilogueSchedule>;
+  using Cutlass3xGemm = cutlass_3x_gemm<InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        ClusterShape,
+                                        KernelSchedule,
+                                        EpilogueSchedule>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm90_int8_dispatch(paddle::Tensor &out,
                                             paddle::Tensor const &a,
@@ -101,25 +134,25 @@ inline void cutlass_gemm_sm90_int8_dispatch(paddle::Tensor &out,
   PD_CHECK(b.dtype() == paddle::DataType::INT8);
 
   using Cutlass3xGemmDefault =
-      typename sm90_int8_config_default<InType, OutType,
-                                        Epilogue>::Cutlass3xGemm;
+      typename sm90_int8_config_default<InType, OutType, Epilogue>::
+          Cutlass3xGemm;
   using Cutlass3xGemmM128 =
       typename sm90_int8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
   using Cutlass3xGemmM64 =
       typename sm90_int8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
   using Cutlass3xGemmM32NBig =
-      typename sm90_int8_config_M32_NBig<InType, OutType,
-                                         Epilogue>::Cutlass3xGemm;
+      typename sm90_int8_config_M32_NBig<InType, OutType, Epilogue>::
+          Cutlass3xGemm;
   using Cutlass3xGemmM32NSmall =
-      typename sm90_int8_config_M32_NSmall<InType, OutType,
-                                           Epilogue>::Cutlass3xGemm;
+      typename sm90_int8_config_M32_NSmall<InType, OutType, Epilogue>::
+          Cutlass3xGemm;
 
   uint32_t const n = out.dims()[1];
   bool const is_small_n = n < 8192;
 
   uint32_t const m = a.dims()[0];
   uint32_t const mp2 =
-      std::max(static_cast<uint32_t>(32), next_pow_2(m)); // next power of 2
+      std::max(static_cast<uint32_t>(32), next_pow_2(m));  // next power of 2
 
   if (mp2 <= 32) {
     // m in [1, 32]
@@ -155,7 +188,8 @@ void cutlass_scaled_mm_sm90_int8_epilogue(paddle::Tensor &out,
   PD_CHECK(b.dtype() == paddle::DataType::INT8);
 
   if (out.dtype() == paddle::DataType::BFLOAT16) {
-    return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
+    return cutlass_gemm_sm90_int8_dispatch<int8_t,
+                                           cutlass::bfloat16_t,
                                            Epilogue>(
         out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
   } else {
@@ -165,4 +199,4 @@ void cutlass_scaled_mm_sm90_int8_epilogue(paddle::Tensor &out,
   }
 }
 
-} // namespace fastdeploy
+}  // namespace fastdeploy

custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cu

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
 
 #include "helper.h"
 #include <stddef.h>
@@ -20,7 +21,8 @@ using namespace fastdeploy;
 
 template <template <typename, typename> typename Epilogue,
           typename... EpilogueArgs>
-void cutlass_scaled_mm_sm75_epilogue(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm75_epilogue(paddle::Tensor& out,
+                                     paddle::Tensor const& a,
                                      paddle::Tensor const& b,
                                      EpilogueArgs&&... epilogue_args) {
   PD_CHECK(a.dtype() == paddle::DataType::INT8);
@@ -36,7 +38,8 @@ void cutlass_scaled_mm_sm75_epilogue(paddle::Tensor& out, paddle::Tensor const&
   }
 }
 
-void cutlass_scaled_mm_sm75(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm75(paddle::Tensor& out,
+                            paddle::Tensor const& a,
                             paddle::Tensor const& b,
                             paddle::Tensor const& a_scales,
                             paddle::Tensor const& b_scales,
@@ -45,7 +48,8 @@ void cutlass_scaled_mm_sm75(paddle::Tensor& out, paddle::Tensor const& a,
   PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-                "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -54,7 +58,8 @@ void cutlass_scaled_mm_sm75(paddle::Tensor& out, paddle::Tensor const& a,
   }
 }
 
-void cutlass_scaled_mm_azp_sm75(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm75(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -75,7 +80,8 @@ void cutlass_scaled_mm_azp_sm75(paddle::Tensor& out, paddle::Tensor const& a,
 
 template <template <typename, typename> typename Epilogue,
           typename... EpilogueArgs>
-void cutlass_scaled_mm_sm80_epilogue(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm80_epilogue(paddle::Tensor& out,
+                                     paddle::Tensor const& a,
                                      paddle::Tensor const& b,
                                      EpilogueArgs&&... epilogue_args) {
   PD_CHECK(a.dtype() == paddle::DataType::INT8);
@@ -91,7 +97,8 @@ void cutlass_scaled_mm_sm80_epilogue(paddle::Tensor& out, paddle::Tensor const&
   }
 }
 
-void cutlass_scaled_mm_sm80(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm80(paddle::Tensor& out,
+                            paddle::Tensor const& a,
                             paddle::Tensor const& b,
                             paddle::Tensor const& a_scales,
                             paddle::Tensor const& b_scales,
@@ -100,7 +107,8 @@ void cutlass_scaled_mm_sm80(paddle::Tensor& out, paddle::Tensor const& a,
   PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-                "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm80_epilogue<c2x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -109,7 +117,8 @@ void cutlass_scaled_mm_sm80(paddle::Tensor& out, paddle::Tensor const& a,
   }
 }
 
-void cutlass_scaled_mm_azp_sm80(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm80(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -130,14 +139,16 @@ void cutlass_scaled_mm_azp_sm80(paddle::Tensor& out, paddle::Tensor const& a,
 
 template <template <typename, typename> typename Epilogue,
           typename... EpilogueArgs>
-void cutlass_scaled_mm_sm89_epilogue(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm89_epilogue(paddle::Tensor& out,
+                                     paddle::Tensor const& a,
                                      paddle::Tensor const& b,
                                      EpilogueArgs&&... epilogue_args) {
   if (a.dtype() == paddle::DataType::INT8) {
     PD_CHECK(b.dtype() == paddle::DataType::INT8);
 
     if (out.dtype() == paddle::DataType::BFLOAT16) {
-      return cutlass_gemm_sm89_int8_dispatch<int8_t, cutlass::bfloat16_t,
+      return cutlass_gemm_sm89_int8_dispatch<int8_t,
+                                             cutlass::bfloat16_t,
                                              Epilogue>(
           out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
     } else {
@@ -151,18 +162,21 @@ void cutlass_scaled_mm_sm89_epilogue(paddle::Tensor& out, paddle::Tensor const&
 
     if (out.dtype() == paddle::DataType::BFLOAT16) {
       return cutlass_gemm_sm89_fp8_dispatch<cutlass::float_e4m3_t,
-                                            cutlass::bfloat16_t, Epilogue>(
+                                            cutlass::bfloat16_t,
+                                            Epilogue>(
           out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
     } else {
       PD_CHECK(out.dtype() == paddle::DataType::FLOAT16);
       return cutlass_gemm_sm89_fp8_dispatch<cutlass::float_e4m3_t,
-                                            cutlass::half_t, Epilogue>(
+                                            cutlass::half_t,
+                                            Epilogue>(
           out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
     }
   }
 }
 
-void cutlass_scaled_mm_sm89(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_sm89(paddle::Tensor& out,
+                            paddle::Tensor const& a,
                             paddle::Tensor const& b,
                             paddle::Tensor const& a_scales,
                             paddle::Tensor const& b_scales,
@@ -171,7 +185,8 @@ void cutlass_scaled_mm_sm89(paddle::Tensor& out, paddle::Tensor const& a,
   PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
   if (bias) {
     PD_CHECK(bias->dtype() == out.dtype(),
-                "currently bias dtype must match output dtype ", out.dtype());
+             "currently bias dtype must match output dtype ",
+             out.dtype());
     return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBias>(
         out, a, b, a_scales, b_scales, *bias);
   } else {
@@ -180,7 +195,8 @@ void cutlass_scaled_mm_sm89(paddle::Tensor& out, paddle::Tensor const& a,
   }
 }
 
-void cutlass_scaled_mm_azp_sm89(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm89(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,

custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
 
 #pragma once
 #include <stddef.h>
@@ -70,34 +71,42 @@ struct enable_sm89_to_sm90 : Kernel {
 #endif
   }
 };
-template <typename Arch, template <typename> typename ArchGuard,
-          typename ElementAB_, typename ElementD_,
-          template <typename, typename> typename Epilogue_, typename TileShape,
-          typename WarpShape, typename InstructionShape, int32_t MainLoopStages,
+template <typename Arch,
+          template <typename>
+          typename ArchGuard,
+          typename ElementAB_,
+          typename ElementD_,
+          template <typename, typename>
+          typename Epilogue_,
+          typename TileShape,
+          typename WarpShape,
+          typename InstructionShape,
+          int32_t MainLoopStages,
           typename FP8MathOperator = cutlass::arch::OpMultiplyAdd>
 struct cutlass_2x_gemm {
   using ElementAB = ElementAB_;
   using ElementD = ElementD_;
 
-  using ElementAcc =
-      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
-                                float>::type;
+  using ElementAcc = typename std::
+      conditional<std::is_same_v<ElementAB, int8_t>, int32_t, float>::type;
 
   using Operator =
       typename std::conditional<std::is_same_v<ElementAB, int8_t>,
                                 cutlass::arch::OpMultiplyAddSaturate,
                                 FP8MathOperator>::type;
 
-  using OutputTileThreadMap =
-      cutlass::epilogue::threadblock::OutputTileThreadLayout<
-          TileShape, WarpShape, float, 4, 1 /* epilogue stages */
-          >;
+  using OutputTileThreadMap = cutlass::epilogue::threadblock::
+      OutputTileThreadLayout<TileShape, WarpShape, float, 4, 1 /* epilogue
+                                                                  stages */
+                             >;
 
   using Epilogue = Epilogue_<ElementD, OutputTileThreadMap>;
   using EVTCompute = typename Epilogue::EVTCompute;
 
   using D = cutlass::epilogue::threadblock::VisitorAuxStore<
-      OutputTileThreadMap, ElementD, cutlass::FloatRoundStyle::round_to_nearest,
+      OutputTileThreadMap,
+      ElementD,
+      cutlass::FloatRoundStyle::round_to_nearest,
       Stride<int64_t, Int<1>, Int<0>>>;
 
   using EVTD = cutlass::epilogue::threadblock::Sm80EVT<D, EVTCompute>;
@@ -129,7 +138,8 @@ struct cutlass_2x_gemm {
 };
 
 template <typename Gemm, typename... EpilogueArgs>
-inline void cutlass_gemm_caller(paddle::Tensor& out, paddle::Tensor const& a,
+inline void cutlass_gemm_caller(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 EpilogueArgs&&... epilogue_params) {
   using ElementAB = typename Gemm::ElementAB;
@@ -183,7 +193,7 @@ inline void cutlass_gemm_caller(paddle::Tensor& out, paddle::Tensor const& a,
   // Launch the CUTLASS GEMM kernel.
   typename Gemm::Op gemm_op;
   size_t workspace_size = gemm_op.get_workspace_size(args);
-  phi::Allocator *allocator = paddle::GetAllocator(a.place());
+  phi::Allocator* allocator = paddle::GetAllocator(a.place());
   auto workspace = allocator->Allocate(workspace_size);
 
   auto stream = a.stream();
@@ -210,11 +220,10 @@ inline void fallback_cutlass_gemm_caller(paddle::Tensor& out,
       sizeof(typename FallbackGemm::KernelType::SharedStorage);
 
   if (gemm_shared_mem_size <= max_shared_mem_per_block_opt_in) {
-    return cutlass_gemm_caller<Gemm>(out, a, b,
-                                     std::forward<EpilogueArgs>(args)...);
+    return cutlass_gemm_caller<Gemm>(
+        out, a, b, std::forward<EpilogueArgs>(args)...);
   } else {
-    PD_CHECK(fallback_gemm_shared_mem_size <=
-                max_shared_mem_per_block_opt_in);
+    PD_CHECK(fallback_gemm_shared_mem_size <= max_shared_mem_per_block_opt_in);
     return cutlass_gemm_caller<FallbackGemm>(
         out, a, b, std::forward<EpilogueArgs>(args)...);
   }

custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm75_dispatch.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm75_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm75_dispatch.cuh
 
 #pragma once
 
@@ -11,8 +12,10 @@
 
 namespace fastdeploy {
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm75_config_default {
   // This config is used in 2 cases,
   // - M in (256, inf]
@@ -22,13 +25,21 @@ struct sm75_config_default {
   using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<8, 8, 16>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 2>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm75,
+                                        enable_sm75_to_sm80,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        2>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm75_config_M256 {
   // M in (128, 256]
   // Shared memory required by this Gemm 65536
@@ -36,13 +47,21 @@ struct sm75_config_M256 {
   using TileShape = typename cutlass::gemm::GemmShape<128, 128, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<8, 8, 16>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 2>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm75,
+                                        enable_sm75_to_sm80,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        2>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm75_config_M64 {
   // M in (32, 64]
   // Shared memory required by this Gemm 49152
@@ -50,13 +69,21 @@ struct sm75_config_M64 {
   using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<8, 8, 16>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 2>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm75,
+                                        enable_sm75_to_sm80,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        2>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm75_config_M32 {
   // M in [1, 32]
   // Shared memory required by this Gemm 49152
@@ -64,13 +91,21 @@ struct sm75_config_M32 {
   using TileShape = typename cutlass::gemm::GemmShape<32, 128, 64>;
   using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<8, 8, 16>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 2>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm75,
+                                        enable_sm75_to_sm80,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        2>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm75_dispatch(paddle::Tensor& out,
                                        paddle::Tensor const& a,
@@ -96,7 +131,8 @@ inline void cutlass_gemm_sm75_dispatch(paddle::Tensor& out,
   // sm75_config_default has the least shared-memory requirements.
   using FallbackGemm = Cutlass2xGemmDefault;
 
-  uint32_t const m = a.dims()[0];;
+  uint32_t const m = a.dims()[0];
+  ;
   uint32_t const mp2 =
       std::max(static_cast<uint32_t>(32), next_pow_2(m));  // next power of 2
   if (mp2 <= 32) {

custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm80_dispatch.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm80_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm80_dispatch.cuh
 
 #pragma once
 
@@ -11,8 +12,10 @@
 
 namespace fastdeploy {
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm80_config_default {
   // This config is used in 2 cases,
   //  - M in (128, inf)
@@ -22,13 +25,21 @@ struct sm80_config_default {
   using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm80, enable_sm80_to_sm89, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm80,
+                                        enable_sm80_to_sm89,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm80_config_M64 {
   // This config is used in 2 cases,
   // - M in (32, 64]
@@ -38,13 +49,21 @@ struct sm80_config_M64 {
   using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm80, enable_sm80_to_sm89, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm80,
+                                        enable_sm80_to_sm89,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm80_config_M32 {
   // M in (16, 32]
   // Shared Memory required by this Gemm - 61440 bytes
@@ -52,13 +71,21 @@ struct sm80_config_M32 {
   using TileShape = typename cutlass::gemm::GemmShape<32, 64, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm80, enable_sm80_to_sm89, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm80,
+                                        enable_sm80_to_sm89,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm80_config_M16 {
   // M in [1, 16]
   // Shared Memory required by this Gemm - 51200 bytes
@@ -66,13 +93,21 @@ struct sm80_config_M16 {
   using TileShape = typename cutlass::gemm::GemmShape<16, 64, 128>;
   using WarpShape = typename cutlass::gemm::GemmShape<16, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm80, enable_sm80_to_sm89, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm80,
+                                        enable_sm80_to_sm89,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm80_dispatch(paddle::Tensor& out,
                                        paddle::Tensor const& a,
@@ -104,7 +139,8 @@ inline void cutlass_gemm_sm80_dispatch(paddle::Tensor& out,
   using FallbackGemm =
       typename sm80_config_M32<InType, OutType, Epilogue>::Cutlass2xGemm;
 
-  uint32_t const m = a.dims()[0];;
+  uint32_t const m = a.dims()[0];
+  ;
   uint32_t const mp2 =
       std::max(static_cast<uint32_t>(16), next_pow_2(m));  // next power of 2
   if (mp2 <= 16) {
@@ -121,7 +157,8 @@ inline void cutlass_gemm_sm80_dispatch(paddle::Tensor& out,
         out, a, b, std::forward<EpilogueArgs>(args)...);
   } else if (mp2 <= 128) {
     // M in (64, 128]
-    uint32_t const n = out.dims()[1];;
+    uint32_t const n = out.dims()[1];
+    ;
     bool const small_n = n < 8192;
     if (small_n) {
       return fallback_cutlass_gemm_caller<Cutlass2xGemmM128SmallN,

custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_fp8_dispatch.cuh
 
 #pragma once
 
@@ -12,8 +13,10 @@
 
 namespace fastdeploy {
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm89_fp8_fallback_gemm {
   // Shared Memory required by this Gemm - 61440 bytes
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
@@ -21,10 +24,16 @@ struct sm89_fp8_fallback_gemm {
   using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAdd;
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm89, enable_sm89_to_sm90, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5,
-                      FP8MathOperator>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm89,
+                                        enable_sm89_to_sm90,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5,
+                                        FP8MathOperator>;
 };
 
 struct sm89_fp8_config_default {
@@ -33,17 +42,21 @@ struct sm89_fp8_config_default {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -52,26 +65,47 @@ struct sm89_fp8_config_default {
       using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 8192) {
       using TileShape = typename cutlass::gemm::GemmShape<256, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
 
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -83,17 +117,21 @@ struct sm89_fp8_config_M256 {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -102,17 +140,31 @@ struct sm89_fp8_config_M256 {
       using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -124,17 +176,21 @@ struct sm89_fp8_config_M128 {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -143,26 +199,47 @@ struct sm89_fp8_config_M128 {
       using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
 
     } else if (np2 <= 16384) {
       using TileShape = typename cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<128, 64, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -172,17 +249,21 @@ struct sm89_fp8_config_M64 {
   // M in (32, 64]
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -193,9 +274,16 @@ struct sm89_fp8_config_M64 {
       using FP8MathOperator = typename cutlass::arch::OpMultiplyAdd;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = typename cutlass::gemm::GemmShape<64, 128, 128>;
@@ -203,9 +291,16 @@ struct sm89_fp8_config_M64 {
       using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<64, 64, 128>;
@@ -213,9 +308,16 @@ struct sm89_fp8_config_M64 {
       using FP8MathOperator = typename cutlass::arch::OpMultiplyAdd;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -226,17 +328,21 @@ struct sm89_fp8_config_M32 {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -246,27 +352,48 @@ struct sm89_fp8_config_M32 {
       using WarpShape = typename cutlass::gemm::GemmShape<16, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = typename cutlass::gemm::GemmShape<32, 128, 128>;
       using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 4, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      4,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<32, 64, 128>;
       using WarpShape = typename cutlass::gemm::GemmShape<16, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5, FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -279,17 +406,21 @@ struct sm89_fp8_config_M16 {
   using FP8MathOperator = typename cutlass::arch::OpMultiplyAddFastAccum;
   static const int32_t MainLoopStages = 5;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
     PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
     using FallbackGemm =
-        typename sm89_fp8_fallback_gemm<InType, OutType,
-                                        Epilogue>::Cutlass2xGemm;
+        typename sm89_fp8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -298,35 +429,55 @@ struct sm89_fp8_config_M16 {
       using TileShape = typename cutlass::gemm::GemmShape<16, 64, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, MainLoopStages,
-                                FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      MainLoopStages,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 24576) {
       using TileShape = typename cutlass::gemm::GemmShape<16, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, MainLoopStages,
-                                FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      MainLoopStages,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = typename cutlass::gemm::GemmShape<32, 64, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, MainLoopStages,
-                                FP8MathOperator>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      MainLoopStages,
+                                      FP8MathOperator>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm89_fp8_dispatch(paddle::Tensor& out,
                                            paddle::Tensor const& a,
@@ -336,7 +487,8 @@ inline void cutlass_gemm_sm89_fp8_dispatch(paddle::Tensor& out,
   PD_CHECK(a.dtype() == paddle::DataType::FLOAT8_E4M3FN);
   PD_CHECK(b.dtype() == paddle::DataType::FLOAT8_E4M3FN);
 
-  uint32_t const m = a.dims()[0];;
+  uint32_t const m = a.dims()[0];
+  ;
   uint32_t const mp2 =
       std::max(static_cast<uint32_t>(16), next_pow_2(m));  // next power of 2
 

custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh
 
 #pragma once
 
@@ -11,8 +12,10 @@
 
 namespace fastdeploy {
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue>
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue>
 struct sm89_int8_fallback_gemm {
   // Shared mem requirement : 61440
   static_assert(std::is_same<InType, int8_t>());
@@ -21,9 +24,15 @@ struct sm89_int8_fallback_gemm {
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
   static int32_t const MainLoopStages = 5;
 
-  using Cutlass2xGemm =
-      cutlass_2x_gemm<cutlass::arch::Sm89, enable_sm89_to_sm90, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
+  using Cutlass2xGemm = cutlass_2x_gemm<cutlass::arch::Sm89,
+                                        enable_sm89_to_sm90,
+                                        InType,
+                                        OutType,
+                                        Epilogue,
+                                        TileShape,
+                                        WarpShape,
+                                        InstructionShape,
+                                        5>;
 };
 
 struct sm89_int8_config_default {
@@ -31,17 +40,21 @@ struct sm89_int8_config_default {
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -50,33 +63,57 @@ struct sm89_int8_config_default {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 8192) {
       using TileShape = cutlass::gemm::GemmShape<256, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<256, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -87,17 +124,21 @@ struct sm89_int8_config_M256 {
   using WarpShape = typename cutlass::gemm::GemmShape<64, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -106,33 +147,57 @@ struct sm89_int8_config_M256 {
       using TileShape = cutlass::gemm::GemmShape<64, 128, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 8192) {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = cutlass::gemm::GemmShape<256, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -142,17 +207,21 @@ struct sm89_int8_config_M128 {
   // M in (64, 128]
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -162,27 +231,45 @@ struct sm89_int8_config_M128 {
       using WarpShape = cutlass::gemm::GemmShape<64, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else if (np2 <= 16384) {
       using TileShape = cutlass::gemm::GemmShape<128, 128, 64>;
       using WarpShape = cutlass::gemm::GemmShape<64, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<64, 64, 128>;
       using WarpShape = cutlass::gemm::GemmShape<32, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -192,17 +279,21 @@ struct sm89_int8_config_M64 {
   // M in (32, 64]
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -212,18 +303,30 @@ struct sm89_int8_config_M64 {
       using WarpShape = cutlass::gemm::GemmShape<32, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<64, 128, 128>;
       using WarpShape = cutlass::gemm::GemmShape<64, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 3>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      3>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -233,17 +336,21 @@ struct sm89_int8_config_M32 {
   // M in (16, 32]
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[1];
     uint32_t const np2 = next_pow_2(n);
@@ -253,18 +360,30 @@ struct sm89_int8_config_M32 {
       using WarpShape = cutlass::gemm::GemmShape<16, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<32, 128, 128>;
       using WarpShape = cutlass::gemm::GemmShape<32, 64, 64>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 4>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      4>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
@@ -275,17 +394,21 @@ struct sm89_int8_config_M16 {
   using WarpShape = typename cutlass::gemm::GemmShape<16, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
 
-  template <typename InType, typename OutType,
-            template <typename, typename> typename Epilogue,
+  template <typename InType,
+            typename OutType,
+            template <typename, typename>
+            typename Epilogue,
             typename... EpilogueArgs>
-  static void dispatch(paddle::Tensor& out, paddle::Tensor const& a,
-                       paddle::Tensor const& b, EpilogueArgs&&... args) {
+  static void dispatch(paddle::Tensor& out,
+                       paddle::Tensor const& a,
+                       paddle::Tensor const& b,
+                       EpilogueArgs&&... args) {
     static_assert(std::is_same<InType, int8_t>());
     PD_CHECK(a.dtype() == paddle::DataType::INT8);
 
     using FallbackGemm =
-        typename sm89_int8_fallback_gemm<InType, OutType,
-                                         Epilogue>::Cutlass2xGemm;
+        typename sm89_int8_fallback_gemm<InType, OutType, Epilogue>::
+            Cutlass2xGemm;
 
     uint32_t const n = out.dims()[0];
     uint32_t const np2 = next_pow_2(n);
@@ -294,24 +417,38 @@ struct sm89_int8_config_M16 {
       using TileShape = cutlass::gemm::GemmShape<16, 64, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 5>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      5>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     } else {
       using TileShape = cutlass::gemm::GemmShape<16, 128, 128>;
 
       return fastdeploy::fallback_cutlass_gemm_caller<
-          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89, fastdeploy::enable_sm89_to_sm90,
-                                InType, OutType, Epilogue, TileShape, WarpShape,
-                                InstructionShape, 4>,
+          fastdeploy::cutlass_2x_gemm<cutlass::arch::Sm89,
+                                      fastdeploy::enable_sm89_to_sm90,
+                                      InType,
+                                      OutType,
+                                      Epilogue,
+                                      TileShape,
+                                      WarpShape,
+                                      InstructionShape,
+                                      4>,
           FallbackGemm>(out, a, b, std::forward<EpilogueArgs>(args)...);
     }
   }
 };
 
-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
+template <typename InType,
+          typename OutType,
+          template <typename, typename>
+          typename Epilogue,
           typename... EpilogueArgs>
 inline void cutlass_gemm_sm89_int8_dispatch(paddle::Tensor& out,
                                             paddle::Tensor const& a,

custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_c3x_sm90.cu

@@ -1,4 +1,5 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm90.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm90.cu
 
 #include "c3x/scaled_mm_helper.hpp"
 #include "c3x/scaled_mm_kernels.hpp"
@@ -10,17 +11,24 @@
 
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
 
-void cutlass_scaled_mm_sm90(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias) {
-  dispatch_scaled_mm(c, a, b, a_scales, b_scales, bias,
+void cutlass_scaled_mm_sm90(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias) {
+  dispatch_scaled_mm(c,
+                     a,
+                     b,
+                     a_scales,
+                     b_scales,
+                     bias,
                      fastdeploy::cutlass_scaled_mm_sm90_fp8,
                      fastdeploy::cutlass_scaled_mm_sm90_int8);
 }
 
-void cutlass_scaled_mm_azp_sm90(paddle::Tensor& out, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm90(paddle::Tensor& out,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -30,8 +38,8 @@ void cutlass_scaled_mm_azp_sm90(paddle::Tensor& out, paddle::Tensor const& a,
   PD_CHECK(a_scales.dtype() == paddle::DataType::FLOAT32);
   PD_CHECK(b_scales.dtype() == paddle::DataType::FLOAT32);
 
-  fastdeploy::cutlass_scaled_mm_azp_sm90_int8(out, a, b, a_scales, b_scales, azp_adj,
-                                        azp, bias);
+  fastdeploy::cutlass_scaled_mm_azp_sm90_int8(
+      out, a, b, a_scales, b_scales, azp_adj, azp, bias);
 }
 
 #endif

custom_ops/gpu_ops/cutlass_kernels/w8a8/scaled_mm_entry.cu

@@ -1,36 +1,42 @@
-// adapted from: https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+// adapted from:
+// https://github.com/vllm-project/vllm/blob/118ff921118cc81061a2af865a1e13840ceb6792/csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
 
 #pragma once
 #include "helper.h"
 #include <iostream>
 
-void cutlass_scaled_mm_sm75(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias);
+void cutlass_scaled_mm_sm75(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias);
 
-void cutlass_scaled_mm_sm80(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias);
+void cutlass_scaled_mm_sm80(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias);
 
-void cutlass_scaled_mm_sm89(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias);
+void cutlass_scaled_mm_sm89(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias);
 
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
-void cutlass_scaled_mm_sm90(paddle::Tensor &c, paddle::Tensor const &a,
-                            paddle::Tensor const &b,
-                            paddle::Tensor const &a_scales,
-                            paddle::Tensor const &b_scales,
-                            paddle::optional<paddle::Tensor> const &bias);
+void cutlass_scaled_mm_sm90(paddle::Tensor& c,
+                            paddle::Tensor const& a,
+                            paddle::Tensor const& b,
+                            paddle::Tensor const& a_scales,
+                            paddle::Tensor const& b_scales,
+                            paddle::optional<paddle::Tensor> const& bias);
 #endif
 
-void cutlass_scaled_mm_azp_sm75(paddle::Tensor& c, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm75(paddle::Tensor& c,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -38,7 +44,8 @@ void cutlass_scaled_mm_azp_sm75(paddle::Tensor& c, paddle::Tensor const& a,
                                 paddle::optional<paddle::Tensor> const& azp,
                                 paddle::optional<paddle::Tensor> const& bias);
 
-void cutlass_scaled_mm_azp_sm80(paddle::Tensor& c, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm80(paddle::Tensor& c,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -46,7 +53,8 @@ void cutlass_scaled_mm_azp_sm80(paddle::Tensor& c, paddle::Tensor const& a,
                                 paddle::optional<paddle::Tensor> const& azp,
                                 paddle::optional<paddle::Tensor> const& bias);
 
-void cutlass_scaled_mm_azp_sm89(paddle::Tensor& c, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm89(paddle::Tensor& c,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -55,7 +63,8 @@ void cutlass_scaled_mm_azp_sm89(paddle::Tensor& c, paddle::Tensor const& a,
                                 paddle::optional<paddle::Tensor> const& bias);
 
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
-void cutlass_scaled_mm_azp_sm90(paddle::Tensor& c, paddle::Tensor const& a,
+void cutlass_scaled_mm_azp_sm90(paddle::Tensor& c,
+                                paddle::Tensor const& a,
                                 paddle::Tensor const& b,
                                 paddle::Tensor const& a_scales,
                                 paddle::Tensor const& b_scales,
@@ -80,10 +89,12 @@ bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability) {
   return false;
 }
 
-void CutlassScaledMm(paddle::Tensor &c, paddle::Tensor const &a,
-                     paddle::Tensor const &b, paddle::Tensor const &a_scales,
-                     paddle::Tensor const &b_scales,
-                     paddle::optional<paddle::Tensor> const &bias) {
+void CutlassScaledMm(paddle::Tensor& c,
+                     paddle::Tensor const& a,
+                     paddle::Tensor const& b,
+                     paddle::Tensor const& a_scales,
+                     paddle::Tensor const& b_scales,
+                     paddle::optional<paddle::Tensor> const& bias) {
   // Checks for conformality
   PD_CHECK(a.dims().size() == 2 && b.dims().size() == 2 &&
            c.dims().size() == 2);
@@ -91,10 +102,10 @@ void CutlassScaledMm(paddle::Tensor &c, paddle::Tensor const &a,
            b.dims()[0] == c.dims()[1]);
 
   // Check for strides and alignment
-  PD_CHECK(a.strides()[1] == 1 && c.strides()[1] == 1); // Row-major
-  PD_CHECK(b.strides()[1] == 1);                        // Column-major
+  PD_CHECK(a.strides()[1] == 1 && c.strides()[1] == 1);  // Row-major
+  PD_CHECK(b.strides()[1] == 1);                         // Column-major
   PD_CHECK(c.strides()[0] % 16 == 0 &&
-           b.strides()[0] % 16 == 0); // 16 Byte Alignment
+           b.strides()[0] % 16 == 0);  // 16 Byte Alignment
 
   if (bias) {
     PD_CHECK(bias->numel() == b.dims()[0] && bias->is_contiguous() &&
@@ -138,13 +149,14 @@ void CutlassScaledMm(paddle::Tensor &c, paddle::Tensor const &a,
       version_num));
 }
 
-void CutlassScaledMmAzp(paddle::Tensor& c, paddle::Tensor const& a,
-                           paddle::Tensor const& b,
-                           paddle::Tensor const& a_scales,
-                           paddle::Tensor const& b_scales,
-                           paddle::Tensor const& azp_adj,
-                           paddle::optional<paddle::Tensor> const& azp,
-                           paddle::optional<paddle::Tensor> const& bias) {
+void CutlassScaledMmAzp(paddle::Tensor& c,
+                        paddle::Tensor const& a,
+                        paddle::Tensor const& b,
+                        paddle::Tensor const& a_scales,
+                        paddle::Tensor const& b_scales,
+                        paddle::Tensor const& azp_adj,
+                        paddle::optional<paddle::Tensor> const& azp,
+                        paddle::optional<paddle::Tensor> const& bias) {
   // Checks for conformality
   PD_CHECK(a.dims().size() == 2 && b.dims().size() == 2 &&
            c.dims().size() == 2);
@@ -154,10 +166,10 @@ void CutlassScaledMmAzp(paddle::Tensor& c, paddle::Tensor const& a,
   PD_CHECK(b_scales.numel() == 1 || b_scales.numel() == b.dims()[0]);
 
   // Check for strides and alignment
-  PD_CHECK(a.strides()[1] == 1 && c.strides()[1] == 1); // Row-major
-  PD_CHECK(b.strides()[1] == 1);                        // Column-major
+  PD_CHECK(a.strides()[1] == 1 && c.strides()[1] == 1);  // Row-major
+  PD_CHECK(b.strides()[1] == 1);                         // Column-major
   PD_CHECK(c.strides()[0] % 16 == 0 &&
-           b.strides()[0] % 16 == 0); // 16 Byte Alignment
+           b.strides()[0] % 16 == 0);  // 16 Byte Alignment
   PD_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
 
   // bias, azp, azp_adj are all 1d
@@ -174,7 +186,8 @@ void CutlassScaledMmAzp(paddle::Tensor& c, paddle::Tensor const& a,
   PD_CHECK(azp_adj.dtype() == paddle::DataType::INT32);
   PD_CHECK(!azp || azp->dtype() == paddle::DataType::INT32);
   PD_CHECK(!bias || bias->dtype() == c.dtype(),
-              "currently bias dtype must match output dtype ", c.dtype());
+           "currently bias dtype must match output dtype ",
+           c.dtype());
 
   int32_t version_num = GetGPUComputeCapability(a.place().GetDeviceId());
 
@@ -210,7 +223,6 @@ void CutlassScaledMmAzp(paddle::Tensor& c, paddle::Tensor const& a,
       version_num));
 }
 
-
 PD_BUILD_STATIC_OP(cutlass_scaled_mm)
     .Inputs({"c", "a", "b", "a_scales", "b_scales", paddle::Optional("bias")})
     .Outputs({"c_out"})
@@ -218,7 +230,14 @@ PD_BUILD_STATIC_OP(cutlass_scaled_mm)
     .SetKernelFn(PD_KERNEL(CutlassScaledMm));
 
 PD_BUILD_STATIC_OP(cutlass_scaled_mm_azp)
-    .Inputs({"c", "a", "b", "a_scales", "b_scales", "azp_adj", paddle::Optional("azp"), paddle::Optional("bias")})
+    .Inputs({"c",
+             "a",
+             "b",
+             "a_scales",
+             "b_scales",
+             "azp_adj",
+             paddle::Optional("azp"),
+             paddle::Optional("bias")})
     .Outputs({"c_out"})
     .SetInplaceMap({{"c", "c_out"}})
     .SetKernelFn(PD_KERNEL(CutlassScaledMmAzp));