修炼AI技能
AI 教程 / 大模型应用 / KV 量化：显存再压缩

KV 量化：显存再压缩

自回归生成时，模型要缓存每个已生成 token 的键值（KV），这部分 KV cache 会随上下文长度线性增长。KV 量化（KV Cache Quantization）把这些缓存从 FP16 压到 INT8 甚至更低，从而省下显存、提升可服务的上下文长度与并发。它是与权重量化互补的另一条压缩路径。

是什么

注意力计算需要历史所有 token 的 K 与 V。把它们以更低位数存储，读取时再反量化参与计算，就是 KV 量化。它与权重量化（如 AWQ、GPTQ）是正交的两件事：一个压权重，一个压缓存。两者可叠加，但从质量与收益上需分别评估。

为什么重要

长上下文场景下，KV cache 常常比模型权重更占显存。例如一个长文档问答服务，上下文越长，KV 占用越大，可能成为显存瓶颈。量化后同样显存能容纳更长上下文或更高并发，对推理成本影响直接。

怎么做

在 vLLM 中可指定 KV cache 的数据类型，例如用 FP8 量化：

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --kv-cache-dtype fp8_e5m2

更激进的方案需结合专用量化内核，并关注精度损失。部分框架还支持按层或按通道选择不同精度。

注意点

  • 量化位数越低，精度损失越大，需以任务效果实测为准。
  • 并非所有层都同样敏感，部分研究对异常值通道做特殊处理。
  • KV 量化与权重量化可叠加，但叠加后需整体评估质量。
  • 某些算子对低精度支持不全，需确认硬件与内核可用性。

实践建议

开启 KV 量化前，先在目标任务上测基线精度，再逐步降低位宽，观察质量拐点。并非所有层同样敏感，可优先对不敏感层做更激进量化。KV 量化与权重量化正交，可叠加但需整体评估。注意部分硬件对低精度内核支持不全，需确认算子可用性。长上下文服务中，KV 量化往往比权重量化带来更直接的并发提升，值得优先验证。建议建立量化档位对照表，记录每位宽下的显存节省与准确率变化，方便按场景切换。若精度下降明显，可只对键值中某一侧量化，或保留部分高精度层。生产环境应保留回退到 FP16 的开关。对于显存受限的边缘部署，KV 量化常常是先于权重量化的首选优化，应作为默认开启项评估。同时监控量化后的困惑度变化，把它纳入上线门禁指标，防止精度悄然越界。

小结

KV 量化通过压缩注意力缓存的位宽，缓解长上下文带来的显存压力，与权重量化互补。落地时要权衡压缩率与生成质量，并以实际任务评测验证。

参考与延伸阅读

本文累计阅读 — 次