Revert flashmask bidirectional attention optimization (#67811)

GuoxiaWang · web-flow · commit 596e50be7666 · 2024-08-30T14:44:16.000+08:00
* Revert flashmask bidirectional attention optimization

* update

* fix unitest
diff --git a/python/paddle/nn/functional/flash_attention.py b/python/paddle/nn/functional/flash_attention.py
@@ -1475,6 +1475,9 @@ def flashmask_attention(
                 has_end = False
             elif startend_row_indices.shape[-1] == 4:
                 has_end = True
+                raise NotImplementedError(
+                    "ending row index is not implemented yet."
+                )
             else:
                 raise ValueError(
                     f"Invalid shape of startend_row_indices, when causal is False, the last dimension should be either 2 or 4 but got {startend_row_indices.shape[-1]}"
diff --git a/test/legacy_test/test_flash_attention.py b/test/legacy_test/test_flash_attention.py
@@ -477,7 +477,7 @@ def setUp(self):
         self.enable_mem_efficient = False
 
 
-class TestFlashAttenionWithMaskAPITest(TestFlashAttentionWithMaskAPI):
+class TestFlashAttentionWithMaskAPITest(TestFlashAttentionWithMaskAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (8, 1024, 16, 128)
@@ -869,7 +869,6 @@ def generate_mask_matrix_from_mask_indices(start_rows):
             for j in range(seq_len):
                 start_row = start_rows[bz_idx, head_idx, j]
                 matrix[bz_idx, head_idx, start_row:, j] = -np.inf
-                matrix[bz_idx, head_idx, j, j] = 0.0
     return matrix
 
 
@@ -945,7 +944,7 @@ def test_dot_scale_product(self):
         np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03)
 
 
-class TestFlashAttenionWithSparseMaskAPITest(
+class TestFlashAttentionWithSparseMaskAPITest(
     TestFlashAttentionWithSparseMaskAPI
 ):
     def setUp(self):
@@ -956,7 +955,7 @@ def setUp(self):
         self.causal = True
 
 
-class TestFlashAttenionWithSparseMaskBF16APITest(
+class TestFlashAttentionWithSparseMaskBF16APITest(
     TestFlashAttentionWithSparseMaskAPI
 ):
     def setUp(self):
diff --git a/third_party/flashattn b/third_party/flashattn
@@ -1 +1 @@
-Subproject commit 666d5c56bba917f9f5197cb4a2fdee123f0f6c38
+Subproject commit 21f25e86ba0765d574959b39cb370e3a43798972
