add flash_attn_varlen_func

Author: umiswing

python/paddle/nn/functional/flash_attention.py

@@ -1049,6 +1049,52 @@ def flash_attention_v3_varlen(
     num_splits=1,
     pack_gqa=None,
     sm_margin=0,
+):
+    return flash_attn_varlen_func(
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        seqused_q,
+        seqused_k,
+        softmax_scale,
+        causal,
+        qv,
+        q_descale,
+        k_descale,
+        v_descale,
+        window_size,
+        softcap,
+        num_splits,
+        pack_gqa,
+        sm_margin,
+    )
+
+
+def flash_attn_varlen_func(
+    query,
+    key,
+    value,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    seqused_q=None,
+    seqused_k=None,
+    softmax_scale=None,
+    causal=False,
+    qv=None,
+    q_descale=None,
+    k_descale=None,
+    v_descale=None,
+    window_size=(-1, -1),
+    softcap=0.0,
+    num_splits=1,
+    pack_gqa=None,
+    sm_margin=0,
 ):
     r"""
     The equation is:
@@ -1097,24 +1143,24 @@ def flash_attention_v3_varlen(
     """
     assert (
         "xpu" not in paddle.get_device()
-    ), "flash_attention_v3_varlen is not supported on xpu"
+    ), "flash_attn_varlen_func is not supported on xpu"
 
     assert not paddle.get_flags(["FLAGS_cudnn_deterministic"])[
         "FLAGS_cudnn_deterministic"
-    ], "flash_attention_v3_varlen does not support deterministic"
+    ], "flash_attn_varlen_func does not support deterministic"
 
     assert (
         paddle.base.framework.get_flags(["FLAGS_flash_attn_version"])[
             "FLAGS_flash_attn_version"
         ]
         == 3
-    ), "FLAGS_flash_attn_version is 2, conflits with flash_attn_varlen_v3"
+    ), "FLAGS_flash_attn_version is 2, conflits with flash_attn_varlen_func"
 
     assert (
         in_dynamic_or_pir_mode()
-    ), "flash_attention_v3_varlen only support dynamic or pir mode"
+    ), "flash_attn_varlen_func only support dynamic or pir mode"
 
-    assert qv is None, "flash_attention_v3_varlen does not support setting qv"
+    assert qv is None, "flash_attn_varlen_func does not support setting qv"
 
     if softmax_scale is None:
         softmax_scale = (