scaled_dot_product_attention函数兼容3D输入 (#73336)

Qin-sx · web-flow · commit c83a7c12ca86 · 2025-06-17T12:10:05.000+08:00
* compatible with 3D input

	modified:   python/paddle/nn/functional/flash_attention.py
	modified:   test/legacy_test/test_scaled_dot_product_attention.py

* updated the test

	modified:   test/legacy_test/test_scaled_dot_product_attention.py

* updated the test

	modified:   test/legacy_test/test_scaled_dot_product_attention.py
diff --git a/python/paddle/nn/functional/flash_attention.py b/python/paddle/nn/functional/flash_attention.py
@@ -1380,14 +1380,20 @@ def scaled_dot_product_attention(
         query(Tensor): The query tensor in the Attention module.
                         4-D tensor with shape:
                         [batch_size, seq_len, num_heads, head_dim].
+                        3-D tensor with shape:
+                        [seq_len, num_heads, head_dim].
                         The dtype can be float16 or bfloat16.
         key(Tensor): The key tensor in the Attention module.
                         4-D tensor with shape:
                         [batch_size, seq_len, num_heads, head_dim].
+                        3-D tensor with shape:
+                        [seq_len, num_heads, head_dim].
                         The dtype can be float16 or bfloat16.
         value(Tensor): The value tensor in the Attention module.
                         4-D tensor with shape:
                         [batch_size, seq_len, num_heads, head_dim].
+                        3-D tensor with shape:
+                        [seq_len, num_heads, head_dim].
                         The dtype can be float16 or bfloat16.
         attn_mask(Tensor, optional): A float mask of the same type as query,
                         key, value that is added to the attention score.
@@ -1413,6 +1419,14 @@ def scaled_dot_product_attention(
             >>> print(output)
             >>> # doctest: -SKIP
     """
+    if query.ndim == 3:
+        query = paddle.unsqueeze(query, axis=0)
+
+    if key.ndim == 3:
+        key = paddle.unsqueeze(key, axis=0)
+
+    if value.ndim == 3:
+        value = paddle.unsqueeze(value, axis=0)
 
     if attn_mask is None:
         # downgraded to ordinary flash attention implementation
diff --git a/test/legacy_test/test_scaled_dot_product_attention.py b/test/legacy_test/test_scaled_dot_product_attention.py
@@ -181,5 +181,42 @@ def test_dot_scale_product_float_mask(self):
         np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03)
 
 
+class TestAttentionWith3DInput(unittest.TestCase):
+    def setUp(self):
+        self.dtype = 'float32'
+        self.dropout = 0.0
+        self.causal = False
+
+    def test_3d_input(self):
+        """Test scaled_dot_product_attention with 3D input tensors."""
+        # test dynamic
+        paddle.disable_static()
+
+        shape_3d = (8, 1, 8)
+
+        query = np.random.random(shape_3d).astype(np.float32)
+        key = np.random.random(shape_3d).astype(np.float32)
+        value = np.random.random(shape_3d).astype(np.float32)
+
+        q = paddle.to_tensor(query, dtype=self.dtype, stop_gradient=False)
+        k = paddle.to_tensor(key, dtype=self.dtype, stop_gradient=False)
+        v = paddle.to_tensor(value, dtype=self.dtype, stop_gradient=False)
+
+        q_ref = paddle.unsqueeze(q, axis=0)
+        k_ref = paddle.unsqueeze(k, axis=0)
+        v_ref = paddle.unsqueeze(v, axis=0)
+
+        with sdp_kernel(
+            enable_math=True, enable_flash=False, enable_mem_efficient=False
+        ):
+            out = scaled_dot_product_attention(
+                q, k, v, None, self.dropout, self.causal
+            )
+
+        out_ref = attention_naive(q_ref, k_ref, v_ref, self.causal)
+
+        np.testing.assert_allclose(out.numpy(), out_ref, rtol=5e-03, atol=1e-03)
+
+
 if __name__ == '__main__':
     unittest.main()
