Add flex_attn to diffllama (#35601)

Add sdpa to diffllama

Author: muellerzr

src/transformers/models/diffllama/modeling_diffllama.py

@@ -595,6 +595,7 @@ class DiffLlamaPreTrainedModel(PreTrainedModel):
     _skip_keys_device_placement = ["past_key_values"]
     _supports_flash_attn_2 = True
     _supports_sdpa = True
+    _supports_flex_attn = True
     _supports_cache_class = True
     _supports_quantized_cache = True
     _supports_static_cache = True