fix loop cond for num_split=1 (PaddlePaddle#5)

kircle888 · Apr 12, 2023 · 5ff4bbf · 5ff4bbf
1 parent f0edf24
commit 5ff4bbf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/csrc/flash_attn/flash_attn.cpp b/csrc/flash_attn/flash_attn.cpp
@@ -426,7 +426,7 @@ bool flash_attn_bwd(
         // 1) num_splits == 1
         // 2) num_splits == 0 for auto calculation, result to num_splits == 1
         // we do allocation for case 2 for simplicity
-        if (num_splits == 1) {
+        if (num_splits == 1 && !loop) {
             *workspace_size = 0;
         } else {
             *workspace_size = uint64_t(total_q) * num_heads * head_size * sizeof(float);