[XPU] support flashmask_attention (PaddlePaddle#71573)

Author: runzhech

paddle/phi/backends/xpu/xpu3_op_list.cc

@@ -615,6 +615,10 @@ XPUOpMap& get_kl3_ops() {
  XPUKernelSet({phi::DataType::BFLOAT16,
  phi::DataType::FLOAT32,
  phi::DataType::FLOAT16})},
+ {"flashmask_attention",
+ XPUKernelSet({phi::DataType::BFLOAT16, phi::DataType::FLOAT16})},
+ {"flashmask_attention_grad",
+ XPUKernelSet({phi::DataType::BFLOAT16, phi::DataType::FLOAT16})},
  {"flash_attn_unpadded",
  XPUKernelSet({phi::DataType::BFLOAT16,
  phi::DataType::FLOAT32,

paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc

@@ -17,36 +17,39 @@
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
 #ifdef PADDLE_WITH_XPU_XRE5
+#include "paddle/phi/kernels/slice_kernel.h"
 #include "paddle/phi/kernels/xpu/flash_attn_utils.h"
 #include "xfa/flash_api.h"
 #endif
 namespace phi {
 #ifdef PADDLE_WITH_XPU_XRE5
 template <typename T, typename Context>
-void FlashAttnGradKernelBase(const Context& ctx,
- const DenseTensor& q,
- const DenseTensor& k,
- const DenseTensor& v,
- const api::VectorParam<int>& lod_seqlen_q,
- const api::VectorParam<int>& lod_seqlen_k,
- const DenseTensor& out,
- const DenseTensor& softmax_lse,
- const DenseTensor& seed_offset,
- const paddle::optional<DenseTensor>& attn_mask,
- const DenseTensor& dout,
- const int batch_size,
- const Scalar& max_seqlen_q_,
- const Scalar& max_seqlen_k_,
- const int num_heads,
- const int num_heads_k,
- const int head_size,
- const int head_size_v,
- float scale,
- float dropout,
- bool causal,
- DenseTensor* dq,
- DenseTensor* dk,
- DenseTensor* dv) {
+void FlashAttnGradKernelBase(
+ const Context& ctx,
+ const DenseTensor& q,
+ const DenseTensor& k,
+ const DenseTensor& v,
+ const api::VectorParam<int>& lod_seqlen_q,
+ const api::VectorParam<int>& lod_seqlen_k,
+ const DenseTensor& out,
+ const DenseTensor& softmax_lse,
+ const DenseTensor& seed_offset,
+ const paddle::optional<DenseTensor>& attn_mask,
+ const paddle::optional<DenseTensor>& startend_row_indices,
+ const DenseTensor& dout,
+ const int batch_size,
+ const Scalar& max_seqlen_q_,
+ const Scalar& max_seqlen_k_,
+ const int num_heads,
+ const int num_heads_k,
+ const int head_size,
+ const int head_size_v,
+ float scale,
+ float dropout,
+ bool causal,
+ DenseTensor* dq,
+ DenseTensor* dk,
+ DenseTensor* dv) {
  xpu::ctx_guard RAII_GUARD(ctx.x_context());
 
  using XPUType = typename XPUTypeTrait<T>::Type;
@@ -62,7 +65,52 @@ void FlashAttnGradKernelBase(const Context& ctx,
 
  const float* bias_data = nullptr;
  int64_t fa_layout = AttnQKVLayout_t::ATTN_BLHD;
- if (attn_mask.get_ptr() != nullptr) {
+ DenseTensor downstart_row_indices, upend_row_indices, downend_row_indices,
+ upstart_row_indices;
+ void *downstart_row_indices_data = nullptr, *upend_row_indices_data = nullptr,
+ *downend_row_indices_data = nullptr, *upstart_row_indices_data = nullptr;
+ bool is_flashmask = startend_row_indices.get_ptr() != nullptr;
+ XPUStream flashmask_stream;
+ if (is_flashmask) {
+ xpu_stream_create(&flashmask_stream);
+ PADDLE_ENFORCE_EQ(
+ startend_row_indices->dims().size(),
+ 4,
+ common::errors::InvalidArgument(
+ "flashmask_attention receive startend_row_indices with dim "
+ "[batch_size, num_heads,seq_len, mask_bounds]"));
+ PADDLE_ENFORCE_EQ(startend_row_indices->dims()[3] == 1 ||
+ startend_row_indices->dims()[3] == 2 ||
+ startend_row_indices->dims()[3] == 4,
+ true,
+ common::errors::InvalidArgument(
+ "flashmask_attention startend_row_indices "
+ "mask_bounds must in [1,2,4]"));
+ downstart_row_indices =
+ phi::Slice<int32_t>(ctx, startend_row_indices.get(), {3}, {0}, {1});
+ downstart_row_indices_data = downstart_row_indices.data();
+ if (startend_row_indices->dims()[3] == 2) {
+ if (!causal) {
+ upend_row_indices =
+ phi::Slice<int32_t>(ctx, startend_row_indices.get(), {3}, {1}, {2});
+ upend_row_indices_data = upend_row_indices.data();
+ } else {
+ downend_row_indices =
+ phi::Slice<int32_t>(ctx, startend_row_indices.get(), {3}, {1}, {2});
+ downend_row_indices_data = downend_row_indices.data();
+ }
+ } else if (startend_row_indices->dims()[3] == 4) {
+ upend_row_indices =
+ phi::Slice<int32_t>(ctx, startend_row_indices.get(), {3}, {3}, {4});
+ upend_row_indices_data = upend_row_indices.data();
+ downend_row_indices =
+ phi::Slice<int32_t>(ctx, startend_row_indices.get(), {3}, {1}, {2});
+ downend_row_indices_data = downend_row_indices.data();
+ upstart_row_indices =
+ phi::Slice<int32_t>(ctx, startend_row_indices.get(), {3}, {2}, {3});
+ upstart_row_indices_data = upstart_row_indices.data();
+ }
+ } else if (attn_mask.get_ptr() != nullptr) {
  const auto& mask_dims = attn_mask->dims();
  if (mask_dims.size() == 3 || (mask_dims[1] == 1 && mask_dims.size() == 4)) {
  fa_layout |= AttnQKVLayout_t::BIAS_BLL;
@@ -175,15 +223,17 @@ void FlashAttnGradKernelBase(const Context& ctx,
  -1, // window_size_left
  -1, // window_size_right
  head_size_v, // v_head_dim
- nullptr, // downstart_row_indices_data
- nullptr, // downend_row_indices_data
- nullptr, // upstart_row_indices_data
- nullptr, // upend_row_indices_data
- 0, // flash_mask_head_num
- nullptr, // flashmask_maxmin
- nullptr // side_stream
- );
+ (const int*)downstart_row_indices_data,
+ (const int*)downend_row_indices_data,
+ (const int*)upstart_row_indices_data,
+ (const int*)upend_row_indices_data,
+ is_flashmask ? startend_row_indices->dims()[1] : 0,
+ nullptr,
+ is_flashmask ? flashmask_stream : nullptr);
  PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_bwd");
+ if (is_flashmask && flashmask_stream != nullptr) {
+ xpu_stream_destroy(flashmask_stream);
+ }
 }
 #endif
 
@@ -236,6 +286,7 @@ void FlashAttnUnpaddedGradKernel(const Context& ctx,
  softmax_lse,
  seed_offset,
  attn_mask,
+ paddle::none,
  dout,
  batch_size,
  max_seqlen_q,
@@ -316,6 +367,7 @@ void FlashAttnGradKernel(const Context& ctx,
  softmax_lse,
  seed_offset,
  attn_mask,
+ paddle::none,
  dout,
  batch_size,
  seqlen_q,
@@ -336,6 +388,85 @@ void FlashAttnGradKernel(const Context& ctx,
 #endif
 }
 
+template <typename T, typename Context>
+void FlashMaskGradKernel(const Context& ctx,
+ const DenseTensor& q,
+ const DenseTensor& k,
+ const DenseTensor& v,
+ const DenseTensor& startend_row_indices,
+ const DenseTensor& out,
+ const DenseTensor& softmax_lse,
+ const DenseTensor& seed_offset,
+ const DenseTensor& dout,
+ float dropout,
+ bool causal,
+ 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);
+
+ // q, k, v [batch_size, seq_len, num_heads, head_dim]
+ const auto& dims = q.dims();
+
+ const int64_t batch_size = dims[0];
+ const int64_t seqlen_q = dims[1];
+ const int64_t num_heads = dims[2];
+ const int64_t head_size_og = dout.dims()[3];
+ const int64_t head_size = dims[3];
+ const int64_t head_size_v = v.dims()[3];
+ const int64_t seqlen_k = k.dims()[1];
+ const int64_t num_heads_k = k.dims()[2];
+
+ PADDLE_ENFORCE_EQ(
+ head_size_og,
+ head_size_v,
+ common::errors::InvalidArgument(
+ "flash_attn_bwd receive input with head_size_og == head_size_v"));
+
+ // lod info
+ std::vector<int> qlod_vec = {0};
+ std::vector<int> kvlod_vec = {0};
+ for (int batch_idx = 1; batch_idx <= batch_size; ++batch_idx) {
+ qlod_vec.push_back(seqlen_q * batch_idx);
+ kvlod_vec.push_back(seqlen_k * batch_idx);
+ }
+ api::VectorParam<int> qlod{
+ qlod_vec.data(), static_cast<int64_t>(qlod_vec.size()), nullptr};
+ api::VectorParam<int> kvlod{
+ kvlod_vec.data(), static_cast<int64_t>(kvlod_vec.size()), nullptr};
+ FlashAttnGradKernelBase<T>(ctx,
+ q,
+ k,
+ v,
+ qlod,
+ kvlod,
+ out,
+ softmax_lse,
+ seed_offset,
+ paddle::none,
+ startend_row_indices,
+ dout,
+ batch_size,
+ seqlen_q,
+ seqlen_k,
+ num_heads,
+ num_heads_k,
+ head_size,
+ head_size_v,
+ 0.0,
+ dropout,
+ causal,
+ dq,
+ dk,
+ dv);
+#else
+ PADDLE_THROW(common::errors::Unimplemented(
+ "re-compile using -DWITH_XPU_XRE5=ON to use FlashMaskGradKernel"));
+#endif
+}
 } // namespace phi
 
 PD_REGISTER_KERNEL(flash_attn_unpadded_grad,
@@ -359,3 +490,13 @@ PD_REGISTER_KERNEL(flash_attn_grad,
  phi::dtype::float16) {
  kernel->InputAt(5).SetBackend(phi::Backend::ALL_BACKEND); // seed_offset
 }
+
+PD_REGISTER_KERNEL(flashmask_attention_grad,
+ XPU,
+ ALL_LAYOUT,
+ phi::FlashMaskGradKernel,
+ phi::dtype::float16,
+ phi::dtype::bfloat16) {
+ kernel->InputAt(6).SetBackend(
+ phi::Backend::ALL_BACKEND); // fixed_seed_offset
+}