From ef52526fc752318aec16bbb99f1b44ad6d9c8190 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Fri, 8 May 2026 15:39:44 +0000 Subject: [PATCH 1/2] Improve attn refactor --- tests/test_attention_implementation.py | 178 +++++++++++++++++++++++++ unsloth/models/_utils.py | 65 +++++---- 2 files changed, 220 insertions(+), 23 deletions(-) create mode 100644 tests/test_attention_implementation.py diff --git a/tests/test_attention_implementation.py b/tests/test_attention_implementation.py new file mode 100644 index 00000000000..1832c0b088f --- /dev/null +++ b/tests/test_attention_implementation.py @@ -0,0 +1,178 @@ +from types import SimpleNamespace + +import unsloth # noqa: F401 +from transformers.utils import import_utils + +from unsloth.models import _utils + + +class SupportsAllAttention: + _supports_flash_attn_2 = True + _supports_flex_attn = True + _supports_sdpa = True + + +class SupportsFlexAndSdpa: + _supports_flash_attn_2 = False + _supports_flex_attn = True + _supports_sdpa = True + + +def _config(model_type, **kwargs): + values = {"model_type": model_type, "attention_dropout": 0} + values.update(kwargs) + return SimpleNamespace(**values) + + +def _enable_attention_backends(monkeypatch, flex_available=True): + monkeypatch.setattr(_utils, "HAS_FLASH_ATTENTION", True) + monkeypatch.setattr( + import_utils, + "is_torch_flex_attn_available", + lambda: flex_available, + raising=False, + ) + monkeypatch.setenv("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") + + +def test_gemma3_family_prefers_flex_attention_over_flash_attention(monkeypatch): + _enable_attention_backends(monkeypatch) + + for model_type in ("gemma3", "gemma3_text", "shieldgemma2"): + config = _config(model_type) + + impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) + + assert impl == "flex_attention" + assert config._attn_implementation == "flex_attention" + + +def test_nested_gemma3_text_prefers_flex_attention(monkeypatch): + _enable_attention_backends(monkeypatch) + text_config = _config("gemma3_text") + config = _config("gemma3", text_config=text_config) + + impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) + + assert impl == "flex_attention" + assert config._attn_implementation == "flex_attention" + + +def test_non_gemma3_keeps_flash_attention_priority(monkeypatch): + _enable_attention_backends(monkeypatch) + config = _config("llama") + + impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) + + assert impl == "flash_attention_2" + assert config._attn_implementation == "flash_attention_2" + + +def test_generic_model_uses_flex_when_flash_attention_is_not_safe(monkeypatch): + _enable_attention_backends(monkeypatch) + config = _config("llama", head_dim=320) + + impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) + + assert impl == "flex_attention" + assert config._attn_implementation == "flex_attention" + + +def test_gemma3n_defaults_to_eager_unless_explicitly_overridden(monkeypatch): + _enable_attention_backends(monkeypatch) + config = _config("gemma3n") + + impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) + + assert impl == "eager" + assert config._attn_implementation == "eager" + + impl = _utils.resolve_attention_implementation( + SupportsAllAttention, + config, + requested_attn_implementation="sdpa", + ) + + assert impl == "sdpa" + assert config._attn_implementation == "sdpa" + + +def test_flex_excluded_model_does_not_auto_select_flex(monkeypatch): + _enable_attention_backends(monkeypatch) + config = _config("gpt_oss") + + impl = _utils.resolve_attention_implementation(SupportsFlexAndSdpa, config) + + assert impl == "sdpa" + assert config._attn_implementation == "sdpa" + + +def test_explicit_gemma3_flash_attention_request_is_preserved(monkeypatch): + _enable_attention_backends(monkeypatch) + config = _config("gemma3") + + impl = _utils.resolve_attention_implementation( + SupportsAllAttention, + config, + requested_attn_implementation="flash_attention_2", + ) + + assert impl == "flash_attention_2" + assert config._attn_implementation == "flash_attention_2" + + +def test_explicit_gemma3_sdpa_and_eager_requests_are_preserved(monkeypatch): + _enable_attention_backends(monkeypatch) + + for requested in ("sdpa", "eager"): + config = _config("gemma3") + + impl = _utils.resolve_attention_implementation( + SupportsAllAttention, + config, + requested_attn_implementation=requested, + ) + + assert impl == requested + assert config._attn_implementation == requested + + +def test_explicit_flash_attention_request_falls_back_on_unsupported_head_dim( + monkeypatch, +): + _enable_attention_backends(monkeypatch) + config = _config( + "gemma3", + head_dim=320, + ) + + impl = _utils.resolve_attention_implementation( + SupportsAllAttention, + config, + requested_attn_implementation="flash_attention_2", + ) + + assert impl == "sdpa" + assert config._attn_implementation == "sdpa" + + +def test_gemma3_uses_flash_attention_when_flex_is_disabled(monkeypatch): + _enable_attention_backends(monkeypatch) + monkeypatch.setenv("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") + config = _config("gemma3") + + impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) + + assert impl == "flash_attention_2" + assert config._attn_implementation == "flash_attention_2" + + +def test_flex_attention_requires_zero_dropout_in_nested_config(monkeypatch): + _enable_attention_backends(monkeypatch) + text_config = _config("gemma3_text", attention_dropout=0.1) + config = _config("gemma3", text_config=text_config) + + impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) + + assert impl == "flash_attention_2" + assert config._attn_implementation == "flash_attention_2" diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index d3eee033256..ea78d574c99 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -234,6 +234,7 @@ def apply_unsloth_gradient_checkpointing( # ModernBERT: create_block_mask with _compile=True hits CUDA illegal memory # access on some GPU architectures (B200). Falls back to eager safely. _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") +_FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2") _EAGER_ONLY_PREFIXES = ("gemma3n",) _FLASH_ATTENTION_MAX_HEAD_DIM = 256 _FLASH_ATTENTION_DISABLED_WARNED = set() @@ -243,10 +244,37 @@ def _is_flex_excluded(model_type): return model_type in _FLEX_EXCLUDED_MODELS +def _config_prefers_flex_attention(config): + return any( + _config_get(attention_config, "model_type", "").lower() + in _FLEX_PREFERRED_MODELS + for attention_config in _iter_attention_configs(config) + ) + + def _is_eager_only(model_type): return any(model_type.startswith(p) for p in _EAGER_ONLY_PREFIXES) +def _supports_flex_attention(model_class, config, model_type): + if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0": + return False + if model_class is None or not getattr(model_class, "_supports_flex_attn", False): + return False + if _is_flex_excluded(model_type): + return False + for attention_config in _iter_attention_configs(config): + attention_dropout = _config_get(attention_config, "attention_dropout", 0) or 0 + if attention_dropout != 0: + return False + try: + from transformers.utils.import_utils import is_torch_flex_attn_available + + return is_torch_flex_attn_available() + except Exception: + return False + + def _config_items(config): if isinstance(config, dict): return config.items() @@ -459,6 +487,9 @@ def resolve_attention_implementation( getattr(model_class, "_supports_flash_attn_2", False) or getattr(model_class, "_supports_flash_attn", False) ) + supports_flex_attention = _supports_flex_attention( + model_class, config, model_type + ) disable_reason = _get_flash_attention_disable_reason(config) flash_attention_disabled = disable_reason is not None @@ -467,6 +498,16 @@ def resolve_attention_implementation( else: if _is_eager_only(model_type): attn_impl = _set_attn_impl(config, "eager") + elif _config_prefers_flex_attention(config) and supports_flex_attention: + attn_impl = _set_attn_impl(config, "flex_attention") + elif ( + not flash_attention_disabled + and HAS_FLASH_ATTENTION + and supports_flash_attention + ): + attn_impl = _set_attn_impl(config, "flash_attention_2") + elif supports_flex_attention: + attn_impl = _set_attn_impl(config, "flex_attention") elif flash_attention_disabled: attn_impl = _disable_flash_attention_if_needed( config, @@ -476,32 +517,10 @@ def resolve_attention_implementation( ), disable_reason = disable_reason, ) - elif HAS_FLASH_ATTENTION and supports_flash_attention: - attn_impl = _set_attn_impl(config, "flash_attention_2") elif supports_sdpa: attn_impl = _set_attn_impl(config, "sdpa") else: - attn_impl = "eager" - 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 getattr(model_class, "_supports_flex_attn", False) - and not _is_flex_excluded(model_type) - ): - attention_dropout = ( - _config_get(config, "attention_dropout", 0) or 0 - ) - if attention_dropout == 0: - attn_impl = _set_attn_impl(config, "flex_attention") - except Exception: - pass - if attn_impl == "eager": - attn_impl = _set_attn_impl(config, "eager") + attn_impl = _set_attn_impl(config, "eager") if requested_attn_implementation is None: final_attn_impl = attn_impl From 426bf74df9052da69b85a01012dbdcc7a4f40a8b Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Sat, 9 May 2026 13:03:10 +0000 Subject: [PATCH 2/2] refactor attention impl determine Signed-off-by: Datta Nimmaturi --- tests/test_attention_implementation.py | 178 ------------------------- unsloth/models/_utils.py | 14 +- 2 files changed, 8 insertions(+), 184 deletions(-) delete mode 100644 tests/test_attention_implementation.py diff --git a/tests/test_attention_implementation.py b/tests/test_attention_implementation.py deleted file mode 100644 index 1832c0b088f..00000000000 --- a/tests/test_attention_implementation.py +++ /dev/null @@ -1,178 +0,0 @@ -from types import SimpleNamespace - -import unsloth # noqa: F401 -from transformers.utils import import_utils - -from unsloth.models import _utils - - -class SupportsAllAttention: - _supports_flash_attn_2 = True - _supports_flex_attn = True - _supports_sdpa = True - - -class SupportsFlexAndSdpa: - _supports_flash_attn_2 = False - _supports_flex_attn = True - _supports_sdpa = True - - -def _config(model_type, **kwargs): - values = {"model_type": model_type, "attention_dropout": 0} - values.update(kwargs) - return SimpleNamespace(**values) - - -def _enable_attention_backends(monkeypatch, flex_available=True): - monkeypatch.setattr(_utils, "HAS_FLASH_ATTENTION", True) - monkeypatch.setattr( - import_utils, - "is_torch_flex_attn_available", - lambda: flex_available, - raising=False, - ) - monkeypatch.setenv("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") - - -def test_gemma3_family_prefers_flex_attention_over_flash_attention(monkeypatch): - _enable_attention_backends(monkeypatch) - - for model_type in ("gemma3", "gemma3_text", "shieldgemma2"): - config = _config(model_type) - - impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) - - assert impl == "flex_attention" - assert config._attn_implementation == "flex_attention" - - -def test_nested_gemma3_text_prefers_flex_attention(monkeypatch): - _enable_attention_backends(monkeypatch) - text_config = _config("gemma3_text") - config = _config("gemma3", text_config=text_config) - - impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) - - assert impl == "flex_attention" - assert config._attn_implementation == "flex_attention" - - -def test_non_gemma3_keeps_flash_attention_priority(monkeypatch): - _enable_attention_backends(monkeypatch) - config = _config("llama") - - impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) - - assert impl == "flash_attention_2" - assert config._attn_implementation == "flash_attention_2" - - -def test_generic_model_uses_flex_when_flash_attention_is_not_safe(monkeypatch): - _enable_attention_backends(monkeypatch) - config = _config("llama", head_dim=320) - - impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) - - assert impl == "flex_attention" - assert config._attn_implementation == "flex_attention" - - -def test_gemma3n_defaults_to_eager_unless_explicitly_overridden(monkeypatch): - _enable_attention_backends(monkeypatch) - config = _config("gemma3n") - - impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) - - assert impl == "eager" - assert config._attn_implementation == "eager" - - impl = _utils.resolve_attention_implementation( - SupportsAllAttention, - config, - requested_attn_implementation="sdpa", - ) - - assert impl == "sdpa" - assert config._attn_implementation == "sdpa" - - -def test_flex_excluded_model_does_not_auto_select_flex(monkeypatch): - _enable_attention_backends(monkeypatch) - config = _config("gpt_oss") - - impl = _utils.resolve_attention_implementation(SupportsFlexAndSdpa, config) - - assert impl == "sdpa" - assert config._attn_implementation == "sdpa" - - -def test_explicit_gemma3_flash_attention_request_is_preserved(monkeypatch): - _enable_attention_backends(monkeypatch) - config = _config("gemma3") - - impl = _utils.resolve_attention_implementation( - SupportsAllAttention, - config, - requested_attn_implementation="flash_attention_2", - ) - - assert impl == "flash_attention_2" - assert config._attn_implementation == "flash_attention_2" - - -def test_explicit_gemma3_sdpa_and_eager_requests_are_preserved(monkeypatch): - _enable_attention_backends(monkeypatch) - - for requested in ("sdpa", "eager"): - config = _config("gemma3") - - impl = _utils.resolve_attention_implementation( - SupportsAllAttention, - config, - requested_attn_implementation=requested, - ) - - assert impl == requested - assert config._attn_implementation == requested - - -def test_explicit_flash_attention_request_falls_back_on_unsupported_head_dim( - monkeypatch, -): - _enable_attention_backends(monkeypatch) - config = _config( - "gemma3", - head_dim=320, - ) - - impl = _utils.resolve_attention_implementation( - SupportsAllAttention, - config, - requested_attn_implementation="flash_attention_2", - ) - - assert impl == "sdpa" - assert config._attn_implementation == "sdpa" - - -def test_gemma3_uses_flash_attention_when_flex_is_disabled(monkeypatch): - _enable_attention_backends(monkeypatch) - monkeypatch.setenv("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") - config = _config("gemma3") - - impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) - - assert impl == "flash_attention_2" - assert config._attn_implementation == "flash_attention_2" - - -def test_flex_attention_requires_zero_dropout_in_nested_config(monkeypatch): - _enable_attention_backends(monkeypatch) - text_config = _config("gemma3_text", attention_dropout=0.1) - config = _config("gemma3", text_config=text_config) - - impl = _utils.resolve_attention_implementation(SupportsAllAttention, config) - - assert impl == "flash_attention_2" - assert config._attn_implementation == "flash_attention_2" diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index ea78d574c99..ccd75aa0006 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -259,7 +259,7 @@ def _is_eager_only(model_type): def _supports_flex_attention(model_class, config, model_type): if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0": return False - if model_class is None or not getattr(model_class, "_supports_flex_attn", False): + if not getattr(model_class, "_supports_flex_attn", False): return False if _is_flex_excluded(model_type): return False @@ -487,21 +487,23 @@ def resolve_attention_implementation( getattr(model_class, "_supports_flash_attn_2", False) or getattr(model_class, "_supports_flash_attn", False) ) - supports_flex_attention = _supports_flex_attention( - model_class, config, model_type - ) disable_reason = _get_flash_attention_disable_reason(config) flash_attention_disabled = disable_reason is not None if model_class is None: attn_impl = _set_attn_impl(config, "sdpa" if supports_sdpa else "eager") else: + supports_flex_attention = _supports_flex_attention( + model_class, config, model_type + ) + prefers_flex_attention = _config_prefers_flex_attention(config) if _is_eager_only(model_type): attn_impl = _set_attn_impl(config, "eager") - elif _config_prefers_flex_attention(config) and supports_flex_attention: + elif prefers_flex_attention and supports_flex_attention: attn_impl = _set_attn_impl(config, "flex_attention") elif ( - not flash_attention_disabled + not prefers_flex_attention + and not flash_attention_disabled and HAS_FLASH_ATTENTION and supports_flash_attention ):