KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache

发表时间: 2024-02 · arXiv:2402.02750 (ICML 2024)

原文: https://arxiv.org/abs/2402.02750

Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen (Henry) Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, Xia Hu

速读

一句话结论 本文提出了一种免微调的非对称 2bit KV Cache 量化算法 KIVI，通过对 Key 缓存按通道量化、对 Value 缓存按 Token 量化，在几乎不损失生成质量的前提下将推理显存峰值降低了 2.6 倍，并带来最高 3.47 倍的吞吐量提升。

要解决什么问题 在大语言模型推理中，为了降低单次请求成本，通常会采用大批量处理。但随着批量大小和上下文长度的增加，用于存储注意力中间状态以避免重复计算的 KV Cache，正成为全新的显存与速度瓶颈。以 540B 参数模型为例，在批量大小为 512、上下文长度为 2048 时，仅 KV Cache 就会占用 3TB，达到模型权重的 3 倍。除了空间占用，推理速度也受严重制约。在生成新 Token 时，GPU 必须将庞大的 KV Cache 从显存加载到 SRAM 中，受限于显存带宽，计算核心在此期间只能处于空闲等待状态。最直接的缩减方案是量化，但现有常规做法（如将 Key 和 Value 直接按 Token 进行 4bit 量化）在尝试压缩到 2bit 时，模型精度会发生断崖式下跌。此前业界缺乏对 KV Cache 内部元素数值分布的深入剖析，导致极低比特量化的误差来源不够清晰，这正是本研究所要解决的卡点。

怎么做的 核心思路是采用非对称的量化维度：对 Key 缓存按通道（Channel）量化，对 Value 缓存按 Token 量化。作者观察到，Key 缓存的数值分布在某些固定通道上存在极大的异常值。按通道分组量化能将误差严格限制在独立通道内，避免异常值破坏正常通道。相比之下，Value 缓存没有明显的异常值规律，但它在注意力机制中是作为被加权求和的项。由于注意力分数矩阵高度稀疏，输出本质上是少数重要 Token 的组合。按 Token 量化能将误差隔离在单个 Token 内部，确保普通 Token 的量化不会污染关键 Token。标准的量化与反量化过程定义为：

$$Q(X) = \lfloor \frac{X - z_X}{s_X} \rceil, X' = Q(X) \cdot s_X + z_X$$

在流式推理场景下，新生成的 Value 缓存能直接按 Token 维度追加到量化张量末尾。然而，Key 缓存按通道量化需要跨越多个 Token，无法直接流式处理。为此，KIVI 设计了分组与残差机制。它将 KV Cache 拆分为两部分：包含若干完整组的“分组缓存”采用 2bit 量化；未凑满一组的“残差缓存”保持 16bit 全精度。新生成的 Token 先以全精度加入残差缓存，积累到预设长度时，再打包量化拼接到分组缓存中。计算注意力分数时，模型分别计算量化部分和全精度部分的点积再拼接：

$$\boldsymbol{A}_g = \boldsymbol{t}_Q Q(\boldsymbol{X}_{K_g}^\top), \boldsymbol{A}_r = \boldsymbol{t}_Q \boldsymbol{X}_{K_r}^\top, \boldsymbol{A} = \text{Concat}([\boldsymbol{A}_g, \boldsymbol{A}_r])$$
这种设计不仅绕开了流式处理的工程卡点，还相当于为最近生成的 Token 保留了一个全精度滑动窗口，这对维持复杂推理任务的准确性起到了决定性作用。

效果如何 实验在单张 80GB 的 NVIDIA A100 硬件上搭建，评测了 Llama、Llama-2、Falcon 和 Mistral 模型家族（包含 7B 和 13B 规模）。测试任务涵盖了常规长度的 LM-Eval（包含 CoQA、TruthfulQA 及 GSM8K 数学推理）、LongBench 长文本基准，以及大海捞针测试。对比基线包含了 16bit 全精度基线、代表现有常规路线的 4bit 双 Token 量化基线，以及多种模拟的 2bit 量化组合路线（如双通道量化、双 Token 量化等）。结果表明，在 Llama-2-7B 模型上，KIVI 能够将包含模型权重在内的峰值显存占用降低 2.6 倍。得益于显存释放，系统能支持高达 4 倍的批量大小，在基于 ShareGPT 数据的真实负载下实现了 2.35 倍到 3.47 倍的吞吐量提升。精度方面，对于 Llama 和 Mistral 模型，即使 KV Cache 压缩到 2bit，在 GSM8K 任务上的准确率下降也仅在 2% 左右，且在大海捞针测试中依然保持了精准的检索能力。作者也指出了方法局限：对于采用多查询注意力（Multi-Query Attention）的 Falcon-7B 模型，因其 KV Cache 仅有单个注意力头，信息已被高度压缩，强制应用 2bit 会导致明显精度损失，必须退回 4bit 才能维持效果。

主要贡献

高效服务大型语言模型（LLMs）需要对大量请求进行批处理，以降低每个请求的成本。然而，随着批处理规模的扩大和上下文长度的增加，用于存储注意力键（Key）和值（Value）以避免重复计算的键值缓存（KV Cache），显著增加了内存需求，并成为推理速度和内存使用的新瓶颈。此外，加载 KV 缓存会导致计算核心处于空闲状态，从而限制了推理速度。

为了减少 KV 缓存的大小，量化是一种直接且有效的解决方案，它可以减少 KV 缓存占用的总字节数。然而，当前缺乏对 KV 缓存元素分布的深入研究，难以理解 KV 缓存量化的难度和局限性。为填补这一空白，本文对流行 LLM 的 KV 缓存元素分布进行了全面研究。

本文的核心创新点与贡献如下：
* 深度的分布分析：对常用 LLM 中 KV 缓存的离群值（outlier）模式和量化误差进行了广泛分析。研究表明，键缓存（Key cache）应按通道（per-channel）进行量化，而值缓存（Value cache）应按词元（per-token）进行量化，并深入解释了为何这两种缓存需要不同的量化方法。
* KIVI 量化算法：基于上述发现，开发了一种名为 KIVI 的免微调（tuning-free）、即插即用的 2bit 极低比特 KV 缓存量化算法。
* 卓越的性能与系统实现：通过硬件友好的实现，KIVI 使 Llama、Falcon 和 Mistral 模型在保持几乎相同生成质量的同时，峰值内存（包含模型权重）降低了 $2.6 \times$。内存使用的减少使得批处理大小（batch size）最多可提升 $4 \times$，在实际的 LLM 推理工作负载中带来了 $2.35 \times \sim 3.47 \times$ 的吞吐量提升。

背景知识与设计原则

注意力机制推理工作流。大模型注意力推理过程包含两个阶段：预填充阶段（Prefill Phase）和解码阶段（Decoding Phase）。在预填充阶段，输入提示用于为每一层生成 KV 缓存。设输入张量为 $X \in \mathbb{R}^{b \times l_{\mathrm{prompt}} \times d}$，其中 $b$ 为批大小，$l_{\mathrm{prompt}}$ 为提示长度，$d$ 为隐藏层维度。键和值张量计算为 $\dot{X}_K = XW_K$ 和 $X_V = XW_V$，计算后将其缓存在内存中。在解码阶段，模型逐个生成词元。设当前输入词元嵌入为 $t \in \mathbb{R}^{b \times 1 \times d}$，输出为 $t_K = tW_K$ 和 $t_V = tW_V$。系统首先更新缓存：$X_K = \mathrm{Concat}(X_K, t_K)$， $X_V = \mathrm{Concat}(X_V, t_V)$。接着计算注意力输出：$t_Q = tW_Q$， $A = \mathrm{Softmax}(t_Q X_K^\top)$， $t_O = A X_V$。随着生成过程重复，KV 缓存的形状变为 $b \times (l_{\mathrm{prompt}} + l_{\mathrm{gen}}) \times d$，导致巨大的内存占用和频繁的显存到 SRAM 数据加载开销。

KV 缓存量化的初步研究。由于 KV 缓存具有流式（streaming）特性，新张量是顺序到达的，因此基于优化的量化方法（如 GPTQ）因开销过大而不适用。最灵活的方法是最近整数量化（round-to-nearest quantization），其过程为 $Q(X) = \lfloor \frac{X - z_X}{s_X} \rceil$，$X^\prime = Q(X) \cdot s_X + z_X$。考虑到流式特性，以往研究通常对键和值缓存都采用按词元（per-token）量化。为探索不同维度量化的影响，作者设计了填充方法来模拟按通道（per-channel）量化。实验观察到三个核心现象：第一，当对键和值缓存均使用按词元量化时，INT4 精度可保持准确性，但降至 INT2 时精度显著下降；第二，当值缓存按通道量化时，无论键缓存如何量化，精度都会急剧恶化；第三，在使用 INT2 时，最准确的策略是键缓存按通道量化，值缓存按词元量化。

按词元与按通道量化的定义
按词元与按通道量化的定义

键缓存量化维度的原理分析。键缓存之所以应该按通道量化，是因为其在特定固定通道上表现出极大的幅度（离群值），这与之前的研究发现（Lin et al., 2023; Xiao et al., 2023a）一致。由于离群值在通道内持续存在，按通道量化能够将量化误差限制在每个单独的通道内，而不会影响其他正常通道。误差统计证实，按词元量化键缓存会导致比按通道量化大近 $5 \times$ 的注意力分数误差。

值缓存量化维度的原理分析。与键缓存不同，值缓存没有明显的通道维度离群值模式。值缓存必须按词元量化的原因在于其物理用途：值缓存用于计算注意力输出 $t_O$，该输出本质上是值缓存的混合器。注意力输出 $[A X_V]_{i*} = \sum_{j=1}^{l_{\mathrm{prompt}}} A_{ij} [X_V]_{j*}$ 是不同词元值缓存的加权求和。由于注意力分数高度稀疏（Tian et al., 2023），输出仅是少数重要词元值缓存的组合。按词元量化可将误差限制在每个单独的词元内部，量化其他词元不会影响重要词元的准确性。实验数据表明，值缓存按词元量化的相对误差比按通道量化小近 $15 \times$。

Llama-2-13B与Falcon-7B的键值缓存幅度分布
Llama-2-13B与Falcon-7B的键值缓存幅度分布

方法细节

流式场景下的非对称量化策略。基于原理分析，KIVI 算法对键缓存采用按通道量化，对值缓存采用按词元量化。从实现角度看，按词元量化与流式自回归推理完美契合，新量化的张量可以直接在词元维度上追加到现有的量化值缓存中。然而，对于键缓存的按通道量化，量化过程需要跨越不同的词元，这无法直接在流式设置中实现。

键缓存的分组与残差切分机制。为了解决按通道量化在流式场景下的实现难题，KIVI 的核心思想是将键缓存每隔 $G$ 个词元进行分组并分别量化。由于键缓存 $X_K$ 中的词元总数是任意的，系统将 $X_K$ 切分为两部分：第一部分是分组键缓存 $X_{K_g} = X_K[:l-r]$，包含若干个完整的词元组；第二部分是残差键缓存 $X_{K_r} = X_K[l-r:]$，其词元数量不足以构成一个完整的组。其中 $l$ 是当前键缓存的总词元数，$r$ 是残差词元数，$l-r$ 必须能被组大小 $G$ 整除。

键缓存的流式更新与注意力计算。系统仅对分组键缓存 $X_{K_g}$ 进行分组量化并存储为 $Q(X_{K_g})$，而残差键缓存 $X_{K_r}$ 保持全精度。在解码过程中，每个新到达的键缓存张量 $t_K$ 首先被追加到全精度的 $X_{K_r}$ 中。一旦 $X_{K_r}$ 积累了 $R$ 个词元（$R$ 为预设的残差长度超参数，且必须被 $G$ 整除），系统便对其进行量化，并在词元维度上与之前已量化的 $Q(X_{K_g})$ 拼接，随后将 $X_{K_r}$ 重置为空张量。在计算注意力分数时，系统利用平铺矩阵乘法（tiled matrix multiplication）分别计算两部分：首先计算分组部分的逻辑值 $A_g = t_Q Q(X_{K_g}^\top)$，接着计算残差部分的逻辑值 $A_r = t_Q X_{K_r}^\top$，最后将两者拼接得到完整的注意力分数 $A = \mathrm{Concat}([A_g, A_r])$。

值缓存的队列式更新机制。对于值缓存，KIVI 同样将其拆分为分组部分 $X_{V_g}$ 和残差部分 $X_{V_r}$。具体而言，系统维护一个队列，每个新到达的值缓存张量被推入该队列，作为全精度的残差部分。一旦队列长度达到预设的残差长度 $R$，最旧的值缓存张量将被弹出。系统对弹出的值缓存进行按词元量化，并将其沿词元维度拼接到之前已量化的值缓存 $Q(X_{V_g})$ 中。需要强调的是，在预填充（Prefill）阶段，尽管内存中最终只保留量化后的 KV 缓存，但精确的全精度键值张量会被传递到下一层以保证计算精度。

KIVI算法概览
KIVI算法概览

全精度滑动窗口的内存与精度权衡。在 KIVI 的设计中，残差键缓存 $X_{K_r}$ 和残差值缓存 $X_{V_r}$ 始终保持全精度，且其包含的词元数量最多不超过 $R$（实践中 $R \le 128$）。由于实际推理的序列长度 $l_{\mathrm{prompt}} + l_{\mathrm{gen}}$ 通常远大于 $R$，因此在极低比特量化带来的巨大内存收益面前，这部分全精度残差所产生的内存开销微乎其微。更重要的是，由于新生成的词元始终以全精度加入残差部分，KIVI 实际上为局部相关的词元维护了一个全精度的 KV 缓存滑动窗口（键缓存的期望窗口大小为 $\frac{R}{2}$，值缓存为 $R$）。这种局部全精度的保留对于在复杂任务（如 GSM8K 的数学推理）中维持模型的高性能起到了决定性作用。

硬件友好的系统级优化。为了在 GPU 上高效运行 KIVI 并最小化量化开销，系统提供了硬件友好的底层实现。利用 CUDA，系统将反量化（dequantization）过程与平铺级别的矩阵乘法深度融合（对应图中的 Q_MatMul 操作）。同时，分组量化内核使用 Triton 进行编写。该系统实现完全兼容仅权重（weight-only）量化方案。

实验环境

  • 数据集：

    • 常规上下文评估（LM-Eval）：CoQA（精确匹配准确率）、TruthfulQA（BLEU 分数）、GSM8K（精确匹配准确率）。
    • 长上下文评估（LongBench）：Qasper（单文档问答）、QMSum、MultiNews（摘要生成）、TREC、TriviaQA、SAMSum（少样本学习）、LCC、RepoBench-P（代码补全）。
    • 长上下文检索：Needle-in-a-haystack (NIAH) 大海捞针测试。
    • 系统效率评估：基于 ShareGPT 的真实 LLM 服务输入输出文本构建工作负载。

  • 模型架构：Llama-2（7B, 13B, Chat版本）、Falcon-7B（采用多查询注意力 MQA）、Mistral-7B（采用多头注意力 MHA）、Llama-3-8B-Instruct、Mistral-7B-Instruct-v0.2、LongChat-7B-v1.5。

  • 硬件配置：单张 NVIDIA A100 GPU (80GB)。
  • 软件配置：基于 Hugging Face Transformers 代码库实现 KIVI 算法。量化组大小 $G$ 设为 32，键和值缓存的残差长度 $R$ 设为 128（部分消融实验设为 32）。

实验结果

不同量化配置的精度对比。在 LM-Eval 数据集上，对多种模拟量化配置进行对比。结果表明，“2bit（键按通道，值按词元）”始终优于其他量化配置。值得注意的是，完全模拟的 2bit 量化在 GSM8K 等困难任务上会发生精度崩溃，而保留全精度滑动窗口的 KIVI-2 在 Llama 和 Mistral 模型上，尽管 KV 缓存被压缩至 2bit，精度下降仅约 2%。对于 Falcon-7B，由于其采用 MQA（仅有 1 个 KV 头），缓存已被高度压缩，因此需要使用 4bit 的 KIVI-4 才能维持精度。

长上下文生成任务性能。在 LongBench 数据集上，对 Llama2-7B/13B 及其 Chat 版本、Falcon-7B 和 Mistral-7B 进行评估（最大序列长度 Mistral 为 8192，其余为 4096）。实验结果证实，KIVI-2 和 KIVI-4 在各种困难的长上下文生成任务中，平均性能得分与 16bit 全精度基线几乎完全一致，证明了其在长文本场景下的有效性和极小的精度损失。

长上下文检索能力（NIAH）。在大海捞针测试中，对 Llama-3-8B-Instruct 和 Mistral-7B-Instruct-v0.2 进行了评估。结果表明，即使将 KV 缓存压缩至 2bit，KIVI 依然能够完美保持 LLM 在海量背景文本中精准检索关键信息的能力。

Llama-3-8B-Instruct与Mistral-7B-Instruct-v0.2的大海捞针结果
Llama-3-8B-Instruct与Mistral-7B-Instruct-v0.2的大海捞针结果

系统效率与吞吐量提升。在 ShareGPT 负载（平均输入 161 词元，输出 338 词元）下，对比了 Llama-2-7B 的 FP16 基线与 KIVI 的显存占用和吞吐量。结果显示，在相似的最大内存使用量下，2bit KIVI 允许批处理大小（Batch Size）扩大至原来的 $4 \times$，并带来了 $2.35 \times \sim 3.47 \times$ 的吞吐量提升。随着上下文和输出长度的增加，该吞吐量增益将更为显著。

2bit KIVI与16bit基线的内存使用与吞吐量对比
2bit KIVI与16bit基线的内存使用与吞吐量对比

结论与展望

本文系统分析了流行 LLM 中 KV 缓存元素的分布，得出结论：键缓存应按通道量化，而值缓存应按词元量化。基于这些观察，提出了 KIVI，一种无需微调、即插即用的 2bit KV 缓存量化算法。在真实的 LLM 工作负载中，KIVI 支持高达 $4 \times$ 的批处理大小，并实现 $3.47 \times$ 的吞吐量提升。未来的工作将进一步优化底层实现，以减少预填充和解码阶段量化过程的开销，并探索与系统级优化（如 PagedAttention）的结合。

附录细节

KIVI 算法伪代码实现。在预填充阶段，输入张量 $X$ 经过线性映射得到 $X_K$ 和 $X_V$。对 $X_V$，系统截取前 $l_{\mathrm{prompt}}-R$ 个词元作为 $X_{V_g}$，按词元维度进行分组量化得到 $Q(X_{V_g})$；剩余的 $R$ 个词元作为 $X_{V_r}$ 保留全精度。对 $X_K$，调用 KeyQuant 函数，将前 $l-r$ 个词元作为 $X_{K_g}$ 按通道进行分组量化得到 $Q(X_{K_g})$，尾部剩余的 $r$ 个词元作为 $X_{K_r}$ 保留全精度。

在解码阶段，新词元的键 $t_K$ 和值 $t_V$ 分别追加到全精度的 $X_{K_r}$ 和 $X_{V_r}$ 中。如果 $X_{K_r}$ 的长度达到 $R$，则对其按通道量化，拼接到 $Q(X_{K_g})$ 后，将 $X_{K_r}$ 清空。如果 $X_{V_r}$ 的长度超过 $R$，则将超出部分（最旧的词元）按词元量化并拼接到 $Q(X_{V_g})$，保持 $X_{V_r}$ 长度为 $R$。注意力分数 $A$ 由全精度的查询 $t_Q$ 分别与量化的 $Q(X_{K_g})$ 和全精度的 $X_{K_r}$ 相乘后拼接，经过 Softmax 计算得到 $A_g$ 和 $A_r$，最终输出 $t_O = A_g Q(X_{V_g}) + A_r X_{V_r}$。

# Algorithm 1: KIVI 预填充与解码算法逻辑简述
# 参数: group size G, residual length R
def Prefill(X):
    X_K = X @ W_K; X_V = X @ W_V
    X_Vg = X_V[:l_prompt - R]; X_Vr = X_V[l_prompt - R:]
    Q_X_Vg = GroupQuant(X_Vg, dim='token', numGroup=d//G)
    Q_X_Kg, X_Kr = KeyQuant(X_K)
    KV_cache = (Q_X_Kg, X_Kr, Q_X_Vg, X_Vr)
    return X_K, X_V

def Decoding(KV_cache, t):
    t_Q = t @ W_Q; t_K = t @ W_K; t_V = t @ W_V
    Q_X_Kg, X_Kr, Q_X_Vg, X_Vr = KV_cache
    
    X_Kr = Concat([X_Kr, t_K], dim='token')
    X_Vr = Concat([X_Vr, t_V], dim='token')
    
    if len(X_Kr) == R:
        Q_X_Kr = KeyQuant(X_Kr)
        Q_X_Kg = Concat([Q_X_Kg, Q_X_Kr], dim='token')
        X_Kr = empty_tensor
        
    if len(X_Vr) > R:
        Q_X_Vr_prime = GroupQuant(X_Vr[:-R], dim='token', numGroup=d//G)
        Q_X_Vg = Concat([Q_X_Vg, Q_X_Vr_prime], dim='token')
        X_Vr = X_Vr[-R:]
        
    A = Concat([t_Q @ Q_X_Kg.T, t_Q @ X_Kr.T], dim='token')
    A_g = Softmax(A)[:-R]; A_r = Softmax(A)[-R:]
    t_O = A_g @ Q_X_Vg + A_r @ X_Vr
    
    KV_cache = (Q_X_Kg, X_Kr, Q_X_Vg, X_Vr)
    return t_O

def KeyQuant(X_K):
    r = l % R
    X_Kg = X_K[:l-r]; X_Kr = X_K[l-r:]
    Q_X_Kg = GroupQuant(X_Kg, dim='channel', numGroup=l//G)
    return Q_X_Kg, X_Kr

超参数的消融实验分析。针对 GSM8K 任务对组大小 $G$ 和残差长度 $R$ 进行了消融实验。固定 $R=128$ 时，组大小为 32 和 64 性能相似（约 20.77~21.00），但增至 128 时性能显著下降至 17.29，因为量化的零点和缩放因子基于组大小计算，过大的组在长输入下会影响压缩效果。固定 $G=32$ 时，残差长度为 32、96、128 时均取得相似的良好结果，但 64 时性能最差。总体而言，保持合理大小的残差长度对复杂任务的性能提升至关重要。

极致压缩下的残差长度探索。在效率评估中发现，当残差长度 $R=32$ 时，KIVI 能实现更高的内存压缩率从而进一步提升吞吐量。补充实验表明，将 $R$ 从 128 降至 32，在 CoQA、TruthfulQA、GSM8K 以及 LongBench 上的性能与 16bit 全精度模型和 $R=128$ 的配置相比，均未出现显著的性能下降，证明了 KIVI 算法在小残差窗口下的鲁棒性。

核心参考文献引述

  • 在分析键缓存通道维度存在极大离群值时，作者引述了以往在激活值量化领域的发现，指出特定固定列存在较大离群值是常见现象【[14] Lin et al., AWQ: Activation-aware weight quantization for llm compression and acceleration, 2023, arXiv】以及【[25] Xiao et al., Smoothquant: Accurate and efficient post-training quantization for large language models, 2023, ICML】。
  • 在解释为何值缓存按词元量化误差极小且有效时，作者引述了注意力分数高度稀疏的特性，指出输出仅是少数重要词元的组合【[22] Tian et al., Scan and snap: Understanding training dynamics and token composition in 1-layer transformer, 2023, arXiv】。