Add flex_attn to diffllama (#35601)

Add sdpa to diffllama
huggingface · Jan 9, 2025 · 8de7b1b · 8de7b1b
1 parent 1e3ddcb
commit 8de7b1b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/models/diffllama/modeling_diffllama.py b/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