ds4.c 学习笔记

Appearance

缓存注意力机制的历史 Key/Value，避免每步重复计算，是自回归推理的核心优化。

为什么需要 KV Cache ​

自回归生成每步产生一个 token，下一步需要用所有已生成的 token（包括历史）做注意力计算。如果不缓存，第 N 步要重新计算 N 组 K 和 V，总计 1+2+3+…+N = N(N+1)/2 次 KV 计算。生成 1000 个 token 就是约 50 万次重复计算。

缓存后每步只算 1 组新 KV，总计 N 次。

核心原理 ​

步骤 1: 输入 [A]       → 算 K₁V₁, 缓存 → 输出 token B
步骤 2: 输入 [B]       → 算 K₂V₂, 追加  → 用 K₁V₁K₂V₂ 做注意力 → 输出 token C
步骤 3: 输入 [C]       → 算 K₃V₃, 追加  → 用 K₁V₁K₂V₂K₃V₃ 做注意力 → 输出 token D

ds4.c 的 KV Cache 由两部分组成：

类型作用大小
原始 KV（raw）最近 128 个 token 的完整 KV每层 128 × 512 × 4B = 256KB
压缩 KV（compressed）128 之前的历史，经 MLA 压缩按 ratio 缩减，可低至原始的 1/64

窗口满时通过 memmove 滑动移除最旧行，保证固定内存占用。

在 ds4.c 中的实现 ​

c
// KV cache push（ds4.c 约 7273 行）
kv_cache_push_raw(cache, scratch->kv);

关键函数：

  • kv_cache_push_raw() — 追加新 KV 行，窗口满时滑动
  • kv_cache_push_comp() — 追加压缩 KV 行
  • layer_attention_rows_one() — 在 raw + compressed 上混合做注意力

DeepSeek 使用 Multi-head Latent Attention (MLA)，KV 通过低秩投影压缩存储，且所有 Q 头共享 1 组压缩 KV，KV Cache 内存比标准 Multi-Head 大幅缩小。

相关概念 ​

  • softmax — 注意力分数的归一化
  • moe — MoE 层和注意力层各自使用 KV Cache
  • rope — K 在缓存前需先经 RoPE 旋转编码位置
  • mmap — KV Cache 可持久化到磁盘，复用历史前缀

磁盘持久化 ​

KV Cache 可以序列化到磁盘文件（<sha1>.kv），下次请求时用文本前缀匹配复用历史计算。持久化逻辑已从 ds4_server.c 提取为独立模块 ds4_kvstore.c/h。

磁盘缓存有容量预算（--kv-cache-budget-mb），满了通过评分公式驱逐：

score = (effective_hits + 1) × tokens / file_size

effective_hits 按 6 小时半衰期衰减，防止过时的热门条目挤占新条目。驱逐时保护刚写入的文件不被立即删除。

预存储驱逐 ​

驱逐在写入新条目之前执行，并传入待写入条目的上下文（model_id、quant_bits、ctx_size、text），避免新条目成为自己的受害者。当新条目是旧条目的严格前缀扩展时，旧条目自动获得评分惩罚。

兼容性加固 ​

加载磁盘缓存时严格校验 model_id、quant_bits、ctx_size 以及文件内容的 SHA1 完整性。不兼容的文件自动删除重建，避免不干净关机后缓存"看似正常但 prefill 总失败"的循环。

分布式 KV 快照 ​

分布式推理模式下，KV 快照是拓扑无关的：保存时聚合所有 worker 的层张量为标准 DSV4 载荷，加载时按当前路由分发。保存的文件在任何 worker 配置下都可复用。

Metal 后端支持 comp_kv_f16 模式——压缩 KV cache 以 float16 存储，每行内存减半，详见 Part 6 — GPU 加速。

详见 Part 5 — KV Cache 持久化。

详见 Part 3 — 注意力机制 和 端到端推理流程。

出链