Merge remote-tracking branch 'datta0/flexattn_refactor' into flexattn_refactor

unsloth/models/_utils.py

@@ -229,7 +229,9 @@ def determine_attention_implementation(model_class, config):
     if HAS_FLASH_ATTENTION and model_type not in ("gpt_oss", "mllama") and not model_type.startswith("gemma3n"):
         supports_fa2 = False
         if model_class is not None:
-            supports_fa2 = getattr(model_class, "_supports_flash_attn_2", False) or getattr(model_class, "_supports_flash_attn", False)
+            supports_fa2 = getattr(
+                model_class, "_supports_flash_attn_2", False
+            ) or getattr(model_class, "_supports_flash_attn", False)
 
         if supports_fa2:
             if config is not None:
@@ -243,8 +245,10 @@ def determine_attention_implementation(model_class, config):
         try:
             from transformers.utils.import_utils import is_torch_flex_attn_available
 
-            if is_torch_flex_attn_available() and (model_class is not None) and getattr(
-                model_class, "_supports_flex_attn", False
+            if (
+                is_torch_flex_attn_available()
+                and (model_class is not None)
+                and getattr(model_class, "_supports_flex_attn", False)
             ):
                 # GPT-OSS, Mllama and Gemma3 use eager/sdpa attention during
                 # inference since flex attention returns incorrect results or errors out.

unsloth/models/llama.py

@@ -2341,7 +2341,9 @@ 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 = determine_attention_implementation(model_function, model_config)
+        preferred_attn_impl = determine_attention_implementation(
+            model_function, model_config
+        )
 
         has_rope_scaling = False
         try:

unsloth/models/vision.py

@@ -608,7 +608,7 @@ class FastBaseModel:
             model_class = auto_model._model_mapping[auto_config.__class__]
         except Exception:
             model_class = None
-            
+
         attn_impl = determine_attention_implementation(model_class, auto_config)
 
         # Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with