libertywing/FlashMemory-Deepseek-V4

来源：https://github.com/libertywing/FlashMemory-Deepseek-V4

FlashMemory DS-V4 Retriever

一个轻量级检索器，用于稀疏化 DeepSeek-V4 压缩稀疏注意力 (CSA) KV 缓存。

给定解码 token 的隐藏状态，该检索器可以预测接下来的约 64 个 token 将关注哪些 CSA KV 缓存块。只有得分最高的块保留在 GPU 上；其余部分可以卸载到 CPU/磁盘。在 downstream 评估中，它在保持 约 10–15% 的 KV 缓存 在设备上的同时，达到或超过了全注意力基线。

模型权重已上传至 Hugging Face (https://huggingface.co/libertywing/FlashMemory-Deepseek-V4)

快速开始

pip install torch safetensors

# 使用模拟输入进行演示
python demo.py --ckpt weights/flashmemory_ds_v4.safetensors

# 玩具级稀疏解码循环
python toy_flashmemory_inference.py --ckpt weights/flashmemory_ds_v4.safetensors

使用方法

from retriever import FlashMemoryRetriever

model = FlashMemoryRetriever.from_checkpoint(
    "weights/flashmemory_ds_v4.safetensors", device="cuda"
)

# hidden:     [B, 4096] 解码 token 的隐藏状态
# comp_k:     [B, N, 132] uint8 类型的压缩 CSA 键
# positions:  [B] int64 类型的 token 位置

# 逐层 sigmoid 分数：{"l10": [B,N], "l12": [B,N], "l20": [B,N]}
per_layer = model(hidden, comp_k, positions)

# 跨层集成（mode="max" 或 "mean"）
scores = model.ensemble(hidden, comp_k, positions, mode="max")       # [B, N]

# 布尔保留掩码
keep = model.select_topk(hidden, comp_k, positions, top_k=512)       # 前 K 个
keep = model.select_topk(hidden, comp_k, positions, threshold=0.5)   # 阈值

compressed_k 格式

每个块 = HEAD_DIM + 4 = 132 个 uint8 字节：

反量化：fp8_values.view(float8_e4m3).float() * scale。
参见 demo.py 中的 make_mock_compressed_k()。

架构

每个 CSA 层，分数计算如下：

hidden [B, 4096]
    → wq_a        (4096 → Q_LORA_RANK)
    → RMSNorm     (q_norm_weight, eps=1e-6)
    → wq_b        (Q_LORA_RANK → N_HEADS * HEAD_DIM)
    → reshape     [B, N_HEADS, HEAD_DIM]
    → RoPE        (YaRN, 最后 ROPE_DIM=64 维, base=160000)
    → Hadamard    (归一化的 Walsh-Hadamard 变换)
    → q           [B, N_HEADS, HEAD_DIM]

hidden [B, 4096]
    → weights_proj (4096 → N_HEADS)
    → × weight_scale  (= HEAD_DIM^-0.5 * N_HEADS^-0.5)
    → fused_w     [B, N_HEADS]

compressed_k [B, N, HEAD_DIM + 4] (uint8)
    → bytes[:HEAD_DIM]  视为 float8_e4m3 → 反量化
    → × bytes[HEAD_DIM:]  视为 float32   → k [B, N, HEAD_DIM]

score = sigmoid( sum_heads( relu(k @ q^T) * fused_w ) )   in [0, 1]

联合检查点 + 集成

检查点包含 三个独立的 CSA 层（l10、l12、l20），每层有自己的权重。在推理时，逐层的 sigmoid 分数会按块集成 — 使用 max（并集，默认）或 mean — 以产生单个保留/丢弃决策。

超参数

玩具级推理参考（toy_flashmemory_inference.py）

一个自包含的示例，说明检索器如何在解码过程中驱动内存召回 — 即 DeepSeek-V4-FlashMemory 内部使用的实际控制流。

推理流程

 ┌──────────┐  压缩并存储       ┌────────────────────────────┐
 │ PREFILL  │  历史 K/V         │  CSA KV 缓存（内存）      │
 │ (密集    │ ──────────────────►  N 个压缩块，              │
 │  注意力) │                    │  每个 = [132] uint8 fp8-K │
 └────┬─────┘                    └──────────────┬─────────────┘
      │ 最后隐藏状态                              │ 每 64 步评分一次
      ▼                                           │
 ┌──────────────────── 解码循环 ──────────────────┼──────────────────┐
 │ 对于每个解码步骤 t：                             │                  │
 │   hidden = toy_decoder.step(token, keep_mask)   │  (稀疏内存注意力) │
 │                                                  │                  │
 │   每 RETRIEVAL_INTERVAL (= 64) 步:                 ▼                  │
 │     scores[N]   = retriever.ensemble(hidden, compressed_k, pos)      │
 │     keep_mask[N] = top-K（或 sigmoid > thresh）的分数                  │
 │     -> 未选中的块在接下来的 64 步中被掩码为 -inf                         │
 └─────────────────────────────────────────────────────────────────────┘

1. 预填充（密集）。短提示通过密集内存注意力运行。其最后隐藏状态为第一个检索周期提供种子。
2. 解码循环。玩具解码器每步产生一个 [B, 4096] 隐藏状态。
3. 检索周期（每 64 步）。真正的 FlashMemoryRetriever 对所有 N 个压缩 K 块进行评分，集成逐层分数，选择保留的块。
4. 稀疏注意力。未选择的块的注意力 logits 被设置为 -inf。

此模拟的内容

• 此玩具示例不执行 真实的 CPU↔GPU KV 缓存传输。交换引擎是 FlashMemory 内部基础设施的一部分，未包含在此处。
• 我们通过掩码注意力 logits 为 -inf 来模拟内存召回。被掩码的块对注意力没有贡献 — 效果等同于未加载其 KV。
• 目的是使解码时的控制流具体化。

它是 / 不是

生产版本依赖于内部的 sglang + DeepSeek-V4 CSA 框架（原生 FP8 索引器、真实压缩 KV 缓存、注意力下沉、阈值回退、按请求路由、实际 KV 交换）并且无法发布。

Downstream 评估

FlashMemory DS-V4 在推理密集的长上下文任务上达到或超过全注意力基线，同时仅保留 约 10–15% 的 CSA KV 缓存 在设备上：

精确的针检索任务（MRCR）需要额外的阈值回退机制 — 这不属于独立发布的一部分。

文件

许可证

MIT

引用

如果在你的研究中使用 FlashMemory，请引用：

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