Repository: huggingface/transformers
File: src/transformers/modeling_utils.py (Python)

Code before:
            is_init_check (`bool`, *optional*):
                Whether this check is performed early, i.e. at __init__ time, or later when the model and its weights are
                fully instantiated. This is needed as we also check the devices of the weights, which are only available
                later after __init__. This allows to raise proper exceptions early before instantiating the full models
                if we know that the model does not support the requested attention.

        Returns:
            `str`: The final attention implementation to use, including potential fallbacks from sdpa to eager, or from
            None to sdpa (to potentially eager).
        """
        applicable_attn_implementation = attn_implementation

        is_paged = attn_implementation is not None and attn_implementation.startswith("paged|")

        # If FA not installed, do not fail but use kernels instead
        requested_original_flash_attn = attn_implementation is not None and (
            attn_implementation.removeprefix("paged|") == "flash_attention_2"
            or attn_implementation.removeprefix("paged|") == "flash_attention_3"
        )
        if (
            requested_original_flash_attn
            and self._supports_flash_attn
            and not (is_flash_attn_2_available() or is_flash_attn_3_available())
            and is_kernels_available()
            and not is_torch_npu_available()
        ):
            applicable_attn_implementation = FLASH_ATTN_KERNEL_FALLBACK[attn_implementation.removeprefix("paged|")]

            if is_torch_xpu_available() and attn_implementation.removeprefix("paged|") == "flash_attention_2":
                # On XPU, kernels library is the native implementation
                # Disabling this flag to avoid giving wrong fallbacks on errors and warnings
                requested_original_flash_attn = False

            if is_paged:
                applicable_attn_implementation = f"paged|{applicable_attn_implementation}"

        if is_kernel(applicable_attn_implementation):
            try:
                # preload flash attention here to allow compile with fullgraph
                if is_paged:
                    lazy_import_paged_flash_attention(applicable_attn_implementation)
                else:
                    lazy_import_flash_attention(applicable_attn_implementation)

                # log that we used kernel fallback if successful
                if requested_original_flash_attn:
                    logger.warning_once(
                        f"You do not have `flash_attn` installed, using `{applicable_attn_implementation}` "
                        "from the `kernels` library instead!"
                    )
            except Exception as e:

Review comment (suggestion):
Sorry a bit late but wdyt about changing this to something along `_compatible_flash_implementations` (and making it a list)? I'm expecting that at some point other FA implementations will also work e.g. sinks for FA4 + we should only warn if we encounter that the requested FA is not one of the compatible ones

Code after:
        Args:
            attn_implementation (`str` or `None`):
                The attention implementation to check for existence/validity.
            is_init_check (`bool`, *optional*):
                Whether this check is performed early, i.e. at __init__ time, or later when the model and its weights are
                fully instantiated. This is needed as we also check the devices of the weights, which are only available
                later after __init__. This allows to raise proper exceptions early before instantiating the full models
                if we know that the model does not support the requested attention.

        Returns:
            `str`: The final attention implementation to use, including potential fallbacks from sdpa to eager, or from
            None to sdpa (to potentially eager).
        """
        # Auto-correct flash_attention_2/3 to model's default flash implementation if specified
        if attn_implementation is not None:
            is_paged = attn_implementation.startswith("paged|")
            base_impl = attn_implementation.removeprefix("paged|")

            if base_impl in ("flash_attention_2", "flash_attention_3"):
                default_flash = getattr(self, "_default_flash_implementation", None)
                if default_flash is not None:
                    logger.warning_once(
                        f"This model is optimized for '{default_flash}'. "
                        f"Automatically using '{default_flash}' instead of '{attn_implementation}'."
                    )
                    attn_implementation = f"paged|{default_flash}" if is_paged else default_flash

        applicable_attn_implementation = attn_implementation

        is_paged = attn_implementation is not None and attn_implementation.startswith("paged|")

        # If FA not installed, do not fail but use kernels instead
        requested_original_flash_attn = attn_implementation is not None and (
            attn_implementation.removeprefix("paged|") == "flash_attention_2"
            or attn_implementation.removeprefix("paged|") == "flash_attention_3"
        )
        if (
            requested_original_flash_attn
            and self._supports_flash_attn
            and not (is_flash_attn_2_available() or is_flash_attn_3_available())
            and is_kernels_available()
            and not is_torch_npu_available()
        ):
            applicable_attn_implementation = FLASH_ATTN_KERNEL_FALLBACK[attn_implementation.removeprefix("paged|")]

            if is_torch_xpu_available() and attn_implementation.removeprefix("paged|") == "flash_attention_2":
                # On XPU, kernels library is the native implementation
                # Disabling this flag to avoid giving wrong fallbacks on errors and warnings
                requested_original_flash_attn = False

            if is_paged: