修炼AI技能
AI 教程 / 大模型应用 / KV Cache 量化与推理显存优化

KV Cache 量化与推理显存优化

自回归大模型在生成每个新 token 时，理论上需要重新计算之前所有 token 的注意力状态。为了避免这种昂贵且重复的运算，推理框架会把每一层注意力已经算好的 Key 和 Value 张量缓存下来，这就是 KV Cache。随着上下文变长、并发请求变多，KV Cache 的显存占用会线性甚至超线性膨胀，逐步取代模型权重，成为推理时的新显存瓶颈。量化（Quantization）把 KV Cache 从 FP16 压到 INT8、FP8 等低比特表示，是缓解这一瓶颈最常用、也最立竿见影的手段。本文先讲清 KV Cache 是什么、显存怎么算，再拆解 INT8、FP8 以及权重量化方法（AWQ、GPTQ）与 KV Cache 的关系，最后给出 vLLM、llama.cpp、AutoAWQ 三类可运行配置。

一、KV Cache 是什么

标准自注意力对每个 token 计算 Query、Key、Value 三组投影。解码阶段生成第 t 个 token 时，第 t 个 token 的 Query 要与此前第 1 到 t 个 token 的 Key、Value 做注意力运算。如果每次都从头算一遍全部 Key、Value，计算量会随着序列长度平方级增长。

KV Cache 的做法是：在生成第一个 token 时把各层算出的 Key、Value 缓存起来，之后每生成一个新 token，只计算当前这一步的 Key、Value，并把它们追加进缓存，再与缓存中历史所有 Key、Value 拼接做注意力。这样每一步的计算量只与新增的 1 个 token 相关，历史部分无需重算。

代价就是显存：缓存要一直保留到整个请求结束，且序列越长、并发越多，缓存占用越大。对于长上下文（32K、128K）或高并发批处理场景，KV Cache 的总量常常超过模型权重本身，这也是为什么「显存优化」的重心正从权重逐步转移到 KV Cache。

二、KV Cache 显存如何计算

KV Cache 的总字节数由层数、KV 头数、头维度、序列长度、批大小和数据精度共同决定。设：

  • B：批大小（并发请求数）
  • S：序列长度（上下文 token 数）
  • L：模型层数
  • h_kv：每层的 KV 头数（MHA 下等于注意力头数 h；GQA、MQA 下小于 h）
  • d：每个注意力头的维度
  • b：每个元素的字节数（FP16 为 2，FP8、INT8 为 1，INT4 为 0.5）

则 KV Cache 的字节数为：

KV Cache 字节数 = 2 × B × S × L × h_kv × d × b

其中因子 2 表示 Key 和 Value 两套缓存。由公式可见，KV Cache 随序列长度 S 和批大小 B 线性增长，这正是长上下文与高并发场景下它最危险的原因。把精度从 FP16（b=2）降到 FP8/INT8（b=1），显存直接减半。

以 Llama-2-7B 为例（L=32，MHA 故 h_kv=32，d=128，FP16 下 b=2）：

每生成一个 token 的边际占用 = 2 × 32 × 32 × 128 × 2 = 524288 字节 ≈ 512 KiB
若上下文 S = 4096，则单请求 KV Cache ≈ 512 KiB × 4096 ≈ 2 GB

也就是说，仅一个 4K 上下文的请求，KV Cache 就接近 2 GB。若把上下文拉长到 32K，单请求 KV Cache 就会超过 16 GB，已经超过 7B 模型权重本身的约 14 GB。再看 Llama-2-70B（采用 GQA，h_kv=8）：

每 token 边际占用 = 2 × 80 × 8 × 128 × 2 = 327680 字节 ≈ 320 KiB
S = 4096 时单请求 ≈ 320 KiB × 4096 ≈ 1.25 GB

对比可见，GQA 通过减少 KV 头数显著压低了 KV Cache，这也解释了为什么现代大模型普遍采用 GQA、MQA 来天然降低 KV 显存。但即便如此，量化仍是进一步压缩、换取更长上下文或更高并发的关键手段。

三、量化方法：INT8、FP8、AWQ、GPTQ 对 KV Cache 的应用

3.1 INT8 KV Cache

INT8 是最直观的均匀量化：用 8 比特整数（256 个级别）表示原本 FP16 的 Key、Value 元素，每个元素字节数从 2 降到 1，显存直接减半。难点在于量化粒度的选择。对 KV Cache 分布的研究（见下文 KIVI）发现：Key cache 沿通道（channel）维度的数值分布差异大，适合按通道分组量化；Value cache 沿 token 维度的分布差异大，适合按 token 量化。因此实践中对 K、V 往往采用不同的量化轴，而不是简单对整个张量做统一缩放。

INT8 的缺点是：每个量化组都需要额外存储缩放因子（scale），组越细、scale 的额外开销越大，会部分抵消显存收益。这也是为什么新一代框架更倾向于用 FP8。

3.2 FP8 KV Cache

FP8 是 8 比特浮点格式，OCP（开放计算项目）定义了两种常见形态：E5M2（5 位指数、2 位尾数）和 E4M3（4 位指数、3 位尾数，精度更高，但动态范围较小，常需配合 FP32 缩放因子）。与 INT8 相比，FP8 本身就是浮点，对注意力中数值范围较宽的部分更友好，且当前主流推理框架对 FP8 KV Cache 已提供成熟支持。

以 vLLM 为例，FP8 KV Cache 通过 kv_cache_dtype 开启，并采用逐张量（per-tensor，标量）缩放因子。官方文档指出，FP8 可让 KV Cache 显存约下降一半，从而容纳约两倍的缓存空间，提升吞吐；且 FP8 E4M3 通常只带来极小的精度退化。FP8 已成为当前生产环境压缩 KV Cache 的首选路径。

3.3 AWQ、GPTQ 与 KV Cache 的关系

这里要先澄清一个常见误解：AWQ 与 GPTQ 本质上是模型权重的量化方法，并不直接量化 KV Cache。

  • GPTQ（arXiv:2210.17323）用二阶误差补偿做 one-shot 权重量化，把权重压到 3 到 4 bit；
  • AWQ（arXiv:2306.00978）依据激活分布保护约 1% 的显著权重通道，做 4-bit 权重量化。

它们量化的是「静态、加载时已知」的权重，而不是「动态、推理时生成、随序列增长」的 KV Cache。不过，这两个方法的工程思想确实启发了 KV Cache 专用量化：AWQ 的「激活感知、保护显著通道」思路，对应到 KV Cache 上就是「按通道量化 Key、按 token 量化 Value」；GPTQ 的「校准集求最优量化参数」思路，则对应到用代表性数据校准 KV Cache 的缩放因子（如 vLLM 推荐用 llm-compressor 做校准）。

因此，专门面向 KV Cache 的低比特方法应运而生，代表工作是 KIVI（arXiv:2402.02750，ICML 2024）：一种免调参、非对称的 2-bit KV Cache 量化，对 Key 按通道、对 Value 按 token 量化，使 Llama、Falcon、Mistral 在几乎不损质量的前提下，峰值显存（含权重）下降约 2.6 倍，可支撑约 4 倍的更大批大小，带来 2.35 到 3.47 倍的吞吐提升。另一项工作 KVQuant（arXiv:2401.18079）把 KV Cache 压到 2 bit 量级并强调校准。归纳起来：权重量化靠 GPTQ、AWQ，KV Cache 量化则靠 FP8/INT8 均匀量化与 KIVI 这类专用方法，二者是两条互补的路线。

四、KV Cache 量化与模型权重量化的区别

二者都叫「量化」，但对象、时态与约束差异很大：

  • 量化对象不同：权重量是静态张量，加载模型时就固定；KV Cache 是动态张量，在解码过程中随序列不断生成、不断追加，生命周期随请求结束而释放。
  • 量化粒度不同：权重量化多按通道（per-channel）或分组（group-wise，如 AWQ 的 q_group_size=128）进行；KV Cache 量化因 Key、Value 分布特性不同，通常对 Key 按通道、对 Value 按 token，二者轴不一致。
  • 缩放因子开销不同：权重的缩放因子被海量参数分摊，几乎可忽略；KV Cache 按 token 追加，若采用细粒度 INT8/INT4，缩放因子会持续累积，反而侵蚀显存收益，所以 KV Cache 更适合 FP8 这类自带浮点、无需额外大开销的方案。
  • 精度影响不同：权重量化影响所有层的前向计算；KV Cache 量化只影响注意力中历史信息的读取精度，长上下文下若量化过激，更容易出现远处信息丢失、长程依赖退化。
  • 校准方式不同：权重量化靠离线校准集求固定参数；KV Cache 因内容随输入变化，常采用在线（on-the-fly）动态缩放，例如 vLLM 的 calculate_kv_scales=True，或用代表性数据预先校准后固化缩放因子。
  • 对应工具不同：权重量化工具以 GPTQ、AWQ、llm-compressor 为代表；KV Cache 量化工具以 vLLM 的 FP8 KV Cache、llama.cpp 的 cache-type、以及 KIVI 等专用内核为代表。

一句话：权重量化解决「模型能不能装下、跑得快不快」，KV Cache 量化解决「长上下文和高并发下显存撑不撑得住」，两者叠加才能把一张卡的利用率推到极限。

五、实战：可运行的配置示例

下面三个示例分别演示在生产级框架（vLLM）、本地 C/C++ 引擎（llama.cpp）与权重量化工具（AutoAWQ）下，如何围绕 KV Cache 做显存优化。

5.1 vLLM：开启 FP8 KV Cache

vLLM 通过 LLM 构造器的 kv_cache_dtype 控制 KV Cache 精度。设为 "fp8"（等价于 "fp8_e4m3"）即开启 FP8 KV Cache；calculate_kv_scales=True 表示在线动态计算缩放因子，免去预先校准。

所需包：vllm。

from vllm import LLM, SamplingParams

sampling_params = SamplingParams(temperature=0.7, top_p=0.8)  # 温度与 top_p 控制生成多样性

llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",
    kv_cache_dtype="fp8",            # 把 KV Cache 存为 FP8
    calculate_kv_scales=True,        # 在线动态计算缩放因子
)

prompt = "请简要介绍一下 KV Cache 的作用。"
outputs = llm.generate(prompt, sampling_params)
print(outputs[0].outputs[0].text)

若已用 llm-compressor 等工具对缩放因子做了离线校准，可把校准结果通过 quantization_param_path 指给引擎，以获得更稳的质量：

llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",
    kv_cache_dtype="fp8",
    quantization_param_path="./kv_cache_scales.json",
)

说明：当前 vLLM 的 FP8 KV Cache 仅支持逐张量（标量）缩放因子，更细粒度（如逐通道）仍在开发中；此外历史版本中 FP8 KV Cache 与 prefix caching 存在不兼容，使用前请核对所用版本的发布说明。

5.2 llama.cpp：用 cache-type 量化 KV Cache

llama.cpp 通过启动参数直接指定 K、V 缓存的数据类型。键、值分别由 -ctk/--cache-type-k 与 -ctv/--cache-type-v 控制，可接受 f16、q8_0、q4_0、q4_1、iq4_nl、q5_0、q5_1 等，默认均为 f16。把二者都设为 q8_0 即可让 KV Cache 以 8-bit 量化格式驻留，显存相较默认 f16 减半。

先准备一个 GGUF 模型文件（如从 Hugging Face 下载对应量化版），随后启动 OpenAI 兼容服务：

llama-server -m models/7B/ggml-model.gguf -ctk q8_0 -ctv q8_0

若使用旧版命令行入口或偏好长参数写法：

./llama-cli -m models/7B/ggml-model.gguf --cache-type-k q8_0 --cache-type-v q8_0

服务默认监听 127.0.0.1:8080，之后即可像调用普通 OpenAI 接口一样发送请求。若想进一步压低显存，可把 q8_0 换成 q4_0、q4_1 等更低位宽，但需留意长上下文下精度退化会更明显。

5.3 AutoAWQ：量化权重，KV Cache 交给推理引擎

需要先说明一个事实：AutoAWQ 只量化模型权重（4-bit AWQ，W4A16 存储权重 INT4、推理 FP16），并不提供 KV Cache 量化选项；其官方仓库也声明已进入弃用状态，后续由 vLLM 的 llm-compressor 接管。因此下面的示例演示的是「用 AutoAWQ 产出 4-bit 权重」，而 KV Cache 的压缩要在加载该权重时，由 vLLM 的 FP8 KV Cache 叠加完成。

所需包：autoawq、transformers。

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "mistralai/Mistral-7B-Instruct-v0.2"  # 待量化模型
quant_path = "mistral-instruct-v0.2-awq"            # 量化产物路径

quant_config = {            # AWQ 量化配置
    "zero_point": True,     # 启用零点
    "q_group_size": 128,    # 按 128 分组
    "w_bit": 4,             # 4-bit 权重
    "version": "GEMM",
}

model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

model.quantize(tokenizer, quant_config=quant_config)  # 执行量化
model.save_quantized(quant_path)                      # 保存量化权重
tokenizer.save_pretrained(quant_path)

量化完成后，用 vLLM 加载这份 AWQ 权重，并同时开启 FP8 KV Cache，便同时压住了「权重」与「KV Cache」两块显存：

from vllm import LLM

llm = LLM(   # 权重已 4-bit（AutoAWQ），KV Cache 再叠加 FP8 压缩
    model="mistral-instruct-v0.2-awq",
    kv_cache_dtype="fp8",
    calculate_kv_scales=True,
)

这一组合也印证了前文结论：权重量化与 KV Cache 量化是两条互补路线，需分别由对应工具完成。

小结

KV Cache 通过缓存历史 Key、Value 避免重复计算，是解码加速的关键，但其显存随序列长度与批大小线性增长，在长上下文、高并发下会超过权重成为新瓶颈。显存公式为 2 × B × S × L × h_kv × d × b，把精度 b 从 FP16（2 字节）降到 FP8/INT8（1 字节）即可让 KV Cache 减半。方法上，INT8 适合做按通道（Key）/按 token（Value）的均匀量化，FP8 因自带浮点、缩放开销小，已成为生产首选；AWQ、GPTQ 虽是权重量化方法，不直接作用于 KV Cache，但其激活感知、校准求参的思想启发了 KIVI 这类 2-bit KV Cache 专用量化。KV Cache 量化与权重量化的本质区别在于对象（动态 vs 静态）、粒度（K/V 异轴 vs 通道/分组）、缩放开销（累积 vs 分摊）与精度影响面（仅注意力历史 vs 全部前向）。实战中可用 vLLM 的 kv_cache_dtype="fp8"、llama.cpp 的 -ctk/-ctv q8_0 直接压缩 KV Cache，而 AutoAWQ 负责把权重压到 4-bit，二者组合方能最大化单卡可服务的上下文长度与并发数。

参考与延伸阅读

本文累计阅读 — 次