Prefer flex attention when available (#3979)

* Enable flex attention by default

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Avoid dropping flex attention when SDPA unsupported

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

unsloth/models/_utils.py

@@ -63,6 +63,7 @@ __all__ = [
     "patch_compiled_autograd",
     "process_vision_info",
     "unsloth_compile_transformers",
+    "prefer_flex_attn_if_supported",
     "patch_fast_lora",
     "validate_loftq_config",
     "RaiseUninitialized",
@@ -184,6 +185,27 @@ def apply_unsloth_gradient_checkpointing(
     return use_gradient_checkpointing
 
 
+def prefer_flex_attn_if_supported(model_class, config):
+    if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0":
+        return None
+    try:
+        from transformers.utils.import_utils import is_torch_flex_attn_available
+
+        if not is_torch_flex_attn_available():
+            return None
+        if model_class is None or not getattr(
+            model_class, "_supports_flex_attn", False
+        ):
+            return None
+        if config is not None:
+            setattr(config, "_attn_implementation", "flex_attention")
+            if hasattr(config, "attn_implementation"):
+                setattr(config, "attn_implementation", "flex_attention")
+        return "flex_attention"
+    except Exception:
+        return None
+
+
 for temporary_patch in TEMPORARY_PATCHES:
     temporary_patch()
 

unsloth/models/llama.py

@@ -2288,6 +2288,10 @@ class FastLlamaModel:
         model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__]
         IS_FALCON_H1 = model_config.model_type.startswith("falcon_h1")
 
+        preferred_attn_impl = (
+            prefer_flex_attn_if_supported(model_function, model_config) or "eager"
+        )
+
         has_rope_scaling = False
         try:
             with open(inspect.getfile(model_function), "r", encoding = "utf-8") as file:
@@ -2366,7 +2370,7 @@ class FastLlamaModel:
                 token = token,
                 max_position_embeddings = max_position_embeddings,
                 trust_remote_code = trust_remote_code,
-                attn_implementation = "eager",
+                attn_implementation = preferred_attn_impl,
                 **kwargs,
             )
         elif not fast_inference:
@@ -2378,7 +2382,7 @@ class FastLlamaModel:
                 token = token,
                 max_position_embeddings = max_position_embeddings,
                 trust_remote_code = trust_remote_code,
-                attn_implementation = "eager",
+                attn_implementation = preferred_attn_impl,
                 **kwargs,
             )
             model.fast_generate = make_fast_generate_wrapper(model.generate)

unsloth/models/vision.py

@@ -517,9 +517,23 @@ class FastBaseModel:
                 correct_dtype = None
 
         # Stop SDPA for some archs like Pixtral / Mistral3
+        flex_attn_impl = None
+        if auto_config is None:
+            auto_config = AutoConfig.from_pretrained(
+                model_name,
+                token = token,
+                trust_remote_code = trust_remote_code,
+            )
+        try:
+            model_class = auto_model._model_mapping[auto_config.__class__]
+        except Exception:
+            model_class = None
+        flex_attn_impl = prefer_flex_attn_if_supported(model_class, auto_config)
+
+        default_attn_impl = "flex_attention" if flex_attn_impl else "sdpa"
         if not ("attn_implementation" in kwargs):
-            kwargs["attn_implementation"] = "sdpa"
-        if not supports_sdpa:
+            kwargs["attn_implementation"] = default_attn_impl
+        if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa":
             if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") == "0":
                 print(
                     f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager."
@@ -651,12 +665,19 @@ class FastBaseModel:
 
         kwargs = add_dtype_kwargs(torch_dtype, kwargs)
 
-        model_config = AutoConfig.from_pretrained(
-            model_name,
-            token = token,
-            attn_implementation = "sdpa" if supports_sdpa else "eager",
-            trust_remote_code = trust_remote_code,
-        )
+        config_attn_impl = kwargs.get("attn_implementation", None)
+        if config_attn_impl is None:
+            config_attn_impl = "sdpa" if supports_sdpa else "eager"
+        if auto_config is None:
+            auto_config = AutoConfig.from_pretrained(
+                model_name,
+                token = token,
+                trust_remote_code = trust_remote_code,
+            )
+        setattr(auto_config, "_attn_implementation", config_attn_impl)
+        if hasattr(auto_config, "attn_implementation"):
+            setattr(auto_config, "attn_implementation", config_attn_impl)
+        model_config = auto_config
         verify_fp8_support_if_applicable(model_config)
 
         raise_handler = RaiseUninitialized()