FARE: A F ine-gr a ined Pipelined Re configurable FlashAttention Kernel

Transformer algorithm-based large machine learning (ML) models are widely used in various application domains [3]. We observe that the attention operation of a transformer algorithm consumes a significant amount of execution time when executed on a GPU, affecting real-time latency for inference. While other major operations in the transformer algorithm (e.g., feed-forward operation) consist of weight stationary matrix multiplications, attention consists of non-linear operations and a sequence of matrix multiplications. This sequence of operations in attention results in memory access bottlenecks in GPUs, directly affecting the latency. While there exists an implementation of an attention kernel to minimize memory access bottlenecks -- FlashAttention [1] -- it lacks reconfigurability. Moreover, GPUs provide limited knobs for customizing the attention datapath. There exist application-specific accelerators for ML models with transformer algorithms, but they too lack reconfigurability [2]. FPGAs, in contrast, enable application-specific dataflows, giving us full control over the design. Therefore, in this work, we propose a fine-grained pipelined reconfigurable flash-attention kernel targeting FPGA platforms. The fine-grained pipelining allows us to execute different stages of the attention mechanism in parallel. It carefully constructs different stages of a flash-attention kernel to reduce resource usage as well as improve the performance. As ML models scale and energy demands increase, this makes FPGA-based attention the sustainable choice. Experimental evaluation on Xilinx Zedboard (Zynq 7000) with different parameters of the attention kernel shows that our proposed design improves the performance by up to 2.9× while reducing the resource usage by up to 16.50× over state-of-the-art [4].