use flash-attn via xformers

[tmm1]

fixes #485 (comment)

[WoosukKwon]

Hi @tmm1 thanks for letting us know the performance issue and submitting the PR.

While using FA2 might improve the performance, we have concerns in using it because it does not support attention bias like ALiBi, V100 GPUs, FP32 data type, and head_size 256 (which is used for GPT-J). So, to use FA2, I believe we should make a fallback option to xformers cutlass backend.

[tmm1]

So, to use FA2, I believe we should make a fallback option to xformers cutlass backend.

thanks @WoosukKwon, I updated the PR to allow xformers to fallback

cc @danthe3rd

[zhaoyang-star]

So, to use FA2, I believe we should make a fallback option to xformers cutlass backend.

thanks @WoosukKwon, I updated the PR to allow xformers to fallback

cc @danthe3rd

Hi @tmm1 , I am very interested in your PR. I see the PR does not allow xformers to fallback.

[danthe3rd]

Hi, xformers maintainer here
The way it's setup in this PR allows xformers to decide which backend to use - this is also what we recommend. It will use by default Flashv2 if available, but will fallback to cutlass if it's not available (eg if using a custom bias, fp32 or v100).

[zhaoyang-star]

Hi @danthe3rd , thanks for your explanation. I just wonder why xformers.ops.memory_efficient_attention_forward does not been used in the decoding stage? The API has used in the prefill stage in vLLM. Hand-written kernel may have lower perf than xformers if we are not an expert of CUDA.

[danthe3rd]

At the moment, memory_efficient_attention is optimized the training mostly - which is similar to the prefilling stage in terms of problem sizes. We have a backend for next token decoding, but it's not fully optimized for all cases, but we're working on it :)

[zhaoyang-star]

At the moment, memory_efficient_attention is optimized the training mostly - which is similar to the prefilling stage in terms of problem sizes. We have a backend for next token decoding, but it's not fully optimized for all cases, but we're working on it :)

Thanks a lot. So it makes sense that most llm inference framworks have hand-written cuda kernel for fused attention impl.

[zhuohan123]

LGTM! Thank you for your contribution!

[KexinFeng]

Hi @tmm1 @WoosukKwon
I have a follow up question to the question

I just wonder why xformers.ops.memory_efficient_attention_forward does not been used in the decoding stage?

Is flash attention (or similar algorithms where softmax is calculated streamingly with fused kernal) also implemented inside vllm.attention_ops.single_query_cached_kv_attention? It looks to me that, in principle, flash attention algorithm is very compatible with paged attention, meaning that the softmax can in principle be computed streamingly and with a fused kernal, in paged storage of tensor, too.

[Lvjinhong]

For allow xformers to pick the best available implementation, I don't quite understand this change, so how should I use flash?