Upload README.md with huggingface_hub

README.md

@@ -1,96 +1,92 @@
----
-datasets:
-- ruler
-- longmemeval
-- longbench-v2
-- mrcr
-language: en
-license: mit
-pipeline_tag: feature-extraction
-tags:
-- deepseek-v4
-- retrieval
-- kv-cache
-- sparse-attention
-- long-context
-- flashmemory
----
-
-# FlashMemory DS-V4 Retriever
-
-A lightweight retriever that sparsifies **DeepSeek-V4 Compressed-Sparse-Attention (CSA) KV-cache**. Given a decode-token hidden state, it predicts which compressed-K chunks the next ~64 tokens will attend to — keeping only those on GPU, offloading the rest.
-
-This model is the Neural Memory Indexer presented in [FlashMemory-DeepSeek-V4: Lightning Index Ultra-Long Context via Lookahead Sparse Attention](https://huggingface.co/papers/2606.09079).
-
-Detailed code and inference scripts can be found on GitHub: [libertywing/FlashMemory-Deepseek-V4](https://github.com/libertywing/FlashMemory-Deepseek-V4).
-
-## Performance
-
-In downstream evaluation, it matches or beats the full-attention baseline on reasoning-heavy long-context tasks (**RULER, LongMemEval, LongBench V2**) while reducing KV-cache usage by **~85–90%**. Precise needle-retrieval tasks require an additional threshold-fallback mechanism (not in this release).
-
-## Quick start
-
-To use this model, you will need the `retriever.py` file from the [official repository](https://github.com/libertywing/FlashMemory-Deepseek-V4).
-
-```bash
-pip install torch safetensors
-python demo.py --ckpt weights/flashmemory_ds_v4.safetensors
-```
-
-## Usage
-
-```python
-from retriever import FlashMemoryRetriever
-
-model = FlashMemoryRetriever.from_checkpoint(
-    "weights/flashmemory_ds_v4.safetensors", device="cuda"
-)
-
-# hidden:     [B, 4096] decode-token hidden state
-# comp_k:     [B, N, 132] uint8 compressed CSA keys
-# positions:  [B] int64 token positions
-
-# Per-layer sigmoid scores: {"l10": [B,N], "l12": [B,N], "l20": [B,N]}
-per_layer = model(hidden, comp_k, positions)
-
-# Cross-layer ensemble (mode="max" or "mean")
-scores = model.ensemble(hidden, comp_k, positions, mode="max")       # [B, N]
-
-# Boolean keep mask
-keep = model.select_topk(hidden, comp_k, positions, top_k=512)       # top-K
-keep = model.select_topk(hidden, comp_k, positions, threshold=0.5)   # threshold
-```
-
-**`compressed_k` format:** each chunk = 128 bytes `float8_e4m3` values + 4 bytes `float32` scale. See `make_mock_compressed_k()` in `demo.py`.
-
-## Architecture
-
-3-layer joint model (`l10`, `l12`, `l20`), 128 heads, 2048 LoRA rank. Per-layer sigmoid scores are ensembled (`max` or `mean`) per chunk.
-
-```
-hidden [B,4096] → q-proj → RoPE(YaRN) → Hadamard → q [B,128,128]
-               → weights_proj → fused_w [B,128]
-compressed_k    → FP8 dequant → k [B,N,128]
-
-score = sigmoid( Σ( relu(k @ qᵀ) · fused_w ) )  ∈ [0,1]
-```
-
-## Citation
-
-If you use FlashMemory in your research, please cite:
-
-```bibtex
-@article{wang2026flashmemory,
-  title   = {FlashMemory-DeepSeek-V4: Lightning Index Ultra-Long Context via Lookahead Sparse Attention},
-  author  = {Yan Wang and Qifan Zhang and Jiachen Yu and Tian Liang and Dongyang Ma and
-             Xiang Hu and Zibo Lin and Chunyang Li and Zhichao Wang and Jia Li and
-             Yujiu Yang and Haitao Mi and Dong Yu},
-  year    = {2026},
-  journal = {arXiv preprint arXiv:2606.09079},
-  url     = {https://huggingface.co/papers/2606.09079},
-}
-```
-
-## License
-
-MIT
+---
+license: mit
+library_name: sglang
+pipeline_tag: text-generation
+tags:
+  - deepseek-v4
+  - long-context
+  - kv-cache
+  - sparse-attention
+  - retriever
+  - sglang
+base_model: deepseek-ai/DeepSeek-V4
+arxiv: 2606.09079
+---
+
+# FlashMemory-DeepSeek-V4
+
+**FlashMemory** is a trained **Memory-Indexer** retriever for **DeepSeek-V4 Compressed-Sparse-Attention (CSA)** KV-cache. Instead of scoring the full history every decode step, it predicts which ~10–15% of chunks the next 64 tokens will attend to — only those stay on GPU, the rest are offloaded to CPU.
+
+- 📄 **[Paper (arXiv:2606.09079)](https://arxiv.org/abs/2606.09079)**
+- 🤗 **[Model weights (this repo)](https://huggingface.co/libertywing/FlashMemory-Deepseek-V4)**
+- 💻 **[Code on GitHub](https://github.com/libertywing/FlashMemory-Deepseek-V4)**
+
+## Highlights
+
+FM-DS-V4 **matches or beats the DS-V4-Flash full-attention baseline** on long-context accuracy while dramatically cutting serving cost:
+
+- At **1M context**, GPU KV cache shrinks by **~90%** (3.73 → 0.37 GB).
+- Per-decode-token compute drops from **118.9 → 35.4 GFLOP (0.30×)**.
+- **2.7× aggregate throughput**, and max concurrency rises from **11 → 40 (3.6×)**.
+
+| Context | KMAX | Ours conc / throughput | Baseline conc / throughput | Gain conc | Gain throughput |
+|---------|------|------------------------|----------------------------|-----------|-----------------|
+| 256K | 96 | 76 / ~2759 tok/s | 47 / ~1584 tok/s | 1.6× | 1.7× |
+| 512K | 192 | 60 / ~2008 tok/s | 25 / ~1028 tok/s | 2.4× | 2.0× |
+| 1M | 384 | 30 / ~1266 tok/s | 11 / ~455 tok/s | 2.7�� | 2.8× |
+
+On LongBench-v2 / LongMemEval / RULER it matches or exceeds the full-attention baseline while keeping only ~10–15% of CSA KV on GPU. Ablations (**Recency-10%**, **Random-10%**) with the same KV budget confirm the gains come from **learned relevance**, not a positional/budget artifact.
+
+## How it works (in brief)
+
+Two indexers cooperate:
+
+- **Level 1 — Memory Indexer (every 64 steps):** a trained retriever scores the history's compressed keys and selects query-critical chunks → a `resident_set` recalled from CPU to GPU.
+- **Level 2 — Lightning Indexer (every step):** the native top-512 runs confined to the `resident_set`, at full speed inside cuda-graph.
+
+The GPU/CPU split is the key: only recalled chunks + compressed indexer keys live on GPU; the bulk KV cache sits on CPU and is pulled on demand.
+
+> For the full architecture, retriever math, hyperparameters and training details, see the **[paper](https://arxiv.org/abs/2606.09079)** and the **[GitHub code](https://github.com/libertywing/FlashMemory-Deepseek-V4)**.
+
+## Inference (PD-disaggregated)
+
+Three servers — **P** (prefill), **D** (decode, all offload/recall), **router**. Clients hit the router at `http://<router>:31503/v1/chat/completions`.
+
+```bash
+pip install -e sglang/python   # also: pip install sgl_kernel==0.3.21
+# download checkpoints/ from this HF repo
+
+# D (decode) — GPU 0–7, 512K example:
+TGT_CONC=60 TGT_CTX=524288 CTX_LEN=1100000 bash launch_decode.sh
+# P (prefill) — GPU 0–7:
+CTX_LEN=1100000 SWA_RATIO=0.1 HOST=<P_IP> bash launch_prefill.sh
+# router:
+PREFILL_IP=<P_IP> DECODE_IP=<D_IP> bash launch_router.sh
+```
+
+All optimizations are env-gated (gate-off = exact DS-V4 baseline). Key switches:
+`SGLANG_DECODE_SWAP_P`, `SGLANG_PATHP_INDEX_K_OFFLOAD`, `SGLANG_PATHP_SCORE_RESIDENT`,
+`SGLANG_PATHP_PAGE_RECALL`, `SGLANG_PATHP_CUDAGRAPH`, `SGLANG_PATHP_ASYNC_RECALL`,
+`SGLANG_PATHP_FUSED_REMAP`.
+
+## Checkpoint
+
+Download `checkpoints/` into the repo root (default `top3_R930_joint.pt`, CSA layers 10/12/20).
+
+## License
+
+MIT
+
+## Citation
+
+```bibtex
+@article{wang2026flashmemory,
+  title   = {FlashMemory-DeepSeek-V4: Lightning Index Ultra-Long Context via Lookahead Sparse Attention},
+  author  = {Yan Wang and Qifan Zhang and Jiachen Yu and Tian Liang and Dongyang Ma and
+             Xiang Hu and Zibo Lin and Chunyang Li and Zhichao Wang and Jia Li and
+             Yujiu Yang and Haitao Mi and Dong Yu},
+  year    = {2026},
+  journal = {arXiv preprint arXiv:2606.09079},
+  url     = {https://arxiv.org/abs/2606.09079},
+}
+```