Honor an explicit sdpa or flex_attention request when flash is disabled#6847
Conversation
When flash attention is disabled for a model, the fallback selection could downgrade a caller who explicitly passed attn_implementation='sdpa' or 'flex_attention' to a different backend, because the disable reason is flash-specific. Keep an explicit non-flash request as-is; flash requests still fall back as before.
for more information, see https://pre-commit.ci
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request ensures that explicitly requested non-flash attention implementations, such as "sdpa" or "flex_attention", are honored and not downgraded to fallback options when flash attention is disabled. A comprehensive set of unit tests has been added to verify this behavior. There are no review comments, and I have no additional feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@codex review |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request ensures that explicit requests for non-flash attention implementations, specifically 'sdpa' and 'flex_attention', are honored and not downgraded when flash attention is disabled. A new test suite has been added to verify this behavior. There are no review comments, and I have no additional feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5cd1ab7706
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if requested_attn_implementation in ("sdpa", "flex_attention"): | ||
| return _set_attn_impl(config, requested_attn_implementation) |
There was a problem hiding this comment.
When this helper is reached without an explicit attn_implementation argument (the default resolver path), requested_attn_implementation is populated from the config, so a stale/default _attn_implementation='flex_attention' is now returned even when supports_flex_attention is false (for example because flex was disabled by env, dropout, model exclusions, or unavailable imports). Prior code would fall back to SDPA/eager in that scenario; gate this new branch on the argument actually being passed explicitly, or unsupported flex can be selected during model loading.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Transformers never defaults _attn_implementation to flex_attention (the default is sdpa), and neither of our callers seeds it from config: the language path loads the config with attn_implementation="sdpa", and the vision path passes the user request as the explicit requested_attn_implementation argument. So an implicit flex_attention in config only appears when it was actually chosen, in which case honoring it is correct. The flash disable reason is head-dim / sink specific and is independent of whether flex is usable, which is computed separately. Left as is.
| attn_implementation = "sdpa", | ||
| supports_sdpa = False, # conservative flag would have skipped sdpa |
There was a problem hiding this comment.
This test covers the private helper with supports_sdpa=False, but real callers go through resolve_attention_implementation; after the helper returns sdpa, the resolver's final if not supports_sdpa and final_attn_impl == "sdpa" guard immediately switches it back to eager. In the flash-disabled + conservative supports_sdpa=False scenario described here, an explicit attn_implementation="sdpa" still is not honored despite this test passing, so the resolver path needs to be fixed or covered directly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The resolver switch from sdpa to eager only fires when supports_sdpa is False, meaning the model genuinely does not support SDPA, so eager is the correct target and forcing sdpa there would just fail at load. In the realistic flash-disabled cases (oversized head dims or sink-attention models) supports_sdpa is True, so an explicit attn_implementation="sdpa" is honored end to end. The unit test intentionally pins the helper contract with the conservative flag; the eager downgrade is a separate SDPA-capability guard, not the flash disable path. Left as is.
|
The helper already honors an explicit |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 084ec4b809
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if requested_attn_implementation in ("sdpa", "flex_attention"): | ||
| return _set_attn_impl(config, requested_attn_implementation) |
There was a problem hiding this comment.
When flash is disabled and a caller requests flex_attention, this early return ignores supports_flex_attention. That flag is false not only for conservative fallback choices, but also for known-broken or unavailable flex paths: _supports_flex_attention rejects _FLEX_EXCLUDED_MODELS such as gpt_oss, disabled flex envs, missing Torch flex support, and model classes that do not advertise flex. In a GPT-OSS load with attn_implementation="flex_attention", the resolver now selects a backend the repo explicitly excludes instead of falling back to eager, which can produce incorrect decoding.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. supports_flex_attention already rejects the excluded/broken/unavailable configs, so I now honor an explicit flex_attention only when that flag is true; gpt_oss and friends fall back instead of selecting a broken backend. Explicit sdpa stays honored unconditionally.
| if requested_attn_implementation in ("sdpa", "flex_attention"): | ||
| return _set_attn_impl(config, requested_attn_implementation) |
There was a problem hiding this comment.
When attn_implementation is None, requested_attn_implementation can come from the config rather than from the caller; the language loader creates configs with attn_implementation="sdpa" and then calls resolve_attention_implementation without a requested value. For a flash-disabled model that supports flex but not SDPA, this branch returns that synthesized SDPA instead of taking the supports_flex_attention fallback below, after which the resolver's not supports_sdpa guard downgrades all the way to eager and skips the intended flex path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Agreed. The honor branch now only fires for an implementation passed by the caller, not for the sdpa the language loader seeds into the config, so a flash-disabled flex-but-not-sdpa model keeps the flex fallback instead of downgrading to eager.
Only honor an explicit non-flash attention request when it comes from the caller argument, not from a config value the loaders synthesize (the language path seeds attn_implementation=sdpa). Honor explicit flex_attention only when supports_flex_attention is True so excluded/broken configs (e.g. gpt_oss) fall back instead of selecting a known-broken backend. Explicit sdpa stays honored.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 38913751e8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if explicit_request == "sdpa": | ||
| return _set_attn_impl(config, "sdpa") |
There was a problem hiding this comment.
When the caller requests attn_implementation="sdpa" for a flash-disabled config while supports_sdpa is false, this helper now returns sdpa, but the public resolve_attention_implementation path immediately applies its existing if not supports_sdpa and final_attn_impl == "sdpa" guard and rewrites it to eager. That means the loader still does not honor the explicit SDPA request in the conservative-support scenario this change is testing; only direct calls to the private helper see the new behavior.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. The resolver captured the helper output but then its own not-supports_sdpa guard rewrote an honored explicit sdpa back to eager. It now records whether sdpa was an explicit caller request and skips that downgrade only for explicit requests; a synthesized or default sdpa still downgrades to eager. Added two resolver-path tests (explicit sdpa stays sdpa, non-explicit sdpa downgrades).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 78a6f17522
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # sdpa (requested_attn_implementation is None, so the value came from the model | ||
| # resolution above or the config) downgrades to eager. | ||
| explicit_sdpa_request = requested_attn_implementation == "sdpa" | ||
| if not supports_sdpa and final_attn_impl == "sdpa" and not explicit_sdpa_request: |
There was a problem hiding this comment.
Please keep the unsupported-SDPA downgrade when supports_sdpa was cleared by model-specific exclusions. For model_type='gpt_oss', flash is disabled via _FLASH_EXCLUDED_MODELS and SDPA is also explicitly disabled via _SDPA_EXCLUDED_MODELS, but requested_attn_implementation='sdpa' now makes this guard skip the fallback to eager and re-enables an attention path the repo marks as wrong, which can silently produce incorrect outputs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. supports_sdpa is already cleared for _SDPA_EXCLUDED_MODELS, so I now honor an explicit sdpa only when the model is not sdpa-excluded; gpt_oss stays on eager instead of re-enabling the known-broken sdpa path, matching how an explicit flex request falls back for _FLEX_EXCLUDED_MODELS. The conservative flash-disabled case (large head dim or attention-sink models where sdpa works) keeps its explicit sdpa. Added a resolver test for gpt_oss plus explicit sdpa resolving to eager.
An explicit attn_implementation="sdpa" was re-enabling sdpa for models in _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) where sdpa is known-broken: the helper honored the request and the resolver's final not-supports_sdpa guard skipped the eager downgrade for any explicit request. Honor an explicit sdpa only when the model is not sdpa-excluded, mirroring the flex guard that already falls back for _FLEX_EXCLUDED_MODELS via supports_flex_attention. Conservative supports_sdpa=False (large head dim / attention-sink models) still honors an explicit sdpa; a synthesized/default sdpa still downgrades to eager.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0e9812381d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| honor_explicit_sdpa = requested_attn_implementation == "sdpa" and not _is_sdpa_excluded( | ||
| model_type | ||
| ) | ||
| if not supports_sdpa and final_attn_impl == "sdpa" and not honor_explicit_sdpa: |
There was a problem hiding this comment.
For Gemma3/Gemma3Text, the loader deliberately forces supports_sdpa = False via DISABLE_SDPA_MODEL_NAMES because those SDPA modules are marked wrong (unsloth/models/loader.py:198-204, applied at loader.py:1558-1560). This new exception makes requested_attn_implementation="sdpa" bypass the final downgrade for every model not in _SDPA_EXCLUDED_MODELS, which currently only covers gpt_oss, so an explicit SDPA request now re-enables a known-bad backend for Gemma3 loads and can produce incorrect results instead of falling back to eager/flex.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. _is_sdpa_excluded now also treats DISABLE_SDPA_MODEL_NAMES membership as excluded, using the same trailing-comma substring match as the loader, so an explicit sdpa on gemma3/gemma3_text downgrades to eager while a genuinely conservative supports_sdpa=False model (not in either list) still keeps sdpa. Moved the constant into _utils.py (re-exported from loader.py) to avoid an import cycle, and added resolver tests.
The honor-explicit-sdpa guard only skipped the sdpa->eager downgrade for models in _SDPA_EXCLUDED_MODELS (gpt_oss). Gemma3/Gemma3Text disable SDPA through the loader's DISABLE_SDPA_MODEL_NAMES (their bundled SDPA modules are wrong), so an explicit sdpa request bypassed the downgrade and re-enabled a known-wrong path. Extend _is_sdpa_excluded to also treat DISABLE_SDPA_MODEL_NAMES membership as excluded, replicating the loader's trailing-comma substring match so gemma3 and gemma3_text match but gemma3n does not. Move the constant into _utils.py (single source of truth, re-exported from loader.py) to avoid a loader -> _utils cycle. Conservative supports_sdpa=False models not in either list still honor explicit sdpa.
|
@codex review |
|
Codex Review: Didn't find any major issues. Delightful! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
|
Codex Review: Didn't find any major issues. Chef's kiss. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
Summary
When flash attention is disabled for a model, the fallback selection could downgrade a caller who explicitly passed
attn_implementation="sdpa"or"flex_attention"to a different backend, because the disable reason is flash-specific.What this does
Keeps an explicit non-flash request as-is; flash requests still fall back as before. Includes a unit test.