Revert "optimize ds3 attention impl (#10200)" (#10208)

This reverts commit 58edb00.

Author: phlrain

paddlenlp/transformers/deepseek_v2/modeling.py

@@ -245,6 +245,11 @@ def scaled_dot_product_attention(
         q_head_dim = query_states.shape[-1]
         softmax_scale = softmax_scale * (q_head_dim**0.5)
         query_states = query_states * softmax_scale
+        value_padding = paddle.zeros(
+            [bsz, kv_seq_len, v_num_heads, head_dim - v_head_dim],
+            dtype=value_states.dtype,
+        )
+        value_states = paddle.concat([value_states, value_padding], axis=-1)
 
         outputs = fusion_ops.fusion_flash_attention(
             query_states,
@@ -257,6 +262,15 @@ def scaled_dot_product_attention(
             sequence_parallel=sequence_parallel,
         )
 
+        if isinstance(outputs, tuple):
+            outputs[0] = outputs[0].reshape([bsz, q_len, v_num_heads, head_dim])
+            outputs[0] = outputs[0][..., :v_head_dim]
+            outputs[0] = outputs[0].reshape([bsz, q_len, -1])
+        else:
+            outputs = outputs.reshape([bsz, q_len, v_num_heads, head_dim])
+            outputs = outputs[..., :v_head_dim]
+            outputs = outputs.reshape([bsz, q_len, -1])
+
         if sequence_parallel:
             outputs = outputs.reshape([bsz * q_len, v_head_dim * num_heads])