from transformers import AutoModelForCausalLM, AutoTokenizer
m = AutoModelForCausalLM.from_pretrained("seconds-0/nsa-117m-byte", trust_remote_code=True)
t = AutoTokenizer.from_pretrained("seconds-0/nsa-117m-byte")
out = m.generate(**t("Hello", return_tensors="pt"), max_new_tokens=16)

What is NSA

Native Sparse Attention (NSA) combines three branches — compressed (cmp), selected (sel), and sliding window (win) — mixed by a learned gate. The 117M configuration uses SDPA everywhere and keeps strict causality.

Architecture (overview):

cmp: compressed blocks (tile length l, stride d) attended with causal masks
sel: top-n selection over blockized keys (block l′, n ranges per step)
win: sliding window attention of size w
gate: small MLP (zero-initialized last layer), softmax(τ=1.0)

Defaults: l=32, d=16, l′=64, n=16, w=512; GQA groups=2.

Performance & Metrics (example targets)

A100 40GB: ≥600 tok/s; TTFT ≤ 350 ms (batch=1, seq=128)
RTX 4090: ≥400 tok/s; TTFT ≤ 450 ms
CPU: ≥10 tok/s; TTFT ≤ 2.0 s

Intended Use / Limitations

Toy assistant and demos; not suitable for high-stakes use.

Memory Budget (KV Cache)

Standard LM approx: Mem ≈ t × H × (d_k + d_v) × bytes_per_elem
NSA decode (M0): Mem ≈ (min(w, t) + n × l′) × H × (d_k + d_v) × bytes_per_elem
Example (w=512, n=16, l′=64): tokens_cached ≈ min(512, t) + 1024 (FP16 → a few MiB for 117M dims)

Notes

Tokenizer: byte-level tokenizer (vocab=256). This is not GPT‑2/BPE; input/output are raw UTF‑8 bytes.
Generation cache: no KV cache in v1 (slower decode for long sequences). Planned follow‑up.
Gate: initialized to uniform mixing by design (zero‑init last layer); differs from trained gate topology.
Remote code uses SDPA-only paths and includes a safe fallback block if NSA is forcibly disabled via env.