Future support for attention bias or masking

[subercui]

Hi, I noticed this plan for customized attention bias was recently removed. Do you still have this in plan anytime? I feel Flash-attention is such a great project, and having this feature will make it perfect 😄

40a25c8#diff-b335630551682c19a781afebcf4d07bf978fb1f8ac04c6bf87428ed5106870f5L80

[tridao]

As mentioned in the README, we have an experimental implementation in Triton that support attention
bias (e.g. ALiBi):
https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attn_triton.py

[subercui]

Thank you! I'll give it a try

[robflynnyh]

are there plans to implement gradients for the bias? i.e for learnt attention biases. Or how difficult would this be too implement do you think?

[tridao]

are there plans to implement gradients for the bias? i.e for learnt attention biases. Or how difficult would this be too implement do you think?

I'm not planning to work on it as I don't use attention bias in my work. Implementing it in Triton is probably not hard, all the necessary ingredients are there in the Triton backward pass implementation. I suspect one just has to add some code to save the gradient.

[Zhang-Shubo]

Hi, I find flashattn of triton version is very experimental, only support 64 head dimension size. Does this condition have any changes？