FlashMemory-Deepseek-V4

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)
• 🤗 Model weights (this repo)
• 💻 Code on GitHub

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×).

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 and the GitHub code.

Inference (PD-disaggregated)

Three servers — P (prefill), D (decode, all offload/recall), router. Clients hit the router at http://<router>:31503/v1/chat/completions.

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

@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},
}