update flash attn select (#54630) (#54716)

feff99f5 · Yuang Liu · GitHub · 570daa19 · feff99f5
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

python/paddle/nn/functional/flash_attention.py python/paddle/nn/functional/flash_attention.py +1 -1

未找到文件。
--- a/python/paddle/nn/functional/flash_attention.py
+++ b/python/paddle/nn/functional/flash_attention.py
@@ -81,7 +81,7 @@ def _math_attention(


 def _select_sdp_cuda(head_dim):
-    if head_dim < 128:
+    if head_dim <= 128:
        return "flash_attn"
    else:
        return "mem_efficient"