《GPU Mode》 L065 2025 · Local attention High priority transcript · failed

Neighborhood Attention

모든 query 가 모든 key 를 볼 필요 없다 — vision · video · diffusion 의 token 들은 공간적으로 인접한 이웃과의 attention 만으로도 충분하다는 관찰. NATTEN (Neighborhood Attention) 은 Ali Hassani 등이 만든 라이브러리 — sliding window 의 vision 친화적 변형 + GPU-optimized 구현. SAM, DiT, CLIP 같은 모델에서 attention 비용을 한 자리수 줄이면서 같은 정확도를 유지한다. 이 노트는 NA 의 수식, 구현, 그리고 다른 sparse attention 과의 차이를 정리한다.

neighborhood attention sliding window NATTEN library vision transformer SAM DiT FlashAttention backward memory

Speaker

Ali Hassani

SHI Labs / 메타 등 (확인 필요) · NATTEN / NA / NA-Transformer 저자

강의 번호

L065

스피커

Ali Hassani

학습 우선순위

High · 정독

자료 상태

transcript 없음 · 슬라이드 있음

하위 목차 · 12개 섹션↓ 클릭해서 이동

01강의가 풀려는 문제why local attn for vision 02sliding-window 의 동기vision token locality 03NA 의 수식k-nearest neighbor attention 04NATTEN 구현CUDA · CUTLASS 05backward 메모리 회계grad of windowed 06SAM · DiT 사례channel adopters 07다른 sparse attention 과 비교SWin / Longformer / STA 08hardware 활용Tensor core utilization 09다음 방향3D · dilated · learned 10기억할 메모와 코드NATTEN repo 11다른 강의로 이어지는 길connections 12열린 질문open questions

§ 01강의가 풀려는 문제· why local attn for vision

vision token 에 quadratic attention 은 너무 비싸다 — 그리고 거의 항상 불필요하다

vision transformer 는 image / video 를 patch 로 잘라 attention 을 도린다. 768×768 image 가 patch 16×16 면 — N = 48 × 48 = 2304 tokens. quadratic attention 의 N² ≈ 5M scores. 작은 것 같지만 — 4K image (256×256 patches) 면 N² = 4×10⁹. 비현실. NATTEN 의 답 — 각 query 가 자기 주변 k×k 이웃만 attend. quadratic 을 linear-in-N (kernel-size 에 의존) 으로 바꿈. 이 아이디어를 GPU 위에서 효율적으로 돌리는 라이브러리.

강의 transcript 가 비어 있으나 슬라이드는 존재 (Google Slides). 본 노트는 NATTEN 의 GitHub repo, NA Transformer 의 paper (Hassani et al. 2022), DiNAT / NA paper (CVPR 2023), 그리고 NATTEN 2.0+ 자료 위에서 재구성. 슬라이드의 정확한 해석은 직접 확인 필요.

강의의 frame

NA 의 핵심 명제 — vision 의 token 은 공간적으로 인접한 token 과의 관계가 중요하다. 멀리 떨어진 token 은 거의 안 본다 (또는 hierarchical 로 충분). 그래서 — 각 query 가 자기 위치 (h, w) 의 k×k 윈도우 안 token 만 보면 된다. 이게 convolution 의 receptive field 와 같은 idea — 그러나 attention 으로 weighting 학습.

“neighborhood attention 은 — convolution 의 inductive bias 와 attention 의 weight learning 의 결합이다. vision 에 가장 자연스러운 sparse pattern.”Ali Hassani / 확인 필요

§ 02sliding-window 의 동기· vision token locality

vision token 의 importance 는 거리에 빠르게 decay 한다

표준 ViT 의 attention map 을 분석하면 — 대부분의 attention 가중치가 인접한 몇 개 token 에 집중. 멀리 떨어진 token 의 attention weight 는 거의 0. 이게 NA 의 empirical 동기.

이 관찰은 — vision 도메인의 자연스러운 inductive bias 와 일치. 이미지의 한 patch 의 의미는 (대부분) 자기 주변 patch 로 결정. 멀리 있는 patch 는 hierarchical (downsample) 로 정보가 모여서 도달.

CNN 의 receptive field: 각 layer 가 인접 영역만 본다. 그게 vision 에 잘 맞는다는 게 deep learning 의 큰 lesson.
Swin Transformer (2021): window 를 7×7 로 자르고 그 안에서만 attention. shifted window 로 cross-window 정보 흐름.
Neighborhood Attention (2022): window 를 더 자연스럽게 — 각 query 의 자기 중심으로 sliding window. “query-centric” 인 게 핵심 차이.

즉 NA 는 — convolution 의 sliding window 를 attention 의 form 으로 가져온 것. 모든 query 가 자기 주변 k×k 만 본다. 다른 attention 가중치는 0.

FIG · Swin vs NA 의 window 정의같은 N tokens · k=3

Swin — 고정 window

NA — query 중심 sliding

Swin 은 같은 window 안 모든 query 가 같은 영역. NA 는 각 query 가 자기 중심 영역. cross-window 흐름이 자동으로 일어남 (인접한 query 의 window 가 겹침). shifted window 트릭이 필요 없음.

§ 03NA 의 수식· k-nearest neighbor attention

각 query position 의 k×k 이웃만 attend — 정확한 정의

Neighborhood Attention
// 위치 (i,j) 의 query Q[i,j] 가 attend 하는 K, V 의 집합

    N(i,j) = {(i', j') : |i' − i| ≤ k/2,

                       |j' − j| ≤ k/2}


    O[i,j] = softmax_k∈N(i,j)( Q[i,j] · K[k] / √d ) · V[k]
    → 매 query 가 자기 위치 ± k/2 의 정사각 윈도우만 본다. boundary 에서는 가장 가까운 valid 영역으로 shift (= 항상 k² 개 attend).

이 수식의 두 가지 미묘한 점.

boundary handling: image 의 corner / edge 에서는 (i ± k/2, j ± k/2) 가 image 밖. 두 가지 선택 — (a) padding (zero 또는 reflection), (b) shift — corner query 도 항상 valid 한 가장 가까운 k×k 영역을 attend. NATTEN 은 shift 채택. 모든 query 가 정확히 k² attend.
비용: query N 개 × attention k² × hidden d. 총 O(N · k² · d). quadratic 의 O(N² · d) 와 대비 — k 가 N 보다 훨씬 작으면 N 으로 linear.

예: 4K image, patch=16, N=65K. quadratic attention 비용 = 4.3×10⁹. NA k=7 = N × 49 = 3.2×10⁶. 1000× 차이.

k 의 선택

NA paper 들이 보여준 일반적 권장값 — k=7 (49개 attend) 또는 k=11 (121개 attend). 더 작으면 receptive field 부족, 더 크면 비용 폭발 (k² 이므로). vision transformer 는 — 여러 layer 를 쌓으면서 effective receptive field 가 자연스럽게 커지므로, 단일 layer 의 k 는 작아도 OK.

variant — Dilated NA (DiNAT, CVPR 2023): k×k 윈도우를 그냥이 아니라 dilation 1, 2, 4 처럼 띄엄띄엄. multi-scale receptive field 를 한 layer 안에서.

§ 04NATTEN 구현· CUDA · CUTLASS

FlashAttention 의 tile-based softmax 를 — sliding window 안에서 한다

NA 의 빠른 구현은 FlashAttention 2/3 의 tile-fused 패턴을 sliding window 의 mask 안에서 도리는 것. NATTEN 라이브러리가 정확히 이 일을 한다 — CUDA + CUTLASS 기반의 fused attention 커널, k×k mask 를 자동으로 생성, backward 도 같은 형태.

구현의 핵심 디테일:

tile 단위 처리: query 를 8×8 또는 16×16 tile 로 묶어 한 SM 이 처리. 같은 tile 안 query 는 비슷한 K/V 영역을 공유 — KV reload 가 줄어듦.
mask 생성 자동: window 가 image boundary 와 어떻게 겹치는지를 SM 이 자기 tile 좌표에서 계산. sparse pattern 의 indexing overhead 없음.
tensor core 활용: attention score Q·K^T 와 softmax-V 가 모두 tensor core 의 matmul. Hopper 의 wgmma 까지 활용 (NATTEN 2.x · 확인 필요).
FA3 위 build: NATTEN 2.x 가 FlashAttention 3 의 internal 을 활용. softmax online + warp specialization.

FIG · NATTEN 의 fused tile pattern2D image · k=7

query tile (8×8)
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
KV tile (14×14)
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·
·

8×8 query 의 union neighborhood = 14×14 (k=7 의 경우). 한 SM 이 14×14 KV 를 SRAM 에 올리고, 그 위에서 8×8 query 의 모든 attention 을 한 번에. KV 의 tensor core matmul.

§ 05backward 메모리 회계· grad of windowed

backward 가 까다롭다 — 같은 K/V 가 여러 query 의 gradient 에 contrib 한다

forward 의 NA 는 직관적이지만 — backward 는 미묘하다. 같은 (i, j) 의 K 가 — 자기 위치를 둘러싼 k×k query 들의 gradient 에 contribute. K 의 gradient = sum over k² queries that attend to it. naive 하게 짜면 race condition / 메모리 폭발.

backward 의 메모리 / computation 패턴:

dQ: 각 query 의 gradient 는 자기 window 안 K 들의 gradient 의 weighted sum. 자기 query 위치 별로 독립 → forward 와 거의 같은 비용.
dK, dV: 한 K[i,j] 가 contribute 하는 query 의 set 은 — 자기 위치를 둘러싼 k×k query. 역전된 NA. forward 와 같은 형태로 계산 가능.
atomic 충돌: naive 한 backward 는 같은 K 위치에 여러 thread 가 atomic add. 성능 폭락. NATTEN 의 backward 는 — query-tile 단위로 묶어 atomic 없이 reduction.

FlashAttention 의 backward (recomputation 으로 attention matrix 재생성) 와 같은 트릭이 NATTEN 에도 적용. activation 메모리 폭발 안 함.

메모리 산수

standard attention 의 backward 는 N² 의 score / probability 를 보관하거나 재계산. NA 는 — N×k² 만 보관하면 됨. N² 이 1000× 큰 영역에서는 메모리 saving 도 1000×. 같은 GPU 에서 1000× 큰 image 를 학습 가능. 이게 NA 의 진짜 가치 중 하나.

NATTEN 의 grad checkpoint 와의 결합 — 각 layer 의 attention output 만 저장, backward 시 forward 재계산. 일반 ViT 에 비해 메모리 요구가 한 단계 더 작음.

§ 06SAM · DiT 사례· channel adopters

NATTEN 을 채택한 큰 모델들

SAM 2 / Hiera

Meta 의 segment-anything. Hiera backbone 이 NATTEN-style local attention. 큰 image 처리에 필수 (확인 필요).

DiT (Diffusion Transformers)

image / video diffusion 의 backbone. 일부 implementation 이 NATTEN 으로 attention 비용 감소 (확인 필요).

NA-Transformer (원조)

Hassani et al. 2022. classification / detection / seg 에서 ViT / Swin 과 비슷한 정확도 + 더 빠름.

DiNAT (Dilated NA)

CVPR 2023. dilation 으로 multi-scale receptive field. dense prediction 에서 강함.

FastVideo의 STA

Sliding Tile Attention — NA 의 3D 일반화. video diffusion 에 채택 (L059 참조).

MAE / DINO 같은 self-supervised

큰 image 에서 self-supervised pretrain 의 base 로 NATTEN 활용 가능 (확인 필요).

중요한 패턴 — NA 가 잘 맞는 자리는 큰 image / video. 작은 image (224×224) 에서는 quadratic attention 이 여전히 OK. 큰 해상도 (1024×1024+) 또는 video 에서 NA 의 효율 우위가 폭발적.

§ 07다른 sparse attention 과 비교· SWin / Longformer / STA

같은 “local attention” 가족 안의 다른 변형들

방법

window 정의

cross-window

domain

Swin (2021) 고정 window 7×7 shifted window vision · 2D

NA (2022) query-centric k×k 자동 (window 겹침) vision · 2D

DiNAT (2023) NA + dilation multi-scale dense prediction

Longformer (2020) 1D sliding + global token global token language · 1D

BigBird (2020) local + global + random random sampling language · 1D

STA (2025) 3D query-centric tile 자동 video diffusion · 3D

Native Sparse Attn 학습된 sparsity token-level language · DeepSeek

NA 의 차별점

(1) query-centric: Swin 의 fixed window 와 달리 query 마다 자기 중심 window. shifted window trick 불필요. (2) vision-native: 2D / 3D structure 를 first-class 로. 1D Longformer 와 다름. (3) FA-style fused 구현: NATTEN 라이브러리가 GPU 에서 dense attention 만큼 빠름 — Swin 의 일반적 구현보다 빠를 수 있음.

§ 08hardware 활용· Tensor core utilization

sparsity 가 있어도 — Tensor Core 가 saturate 한다

naive sparse attention 은 GPU 위에서 dense 보다 느린 경우가 많다 — sparsity 의 indexing overhead, irregular access pattern 이 tensor core 활용을 깨뜨림. NATTEN 은 이 문제를 — regular sparsity (sliding window) 와 tile-based fusion 의 결합으로 푼다.

naive PyTorch (sparse)

~80 TF

FlashAttention (dense)

~640 TF

NATTEN 1.x (CUDA)

~540 TF

NATTEN 2.x (FA3)

~790 TF

peak BF16

990 TF

예시 수치 — H100 BF16, 일반적 vision transformer setting. NATTEN 2.x 가 dense FA 와 거의 동등한 utilization. sparsity 의 비용 거의 0. (확인 필요)

NATTEN 의 hardware 친화적 디자인 결정:

regular pattern: query position 의 함수로 KV index 가 정확히 결정. branch 없이 mask 가 mathematical 로 표현.
tile fusion: 같은 tile 안 query 가 비슷한 KV 영역을 share. SRAM 안에서 KV 한 번 load 로 여러 query 처리.
boundary 처리 = shift: edge query 도 항상 정확히 k² 개 attend. dynamic mask 없음 — 모든 SM 이 같은 양의 일.
FA3 통합: warp specialization, async TMA copy. Hopper 의 native feature 모두 활용.

irregular sparsity 의 함정

BlockSparse / random sparse attention 같은 게 — sparsity ratio 90% 라도 GPU 위에서는 dense 보다 별로 안 빠른 경우 많음. 이유 — irregular access, branch divergence, mask gather. NA 의 regular pattern 이 이 함정을 피하는 가장 큰 차이.

이 hardware 친화성 덕에 — NATTEN 이 “sparse attention = 느리다” 의 통념을 깨고 production 채택이 늘어남.

§ 09다음 방향· 3D · dilated · learned

NATTEN 의 진화 축

3D NA / Sliding Tile Attention — video / volumetric data. (T, H, W) 3D 윈도우. FastVideo 의 STA 가 이 가족.
Dilated NA — DiNAT. 같은 layer 안 multi-scale receptive field. dense prediction 에서 강함.
Learned NA — window size / dilation 을 layer / token 별로 학습. data-dependent. 진행 중 — 확인 필요.
Hopper / Blackwell 전용 최적화 — TMA, wgmma, TMEM 등 새 hardware 기능 활용. NATTEN 2.x 가 이 방향 진행 중.
autograd 통합 향상 — torch.compile 친화적 구현. 자동 fusion 과 결합.
ALiBi / RoPE 결합 — positional encoding 과 NA 의 통합. local attention 의 positional bias 자연스럽게.

“NA 의 진화는 — vision 의 inductive bias 를 attention 의 한 형태로 표현하는 깊은 길이다. CNN 과 transformer 의 합류 지점.”학습 노트

§ 10기억할 메모와 코드· NATTEN repo

다시 열었을 때 빠르게 잡혀야 할 것

NA 의 정의

각 query 가 자기 (i,j) 의 k×k 이웃만 attend. 비용 N×k² (vs N²).

query-centric vs window-fixed

Swin 은 fixed window. NA 는 query 마다 자기 중심. shift trick 불필요.

k 의 권장값

7 또는 11. 너무 작으면 receptive field 부족, 너무 크면 비용.

backward 의 까다로움

같은 K 가 여러 query 의 gradient 에 contrib. tile 단위 reduction 으로 풀림.

메모리 saving

N² 이 N×k². 큰 image 에서 1000× saving — backward 메모리도 함께.

NATTEN 라이브러리

CUDA + CUTLASS. FA3 통합. natten.functional.na2d_qk 등.

DiNAT

dilation 추가. multi-scale receptive field. dense prediction 에서 SOTA.

vs 다른 sparse

Swin / Longformer / BigBird 와 가족. regular pattern 으로 tensor core 활용.

YouTube youtube.com/watch?v=y5r2asbfNcs

Slides Google Slides

NATTEN repo github.com/SHI-Labs/NATTEN

NA paper Neighborhood Attention Transformer · arXiv:2204.07143

DiNAT paper Dilated NA · arXiv:2209.15001

FlashAttention 3 Dao-AILab/flash-attention

NATTEN docs shi-labs.com/natten (확인 필요)

손에 새기기 — 실습 시퀀스

NATTEN 설치 + 기본 forward — pip install natten. 같은 (Q,K,V) 에 대해 standard attention 과 NA 의 forward 비교. 결과 모양 확인.
속도 비교 — H=W=64, 128, 256, 512 sweep. quadratic attention vs NA k=7 의 ms. crossover 가 어디서.
k sweep — k ∈ {3, 5, 7, 11, 17} 에서의 정확도와 속도 trade-off. 다운스트림 task accuracy 비교.
큰 이미지 backward — 1024×1024 image 의 backward 메모리 측정. quadratic attention 은 OOM, NA 는 통과.
DiNAT 시도 — dilation 1, 2, 4 의 stack. multi-scale receptive field 의 학습 효과.
SAM 또는 DiT 에 NATTEN 끼워넣기 — 표준 attention 을 NA 로 교체. 정확도 회복 / 속도 향상 확인.
nsys 로 timeline — NATTEN 커널의 SM 활용 / Tensor Core utilization 측정. dense FA 와 비교.

§ 11다른 강의로· connections

이 강의의 frame 이 다른 강의에서 어떻게 다시 등장하는지

L059

FastVideo

— Sliding Tile Attention 이 NA 의 3D 일반화. video diffusion 핵심

L012

Flash Attention

— NATTEN 의 base. tile-fused attention 의 dense 형제

L011

Sparse Attention

— sparse attention 가족 전체. NA 가 그 안의 한 갈래

L060

Linear Attention

— quadratic 줄이는 다른 길. local vs linear 의 trade-off

L021

CUTLASS / CuTe

— NATTEN 이 build 한 base. CUTLASS 의 fused attention

L036

Diffusion 가속

— DiT / image diff 의 backbone 으로 NA 의 채택

§ 12열린 질문· open questions

transcript 가 비어 있어 직접 검증해야 할 것들

슬라이드 직접 확인 — 강의의 슬라이드는 존재 (Google Slides). transcript 없는 상태에서 슬라이드의 그림 / 수식 / 벤치마크 수치 직접 확인.
NATTEN 2.x 의 정확한 release 시점 — FA3 통합이 어디까지 왔는지, Hopper-only 인지 — 확인 필요.
실제 production 채택 list — SAM 2 / Hiera / DiT 등에서의 정확한 NATTEN 사용 — 강의에서 어떤 사례를 다뤘는지 미확인.
backward 의 정확한 trick — atomic 회피와 query-tile reduction 의 정확한 디테일이 NATTEN repo 의 CUDA 코드에 있음. 강의 깊이 미확인.
3D NA 의 status — STA (FastVideo) 와 NATTEN 의 자체 3D 지원의 관계 — 같은 가족인지 별도 라이브러리인지 미확인.
사용자 정의 mask 지원 — NATTEN 이 fixed sliding window 외 user-defined sparse pattern 도 지원하는지 (확인 필요).
distillation / training 시 효과 — quadratic 으로 학습된 모델을 NA 로 fine-tune 시 정확도 회복 곡선.

검증 메모

본 노트는 NA / DiNAT 의 paper 들과 NATTEN repo 의 README / docs 의 일반적 framework 위에서 재구성. 강의의 Google Slides 가 존재하므로 — 슬라이드를 직접 확인하면 이 노트의 디테일을 빠르게 보강 가능. 특히 hardware 활용 수치와 backward 의 정확한 algorithm 은 슬라이드 / 영상 검증 필요.

← Lecture 064 Multi-GPU programming — single kernel sparsity 와 다른 axis