[XPU] qwen2 supports flash_attn on XPU (#9549)

Author: will-jl944

paddlenlp/transformers/qwen2/modeling.py

@@ -40,6 +40,7 @@
     create_skip_config_for_refined_recompute,
     recompute,
 )
+from paddlenlp.utils.tools import get_env_device
 
 from .. import linear_utils
 from ..activations import ACT2FN
@@ -1020,7 +1021,14 @@ def _prepare_decoder_attention_mask(attention_mask, input_shape, past_key_values
                 past_key_values_length=past_key_values_length,
             )
         # Convert bool attention_mask to float attention mask, which will be added to attention_scores later
-        expanded_attn_mask = paddle.where(expanded_attn_mask.to("bool"), 0.0, paddle.finfo(dtype).min).astype(dtype)
+        if get_env_device() == "xpu":
+            x = paddle.to_tensor(0.0, dtype="float32")
+            y = paddle.to_tensor(-1.7005809656952787e38, dtype="float32")
+            expanded_attn_mask = paddle.where(expanded_attn_mask, x, y)
+        else:
+            expanded_attn_mask = paddle.where(expanded_attn_mask.to("bool"), 0.0, paddle.finfo(dtype).min).astype(
+                dtype
+            )
         return expanded_attn_mask
 
     @paddle.jit.not_to_static