From 256c6e4884ae7385702560771efae3232fde17d9 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 1 Apr 2026 13:00:21 +0530 Subject: [PATCH] Refactor flex attn to prefer flash if possible (#4734) Replaces prefer_flex_attn_if_supported (which only returned flex_attention or None) with determine_attention_implementation, a centralized hierarchy: FA2 > Flex > SDPA > Eager. Changes: - New determine_attention_implementation function in _utils.py with clear priority chain - _set_attn_impl helper to stamp config consistently - _FLEX_EXCLUDED_MODELS / _FLEX_EXCLUDED_PREFIXES for model-specific exclusions - Gemma3N explicit eager override in vision.py (timm vision towers) - Preserved sdpa fallback for unmapped/remote-code vision configs - Config re-stamped to eager when supports_sdpa guard fires Co-authored-by: Datta Nimmaturi --- unsloth/models/_utils.py | 105 +++++++++++++++++++++++++-------------- unsloth/models/llama.py | 4 +- unsloth/models/vision.py | 29 ++++++----- 3 files changed, 84 insertions(+), 54 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index d296ac7e74..49168eb8b6 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -64,7 +64,8 @@ __all__ = [ "patch_compiled_autograd", "process_vision_info", "unsloth_compile_transformers", - "prefer_flex_attn_if_supported", + "determine_attention_implementation", + "_set_attn_impl", "patch_fast_lora", "validate_loftq_config", "RaiseUninitialized", @@ -222,44 +223,74 @@ 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 +# Models that don't work with flex_attention: +# GPT-OSS: left padding issues cause incorrect outputs. +# Mllama: BlockMask Q_LEN!=KV_LEN ValueError on decode. +# NemotronH: hybrid Mamba-2 + Transformer, raises NotImplementedError. +# Gemma3N: timm vision wrappers don't support flex_attention. +_FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h") +_EAGER_ONLY_PREFIXES = ("gemma3n",) - 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 - attention_dropout = getattr(config, "attention_dropout", 0) or 0 - if attention_dropout > 0: - return None - # GPT-OSS, Mllama and Gemma3N use eager/sdpa attention during - # inference since flex attention returns incorrect results or errors out. - # GPT-OSS: left padding issues cause incorrect outputs. - # Mllama: _update_causal_mask uses make_flex_block_causal_mask which - # creates BlockMask with Q_LEN=KV_LEN=total_seq_len, but during - # decode q_len=1, causing ValueError. Needs transformers update. - # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not - # support flex_attention. - # NemotronH: hybrid Mamba-2 + Transformer model that does not - # support flex_attention (raises NotImplementedError from transformers). - model_type = getattr(config, "model_type", "") if config else "" - if model_type in ("gpt_oss", "mllama", "nemotron_h") or str( - model_type - ).startswith("gemma3n"): - 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 +def _is_flex_excluded(model_type): + return model_type in _FLEX_EXCLUDED_MODELS + + +def _is_eager_only(model_type): + return any(model_type.startswith(p) for p in _EAGER_ONLY_PREFIXES) + + +def _set_attn_impl(config, impl): + """Helper function to set attention implementation on config and return it.""" + if config is not None: + setattr(config, "_attn_implementation", impl) + if hasattr(config, "attn_implementation"): + setattr(config, "attn_implementation", impl) + return impl + + +def determine_attention_implementation(model_class, config): + model_type = getattr(config, "model_type", "").lower() + + # Eager-only models (e.g. gemma3n timm vision towers) + if _is_eager_only(model_type): + _set_attn_impl(config, "eager") + return "eager" + + # Flash Attention 2 + if HAS_FLASH_ATTENTION and model_class is not None: + supports_fa2 = getattr(model_class, "_supports_flash_attn_2", False) or getattr( + model_class, "_supports_flash_attn", False + ) + if supports_fa2: + _set_attn_impl(config, "flash_attention_2") + return "flash_attention_2" + + # Flex Attention + if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") != "0": + 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) + and not _is_flex_excluded(model_type) + ): + attention_dropout = getattr(config, "attention_dropout", 0) or 0 + if attention_dropout == 0: + _set_attn_impl(config, "flex_attention") + return "flex_attention" + except Exception: + pass + + # SDPA + if model_class is not None and getattr(model_class, "_supports_sdpa", False): + _set_attn_impl(config, "sdpa") + return "sdpa" + + _set_attn_impl(config, "eager") + return "eager" def _run_temporary_patches(phase): diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 93d93e26d6..8e3892c712 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2341,8 +2341,8 @@ 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" + preferred_attn_impl = determine_attention_implementation( + model_function, model_config ) has_rope_scaling = False diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index f558aa3f00..a961225827 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -597,8 +597,6 @@ class FastBaseModel: custom_datatype = None 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, @@ -609,7 +607,14 @@ class FastBaseModel: 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) + if model_class is None: + # When model_class cannot be resolved (remote-code or unmapped + # configs), preserve the old fallback of sdpa when supported. + attn_impl = _set_attn_impl( + auto_config, "sdpa" if supports_sdpa else "eager" + ) + else: + 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 # FP8 weights. We just need to update it here for sanity. @@ -620,21 +625,15 @@ class FastBaseModel: except Exception: model_class = None - model_type = str(getattr(auto_config, "model_type", "")).lower() - if model_type.startswith("gemma3n"): - # Gemma3N variants initialize timm-based vision towers which do - # not support flex_attention, so default to eager unless overridden. - default_attn_impl = "eager" - else: - default_attn_impl = "flex_attention" if flex_attn_impl else "sdpa" if not ("attn_implementation" in kwargs): - kwargs["attn_implementation"] = default_attn_impl + kwargs["attn_implementation"] = 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." - ) + print( + f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." + ) del kwargs["attn_implementation"] + # Re-stamp config so it stays consistent with the actual impl + _set_attn_impl(auto_config, "eager") bnb_config = None user_quantization_config = kwargs.get("quantization_config", None)