refactor attn impl replacement

[Datta0]

The previous impl was setting it to flex_attn for almost every model causing the models to not use FA2. This change creates a preference hierarchy

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly refactors the attention mechanism selection logic within the Unsloth framework. The primary goal is to establish a clear preference hierarchy for attention implementations, ensuring that models leverage the most efficient available option, such as Flash Attention 2, before falling back to Flex Attention, SDPA, or Eager attention. This change addresses a previous issue where models were often defaulting to Flex Attention even when more performant alternatives were supported, thereby improving overall model performance and resource utilization.

Highlights

• Refactored Attention Implementation Selection: Introduced a new, unified function determine_attention_implementation to replace the previous prefer_flex_attn_if_supported function, centralizing the logic for selecting the optimal attention mechanism.
• Hierarchical Preference for Attention Mechanisms: Implemented a clear preference hierarchy for attention implementations: Flash Attention 2 (FA2) is now prioritized, followed by Flex Attention, then SDPA, and finally Eager attention as a fallback. This ensures models utilize the most performant option available.
• Corrected Attention Mechanism Defaulting: Addressed an issue where many models were inadvertently defaulting to Flex Attention, preventing them from utilizing more efficient alternatives like Flash Attention 2. The new hierarchy ensures FA2 is chosen when supported.
• Improved Model-Specific Handling: Consolidated model-specific exclusions (e.g., for gpt_oss, mllama, gemma3) within the new determination function, making the logic more robust and easier to maintain across different attention types.

Changelog

• unsloth/models/_utils.py
  • Renamed the function prefer_flex_attn_if_supported to determine_attention_implementation.
  • Implemented a new hierarchical logic within determine_attention_implementation to select attention based on availability and model support, prioritizing Flash Attention 2, then Flex Attention, SDPA, and finally Eager attention.
• unsloth/models/llama.py
  • Updated the call for determining the preferred attention implementation to use the new determine_attention_implementation function.
• unsloth/models/vision.py
  • Replaced the previous conditional logic for flex_attn_impl and default_attn_impl with a single call to the new determine_attention_implementation function.
  • Removed redundant model-specific handling for gemma3n attention defaults, as this is now managed by the unified determination function.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The pull request refactors the attention mechanism selection logic by introducing a new determine_attention_implementation function. This function centralizes the process of selecting the optimal attention implementation (Flash Attention 2, Flex Attention, SDPA, or Eager) based on a hierarchical fallback strategy, model capabilities, and specific model type exclusions. This change replaces the previous prefer_flex_attn_if_supported function and simplifies the attention configuration across various model types, including Llama and Vision models.

_{Note: Security Review is unavailable for this PR.}