Performance of the C++ interface of flash attention and flash attention v2 in large language model (LLM) inference scenarios. The calculation expression is as follows, where the precision of tensor Q, K, V and O is FP16. Remove redundant code from flash attention that has nothing to do with inference, such as backward, dropout, bf16 and torch dependencies, so you can easily integrate flash attention into LLM inference programs.
O = Softmax(Q * K^T) * V
- GQA/MQA Inference: Group Query Attention / Multi Query Attention Inference
- Hybrid Inference: Hybrid Inference by Prefill and Decoding
- ALiBi Inference: Attention with Linear Biases
- OS: Linux
- Cmake Version: >= 3.12
- GCC Version: >= 5
- CUDA Version: >= 11.4
- Gflags: install on ubuntu as follows
sudo apt-get install libgflags-dev
git clone https://github.com/Bruce-Lee-LY/flash_attention_inference.git
cd flash_attention_inference
./build.sh -a 86 -t Release -b OFF
./build.sh -a 86 -t Debug -b OFF
cd flash_attention_inference
./build.sh -a 80 -t Release -b OFF
./build.sh -a 80 -t Debug -b OFF
./run_sample.sh
Process the data in the log and plot it as a line chart.
cd tools/performance
./performance.sh
- CUDA Version: 11.8
- Head Num: 32
- Head Dim: 128
The performance of both is similar for short sequences and Flash Attention v2 performs well in long sequences. It can increase by about 50%.
- Batch Size: 128
- Seq Q: Seq
- Seq K: Seq
When the Batch is small, the Flash Attention v2 performance is better. When the Batch is large, the performance of the two kernels is comparable.
- Batch Size: Batch
- Seq Q: 128
- Seq K: 128
The performance of both is similar for short sequences and Flash Attention performs well in long sequences.
- Batch Size: 128
- Seq Q: 1
- Seq K: Seq
The Flash Attention performance is better regardless of the size of the Batch.
- Batch Size: Batch
- Seq Q: 1
- Seq K: 128
Regardless of the ratio of Prefill to Decoding, Flash Attention and Flash Attention v2 are similar in performance.
- Batch Size: 100
- Seq Q: 128
- Seq K: 128
- flash attention: v1.0.9
- flash attention v2: v2.1.0
- cutlass: v3.1.0