fix block_multi_head_attention and add shape enforce in flash_attn (#64416)

kircle888 · web-flow · commit 3e275d54582b · 2024-05-19T23:03:47.000-07:00
* fix block_multi_head_attention and add shape enforce in flash_attn

* add comment
diff --git a/paddle/phi/kernels/fusion/gpu/block_multi_head_attention_kernel.cu b/paddle/phi/kernels/fusion/gpu/block_multi_head_attention_kernel.cu
@@ -400,6 +400,10 @@ void DispatchWithDtype(
                              max_seq_len,
                              dim_head);
       VLOG(3) << "qkv split end";
+      // Reshape fmha_buf to 3-D because FlashAttnUnpaddedKernel requries
+      // q,k,v,out all in 3-D [token_num, num_head, dim_head].
+      auto fmha_shape = fmha_buf.dims();
+      fmha_buf.Resize({token_num, num_head, dim_head});
       phi::FlashAttnUnpaddedKernel<T>(dev_ctx,
                                       unpadding_q,
                                       unpadding_k,
@@ -420,6 +424,8 @@ void DispatchWithDtype(
                                       &softmax_out,
                                       &softmax_lse,
                                       &seed_offset);
+      // Reshape fmha_buf back (to 2-D), to not affect following codes.
+      fmha_buf.Resize(fmha_shape);
     } else {
       qkv_transpose_split<T>(
           dev_ctx,
diff --git a/paddle/phi/kernels/gpu/flash_attn_kernel.cu b/paddle/phi/kernels/gpu/flash_attn_kernel.cu
@@ -56,7 +56,7 @@ void FlashAttnUnpaddedBaseKernel(
     DenseTensor* seed_offset,
     bool varlen_padded) {
 #ifdef PADDLE_WITH_FLASHATTN
-  ctx.template Alloc<T>(out);
+  if (!out->IsInitialized()) ctx.template Alloc<T>(out);
   if (varlen_padded) {
     std::vector<const DenseTensor*> inputs{};
     std::vector<DenseTensor*> outputs{out};
@@ -72,6 +72,21 @@ void FlashAttnUnpaddedBaseKernel(
       3,
       phi::errors::InvalidArgument("flash_attn_raw receive input with dim "
                                    "[total_seq_len, num_heads, head_dim]"));
+  PADDLE_ENFORCE_EQ(
+      k.dims().size(),
+      3,
+      phi::errors::InvalidArgument("flash_attn_raw receive input with dim "
+                                   "[total_seq_len, num_heads, head_dim]"));
+  PADDLE_ENFORCE_EQ(
+      v.dims().size(),
+      3,
+      phi::errors::InvalidArgument("flash_attn_raw receive input with dim "
+                                   "[total_seq_len, num_heads, head_dim]"));
+  PADDLE_ENFORCE_EQ(
+      out->dims().size(),
+      3,
+      phi::errors::InvalidArgument("flash_attn_raw receive input with dim "
+                                   "[total_seq_len, num_heads, head_dim]"));
 
   const int64_t batch_size = cu_seqlens_q.numel() - 1;
   const int64_t num_heads = dims[1];
@@ -314,7 +329,21 @@ void FlashAttnBaseKernel(
                     phi::errors::InvalidArgument(
                         "flash_attn receive input with dim "
                         "[batch_size, seq_len, num_heads, head_dim]"));
-
+  PADDLE_ENFORCE_EQ(k.dims().size(),
+                    4,
+                    phi::errors::InvalidArgument(
+                        "flash_attn receive input with dim "
+                        "[batch_size, seq_len, num_heads, head_dim]"));
+  PADDLE_ENFORCE_EQ(v.dims().size(),
+                    4,
+                    phi::errors::InvalidArgument(
+                        "flash_attn receive input with dim "
+                        "[batch_size, seq_len, num_heads, head_dim]"));
+  PADDLE_ENFORCE_EQ(out->dims().size(),
+                    4,
+                    phi::errors::InvalidArgument(
+                        "flash_attn receive input with dim "
+                        "[batch_size, seq_len, num_heads, head_dim]"));
   const int64_t batch_size = dims[0];
   const int64_t seqlen_q = dims[1];
   const int64_t num_heads = dims[2];
@@ -360,8 +389,7 @@ void FlashAttnBaseKernel(
     VLOG(10) << "[FlashAttn Forward] attn_mask.shape=["
              << (attn_mask.get_ptr())->dims() << "]";
   }
-
-  ctx.template Alloc<T>(out);
+  if (!out->IsInitialized()) ctx.template Alloc<T>(out);
 
   cudaStream_t stream = ctx.stream();
 
