4.4 KV Cache 量化:KIVI 2-bit 与 FP8 KV Cache
长上下文时代 KV Cache 才是显存大头:KIVI 的 K-per-channel/V-per-token 非对称设计,以及 vLLM 中 FP8 KV Cache 的三种 scale 校准路线
第 4.1 节最后算过一笔账:当权重被 INT4 压到约 37GB 后,KV Cache 成了显存的下一个大头。KV Cache 有多大?以 Llama-2-70B 为例,每层 8 个 attention head、每 head 128 维,32 层,8K 上下文,单请求的 KV Cache 约 2GB;如果把上下文拉到 128K(现在的大模型都这么干),就是 32GB——比 INT4 后的整个模型还大。而且 KV Cache 是每个并发请求一份,200 个并发请求就是 6.4TB,根本没有 GPU 装得下。
KV Cache 量化就是在这个背景下成为刚需的。这一节先讲学术界的标杆方案 KIVI(2-bit,理解”怎么量化才不掉点”的分布学),再讲工业界的实际主流 FP8 KV Cache(vLLM 原生支持,怎么配、怎么校准)。
📑 目录
- 1. KV Cache 为什么成为量化目标
- 2. 先做分布研究:K 和 V 的性格不一样
- 3. KIVI:2-bit 的非对称量化方案
- 4. KIVI 的效果:2.6x 显存、3 倍吞吐
- 5. 工业界主流:FP8 KV Cache
- 6. vLLM 实战:三种 scale 校准路线
- 7. 敏感层跳过与量化组合拳
- 总结
- 自我检验清单
- 参考资料
1. KV Cache 为什么成为量化目标
先回顾第 1 章的 KV Cache 显存账本。对 MHA 模型( 层、 个 head、 维、上下文 、批量 、FP16):
关键特性:
- 与上下文长度成正比:长上下文(128K+)时 KV Cache 线性膨胀,权重是固定的、它不固定
- 与并发请求数成正比:每多一个并发请求就多一份完整 KV Cache,而 Decode 阶段要把并发请求拼成大 batch 才有吞吐——并发能力直接受 KV Cache 显存限制
- 它是 Decode 阶段反复读取的热数据:每生成一个 Token,整个序列的 KV 都要读一遍做 Attention(这也是 Decode Memory Bound 的第二大来源)
📌 关键点:量化 KV Cache 直接解决三件事——同一张卡能塞更多并发(吞吐↑)、能支持更长的上下文(能力↑)、Decode 每步读取的 KV 字节数变少(延迟↓)。所以 KV Cache 量化是”长上下文 + 高并发”场景下收益最大的单项优化。
2. 先做分布研究:K 和 V 的性格不一样
KIVI(ICML 2024,arXiv 2402.02750)开篇做了一件此前没人系统做过的事:仔细研究 KV Cache 里元素的分布,而不是想当然地套用一般量化方法。结论有三条:
- K 的分布是 per-channel 结构的:K Cache 按通道(hidden 维度)统计时,每个通道的数值范围有稳定的统计结构——某些通道天生偏大/偏小,这种”通道结构性”使得按通道(per-channel)量化 K 最有效。
- V 的分布是 per-token 结构的:V Cache 不同,它的数值范围主要随 Token 变化,同一个 Token 的各通道值相近——按 Token(per-token)量化 V 最有效。
- 对称的直觉是错的:如果”想当然”地对 K 和 V 都用 per-token 量化,INT4 没问题,但降到 INT2 精度崩;对 V 用 per-channel 量化则任何位宽都明显变差(表里最差的一档)。
一句话:K 按通道量化、V 按 Token 量化,两者必须”不对称”对待。这是 KIVI 名字的来源(KIVI = asymmetric KV cache quantization)。
💡 提示:背后的直觉可以这样理解——K 与权重类似,是”静态特征”的载体(通道决定特征方向),所以和权重量化一样适合 per-channel;V 更像”即时激活”(随输入内容逐 Token 变化),所以和激活量化一样适合 per-token。KV Cache 的量化设计,本质是”权重-激活”两种性格的混合体。
3. KIVI:2-bit 的非对称量化方案
3.1 工程约束:流式写入
KIVI 要面对一个 K/V 都没有的难题——KV Cache 是流式写入的:Prefill 一次写一批,Decode 每步追加一个 Token。per-token 的 V 好办(每个新 Token 的 V 量化后直接追加);但 per-channel 的 K 需要”攒够一组才能算 scale”——只来一个 Token 时,无法确定这一组(group)的统计量。
3.2 分块 + 残差:KIVI 的结构
KIVI 的解法是把 KV Cache 分成”分组部分”和”残差部分”:
- 分组部分(Grouped):按 group size 32 个 Token 一组,整组量化(K 按通道、V 按 Token),2-bit 存储
- 残差部分(Residual):最近的 个 Token(论文默认 128)保持 FP16 不量化——它们是”还没攒够一组”的尾巴
Decode 时,新 Token 先进残差区;残差区攒到 个后,整体量化并并入分组区,残差区清空重来。这样一个流式场景就被平滑地转成了”批量量化”。
3.3 为什么残差区有用
除了工程上的”攒组”需求,残差区还有一个精度收益:最近 个 Token 对当前 Attention 的贡献最大(它们在 softmax 里与当前 query 的位置最相关),保持全精度相当于给模型留了一个”高精度滑窗”。论文在 GSM8K 等困难任务上验证:没有残差区(R=0)时 2-bit 精度明显下降,R=128 时与 FP16 几乎持平。
3.4 配置参数
group_size = 32:量化分组大小(论文默认,32/64 效果接近,128 会明显变差)residual_length = 128:残差窗口长度- K/V 各 2-bit 或 4-bit 可配
4. KIVI 的效果:2.6x 显存、3 倍吞吐
论文在 Llama-2、Falcon、Mistral 上的实测:
| 指标 | FP16 KV Cache | KIVI 2-bit |
|---|---|---|
| 峰值显存(含权重) | 基线 | 减少 2.6x |
| 可支持的最大 batch | 基线 | 最多 4 倍 |
| 端到端吞吐 | 基线 | 2.35x ~ 3.47x |
| 模型质量(困惑度/下游任务) | 基线 | 几乎持平 |
注意两个细节:
- 2.6x 是包含权重的整体峰值显存下降(不只是 KV 部分),因为 KV 部分下降了约 8 倍
- 吞吐提升来自”batch 变大”(并发能力 ↑)而不是单请求变快——KIVI 的 2-bit 数据在 GPU 上仍然是按块读取、解包计算,单步延迟改善有限,它的主战场是显存和并发
📌 关键点:KV Cache 量化的收益模式与权重量化不同——权重量化直接减单请求延迟(带宽),KV Cache 量化主要加并发容量(吞吐)。生产上两者的组合(INT4 权重 + 2-bit/FP8 KV)才是把一张卡吃干榨净的完整配方。
5. 工业界主流:FP8 KV Cache
KIVI 证明了 2-bit 是可行的,但工业界(vLLM、TensorRT-LLM、SGLang)的默认选择却是 FP8 KV Cache,原因很实际:
- 精度几乎无损:FP8 的 E4M3 格式(1 符号 + 4 指数 + 3 尾数)动态范围大,8K 上下文内的 KV Cache 量化到 FP8 通常与 FP16 无可测量差异——不需要 KIVI 那种精细的非对称设计
- 生态成熟:Hopper 之后 GPU 原生支持 FP8,FlashAttention 等核心 kernel 对 FP8 KV Cache 是直接内建支持的(量化后的 KV 直接参与低精度计算,不需要”量化→反量化”绕一圈)
- 实现简单:vLLM 里一个
kv_cache_dtype="fp8"参数就搞定,KIVI 2-bit 反而因为 kernel 特殊性难以通用落地
FP8 KV Cache 的收益:KV 字节数减半(FP16 2 字节 → FP8 1 字节),长上下文/高并发能力提升约 1 倍(配合权重 INT4 时 KV 成为显存大头,这笔收益非常可观)。
⚠️ 注意:FP8 减半 vs 2-bit 减 8 倍,幅度差距很大,但”能直接进生产”的价值在工程上常常大于”省更多显存”。KIVI 的思想(K 与 V 不对称、残差滑窗)更多被后续的 4-bit/2-bit KV 研究(如 KVQuant、WKVQuant、NVFP4 KV)继承。
6. vLLM 实战:三种 scale 校准路线
vLLM(截至 v0.26.0)的 FP8 KV Cache 用 kv_cache_dtype="fp8"(等价于 fp8_e4m3,CUDA 11.8+;另有 fp8_e5m2)。真正需要决定的是 scale 怎么来,有三种路线:
6.1 路线一:不校准(默认 scale=1.0)
llm = LLM(model="...", kv_cache_dtype="fp8", calculate_kv_scales=False)
所有 scale 固定为 1.0。几乎不推荐——除非你的模型激活范围本来就小,否则量化误差大。但它零成本、零依赖,适合快速验证流程。
6.2 路线二:随机 Token 在线校准
llm = LLM(model="...", kv_cache_dtype="fp8", calculate_kv_scales=True)
启动时用一批随机 Token 跑一遍预热,统计出每个 tensor 的 scale 后固定下来。不用准备校准集,精度通常够用——生产里”不想折腾校准数据”时的默认选择。
6.3 路线三(官方推荐):llm-compressor 数据集校准
用校准数据集离线算好 scale,按 attention head 粒度(per-head scale)存进模型,vLLM 加载时直接用:
# llm-compressor 端(量化时)
from llmcompressor.modifiers.quantization import QuantizationModifier
fp8_args = QuantizationArgs(num_bits=8, type="float", strategy="attn_head") # 或 "tensor"
recipe = QuantizationModifier(
config_groups={"attention": QuantizationScheme(
targets=["LlamaAttention"], input_activations=fp8_args)},
kv_cache_scheme=fp8_args,
)
# 启动时 vLLM 自动读取保存的 scale
llm = LLM(model="量化后的模型目录", kv_cache_dtype="fp8")
三种路线的精度排序:路线三 ≈ 路线二 > 路线一;成本排序正好反过来。per-head 量化比 per-tensor 更准,但仅 FlashAttention 后端支持 per-head scale(vLLM 文档明确说明)。
💡 提示:vLLM 的 FA3(FlashAttention 3)后端配合 FP8 KV Cache 时,Attention 计算会在量化域内进行(query 也会被量化成 FP8)——这比”FP8 存取 + FP16 计算”更省带宽,是长上下文场景下的完整形态。
7. 敏感层跳过与量化组合拳
7.1 有些层就是不能量化
不是所有 attention 层都适合 KV 量化:sliding-window(滑动窗口)层、MLA(DeepSeek 的 Multi-head Latent Attention)等结构对 KV 精度更敏感。vLLM v0.26.0 提供:
# 跳过所有 sliding-window 层
vllm serve <model> --kv-cache-dtype fp8 --kv-cache-dtype-skip-layers sliding_window
# 跳过指定层
vllm serve <model> --kv-cache-dtype fp8 --kv-cache-dtype-skip-layers 0 1 23
--kv-cache-dtype-skip-layers 让敏感层保持原精度,其余层照常量化——“全部量化”从来不是最优解,“选择性量化”才是。
7.2 组合拳:一份完整的显存优化配方
| 组件 | 方法 | 显存效果 |
|---|---|---|
| 权重 | INT4(GPTQ/AWQ) | 减 4 倍 |
| KV Cache | FP8 | 减 2 倍 |
| 优化器状态/中间激活 | 无需(推理无优化器) | - |
| 输出 | fp32 lm_head(精度关键,vLLM 0.26 新增 head_dtype) | 少量增加 |
组合后,一张 80GB H100 跑 70B 模型的典型配置:INT4 权重约 37GB + FP8 KV Cache——剩余显存能支撑的并发量是”只量化权重”方案的两倍。
📝 总结
- KV Cache 量化是长上下文/高并发时代的刚需:权重量化完之后它成为显存大头,且与上下文长度、并发数线性相关
- KIVI 的分布学:K 按通道量化(静态特征性格)、V 按 Token 量化(动态激活性格),2-bit 靠”分组+残差”结构实现流式写入
- KIVI 效果:峰值显存减 2.6x、batch 上限 4 倍、吞吐 2.35-3.47x,精度几乎无损
- 工业界选 FP8 而非 2-bit:精度几乎无损 + 硬件/生态原生支持 + 实现简单,vLLM 一行参数开启
- 三种 scale 路线:不校准(快但糙)/ 随机 Token 校准(默认)/ llm-compressor 数据集校准(最准,支持 per-head)
- 选择性量化:
kv_cache_dtype_skip_layers跳过敏感层(sliding-window、MLA) - 组合拳:INT4 权重 + FP8 KV Cache,一张卡吃干榨净的完整配方
🎯 自我检验清单
- KV Cache 显存为什么与上下文长度和并发数都线性相关?
- KIVI 为什么对 K 用 per-channel、对 V 用 per-token?
- 残差区(R=128)在 KIVI 里的两个作用?
- KV Cache 量化为什么主要提升吞吐(并发)而不是单请求延迟?
- 工业界为什么选 FP8 而不是 2-bit?
- vLLM 三种 scale 校准路线的精度/成本排序?
- per-head scale 需要什么后端支持?
📚 参考资料
- Liu et al., KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache(https://arxiv.org/abs/2402.02750)
- vLLM 文档:Quantized KV Cache(https://docs.vllm.ai/en/latest/features/quantization/quantized_kvcache/)
- LLM Compressor 示例:KV Cache 量化(https://github.com/vllm-project/llm-compressor/tree/main/examples/quantization_kv_cache)