Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Level
Authors: Ali Hassani
Wen-mei Hwu
Humphrey Shi
Venue: NeurIPS 2024 poster (2024)
Keywords: fused attention kernel
neighborhood attention
sliding window attention
Primary area: infrastructure

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 aim to massively improve upon existing infrastructure by providing two new methods for implementing neighborhood attention. 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 runtime compared to existing naive CUDA kernels for 1-D and 2-D neighborhood attention respectively. We find that aside from being heavily bound by memory bandwidth, certain inherent inefficiencies exist in all unfused implementations of neighborhood attention, which in most cases undo their theoretical efficiency gain. Motivated by the progress made into fused dot-product attention kernels, we 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 runtime. We observe that our fused implementation successfully circumvents 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 548% and 193% in 1-D and 2-D problems respectively, our fused kernels improve naive kernels by an average of 1759% and 958% in 1-D and 2-D problems respectively. These improvements translate into up to 104% improvement in inference and 39% improvement in training existing models based on neighborhood attention, and additionally extend its applicability to image and video perception, as well as other modalities. Our work is open-sourced at https://github.com/SHI-Labs/NATTEN/.

Faster Neighborhood Attention: Reducing the O Cost of Self Attention at the Threadblock Level
Ali Hassani', Wen-mei 2w Humphrey Shil.3 SHI Labs Georgia Tech, 'NVIDI,, 'IIUC
Abstract
Neighborhood attention reduces the cost of self attention by restricting each to- ken's 's attention span to its nearest neighbors. This restriction, parameterized by window size and dilation factor, draws spectrum of possible attention pat- terns 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 aim to massively improve upon existing infrastructure by providing two new methods for implementing neighbor- hood attention. We first show that neighborhood attention can be represented as batched GEMM problem, similar 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 runtime compared existing naive CUDA kernels for 1-D and 2-D neighborhood attention respectively. We find that aside from being heavily bound by memory bandwidth, certain inherent inefficiencies exist in all unfused implementations of neighborhood attention, which in most cases undo their theoretical efficiency gain. Motivated by the progress made into fused dot-produce attention kernels, we developed fused neighborhood attention; an adaptation of fused dot-product attention kernels that allow fine-grate ene-grained control over attention across different spatial axes. Known for reducing the quadratic time complexity of self attention to linear complexity, neighborhood attention can now enjoy reduced and constant memory footprint, and record- breaking half precision runtime. We observe that our fused implementation successfully circumvents some of the unavoidable inefficiencies in unfused implementations. While unfused GEMM based kernels only improve half precision performance compared to naive kernels by average of 548% and 193% in 1-D and 2-D problems respectively, our fused kernels improve naive kernels by an average of 1759% and 958% 1-D and 2-D problems respectively. These improvements translate into up 104% improvement in inference and 39% improvement in training existing models based neighborhood attention, and additionally extend its applicability to image and video perception, as well as other modalities. Our work is open-sourced at attps://githin com SHw-saHsSsa/////.eee SHI-Labs/NATTEN
Introduction
Inarguably among the most highly utilized and influential primitives in modern deep learning, attention has long been cited for its complexity and memory footprint, especially when the query and context sets are identical (self attention). For years since its adoption in deep learning [23], the most common implementation of attention was through two batched GEMM (General Matrix-Matre Multiplication) operations, sometimes referred as "BMM-style" attention. This implementation stores attention weights to global memory, which can become bottleneck in both speed and memory
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

1800%
1,739
Nalve
GEMM
1400%
Fused
1200%
Speedup (96)
1,135
1000%
958
844
600%
400%
385
467
200%
193
Baseline
NA-1D
NA-20
MA-30
NA-10
NA-20
NA-30
Forward only (FP16)
Forward backward pass (FP16)
Figure 1: Overview of average improvement in speed on A100 from our proposed imple- mentation. Baseline the set of naive CUDA kemnels introduced in Neighborhood Attention Transformer [9] GEMM-based NA improves 1-D problems by average of 548% (forward pass) and 502% (forward backward), and 2-D problems by an average of 193% (forward pass) and 92% (forward backward). GEMM-based NA does not implement 3-D problems yet. Fused NA boosts performance further and improves 1-D problems by an average 1759% (forward pass) and 844% (forward backward), and 2-D problems by average of 958% (forward pass) and 385% (forward + backward), and 3-D problems by an average of 1135% (forward pass) and 447% (forward + backward).
footprint. As the number of tokens grow, the number of attention weights grow as well, and the problem gets bounded by global memory bandwidth and capacity.
Over the past few years, some works proposed attention implementations in which attention weights are kept in on-chip memory (shared memory register file) instead, until the second matrix mul- tiplication is performed and the resulting attention outputs are written directly to global memory These implementations, known as fused memory-efficient attention, reduce the number of global memory accesses in addition to global memory usage, and successfuly turn dot product attention into compute-boure problem scale. Thanks the first open-source implementation, Flash Attention [6], these fused attention kernels have started replacing the standard BMM-style implementations in many deep learning frameworks and inference engines such as Py Torch [16].
[8 6
Orthogonal to these efforts, many have sought to address the quadratic complexity of self attention, which can become a significant bottleneck vision models more quickly. Neighborhood attention p is one such method in which each query token is restricted to only interact with its nearest neighboring context tokens. In most cases, this pattern creates sliding window pattern, like that the discrete convolution operator heavily employed in vision models. This restriction can similarly be parame- terized by a window size and dilation factor, and reduces the quadratic complexity of self attention down to a linear complexity. This approach is, however, very difficult to implement at the tensor library or deep learning framework level. Tensor views can represent sliding window attention [1)] but not the neighborhood attention pattern. In addition, standard GEMM implementations typically do not support such tensor views in higher-ran//multiddiieeen
spaces (2-D and 3-D) without explicit copying into contiguous tensors, which in practice undoes the theoretical efficiency gain from the reduced attention complexity. As a result, neighborhood attention was proposed along with an extension carrying naive CUDA kernels [9] implementing the operation. While those kernels provided competitive FP32 performance in eager mode inference, and some cases even FP FP16/BFI6 performance, they fall short of general adoption in larger scale experiments. In addition, fused attention implementations, such as Flash Attention, effectively eliminate the (n memory footprint in self attention, while also reducing runtime significantly [6), making subquadratic attention patterns that are only possible to implement "BMM-style" less practical.
In this work, we present two new classes of neighborhood attention kernels: GEMM-based BMM- style kernels (GEMM NA), and fused kernels (Fused NA), which are aimed at providing significantly improved infrastructure for neighborhood attention. We first show that neighborhood attention, and by

---

[3
Causal NA x3]
Cousal NA (7-77
Cousal NA (9-x8)
HA11
M6(T-T)
DINA 13 +3, +2)
DIMA 3-]
DINA 4
DINA 15 +5,2
Figure 2: Illustration of the spectrum of possible attention patterns provided by neighborhood attention. Neighborhood attention only attempts to center the query token (red) within the context window (blue), unlike sliding window attention [19] which forces it. Neighborhood attention with window size is equivalent to linear projection ("no attention"). Neighborhood attention approaches self attention as window size grows, and matches it when equal to input size. Dilation introduces sparse global context, and causal masking prevents interaction between query tokens that have smaller coordinate than neighboring context tokens along the corresponding mode. Window size, dilation, and whether or not causally masked, can be defined per mode/axis.
extension sliding window attention, both of which are GEMV (General Matrix-Vector Multiplication) problems, can be expressed as GEMM problems with space-aware tiling and gather/scatter fusion. This would allow implementing such attention patterns with cefformance-optimizz GEMM primi- tives, which can also utilize specialized hardware components such as NVIDIA's Tensor Cores. We then extend the same logic to fused attention kernels by removing all assumptions that the token mode ("space") is rank-1 1 (single-axis). We write specializations that support spaces, such as 2-D and 3-D. This, in theory, allows any fused attention kernel to be modified to accommodate token spaces of any rank. In addition, part of the logic evaluated at compile time, resulting in less overhead. Finally, the structural simplicity of the resulting fused neighborhood attention kernlss allows for easily adding features such as varying window sizes / dilation values across ranks/axes, causal masking, and more.
Related works
Attention being adopted as deep learning primitive is largely owed to the Transformer architec- ture (23). which despite its original application in machine translation
the position of being the predominant deep learning architecture. Its design and use of the attention operator have been extended many other applications and modalities [15 177. Attention defined an operation between two sets of vectors: a query set and a context set. The two undergo linear projections, with the latter projected into set of key and value pairs. Scaled dot product of query and key vectors. A, is mapped into a probability distribution through the softmax operator, which produces the final attention weights, The output is set of vectors, each derived from the weighted sum all value vectors according to the query vector's attention weights. It can be expressed as follows:
QKT Vd
Attention K, V) = softmax
V.
(1)

---

where Q, K, and V are matrices of query, key, and value vectors as rows respectively, id the scale term, and d is the number of dimensions for which the dot product is computed (number of columns in Q and K). Dot product self attention, or simply, self attention, is special case in which the query and context sets are identical. This means for a set of n input vectors, the attention weights matrix, P. is E |R**", incurring an (n²² time and space complexity. In addition, the softmax term requires column-wise reduction over the attention weight matrix, making kernel fusion more challenging.
Nevertheless, fused attention kernels successfully eliminate the (nn²) global memory footprint, which makes self attention finally bound by compute and not memory bandwidth. These two achievements paved the way for the scaling and application of attention across modalities. To our knowledge, the first open-source implementation of fused multi-headed attention (FMHA) kernel was contributed to the NVIDIA Apex project by Young-Jun Ko, which was primarily used for accelerating inference of Transformer- based language models. As a result of that, it was heavily limited in terms of supported models and problem sizes, as it was performing a full softmax reduction step within the kernel. On the other hand, Milakov and Gimelshein [14 presented technique for computing partial softmax statistics, over which we can perform a final reduction step and derive exact softmax results. This method makes the fusion of attention kernels more practical, because they would no longer be required to compute full row of attention weights before proceeding to perform the second matrix multiplication. Dao et al. 66] presented and spen-sourced Flash Attention, which utilizes online softmax in order create a performant and generic fused attention implementation. Outperforming Mtye- implementations available in both training and inference, Flash Attention was quickly adopted by many frameworks such PyTorch (16), and further improved for the NVIDIA Ampere [5] and Hopper architectures (20).
Parallel to these efforts, many proposed restricted self attention patterns, in which context is restricted to a subset in order to generate fewer attention weights, which in turn reduces the On²²) time and space complexity. Stand-alone self attention (SASA) [9 is a simple 2 2-dimensional sliding window attention pattern, which was shown to effectively replace convolution operations in ResNet [TO] variants. Noting challenges in implementing such patterns without incurring additional overhead from tensor copics and expansion, the authors later moved away from explicit sliding window attention patterns alternatives that relaxed the sliding window movement in HaloNet (22). In addition to these works, sliding window attention patterns in I-dimensional spaces has been explored in language, in works such Sparse Transformers [4) Longformer (2), BigBird (26). and more recently, Mistral [IT]. Neighborhood attention [ 8 is the practice of restricting the context of each token to its nearest neighbors, which in many cases behaves like sliding window pattern, with the exception of corner cases in which the query cannot be centered in a sliding window. Per definitions from SASA [19] and Longformer (2). the sliding context window can go out of bounds, in which case the attention weights corresponding to out-of-bounds tokens are masked. This means tokens close to spatial bounds interact with fewer context tokens. This difference allows neighborhood attention to approach self attention as window size grows. In addition, neighborhood attention defines a dilation factor (8), where the number of such corner cases only increase. Fig. depicts possible attention patterns for a single token under different neighborhood attention parameters. Facing similar implementation challenges as previous works 19], neighborhood attention was implemented with naive CUDA kernels packaged as PyTorch extension, named N NATTEN. While those kernels have accelerated research in this direction, they were simply not intended to fully utilize the underlying hardware. The only exception is the tiled kernels, which are somewhat better optimized, but only apply to a fraction of common use cases, and are not extensible. In addition, with the rise of fused attention kernels such as Flash Attention (6. such implementations which are not performance-optimizee and heavily bandory-anndiithh bound, can hardly compete in terms of performance and memory footprint.
To address these challenges, we present two new implementations and integrate them into N ATT EN, aiming to accelerate all neighborhood attention applications, reduce their existing memory overhead, and extend existing functionality. We first simplify the operations that implement neighborhood attention's 's forward and backward pass into 3 primary operators, and show each can be implemented as batched GEMM kernels with a fused gather/scatter operation. We then point out key limita- tions in unfused neighborhood attention implementations that would prevent them from achieving competitive performance compared to standard BMM-style attention implementations (in more memooy-bandwidthh bound cases.) Motivated by this, and the progress made in fused attention kernels, we propose fused neighborhood attention, which directly extends our batched GEMM
bbapss/giitubb.oo//////W///

---

K
TTT d
Tw x An
Flathoned ble
Select
ThT. x.
A
GEMM output
T.T.
T, Tgp x
global memory
shared memory register
global memory
Figure 3: Illustration of GEMM-based implementation of the 2-D PN operation. Input tensors Q and K are tiled according to their 2-D spatial layout. Q is tiled with a static tile shape, Th x Tw. K tiled with a haloed shape of the Q tile, тV x T which is function of the attention window size (k, X kw)) and the Q tile coordinates. Once tiles are moved into local memory, they are viewed in matrix layout, and a Tw x TH d shaped GEMM is computed
embedding dim). Once done, the tile of dot products with shape TT, x T T is scattered into valid attention weights of shape Th Tw kn
methodology. Since our main objectives are efficiency, Tensor Core utilization, and performance optimization, and our methodology requires significant flexibility in the programming model, we implement both approaches in CUDA C++ using NVIDIA's CUTLASS (21) framework. We show that the batched GEMM kernels can successfully outperform most existing / ATTEN kernels in performance, and that our fused kernels can outperform our batched GEMM kernels while reducing the memory footprint.
Methodology
Herein describe three primary operations (excluding softmax) that required to implement full neighborhood attention forward and backward pass. We then show that each operation can be expressed a batched GEMM problem, long tiling is done according to the underlying spatial rank, and attention weights are scatter/gatt geatter/ However, we find that scatter/gather is a major bottleneck for all unfused implementations of neighborhood attention, limiting their low-precision performance specifically on more recent architectures (Ampere and later.) We then introduce our fused neighborhood attention (FNA) formulation, which builds tiles according to the underlying spatial rank. This approach no longer requires scatter/ paatterr/athrringg of attention weights to/from global memory by definition, and thereby circumvents the aforementioned bottleneck and successfully boosts lower- power-precision performance on modern architectures.
batched GEMM formulation and
3.1 Operators
standard BMM-style attention forward pass (excluding softmax) is comprised of two operations: QK", which produces pre-softmax attention weights (A), and PV. which applies post-softmax attention weights (P) values (V). These operations are different due to layout differences in the matrix multiplications (note that is transposed, V not).
In the case of neighborhood attention, and sliding window attention in general, these will become Matrix-Veetor General Matrix- Vector Multiplication (GEMV) problems. QKT, each query token (vector) is multiplied by its neighboring surrounding key tokens (matrix), and in PV, the set of attention
RQK is layout and PV is TT-layout GEMM in BLAS.
5

---

weights corresponding to each query token (vector) is multiplied by corresponding value tokens (matrix). Given that some of these operations can be reused in the backward pass, we dub the QKT
(PN) and the PV operation (NN). PN can compute the gradient for post-softmax attention weights PVP) when operating on the output gradient instead of Q, and V instead of K. Similarly, NN can compute the gradient for Q (VQ) when operating on the pre-softmax attention gradient (V A) instead of A and K instead V. We define a third operator, which can compute gradients for both K and V: Inverse-Neighborte (IN), This operation is very similar to NN, but differs gather pattern, as well as the number of attention weights. IN may require loading more attention weights for every token, because unlike in self attention, the relationship between query and context tokens in neighborhood attention is not commutative. In other words, query token at coordinate attending to context token coordinate does not imply that query token at coordinate j attends to context token at coordinate i.
operation
BMM-st implementations of standard self attention have clear edge over neighborhood and sliding window attention implementations, because they are GEMM problems and by extension not as bound by memory bandwidth as the latter, all of which are GEMV problems. In addition, GEMV problems cannot effectively utilize matrix multiply and accumulate (MMA) accelerators, such as Tensor Cores. We aim to minimize this issue by formulating all three operators as batched GEMM problems with scatter/gather fusion, in order to better utilize modern hardware accelerators.
3.2 Batched GEMM NA
We transform the aforementioned GEMV problems into batched GEMMs with scatter/gather fusion. At an abstract level, implementations of GEMM- based neighborhood attention predicate the execution of tiled MMAs whether any of the rows in the query tile interact with at least of the rows in the context tile, given the context window size, dilation, and other masking-related parameters. We propose modifying CUTLASS GEMM as follows in order to implement PN, NN, and IN:
1. GEMM tiling is done according to the original multi-dimensional layout of the token mode in QKV. For example, if the attention problem is 1-D, query and context tensors are tiled along the sequence into tiles of size 64, for a 2-D problem, the token mode, which is comprised of height and width, are tiled by a 2-D tiler of the same size, like 8 x 8
2. Predication logic, and global pointer iterators and accessors are modified to iterate according to the original layout in global memory instead of assuming standard rank-2 matrix layout.
3. Attention weights are required to be scattered to and gathered from global memory, which in 16-bit or lower precision cannot be copied asynchronously (with LDGSTS), which breaks pipelining in those kernels on modern architectures. This is because the minimum transaction size for LDGSTS is 32 bits.
We implemented these concepts by extending implicit GEMM (convolution) in CUTLASS (2.X API) into kernels that compute the three neighborhood attention operators in 1-D and 2-D. Fig. shows an illustration of the 2-D GEMM-based PN kernel. The first change is relatively inexpensive, but the second change incurs additional predication and indexing logic that can result in additional overhead and register pressure. The final change is a major bottleneck, and leads to lower-precision kernels (FP16/BF16) providing little no improvement compared to their full precision (FP32/TF32) counterparts, NN and IN suffer from this issue more significantly. because gathering attention weights (LDG) breaks pipelined kernels on Ampere, since they load GEMM operands asynchronously (LDGSTS), which has minimum transaction size of 32 bits. This forces our FP GEMM- based kernels to fall back to global loads (LDG), which significantly impacts achievable runtime. To our knowledge, this issue is unavoidable in most cases, and will continue to be bottleneck as long as attention weights are stored in global memory.
3.3 Fused NA
We extend our methodology for implementing neighborhood attention operators using batched GEMM kernels to fused attention kernels like Flash Attention (6). This is not only motivated by the potential to reduce runtime and memory footprint, and potentially making neighborhood attention actually bound by compute, but also to circumvent the bottleneck in the batched GEMM and naive kernels: scatter/gathering attention weights to/from global memory. Since attention weights are only computed

---

K/V
Output
be computed within Hattened sub-tle
masking
weights
Flartened
Flattered
global memory
shared memory register file
global memory
Figure simplified illustration of fused neighborhood attention. Q and KV tensors are tiled according to their spatial layout (I-D, 2-D, 3-D), with the latter haloed to include the entire neighborhood corresponding queries in the query tile. Resulting attention weights from first GEMM are masked according to neighborhood attention parameters, before undergoing online softmax scaling, and going through the second GEMM with the corresponding value sub-tile.
at the threadblock level and never fully stored in global memory in fused kernels, the bottleneck will simply cease exist. We started off with xFormers FMHA [13], fused multi-heades attention kernel based on the CUTLASS 2.X API, which can target architectures even older than Ampere (Maxwell, SM50; Volta, SM70; and Turing, SM75.) By carefully applying space-aware tiling, neighborhood attention masking, and software predication for multi-dimensional tensor layouts, successfully implemented neighborhood attention for 1-D, 2-D, and 3-D problems. Fig. 4 presents an overview of how our fused kernels function when dealing with multi-dimensional (multi-axis) data.
methodology for
3.4 Dilation and causal masking
Our methodology allows for dilation support trivially, through simple partitioning and slicing ahead of time. A dilated neighborhood attention problem can be mapped to a set of non-dilated neighborhood attention problems the same kernel call, simply by issuing more CTAs in the grid. We additionally define and implement causal neighborhood attention into fused kernel, which can be crucial certain applications where only one spatial dimension requires causal masking (i.e. video embeddings may benefit from causally masked attention across the time axis and standard attention across height and width axes, which would be an exact 3-D tento-l attention module.)
non-overlapping tiles of the input. All sub- sub-problems
3.5 Notes arithmetic intensity
Arithmetic intensity the ratio of floating point operations over bytes of memory transactions, as defined by the Roofline model 25
Nops
Arithmetic Intensity =
Nbytes
Arithmetic intensity is typically used to determine whether an implemeetatiio/algoriim is bound by memory bandwidth or computational capacity, on a given problem size and hardware. Let's consider a simplified representation of self attention, where we only look at pure matrix multiplication FLOPPs and bytes. Self attention comprised of two back-to-back BMMs, which would be 2h/nnd FLOPss for each of the BMMs, where b, h, n, and d denote batch size, number of attention heads, sequence

---

length, and per-head dimension respectively. In total, that would be 4bhn'd FLOPs. In the unfused implementation, tensors with size bhnd (Q, K, V and output) are accessed in global memory, along with one intermediary tensor with size bhn² (attention weights or P), which is accessed twice. In total, that is bhmd + x bhn²) x &dtype bytes, where Sdtype is the byte size of the tensor element type. When implemented with fused attention, however, the number of bytes accessed for matrix multiplication from global memory is reduced to only reads and writes for Q. K, V, and attention outputs, or 4 x blndd X ddtype.
Unfused implementations of attention are typically memory memory-bandwidthe wbw.dwddth-bbummi at scale, given that their arithmetic intensity approaches a constant value as sequence length grows. If we take the limit of their intensity according to the aforementioned approximation of FLOPs and transaction bytes, as + oo with everything else as constants, we see that:
4bhm?d
2nd
2d
lim
lim
(3)
10-00 (4bhnd + 22hn 2bhn2)s dttppe n-000 (2d + &
8 ddtype
Fused attention therefore solves key problem here, by reducing the number of memory transactions from O(n²) to (n), which means as sequence length grows, the limit of arithmetic intensity in fused attention does not converge, and it will therefore only become bound by computational capacity. This means that optimal fused attention kernels can almost fully utilize the underlying computational power of modern GPUs, and is the reason behind FP8 attention kernels for the Hopper architecture exceeding the POtaFLOPPs threshold [3] 20).
A natural question to ask is what happens to local attention patterns such as neighborhood attention, which promised to deliver more efficiency. fused implementations of neighborhood attention (i.e. our proposed FNA). we can look at the growth of arithmetic intensity similar to self attention. If we consider the FLOPs for neighborhood attention to be 4bhnLd, where É is the size of the attention window. and that the number of global memory transaction bytes is the same as fused self attention (worst case). 4 x bhnd x Sdtype, then we see that as n - oc, we converge towards a constant again, therefore making neighborhood attention more memory-bandwidee hounn as alone scales:
4bhmnd
lim 000
4bhndsatgpe
8dtype
However, the constant here is a function of É, the size of our attention window, which means that as we scale the sequence length or feature map size, attention window size will determine whether or not the problem is bound by memory bandwidth or computational power. Since smaller window sizes are closer to linear projections, and larger window sizes are closer to self attention, the fact that neighborhood attention can be bound by compute memory bandwidth depending on window size is not a surprise. Therefore, it is highly recommended to choose neighborhood attention window sizes according to the input size and even hardware maximize efficiency gain.
3.6 Limitations
Our formulation of GEMM-based and fused neighborhood attention kernels poses a critical question: how much overhead can one expect from switching from standard self attention kernel to neighbor- hood attention? As pointed out in Sec. 3.2 our GEMM-based kernels suffer from a major bottleneck, especially in pweer-precsioon, which stems from scatter/gathering of attention weights. We consider this to be an unavoidable issue in unfused implementations of neighborhood and sliding window attention. Unsurprisingly, our proposed changes to fused implementations are also not free. Changes that we find unavoidable, which in some cases can cause our fused kernels to incur higher runtime than the self attention baseline (xFormers FMHA) are the following (ordered by most significant to least significant):
1. Kernels specialized for 2-D and 3-D problems are no longer GEMMs, they are General Tensor- contractions (GETTs)! Similar to convolution, if the input layout is multi- dimensional, then the GEMM is converted to special case of GETT. On older GPU architectures, this requires more complicated software predication, which will incur more instructions and heavier register usage, whereas on modern architectures like Hopper, the Tensor Memory Accelerator (TMA) can easily provide hardware predication. Our software predication logic is similar to standard practice for such cases in CUTLASS 2.X GEMMs, and similarly less performant than predication for contiguous matrix layouts. We find this to
8

---

Table I: FPI6 forward pass benchmark overview. We benchmark naive neighborhood attention ker- nels against our proposed GEMM and fused kernels in half precision, over large set of problem sizes varying in batch size, spatial size, number of atten- tion heads, and dimensions per head, and over dif- ferent window sizes and dilation values. For every problem size, we also benchmarked self attention running with the xFormers FMHA (our baseline) and Flash Attention V2.
Table 2: FP32 forward pass benchmark overview. We benchmark naive neighborhood attention kernels against our proposed GEMM and fused kernels in full precision, over a large set of problem sizes varying batch size, spa- tial size, number of attention heads, and dimen- sions per head, and over different window sizes and dilation values. For every problem size, we also benchmarked self attention running with the xFormers FMHA (our baseline).
NA Kernel
of problems matched or outperformed
NA Kernel % of problems matched or outperformed
neighborhood atin GEMM Fused
self atin
neighborhood attn GEMM
self attn FMHA
Naive
FMHA
FAv2
Nalve
Fused
I dimensionad neighhorhood attention
I-dimensional neighborhood attention
Naive GEMM
0.0 % 0.0 %
21,8
8.8 %
Naive GEMM Fused
0.0 37.7 %
34.6 %
98.7 % 100.0 %
72.0%
54.2%
99.9 % 100.0 %
98.4 %
Fused
100.0 %
100.0 %
98.2
64.8% %
99.9 %
2-dimensional neighhorhood attention
2-dimensional neighborhood attention
Naive GEMM
16.4% %
0.0 %
32.9 %
15.8 %
Naive GEMM Fused
%
$2.0 %
84.0 % 100.0 %
0.0 %
59.3 %
29.8 %
89.5 % 96.0 %
28.1 1 %
92.4 %
Fused
98.6
92.4
74.0% %
99.3 %
3 dimensional neighborhood attention
3-dimensional neighborhood attention
Naive Fused
0.0 %
43,5 %
20.2 %
Naive Fused
0.0
61.1 %
100.0 %
97.3 87.0
100.0 %
98.6 %
be the most significant contributor to additional runtime in our fused kernels, when compared to the baseline fused self attention kernel, FMHA. However, FNA is perfectly capable of hiding this additional overhead in many cases, and only falls behind in cases close self attention (window size is approximately the same as input size.)
2. The attention masking logie, which depends on corresponding query and context token coordinates, original layout, and window size, introduces additional indexing logic in order to map linear indices to coordinates (unlike in 1-D problems where the mapping is the identity function), and it gets more complicated with more dimensions. This, along with additional statements in the masking condition, contributes to runtime, and is expected to worsen with more dimensions. Together, these contribute to more serious register spilling than the original 1-D kernel.
Despite these issues, we find that our fused kernels can still match or outperform our self attention baseline in approximately 100% of I-D, 98.6% of 2-D, and 97.3% of 3-D problem sizes that we benchmarked.
4 Experiments
We evaluate the performance of our proposed methods by measuring their runtime against existing kernels in A ATTEN. Most use cases in N ATTEN target naive CUDA kernels, with the exception of 2-D neighborhood attention with 32-dimensional attention heads. WATTEN implements tiled kernels for those cases for up to and including window size 13 x 13, and only for the QK operation. However, we treat all kernels in N AATTEN as our baseline, and will refer to them as naive kernels. We use fixed set of problem sizes that vary in batch size. spatial size, number of attention heads, and dimensions per attention head, and run them through every implementation on an NVIDIA A100 GPU and measure their runtime using CUDA events. We iterate through multiple neighborhood attention window sizes and dilation values for every problem size. A summary of these benchmarks presented in Tab. (FP16) and Tab. (FP32). We find that our GEMM- based kernels can improve or match the naive runtime in approximately 99% of 1-D problems (of 6150), and 84% of 2-D problems (of 5676) in half precision, and approximately 100% of the 1-D problems and 96% of the 2-D problems in full precision. Note that over 40% of the 2-D problems target tiled kernels in N ATT E.N, which we find can sometimes outperform our GEMM-based kernels. Another point of disadvantage in the FP16/BF16 variants of our GEMM-based kernels is using LDGs in pipelined

---

kernels, noted in Sec. 3.2 On the other hand, our fused kernels improve or match the naive runtime in approximately 100% of both 1-D (of 6150) and 3-D problems (of 2448) in both half precision and full precision, an 100% of 2-D problems in half precision, while only improving approximately 96% of 2-D problems full precision. We also find that our fused kernels match or outperform our GEMM kernels in 100% of both 1-D and 2-D problems in half precision, while only doing so in approximately 65% of I-D problems and 74% of 2-D problems full precision, which is not very surprising given that full precision is typically more seanory-bandwidth-BBe
In both Tab. and Tab. 2 we also inspect the percentage problem sizes which using our fused neighborhood attention kernel is outperformed by the FMHA kernel. This is only inspect additional overhead caused by our implementation, which we expect to be more noticeable in 2-D and 3-D problems. Some of the overhead may be avoidable, but our takeaway is that it is unlikely to be fully avoidable, as pointed out in Sec. 3.6
We further present breakdown of our benchmarks in Tab. where we report the average, minimum, and maximum improvement observed from switching from naive GEMM-based, naive to fused, and GEMM-based to fused kernels. GEMM-based kernels exhibit strong performance compared to both naive and fused kernels in full precision, where fused kernels only have a very minimal edge over unfused. GEMM-based kernels also outperform naive kernels in half precision, especially in cases where tiled kernels are not available. While the tiled kernels are sometimes the better choice, we note that they simply cannot generalize to all problem sizes GEMM-based kernels can, nor are they easily extensible.
Table 3: Forward pass benchmark breakdown. Both GEMM-based and fused NA improve the baseline naive kernels on average. However, there exist cases which naive kernels may be preferable to GEMM- in both FP16 and FP32, but naive is rarely a good choice in half precision where both naive and GEMM are more memory bandwidth bound than fused.
Dim
GEMM over naive Min
Fused over naive
Fused over GEMM
Average
Max
Average
Min
Max
Average
Min
Max
FP16
1-D 2-D 3-D
t 548 % t 193 %
-53 %
1 3025 % 1 862 %
t 1759 %
60 %
t 11885 %
f 180 % t 257 %
t 71
1 466 %
4 -57
+ 958 %
%
1 7169 %
1 38
1 1199
t 1135
t 118
1 5497 %
FP32
1-D
t 874 %
-31 %
1 3565 % 1 1933 %
978 %
13 % 1
1 4419 %
1 17 % 1 %
J -54 % 4 -53
1 133 % 1 451 %
2-D
t 386
4 -43
564 %
-30
1. 4:433 %
3-D
T 712 %
25 %
1 3029 %
5
Future work & Conclusion
In this work, we formulated the neighborhood attention problem, and by extension multi-dimensional sliding window attention, which are inherently GEMV problems, GEMM/GETT problems. Through this finding, we implemented extensible GEMM-based and fused CUDA kemels that implement neighborhood attention, which can significantly improve upon existing kernels in the WATTEN project. These kernels will not only speed up previously-propon proposed models based on neighborhood attention, but can also significantly enhance ongoing research efforts in this direction. In addition, our fused kemels are the most flexible in terms of parameterization, by supporting varying window sizes, dilation factors, and causal masking across different axes, which enable unique applications such as 3-D spatio-temporal attention with causal masking across time. They also enjoy reduced memory footprint, and can avoid being bound by memory bandwidth scale.
Future directions in this area include but are not limited to: support for Context Parallelism (CP), implementations using more efficient predication (i.e. with the Hopper TMA), extension to more modern architectures (warp-specialized kernels in Hopper and Blackwell), extension to other AI accelerators, and better auto-tuning (or alternatives involving graph compilation).
We've ve shown that multi-dimensional local attention can indeed serve as solutions for scaling future large-scale long-context architectures, when provided with suitable software infrastructure. We hope that this inspires more research into multi-dimennsonal attention, as deep learning systems continue to grow larger in both model and input size.
10

---

Acknowledgements. We would like to thank NVIDIA and members of the CUTLASS project, in particular Haicheng Wu, for his valuable feedback and comments which led to the creation of GEMM-baser NA. We also thank Meta xFormers team for developing FMHA, which is what our fused neighborhood attention kernels are based on. A. H. thanks Michael Isaev, Aditya Kane, and Kai Wang for their feedback on the paper. A. H. also thanks Bing Xu, Hao Lu, Michael lovine, and Terry Chen for the invaluable learning experience while interning at HippoML, which helped accelerate the timeline of this project. This research was supported in part by National Science Foundation under Award #2427478 - CAREER Program, and by National Science Foundation and the Institute of Education Sciences, U.S. Department of Education under Award #2229873 National AI Institute for Exceptional Education. This project was also partially supported by cyberinfrastructure resources and services provided by the Partnership for Advanced Computing Environment (PACE) at the Georgia Institute of Technology, Atlanta, Georgia, USA.
References
[1] Anurag Arnah, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučié, and Cordelia Schmid. Vivit: A video vision transformer. In Proceedings of the IEEF/CVF International Conference Computer Vision (ICCV), 2021.
[2] Iz Beltagy, Matthew E Peters. and Arman Cohan.
Longformer: The long-document transformer. arXiv
preprint arXiv:2004.05 2020.
[3] HippoMI. Blog. Petaflops inference era: 1 pflops attention, and preliminary end-to-end results. https 2024.
//medium. http:////221682e2eeeeem
[4] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: /904. /0509, 2019.
[5] Tri Dao. Flashattention-2:: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2023.
[6] Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[7] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbom, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2020.
[8] Ali Hassani and Humphrey Shi. Dilated neighborhood attention transformer. arXiv preprint
arXiv:2209. 15001, 2022.
Ali Hassani, Steven Walton, Jiachen Li, Shen Li, and Humphrey Shi. Neighborhood attention transformer.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), 2023.
[10] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2016. [11] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel. Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.00 06825, 2023.
[12] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep convolutional neural networks. Advances Neural Information Processing Systems (NeuriPS), 2012.
[13] Benjamin Lefaudeux, Francisco Massa, Diana Liskovich. Wenhan Xiong. Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xformers: modular and hackable transformer modelling library. com/ facebookreaearchh x xfomerrs] 2022.
https://gi
[14] Maxim Milakov and Natalia Ginelshein. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.0288 0808.7887 2018
[15] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning (ICMI.) 2018.
[16] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan. Trevor Killeen, et al. Pytorch: An imperative style, high- gigg-performance deep learning library. In Advances in Neural Information Processing Systems (NeuriPS), 2019.
[17] William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint
mrXiv:222..09744 2022.
[18] Markus N Rabe and Charles Staats. Self-attention does not need O(m²) memory. arXiv preprint
2 2021.
[19] Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jon Shlens. Stand-alone self attentioo in vision models. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
11

---

[20]
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani. and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407 0006008,8, 2024.
[21] Vijay Thakkar, Pradeep Ramani, Cris Cecka, Aniket Shivam, Honghao Lu, Ethan Yan, Jack Kosaian, Mark Hoemmen, Haicheng Wu, Andrew Kerr, et al. Cutlass, 2023.
[22]
Ashish Vaswani, Prajit Ramachandran, Aravind Srinivas, Niki Parmar, Blake Hechtman, and Jonathon Shlens. Scaling local self-attention for parameter efficient visual backbones. In Proceedings of the
IEEE/CVF Conference Computer Vision and Puttern Recognition (CVPR). 2021.
[23]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Prucessing Systems (NeurIPS), 2017.
[24] Steven Walton, Ali Hassani. Xingqian Xu, Zhangyang Wang, and Humphrey Shi. Stylenat: Giving each
head a new perspective. arXiv preprint arXiv:2211.05777 2022.
[25] Samuel Williams, Andrew Waterman, and David Patterson. Roofline: insightful visual performance model multicore architectures. Communications of ACM, 52(4)x65-76, 2009.
[26] Manzil Zaheer. Guru Guruganesh, Kumar Avinava Dubey. Joshua Ainslie. Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, al. Big bird: Transformers for longer sequences. In Advances Neural Information Processing Systems (NeurlPS). 2020.
12

---

A Auto-tuner
GEMM kernels are, among other settings, parameterized by their tiling shape. Multi-dimensionn variants (2-D and 3-D neighborhood attention) can also be parameterized by their fine-grained tile sizes, introduced by our formulation. As mentioned earlier, a GEMM with row tile size 64 can be reinterpreted as number of 2-D and 3-D tiles (i.e. X x y for all positive integers x and where ry 64, and y X z for all positive integers x, y, and z where zyz 64.) As a result, selecting tiling sizes based on factors such as problem size, hardware, and environment can further decrease achievable runtime. We therefore implement a very simple auto-tuning method as a proof of concept. Auto-tuning creates and maintains a cache for the lifetime of the application, which maps problems (defined by problem size, data type, and other such factors) to a tiling configuration. On a cache miss, the problem is benchmarked over a set of tiling configurations, and the best configuration gets cached.
While the auto-tuner noticeably improve performance even further, we note that it is presently limited in the following:
1. Distributed training. auto-tuner context is limited single process, meaning jobs involv- ing distributed training or inference will run the auto-tuner separately in each individual process. Aside from the possibility of different processes choosing different settings, which can slightly impact numerical determinism, this behavior is counter-intuitive. A more advanced auto-tuner would distribute possible settings over available processes and reduce auto-tuning time in the process, and guarantee the same settings across processes.
2. Vast search space. there exist in the order of thousands of valid settings for any given problem size, and searching over all of them is intractable. Our solution so far has been to generate far fewer possible settings, and even reduce the number of settings further by introducing a "thorough mode", which is disabled by default, but when enabled, will allow users to search over more settings and potentially gain more in speed. This issue is common problem in modern computational packages, and we hope to alleviate it by common practices such as reducing benchmark time, distributing the process, caching to disk, lazy benchmarking, and approximate performance models.
B Additional experiments
Herein we present some additional performance metrics from our GEMM-based and fused kernels compared against the baseline.
In Tab. we break down expected performance improvements at the operation level from single for- ward and backward pass. Both our GEMM-based and fused kernels provide significant improvement on average over the baseline. while there still exist cases where naive could potentially perform better, especially compared to our GEMM- -ased kernels. As pointed out in Sec. 3.2) the scatter and gather operation in our GEMM kernels are a sigmificant bottleneck, especially lower-precision and in NN and IN operations. In lower precision, NN and IN, which account for 75% of the backward pass operations (excluding softmax) will fail to hide their prefetch runtime from global reads, which are not asynchronous, and this will essentially impact the backward pass more than it does the forward pass. This issue, however, is limited to our unfused variant, and our fused kernels maintain their superior performance levels, offering up to an order of magnitude improvement in all variants (1-D, 2-D, and 3-D).
Table 4: Forward backward pass benchmark breakdown. Improvements over naive, while not as significant as in the forward pass, are still significant. We report benchmark the full forward and backward pass in half precision only, because most training is done in lower precision.
Dim
GEMM over naive
Fused over naive
Fused over GEMM
Average
Min
Max
Average
Min
Max
Average
Min
Max
1-D
t 502 % t 92
1 -48 % 1 -70 %
1 3017 %
1 844 %
4 -20 %
t 7605 %
1 57 % 150 %
-50 % $ .49 %
t 222 %
2-D
1 474 %
t 385 %
-61 %
1 3723 %
1 1855 %
3-D
T 447 %
4 -45 %
1 2824 %
In addition to our operation-level benchmarks, we also evaluate the effect of our proposed methodol- ogy existing models that use neighborhood attention as primitive, NAT [9] and DiNAT [8]. We
13

---

benchmark the throughput from all variants according to ImageNet-IK [2] specifications, and report FP16 and FP32 measurements in Tab. annd Tab. rrespecciivv We also benchmark a style-based generative adversarial (GAN) model based neighborhood attention, StyleNAT (24) and report performance improvements in Tab. 7 We find that least in problem sizes that the ImageNet classification models NAT and DINAT typically require, which are typically smaller in spatial size and window size, and larger batch size, our GEMM-based approach fails to improve the baseline in half precision, and only minimally improves it in full precision. Our fused kernels the other hand never fail to improve upon the baseline, but they only provide significant improvement in half precision, and cases that use dilation frequently (DiNAT [8] variants). Improvements in the generative model, StyleGAN [24], are only observed in full precision (half precision is not recommended in this application), where again we find that both GEMM based and fused kernels improve inference speed compared to existing naive kernels, with our fused kernels having a much more noticeable edge.
Table 5: Model-level throughput changes when using our proposed GEMM-based and fused kernels in ImageNet classification. Hierarchical vision transformers NAT and DiNAT can see between 26% to 104% improvement in FP16 throughput on an A100 (batch size 128) with our proposed fused kernel. Suffering from the memory alignment issue, half precision GEMM kernels usually result in a much smaller improvement over naive kernels, particularly the tiled variants. The same measurements with FP32 precision presented in Tab.
Model
# of FLOPs
Throughput
Top-1 Accuracy (%)
Params
Naive
GEMM
Fused
(M)
(G)
(imgs/sec)
NAT-M DINAT-M
20
2.7
2975
2660 4 -111 %
3742 ( 1
26 % )
81.8
20
2.7
2672
2548 ( + -5 %
) 3930 f
47 % )
81.8
DINAT, -T
28
4.5
2850
2504
4-12 %
3847 ( 1 35 %
)
81.8
NAT-T
28
4.3
2167
1939
4-11 %
2772 ( t 28 %
83.2
DINAT-T
28
4.3
1910
1845 ( 4 -3 %
) 2909 ( 1 52 %
82.7
DINAT -S
50
8.7
1800
1571
4-13 %
2445 ( 1 36 %
83.5
NAT-S
7.8
1457
1309
4-10 %
) 1879 ( 1 29 %
83.7
DINAT-S
51
7.8
1360
1313
-3
) 2145
1
58 %
83.8
DINAT -B
88
15.4
1351
1178
4-13
1837 ( 1 36 %
83.8
NAT-B
90
13.7
1110
997
( 4-10 %
1448 ( T 30 %
84.3
DINAT-B
90
13.7
982
(
1517 ( 1 54 %
84.4
DINAT, -L
197
34.5
846
744 ( 4-12 %
1119 ( 1 32 %
)
86.5
DINAT-L
200
30.6
669
647 ( ( 4 %
)
1042 ( 1 56 %
86.6
DiNAT -L (384 384)
197
101.5
295
239 ( 4-19 %
391 ( 1 33 % )
87.4
DINAT-L (384x 384)
200
92.4
153
134
4-12
312 ( 104 %
87.5
Finally, we also attempted estimate improvements in training time compared to baseline. As suggested by our earlier findings regarding the limit of our GEMM-based implementation in the backward pass, we do not see any improvement training time compared to the naive baseline. However, we find that our fused kernels deliver on the promise of improved half precision training time. We present our estimates Tab. 8 which are based on measurements from training NAT 9 and DiNAT [8] variants according to their original specifications. We ran each model for warmup epoch, and benchmark epoch, the average throughput of which used estimate training time for 300 epochs.
14

---

Table 6: Model-level throughput changes when using our proposed GEMM-based and fused kernels in ImageNet classification (full precision). While fused attention kernels not expected to have as large of an edge over BMM-stylc attention kernels FP32, our fused kernels still happen outperform naive kernels in full precision. It is also visible that GEMM kemelss outperform naive kernels when we eliminate the memory alignment issue. That said, our FP32 GEMM kernels still impose maximum alignment of 1 element the attention weights tensor, which limits its ability to compete with other B.Me-style attention kernels.
Model
# of FLOPs
Throughput
Top-1 Accuracy (%)
Params
Naive
GEMM
Fused
(M)
(G)
(imgs/sec)
NAT-M
20
2.7
2416
2481 ( 1 3 %
2658
f 10 %
)
81.8
DiNAT-M
20
2.7
2217
2364
f 7 %
2905
f 31 %
81.8
DiNAT,-T
28
4.5
2270
2255
4-1 %
2771
22 %
81.8
NAT-T
28
4.3
1739
1802
4 %
1942
12 %
83.2
DiNAT-T
28
4.3
1591
1706
f 7 %
2123
f 33
%
82.7
DiNAT, --S
50
8.7
1403
1393
4-1 %
1717
f 22 %
83.5
NAT-S
51
7.8
1160
1199
%
)
1293
11 %
83.7
DiNAT-S
51
7.8
1102
1183
f 7.%
1490
1 35 %
83.8
DiNAT,
88
15.4
1020
1009
4-1 %
1240
1 22 % )
83.8
NAT-B
90
13.7
867
897 ( ( %
966
1 11 %
84.3
DiNAT-B
90
13.7
795
851 ( 717%
1059
33 %
84.4
DiNAT,-L -L
197
34.5
609
601 ( 4-1 )
721
18 %
86.5
DiNAT-L
200
30.6
506
540 ( 7%%
669
1 32 %
86.6
DiNAT, -I. -I,3884 384)
197
101.5
211
193 ( 4.
245
16
%
87.4
DiNAT-L (384 384)
200
92.4
116
115 ( 4
179
54 %
87.5
Table 7: Model-level throughput changes when using our proposed GEMM-based and fused kernels in style-based image generation. We benchmark StyleNAT 224, a style-based generative adversarial model based on neighborhood attention under different kernels. We experimented with different batch sizes order to achieve peak performance, and settled for 64 for the 256 256 variant, and 8 for the 1024 x 1024. StyleNAT does not recommend lower-precision, therefore these measurements are only done in FP32.
Dataset
# of Params
Throughput GEMM (imgs/sec)
FID
Naive
Fused
FFHQ (256 256)
48,9 M
36.7
40.6 t 11 %
45.5
1 24 % )
2.05
FFHQ (1024 x 1024)
49.4 M
8.2
8.5 ( f 3 %
11.5
40 %
4.17
15

---

Table 8: Training time improvement when using fused neighborhood attention kernels. We each of the classification models based on neighborhood attention for one warmup epoch and one benchmark epoch, all with half precision (the typical training scenario), and report the estimated training time. Note that these numbers exclude positional biases, as our fused backward kernel does not support it.
Model
# of
FLOPs
Training time estimate GEMM
Params
Naive
Fused
(M)
(G)
(hours)
NAT-M DiNAT-M
20
2.7
19.4
20.4
4 -5 %
16.6 ( t 17 %
20
2.7
20.4
21.2
$ -4 %
17.4 (
f 17
DiNAT -T
28
4.5
21.1
22.0 (
4 %
17.4 ( 1 21 % 10 %
)
NAT-T
28
4.3
26.5
28.2
$ -6 %
24.0 ( 1
DINAT-T
28
4.3
27.4
28.5
4 -4 %
21.9 ( ( %
DiNAT, -S
50
8.7
33.3
33.2 ( 0 %
25.1 ( 1 %
NAT-S
7.8
39.2
41.8
466 %
33.7 ( 16 % )
DINAT-S
51
7.8
38.0
40.1
4 -5 %
30.8 ( 2. %
DiNAT, -B
88
15.4
45.4
46.1
4 %
32.6 ( 39 %
NAT-B
90
13.7
51.1
54.6
J. -6 %
47.7 (1 7 7
DINAT-B
13.7
54.4
56.0
4-3 %
41.0
13 %
16

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes].
Justification: All claims made are with regard to earlier implementations of neighborhood attention, and how our approach aims to accelerate it. Claims are also based on experiments presented in the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes].
Justification: Limitations in terms of performance upper bounds and utilization of newer hardware accelerators are discussed in full.
Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: [NA]
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: any reported performance numbers can be reproduced on the same hardware and with the specified versions of software, alongside our our released software distribution.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: All of our implementations open sourced. Release of anonymized code is not possible.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification: [NA]
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: [NA]
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [NA]
17

---

Justification: [NA]
9. Code Of Ethies
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpgg/tttts////aapooo..ooooooot
Answer: [Yes].
Justification:
Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The work introduces set of implementations, and does not of itself pose any societal impacts positive or negative.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes].
Justification: all prior arts have been cited, and packages open source projects and software used credited and if possible cited.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes].
Justification: The open source project associated is well documented for all uses.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: [NA]
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
18


Reviews:
[{"review_id": "uOpU0vimgO", "reviewer": "Reviewer_RCPh", "summary": "The paper introduce a method to improve the performance of neighborhood attention mechanisms. The authors present two new implementations: GEMM-based kernels and fused kernels. These implementations aim to reduce the latency and memory footprint of neighborhood attention in deep learning models, particularly in higher-dimensional spaces. The proposed methods show significant speedup in both full and half precision compared to existing naive CUDA kernels.", "questions": "1. What is the implications of the fp16 acceleration to more advanced low precision format such as fp8 or even fp6?\n2. While the paper provides significant speedup of neighborhood attention, more discussions on the overall bottleneck of it would be interesting, whether the improvements on the speed can be further translated to better accuracy.", "limitations": "No significant negative societal impacts.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The proposed GEMM-based and fused kernels significantly reduces latency, average improvements of 895% and 272% in full precision latency for 1-D and 2-D neighborhood attention\n2. The implementations can also reduce the memory consumption of neighborhood attention.\n3. The methods are designed to work efficiently across different spatial ranks (1-D, 2-D, and 3-D), which makes it more generally applicable.", "weaknesses": "1. The performance improvements in half precision are limited due to the inefficiencies in gather/scatter operations。\n2. The performance gains are more significant on newer architectures like Ampere and Hopper, limiting the generalizability of the results to other hardware setups.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your feedback.\nTo respond to your questions:\n\n1. To clarify, our GEMM-based approach's FP16 performance is significantly affected by the gathering and scatter of attention weights, not Fused Neighborhood Attention. FNA does not store or load attention weights to global memory, and hence has no need for a gather scatter on those, and as a result will not run into the \"under-alignment\" issues on modern hardware. Because of this, we spent more time on the fused approach which does not suffer from this issue, and massively improves FP16 performance compared to both the original implementation and our GEMM-based approach.\nWith regard to FP8 and lower precision, we do not believe there are any significant blockers, mainly because Fused Neighborhood Attention as a concept can be applied to any fused attention implementation, and that includes open source FP8 implementations like Colfax's FP8 Flash Attention, and the recent Flash Attention v3.\n\n2. Thank you for mentioning this; we would be happy to add those discussions and expand further on the bottlenecks and implications on accuracy. Improvements on the speed of neighborhood attention alone won't directly affect accuracy, since functionality is unaffected. However, using faster implementations like fused neighborhood attention unlocks many new features and provides much better scalability, both of which will provide more flexibility to researchers when building their models, and that will definitely help accuracy in the long term.\n\n\nWith regard to the weaknesses, we would like to clarify that:\n\n1. The performance improvements of the GEMM-based approach are limited in FP16 due to the gather scatter operations, and for that we actually \"recommend\" researchers working on local attention kernels to move away from BMM-style implementations and work on fused attention directly instead. The memory alignment issue will grapple any multi-axis local attention kernel, and the limitation is hardware related.\n\n\n2. To clarify, FNA supports all NVIDIA architectures since Pascal, and natively targets all   architectures up to and including Ampere, and in theory can also target Ada Lovelace.    Extension to natively target Hopper Tensor Cores, TMA, and programming model, along with other architectures is in our list of future works.\nOur intention is to provide the best possible implementation to accelerate research in this direction, and thus far we have only had experience working with NVIDIA hardware, and only had NVIDIA hardware at our disposal. \nThat said, it is possible to extend our implementations to ROCm (AMD GPUs), Metal (Apple Silicon), and more, but as other reviewers have suggested, we may even look into Triton based implementations as well in order to get there more quickly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D2ALOAtkVC", "reviewer": "Reviewer_9m2o", "summary": "They implemented a fused neighborhood attention kernel (N Atten). N Atten is very useful in reducing the computational cost of various tasks because sequences usually attend to nearby (e.g., Mistral and StreamingLLM). \n\nHowever, previous implementations of N-Atten kernels are very inefficient because of a lack of \n(1) utilizing HW tensor acceleration (matmul units such as TensorCore), \n(2) utilizing fused softmax attention scheme (flash-attention).\n\nThey provide two solutions for each problem.\n(1) Using GEMM by padding the attention matrix (masking the QK matmul), the implementation could be much more easily accelerated by modern HW.\n(2) using fused softmax attention (flash-attention) to compute the attention output.", "questions": "1. Can you provide the code?\n2. This could be easily implemented in tiled acceleration languages like OpenAI Triton and Apache TVM Tensor Expression. How about implementing this method in those languages and checking the performance in heterogeneous architectures (e.g., AMD ROCm, Intel Gaudi)?", "limitations": "The proposed method only supports the N-Atten. In recent LLM research (e.g., StreamingLLM), many works propose novel methodologies to perform linear or sub-quadratic attention mechanisms using sparse masks. The current state of implementation does not consider such global attention or regional attention.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "They provide clear procedures for building the N-Atten efficiently. The writing and presentation are very clear, especially in Figures 3 and 4. Their implementation is very general to handling 1D-2D-3D attention.", "weaknesses": "Their scientific contribution is not novel enough.\n(1) GEMM-based acceleration is not a unique technique. In many sparse matrix multiplication applications, partially padding and using GEMM units is quite a natural solution (e.g., structural sparsity) \n(2) Fused softmax attention was a common acceleration technique now (after the release of flash-attention)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your feedback.\nTo answer your questions:\n\n1. Yes; we can link to or upload an anonymized version of our code if you would like, but our intention is to open source all of our code and integrate into the existing neighborhood attention package. Given the volume of the code it did not seem appropriate/necessary, but we would be happy to scrub the core implementation and share it here.\n\n\n2. Yes; you are correct that other frameworks like Triton and TVM will help implement these kernels across architectures, and we fully intend to use such solutions to expand our implementations to more architectures and platform. The reason why we chose CUTLASS and CUDA specifically was mostly because of the level of freedom in optimization and customization that CUDA C++ provides, in addition to being more familiar with those platforms and the process of performance optimization.\nOur eventual goal is to provide fast generic and cross-platform solutions for neighborhood attention and more generally multi-axis attention, and we certainly would not limit ourselves to just one platform or one case.\nFor this work, we mainly wished to illustrate with a proof of concept that the neighborhood attention family of patterns can be re-formulated as a more generic multi-axis attention, and achieve performance levels close to a practical baseline that is used in production (namely xFormers' FMHA.)\nReaching performance levels of newer fused attention kernels, extension to other platforms and architectures, and other attention patterns are definitely on our list for future work.\nWith that said, we'd note that most performant Flash Attention implementations (revision of v1, v2, and the recent v3), were also based on CUTLASS and CUDA C++, similarly to our GEMM-based and fused kernels introduced in the paper.\nIn fact, the more recent FAv3's key difference in terms of using persistent warp-specialized kernels were all concepts that naturally existed in CUTLASS before Triton, given that CUTLASS is an open source solution created by NVIDIA in the first place.\nWhile these are specific examples, we merely wish to illustrate another advantage of CUTLASS, which is quicker adoption of new architectural designs from NVIDIA for NVIDIA hardware.\nOur aim is to provide the best implementation for different hardware and architectures in order to advance research in this direction, and the hardware we had at our disposal happens to be NVIDIA hardware.\nBut again, your point stands that we should not be limiting our implementations to specific platforms and hardware, and we will definitely acknowledge this in the paper.\n\nRegarding the weaknesses, we understand that the techniques used in implementing structured sparsity have existed, and we wish to clarify that we absolutely did not intend to take claim for those, and will definitely revise the paper to reflect this.\nHowever, the work is novel in presenting, to our knowledge, the first multi-dimensional fused attention kernel (one that is back to back GETTs and not GEMMs.)\n\nWith regard to limitations, we actually can support any non-explicit attention masks (in theory). Our implementation of neighborhood attention masking in FNA provides a very simple interface through which it is easy to implement new arbitrary sparse masks.\n\nIt is perfectly feasible to connect our C++ template interface to JIT engines (torch compile) and AOT engines (AITemplate), which can \"translate\" user-specified masks into FNA masks, somewhat similar to FlexAttention. We did not mention this in the paper, but are happy to do so.\n\nIn the future, we can and will support arbitrary non-explicit masks. Lack of implementations of explicit masks was mainly for simplicity, but there isn't anything that fundamentally blocks us from doing so. We will clarify this further in our revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wuL8yn8nIi", "reviewer": "Reviewer_rBem", "summary": "This paper proposes two CUDA kernel optimization techniques: batched GEMM and fusion for neighborhood attention. On average, batched GEMM optimization provides 895% (548%) and 272% (193%) improvement in full precision (half precision) latency compared to existing naive CUDA kernels for 1-D and 2-D neighborhood attention, respectively. Fusion optimization improves naive kernels by an average of 1759% and 958% in 1-D and 2-D problems respectively. These optimizations translate into up to 104% improvement in inference and 39% improvement in training existing models based on neighborhood attention. The fused kernels can match or outperform the authors' self-attention baseline in approximately 100% of 1-D, 98.6% of 2-D, and 97.3% of 3-D problem sizes that they benchmarked.", "questions": "1. The paper does not quantitively explain where the performance benefits come from. It would help the readers assess the kernel better if you could show the arithmetic intensity of the naive, GEMM, and fused kernels under different inputs and configurations. Meanwhile, could you show some measurements on the memory, cache, and occupancy that would help explain the performance improvement?\n2. Could you show the best performance achieved by torch.compile in PyTorch 2? That would help motivate the necessity of developing CUDA kernels.\n3. Could you show how the optimized kernels impact the algorithm, such as enabling larger models or longer contexts?\n4. Could you explain whether the new tensor memory accelerator introduced in H100 can help with the neighborhood attention kernel?", "limitations": "1. The paper only shows that optimized kernels can accelerate current algorithms, but not extend their ability to more applications.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper clearly explains the overhead one can expect from switching from a standard self-attention kernel to neighborhood attention.\n2. The paper clearly illustrates the CUDA kernel optimization techniques they use, which is challenging because the kernels are complex.", "weaknesses": "1. The paper does not quantitively compare the neighborhood attention kernel with the state-of-the-art fused dot-product attention kernel.\n2. The paper only evaluates one type of GPU, the NVIDIA A100, which questions the generality of the proposed kernel optimization technique.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your feedback.\nTo answer your questions:\n\n1. Our performance evaluations are limited to runtime, mainly because the naive kernel is already not utilizing anything other than occupancy (through launching as large of a wave as possible and assigning a single dot-product’s worth of work to individual threads). It does not use tensor cores, and therefore there is no effective way of measuring math utilization in that case, at least not to our knowledge. The naive kernel also doesn’t explicitly use any levels of cache.\nThere are only two cases in which the naive kernel outperforms ours, which are:\na) GEMM and FNA hit tile quantization; or\nb) Very small wave, which means the prologue, epilogue, and gmem to smem overheads will bottleneck GEMM and FNA, which will happen in any naive vs tiled implementation.\nThe naive kernel, which is our baseline, was as we understand it a proof of concept implementation, and not a performance optimized baseline to begin with.\nAs for GEMM vs FNA, implementations are very similar, and the only bottleneck with GEMM is the memory alignment issue which we described in the paper, which FNA does not run into since the first GEMM in FNA will skip the epilogue and move results to shared memory instead of gmem.\nIn addition, prior research such as Flash Attention clearly shows that BMM-style implementations are typically memory bound problems, which includes both the naive baseline and our GEMM-based implementations. Finally, neighborhood attention itself is inherently a GEMV problem, which is also heavily memory bandwidth bound, therefore runtime / achievable FLOPS are really the only metrics we can use.\nWe would be happy to add more benchmarks if you happen to have any other metrics in mind.\n\n\n2. Unfortunately there is no way of implementing this idea with torch.compile, other than the very new FlexAttention, which to our knowledge is still in prototype. torch.compile uses an induction engine to attempt and find patterns in the graph for which a fused template kernel exists. These are typically limited to a single fusion of an elementwise or at best a reduction into a GEMM kernel, and do not extend beyond that.\nAs a result, torch.compile is great for more common fusions, such as elementwise or reduction fusion into GEMMs, and the like, there still aren't too many specific templates for attention.\nThe only functionally correct neighborhood attention implementations in python were done using im2col and padding, which cannot be inducted as traversals/views of the same tensor so that they could be fused into an attention kernel.\n\n\n3. We would note that our implementation also adds features that were previously not implemented for neighborhood attention. To name a few, causal masking, and varying per-axis parameters unlock many attention patterns that were previously non-existent. A very important one of those is spatio-temporal attention (with causal masking along time and not space), which we foresee will have great applications in video generation and video recognition.\nAside from the new features, as you pointed out, acceleration to larger context / larger inputs is one key aspect; unfortunately we have not found too many applications of neighborhood attention at very large scales, mainly because most NA applications are vision-focused, and applications with very large contexts are limited.\nHowever, we would be happy to add inference and even training results on higher-resolution tasks such as object detection and segmentation, which can better illustrate model-level improvements than image classification.\n\n4. Thank you for asking! Tensor memory accelerator (TMA) is a hardware engine for bulk data movement from global memory into shared memory. It handles some of the layout transformation, but primarily the pointer movement and predication, among other things.\nOne interesting property of the TMA is that the bulk memory accesses do not necessarily have to be in contiguous memory, and can be according to any up-to-rank-5 layout. This means that our implementation's primary bottleneck, which is data movement in the presence of non-trivial sequence \"modes\", will be handled natively through the TMA in a Hopper-oriented implementation.\nOne open source example is the CUTLASS GETT (for General Tensor-Tensor Contraction) example, in which it is illustrated how the same GEMM kernels written for Hopper can be manipulated through layout transformations and use of the TMA, to perform tensor contractions.\nAll fused attention kernels so far have been the fusion of back-to-back GEMMs with a softmax in between, but FNA is actually the fusion of back-to-back GETTs.\nThe same logic can be extended to FNA, in which our sequence \"mode\" can be 1-D, 2-D, or 3-D, which we can tile in different shapes, all through host-side layout transformations and the TMA.\nTo our knowledge, the details of the TMA copy (given a fixed number of transaction bytes) do not affect its performance as much, and certainly do not lead to unavoidable  branching (since it is a single instruction), which is what our issue is in FNA.\n\n\nWith regard to weaknesses, it is true that our comparison is only to FMHA, which is our baseline, and the reason for that is precisely the fact that we know other optimization techniques used in other state of the art methods such as Flash Attention v2 and v3 that are not used in either FMHA or FNA. However, this does not mean that our methodology is limited to FMHA or a specific architecture; it was just more generic.\n\nIn addition, our comparisons are done on Ampere (A100 specifically) mostly because that is the\nonly data-center class card we have at our disposal, and also because it is still a widely used\ncard for many in the field. However our implementation is multi-architecture and supports all architectures since Pascal, and architectures newer than Ampere can still run the kernels.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their time and their valuable feedback and suggestions.

We've posted individual rebuttals, and hope to have answered their questions and concerns.

Please let us know if there are any more questions, and we would be happy to elaborate further.

Meta-review: The paper locates at the interface of systems hardware and neural architectures, and provides optimized fused CUDA kernels that significantly speed up neighborhood attention (such as popular sliding window attention) variants, a key element of modern transformer models. Reviewers valued the principled engineering work performed here and the level of efficiency gains. 

For the camera-ready version, we urge the authors to incorporate the mentioned feedback by the reviewers.
Decision: Accept (poster)