Multi-head causal flash attention support? about flash-attention-jax HOT 8 CLOSED

Is there any plan to implement this feature?

I want to apply it to my custom jax code.

from flash-attention-jax.

^0^ great. Thanks for your support! Take care.

from flash-attention-jax.

I am looking into this code more carefully, and it seems that the unwanted computation (upper triangular region in the causal attention) is not excluded in the computational process. (I don't expect that the compiler also handles this aspect..)

I think it is intended for the easy understanding of flash attention, but it could be 2x faster if the length of query and key is the same and the process is compute bounded.

The issue is closed.

from flash-attention-jax.

yea, it should, as it is agnostic to how many preceding dimensions there are (whether it is batch, heads, etc)

from flash-attention-jax.

oh shoot, i never built it

believe at the time i thought vmap would suffice

from flash-attention-jax.

sure! I can add it tomorrow morning, California time

from flash-attention-jax.

provided I don't drink too much tonight :)

from flash-attention-jax.

@sh0416 ok its done, you can test it by running

from flash_attention_jax import causal_attention, causal_flash_attention, value_and_grad_difference

diff, (dq_diff, dk_diff, dv_diff) = value_and_grad_difference(
    causal_attention,
    causal_flash_attention,
    seed = 42,
    add_key_mask = False
)

print('shows differences between normal and flash attention for output, dq, dk, dv')
print(f'o: {diff}')       # < 1e-4
print(f'dq: {dq_diff}')   # < 1e-6
print(f'dk: {dk_diff}')   # < 1e-6
print(f'dv: {dv_diff}')   # < 1e-6

from flash-attention-jax.