Fused Neighborhood Attention

[AndreSlavescu]

🚀 The feature, motivation and pitch

Interesting work around efficient attention and general sparse attention.

Reference paper with fused NATTEN implementation in cutlass:

https://arxiv.org/pdf/2504.16922

Relevant code: https://github.com/SHI-Labs/NATTEN/tree/main/csrc/include/natten/cuda/fna https://github.com/SHI-Labs/NATTEN/blob/main/csrc/include/natten/cuda/fna/kernel_forward.h https://github.com/SHI-Labs/NATTEN/blob/main/csrc/include/natten/cuda/fna/kernel_backward.h

Alternatives

No response

Additional context

No response