Research map / Architectures and efficiency / Attention and transformers
Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Level
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes GEMM-based and fused CUDA kernels (built on CUTLASS and xFormers FMHA) for 1-D/2-D/3-D neighborhood attention, yielding large speedups over the naive NATTEN kernels on A100.
Problem
Neighborhood and sliding-window attention reduce self-attention cost but lack efficient implementations in higher-rank (2-D/3-D) token spaces; existing naive CUDA kernels are memory-bound and slow.
Approach
Reformulates neighborhood attention's GEMV-like operators (PN, NN, IN) as batched GEMMs with space-aware tiling and gather/scatter fusion, and extends fused FMHA kernels to multi-dimensional token spaces with per-axis window, dilation, and causal masking.
Key results
- Fused NA improves naive kernels by an average of 1759% (1-D) and 958% (2-D) in FP16 forward pass on A100.
- GEMM-based NA improves naive kernels in FP16 forward by 548% (1-D) and 193% (2-D) on average.
- Fused kernels give up to 104% FP16 inference throughput improvement for DiNAT-L at window size 384 and up to 39% training time improvement.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use fused neighborhood attention instead of naive kernels | Neighborhood attention kernels · 1-D neighborhood attention, FP16 forward | average speed improvement over naive: 1759% | Helped |
| Use fused neighborhood attention instead of naive kernels | Neighborhood attention kernels · 2-D neighborhood attention, FP16 forward | average speed improvement over naive: 958% | Helped |
| Use fused neighborhood attention instead of naive kernels | Neighborhood attention kernels · 3-D neighborhood attention, FP16 forward | average speed improvement over naive: 1135% | Helped |
| Use GEMM-based neighborhood attention instead of naive kernels | Neighborhood attention kernels · 1-D neighborhood attention, FP16 forward | average speed improvement over naive: 548% | Helped |
| Use GEMM-based neighborhood attention instead of naive kernels | Neighborhood attention kernels · 2-D neighborhood attention, FP16 forward | average speed improvement over naive: 193% | Helped |
| Use GEMM-based neighborhood attention instead of naive kernels | Neighborhood attention kernels · 1-D neighborhood attention, FP16 forward + backward | average speed improvement over naive: 502% | Helped |
| Use GEMM-based neighborhood attention instead of naive kernels | Neighborhood attention kernels · 2-D neighborhood attention, FP16 forward + backward | average speed improvement over naive: 92% | Helped |
| Use fused neighborhood attention instead of naive kernels | Neighborhood attention kernels · 1-D neighborhood attention, FP16 forward + backward | average speed improvement over naive: 844% | Helped |
| Use fused neighborhood attention instead of naive kernels | Neighborhood attention kernels · 2-D neighborhood attention, FP16 forward + backward | average speed improvement over naive: 385% | Helped |
| Use fused neighborhood attention instead of naive kernels | Neighborhood attention kernels · 3-D neighborhood attention, FP16 forward + backward | average speed improvement over naive: 447% | Helped |
| Replace naive NAT-M kernels with GEMM-based kernels | NAT-M · ImageNet-1K, FP16 | throughput (imgs/sec): 2975 -> 2660 | Hurt |
| Replace naive NAT-M kernels with fused kernels | NAT-M · ImageNet-1K, FP16 | throughput (imgs/sec): 2975 -> 3742 | Helped |
Surprising findings
- GEMM-based (unfused) kernels underperform naive kernels for NAT/DiNAT ImageNet models in FP16 throughput, due to attention-weight gather/scatter breaking asynchronous pipelining.
- GEMM-based kernels show no training-time improvement over the naive baseline.
Limitations
- Unfused GEMM kernels suffer from scatter/gather of attention weights, especially in FP16/BF16, limiting gains.
- Fused kernels can be slower than the FMHA self-attention baseline in cases close to self attention; 2-D/3-D predication adds overhead.
- Auto-tuner is single-process and the search space is large.
Experimental setup
Datasets: ImageNet-1K FFHQ
Benchmarks: ImageNet-1K throughput benchmarks FFHQ StyleNAT throughput benchmarks operator-level FP16/FP32 kernel benchmarks on A100
Metrics: runtime throughput (imgs/sec) training time estimate (hours) FID
Baselines: NATTEN naive CUDA kernels xFormers FMHA FlashAttention V2
Builds on: Neighborhood Attention FlashAttention xFormers FMHA CUTLASS NATTEN
Closest papers by meaning
- Neighborhood Attention Transformer (CVPR 2023 · similarity 0.91)
- The Kernel Reality Check: Benchmarking and Distilling Efficient Attention at Scale (NeurIPS 2026 · similarity 0.89)
- Bridging the Divide: Reconsidering Softmax and Linear Attention (NeurIPS 2024 · similarity 0.88)
- Towards Interpretable and Efficient Attention: Compressing All by Contracting a Few (NeurIPS 2025 · similarity 0.87)
- Faster Video Diffusion with Trainable Sparse Attention (NeurIPS 2025 · similarity 0.87)
- Attend Locally, Remember Linearly: Linear Attention as Cross-Frame Memory for Autoregressive Video Diffusion (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.