Decoding Attention is specially optimized for multi head attention (MHA) using CUDA core for the decoding stage of LLM inference. It mainly refers to OpenPPL and Flash Attention, which can solve the problem of low tensor core utilization of Flash Attention in the decoding stage of LLM inference and support more types of attention and kv cache quantization optimization. The calculation expression is as follows, where the precision of tensor Q, K, V and O is FP16 or BF16. In some LLM inference decoding scenarios, the performance of Decoding Attention is better than Flash Decoding (Flash Attention) and FlashInfer. In addition, Decoding Attention also supports variable length, GQA / MQA and ALiBi inference scenarios.
O = Softmax(Q * K^T) * V
- Variable Length: Variable kv length inference
- GQA / MQA: Group query attention / multi query attention inference
- ALiBi: Attention with linear biases inference
- OS: Linux
- Cmake Version: >= 3.16
- GCC Version: >= 5.0
- CUDA Version: >= 11.4
- Others: gflags, ccache, pytest
sudo apt-get install libgflags-dev ccache
pip install pytest
git clone https://github.com/Bruce-Lee-LY/decoding_attention.git
cd decoding_attention
./build_cpp.sh -a 80 -t Release -b OFF
./build_cpp.sh -a 80 -t Debug -b OFF
cd decoding_attention
./build_cpp.sh -a 86 -t Release -b OFF
./build_cpp.sh -a 86 -t Debug -b OFF
./run_cpp.sh
./run_cpp.sh
Process the cpp result in the log and plot it as a line chart.
cd tools/performance/cpp
./performance.sh
cd decoding_attention
./install_python.sh
./run_python.sh
./run_python.sh
Process the python result in the log and plot it as a line chart.
cd tools/performance/python
./performance.sh
- CUDA Version: 12.1
- Head Num: 32
- Head Dim: 128
- Data Type: FP16
The performance of Decoding Attention is better when the sequence length is below 1536, while the performance of Flash Decoding (Flash Attention) and FlashInfer is better when the sequence length is above 1536.
- Batch Size: 1
- Seq Q: 1
- Seq K: Seq Len
Regardless of bacth size, Decoding Attention has better performance than Flash Decoding (Flash Attention) and FlashInfer.
- Batch Size: Batch Size
- Seq Q: 1
- Seq K: 128
- ppl.llm.kernel.cuda
- flash-attention: v2.6.3
- flashinfer: v0.1.6
- Kernel Optimization
- KV Cache Quantization: FP8、Int8、Int4