Add attn_mask supported for FlashAttnKernel.

paddle/phi/api/yaml/backward.yaml

@@ -818,8 +818,9 @@
   inplace : (out_grad -> x_grad)
 
 - backward_op : flash_attn_grad
-  forward : flash_attn (Tensor q, Tensor k, Tensor v, Tensor fixed_seed_offset, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "") -> Tensor(out), Tensor(softmax), Tensor(softmax_lse), Tensor(seed_offset)
-  args : (Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, Tensor seed_offset, Tensor out_grad, float dropout = 0.0, bool causal = false)
+  forward : flash_attn (Tensor q, Tensor k, Tensor v,  Tensor fixed_seed_offset, Tensor attn_mask, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "") -> Tensor(out), Tensor(softmax), Tensor(softmax_lse), Tensor(seed_offset)
+  args : (Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, Tensor seed_offset, Tensor attn_mask, Tensor out_grad, float dropout = 0.0, bool causal = false)
+  optional : attn_mask
   output : Tensor(q_grad), Tensor(k_grad), Tensor(v_grad)
   infer_meta :
     func : FlashAttnGradInferMeta
@@ -829,8 +830,9 @@
     data_type: q
 
 - backward_op : flash_attn_unpadded_grad
-  forward : flash_attn_unpadded (Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, Tensor fixed_seed_offset, int64_t max_seqlen_q, int64_t max_seqlen_k, float scale, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "") -> Tensor(out), Tensor(softmax), Tensor(softmax_lse), Tensor(seed_offset)
-  args : (Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, Tensor out, Tensor softmax_lse, Tensor seed_offset, Tensor out_grad, int64_t max_seqlen_q, int64_t max_seqlen_k, float scale, float dropout = 0.0, bool causal = false)
+  forward : flash_attn_unpadded (Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, Tensor fixed_seed_offset, Tensor attn_mask, int64_t max_seqlen_q, int64_t max_seqlen_k, float scale, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "") -> Tensor(out), Tensor(softmax), Tensor(softmax_lse), Tensor(seed_offset)
+  args : (Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, Tensor out, Tensor softmax_lse, Tensor seed_offset, Tensor attn_mask, Tensor out_grad, int64_t max_seqlen_q, int64_t max_seqlen_k, float scale, float dropout = 0.0, bool causal = false)
+  optional : attn_mask
   output : Tensor(q_grad), Tensor(k_grad), Tensor(v_grad)
   infer_meta :
     func : FlashAttnGradInferMeta

paddle/phi/api/yaml/ops.yaml

@@ -910,9 +910,9 @@
   backward : fill_diagonal_tensor_grad
 
 - op : flash_attn
-  args : (Tensor q, Tensor k, Tensor v, Tensor fixed_seed_offset, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "")
+  args : (Tensor q, Tensor k, Tensor v, Tensor fixed_seed_offset, Tensor attn_mask, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "")
   output : Tensor(out), Tensor(softmax), Tensor(softmax_lse), Tensor(seed_offset)
-  optional : fixed_seed_offset
+  optional : fixed_seed_offset, attn_mask
   infer_meta :
     func : FlashAttnInferMeta
     param : [q, k, v]
@@ -923,9 +923,9 @@
   backward : flash_attn_grad
 
 - op : flash_attn_unpadded
-  args : (Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q,  Tensor cu_seqlens_k, Tensor fixed_seed_offset, int64_t max_seqlen_q, int64_t max_seqlen_k, float scale, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "")
+  args : (Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q,  Tensor cu_seqlens_k, Tensor fixed_seed_offset, Tensor attn_mask, int64_t max_seqlen_q, int64_t max_seqlen_k, float scale, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "")
   output : Tensor(out), Tensor(softmax), Tensor(softmax_lse), Tensor(seed_offset)
-  optional : fixed_seed_offset
+  optional :  fixed_seed_offset , attn_mask
   infer_meta :
     func : FlashAttnInferMeta
     param : [q, k, v]

paddle/phi/kernels/flash_attn_grad_kernel.h

@@ -29,6 +29,7 @@ void FlashAttnUnpaddedGradKernel(const Context& ctx,
                                  const DenseTensor& out,
                                  const DenseTensor& softmax_lse,
                                  const DenseTensor& seed_offset,
+                                 const paddle::optional<DenseTensor>& attn_mask,
                                  const DenseTensor& dout,
                                  int64_t max_seqlen_q,
                                  int64_t max_seqlen_k,
@@ -47,6 +48,7 @@ void FlashAttnGradKernel(const Context& ctx,
                          const DenseTensor& out,
                          const DenseTensor& softmax_lse,
                          const DenseTensor& seed_offset,
+                         const paddle::optional<DenseTensor>& attn_mask,
                          const DenseTensor& dout,
                          float dropout,
                          bool causal,

paddle/phi/kernels/flash_attn_kernel.h

@@ -19,6 +19,74 @@
 
 namespace phi {
 
+template <typename T, typename Context>
+void FlashAttnFwdWithBiasAndMask(
+    const Context& ctx,
+    const void*
+        q,  // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
+    const void*
+        k,  // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+    const void*
+        v,  // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+    void*
+        out,  // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+    const int32_t*
+        cu_seqlens_q,  // int32, batch_size+1, starting offset of each sequence
+    const int32_t*
+        cu_seqlens_k,  // int32, batch_size+1, starting offset of each sequence
+    const int total_q,
+    const int total_k,
+    const int batch_size,
+    const int num_heads,
+    const int head_size,
+    const int max_seqlen_q,
+    const int max_seqlen_k,
+    const float dropout,
+    const float scale,
+    const bool zero_tensors,
+    const bool is_bf16,
+    const int num_splits,   // SMs per attention matrix, can be 1
+    void* softmax_lse_ptr,  // softmax log_sum_exp
+    cudaStream_t stream,
+    uint64_t seed,
+    uint64_t offset,
+    const void* attn_mask,
+    const int64_t* mask_dims);
+
+template <typename T, typename Context>
+void FlashAttnFwd(
+    const Context& ctx,
+    const void*
+        q,  // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
+    const void*
+        k,  // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+    const void*
+        v,  // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+    void*
+        out,  // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+    const void*
+        cu_seqlens_q,  // int32, batch_size+1, starting offset of each sequence
+    const void*
+        cu_seqlens_k,  // int32, batch_size+1, starting offset of each sequence
+    const int total_q,
+    const int total_k,
+    const int batch_size,
+    const int num_heads,
+    const int head_size,
+    const int max_seqlen_q,
+    const int max_seqlen_k,
+    const float dropout,
+    const float scale,
+    const bool zero_tensors,
+    const bool causal,
+    const bool is_bf16,
+    const int num_splits,   // SMs per attention matrix, can be 1
+    void* softmax_lse_ptr,  // softmax log_sum_exp
+    const bool return_softmax,
+    cudaStream_t stream,
+    uint64_t seed,
+    uint64_t offset);
+
 template <typename T, typename Context>
 void FlashAttnUnpaddedKernel(
     const Context& ctx,
@@ -28,6 +96,7 @@ void FlashAttnUnpaddedKernel(
     const DenseTensor& cu_seqlens_q,
     const DenseTensor& cu_seqlens_k,
     const paddle::optional<DenseTensor>& fixed_seed_offset,
+    const paddle::optional<DenseTensor>& attn_mask,
     int64_t max_seqlen_q,
     int64_t max_seqlen_k,
     float scale,
@@ -47,6 +116,7 @@ void FlashAttnKernel(const Context& ctx,
                      const DenseTensor& k,
                      const DenseTensor& v,
                      const paddle::optional<DenseTensor>& fixed_seed_offset,
+                     const paddle::optional<DenseTensor>& attn_mask,
                      float dropout,
                      bool causal,
                      bool return_softmax,