From-scratch reimplementation of DeepSeek's Native Sparse Attention (Yuan et al., February 2025, arXiv:2502.11089): a sparse-attention design that holds language-model quality at long context while beating FlashAttention-3 in wall-clock time as the sequence grows.

Full writeup: writeup/post.md

Throughput at 1k to 64k, perplexity at 2k to 256k, training-loss curves for a five-model fleet (NSA-100M, NSA-150M, NSA-300M, dense-100M baseline, NSA-100M at 64k context), LongBench v2 likelihood across six tasks, and a MoBA cross-comparison are all in the writeup.

• All three NSA branches (compressed, selected, sliding window) as Triton kernels.
• The selected branch additionally implemented in CUDA C++ on Hopper using inline-PTX WGMMA atoms (m64n64k16 and m64n128k16 at D=64 and D=128).
• Multi-precision: FP16, BF16, FP8 (per-tensor absmax, bf16 dequant in the wrapper).
• Hand-written Triton selected backward (FA-2 style streaming softmax adapted to the gather pattern, atomic-add into fp32 dK/dV buffers).
• 48-config autotune sweep (BLOCK_M, BLOCK_N, num_warps, num_stages) for the selected branch on Hopper.
• A five-model training fleet on FineWeb-Edu at 32k to 64k context with bf16 mixed precision.
• Long-context perplexity sweep across six evaluation lengths (2k to 256k) with straight and NTK-aware RoPE extension.
• LongBench v2 gold-answer-likelihood subset across six tasks.
• MoBA (Liu et al., February 2025) cross-comparison on the same Triton bench harness.

Python, PyTorch, Triton, CUDA C++ (Hopper sm_90a, WGMMA inline-PTX), torch.utils.cpp_extension, pybind11, wandb, matplotlib. RunPod for GPU rental (H100 NVL, A100 SXM, A100 PCIe).

pip install -e '.[dev,train]'
pytest tests/ -q                              # CPU-only runs test_reference.py alone; everything else needs CUDA (the CUDA-C++ WGMMA and FP8 tests additionally need Hopper sm_90)
# Headline throughput is a 3-seed protocol (100 iters, 25 warmup each),
# pooled into a 95% bootstrap CI (this is what reports 7.07x at 64k):
for s in 0 1 2; do
  python -m nsa.bench.throughput --seq-lens 1024,2048,4096,8192,16384,32768,65536 \
      --impls nsa,fa3,full_sdpa --dtype bf16 --warmup 25 --iters 100 --seed $s \
      --out runs/throughput_seed$s.json
done
python scripts/multiseed_combine.py \
    --seeds runs/throughput_seed0.json,runs/throughput_seed1.json,runs/throughput_seed2.json \
    --out runs/throughput_multiseed.json

The throughput benchmark needs a Hopper card (the WGMMA selected forward and the FA-3 baseline both require sm_90a). On a non-Hopper CUDA card (sm_89) the Triton, training, and reference tests run while the CUDA-C++ WGMMA and FP8 tests skip; on a CPU-only runner (the CI Reference tests (CPU) job) only tests/test_reference.py runs and everything that needs CUDA skips.

Configurations live in nsa/train/config_*.yaml. The five recipes shipped:

# NSA scaling points, all 32k context
python -m nsa.train.train --config nsa/train/config_100m.yaml         # NSA-100M, 1B tokens
python -m nsa.train.train --config nsa/train/config_150m.yaml         # NSA-150M, 1B tokens
python -m nsa.train.train --config nsa/train/config_300m.yaml         # NSA-300M, 500M tokens

# Dense baseline at the largest workable context (8k on A100 PCIe at this batch shape)
python -m nsa.train.train --config nsa/train/config_dense_100m.yaml

# Long-context training demonstration (64k context, 100M parameters)
python -m nsa.train.train --config nsa/train/config_100m_64k.yaml

Each run logs to wandb project nsa-from-scratch and saves checkpoints to runs/<run_id>/. The token counts above are below Chinchilla-optimal because the autograd-aware bwd through the compressed and sliding branches dominates training throughput at 32k context; the writeup discusses this and the natural follow-up (hand-written Triton sliding and compressed backward kernels).

python -m nsa.eval.perplexity \
    --run_dir runs/nsa-100m-1b-32k \
    --seq_lens 2048,4096,8192,16384,32768 \
    --num_sequences_per_len 32 \
    --out writeup/figures/data/perplexity_nsa-100m-1b-32k.json

The eval streams a deterministic offset of FineWeb-Edu, packs into the target seq_len, and reports cross-entropy and perplexity per cell.

nsa/
  reference.py          plain-torch reference for the three branches (correctness oracle)
  triton/               compressed.py, selected.py, sliding.py, gating.py, forward.py, backward.py, fp8.py, moba.py
  cuda/                 selected_fwd.cu (Hopper WGMMA), selected_bwd.cu (in-progress), bindings.cpp, setup.py
  model/                llama_nsa.py, llama_dense.py, config.py
  train/                train.py, data.py, config_*.yaml (NSA-100M, 150M, 300M, dense, 64k, sanity)
  eval/                 perplexity.py, long_context_probe.py, longbench.py
  bench/                throughput.py, memory.py, memory_sweep.py, correctness.py, branch_breakdown.py, autotune.py, plots.py
tests/                  test_reference, test_{compressed,selected,sliding,combined}_forward, test_selected_backward, test_{fp8,moba,training_step}, test_cuda_selected, test_cuda_selected_bwd
writeup/                post.md, decisions/*.md, figures/{01..08}.png
notes/                  refs.bib
scripts/                fetch_wandb.py, make_plots.py, multiseed_combine.py

MIT. See LICENSE.