qflen/nsa-from-scratch

nsa-from-scratch is a from scratch reimplementation of a technique called Native Sparse Attention, based on a February 2025 research paper from DeepSeek (arXiv 2502.11089). The project focuses on making AI language models handle long pieces of text faster and with less memory. Background: attention is the part of a language model that figures out how words relate to each other. Standard, full attention gets slower and more memory hungry as the amount of text grows, roughly by the square of the length. FlashAttention-3 is a well known fast implementation of full attention. Native Sparse Attention takes a different approach: instead of comparing every piece of text to every other piece, it skips over sections that are unlikely to matter for the current calculation. This repository writes all three parts of Native Sparse Attention, called compressed, selected, and sliding window, as custom GPU code using Triton and CUDA. The selected part also has a separate CUDA C++ version built specifically for Hopper generation GPUs, using a low level instruction called WGMMA. On an H100 NVL GPU, the selected forward pass runs 7.4 times faster than FlashAttention-3 at a 64k token context length. Beyond the core kernels, the project includes a set of five trained models ranging from 100 million to 300 million parameters, tests of how well the models predict text, called perplexity, across context lengths from 2k up to 256k tokens, an evaluation called LongBench v2, and a comparison against a rival method called MoBA. The perplexity results stay roughly stable as context length increases, which suggests the sparse method does not hurt language quality much compared to full attention. The project is written in Python, using PyTorch, Triton, and CUDA C++. It comes with an automated test suite and instructions for reproducing the benchmarks, training runs, and evaluations. The license is Apache 2.0, which allows free use, including commercial use.