Generalized Neighborhood Attention: Multi-dimensional Sparse Attention at the Speed of Light

2025/04/23 by Ali Hassani, Hassani, Ali, Aditya Kane +27 · 16 citations
Computer Science · Engineering · Neuroscience · #Artificial Intelligence (cs.AI) #CCD and CMOS Imaging Sensors #Computer Vision and Pattern Recognition (cs.CV) #FOS: Computer and information sciences #Gaze Tracking and Assistive Technology #Machine Learning (cs.LG) #Visual perception and processing mechanisms

paper · pdf · doi:10.48550/arxiv.2504.16922

openalex publication_date 2025/04/23 · openalex created_date 2025/10/11 · openalex updated_date 2026/07/28

Abstract

Many sparse attention mechanisms such as Neighborhood Attention have typically failed to consistently deliver speedup over the self attention baseline. This is largely due to the level of complexity in attention infrastructure, and the rapid evolution of AI hardware architecture. At the same time, many state-of-the-art foundational models, particularly in computer vision, are heavily bound by attention, and need reliable sparsity to escape the O(n²) complexity. In this paper, we study a class of promising sparse attention mechanisms that focus on locality, and aim to develop a better analytical model of their performance improvements. We first introduce Generalized Neighborhood Attention (GNA), which can describe sliding window, strided sliding window, and blocked attention. We then consider possible design choices in implementing these approaches, and create a simulator that can provide much more realistic speedup upper bounds for any given setting. Finally, we implement GNA on top of a state-of-the-art fused multi-headed attention (FMHA) kernel designed for the NVIDIA Blackwell architecture in CUTLASS. Our implementation can fully realize the maximum speedup theoretically possible in many perfectly block-sparse cases, and achieves an effective utilization of 1.3 petaFLOPs/second in FP16. In addition, we plug various GNA configurations into off-the-shelf generative models, such as Cosmos-7B, HunyuanVideo, and FLUX, and show that it can deliver 28% to 46% end-to-end speedup on B200 without any fine-tuning. We will open source our simulator and Blackwell kernels directly through the NATTEN project.

Citations

• Command A: An Enterprise-Ready Large Language Model
• WeatherMesh-3: Fast and accurate operational global weather forecasting
• Ultra-Resolution Adaptation with Ease
• XAttention: Block Sparse Attention with Antidiagonal Scoring
• Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
• Fast Video Generation with Sliding Tile Attention
• Sparse VideoGen: Accelerating Video Diffusion Transformers with Spatial-Temporal Sparsity
• Cosmos World Foundation Model Platform for Physical AI
• HunyuanVideo: A Systematic Framework For Large Video Generative Models
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
• Faster Neighborhood Attention: Reducing the O(n²) Cost of Self Attention at the Threadblock Level
• Scaling Rectified Flow Transformers for High-Resolution Image Synthesis
• Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers
• VBench: Comprehensive Benchmark Suite for Video Generative Models
• GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignment
• Mistral 7B
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• Micrograph segmentations for DDEVD
• Segment Anything
• Scalable Diffusion Models with Transformers
• Efficient Image Generation with Variadic Attention Heads
• OneFormer: One Transformer to Rule Universal Image Segmentation
• Token Merging: Your ViT But Faster
• Dilated Neighborhood Attention Transformer
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• MANIQA: Multi-dimension Attention Network for No-Reference Image Quality Assessment
• Neighborhood Attention Transformer
• Exploring Plain Vision Transformer Backbones for Object Detection
• Dynamic N:M Fine-grained Structured Sparse Attention Mechanism
• Self-attention Does Not Need O(n²) Memory
• Scatterbrain: Unifying Sparse and Low-rank Attention Approximation
• Scaling Local Self-Attention for Parameter Efficient Visual Backbones
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Big Bird: Transformers for Longer Sequences
• Linformer: Self-Attention with Linear Complexity
• Longformer: The Long-Document Transformer
• Efficient Content-Based Sparse Attention with Routing Transformers
• Blockwise Self-Attention for Long Document Understanding
• Fast Transformer Decoding: One Write-Head is All You Need
• Stand-Alone Self-Attention in Vision Models
• Generating Long Sequences with Sparse Transformers
• Online normalizer calculation for softmax
• Image Transformer
• Attention Is All You Need
• Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
• Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

Cited by

• NAF: Zero-Shot Feature Upsampling via Neighborhood Attention Filtering
• Hilbert-Guided Sparse Local Attention
• World Simulation with Video Foundation Models for Physical AI
• AnyUp: Universal Feature Upsampling
• Bidirectional Sparse Attention for Faster Video Diffusion Training
• Compact Attention: Exploiting Structured Spatio-Temporal Sparsity for Fast Video Generation
• VSA: Faster Video Diffusion with Trainable Sparse Attention
• Light Forcing: Accelerating Autoregressive Video Diffusion via Sparse Attention

Related

• Faster Neighborhood Attention: Reducing the O(n²) Cost of Self Attention at the Threadblock Level
• Neighborhood Attention Transformer
• Hilbert-Guided Sparse Local Attention
• Dynamic N:M Fine-grained Structured Sparse Attention Mechanism
• HilbertA: Hilbert Attention for Image Generation with Diffusion Models