[XPU] fix kl2 build on new version of xhpc (#65221)

houj04 · web-flow · commit 0d07febf5f75 · 2024-06-24T10:58:36.000+08:00
* [XPU] fix kl2 build on new version of xhpc

* use xdnn fc on kl2
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,7 +30,7 @@ if(NOT DEFINED XPU_XRE_BASE_VERSION)
   set(XPU_XRE_BASE_VERSION "4.32.0.1")
 endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "20240515")
+  set(XPU_XHPC_BASE_DATE "20240601")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.2.1.2")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
@@ -58,7 +58,6 @@ set(XPU_XPTI_BASE_URL
 
 if(WITH_XPU_XRE5)
   set(XPU_XRE_BASE_VERSION "5.0.3.1")
-  set(XPU_XHPC_BASE_DATE "20240601")
   set(XPU_XRE_BASE_URL
       "https://klx-sdk-release-public.su.bcebos.com/xre/kl3-release/${XPU_XRE_BASE_VERSION}"
   )
@@ -79,12 +78,11 @@ elseif(WITH_SUNWAY)
   set(XPU_XCCL_DIR_NAME "") # TODO: xccl has no deepin output at now.
   set(XPU_XFT_DIR_NAME "") # TODO: xft has no deepin output at now.
 elseif(WITH_BDCENTOS)
+  set(XPU_XHPC_DIR_NAME "xhpc-bdcentos7_x86_64")
   if(WITH_XPU_XRE5)
     set(XPU_XRE_DIR_NAME "xre-bdcentos-x86_64-${XPU_XRE_BASE_VERSION}")
-    set(XPU_XHPC_DIR_NAME "xhpc-bdcentos7_x86_64")
   else()
     set(XPU_XRE_DIR_NAME "xre-bdcentos_x86_64")
-    set(XPU_XHPC_DIR_NAME "xhpc-bdcentos_x86_64")
   endif()
   set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}-bdcentos_x86_64")
   set(XPU_XFT_DIR_NAME "xft_bdcentos6u3_x86_64_gcc82")
@@ -99,7 +97,7 @@ else()
     set(XPU_XHPC_DIR_NAME "xhpc-ubuntu2004_x86_64")
   else()
     set(XPU_XRE_DIR_NAME "xre-ubuntu_1604_x86_64")
-    set(XPU_XHPC_DIR_NAME "xhpc-ubuntu_x86_64")
+    set(XPU_XHPC_DIR_NAME "xhpc-ubuntu1604_x86_64")
   endif()
   set(XPU_XCCL_DIR_NAME "${XPU_XCCL_PREFIX}-ubuntu_x86_64")
   set(XPU_XFT_DIR_NAME "xft_ubuntu1604_x86_64")
@@ -206,8 +204,11 @@ include_directories(${XPU_XHPC_INC_DIR})
 set(XPU_XBLAS_INC_DIR "${XPU_INC_DIR}/xhpc/xblas")
 include_directories(${XPU_XBLAS_INC_DIR})
 
-set(XPU_XFA_INC_DIR "${XPU_INC_DIR}/xhpc/xfa")
-include_directories(${XPU_XFA_INC_DIR})
+if(WITH_XPU_XRE5)
+  add_definitions(-DPADDLE_WITH_XPU_XRE5)
+  set(XPU_XFA_INC_DIR "${XPU_INC_DIR}/xhpc/xfa")
+  include_directories(${XPU_XFA_INC_DIR})
+endif()
 
 if(WITH_XPTI)
   message(STATUS "Compile with XPU XPTI!")
@@ -225,8 +226,13 @@ if(WITH_XPTI)
   target_link_libraries(xpulib ${XPU_XPTI_LIB})
 endif()
 
-target_link_libraries(xpulib ${XPU_RT_LIB} ${XPU_BKCL_LIB} ${XPU_XBLAS_LIB}
-                      ${XPU_API_LIB} ${XPU_XFA_LIB})
+if(WITH_XPU_XRE5)
+  target_link_libraries(xpulib ${XPU_RT_LIB} ${XPU_BKCL_LIB} ${XPU_XBLAS_LIB}
+                        ${XPU_API_LIB} ${XPU_XFA_LIB})
+else()
+  target_link_libraries(xpulib ${XPU_RT_LIB} ${XPU_BKCL_LIB} ${XPU_XBLAS_LIB}
+                        ${XPU_API_LIB})
+endif()
 
 add_dependencies(xpulib ${XPU_PROJECT})
 
diff --git a/paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc b/paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc
@@ -16,9 +16,9 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
-
+#ifdef PADDLE_WITH_XPU_XRE5
 #include "xfa/flash_api.h"
-
+#endif
 namespace phi {
 
 template <typename T, typename Context>
@@ -36,6 +36,7 @@ void FlashAttnGradKernel(const Context& ctx,
                          DenseTensor* dq,
                          DenseTensor* dk,
                          DenseTensor* dv) {
+#ifdef PADDLE_WITH_XPU_XRE5
   ctx.template Alloc<T>(dq);
   ctx.template Alloc<T>(dk);
   ctx.template Alloc<T>(dv);
@@ -146,6 +147,10 @@ void FlashAttnGradKernel(const Context& ctx,
       bias_data                                   // bias
   );
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_bwd");
+#else
+  PADDLE_THROW(phi::errors::Unimplemented(
+      "re-compile using -DWITH_XPU_XRE5=ON to use FlashAttnGradKernel"));
+#endif
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/xpu/flash_attn_kernel.cc b/paddle/phi/kernels/xpu/flash_attn_kernel.cc
@@ -16,9 +16,9 @@
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/common/memory_utils.h"
 #include "paddle/phi/core/kernel_registry.h"
-
+#ifdef PADDLE_WITH_XPU_XRE5
 #include "xfa/flash_api.h"
-
+#endif
 namespace phi {
 
 template <typename T, typename Context>
@@ -43,6 +43,7 @@ void FlashAttnUnpaddedKernel(
     DenseTensor* softmax,
     DenseTensor* softmax_lse,
     DenseTensor* seed_offset) {
+#ifdef PADDLE_WITH_XPU_XRE5
   xpu::ctx_guard RAII_GUARD(ctx.x_context());
   // q, k, v [batch_size * seq_len, num_heads, head_dim]
   std::vector<int64_t> dims = common::vectorize(q.dims());
@@ -169,6 +170,10 @@ void FlashAttnUnpaddedKernel(
         nullptr);
     PADDLE_ENFORCE_EQ(r, 0, "xpu::qk_v_attention failed.");
   }
+#else
+  PADDLE_THROW(phi::errors::Unimplemented(
+      "re-compile using -DWITH_XPU_XRE5=ON to use FlashAttnUnpaddedKernel"));
+#endif
 }
 
 template <typename T, typename Context>
@@ -187,6 +192,7 @@ void FlashAttnKernel(const Context& ctx,
                      DenseTensor* softmax,
                      DenseTensor* softmax_lse,
                      DenseTensor* seed_offset) {
+#ifdef PADDLE_WITH_XPU_XRE5
   if (return_softmax == true) {
     PADDLE_THROW(phi::errors::Unimplemented("return_softmax should be false"));
   }
@@ -319,6 +325,10 @@ void FlashAttnKernel(const Context& ctx,
       bias_data                                   // bias
   );
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_fwd");
+#else
+  PADDLE_THROW(phi::errors::Unimplemented(
+      "re-compile using -DWITH_XPU_XRE5=ON to use FlashAttnKernel"));
+#endif
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/xpu/xpu_api_wrapper.h b/paddle/phi/kernels/xpu/xpu_api_wrapper.h
@@ -258,7 +258,7 @@ static void xblas_fc_wrapper(xpu::Context* ctx,
     std::vector<int> axis = {1, 0};
     r = xpu::transpose<XPUType>(ctx, x, l3_addr, shape, axis);
     PADDLE_ENFORCE_XDNN_SUCCESS(r, "transpose");
-
+#ifdef PADDLE_WITH_XPU_XRE5
     r = xblas::fc_fusion<XPUType, XPUType, XPUType, FCT>(ctx,
                                                          l3_addr,
                                                          w,
@@ -283,7 +283,30 @@ static void xblas_fc_wrapper(xpu::Context* ctx,
                                                          scale_x_mode,
                                                          scale_w_mode);
     PADDLE_ENFORCE_XDNN_SUCCESS(r, "xblas_fc_fusion");
+#else
+    r = xpu::fc_fusion<XPUType, XPUType, XPUType, FCT>(ctx,
+                                                       l3_addr,
+                                                       w,
+                                                       y,
+                                                       m,
+                                                       n,
+                                                       k,
+                                                       false,
+                                                       w_trans,
+                                                       x_maxptr,
+                                                       w_maxptr,
+                                                       y_maxptr,
+                                                       k,
+                                                       ldw,
+                                                       ldy,
+                                                       alpha,
+                                                       beta,
+                                                       bias,
+                                                       act);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_fc_fusion");
+#endif
   } else {
+#ifdef PADDLE_WITH_XPU_XRE5
     r = xblas::fc_fusion<XPUType, XPUType, XPUType, FCT>(ctx,
                                                          x,
                                                          w,
@@ -309,6 +332,28 @@ static void xblas_fc_wrapper(xpu::Context* ctx,
                                                          scale_w_mode);
 
     PADDLE_ENFORCE_XDNN_SUCCESS(r, "xblas_fc_fusion");
+#else
+    r = xpu::fc_fusion<XPUType, XPUType, XPUType, FCT>(ctx,
+                                                       x,
+                                                       w,
+                                                       y,
+                                                       m,
+                                                       n,
+                                                       k,
+                                                       x_trans,
+                                                       w_trans,
+                                                       x_maxptr,
+                                                       w_maxptr,
+                                                       y_maxptr,
+                                                       ldx,
+                                                       ldw,
+                                                       ldy,
+                                                       alpha,
+                                                       beta,
+                                                       bias,
+                                                       act);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu_fc_fusion");
+#endif
   }
 }
 
@@ -348,6 +393,11 @@ DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeBF16, XPUTypeFP16)
 DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeFP16, int32_t)
 DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(XPUTypeFP16, tfloat32)
 
+#ifndef PADDLE_WITH_XPU_XRE5
+// disable in kl2
+DECLARE_UNSUPPORTED_XBLAS_FC_WRAPPER(float, XPUTypeFP16)
+#endif
+
 template <typename XPUType, typename FCT, typename TGEMM_OUT>
 static void xblas_fc_batch_wrapper(xpu::Context* xpu_ctx,
                                    int bs,
@@ -366,6 +416,7 @@ static void xblas_fc_batch_wrapper(xpu::Context* xpu_ctx,
                                    int stride_y,
                                    const float* x_maxptr,
                                    const float* w_maxptr) {
+#ifdef PADDLE_WITH_XPU_XRE5
   int r = xblas::fc_batched<XPUType, XPUType, XPUType, FCT, TGEMM_OUT, 0>(
       xpu_ctx,
       bs,
@@ -384,7 +435,28 @@ static void xblas_fc_batch_wrapper(xpu::Context* xpu_ctx,
       stride_y,
       x_maxptr,
       w_maxptr);
-  PADDLE_ENFORCE_XDNN_SUCCESS(r, "xblas_fc_batch_wrapper");
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "xblas_fc_batched");
+#else
+  int r = xpu::fc_batched<XPUType, XPUType, XPUType, FCT>(
+      xpu_ctx,
+      bs,
+      trans_x,
+      trans_w,
+      m,
+      n,
+      k,
+      alpha,
+      reinterpret_cast<const XPUType*>(x),
+      stride_x,
+      reinterpret_cast<const XPUType*>(w),
+      stride_w,
+      0.0,
+      reinterpret_cast<XPUType*>(y),
+      stride_y,
+      x_maxptr,
+      w_maxptr);
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "xdnn_fc_batched");
+#endif
 }
 
 #define DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUType, FCT, TGEMM_OUT) \
@@ -429,6 +501,18 @@ DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, int32_t, float)
 DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, int16_t, float)
 DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16, int32_t, float)
 
+#ifndef PADDLE_WITH_XPU_XRE5
+// disable in kl2
+DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, tfloat32, float)
+DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeBF16, float, float)
+DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(float, XPUTypeFP16, float)
+DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(float, XPUTypeFP16, XPUTypeFP16)
+DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16, XPUTypeFP16, float)
+DECLARE_UNSUPPORTED_XBLAS_FC_BATCH_WRAPPER(XPUTypeFP16,
+                                           XPUTypeFP16,
+                                           XPUTypeFP16)
+#endif
+
 template <typename T>
 static void MatMulXPUFunction(
     xpu::Context* xpu_ctx,
