Accurate KV Cache Quantization with Outlier Tokens Tracing
在大规模语言模型(LLM)的推理阶段,“Accurate KV Cache Quantization with Outlier Tokens Tracing”(OTT)提出一种“动态追踪并剔除异常 token + 2 bit 层次量化 + 多缓存并行解码”的轻量级框架。论文发现只有极少量“异常 token”会破坏通道均匀性、导致量化误差激增,于是通过在线检测这部分 token 并保持其全精度表示,在几乎不损失准确率的前提下,将 KV Cache 显存压缩 6.4 倍、系统吞吐提升 2.3 倍,显著优于先前的 KIVI 等静态 2-bit 方案。(arxiv.org, arxiv.org, arize.com)
1. 研究动机与整体思路
长上下文推理时,KV Cache 成为显存和带宽瓶颈;现有 2 bit 量化(如 KIVI)假设“Key 按通道、Value 按 token”分布均匀,但作者发现少数 “异常 token”在高幅度通道里却呈低幅度,量化后会把注意力权重放大数十倍,引发输出误差。OTT 的核心思路是:实时识别这些异常 token → 将其转入全精度缓冲区 → 其余位置继续用 2 bit 量化。(arxiv.org, researchgate.net, chatpaper.com)
2. 关键观察
2.1 Key 的通道离群现象
实验显示 Key 的模值呈“多数通道均匀 + 少数通道幅度极大”的长尾分布;在这些“离群通道”内,又隐藏着数个幅度极小的 outlier token。量化把整个通道映射到同一动态范围时,这些 token 会被严重放大,最终污染注意力权重。(arxiv.org)
2.2 误差来源分析
附录 D 给出理论推导:Key 量化误差在 softmax 前被\exp(\cdot) 放大,若 Key 本身已处于大幅度通道,量化噪声会导致 attention score 偏差呈指数级膨胀。将异常 token 保留全精度可显著降低 L1/L2 误差。(arxiv.org)
3. OTT 方法细节
3.1 异常 token 判定规则
量化步长:每隔G=128 个 token 对最近一组 KV 进行 2 bit 量化。
评分指标:计算每个候选 token 的 Key 模值\lVert K\rVert_2;若其处于离群通道且幅度位于全局最小的前p%,则标记为异常。(arxiv.org, deeppapers.dev)
3.2 Outlier Pool 机制
维护容量为P 的 Outlier Pool 存放全精度 KV;所有新 token 与池内 token 竞争,若更“离群”则替换进入池子。
被替换出的 KV 写回主缓存(或丢弃在 size-32 的“冷备池”中,论文实验证明已足够)。(arxiv.org)
3.3 多缓存融合解码
解码时 Query 需同时与三类 Key 相乘:
Quantized KV(2 bit,绝大多数 token);
Residual Window(最近R=32 token,保持 FP16);
Outlier Pool(全精度异常 token)。
作者实现了 CUDA fused kernel 把量化矩阵与全精度矩阵拼接计算,避免额外带宽。(arxiv.org, youtube.com)
flowchart LR
subgraph Quantization(每G步)
Q1[计算 |K|] --> Filt{离群?}
Filt -- 是 --> Pool[写入 Outlier Pool]
Filt -- 否 --> Q2[2-bit 量化]
end
subgraph Decoding(每步)
Q[Query] -->|× Quantized K| A
Q -->|× Residual K| B
Q -->|× Outlier K| C
A & B & C --> D[Concat & Softmax] --> E[× 对应 V 得 output]
end
4. 结果与消融
正常上下文:GSM8K/BBH/HumanEval 等基准下,OTT 平均准确率比 KIVI 高 1.8–12.9 pp,接近 FP16。(arxiv.org)
长上下文 LongBench:在 LLaMA-2-13B/Mistral-7B 等模型上,OTT 仅比 FP16 低 0.2–0.6 pp,却稳压 KIVI。(arxiv.org)
效率:A800-80G 上批大小 32 时吞吐提升 2.3×;最坏情况下显存削减 6.4×。(arxiv.org, arize.com)
超参数消融:增大 Residual Window 明显提升准确率;Group Size 对准确率影响不大但能改善压缩率,故默认G=128,R=32。(arxiv.org)
5. 与相关工作的区别
6. 实践要点与启示
步长 G 的选择:若追求极限压缩,可将G 调大到 256,但需同步增大 Residual 以守住准确率。(arxiv.org)
Outlier Pool 大小:论文用P=32 已可覆盖 8 K 序列;对更长序列可适当放宽。(arxiv.org)
与 KV Sparsity 互补:可先用 OTT 量化,再对 Quantized KV 做 Top-K 稀疏存取,进一步降低显存。(deeppapers.dev)
硬件友好:全部运算只需 INT2/FP16 张量核与一次 fused kernel;对 A100/H100、Llama.cpp CPU 路径均易集成。(paperswithcode.com, huggingface.co)
参考文献
原论文 arXiv:2505.10938 (arxiv.org)
Arize AI 解读博客 (arize.com)
YouTube 研讨会录像 (youtube.com)
DeepPapers 播客笔记 (deeppapers.dev)
ChatPaper 中文速读 (chatpaper.com)
PapersWithCode 索引 (paperswithcode.com)
AIMODELS 摘要 (aimodels.fyi)
ResearchGate 条目 (researchgate.net)
KIVI 基线论文 (arxiv.org)
HuggingFace KV quantization 博文 (huggingface.co)