[API Compatibility] Add paddle.compat.nn.functional.sdpa (PaddlePaddle#76446)

* Implement paddle.nn.functional.sdpa

* Enable flash attention test and disable test_compat_attention on Windows

* Refactor sdpa

* check dtype for mem_efficient_attention, support 3d attn_mask, refine tests

* fix test_flash_attn error

* feat: refactor GQA implementation and improve tensor handling

- Move GQA logic from compat module to main scaled_dot_product_attention function
- Add enable_gqa parameter to function signature with proper documentation
- Simplify tensor dimension handling using is_batched check
- Remove duplicate GQA validation and expansion code from compat module
- Improve code organization by centralizing GQA functionality in main implementation

* feat(test): update attention test shape for better alignment

Change the shape parameter in TestSDPAttentionWithScale from (2, 32, 8, 32) to (2, 8, 8, 32) to improve test alignment and ensure proper attention mechanism validation with more realistic tensor dimensions.

* feat(attention): update documentation and mask handling logic

- Update scaled_dot_product_attention documentation to clarify dtype support and remove GQA mode mention
- Simplify mask padding logic in MultiheadAttention to always use input dtype
- Add tensor shape comments for better code readability
- Refactor attention mask generation logic to improve efficiency
- Remove unused device capability checking functions

These changes improve code clarity and maintainability while ensuring consistent behavior across different input types.

* feat(transformer): initialize bias parameters with None and conditionally create bias parameters

Initialize all bias parameters (in_proj_bias, q_proj_bias, k_proj_bias, v_proj_bias) to None at class initialization. Conditionally create bias parameters only when bias=True, moving the bias parameter creation logic to the appropriate conditional branches. This improves code clarity by ensuring bias parameters are always defined and only created when needed.

* feat(nn): remove  __all__ from compat nn module

* feat: fix CUDA availability check in scaled dot product attention

Change `paddle.device.is_available()` to `paddle.cuda.is_available()` in the CUDA availability check function. This ensures proper detection of CUDA availability specifically for GPU operations in the scaled dot product attention implementation.

* feat: update shape output format in docstrings and rename attention module

- Change shape output format from list to paddle.Size in AvgPool1D, AvgPool2D, AvgPool3D, and Unfold docstrings
- Rename attention.py to sdpa.py and update import paths
- Remove debug parameter from check_all_tensors_on_device function
- Replace debug warning with info logging for tensor device placement checks
- Update MultiheadAttention documentation regarding optimized implementation conditions

* feat: reduce log verbosity in attention validation functions

Changed logger calls from info to debug level in SDPA validation functions to reduce noise in production logs. This maintains the same validation logic but only shows detailed validation messages when debug logging is enabled.

* feat: add bfloat16 support check for MHA tests on CUDA

Add paddle.device.is_bf16_supported() check to ensure bfloat16 tests
only run on CUDA devices that support bfloat16. This prevents test
failures on CUDA devices without bfloat16 support by falling back to
float32 dtype in those cases.

* feat: add runtime flags for attention backends and fix bf16 support check

- Add FLAGS_memory_efficient_attention_available and FLAGS_flash_attention_available
  to conditionally enable attention backends at runtime
- Update SDPA backend selection to use runtime flags instead of hardcoded values
- Fix bf16 support detection in multihead attention tests by checking CUDA compute capability
- Remove redundant scale check in flash attention constraints
- Improve test coverage by using consistent bf16 capability checks

* feat: add global flags for attention kernel availability

Add global boolean flags `memory_efficient_attention_available` and `flash_attention_available` to centralize availability checks for memory efficient and flash attention kernels. Move flag definitions from individual kernel files to flags.cc for better maintainability and to avoid code duplication. The flags automatically set to true when corresponding compilation macros (PADDLE_WITH_MEMORY_EFFICIENT_ATTENTION and PADDLE_WITH_FLASHATTN) are defined, allowing runtime detection of available attention implementations.

* Fix compile error on windows

* Fix build error

* feat(nn): use safe dict get for attention backend flags

Replace direct dictionary access with get() method to handle missing flags gracefully. This prevents KeyError exceptions when the global flags dictionary doesn't contain the expected flash attention and memory efficient attention availability flags, providing default False values instead.

Author: LittleHeroZZZX

ci/h-test.sh

@@ -161,7 +161,10 @@ concurrency_list="^test_fp8_deep_gemm$|\
 ^test_dist_fuse_gemm_epilogue_pass$|\
 ^test_fuse_allreduce_split_to_reducescatter_pass$|\
 ^test_ps_server_pass$|\
-^test_white_lists$"
+^test_white_lists$|\
+^test_scaled_dot_product_attention$|\
+^test_compat_scaled_dot_product_attention$|\
+^test_flash_attention$"
 
 cd ${work_dir}/build
 tmp_dir=`mktemp -d`

paddle/phi/core/flags.cc

@@ -32,3 +32,25 @@ PHI_DEFINE_EXPORTED_int64(conv_workspace_size_limit,
                           phi::backends::gpu::kDefaultConvWorkspaceSizeLimitMB,
                           "cuDNN convolution workspace limit in MB unit.");
 #endif
+
+#ifdef PADDLE_WITH_MEMORY_EFFICIENT_ATTENTION
+static const bool kMemEffAttnDefault = true;
+#else
+static const bool kMemEffAttnDefault = false;
+#endif
+
+PHI_DEFINE_EXPORTED_bool(
+    mem_efficient_attn_available,
+    kMemEffAttnDefault,
+    "Weather memory efficient attention is available on the current device.");
+
+#ifdef PADDLE_WITH_FLASHATTN
+static const bool kFlashAttnDefault = true;
+#else
+static const bool kFlashAttnDefault = false;
+#endif
+
+PHI_DEFINE_EXPORTED_bool(
+    flash_attn_available,
+    kFlashAttnDefault,
+    "Weather flash attention is available on the current device.");

python/paddle/compat/nn/__init__.py

@@ -400,9 +400,6 @@ def __setstate__(self, state):
         self.__dict__.setdefault("count_include_pad", True)
 
 
-__all__ = ['Unfold', 'Linear', 'MultiheadAttention']
-
-
 class Unfold(nn.Unfold):
     """
     A compatible version of paddle.nn.Unfold:

python/paddle/compat/nn/functional/__init__.py

@@ -25,6 +25,8 @@
 from paddle.tensor import softmax
 from paddle.utils.decorator_utils import ForbidKeywordsDecorator
 
+from .sdpa import scaled_dot_product_attention
+
 if TYPE_CHECKING:
     from typing_extensions import TypeAlias
 
@@ -39,7 +41,7 @@
     ]
 
 
-__all__ = ['pad', 'softmax', 'linear', 'unfold']
+__all__ = ['pad', 'softmax', 'linear', 'scaled_dot_product_attention', 'unfold']
 
 
 def _check_valid_pad_len(pad_len, x_dim, is_constant):

python/paddle/compat/nn/functional/sdpa.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle.nn.functional as F
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def scaled_dot_product_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attn_mask: Tensor | None = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float | None = None,
+    enable_gqa: bool = False,
+) -> Tensor:
+    r"""
+    The equation is:
+
+    .. math::
+
+        result=softmax(\frac{ Q * K^T }{\sqrt{d}}) * V
+
+    where : ``Q``, ``K``, and ``V`` represent the three input parameters of the attention module.
+    The dimensions of the three parameters are the same.
+    ``d`` represents the size of the last dimension of the three parameters.
+
+
+    Warning:
+        This API only verifies inputs with dtype float16 and bfloat16, other dtypes may fall back to math
+        implementation, which is less optimized.
+
+    Note:
+        This API differs from :ref:`api_paddle_nn_functional_scaled_dot_product_attention` in that:
+        The QKV layout of this API is [batch_size, num_heads, seq_len, head_dim] or [num_heads, seq_len, head_dim].
+
+    Args:
+        query(Tensor): The query tensor in the Attention module.
+                        4-D tensor with shape:
+                        [batch_size, num_heads, seq_len, head_dim].
+                        3-D tensor with shape:
+                        [num_heads, seq_len, 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, num_heads, seq_len, head_dim].
+                        3-D tensor with shape:
+                        [num_heads, seq_len, 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, num_heads, seq_len, head_dim].
+                        3-D tensor with shape:
+                        [num_heads, seq_len, head_dim].
+                        The dtype can be float16 or bfloat16.
+        attn_mask(Tensor, optional): The attention mask tensor. The shape should be broadcastable to
+                        [batch_size, num_heads, seq_len_key, seq_len_query]. The dtype can be bool
+                        or same type of query. The bool mask indicates the positions should take part
+                        in attention. The non-bool mask will be added to attention score.
+
+        is_causal(bool, optional): Whether enable causal mode. If True, the attention masking is a lower
+                        triangular matrix when the mask is a square matrix. The attention masking has the
+                        form of the upper left causal bias when the mask is a non-square matrix.
+                        An error is thrown if both attn_mask and is_causal are set.
+        scale(float, optional): The scaling factor used in the calculation of attention weights.
+                        If None, scale = 1 / sqrt(head_dim).
+        enable_gqa(bool, optional): Whether enable GQA mode. Default False.
+
+    Returns:
+        out(Tensor): The attention tensor.
+                    4-D tensor with shape: [batch_size, num_heads, seq_len, head_dim].
+                    3-D tensor with shape: [num_heads, seq_len, head_dim].
+                    The dtype can be float16 or bfloat16.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +SKIP('bfloat need V100 compile')
+            >>> import paddle
+            >>> q = paddle.rand((1, 2, 128, 16), dtype=paddle.bfloat16)
+            >>> output = paddle.compat.nn.functional.scaled_dot_product_attention(q, q, q, None, 0.9, False)
+            >>> print(output)
+            >>> # doctest: -SKIP
+    """
+    if is_causal and attn_mask is not None:
+        raise RuntimeError(
+            "Explicit attn_mask should not be set when is_causal=True"
+        )
+
+    query, key, value = (
+        query.swapaxes(-3, -2),
+        key.swapaxes(-3, -2),
+        value.swapaxes(-3, -2),
+    )
+    out = F.scaled_dot_product_attention(
+        query,
+        key,
+        value,
+        attn_mask,
+        dropout_p,
+        is_causal,
+        True,  # training
+        None,  # backend
+        scale,
+        enable_gqa,
+        None,  # name
+    )
+    return out.swapaxes(-3, -2)