AttentionLib: A Scalable Optimization Framework for Automated Attention Acceleration on FPGA

Zhenyu Liu, Xilang Zhou, Faxian Sun, Jianli Chen, Jun Yu, Kun Wang

Abstract: The self-attention mechanism is a fundamental component within transformer-based models. Nowadays, as the length of sequences processed by large language models (LLMs) continues to increase, the attention mechanism has gradually become a bottleneck in model inference. The LLM inference process can be separated into two phases: prefill and decode. The latter contains memory-intensive attention computation, making FPGA-based accelerators an attractive solution for acceleration. However, designing accelerators tailored for the attention module poses a challenge, requiring substantial manual work. To automate this process and achieve superior acceleration performance, we propose AttentionLib, an MLIR-based framework. AttentionLib automatically performs fusion dataflow optimization for attention computations and generates high-level synthesis code in compliance with hardware constraints. Given the large design space, we provide a design space exploration (DSE) engine to automatically identify optimal fusion dataflows within the specified constraints. Experimental results show that AttentionLib is effective in generating well-suited accelerators for diverse attention computations and achieving superior performance under hardware constraints. Notably, the accelerators generated by AttentionLib exhibit at least a 25.1 × improvement compared to the baselines solely automatically optimized by Vitis HLS. Furthermore, these designs outperform GPUs in decode workloads, showcasing over a 2× speedup for short sequences.