ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor flex attn to prefer flash if possible (#4734)

Browse source 
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 <Datta0@users.noreply.github.com>
This commit is contained in:
Datta Nimmaturi 2026-04-01 13:00:21 +05:30 • committed by GitHub
commit 256c6e4884
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 84 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

105
unsloth/models/_utils.py
View file

@ -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):

4
unsloth/models/llama.py
View file

@ -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

29
unsloth/models/vision.py
View file

@ -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)