Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers))

Subscriber deck — free during early access

1 cards in “2.4 Kernel Design” · whole deck

2.4 Kernel Design

  1. How does NSA’s kernel differ from FlashAttention’s?
    FlashAttention loads blocks of consecutive queries, but neighboring queries may select different KV blocks. NSA’s kernel instead loads, for one position, all query heads of a GQA group, which share the same selected blocks, then streams those contiguous KV blocks into SRAM, avoiding redundant KV transfers.
    • Yuan et al., 2025, Native Sparse Attention (Proceedings of the 63rd Annual Meeting of ACL (Volume 1: Long Papers)), §2.4, p. 23082

Source

  • Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Yuxing Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang and Wangding Zeng. 2025. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 23078–23097. License: CC BY 4.0. Flashcards adapted from this work. doi:10.18653/v1/2025.acl-long.1126 https://aclanthology.org/2025.acl-long.1126/