fix backend bug (#52526)

380a9bf7 · Chitsing KUI · GitHub · 8ac5a6b6 · 380a9bf7
隐藏空白更改
内联并排

Showing with 2 addition and 2 deletion

paddle/phi/kernels/gpu/flash_attn_grad_kernel.cu paddle/phi/kernels/gpu/flash_attn_grad_kernel.cu +2 -2

未找到文件。
--- a/paddle/phi/kernels/gpu/flash_attn_grad_kernel.cu
+++ b/paddle/phi/kernels/gpu/flash_attn_grad_kernel.cu
@@ -229,7 +229,7 @@ PD_REGISTER_KERNEL(flash_attn_unpadded_grad,
                   phi::FlashAttnUnpaddedGradKernel,
                   phi::dtype::float16,
                   phi::dtype::bfloat16) {
-  kernel->InputAt(7).SetBackend(phi::Backend::CPU);  // seed_offset
+  kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);  // seed_offset
 }

 PD_REGISTER_KERNEL(flash_attn_grad,
@@ -238,5 +238,5 @@ PD_REGISTER_KERNEL(flash_attn_grad,
                   phi::FlashAttnGradKernel,
                   phi::dtype::float16,
                   phi::dtype::bfloat16) {
-  kernel->InputAt(5).SetBackend(phi::Backend::CPU);  // seed_offset
+  kernel->InputAt(5).SetBackend(phi::Backend::ALL_BACKEND);  // seed_offset
 }