Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Level

We thank you for your feedback. To answer your questions:

1. 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. There are only two cases in which the naive kernel outperforms ours, which are: a) GEMM and FNA hit tile quantization; or b) 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. The 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. As 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. In 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. We would be happy to add more benchmarks if you happen to have any other metrics in mind.

2. 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. As 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. The 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.

3. 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. Aside 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. However, 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.

4. 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. One 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. One 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. All 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. The 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. To 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.

With 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.

In addition, our comparisons are done on Ampere (A100 specifically) mostly because that is the only data-center class card we have at our disposal, and also because it is still a widely used card 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.

We are also trying to gain access to H100s, so we can develop future versions of the kernels specifically targeting Hopper. We are hopeful that we will gain access in the coming months, and if time permits we will add those results as well.

Finally, as stated, we fully intend to extend our methodology to not just newer CUDA architectures, but other platforms and hardware as well (ROCm and Metal to name a few.)

We thank you for your feedback. To answer your questions:

1. 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.

2. 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. Our 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. For 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.) Reaching 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. With 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. In 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. While 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. Our 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. But 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.

Regarding 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. However, 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.)

With 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.

It 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.

In 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.

Thank you for detailed and kindly described responses. I read most of your response (including mine), and I want to raise my score, because I understand how detailed engineering considerations are there in behind. I think this paper is good enough to be accpeted because this paper will give many insight about fused N Attention applications. I hope this kind of detailed considerations are described in the further revision (in Appendix).

We thank you for your feedback. To respond to your questions:

1. 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. With 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.

2. 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.

With regard to the weaknesses, we would like to clarify that:

1. 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.

2. 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. Our 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. That 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.