[XPU] support selecting tgemm for xpu flash attention by env variable (#66143)

runzhech · root · web-flow · commit a5b80d05934a · 2024-07-19T15:48:56.000+08:00
Co-authored-by: root &lt;root@szzj-r300-training06.szzj&gt;
diff --git a/paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc b/paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc
@@ -17,6 +17,7 @@
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
 #ifdef PADDLE_WITH_XPU_XRE5
+#include "paddle/phi/kernels/xpu/flash_attn_utils.h"
 #include "xfa/flash_api.h"
 #endif
 namespace phi {
@@ -60,6 +61,7 @@ void FlashAttnGradKernel(const Context& ctx,
 
   // raw pointers
   using XPUType = typename XPUTypeTrait<T>::Type;
+  using XPUTypeFP16 = typename XPUTypeTrait<phi::dtype::float16>::Type;
   const XPUType* q_data = reinterpret_cast<const XPUType*>(q.data<T>());
   const XPUType* k_data = reinterpret_cast<const XPUType*>(k.data<T>());
   const XPUType* v_data = reinterpret_cast<const XPUType*>(v.data<T>());
@@ -119,7 +121,17 @@ void FlashAttnGradKernel(const Context& ctx,
 
   // get seed offset
   const int64_t* seed_offset_data = seed_offset.data<int64_t>();
+  int fa_tgemm = get_flash_attn_tgemm<XPUType>();
 
+  auto flash_attention_grad_kernel =
+      baidu::xpu::xfa::mha_varlen_bwd<XPUType, float, tfloat32, int>;
+  if (fa_tgemm == XPU_FA_TGEMM::FA_FLOAT) {
+    flash_attention_grad_kernel =
+        baidu::xpu::xfa::mha_varlen_bwd<XPUType, float, float, int>;
+  } else if (fa_tgemm == XPU_FA_TGEMM::FA_FLOAT16) {
+    flash_attention_grad_kernel =
+        baidu::xpu::xfa::mha_varlen_bwd<XPUType, float, XPUTypeFP16, int>;
+  }
   // template<typename T, typename TACCUM, typename TGEMM, typename TID = int>
   // int mha_varlen_bwd(xdnn::Context* ctx, const T* dout, const T* q, const T*
   // k, const T* v, const T* out, const TACCUM* softmax_lse, T* dq, T* dk, T*
@@ -132,7 +144,7 @@ void FlashAttnGradKernel(const Context& ctx,
   // k_maxptr = nullptr, const float* v_maxptr = nullptr, const float* o_maxptr
   // = nullptr, float* dq_maxptr = nullptr, float* dk_maxptr = nullptr, float*
   // dv_maxptr = nullptr, const float* do_maxptr = nullptr);
-  int r = baidu::xpu::xfa::mha_varlen_bwd<XPUType, float, tfloat32, int>(
+  int r = flash_attention_grad_kernel(
       ctx.x_context(),
       dout_data,                                  // dout
       q_data,                                     // q
@@ -166,7 +178,11 @@ void FlashAttnGradKernel(const Context& ctx,
       nullptr,                                    // do_maxptr
       false,                                      // is_qkv_fusion
       false,                                      // is_dqkv_fusion
-      fa_layout                                   // qkv_layout
+      fa_layout,                                  // qkv_layout
+      nullptr,                                    // alibi_slopes
+      {},                                         // alibi_slopes_shape
+      -1,                                         // window_size_left
+      -1                                          // window_size_right
   );
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_bwd");
 #else
diff --git a/paddle/phi/kernels/xpu/flash_attn_kernel.cc b/paddle/phi/kernels/xpu/flash_attn_kernel.cc
@@ -14,9 +14,9 @@
 
 #include "paddle/phi/kernels/flash_attn_kernel.h"
 #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 "paddle/phi/kernels/xpu/flash_attn_utils.h"
 #include "xfa/flash_api.h"
 #endif
 namespace phi {
@@ -231,32 +231,13 @@ void FlashAttnKernel(const Context& ctx,
   // generate seed offset
   seed_offset->Resize({2});
   int64_t* seed_offset_data = ctx.template HostAlloc<int64_t>(seed_offset);
-  if (fixed_seed_offset.get_ptr()) {
-    if ((fixed_seed_offset->place()).GetType() == phi::AllocationType::XPU) {
-      memory_utils::Copy(phi::CPUPlace(),
-                         seed_offset_data,
-                         fixed_seed_offset->place(),
-                         fixed_seed_offset->data<int64_t>(),
-                         sizeof(int64_t) * 2);
-    } else {
-      const int64_t* fixed_seed_offset_data =
-          fixed_seed_offset->data<int64_t>();
-      seed_offset_data[0] = fixed_seed_offset_data[0];
-      seed_offset_data[1] = fixed_seed_offset_data[1];
-    }
-  } else {
-    std::pair<uint64_t, uint64_t> seed_offset_pair;
-    uint64_t inc = batch_size * num_heads * 32;
-    if (rng_name != "") {
-      auto gen = phi::GetRandomSeedGenerator(rng_name);
-      seed_offset_pair = gen->IncrementOffset(inc);
-    } else {
-      auto* gen = ctx.GetGenerator();
-      seed_offset_pair = gen->IncrementOffset(inc);
-    }
-    seed_offset_data[0] = static_cast<int64_t>(seed_offset_pair.first);
-    seed_offset_data[1] = static_cast<int64_t>(seed_offset_pair.second);
-  }
+
+  phi::GenerateRNGState(ctx,
+                        fixed_seed_offset,
+                        seed_offset_data,
+                        rng_name,
+                        batch_size,
+                        num_heads);
 
   // raw pointers
   using XPUType = typename XPUTypeTrait<T>::Type;
@@ -309,33 +290,47 @@ void FlashAttnKernel(const Context& ctx,
   // const TACCUM* bias = nullptr, const float* q_maxptr = nullptr, const float*
   // k_maxptr = nullptr, const float* v_maxptr = nullptr, float* o_maxptr =
   // nullptr);
-  int r = baidu::xpu::xfa::mha_varlen_fwd<XPUType, float, tfloat32, int>(
-      ctx.x_context(),
-      q_data,                                     // q
-      k_data,                                     // k
-      v_data,                                     // v
-      out_data,                                   // out
-      softmax_lse_data,                           // softmax_lse
-      qlod,                                       // lod_seqlens_q
-      kvlod,                                      // lod_seqlens_k
-      seqlen_q,                                   // max_seqlen_q
-      seqlen_k,                                   // max_seqlen_k
-      num_heads,                                  // head_num
-      num_heads_k,                                // head_num_k
-      head_size,                                  // head_dim
-      1.0f / std::sqrt(head_size),                // softmax_scale
-      dropout,                                    // p_dropout
-      static_cast<int32_t>(seed_offset_data[0]),  // seed
-      causal,                                     // is_causal
-      nullptr,                                    // attn_mask
-      bias_data,                                  // bias
-      nullptr,                                    // q_maxptr
-      nullptr,                                    // k_maxptr
-      nullptr,                                    // v_maxptr
-      nullptr,                                    // o_maxptr
-      false,                                      // is_qkv_fusion
-      fa_layout                                   // qkv_layout
-  );
+  int fa_tgemm = get_flash_attn_tgemm<XPUType>();
+  auto flash_attention_kernel =
+      baidu::xpu::xfa::mha_varlen_fwd<XPUType, float, tfloat32, int>;
+  if (fa_tgemm == XPU_FA_TGEMM::FA_FLOAT) {
+    flash_attention_kernel =
+        baidu::xpu::xfa::mha_varlen_fwd<XPUType, float, float, int>;
+  } else if (fa_tgemm == XPU_FA_TGEMM::FA_FLOAT16) {
+    flash_attention_kernel =
+        baidu::xpu::xfa::mha_varlen_fwd<XPUType, float, XPUTypeFP16, int>;
+  }
+  int r =
+      flash_attention_kernel(ctx.x_context(),
+                             q_data,                       // q
+                             k_data,                       // k
+                             v_data,                       // v
+                             out_data,                     // out
+                             softmax_lse_data,             // softmax_lse
+                             qlod,                         // lod_seqlens_q
+                             kvlod,                        // lod_seqlens_k
+                             seqlen_q,                     // max_seqlen_q
+                             seqlen_k,                     // max_seqlen_k
+                             num_heads,                    // head_num
+                             num_heads_k,                  // head_num_k
+                             head_size,                    // head_dim
+                             1.0f / std::sqrt(head_size),  // softmax_scale
+                             dropout,                      // p_dropout
+                             static_cast<int32_t>(seed_offset_data[0]),  // seed
+                             causal,     // is_causal
+                             nullptr,    // attn_mask
+                             bias_data,  // bias
+                             nullptr,    // q_maxptr
+                             nullptr,    // k_maxptr
+                             nullptr,    // v_maxptr
+                             nullptr,    // o_maxptr
+                             false,      // is_qkv_fusion
+                             fa_layout,  // qkv_layout
+                             nullptr,    // alibi_slopes
+                             {},         // alibi_slopes_shape
+                             -1,         // window_size_left
+                             -1          // window_size_right
+      );
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_fwd");
 #else
   PADDLE_THROW(phi::errors::Unimplemented(
diff --git a/paddle/phi/kernels/xpu/flash_attn_utils.h b/paddle/phi/kernels/xpu/flash_attn_utils.h
@@ -0,0 +1,90 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#ifdef PADDLE_WITH_XPU
+
+#include <vector>
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "xfa/flash_api.h"
+
+namespace xfa = baidu::xpu::xfa;
+namespace phi {
+
+#ifdef PADDLE_WITH_XPU_XRE5
+using XPUTypeFP16 = typename XPUTypeTrait<phi::dtype::float16>::Type;
+using XPUTypeBF16 = typename XPUTypeTrait<phi::dtype::bfloat16>::Type;
+
+enum XPU_FA_TGEMM {
+  FA_FLOAT = 0,
+  FA_TFLOAT32,
+  FA_FLOAT16,
+};
+
+template <typename T>
+XPU_FA_TGEMM get_flash_attn_tgemm() {
+  const char* xpu_paddle_fa_float16 =
+      std::getenv("XPU_PADDLE_FA_TGEMM_FLOAT16");
+  if (xpu_paddle_fa_float16 != nullptr &&
+      (std::is_same<phi::dtype::float16, T>::value ||
+       std::is_same<XPUTypeFP16, T>::value)) {
+    return XPU_FA_TGEMM::FA_FLOAT16;
+  } else if (std::getenv("XPU_PADDLE_FA_TGEMM_FLOAT")) {
+    return XPU_FA_TGEMM::FA_FLOAT;
+  } else {
+    return XPU_FA_TGEMM::FA_TFLOAT32;
+  }
+}
+
+static void GenerateRNGState(
+    const XPUContext& ctx,
+    const paddle::optional<DenseTensor>& fixed_seed_offset,
+    int64_t* seed_offset_data,
+    const std::string& rng_name,
+    const int64_t batch_size,
+    const int64_t num_heads) {
+  if (fixed_seed_offset.get_ptr()) {
+    if ((fixed_seed_offset->place()).GetType() == phi::AllocationType::XPU) {
+      memory_utils::Copy(phi::CPUPlace(),
+                         seed_offset_data,
+                         fixed_seed_offset->place(),
+                         fixed_seed_offset->data<int64_t>(),
+                         sizeof(int64_t) * 2);
+    } else {
+      const int64_t* fixed_seed_offset_data =
+          fixed_seed_offset->data<int64_t>();
+      seed_offset_data[0] = fixed_seed_offset_data[0];
+      seed_offset_data[1] = fixed_seed_offset_data[1];
+    }
+  } else {
+    std::pair<uint64_t, uint64_t> seed_offset_pair;
+    uint64_t inc = batch_size * num_heads * 32;
+    if (rng_name != "") {
+      auto gen = phi::GetRandomSeedGenerator(rng_name);
+      seed_offset_pair = gen->IncrementOffset(inc);
+    } else {
+      auto* gen = ctx.GetGenerator();
+      seed_offset_pair = gen->IncrementOffset(inc);
+    }
+    seed_offset_data[0] = static_cast<int64_t>(seed_offset_pair.first);
+    seed_offset_data[1] = static_cast<int64_t>(seed_offset_pair.second);
+  }
+}
+
+#endif
+
+}  // namespace phi
+#endif
