PaddlePaddle
diff --git a/‎paddle/phi/kernels/flash_attn_kernel.h
Copy file name to clipboard
+15-17Lines changed: 15 additions & 17 deletions b/‎paddle/phi/kernels/flash_attn_kernel.h
Copy file name to clipboard
+15-17Lines changed: 15 additions & 17 deletions
diff --git a/‎paddle/phi/kernels/gpu/calc_reduced_attn_kernel.cu
Copy file name to clipboard
-1Lines changed: 0 additions & 1 deletion b/‎paddle/phi/kernels/gpu/calc_reduced_attn_kernel.cu
Copy file name to clipboard
-1Lines changed: 0 additions & 1 deletion
diff --git a/‎paddle/phi/kernels/gpu/flash_attn_grad_kernel.cu
Copy file name to clipboard
+78-35Lines changed: 78 additions & 35 deletions b/‎paddle/phi/kernels/gpu/flash_attn_grad_kernel.cu
Copy file name to clipboard
+78-35Lines changed: 78 additions & 35 deletions
@@ -60,22 +60,20 @@ void FlashAttnKernel(const Context& ctx,
                      DenseTensor* seed_offset);
 
 template <typename T, typename Context>
-void FlashAttnWithSparseMaskKernel(
-    const Context& ctx,
-    const DenseTensor& q,
-    const DenseTensor& k,
-    const DenseTensor& v,
-    const DenseTensor& attn_mask_start_row_indices,
-    const paddle::optional<DenseTensor>& fixed_seed_offset,
-    float dropout,
-    bool causal,
-    int attn_mask_start_row,
-    bool return_softmax,
-    bool is_test,
-    const std::string& rng_name,
-    DenseTensor* out,
-    DenseTensor* softmax,
-    DenseTensor* softmax_lse,
-    DenseTensor* seed_offset);
+void FlashMaskKernel(const Context& ctx,
+                     const DenseTensor& q,
+                     const DenseTensor& k,
+                     const DenseTensor& v,
+                     const DenseTensor& startend_row_indices,
+                     const paddle::optional<DenseTensor>& fixed_seed_offset,
+                     float dropout,
+                     bool causal,
+                     bool return_softmax,
+                     bool is_test,
+                     const std::string& rng_name,
+                     DenseTensor* out,
+                     DenseTensor* softmax,
+                     DenseTensor* softmax_lse,
+                     DenseTensor* seed_offset);
 
 }  // namespace phi
@@ -42,7 +42,6 @@ struct CalcReducedAttnScoresParams : public FlashAttnParamsBase {
                             _head_size,
                             _scale,
                             /*_causal=*/false,
-                            /*_attn_mask_start_row=*/0,
                             q_dtype,
                             paddle::optional<DenseTensor>{},
                             paddle::optional<DenseTensor>{}) {}
 
@@ -25,6 +25,7 @@
 #include "paddle/phi/kernels/funcs/elementwise_base.h"
 #include "paddle/phi/kernels/gpu/flash_attn_utils.h"
 #include "paddle/phi/kernels/reduce_sum_kernel.h"
+#include "paddle/phi/kernels/slice_kernel.h"
 
 COMMON_DECLARE_bool(cudnn_deterministic);
 
@@ -303,10 +304,9 @@ void FlashAttnUnpaddedGradBaseKernel(
                            dropout,
                            scale,
                            causal,
-                           0,  // attn_mask_start_row
                            q.dtype(),
                            attn_mask,
-                           nullptr,  // attn_mask_start_row_indices
+                           nullptr,  // startend_row_indices,
                            seed_offset.data<int64_t>());
 
   VLOG(10) << "FlashAttn bwd seed: " << params.seed
@@ -632,11 +632,10 @@ void FlashAttnGradBaseKernel(
     const DenseTensor& softmax_lse,
     const DenseTensor& seed_offset,
     const paddle::optional<DenseTensor>& attn_mask,
-    const paddle::optional<DenseTensor>& attn_mask_start_row_indices,
+    const paddle::optional<DenseTensor>& startend_row_indices,
     const DenseTensor& dout,
     float dropout,
     bool causal,
-    int attn_mask_start_row,
     DenseTensor* dq,
     DenseTensor* dk,
     DenseTensor* dv) {
@@ -712,10 +711,9 @@ void FlashAttnGradBaseKernel(
                            dropout,
                            softmax_scale,
                            causal,
-                           attn_mask_start_row,
                            q.dtype(),
                            attn_mask,
-                           attn_mask_start_row_indices,
+                           startend_row_indices,
                            seed_offset.data<int64_t>());
 
   VLOG(10) << "[FlashAttn Forward] q.shape=[" << q.dims() << "], k.shape=["
@@ -731,6 +729,57 @@ void FlashAttnGradBaseKernel(
 
   int num_splits = get_num_split();
 
+  DenseTensor flashmask_maxmin, 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 = params.startend_row_indices != nullptr;
+  if (is_flashmask) {
+    PADDLE_ENFORCE_EQ(
+        startend_row_indices->dims().size(),
+        4,
+        phi::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,
+        phi::errors::InvalidArgument("flashmask_attention startend_row_indices "
+                                     "mask_bounds must in [1,2,4]"));
+    auto flashmask_maxmin_shape = params.startend_row_indices->dims();
+    flashmask_maxmin_shape[2] = (flashmask_maxmin_shape[2] + 31) / 32 * 8;
+    flashmask_maxmin.set_type(phi::DataType::INT32);
+    flashmask_maxmin.Resize(flashmask_maxmin_shape);
+    ctx.template Alloc<T>(&flashmask_maxmin);
+
+    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();
+    }
+  }
+
 #ifdef PADDLE_WITH_HIP
   bool succ = phi::dynload::flash_attn_bwd(
       dout.data(),
@@ -799,13 +848,12 @@ void FlashAttnGradBaseKernel(
       params.offset,
       params.attn_mask_tensor ? params.attn_mask_tensor->data() : nullptr,
       params.attn_mask_tensor ? params.mask_dims.data() : nullptr,
-      params.attn_mask_start_row_indices_tensor
-          ? params.attn_mask_start_row_indices_tensor->data()
-          : nullptr,
-      params.attn_mask_start_row_indices_tensor
-          ? params.attn_mask_start_row_indices_dims.data()
-          : nullptr,
-      params.attn_mask_start_row,
+      is_flashmask ? downstart_row_indices_data : nullptr,
+      is_flashmask ? params.startend_row_indices_dims.data() : nullptr,
+      is_flashmask ? upend_row_indices_data : nullptr,
+      is_flashmask ? downend_row_indices_data : nullptr,
+      is_flashmask ? upstart_row_indices_data : nullptr,
+      is_flashmask ? flashmask_maxmin.data() : nullptr,
       q.strides()[1],
       k.strides()[1],
       v.strides()[1],
@@ -934,7 +982,6 @@ void FlashAttnGradKernel(const Context& ctx,
                                       dout,
                                       dropout,
                                       causal,
-                                      0,
                                       dq,
                                       dk,
                                       dv);
@@ -983,7 +1030,6 @@ void FlashAttnQKVPackedGradKernel(
                                       dout,
                                       dropout,
                                       causal,
-                                      0,
                                       &dq,
                                       &dk,
                                       &dv);
@@ -993,22 +1039,20 @@ void FlashAttnQKVPackedGradKernel(
 }
 
 template <typename T, typename Context>
-void FlashAttnWithSparseGradKernel(
-    const Context& ctx,
-    const DenseTensor& q,
-    const DenseTensor& k,
-    const DenseTensor& v,
-    const DenseTensor& attn_mask_start_row_indices,
-    const DenseTensor& out,
-    const DenseTensor& softmax_lse,
-    const DenseTensor& seed_offset,
-    const DenseTensor& dout,
-    float dropout,
-    bool causal,
-    int attn_mask_start_row,
-    DenseTensor* dq,
-    DenseTensor* dk,
-    DenseTensor* dv) {
+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) {
   if (dq) {
     ctx.template Alloc<T>(dq);
   }
@@ -1026,11 +1070,10 @@ void FlashAttnWithSparseGradKernel(
                                       softmax_lse,
                                       seed_offset,
                                       paddle::none,
-                                      attn_mask_start_row_indices,
+                                      startend_row_indices,
                                       dout,
                                       dropout,
                                       causal,
-                                      attn_mask_start_row,
                                       dq,
                                       dk,
                                       dv);
@@ -1073,10 +1116,10 @@ PD_REGISTER_KERNEL(flash_attn_qkvpacked_grad,
   kernel->InputAt(3).SetBackend(phi::Backend::ALL_BACKEND);  // seed_offset
 }
 
-PD_REGISTER_KERNEL(flash_attn_with_sparse_mask_grad,
+PD_REGISTER_KERNEL(flashmask_attention_grad,
                    GPU,
                    ALL_LAYOUT,
-                   phi::FlashAttnWithSparseGradKernel,
+                   phi::FlashMaskGradKernel,
                    phi::dtype::float16,
                    phi::dtype::bfloat16) {
   kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);  // seed_offset