https://github.com/hinanohart/nsa-eval
MPS-native Native Sparse Attention reference implementation + unified long-context evaluation suite (CUDA / MPS / CPU)
https://github.com/hinanohart/nsa-eval
apple-silicon benchmark deepseek long-context mps pytorch ruler sparse-attention
Last synced: about 1 month ago
JSON representation
MPS-native Native Sparse Attention reference implementation + unified long-context evaluation suite (CUDA / MPS / CPU)
- Host: GitHub
- URL: https://github.com/hinanohart/nsa-eval
- Owner: hinanohart
- License: mit
- Created: 2026-05-17T14:31:25.000Z (2 months ago)
- Default Branch: main
- Last Pushed: 2026-06-10T10:37:40.000Z (about 1 month ago)
- Last Synced: 2026-06-10T11:07:56.643Z (about 1 month ago)
- Topics: apple-silicon, benchmark, deepseek, long-context, mps, pytorch, ruler, sparse-attention
- Language: Python
- Size: 407 KB
- Stars: 0
- Watchers: 0
- Forks: 0
- Open Issues: 90
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
- Citation: CITATION.cff
- Codeowners: .github/CODEOWNERS
- Security: SECURITY.md
Awesome Lists containing this project
README
# nsa-eval
[](https://github.com/hinanohart/nsa-eval/actions/workflows/ci.yml) [](https://opensource.org/licenses/MIT)
**MPS-native Native Sparse Attention (NSA) reference implementation + unified long-context evaluation suite.**
A vanilla-PyTorch NSA reference that dispatches onto Metal kernels on Apple Silicon (MPS), with a cross-device (CUDA / MPS / CPU) benchmark runner and reproducible Kaggle notebooks. To the maintainer's knowledge there is no other OSS NSA implementation that targets MPS as a first-class device — please open an issue if you find a counterexample. The project runs with minimum human intervention through v0.1: a weekly Claude Code Routine drives the issue board, and a Mac contributor is needed only for MPS validation.
## Architecture
## Why
- DeepSeek NSA ([arXiv 2502.11089](https://arxiv.org/abs/2502.11089), ACL 2025 Best Paper) ships as a CUDA/Triton implementation. There is no published path for Apple Silicon users to run it locally.
- Sparse attention research papers use inconsistent benchmark setups, which makes head-to-head comparison and reproduction hard.
- This repo provides three things:
1. An MPS-native NSA reference implementation, with a Linux/CUDA path that mirrors it cell-for-cell.
2. A single runner for RULER / LongBench / Needle-in-a-Haystack / AgentBench against a fixed set of small open LLMs.
3. Kaggle notebooks (T4 / P100 / TPU v3-8) and Apple Silicon scripts that reproduce every reported number end-to-end.
## Quickstart
### Kaggle T4 (zero local setup)
Open `notebooks/_template_kaggle.ipynb` on Kaggle and click **Run All**. The notebook self-installs everything, downloads the configured model, and writes results to `benchmarks/results//`.
### Mac — Apple Silicon (MPS)
```bash
uv venv && source .venv/bin/activate
uv pip install -e .
nsa-eval run --config configs/config.yaml
```
To change the attention backend, model, benchmark, or device, point `--config` at a different YAML file (compose from `configs/attention/`, `configs/model/`, etc.):
```bash
nsa-eval run --config configs/config.yaml # uses YAML defaults
nsa-eval list-backends # show registered attention backends + benchmarks
```
### Linux / CUDA
```bash
uv pip install -e ".[nsa_cuda]"
# vendor the fla-org/native-sparse-attention submodule first:
git submodule update --init --recursive
nsa-eval run --config configs/config.yaml
```
## How it works
1. **Config** — every run is fully described by a YAML file (`configs/config.yaml` composes model + benchmark + attention + device sub-configs from `configs/`).
2. **Registry** — `BenchmarkRunner` looks up attention backends and benchmarks by name from registries populated at import time. Third-party plugins can register their own implementations without modifying core files.
3. **EvalSpec** — a Pydantic model that captures the four-tuple `(model, benchmark, attention, device)` plus a random seed. Its `.slug()` method names the output JSON.
4. **Results** — each run writes a JSON file to `benchmarks/results//.json`. Scaffold-only runs go to `benchmarks/results/_scaffold/` so they cannot pollute real result series used by the leaderboard and paper figures.
5. **MPS backend** — ships as a complete interface + Linux-runnable stub. Apple Silicon hardware is required only for real MPS execution; the stub allows CI to run everywhere.
## v0.1 (3-month) acceptance criteria
| # | criterion |
|---|---|
| 1 | Evaluation matrix: 64 of 80 cells (80%) completed |
| 2 | MPS NSA reference implementation runs Qwen2.5-1.5B + RULER 16k end-to-end on Mac M1/M2 (needs a Mac contributor) |
| 3 | CUDA NSA path reproduces the same numbers on Kaggle T4 |
| 4 | Tech report v0.1 published as a GitHub release (Markdown) |
| 5 | Static benchmark dataset + reference script published on Hugging Face Datasets |
| 6 | Kaggle notebook public and reproducible end-to-end |
| 7 | README documents the MPS path honestly: "minimum human intervention through v0.1" + maintainer-known absence of another MPS NSA OSS (counterexamples invited) |
| 8 | `experiments/_wip/` contains the failed cells, not deleted |
## Supported models and backends
**Models** (via `configs/model/`): `qwen25_0_5b`, `qwen25_1_5b`, `llama32_1b`, `llama32_3b`
**Attention backends** (via `configs/attention/`): `nsa_mps`, `h2o`, `snapkv` (registered and callable); `nsa_cuda`, `full`, `morphkv` config files exist but are not yet registered in the runtime registry
**Benchmarks** (via `configs/benchmark/`): `ruler`, `longbench`, `niah`, `agentbench`
**Devices** (via `configs/device/`): `mps`, `cuda`, `cpu`
## Mac contributor wanted
The maintainer's development machine is WSL2 Linux x86_64, so MPS execution cannot be validated locally. The MPS backend ships as a complete interface + Linux-runnable stub, and we need contributors with Apple Silicon hardware (M1 / M2 / M3) to validate, profile, and report regressions. See issue #1.
## Layout
```
src/nsa_eval/
attention/ # protocol, NSA wrapper, MPS backend stub, MorphKV decode, baselines
eval/ # eval spec, runner, benchmarks (RULER / LongBench / NIAH / AgentBench)
models/ # Qwen2.5, Llama-3.2, HF adapter
config/ # pydantic schema
utils/ # seed, device select, run manifest
cli.py
configs/ # YAML configs: model / benchmark / attention / device
notebooks/ # Kaggle + MPS templates
papers/ # tech report v0.1 source
experiments/_wip/ # failed-experiment archive (rm-forbidden)
.github/workflows/ # ci / bench-dispatch / bench-mps / leaderboard-update / preprint-rebuild / release / nightly-deps
```
## License
MIT.