SWAT: Scalable and Efficient Window Attention-based Transformers Acceleration on FPGAs
Zhenyu Bai, Pranav Dangi, Huize Li, Tulika Mitra
TL;DR
SWAT tackles the quadratic bottleneck of self-attention in long-context transformers by adopting sliding-window (static sparse) attention and implementing it on an FPGA with a dataflow-aware, input-stationary design. The approach fuses kernels and arranges computation in a row-major, data-reuse-friendly flow, exploiting the structured sparsity to achieve linear complexity $O(nw)$ while maximizing on-chip reuse. Key contributions include kernel-fusion of $QK$, SoftMax, and $SV$ into a single row-wise kernel, a row-major dataflow that preserves K/V reuse, and an FPGA architecture with parameterizable window size and attention modes, achieving up to 22x latency and 5.7x energy efficiency over a baseline FPGA and about 15x energy efficiency over GPU for long sequences. The work demonstrates practical impact by enabling scalable long-context transformers on FPGA platforms with competitive accuracy and substantial energy savings, and by extending the design to support global and random attention patterns through parameterization.
Abstract
Efficiently supporting long context length is crucial for Transformer models. The quadratic complexity of the self-attention computation plagues traditional Transformers. Sliding window-based static sparse attention mitigates the problem by limiting the attention scope of the input tokens, reducing the theoretical complexity from quadratic to linear. Although the sparsity induced by window attention is highly structured, it does not align perfectly with the microarchitecture of the conventional accelerators, leading to suboptimal implementation. In response, we propose a dataflow-aware FPGA-based accelerator design, SWAT, that efficiently leverages the sparsity to achieve scalable performance for long input. The proposed microarchitecture is based on a design that maximizes data reuse by using a combination of row-wise dataflow, kernel fusion optimization, and an input-stationary design considering the distributed memory and computation resources of FPGA. Consequently, it achieves up to 22$\times$ and 5.7$\times$ improvement in latency and energy efficiency compared to the baseline FPGA-based accelerator and 15$\times$ energy efficiency compared to GPU-based solution.
