KV cache 原理

阶段: B2 · attention/decoder/KV + judge 校准（Day 11-20） 标签: #kv-cache #autoregressive #inference #memory-bandwidth

今日导引（由浅入深）

Day 11-13 我们把 decoder block 写对、验对了（causal + 数值对齐）。但「正确」之外，自回归生成还有「快」的问题——朴素实现每生成一个 token 都重算整条前缀的注意力，是 O(n²) 的纯浪费。今天 Day 14 引入推理侧第一个、也是最基础的优化：KV cache——把历史 token 的 K/V 存下，每步只算新 token，把每步注意力从 O(n) 总量重算降到 O(n) 单步。这是从「能正确生成」到「能高效生成」的第一步，也是 B12（PagedAttention/MLA/Mooncake）整条推理优化链的源头。它紧接 Day 13，因为 cache 加速的硬约束是「数值结果不变」——正是昨天对齐能验收的；它通向明天 Day 15 的采样策略（cache 让逐 token 采样变得划算）。今天的「最小可判定产出」：cache on/off 的逐 token 提速倍数 + 一行 KV 显存实算（GB）。

1. 机理精读

朴素自回归的浪费：生成第 n 个 token 时，注意力需要全部 1..n 的 K/V。朴素实现每步把整条前缀重新过一遍 Q/K/V 投影和 attention——但历史 token 的 K/V 在每一步都不变（它们只依赖各自的输入与权重，不依赖后来的 token）。于是第 n 步重算了 1..n-1 的 K/V，纯属冗余。整段生成的总计算因此是 O(n²)。

KV cache 的核心 trick：把每个已生成 token 的 K、V 向量缓存下来；生成新 token 时，只对这一个新 token算它的 Q/K/V，把新 K/V append 到 cache，然后用「新 Q」对「cache 里全部 K/V」做一次注意力。这样每步的新增计算从「重算整条前缀」降为「只算一个 token + 一次对 n 长 cache 的注意力」，单步注意力 O(n)、整段从 O(n²) 的重算降为线性增量。代价是用显存换计算。

KV 显存怎么算：cache 要为每层、每个 token、K 和 V 各存一份 d 维向量。公式：

KV 显存 ≈ 2 · L · n_layer · d · dtype_bytes

其中 2 = K 和 V 两份，L = 序列长度，n_layer = 层数，d = 隐藏维，dtype_bytes = 每元素字节（fp16=2）。seed 给的实算例：

L=1024, n_layer=12, d=768, fp16(2B)
→ 2 · 1024 · 12 · 768 · 2 B ≈ 0.036 GB

逐步核对这个数：2·1024·12·768·2 = 37,748,736 B；÷ 1024³ ≈ 0.0352 GB ≈ 0.036 GB ✓。

这只是单条 1024-token 序列、12 层小模型。换成一个真实大模型量级（仍按 MHA、未压缩）感受一下膨胀：

最后一行单条序列就要 256 GB KV——远超单卡显存。再叠 batch 并发，KV cache 会膨胀到「比模型权重还大」——这就是为什么 KV 显存是长上下文推理的头号瓶颈，催生了下面 SOTA 段的一系列压缩/分页方案（GQA/MQA 减 KV 头数、MLA 压缩、PagedAttention 防碎片）。注意上表是 MHA 估算；GQA（多 query 头共享少量 KV 头）能把 KV 项除以「query头/KV头」的比值，是当代标配的省显存手段。

prefill vs decode 两阶段：自回归生成天然分两段——(1) prefill：把整段 prompt 一次性喂入、并行算出全部 token 的 KV 填满 cache（计算密集、可高度并行）；(2) decode：逐 token 生成，每步只算 1 个新 token、读全量 cache（访存密集、并行度低）。KV cache 让 decode 阶段从「每步重算前缀」变成「每步读 cache + 算 1 token」。这两阶段的算力/访存特征截然不同，正是 Mooncake 等做 prefill/decode 分离调度的动机——今天先认识这个二分。

decode 是访存瓶颈，不是算力瓶颈：decode 每步只算 1 个 token 的少量矩阵乘，但要把整个 KV cache（可能几十 GB）从显存读进计算单元。算得快没用，喂不饱——decode 吞吐由显存带宽而非 FLOPs 决定。这就是为什么省 KV 体积（GQA/MLA/量化 KV）能直接提速：读得少 = 跑得快。

cache 不许改变数值：开/关 cache 必须给出逐元素相同的输出（在浮点容差内）——cache 只是「不重算已知不变的量」，数学上是恒等优化。这正好用 Day 13 的对齐方法验收：abs(out_cache - out_nocache) < 容差。若不一致，说明 cache append 顺序或 mask 处理有 bug。参考 Karpathy nanoGPT / llm.c 的 KV cache 讲解。

2. 推导 / 手算 / 代码走读

seed 计划给 decoderBlock.ts 加 kvCache 参数、写 scripts/bench-kv.ts。当前仓库：scripts/bench-kv.ts 不存在（Glob 核实），且 src/agent/transformer/transformer.ts 的 decoderBlock(x, p) 当前没有 kvCache 参数（签名只有 (x: Matrix, p: BlockParams)，第 111 行）——所以 KV cache 改造属 seed 标注的「待建」，下文为设计走读，不编造不存在的参数。

KV cache 改造的最小设计（待建）：

1. 给 decoderBlock 增可选第三参 kvCache?: { K: Matrix; V: Matrix }（每层一份历史 K/V）。
2. 逐 token 生成模式：输入只是新 token 的 [1, d]，算其 Q/K/V；把新 K/V append 进 kvCache.K/kvCache.V（变 [n, d]）。
3. 注意力改成「Q_new([1,d]) · cache_K^T → softmax → · cache_V」；因为 Q 只有 1 行、且对全部历史可见，causal mask 在增量模式下自动满足（新 token 本就该看全部历史）。
4. bench-kv.ts：对一条 N-token 序列，分别跑 cache-off（每步喂整条前缀）与 cache-on（每步只喂新 token），计时每 token，输出提速倍数；并用上面公式打印这条序列的 KV GB。

手算提速直觉（为何会快）：cache-off 第 k 步注意力算 k×k 的 scores；总和 Σ k² = O(N³/3)（含投影则 O(N²·d) 级重算）。cache-on 第 k 步只算 1×k；总和 Σ k = O(N²/2)。比值随 N 增大而拉开——这就是为什么 seed 预期一个明显倍数（如示例 3.2×，该数为待实测占位，非已测）。

把比值写细一点（只看 scores 计算量）：

• cache-off 总量 ∝ Σ_{k=1..N} k² = N(N+1)(2N+1)/6 ≈ N³/3。
• cache-on 总量 ∝ Σ_{k=1..N} k = N(N+1)/2 ≈ N²/2。
• 理论比值 ≈ (N³/3)/(N²/2) = 2N/3，随 N 线性增长。

注意理论比值（∝N）与实测倍数会差很多：实测还含投影、FFN、Python/JS 解释开销、内存分配等固定成本，且小 N 时这些常数项占主导，所以实测倍数（如 3.2×）远低于理论 2N/3。这正是为什么要实测而非套公式——seed 把 3.2× 明确标为「先占位待跑」。

逐步追一条 4-token 生成（看 cache 怎么长、计算怎么省）：

step 1: 输入 token t1
        cache-off: 算 K/V for [t1]            scores 1×1
        cache-on : 算 K/V for [t1], append   cache=[t1]      scores 1×1
step 2: 输入 token t2
        cache-off: 重算 K/V for [t1,t2]       scores 2×2  ← 重算了 t1
        cache-on : 只算 K/V for [t2], append  cache=[t1,t2]  scores 1×2  ← t1 复用
step 3: 输入 token t3
        cache-off: 重算 K/V for [t1,t2,t3]    scores 3×3  ← 重算 t1,t2
        cache-on : 只算 K/V for [t3], append  cache=[t1,t2,t3] scores 1×3
step 4: 同理 cache-off 4×4，cache-on 1×4

cache-off 的 scores 计算量 = 1+4+9+16 = 30（∝Σk²）；cache-on = 1+2+3+4 = 10（∝Σk）。本例比值 3×——与 seed 占位的 3.2× 同量级（但实际倍数仍需实测，含投影/FFN/解释器开销）。每步 cache 只 append 一行新 K/V，历史从不重算——这就是 O(n²)→线性增量的来源。

3. 今日实战

seed：「给 decoderBlock.ts 加可选 kvCache 参数，token-by-token 生成时复用历史 K/V。写 scripts/bench-kv.ts 量 cache on/off 的逐 token 生成耗时，并用公式实算一条序列的 KV GB。」可执行步骤（待建）：

1. 在 transformer.ts 给 decoderBlock 加 kvCache 可选参数 + 增量 attention 路径（保留无 cache 路径不变，便于对齐验收）。
2. 加单测：cache-on 与 cache-off 对同序列逐 token 输出 toBeCloseTo(..., 9)——用 Day 13 对齐精神证明「加速不改数值」。
3. 新建 scripts/bench-kv.ts：循环 N 个 token，两种模式各 performance.now() 计时，打印 speedup = t_off / t_on。
4. 同脚本里用 2·L·n_layer·d·dtype_bytes 打印一行 KV GB（先用 seed 的 L=1024,n_layer=12,d=768,fp16 → ≈0.036 GB 验证公式实现对）。

验收顺序（先正确、后提速，呼应 Day 13）：

1. 数值等价单测先过：cache-on 与 cache-off 对同序列逐 token 输出 toBeCloseTo(..., 9)。不过这一关，提速数字毫无意义（你可能在「更快地算错」）。
2. 再测提速：等价性确立后才跑 bench-kv.ts 量倍数。
3. 再核显存公式：0.036 GB 这个可手算的小例子用来验证 GB 计算实现没写错单位（B/KB/GB 换算）。 这个「正确 → 性能 → 资源」的三段验收，是把 Day 13 的对齐纪律延用到优化场景。

4. 今日实测 / 产出

• 状态（按 seed）：seed 标注「待建/待跑」。目标产出：提速倍数（如 3.2×）+ 一行 KV GB 实算（例：L=1024,n_layer=12,d=768,fp16 → 2·1024·12·768·2 B ≈ 0.036 GB）。
• 诚实标注：scripts/bench-kv.ts 尚未创建，decoderBlock 尚无 kvCache 参数——提速倍数与 bench 报告未产出，保持「待建/待跑」。seed 已明确「倍数为实测，先占位待跑」——3.2× 是占位示例，非已测得的真实数。
• 可立即验证的部分：KV GB 公式 0.036 GB 是纯算术，可在任何环境核对（2·1024·12·768·2 = 37,748,736 B ≈ 0.0352 GB，与 seed 的 ≈0.036 GB 一致）。

bench 设计的两个公平性要点（待建脚本时要守）：

• 同 warmup、同次数：JIT/缓存预热会让首跑偏慢，两种模式都要先空跑几次再计时，取多次中位数。
• 隔离纯生成耗时：别把 prompt 构造、I/O、JSON 序列化算进去，只计 forward 那段，否则倍数被固定开销稀释。

5. 常见误区 / 陷阱

1. 以为 cache 能改变结果：cache 是恒等优化，输出必须逐元素一致；不一致 = bug，别当「精度抖动」放过。
2. 增量模式还套全 seq×seq mask：增量时 Q 只有新 token 一行、对全历史可见，不需要再 mask 未来（已无未来）。套旧 mask 会算错。
3. 低估 KV 显存膨胀：小模型 0.036 GB 没感觉，但 32 层 / d=4096 / 128K / 多并发会到几十甚至几百 GB，是长上下文真瓶颈。
4. 把朴素 cache 当生产终点：生产侧还要 PagedAttention 防碎片、MLA 压 KV、prefill/decode 分离——朴素 cache 只是起点。
5. 用理论比值 2N/3 冒充实测：理论忽略固定开销，小 N 下严重高估；必须实测且诚实标注（seed 的 3.2× 是占位）。
6. cache append 顺序错：新 token 的 K/V 必须按生成顺序追加，错位会让位置对应乱掉、注意力读错历史。

6. 学习资源（每条带 YYYY-MM）

• Karpathy nanoGPT / llm.c KV cache 讲解 (2023-01 起, 2024 维护)——seed 指定来源，朴素 cache 最清晰的实现参考。
• Kwon et al.《Efficient Memory Management for LLM Serving with PagedAttention》(vLLM, 2023-09, arXiv:2309.06180)——KV cache 显存碎片问题与分页解法。
• DeepSeek-V2/V3 Technical Report（MLA 部分, 2024-05 / 2024-12）——KV cache 压缩（Multi-head Latent Attention）的当代 SOTA。
• Mooncake (Moonshot/Kimi, 2024-06, arXiv:2407.00079)——prefill/decode 分离 + KV-centric 推理架构。
• 本仓 src/agent/transformer/transformer.ts（AICAP-180, 2026-06）——待加 kvCache 的目标 decoderBlock（当前签名无 cache 参数，属待建）。

SOTA检查 (2026-06 更新)

• 当前主流：KV cache 是所有推理引擎的基础设施，2026 仍是默认开启的必备优化。
• 是否仍 SOTA：朴素 KV cache 是「地基」不是「天花板」。2026 SOTA 在其上叠 PagedAttention（vLLM, 2023-09）防显存碎片、MLA 压缩（DeepSeek-V3, 2024-12）省 KV 体积、prefill/decode 分离（Mooncake, 2024-06）。教学用朴素 cache 没问题。三种省 KV 思路的定位：GQA（减 KV 头数，质量近 MHA）属「结构层省」，MLA（低秩压缩 KV）属「表示层省」，PagedAttention（分页防碎片）属「内存管理层省」——它们正交可叠加，朴素 cache 是它们共同的底座。
• 过时黑名单：避免宣称朴素 cache 是生产方案——显存仍是瓶颈；避免在增量模式重复套全序列 causal mask。
• 下次复查点：B2 末（Day 20）；B12 推理优化阶段用 vLLM PagedAttention / MLA / Mooncake 真实方案替换教学基线，并按当周版本号重验。

衔接

• 昨天：Day 13 — Reference 对齐（数值对齐验收 decoder block「算得对」，max abs diff < 1e-4）。
• 今天：在已验证正确的 block 上引入 KV cache，把自回归从 O(n²) 重算降为线性增量，且加速不改数值。
• 明天：Day 15 — 采样策略（cache 让逐 token 解码变划算后，研究 greedy/temperature/top-p/top-k 与 structured output）。