Gemma attn

[Datta0]

For gemma3, previously we had the "flex_attn" backend. But after the refactor in #4210 #4201 it seems to have been falling back to SDPA. This is an attempt to restore old behaviour via

1. Set a list of models which PREFER flex attn
2. if the model is among those and we see flex attn available, we use that.

Using SDPA was resulting in errors wrt mask shape. The mask behaviour is addressed in https://github.com/unslothai/unsloth-zoo/pull/635/changes

[gemini-code-assist]

Code Review

This pull request introduces logic to prioritize Flex Attention for specific models, such as Gemma 3 and ShieldGemma 2, by adding preference and support check helpers. The resolve_attention_implementation function was refactored to incorporate these checks. A logic bug was identified in the fallback sequence where models preferring Flex Attention would skip Flash Attention 2 if Flex was unavailable, potentially leading to inefficient performance. It is recommended to remove the restrictive condition to allow Flash Attention 2 as a fallback.

[gemini-code-assist]

The condition not prefers_flex_attention at line 505 creates a logic bug where models that prefer Flex Attention but cannot use it (e.g., due to an older PyTorch version or environment settings) will skip Flash Attention 2 even if it is available and supported. This forces a fallback to SDPA or Eager, which is significantly less efficient. Since the previous elif block already handles the case where Flex is both preferred and supported, this check is unnecessary and harmful for fallback scenarios.

Suggested change

	elif (
	not prefers_flex_attention
	and not flash_attention_disabled
	and HAS_FLASH_ATTENTION
	and supports_flash_attention
	):
	attn_impl = _set_attn_impl(config, "flash_attention_2")
	elif (
	not flash_attention_disabled
	and HAS_FLASH_ATTENTION
	and supports_flash_attention
	):
	attn_impl = _set_attn_impl(config, "flash_attention_2")

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6640371bff

ℹ️ 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".

[chatgpt-codex-connector]

Allow flash fallback when preferred flex backend is unavailable

The new not prefers_flex_attention guard in resolve_attention_implementation prevents all preferred models (e.g. gemma3) from ever selecting flash_attention_2 unless flex attention is already available. In environments where is_torch_flex_attn_available() is false but Flash Attention is installed, this now falls through to SDPA/eager instead of using flash, which is a regression from the previous selection order and can reintroduce the SDPA-path failures this change is trying to avoid.

Useful? React with 👍 / 👎.