add blockmask

Author: starcrown001

paddle/phi/backends/dynload/flashmaskv2.h

@@ -225,6 +225,10 @@ FLASHMASK_V2_HANDLE_ROUTINE(ut_start_ptr)
 FLASHMASK_V2_HANDLE_ROUTINE(ut_end_ptr)
 FLASHMASK_V2_HANDLE_ROUTINE(flashmask_maxmin_ptr)
 
+FLASHMASK_V2_HANDLE_ROUTINE(m_block_dim)
+FLASHMASK_V2_HANDLE_ROUTINE(n_block_dim)
+FLASHMASK_V2_HANDLE_ROUTINE(block_mask_ptr)
+
 #define FLASHMASK_V2_BWD_HANDLE_ROUTINE(type, member)                          \
   DECLARE_DYNAMIC_LOAD_FLASHMASK_V2_WRAP(flashmaskv2_bwd_params_get_##member); \
   DECLARE_DYNAMIC_LOAD_FLASHMASK_V2_WRAP(flashmaskv2_bwd_params_set_##member);

paddle/phi/kernels/gpu/flash_attn_v3_grad_kernel.cu

@@ -850,6 +850,8 @@ void FlashMaskV2GradBaseKernel(
         &seqused_k_,  // b. If given, only this many elements of each batch
                       // element's keys are used.
     const paddle::optional<DenseTensor> &startend_row_indices_,
+    const paddle::optional<DenseTensor>
+        &block_mask_indices_,  // （(b,h,s//128,s//128)
     int max_seqlen_q_,
     int max_seqlen_k_,
     float const softmax_scale,
@@ -1080,6 +1082,50 @@ void FlashMaskV2GradBaseKernel(
     }
   }
 
+  bool const is_blockmask = block_mask_indices_.is_initialized();
+  DenseTensor block_mask_indices;
+  if (is_blockmask) block_mask_indices = block_mask_indices_.get();
+
+  if (is_blockmask) {
+    PADDLE_ENFORCE_EQ(
+        is_flashmask,
+        true,
+        common::errors::InvalidArgument(
+            "blockmask should be used with flashmask at the same time "));
+
+    PADDLE_ENFORCE_EQ(block_mask_indices.dims().size(),
+                      4,
+                      common::errors::InvalidArgument(
+                          "blockmask receive blockmask_indices with dim "
+                          "[batch_size, num_heads, blocklen_q, blocklen_k]"));
+
+    PADDLE_ENFORCE_EQ(block_mask_indices.dims()[2],
+                      (seqlen_q + 127) / 128,
+                      common::errors::InvalidArgument(
+                          "blockmask only supports blockdim_q = 128 now"));
+
+    PADDLE_ENFORCE_EQ(block_mask_indices.dims()[3],
+                      (seqlen_k + 127) / 128,
+                      common::errors::InvalidArgument(
+                          "blockmask only supports blockdim_k = 128 now"));
+
+    PADDLE_ENFORCE_EQ(
+        block_mask_indices.dims()[1],
+        startend_row_indices.dims()[1],
+        common::errors::InvalidArgument(
+            "blockmask only supports same dim num_heads with flashmask now"));
+
+    PADDLE_ENFORCE_LE(seqlen_k,
+                      1024 * 128,
+                      common::errors::InvalidArgument(
+                          "blockmask only supports seqlen <= 128k in bwd now"));
+
+    PADDLE_ENFORCE_LE(seqlen_q,
+                      1024 * 128,
+                      common::errors::InvalidArgument(
+                          "blockmask only supports seqlen <= 128k in bwd now"));
+  }
+
   const bool has_lt_start = lt_start_row_indices.initialized();
   const bool has_lt_end = lt_end_row_indices.initialized();
   const bool has_ut_start = ut_start_row_indices.initialized();
@@ -1284,6 +1330,7 @@ void FlashMaskV2GradBaseKernel(
   if (softmax_lse_log2) {
     dev_ctx.template Alloc<float>(softmax_lse_log2);
   }
+  // std::cout << "dq_accum:" << dq_accum->dims() << std::endl;
   if (dq_accum) {
     if (!is_varlen) {
       dq_accum->Resize(common::make_ddim(
@@ -1292,6 +1339,7 @@ void FlashMaskV2GradBaseKernel(
       dq_accum->Resize(common::make_ddim(
           {num_heads, total_q_padded_rounded * head_size_rounded}));
     }
+    // std::cout << "enter:" << dq_accum->dims() << std::endl;
     dev_ctx.template Alloc<float>(dq_accum);
   }
   if (num_heads_k != num_heads) {  // MQA / GQA
@@ -1457,6 +1505,17 @@ void FlashMaskV2GradBaseKernel(
     dynload::flashmaskv2_bwd_params_set_h_h_flashmask_ratio(params_handle, 0);
   }
 
+  if (is_blockmask) {
+    // xhy: blockmask is now only support blockdim_q k = 128
+    dynload::flashmaskv2_bwd_params_set_m_block_dim(params_handle, 128);
+    dynload::flashmaskv2_bwd_params_set_n_block_dim(params_handle, 128);
+    dynload::flashmaskv2_bwd_params_set_block_mask_ptr(
+        params_handle,
+        const_cast<int32_t *>(block_mask_indices.data<int32_t>()));
+    // phi::funcs::SetConstant<Context, T> set_dq_zero;
+    // // dev_ctx.template Alloc<T>(dq);
+    // set_dq_zero(dev_ctx, dq, T{0});
+  }
 #ifdef FLASHATTENTION_DISABLE_LOCAL
   PADDLE_ENABLE_EQ(
       !dynload::flashmaskv2_bwd_params_get_is_local(params_handle),
@@ -1504,6 +1563,7 @@ void FlashMaskV2GradKernel(
     const DenseTensor &out,
     const DenseTensor &softmax_lse,
     const DenseTensor &startend_row_indices,  // TODO(xiehaoyang): remove this
+    const paddle::optional<DenseTensor> &block_mask_indices,
     const DenseTensor &out_grad,
     float const softmax_scale,
     bool is_causal,
@@ -1540,6 +1600,7 @@ void FlashMaskV2GradKernel(
                                         paddle::none,
                                         paddle::none,
                                         startend_row_indices,
+                                        block_mask_indices,
                                         0,  // max_seqlen_q,
                                         0,  // max_seqlen_k,
                                         softmax_scale,

paddle/phi/kernels/gpu/flash_attn_v3_kernel.cu

@@ -1235,7 +1235,10 @@ void FlashMaskV2BaseKernel(
     const paddle::optional<DenseTensor> &k_descale_,  // (b, h_k)
     const paddle::optional<DenseTensor> &v_descale_,  // (b, h_k)
     const paddle::optional<DenseTensor> &scheduler_metadata_,  // (b + 1)
-    const paddle::optional<DenseTensor> &startend_row_indices_,
+    const paddle::optional<DenseTensor>
+        &startend_row_indices_,  // （b,h,s_1,[1,2,4])
+    const paddle::optional<DenseTensor>
+        &block_mask_indices_,  // （(b,h,s// 128,s // 128)
     const int
         max_seqlen_q_,  // if max_seqlen_q_ is set to 0, it indicates that it is
                         // uninitialized and should not be referenced
@@ -1432,6 +1435,7 @@ void FlashMaskV2BaseKernel(
   }
 
   bool const is_flashmask = startend_row_indices_.is_initialized();
+  bool const is_blockmask = block_mask_indices_.is_initialized();
 
   // This needs to go before kBlockM & kBlockN since we rely on the correct
   // window_size and is_causal to set kBlockM
@@ -2068,6 +2072,8 @@ void FlashMaskV2BaseKernel(
   // flashmask
   DenseTensor startend_row_indices;
   if (is_flashmask) startend_row_indices = startend_row_indices_.get();
+  DenseTensor block_mask_indices;
+  if (is_blockmask) block_mask_indices = block_mask_indices_.get();
   DenseTensor flashmask_maxmin, lt_start_row_indices, lt_end_row_indices,
       ut_start_row_indices, ut_end_row_indices;
   if (is_flashmask) {
@@ -2142,6 +2148,45 @@ void FlashMaskV2BaseKernel(
     }
   }
 
+  if (is_blockmask) {
+    PADDLE_ENFORCE_EQ(
+        is_flashmask,
+        true,
+        common::errors::InvalidArgument(
+            "blockmask should be used with flashmask at the same time "));
+
+    PADDLE_ENFORCE_EQ(block_mask_indices.dims().size(),
+                      4,
+                      common::errors::InvalidArgument(
+                          "blockmask receive blockmask_indices with dim "
+                          "[batch_size, num_heads, blocklen_q, blocklen_k]"));
+
+    PADDLE_ENFORCE_EQ(block_mask_indices.dims()[2],
+                      (seqlen_q + 127) / 128,
+                      common::errors::InvalidArgument(
+                          "blockmask is now only support blockdim_q = 128 "));
+
+    PADDLE_ENFORCE_EQ(block_mask_indices.dims()[3],
+                      (seqlen_k + 127) / 128,
+                      common::errors::InvalidArgument(
+                          "blockmask is now only support blockdim_k = 128 "));
+
+    PADDLE_ENFORCE_EQ(
+        block_mask_indices.dims()[1],
+        startend_row_indices.dims()[1],
+        common::errors::InvalidArgument("blockmask is now only support same "
+                                        "dim num_heads with flashmask "));
+  }
+
+  if (is_blockmask) {
+    // xhy: blockmask is now only support blockdim_q k = 128
+    dynload::flashmaskv2_fwd_params_set_m_block_dim(params_handle, 128);
+    dynload::flashmaskv2_fwd_params_set_n_block_dim(params_handle, 128);
+    dynload::flashmaskv2_fwd_params_set_block_mask_ptr(
+        params_handle,
+        const_cast<int32_t *>(block_mask_indices.data<int32_t>()));
+  }
+
   if (is_flashmask) {
     if (lt_start_row_indices.initialized())
       dynload::flashmaskv2_fwd_params_set_lt_start_ptr(
@@ -2260,6 +2305,7 @@ void FlashMaskV2Kernel(const Context &dev_ctx,
                        const DenseTensor &k,
                        const DenseTensor &v,
                        const DenseTensor &startend_row_indices,
+                       const paddle::optional<DenseTensor> &block_mask_indices,
                        const float softmax_scale,
                        bool is_causal,
                        DenseTensor *out,
@@ -2290,6 +2336,7 @@ void FlashMaskV2Kernel(const Context &dev_ctx,
                                     paddle::none,  // v_descale_
                                     paddle::none,  // scheduler_metadata_
                                     startend_row_indices,
+                                    block_mask_indices,
                                     0,  // max_seqlen_q_
                                     0,  // max_seqlen_k_
                                     softmax_scale,
@@ -2333,4 +2380,7 @@ PD_REGISTER_KERNEL(flashmask_attention_v2,
                    ALL_LAYOUT,
                    phi::FlashMaskV2Kernel,
                    phi::float16,
-                   phi::bfloat16) {}
+                   phi::bfloat16) {
+  kernel->InputAt(4).SetBackend(
+      phi::Backend::ALL_BACKEND);  // block_mask_indices
+}

paddle/phi/ops/yaml/backward.yaml

@@ -1228,8 +1228,9 @@
     data_type: q
 
 - backward_op : flashmask_attention_v2_grad
-  forward : flashmask_attention_v2 (Tensor q, Tensor k, Tensor v, Tensor startend_row_indices, float softmax_scale, bool is_causal) -> Tensor(out), Tensor(softmax_lse)
-  args : (Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, Tensor startend_row_indices, Tensor out_grad, float softmax_scale, bool is_causal)
+  forward : flashmask_attention_v2 (Tensor q, Tensor k, Tensor v, Tensor startend_row_indices,Tensor block_mask_indices, float softmax_scale, bool is_causal) -> Tensor(out), Tensor(softmax_lse)
+  args : (Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, Tensor startend_row_indices, Tensor block_mask_indices, Tensor out_grad, float softmax_scale, bool is_causal)
+  optional : block_mask_indices
   output : Tensor(q_grad), Tensor(k_grad), Tensor(v_grad)
   infer_meta :
     func : FlashAttnGradInferMeta

paddle/phi/ops/yaml/ops.yaml

@@ -2153,8 +2153,9 @@
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : flashmask_attention_v2
-  args : (Tensor q, Tensor k, Tensor v, Tensor startend_row_indices, float softmax_scale, bool is_causal)
+  args : (Tensor q, Tensor k, Tensor v, Tensor startend_row_indices, Tensor block_mask_indices, float softmax_scale, bool is_causal)
   output : Tensor(out), Tensor(softmax_lse)
+  optional : block_mask_indices
   infer_meta :
     func : FlashMaskV2InferMeta
     param : [q, k, v]

python/paddle/nn/functional/flash_attention.py

@@ -1575,6 +1575,7 @@ def flashmask_attention(
     training: bool = True,
     name: str | None = None,
     softmax_scale: float | None = None,
+    block_mask_indices: Tensor | None = None,
 ):
     r"""
     FlashMask: Official Implementation
@@ -1635,6 +1636,26 @@ def flashmask_attention(
         training (bool): Whether the module is in training mode. Default is True.
         name (str, optional): Name of the operation. Default is None. Normally, users do not need to set this property.
             For more information, refer to :ref:`api_guide_Name` .
+        block_mask_indices (tensor, optional):
+            block_mask_indices (Tensor, optional):
+                A 4-D integer mask tensor indicating whether each block in the attention matrix should be kept or masked. Must be used together with flashmask.
+                The shape should be [batch_size, num_heads, blocklen_q, blocklen_k], where:
+
+                blocklen_q = ceil(seqlen_q / 128), i.e., block_mask_indices.shape[2] must be (seqlen_q + 127) // 128
+                blocklen_k = ceil(seqlen_k / 128), i.e., block_mask_indices.shape[3] must be (seqlen_k + 127) // 128
+                block_mask_indices.shape[1] (number of heads) must match the num_heads dimension of the flashmask
+                Both seqlen_q and seqlen_k must be less than or equal to 128 * 1024
+                The dtype should be int32, and each element should be either 0 or 1.
+                A value of 1 indicates that the corresponding block is kept (not masked), while 0 means the block is masked.
+
+                Usage Notes:
+
+                Only supported when blockdim_q = blockdim_k = 128 now.
+                Only supported when headdim = 128 now.
+                This argument must be provided together with flashmask.
+                The mask will be applied at the block level: each [i, j] position in block_mask_indices controls whether the corresponding [128 x 128] block in the attention matrix is masked.
+                Any mismatch in expected shape or head dimension will raise an error.
+
 
     Returns
         Tensor. The computed attention result with the same shape as the input `query`.
@@ -2207,6 +2228,12 @@ def flashmask_attention(
                 startend_row_indices, min=0, max=sq
             ).repeat_interleave(bsz, 0)
 
+    if block_mask_indices is not None:
+        # xhy: can set a full startend_row_indices for block_mask_attn when using block_mask_attn?
+        assert startend_row_indices is not None, (
+            "must provide startend_row_indices when using block_mask_attn"
+        )
+
     if startend_row_indices is None:
         (
             out,
@@ -2248,6 +2275,33 @@ def flashmask_attention(
             "startend_row_indices head_num must be equal to 1(broadcast) or head_num_k."
         )
 
+        if block_mask_indices is not None:
+            assert block_mask_indices.dtype == paddle.int32, (
+                f"block_mask_indices.dtype must be paddle.int32, but got {block_mask_indices.dtype}"
+            )
+
+            assert block_mask_indices.shape[0] == key.shape[0], (
+                f"block_mask_indices.shape[0] must be equal to batch_size, but got {block_mask_indices.shape[0]} and {key.shape[0]}"
+            )
+
+            assert (
+                block_mask_indices.shape[1] == startend_row_indices.shape[1]
+            ), (
+                f"block_mask_indices.shape[1] must be equal to startend_row_indices.shape[1], but got {block_mask_indices.shape[1]} and {key.shape[2]}"
+            )
+
+            assert (
+                block_mask_indices.shape[2] == (query.shape[1] + 127) // 128
+            ), "block_size must be 128 when using block_mask_attn"
+
+            assert block_mask_indices.shape[3] == (key.shape[1] + 127) // 128, (
+                "block_size must be 128 when using block_mask_attn"
+            )
+
+            assert key.shape[3] == 128, (
+                "headdim must be 128 when using block_mask_attn"
+            )
+
         if causal:
             if startend_row_indices.shape[-1] == 1:
                 has_end = False
@@ -2329,7 +2383,13 @@ def flashmask_attention(
                 out,
                 result_softmax_lse,
             ) = _C_ops.flashmask_attention_v2(
-                query, key, value, startend_row_indices, softmax_scale, causal
+                query,
+                key,
+                value,
+                startend_row_indices,
+                block_mask_indices,
+                softmax_scale,
+                causal,
             )
         else:
             raise ValueError(f"Invalid flash attention version: {fa_version}")