Unify fused_dot_product_attention with scaled_dot_product_attention (#65507)

deepllz · Wong4j · web-flow · commit 47461b5984fc · 2024-07-08T17:01:51.000+08:00
* update cudnn flash attention

fix error

fix pir pass

fix ci

add more ut

minor changes

remove workspace_opt env

minor change

update api doc

* polish

* fix static test

* revert graph_pattern_detector.cc modify

---------

Co-authored-by: Shijie Wang &lt;jaywan@nvidia.com&gt;
diff --git a/python/paddle/incubate/nn/functional/fused_dot_product_attention.py b/python/paddle/incubate/nn/functional/fused_dot_product_attention.py
@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-from paddle import _C_ops
+from paddle import Tensor, _C_ops
 from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
 
 
@@ -57,10 +56,6 @@ def cudnn_flash_attention(
         This API is provides the full functionality of the cuDNN flash attention. Such as dbias, padding_causal mask, etc.
     """
 
-    batch_size = q.shape[0]
-    q_seqlen = q.shape[1]
-    k_seqlen = k.shape[1]
-
     if mask_type is None:
         mask_type = "none"
     if bias_type is None:
@@ -78,7 +73,7 @@ def cudnn_flash_attention(
     ], "bias_type should be 'none', 'pre_scale_bias', 'post_scale_bias'"
 
     if in_dynamic_or_pir_mode():
-        out, softmax, _ = _C_ops.fused_dot_product_attention(
+        out, _, _ = _C_ops.fused_dot_product_attention(
             q,
             k,
             v,
@@ -130,31 +125,42 @@ def cudnn_flash_attention(
 
 
 def fused_dot_product_attention(
-    q,
-    k,
-    v,
-    scaling_factor,
-    mask=None,
-    dropout_prob=0.0,
-    is_causal=False,
-    training=True,
-    name=None,
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attn_mask: Tensor = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scaling_factor: float = None,
+    training: bool = True,
+    name: str = None,
 ):
     r"""
     Fused Dot Product Attention. This is a fusion operator to compute scaled dot product attention in transformer
     model architecture. This operator only supports running on Ampere and Hopper GPU and need cudnn version >= 8906.
 
     Args:
-        q (Tensor): The query tensor. The data type is bfloat16, float16.
-        k (Tensor): The key tensor. The data type is bfloat16, float16.
-        v (Tensor): The value tensor. The data type is bfloat16, float16.
-        scaling_factor (float): The scaling factor for the attention scores.
-        mask(Tensor,optional): A float mask of the same type as query,
+        query(Tensor): The query tensor in the Attention module.
+                        4-D tensor with shape:
+                        [batch_size, 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].
+                        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].
+                        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.
-        dropout_prob (float): The dropout probability.
+        dropout_p (float): The dropout probability.
         is_causal (bool): A flag indicating whether it is causal masking or not. If True, the mask will be ignored.
+        scaling_factor (float): The scaling factor for the attention scores.
         training (bool): A flag indicating whether it is in train phrase or not.
-
+        name(str|None, optional): The default value is None. Normally there is no need for user
+                        to set this property. For more information, please refer to
+                        :ref:`api_guide_Name`.
 
     Returns:
         A Tensor representing the fused dot product attention, has same shape and data type as `q` .
@@ -167,64 +173,64 @@ def fused_dot_product_attention(
             by cuDNN, which achieves better performance on the latest GPU architectures (Hopper and After).
 
         The mask is passed as a post scale bias in this API. So this mask can be a arbitrary mask, not just padding mask.
-
-
     """
 
-    cu_seqlen_q = None
-    cu_seqlen_k = None
-
     if is_causal:
-        assert mask is None, "mask must be None when is_causal is True"
+        assert attn_mask is None, "mask must be None when is_causal is True"
         mask_type = "causal"
         bias_type = "none"
-    elif mask is not None:
+    elif attn_mask is not None:
         mask_type = "none"
         bias_type = "post_scale_bias"  # pass mask as a post scale bias
     else:
         mask_type = "none"
         bias_type = "none"
 
-    if mask is not None:
+    if attn_mask is not None:
         assert (
-            mask.dtype == q.dtype
-        ), "mask dtype should be the same as qkv dtype"
+            attn_mask.dtype == query.dtype
+        ), "attn_mask dtype should be the same as qkv dtype"
 
+    cu_seqlen_q = None
+    cu_seqlen_k = None
+    if scaling_factor is None:
+        head_dim = query.shape[3]
+        scaling_factor = head_dim**-0.5
     if in_dynamic_or_pir_mode():
-        out, softmax, _ = _C_ops.fused_dot_product_attention(
-            q,
-            k,
-            v,
-            mask,
+        out, _, _ = _C_ops.fused_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask,
             cu_seqlen_q,
             cu_seqlen_k,
             scaling_factor,
-            dropout_prob,
+            dropout_p,
             training,
             mask_type,
             bias_type,
         )
         return out
     else:
         helper = LayerHelper('fused_dot_product_attention', **locals())
-        out = helper.create_variable_for_type_inference(dtype=q.dtype)
+        out = helper.create_variable_for_type_inference(dtype=query.dtype)
         softmax_out = helper.create_variable_for_type_inference(dtype="float")
         rng_state = helper.create_variable_for_type_inference(dtype='int64')
 
         attrs = {
             "scaling_factor": scaling_factor,
-            "dropout_probability": dropout_prob,
+            "dropout_probability": dropout_p,
             "is_training": training,
             "mask_type_str": mask_type,
             "bias_type_str": bias_type,
         }
         helper.append_op(
             type='fused_dot_product_attention',
             inputs={
-                'q': q,
-                'k': k,
-                'v': v,
-                'bias': mask,
+                'q': query,
+                'k': key,
+                'v': value,
+                'bias': attn_mask,
                 'cu_seqlen_q': cu_seqlen_q,
                 'cu_seqlen_k': cu_seqlen_k,
             },
diff --git a/python/paddle/nn/functional/flash_attention.py b/python/paddle/nn/functional/flash_attention.py
@@ -1016,15 +1016,15 @@ 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].
-                        The dtype can be float61 or bfloat16.
+                        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].
-                        The dtype can be float61 or bfloat16.
+                        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].
-                        The dtype can be float61 or bfloat16.
+                        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.
         dropout_p(float, optional): The dropout ratio.
diff --git a/test/legacy_test/test_fused_dot_product_attention_op.py b/test/legacy_test/test_fused_dot_product_attention_op.py
@@ -197,11 +197,10 @@ def _get_fused_attn_out(self):
             q_tensor,
             k_tensor,
             v_tensor,
-            self.scaling_factor,
             attn_mask,
             self.dropout_prob,
             self.is_causal_masking,
-            True,
+            training=True,
         )
 
         paddle.autograd.backward(
diff --git a/test/legacy_test/test_fused_dot_product_attention_op_static.py b/test/legacy_test/test_fused_dot_product_attention_op_static.py
@@ -71,7 +71,20 @@ def test_static_op(self):
         mask = paddle.to_tensor(
             mask_data, place=self.place, dtype=self.dtype, stop_gradient=True
         )
-        out = fused_dot_product_attention(q, k, v, 1.0, mask)
+        out0 = fused_dot_product_attention(q, k, v, mask)
+        out1 = cudnn_flash_attention(
+            q,
+            k,
+            v,
+            mask,
+            None,
+            None,
+            1.0,
+            0.0,
+            True,
+            None,
+            "post_scale_bias",
+        )
         paddle.enable_static()
         paddle.seed(312)
 
@@ -90,7 +103,7 @@ def test_static_op(self):
                 name="mask", shape=self.mask_shape, dtype=self.dtype
             )
 
-            outs = fused_dot_product_attention(q, k, v, 1.0, mask)
+            outs = fused_dot_product_attention(q, k, v, mask)
 
             exe = paddle.static.Executor(self.place)
             out_s = exe.run(
@@ -102,7 +115,7 @@ def test_static_op(self):
                 },
                 fetch_list=[outs],
             )
-            np.testing.assert_allclose(out_s[0], out)
+            np.testing.assert_allclose(out_s[0], out0)
 
         # call cudnn_flash_attention in static mode
         with paddle.static.program_guard(paddle.static.Program()):
@@ -143,6 +156,7 @@ def test_static_op(self):
                 },
                 fetch_list=[outs],
             )
+            np.testing.assert_allclose(out_s[0], out1)
 
 
 if __name__ == "__main__":
