Compatibility with modeling code relying on `flash-attn` package

[pmalic]

There are widely used models on HF that ship with custom modeling code that relies on the original flash-attn package, either by calling transformers.utils.is_flash_attn_2_available() and/or relying that the AutoModel.from_pretrained() parameter attn_implementation is equal to flash_attention_2 (for example, Jina Embeddings v4: https://huggingface.co/jinaai/jina-embeddings-v4/blob/main/qwen2_5_vl.py).

I wanted to switch to kernels-community/flash-attn2 to avoid building flash-attn or using pre-built wheels from https://github.com/mjun0812/flash-attention-prebuild-wheels (has no wheels index, so you have to hardcode the path to a particular .whl file) and managed to do it with a few lines of code.

Have you thought about adding some documentation/example for this case. In my experience, model authors rarely update the included custom modeling code to support new stuff (e.g., Transformers v5 or kernels), so unless you want to create your own local modeling code, you're stuck with whatever the authors used when they published the model.

Here's the code that "exposes" kernels-community/flash-attn2 as the original flash-attn:

from pathlib import Path
import sys

from kernels import get_kernel


def setup_flash_attn_2() -> None:

    flash_attn = get_kernel("kernels-community/flash-attn2", version=1)

    sys.modules["flash_attn"] = flash_attn

    dist_info = Path(flash_attn.__file__).parent / "flash_attn.dist-info"
    dist_info.mkdir(parents=True, exist_ok=True)

    (dist_info / "METADATA").write_text("Metadata-Version: 2.5\nName: flash-attn\nVersion: 2.8.3")

    sys.path.append(str(dist_info.parent))

[sayakpaul]

You can do by providing attn_implementation="kernels-community/flash-attn2" while loading the model with from_pretrained().

Cc: @vasqu from transformers.

[vasqu]

We also auto fallback to kernels when we detect any base FA request, see kernels in your environment but not the requested FA package! (you will get a logger information for fallback)

[pmalic]

Hi @vasqu, thanks for jumping in. I wasn't aware of this auto fallback in Transformers, otherwise I wouldn't create the issue in this repo, sorry @sayakpaul.

The fallback, you're referring to FLASH_ATTN_KERNEL_FALLBACK?

This is v5 only, can't find it in v4? The thing is, people who are having the problem I described above are also very likely stuck with v4 :)

Also, based on reading this test_kernels_fallback(), is_flash_attn_2_available() isn't "covered" by the auto fallback, so if the custom modeling code uses it (as the model I mentioned above does), you're out of luck.

Is it even realistic that the fallback gets back-ported to v4 and that is_flash_attn_2_available() also uses it in some mode (if not by default)?

[vasqu]

Hey @pmalic, there already is a fallback in v4 - it's just more hidden https://github.com/huggingface/transformers/blob/753d61104116eefc8ffc977327b441ee0c8d599f/src/transformers/modeling_utils.py#L2653-L2664

Otherwise, you can still pass the kernel directly instead and we use that for the attention calculation