update

Author: zhangyuqin1998

paddlenlp/transformers/ring_flash_attention.py

@@ -55,8 +55,14 @@ def wait(self):
 
     def add_to_buffers(self, key, value):
         if key.shape != self._k_buffer[self._next_buffer_idx].shape:
-            self._k_buffer[self._next_buffer_idx][:, : key.shape[1], :, :] += key
-            self._v_buffer[self._next_buffer_idx][:, : key.shape[1], :, :] += value
+            k_buffer_chunk = paddle.slice(
+                self._k_buffer[self._next_buffer_idx], axes=[1], starts=[0], ends=[key.shape[1]]
+            )
+            v_buffer_chunk = paddle.slice(
+                self._v_buffer[self._next_buffer_idx], axes=[1], starts=[0], ends=[value.shape[1]]
+            )
+            k_buffer_chunk += key
+            v_buffer_chunk += value
         else:
             self._k_buffer[self._next_buffer_idx] += key
             self._v_buffer[self._next_buffer_idx] += value
@@ -82,13 +88,13 @@ def update_out_and_lse(old_out, old_lse, block_out, block_lse, second_chunk_only
         return block_out.to("float32"), block_lse.to("float32")
 
     if second_chunk_only:
-        second_chunk_out = old_out[:, old_out.shape[1] // 2 :, :, :]
-        second_chunk_lse = old_lse[:, old_lse.shape[1] // 2 :, :, :]
+        second_chunk_out_ = paddle.slice(old_out, axes=[1], starts=[old_out.shape[1] // 2], ends=[old_out.shape[1]])
+        second_chunk_lse_ = paddle.slice(old_lse, axes=[1], starts=[old_lse.shape[1] // 2], ends=[old_lse.shape[1]])
         second_chunk_out, second_chunk_lse = update_out_and_lse(
-            second_chunk_out, second_chunk_lse, block_out, block_lse
+            second_chunk_out_, second_chunk_lse_, block_out, block_lse
         )
-        old_out[:, old_out.shape[1] // 2 :, :, :] = second_chunk_out
-        old_lse[:, old_lse.shape[1] // 2 :, :, :] = second_chunk_lse
+        paddle.assign(second_chunk_out, second_chunk_out_)
+        paddle.assign(second_chunk_lse, second_chunk_lse_)
         return old_out, old_lse
     else:
         block_out, block_lse = block_out.to("float32"), block_lse.to("float32")
@@ -242,6 +248,9 @@ def balanced_ring_flash_attention_bwd_func(
         )
         lse_second_chunk = paddle.slice(lse, axes=[2], starts=[local_q_seq_len // 2], ends=[local_q_seq_len])
         out_grad_second_chunk = paddle.slice(out_grad, axes=[1], starts=[local_q_seq_len // 2], ends=[local_q_seq_len])
+        query_grad_buffer_second_chunk = paddle.slice(
+            query_grad_buffer, axes=[1], starts=[local_q_seq_len // 2], ends=[local_q_seq_len]
+        )
 
     if attn_mask is not None:
         attn_masks_list = paddle.split(attn_mask, num_or_sections=cp_size * 2, axis=3)
@@ -285,7 +294,7 @@ def balanced_ring_flash_attention_bwd_func(
                     dropout,
                     False,
                 )
-                query_grad_buffer[:, local_q_seq_len // 2 :, :, :] += block_q_grad
+                query_grad_buffer_second_chunk += block_q_grad
             else:
                 block_q_grad, block_k_grad, block_v_grad = flash_attn_bwd(
                     local_query,