[XPU] update xhpc to 0405, solving flashmask error in some cases, and add xpu_wait for side_stream (#71908)

runzhech · web-flow · commit 91d623bf072f · 2025-04-07T15:13:20.000+08:00
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,7 +30,7 @@ set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 set(XPU_XPUDNN_LIB_NAME "libxpu_dnn.so")
 
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "dev/20250318")
+  set(XPU_XHPC_BASE_DATE "dev/20250405")
 endif()
 set(XPU_XCCL_BASE_VERSION "3.0.2.5") # For XRE5
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc b/paddle/phi/kernels/xpu/flash_attn_grad_kernel.cc
@@ -232,6 +232,8 @@ void FlashAttnGradKernelBase(
       is_flashmask ? flashmask_stream : nullptr);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_bwd");
   if (is_flashmask && flashmask_stream != nullptr) {
+    r = xpu_wait(flashmask_stream);
+    PADDLE_ENFORCE_XPU_SUCCESS(r);
     xpu_stream_destroy(flashmask_stream);
   }
 }
diff --git a/paddle/phi/kernels/xpu/flash_attn_kernel.cc b/paddle/phi/kernels/xpu/flash_attn_kernel.cc
@@ -210,12 +210,15 @@ void FlashAttnKernelBase(
       (const int*)upstart_row_indices_data,       // upstart_row_indices_data
       (const int*)upend_row_indices_data,         // upend_row_indices_data
       is_flashmask ? startend_row_indices->dims()[1]
-                   : 0,                          // flash_mask_head_num
-      nullptr,                                   // flashmask_maxmin
-      is_flashmask ? flashmask_stream : nullptr  // side_stream
+                   : 0,                           // flash_mask_head_num
+      nullptr,                                    // flashmask_maxmin
+      is_flashmask ? flashmask_stream : nullptr,  // side_stream
+      0                                           // fixlen_batch_num
   );
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_fwd");
   if (is_flashmask && flashmask_stream != nullptr) {
+    r = xpu_wait(flashmask_stream);
+    PADDLE_ENFORCE_XPU_SUCCESS(r);
     xpu_stream_destroy(flashmask_stream);
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -232,6 +232,8 @@ void FlashAttnGradKernelBase(`
`232`	`232`	`is_flashmask ? flashmask_stream : nullptr);`
`233`	`233`	`PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_bwd");`
`234`	`234`	`if (is_flashmask && flashmask_stream != nullptr) {`
	`235`	`+ r = xpu_wait(flashmask_stream);`
	`236`	`+ PADDLE_ENFORCE_XPU_SUCCESS(r);`
`235`	`237`	`xpu_stream_destroy(flashmask_stream);`
`236`	`238`	`}`
`237`	`239`	`}`