快速多维稀疏注意力
NATTEN 是一个开源项目,致力于为多维稀疏注意力方法提供基础架构,特别是邻域注意力(NA),一种滑动窗口自注意力机制,以及其扩展(延迟 NA、因果 NA、步长 NA)。我们提供融合多头注意力(FMHA)和训练和推理内核,适用于自 Maxwell( SM50) 以来的所有 NVIDIA 架构。我们还提供了Hopper(SM90) 和 Blackwell(SM100、SM103)原生内核,可提供与 cuDNN 和 Flash Attention 3.0 相比的 FLOP 减少比例的加速。
邻域注意力将局部性和稀疏性引入到自注意力中,其方式类似于卷积。这意味着对于任何自注意力问题,您都可以指定 kernel_size、stride 和 dilation。由于这是注意力,您还可以切换因果掩码。
kernel_size=(6,6) |
kernel_size=(6,6) |
dilation=(2,2) |
|
| --- | --- |
kernel_size=(6,6) |
kernel_size=(6,6) |
暂无开放 Issues，或尚未同步最近议题。