Flexattn refactor

[danielhanchen]

No description provided.

[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 streamlines how attention mechanisms are selected and configured across various models within the unsloth library. By introducing a unified, hierarchical function, it ensures that the most performant and compatible attention implementation (Flash Attention 2, Flex Attention, SDPA, or Eager) is automatically chosen based on available hardware, model support, and environmental settings. This change enhances efficiency and consistency in attention handling without altering the core behavior for models with known limitations.

Highlights

• Refactored Attention Implementation Logic: The core logic for determining the attention mechanism has been refactored into a new function, determine_attention_implementation, replacing the previous prefer_flex_attn_if_supported.
• Hierarchical Attention Selection: The new function introduces a clear hierarchy for selecting attention implementations: Flash Attention 2 is prioritized, followed by Flex Attention, then SDPA, and finally falling back to Eager attention.
• Model-Specific Exclusions Maintained: Specific model types like GPT-OSS, Mllama, and Gemma3 variants continue to be excluded from Flash Attention 2 and Flex Attention due to known compatibility issues, ensuring stable operation.
• Unified Attention Configuration: The llama.py and vision.py models now consistently use the new determine_attention_implementation function to configure their attention mechanisms, simplifying the codebase and improving maintainability.

Changelog

• unsloth/models/_utils.py
  • Renamed prefer_flex_attn_if_supported to determine_attention_implementation.
  • Implemented a new hierarchical logic for selecting attention mechanisms (Flash Attention 2, Flex Attention, SDPA, Eager).
  • Added checks for HAS_FLASH_ATTENTION and model-specific support for different attention types.
  • Maintained exclusions for certain model types (gpt_oss, mllama, gemma3) from Flash Attention 2 and Flex Attention.
• unsloth/models/llama.py
  • Updated the from_pretrained method to use the new determine_attention_implementation function for selecting the preferred attention implementation.
• unsloth/models/vision.py
  • Removed previous logic for flex_attn_impl and default_attn_impl.
  • Integrated the new determine_attention_implementation function to set attn_impl.
  • Simplified the logic for setting attn_implementation in kwargs, removing specific gemma3n checks.

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

This pull request refactors the attention implementation selection logic by introducing a new determine_attention_implementation function. This function centralizes the logic for choosing the best available attention backend, prioritizing them in the order of Flash Attention 2, Flex Attention, SDPA, and Eager. This is a great improvement for code clarity and maintainability. I have one suggestion to further refactor the new function to reduce some code duplication.

[gemini-code-assist]

This function has a good structure, but there's a lot of repeated code for setting the _attn_implementation and attn_implementation attributes on the config object. This can be refactored into a helper function to improve maintainability and reduce code duplication. I've also taken the liberty to slightly simplify some of the conditional logic.

def determine_attention_implementation(model_class, config):
    model_type = getattr(config, "model_type", "").lower()

    def _set_attn_impl_and_return(implementation):
        if config is not None:
            setattr(config, "_attn_implementation", implementation)
            if hasattr(config, "attn_implementation"):
                setattr(config, "attn_implementation", implementation)
        return implementation

    # 1. Flash Attention 2
    if (
        HAS_FLASH_ATTENTION
        and model_type not in ("gpt_oss", "mllama")
        and not model_type.startswith("gemma3")
    ):
        supports_fa2 = (model_class is not None) and (
            getattr(model_class, "_supports_flash_attn_2", False)
            or getattr(model_class, "_supports_flash_attn", False)
        )
        if supports_fa2:
            return _set_attn_impl_and_return("flash_attention_2")

    # 2. Flex Attention
    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 (model_class is not None)
                and getattr(model_class, "_supports_flex_attn", False)
                and model_type not in ("gpt_oss", "mllama")
                and not model_type.startswith("gemma3")
            ):
                # GPT-OSS, Mllama and Gemma3 use eager/sdpa attention during
                # inference since flex attention returns incorrect results or errors out.
                # GPT-OSS: left padding issues cause incorrect outputs.
                # Mllama: _update_causal_mask uses make_flex_block_causal_mask which
                # creates BlockMask with Q_LEN=KV_LEN=total_seq_len, but during
                # decode q_len=1, causing ValueError. Needs transformers update.
                # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not
                # support flex_attention.
                return _set_attn_impl_and_return("flex_attention")
        except Exception:
            pass

    # 3. SDPA
    if model_class is not None and getattr(model_class, "_supports_sdpa", False):
        return _set_attn_impl_and_return("sdpa")

    # 4. Eager
    return _set_attn_impl_and_return("eager")

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c0b9cd98e8

ℹ️ 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]

Restore Gemma3N eager default for auto attention selection

This assignment now applies determine_attention_implementation() directly, which removed the previous gemma3n-specific eager override from FastBaseModel.from_pretrained. With the new helper, gemma3n is excluded from flash/flex and then falls through to SDPA when _supports_sdpa is true, so Gemma3N defaults to sdpa instead of the prior eager-safe path. That is a behavior regression for Gemma3N loads (especially when callers do not pass attn_implementation) and can reintroduce the attention incompatibility that the earlier Gemma3N guard in this loader was added to avoid.

Useful? React with 👍 / 👎.

[Datta0]

I guess I'll let @mmathew23 review this as it doesn't make much sense to approve my code myself lol 😭