快速多维稀疏注意力

   

NATTEN 是一个开源项目，致力于为多维稀疏注意力方法提供基础设施，特别是 邻域注意力（NA）， 一种滑动窗口自注意力机制，以及其扩展版本 空洞邻域注意力、 因果邻域注意力、 步进邻域注意力。 我们提供了融合多头注意力（FMHA）和 融合邻域注意力（FNA） 的训练与推理内核，适用于自 Maxwell（SM50）以来的所有 NVIDIA 架构。 此外，我们还发布了针对 Hopper（SM90）和 Blackwell（SM100、SM103） 的原生内核， 相较于 cuDNN 和 Flash Attention 3，能够实现与 FLOPs 减少成比例的速度提升。

邻域注意力以类似于卷积的方式，在自注意力中引入了局部性和稀疏性。 这意味着对于任何自注意力问题，你都可以指定 kernel_size、stride 和 dilation。由于它是注意力机制，你还可以启用或禁用因果掩码。

NATTEN 致力于处理多维布局的 token（即 2-D 和 3-D 特征图）。 用户可以自由探索 NATTEN 提供的巨大参数空间，在其中你可以通过各自的 kernel_size、stride、dilation 和 is_causal 参数来控制输入在任意维度/轴上的注意力范围。

kernel_size=(6,6)	kernel_size=(6,6)
	dilation=(2,2)

kernel_size=(6,6)	kernel_size=(6,6)
is_causal=(True,True)	stride=(2,2)

快速入门

NATTEN 支持 PyTorch >= 2.7 和 Python >= 3.9（即 PyTorch 所支持的一切环境）。 请参阅 安装说明 以获取有关如何安装 NATTEN 的详细信息。

:fire: 发布 0.21.6

自上次发布 (0.17.5) 以来，NATTEN 经历了重大变化，因此我们强烈建议您在升级之前阅读本页面上更新后的文档。

我们最新的版本包含了 Hopper FNA 和 Blackwell FNA 内核，可在 现代数据中心级 NVIDIA GPU 上带来 显著的速度提升，例如 H100 和 B200。 它还将现有 Ampere FNA 内核的推理速度在完全块稀疏的情况下提升至 1.47 倍， 提供了更为清晰的错误报告，并附带我们的 性能分析工具包 等诸多改进！

许可证

NATTEN 根据 MIT 许可证 发布。

引用

如果您在工作中发现 NATTEN 或邻域注意力有所帮助，请考虑引用相关论文：

邻域注意力原始论文

首次提出邻域注意力，并引入 NATTEN。

@inproceedings{hassani2023neighborhood,
  title        = {Neighborhood Attention Transformer},
  author       = {Ali Hassani and Steven Walton and Jiachen Li and Shen Li and Humphrey Shi},
  year         = 2023,
  booktitle    = {IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}
}

空洞邻域注意力

引入了 dilation 参数，用于捕捉稀疏的全局上下文。

@article{hassani2022dilated,
  title        = {Dilated Neighborhood Attention Transformer},
  author       = {Ali Hassani and Humphrey Shi},
  year         = 2022,
  journal      = {arXiv preprint arXiv:2209.15001}
}

基于 GEMM 和融合的邻域注意力

首次提出了多维注意力内核：基于 GEMM 和融合的邻域注意力（FNA）。

引入了因果邻域注意力，并扩展了实现以支持不同维度上的不同参数。

@inproceedings{hassani2024faster,
  title        = {Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Level},
  author       = {Ali Hassani and Wen-Mei Hwu and Humphrey Shi},
  year         = 2024,
  booktitle    = {Advances in Neural Information Processing Systems},
}

广义邻域注意力：迈向光速性能

引入了偶数大小的窗口、步进邻域注意力、邻域注意力的块稀疏形式、NATTEN 模拟器，以及我们全新的 Hopper 和 Blackwell FNA 内核，这些内核采用内核外的 token 排列方式实现。

@article{hassani2025generalized,
  title        = {Generalized Neighborhood Attention: Multi-dimensional Sparse Attention at the Speed of Light},
  author       = {Hassani, Ali and Zhou, Fengzhe and Kane, Aditya and Huang, Jiannan and Chen, Chieh-Yun and Shi, Min and Walton, Steven and Hoehnerbach, Markus and Thakkar, Vijay and Isaev, Michael and others},
  year         = 2025,
  journal      = {arXiv preprint arXiv:2504.16922}
}

贡献者与致谢

详情请参阅 CONTRIBUTORS.md。

NATTEN 快速上手指南

NATTEN 是一个专注于多维稀疏注意力机制的开源项目，核心提供了邻域注意力（Neighborhood Attention, NA）及其变体（如膨胀、因果、步长邻域注意力）的高效训练与推理内核。它支持从 Maxwell 到最新的 Blackwell (SM100) 架构的 NVIDIA GPU，旨在通过引入类似卷积的局部性和稀疏性，显著降低自注意力的计算成本。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python: >= 3.9
• PyTorch: >= 2.7
• GPU: NVIDIA GPU (架构需为 Maxwell/SM50 或更高)。
  • 若需体验极致性能（Hopper/Blackwell 内核），建议使用 H100、B200 等数据中心级显卡。
• 编译器: 需安装与当前 PyTorch 版本兼容的 CUDA Toolkit 和 C++ 编译器（如 g++ 或 nvcc）。

注意：NATTEN 强依赖 PyTorch 版本，请务必先升级 PyTorch 至 2.7 或以上版本。

安装步骤

1. 升级 PyTorch (如未满足版本要求)

请根据您的 CUDA 版本选择对应的安装命令（以下为示例，具体请参考 PyTorch 官网）：

pip install torch>=2.7 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

2. 安装 NATTEN

您可以直接通过 PyPI 安装最新稳定版（当前最新版本为 0.21.6，包含 Hopper 和 Blackwell 内核优化）：

pip install natten

如果您需要安装特定版本或从源码构建以适配特殊环境，请访问 官方安装文档 获取详细指引。

基本使用

NATTEN 的核心优势在于允许用户像配置卷积层一样配置注意力机制，通过 kernel_size、stride、dilation 和 is_causal 参数灵活控制注意力范围。

以下是一个最简单的 2D 邻域注意力 (Na2d) 使用示例：

import torch
import natten

# 1. 准备输入数据
# 假设输入形状为: [Batch, Heads, Height, Width, Dim]
# 例如: Batch=2, Heads=8, H=64, W=64, Dim=64
x = torch.randn(2, 8, 64, 64, 64, device="cuda", dtype=torch.float16)

# 2. 初始化邻域注意力模块
# 设置核大小为 7x7，无膨胀，无步长，非因果掩码
na_op = natten.Na2d(
    kernel_size=7,
    dilation=1,
    stride=1,
    is_causal=False
)

# 3. 执行前向传播
# 输出形状与输入相同
y = na_op(x)

print(f"输入形状：{x.shape}")
print(f"输出形状：{y.shape}")

进阶特性示例

您可以根据需求轻松切换不同的注意力模式：

• 膨胀邻域注意力 (Dilated NA): 扩大感受野

  na_dilated = natten.Na2d(kernel_size=7, dilation=2, stride=1, is_causal=False)
  

• 因果邻域注意力 (Causal NA): 适用于生成式任务

  # 在高度和宽度维度上均启用因果掩码
  na_causal = natten.Na2d(kernel_size=7, dilation=1, stride=1, is_causal=(True, True))
  

• 步长邻域注意力 (Strided NA): 下采样特征图

  na_strided = natten.Na2d(kernel_size=7, dilation=1, stride=2, is_causal=False)
  

NATTEN 会自动根据您当前的 GPU 架构选择最优的后端内核（如 Ampere FNA, Hopper FNA 等），无需手动配置。