Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Level
(2403.04690)Abstract
Neighborhood attention reduces the cost of self attention by restricting each token's attention span to its nearest neighbors. This restriction, parameterized by a window size and dilation factor, draws a spectrum of possible attention patterns between linear projection and self attention. Neighborhood attention, and more generally sliding window attention patterns, have long been bounded by infrastructure, particularly in higher-rank spaces (2-D and 3-D), calling for the development of custom kernels, which have been limited in either functionality, or performance, if not both. In this work, we first show that neighborhood attention can be represented as a batched GEMM problem, similar to standard attention, and implement it for 1-D and 2-D neighborhood attention. These kernels on average provide 895% and 272% improvement in full precision latency compared to existing naive kernels for 1-D and 2-D neighborhood attention respectively. We find certain inherent inefficiencies in all unfused neighborhood attention kernels that bound their performance and lower-precision scalability. We also developed fused neighborhood attention; an adaptation of fused dot-product attention kernels that allow fine-grained control over attention across different spatial axes. Known for reducing the quadratic time complexity of self attention to a linear complexity, neighborhood attention can now enjoy a reduced and constant memory footprint, and record-breaking half precision latency. We observe that our fused kernels successfully circumvent some of the unavoidable inefficiencies in unfused implementations. While our unfused GEMM-based kernels only improve half precision performance compared to naive kernels by an average of 496% and 113% in 1-D and 2-D problems respectively, our fused kernels improve naive kernels by an average of 1607% and 581% in 1-D and 2-D problems respectively.
Overview
-
The paper introduces GEMM-based and Fused CUDA kernels as novel methodologies for implementing neighborhood attention in deep learning models, significantly enhancing performance and functionality.
-
Neighborhood attention is optimized to reduce computational cost by focusing on nearest neighbors, but existing implementations were limited until the introduction of the novel approaches in this study.
-
Through GEMM-based implementation, the study leverages General Matrix-Matrix Multiplication efficiency to improve hardware acceleration and reduce computational overhead in neighborhood attention.
-
The use of fused CUDA kernels for neighborhood attention addresses major bottlenecks in memory and computational efficiency, showing superior performance in benchmarks and practical applications.
Accelerating Neighborhood Attention via GEMM-based and Fused CUDA Kernels
Introduction
Neighborhood Attention has emerged as a pivotal optimization technique for reducing the computational cost of self-attention mechanisms in deep learning models, particularly within the realm of Natural Language Processing and Computer Vision. This technique limits the scope of attention to immediate neighboring tokens, thereby transitioning the complexity from quadratic to linear. Despite its efficiency, implementing neighborhood attention, especially in higher-dimensional spaces, has been challenging due to its inherent dependency on custom CUDA kernels. These kernels often lacked in either performance or functionality, leading to a barrier in its widespread adoption. Addressing this gap, the presented study introduces two novel methodologies for implementing neighborhood attention: GEMM-based and Fused CUDA kernels. These approaches not only offer significant performance enhancements over existing methods but also expand the utility of neighborhood attention across various modalities.
Motivation and Background
Neighborhood attention reduces the computational overhead of traditional self-attention by focusing each token's attentiveness on its closest neighbors. While efficient in theory, its practical application has been hindered by the limitations of existing CUDA kernels, particularly in 2-D and 3-D contexts. The demand for custom, optimized kernels has led to the development of the methodologies outlined in this work. Leveraging the General Matrix-Matrix Multiplication (GEMM) approach and introducing fused CUDA kernels, the study showcases a massive leap in performance and functionality.
Methodological Innovations
GEMM-based Implementation
The paper first introduces a GEMM-based approach to implement neighborhood attention. Recognizing neighborhood attention as a GEMM problem enables leveraging the underlying efficiency of GEMM kernels, thus addressing the primary shortcomings of earlier implementations. By mapping the GEMV problems intrinsic to neighborhood attention to GEMM operations, the approach benefits from improved hardware acceleration and reduced computational overhead. This method demonstrates a significant improvement in latency for both 1-D and 2-D neighborhood attention problems.
Fused CUDA Kernels
Building on the limitations of unfused (BMM-style) implementations, the study introduces fused CUDA kernels for neighborhood attention. These kernels eliminate the need to store attention weights in global memory, addressing a major bottleneck in previous methodologies. The fused approach results in reduced memory footprint and enhanced computational efficiency, particularly notable in half-precision calculations. The adaptability of these kernels across different spatial ranks and their ability to incorporate features like causal masking further solidify their utility.
Experimental Validation
The effectiveness of the proposed GEMM-based and fused CUDA kernels is empirically validated through comprehensive experiments. Benchmarks reveal that these new kernels can significantly outperform existing naive CUDA implementations in terms of latency. The fused kernels, in particular, exhibit superior performance across all tested scenarios, enhancing throughput by up to 97% in certain configurations. The applicability of these methodologies is further demonstrated through implementation in existing models like NAT and DiNAT, where notable improvements in throughput are observed without compromising model accuracy.
Implications and Future Directions
The introduction of GEMM-based and fused CUDA kernels for neighborhood attention holds profound implications for the future of attention-based models. By substantially reducing the computational cost and memory footprint, these methodologies pave the way for more efficient and scalable implementations of attention mechanisms. The observed improvements in throughput and latency not only enhance the performance of existing models but also broaden the horizon for the development of more complex and higher-dimensional attention-based architectures. Looking forward, extending these kernels to support backward passes and integrating additional features will be crucial in maximizing their utility across a wider array of deep learning applications.
Conclusion
The study marks a significant advancement in the implementation of neighborhood attention mechanisms through the introduction of GEMM-based and fused CUDA kernels. These methodologies offer a robust solution to the limitations faced by previous implementations, providing a blend of enhanced performance, reduced memory requirements, and greater adaptability. As attention mechanisms continue to play a central role in deep learning, the innovations presented in this paper are poised to significantly contribute to their evolution and expanded application.
Create an account to read this summary for free: