MSA: Memory Sparse Attention
A scalable, end-to-end trainable latent-memory framework for 100M-token contexts.
Anonymous code release for peer review. This repository contains the inference and evaluation code used to produce the reported benchmark results.
| Code (this repo) | Anoy123423123/MSA-Code |
| Model weights | Anoy123423123/MSA-4B |
| Benchmark data | Anoy123423123/MSA-RAG-BENCHMARKS |
Requirements
Linux x86-64, Python 3.12, CUDA 12.x. 8× 80GB GPUs for the default benchmark suite; the 100M-token
ms_100M setting runs on 2× A800 via Memory Parallel. ~9 GB disk for weights, ~0.6 GB for data.
transformers must be exactly 4.51.3 — the attention implementation targets that API and will not
load on 4.52+ or 5.x. The other pins in requirements.txt are the tested versions.
Quick Start
1. Get the code
pip install -U "huggingface_hub==0.31.4"
# export HF_ENDPOINT=https://hf-mirror.com # uncomment if you need a mirror
huggingface-cli download Anoy123423123/MSA-Code --repo-type=model --local-dir MSA
cd MSA
2. Install
conda create -n msa python=3.12 -y && conda activate msa
pip install -r requirements.txt
pip install flash-attn==2.7.4.post1 --no-build-isolation
flash-attn is separate because it needs --no-build-isolation. To skip the slow source build, see the
prebuilt wheel in QUICK_START.md. Verify:
python -c "import torch, transformers, flash_attn, pynvml; print(torch.__version__, transformers.__version__, flash_attn.__version__)"
# expected: 2.6.0 4.51.3 2.7.4.post1
3. Download the weights
huggingface-cli download --resume-download Anoy123423123/MSA-4B --local-dir ckpt/MSA-4B
This is the path scripts/run_benchmarks.sh expects (model_path=ckpt/MSA-4B).
4. Run inference
bash scripts/run_benchmarks.sh eval_benchmark
Benchmark data is fetched automatically on first use into ./data/. Results land in
src/evaluation/outputs/eval_benchmark/ as one .log and one .json per benchmark.
5. Score the answers
export OPENROUTER_API_KEY=<your key>
bash scripts/calculate_llm_score.sh eval_benchmark
The argument must match the log directory name from step 4.
Benchmarks
Configured in the benchmarks array at the top of scripts/run_benchmarks.sh as name:batch_size:
2wikimultihopqa hotpotqa musique nature_questions triviaqa_06M triviaqa_10M
hipporag_narrative hipporag_popqa msmarco_v1 dureader ms_100M
ms_100M (the 100M-token setting) and triviaqa_06M are commented out by default.
Batch size must be scaled with your GPU count — the memory corpus is partitioned into one bucket per
visible GPU, and the defaults assume the 8 GPUs in CUDA_VISIBLE_DEVICES.
Notes
| Symptom | Fix |
|---|---|
| Errors loading the model | transformers must be exactly 4.51.3. |
import pynvml fails |
The module comes from nvidia-ml-py; the separate deprecated pynvml distribution shadows it. pip uninstall -y pynvml && pip install nvidia-ml-py==12.575.51 |
| CUDA OOM | Lower batch_size in scripts/run_benchmarks.sh. |
| Data download stalls | export HF_ENDPOINT=https://hf-mirror.com and retry; partial files resume. |
See QUICK_START.md for the project layout and the prebuilt flash-attn wheel.
Citation
Withheld during anonymous peer review.
License
MIT.