KVCache
大模型推理
Prefill + Decode 分离
KVCache解决的是什么问题?
回顾Transformer的Decoder结构,在decode阶段,模型进行推理时,预测t位置的token,会用到前t-1个token,序列token的key、value向量会被重复计算,KVCache通过缓存已计算token的key、value向量降低冗余计算。
实现
缓存结构(Pytorch示意)
kv_cache = { layer_0: { "key": Tensor [seq_len, num_heads, head_dim], "value": Tensor [seq_len, num_heads, head_dim] }, ... }
标准推理 vs KVCache
| 特性 | 标准推理 | KVCache缓存 |
|---|---|---|
| 每词计算量 | 每个词都重复相同计算 | 复用历史计算,结果更快 |
| 内存占用 | 每步内存占用少,但文本越长内存增长 | 额外占用内存存储历史信息,但整体更高效 |
| 速度 | 文本越长越慢,因为重复计算 | 避免重复计算,长文本也能保持高速 |
| 效率 | 计算成本高,响应慢 | 记忆历史计算,效率更高,响应更快 |
| 长文本处理 | 长文本时因重复计算容易卡顿 | 适合长文本,能记住历史步骤,处理流畅 |
KV Cache优化
有限的GPU显存条件下,超长上下文推理的KVCache很容易成为显存瓶颈。
标准KV Cache存在以下几种内存浪费:
内碎片(Internal Fragmentation)
预留浪费(Reservation Waste)
外碎片(External Fragmentation)
| 类型 | 定义 | 原因 | 常见场景 | 解决方式 |
|---|---|---|---|---|
| Internal Fragmentation | 预分配空间未被用到 | 估算过长 | 设定 max_length 过大 | 动态分配、估算优化 |
| Reservation Waste | 为未来 token 预留 | 保证连续写入 | 每步生成时存在 | 按需动态增长 |
| External Fragmentation | 多请求之间出现空洞 | 长度不一致 | 多用户并发推理 | 分块调度、defrag |
Paged KV Cache
通过将传统连续缓存(KV Cache)拆分成可管理的小块(pages / blocks),显著提升了内存利用率与推理效率。
MQA / GQA
MQA
每个Head有一个独立的Query,所有Head共享一组Key、Value。减少Key、Value的内存占用和带宽开销。
GQA
每个Head有一个独立的Query,所有Head划分为G组,每组共享一组Key、Value。GQA 是介于 MHA 与 MQA 之间的 中间折中方案。
参考
Transformers KV Caching Explained
KV Caching Explained: Optimizing Transformer Inference Efficiency
Paged KV Cache
PQCache: Product Quantization-based KVCache for Long
Context LLM Inference
H2O: Heavy-Hitter Oracle for Efficient Generative
Inference of Large Language Models