---
license: apache-2.0
language:
- en
tags:
- nsa
- sparse-attention
- 117m
datasets:
- fineweb-edu
library_name: transformers
pipeline_tag: text-generation
base_model: byte-256
---

# NSA 117M (FineWeb-Edu) — Remote Code

This repository contains a 117M NSA decoder-only model with remote code. It exposes `NSAConfig` and `NSAForCausalLM` so you can load via:

```python
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.