feat: add optional flash attention kernels

[art-test-stack]

Hi,

Sorry, I have done something wrong that closed the previous PR related to this #229, and could not reopen it

TLDR; this PR addresses the case of using flash attention in the dot product with attn_mask=None or for causal attention on kernels (not installing flash_attention yet).

Warning

This PR will need to be tested on cuda device (as I don't have my own, I could not; but it should theoretically offer faster performance for the same input-output couple)

flash_attn control

As mentioned in the last PR (#229)

• as proposed by @HennerM, added in AttentionConfig a control on attention implementation (inspired by (axolotl)[https://docs.axolotl.ai/docs/attention.html] as proposed.
• but also added an environment variable NOETHER_ATTN_IMPL to have a higher level control

To disable flash attention mode, just set the environment variable NOETHER_ATTN_IMPL to something different from "fast-attn". For example:

export NOETHER_ATTN_IMPL="sdpa"

module

Added src/noether/modeling/modules/attention/_flash_attention.py.
This module exposes the minimal functions used across the attention modules in this codebase: flash_attn_func, flash_attn_qkvpacked_func, flash_attn_with_kvcache.
The implementation prefers the external kernel when available, but falls back to implementations based on torch.nn.functional.scaled_dot_product_attention so semantics remain identical.

Note

For all flash attention kernel calls, the signatures are from official implementation by Dao-AILab/flash-attention.

comments

To answer the comments that I have not answered in previous PR:

to @Ndles

The current import-time Hub download will bite people in CI/offline

Then I am thinking of implementing both interfaces in one—kernels is easier to install than fa3 or fa4

With that implemented it should be possible to use FA3 and FA4 (hopefully with minimal tweaks 🤞 ).

I'm working on understanding the modules to have these tweaks minimal :)

to @HennerM

although I am a bit surprised that not all kernels support this tbh

yes, but it comes quite from how flash_attention optimizes the dot-product operation. For instance, PyTorch may uses fa2 backend if attn_mask=None.
However, I m looking at how it is possible to "play" with other inputs, as the flash_attention ViT module does here.

todo suggestions

As TODOs (related to faster sdpa kernels) I have identified the following, which could be beneficial:

• Dao-AILab/flash-attention provides an implementation for ViT here. Maybe it can be interesting to think around it to get faster implementation.
• is_causal should be preferred as an argument for causal attention, instead of being given a causal mask
• torch.flex_attention looks like a good option for a custom mask, but maybe not compatible as it just came out with the last release (torch 2.12)
• considering to fuse some of the attention kernels in one, in order to not have duplicated code?

Finally, sorry for the delay, I tried to do this asap. I am motivated to add other flash_attn variants, but without cuda device, I am afraid it is not very reliable for you.

12/6 EDIT: added flash-attention-3 installation description with support

14/6 EDIT:
I have added some tests, and now it should be fully testable on GPUs. Also, it is fully backward compatible as the default mode for attn_implementation is sdpa.

Minor changes that are coming to be pushed—but these changes won't impact the core logic, which is done. So I let this PR as ready to review even if some changes will be needed:

• add benchmark between 'sdpa' and 'flash_attention_3' modes (with pytest)
• merge some functionalities about tests
• maybe test the different models with attn_mode='flash_attention'

NB: one very likely crash cause could be the input tensors that are not contiguous in fast_attn_func.

Also working on another PR from this one (as it will be a big one) for flex_attention. Will be quite big as it would support some kind of attn_mask (actually, it gives as an argument some function that constructs the mask given the position in the matrix). However, it would make a very massive PR if I put it in here. Same for changing the whole tensor coordinate convention— PyTorch's convention (B, H, T, D)— or Flash Attention's convention (B, T, H, D).

Looking forward to your feedback :))

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.