KV cache 与前缀复用 · CS2680 学习手册

为什么保留 K 和 V

在自回归 Transformer 中，新 token 的 query 需要和历史 token 的 key 比较，再用 attention 权重组合历史 value。如果每生成一个 token 都重新计算全部历史 K、V，就会反复做同一部分工作。KV cache 把每一层已处理 token 的 key 和 value 保存下来；下一步只计算新增 token 的状态，再读取历史缓存。

缓存是运行状态，不是模型权重，也不是完整答案的复用。它会随有效上下文增长，并在多请求并发时占用大量内存。通常不保存历史 query，因为未来 token 需要的是历史 K、V。这里讨论标准 Transformer attention 的逻辑容量；滑动窗口、混合状态模型、压缩注意力等结构会改变公式，不能直接照搬。

把每个维度乘出来

对批中长度相同、没有共享的 B 条序列，设层数 L、序列长度 T、KV head 数 Hkv、每个 head 维度 D、每个元素字节数 s，则 KV bytes = 2 × B × L × T × Hkv × D × s。前面的 2 分别代表 K 和 V。长度不同时，把 B×T 换成各序列长度之和；上下文长度包含已保留输入与已生成 token。

MHA 通常每个 query head 对应一个 KV head，因此 Hkv 等于 query head 数；GQA 让多个 query head 共享一组 K、V，Hkv 更少。假设 B=1、L=32、T=4096、D=128、s=2，MHA 的 Hkv=32 时逻辑 KV 为 2 GiB；GQA 的 Hkv=8 时为 0.5 GiB。这是同一组假设下的四倍差异，不代表可以无损修改任意模型的 head 配置。

逻辑容量与实际显存不同

上面的公式不包含权重、激活、临时工作区、内存分配器保留空间和元数据。多卡部署也不能总是直接除以卡数，head 分片方式和复制策略会影响每卡容量。规划时先给模型与运行时留空间，再估算可容纳的活跃 token 总数，而不是只报“支持多少并发请求”。十个短请求和十个超长请求的压力可能相差很大。

如果按每条请求的最大长度一次性预留连续空间，短请求会浪费容量，长短混合还可能产生碎片。PagedAttention 将逻辑 KV 序列映射到固定大小的物理块，允许按需分配和非连续存放。它主要改善内存管理与可共享性，并没有让每个 token 的 K、V 自动变小。最后一个未填满的块仍会产生部分浪费。

还要分清缓存容量和 attention 的临时计算。标准全量 attention 在处理长输入时涉及大量 token 对之间的关系，而 KV cache 存的是每个历史位置的向量，按本章假设随长度线性增长。采用节省 attention 中间存储的内核，并不等于不再需要历史 KV。排查内存时分别列出持久缓存和本次执行的临时空间，才能判断应该减少活跃上下文，还是更换执行方法。

相同开头，才能直接复用历史

自动前缀缓存跨请求保留可复用的 KV 状态。当模型配置和 token 前缀一致时，新请求可以跳过已经计算的前缀 prefill。实践中适合复用固定系统说明、工具定义、公共文档或稳定的对话历史；把每次变化的时间戳放在最前面，则可能让共享前缀很早中断。缓存键还需正确隔离模型、适配器及租户语义。

文字含义相近不等于 token 前缀相同。调整空格、模板或工具顺序都可能影响命中。前缀缓存主要节省 prefill，不会替你生成剩余答案，也不会消除 decode 对历史 KV 的读取。十次重复请求的热缓存结果不能代表首次访问；若内存紧张发生驱逐，之前出现过的前缀也可能不再可用。

本机实验：容量与前缀敏感性

保存为 kv_lab.py 并运行。示例用词编号模拟 token，只验证共同前缀的语义，不实现推理引擎的块哈希和缓存驱逐。

def kv_bytes(layers, tokens, kv_heads, head_dim, size=2, batch=1):
    return 2 * batch * layers * tokens * kv_heads * head_dim * size

def common_prefix(a, b):
    count = 0
    for x, y in zip(a, b):
        if x != y:
            break
        count += 1
    return count

mha = kv_bytes(32, 4096, 32, 128)
gqa = kv_bytes(32, 4096, 8, 128)
print("MHA GiB", mha / 2**30)
print("GQA GiB", gqa / 2**30)
assert mha == 4 * gqa
stable = [10, 11, 12, 13]
a = stable + [100, 101]
b = stable + [200]
print("shared tokens", common_prefix(a, b))
print("timestamp first", common_prefix([1] + a, [2] + b))
assert common_prefix(a, b) == 4
assert common_prefix([1] + a, [2] + b) == 0

结果应是 2.0 GiB、0.5 GiB、共享 4 个 token，以及动态前缀下共享 0 个 token。把长度或并发翻倍，逻辑 KV 同样翻倍。把字节数从二改为一只是在做容量估算；真正的 KV 量化还有数值误差、缩放因子和实现支持问题。