SHI-Labs/NATTEN

What it solves

NATTEN은 일반적으로 token의 수에 따라 제곱으로 증가하는 표준 셀프 어텐션의 높은 계산 비용 문제를 해결합니다. 어텐션 메커니즘에 지역성(locality)과 희소성(sparsity)을 도입하여 이미지(2D) 및 볼륨 데이터(3D)와 같은 다차원 데이터에 대해 더욱 효율적으로 만듭니다.

How it works

NATTEN은 슬라이딩 윈도우 셀프 어텐션 메커니면즘인 Neighborhood Attention (NA)을 구현합니다. 컨볼루션(convolutions)과 유사하게, 각 차원별로 kernel_size, stride, dilation을 정의하여 어텐션 범위를 제격어할 수 있습니다. 또한 토큰이 미래의 토큰을 참조하지 못하도록 하는 인과적 마스킹(causal masking)도 지원합니다.

높은 성능을을 위해, 이 프로젝트는 NVIDIA GPU에 최적화된 Fused Multi-Headed Attention (FMHA) and Fused Neighborhood Attention (FNA) 커널을 제공합니다. 여기에는 Hopper (SM90) 및 Blackwell (SM100, SM103) 아키텍처를 위한 네이티브 커널이 포함됩니다.

Who it’s for

다차원 피처 맵을 위한 Transformer를 구축하는 연구자 및 개발자를 위해 설계되었습니다. 특히 컴퓨터 비전이나 3D 데이터 처리를 다루는 분들을 대상으로 합니다.

Highlights

• Multi-dimensional support: 2D 및 3D 토큰 레이아웃에 대한 네이티브 지원.
• High-performance kernels: Maxwell (SM50) から Blackwell (SM100/103)까지의 모든 NVIDIA 아키텍처에 최적화된 융합 커널.
• Flexible configuration: 차원별 kernel_space kernel_size, stride, dilation 제어.
• Causal masking: 다차원 희소 어텐션에 대해 인과적 어텐션을 토글할 수 있는 기능.