Honor an explicit sdpa or flex_attention request when flash is disabled#6847
Honor an explicit sdpa or flex_attention request when flash is disabled#6847danielhanchen wants to merge 3 commits into
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 👍 / 👎.
| 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 👍 / 👎.
|
The helper already honors an explicit |
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.