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

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.

Strengths

1. The paper clearly explains the overhead one can expect from switching from a standard self-attention kernel to neighborhood attention. 2. 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. 2. The paper only evaluates one type of GPU, the NVIDIA A100, which questions the generality of the proposed kernel optimization technique.

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? 2. Could you show the best performance achieved by torch.compile in PyTorch 2? That would help motivate the necessity of developing CUDA kernels. 3. Could you show how the optimized kernels impact the algorithm, such as enabling larger models or longer contexts? 4. Could you explain whether the new tensor memory accelerator introduced in H100 can help with the neighborhood attention kernel?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

1. The paper only shows that optimized kernels can accelerate current algorithms, but not extend their ability to more applications.

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). However, previous implementations of N-Atten kernels are very inefficient because of a lack of (1) utilizing HW tensor acceleration (matmul units such as TensorCore), (2) utilizing fused softmax attention scheme (flash-attention). They provide two solutions for each problem. (1) Using GEMM by padding the attention matrix (masking the QK matmul), the implementation could be much more easily accelerated by modern HW. (2) using fused softmax attention (flash-attention) to compute the attention output.

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. (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) (2) Fused softmax attention was a common acceleration technique now (after the release of flash-attention)

Questions

1. Can you provide the code? 2. 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)?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

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.

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.

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 2. The implementations can also reduce the memory consumption of neighborhood attention. 3. 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。 2. The performance gains are more significant on newer architectures like Ampere and Hopper, limiting the generalizability of the results to other hardware setups.

Questions

1. What is the implications of the fp16 acceleration to more advanced low precision format such as fp8 or even fp6? 2. 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.

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

No significant negative societal impacts.