Flash Attention V2

[nivibilla]

https://github.com/Dao-AILab/flash-attention

Flash attention v2 was released claiming 2x speedups. Making an issue to remind myself to have a look at it. And also if anyone else wants to try implement it.

[chenrui17]

I use benchmarks/benchmark_throughput.py to test flash attention V2, but it doesn't seem to have any effect. my test step is like this,

• update xformers to latest version
• modify the line to self.attn_op = xops.fmha.flash.FwOp()
• python3 benchmark_throughput.py --dataset=./ShareGPT_V3_unfiltered_cleaned_split.json --model=/huggingface_data/llama-7b-hf/ --tokenizer=hf-internal-testing/llama-tokenizer --num-prompts=500

test time is like this,

Further analysis of performance, i found that the replaced part (flash attention V2) cost is too small, only at the beginning of the execution, i am confused , for flash attention V2, what can we do for vllm？

[tmm1]

• update xformers to latest version

hi, which version specifically? i don't think flash v2 support has been released yet, so you would have to install from git. also there are still some open PRs to bump xformers to flash-attn v2.0.4 bugfix release (facebookresearch/xformers#816).

[tmm1]

• modify the line to self.attn_op = xops.fmha.flash.FwOp()
• python3 benchmark_throughput.py --dataset=./ShareGPT_V3_unfiltered_cleaned_split.json --model=/huggingface_data/llama-7b-hf/ --tokenizer=hf-internal-testing/llama-tokenizer --num-prompts=500

I tried this as well, and there was no improvement in the benchmarks after switching to flash-attn v2.

I will try to profile the benchmark script.

[Zhuqln]

modify the line to self.attn_op = xops.fmha.flash.FwOp()

i dont think this one really works. because flash-attn's another important feature is to decrease the highly gpu-memory usage in super long-context like more than 5k.
when i set that line and run inference . i dont see any changes on memory usage.

[WoosukKwon]

Hi @nivibilla, thanks for submitting the issue. The latest version of xformers now uses the FlashAttention-V2 algorithm, so vLLM also now takes advantage of it. Please upgrade vLLM to v0.1.4.

@tmm1 @Zhuqln To my understanding, the overall speedup should depend on your workload. At the inference time, FlashAttention is only used for the prompt inputs, and never used for the decoding inputs. For many workloads, the decoding stage takes a majority of the total execution time, so changing to FlashAttention V2 may not give a notable speedup. However, for other workload like text summarization where the prompts are very long, I believe computing attention for the prompt inputs will take a significant portion of the execution time, and thus FlashAttention V2 will have a huge impact on the overall performance.

[nivibilla]

@WoosukKwon thanks for the explanation!

[tmm1]

The latest version of xformers now uses the FlashAttention-V2 algorithm, so vLLM also now takes advantage of it. Please upgrade vLLM to v0.1.4.

Hi, this is inaccurate since the code is still forcing xops.fmha.cutlass.FwOp to be used. If you want to take advantage of FA2, you would need to switch to xops.fmha.flash.FwOp

See benchmark results in facebookresearch/xformers#832

[zhaoyang-star]

Hi @nivibilla, thanks for submitting the issue. The latest version of xformers now uses the FlashAttention-V2 algorithm, so vLLM also now takes advantage of it. Please upgrade vLLM to v0.1.4.

@tmm1 @Zhuqln To my understanding, the overall speedup should depend on your workload. At the inference time, FlashAttention is only used for the prompt inputs, and never used for the decoding inputs. For many workloads, the decoding stage takes a majority of the total execution time, so changing to FlashAttention V2 may not give a notable speedup. However, for other workload like text summarization where the prompts are very long, I believe computing attention for the prompt inputs will take a significant portion of the execution time, and thus FlashAttention V2 will have a huge impact on the overall performance.

Thanks for the details @WoosukKwon . I just have a question. Why FlashAttention could not be used for decoding phase?

[learning-chip]

Why FlashAttention could not be used for decoding phase?

Its tiling strategy is not optimized for Q with seqlen=1 Dao-AILab/flash-attention#427 (comment)