浏览: 4

在大模型推理中，如果直接把不重要的 KV 对整块丢弃（eviction），不仅会丢失上下文细节，还可能触发安全失效、幻觉或对话错乱等严重问题(arxiv.org, arxiv.org)。


《No Token Left Behind》提出的 MiKV（Mixed-precision KV cache） 通过“重要度驱动的混合精度量化”让缓存既大幅压缩又可靠：把待“驱逐”的 KV 对保留下来并压到极低位宽（INT2/INT3/INT4），而真正重要的 KV 对继续用较高位宽（FP16/INT8）存放，同时用动态离群通道调节（outlier-aware）解决低比特下的量化误差，并给出了可在 GPU 上高效运行的融合 Kernel。下文分四步解读其设计。

1 问题与动机

  • KV Cache 膨胀瓶颈：缓存大小随批量和序列线性增长，常常比模型权重还大，长上下文场景尤甚(arxiv.org, arxiv.org)。

  • 单纯 Eviction 的风险：H2O 等方法只保留“重击词(Heavy Hitters)”，但彻底丢弃其余 KV 对会引发安全绕过、幻觉和上下文遗漏(arxiv.org)。

  • 关键观察：即便把被驱逐的 KV 对用很低的位宽保存一点点信息，也能显著恢复性能；反之，真正重要的 KV 对若降到过低精度则会劣化生成质量(arxiv.org, arxiv.org)。

2 MiKV 三大核心组件

2.1 保留驱逐 KV 对的低比特量化 (Section 3.1)

  1. 先用现有策略（如 H2O）选出“暂定驱逐集”与“重要集”(arxiv.org, arxiv.org)。

  2. 对驱逐集做 逐 token 非对称 RTN 量化，位宽可选 2/3/4 bit，scale 用公式 α=(max-min)/(2^b-1)、zero = β=min (x) 计算(arxiv.org)。

  3. 实验表明：在 Llama-2-7B line-retrieval 任务里，50 % 驱逐集若直接删除准确率跌到 43 %，若改存 INT3 立刻回升到 99.8 % (arxiv.org)。

2.2 动态离群通道调节 (Section 3.2)

  • 发现 Key 与 Query 通常在少数通道上出现系统性 outlier，极易导致 INT2/3 量化误差暴涨(arxiv.org)。

  • 通道平衡器（balancer）：

    • 预填阶段统计每层-每头-每通道的最大值

      b_{l,h,c}= \sqrt{\frac{\max q_{0:t-1}^{lhc}}{\max k_{0:t-1}^{lhc}}} \]:contentReference[oaicite:8]{index=8}

    • 后续生成时把 Key 乘 b，再写入量化；Query 除以 b，再去参与注意力计算 ⇢ 抵消通道极值。

  • 加上 balancer 后，INT2 缓存的 line-retrieval 准确率从 64 %→ 92.6 %，几乎追平 FP16，而额外显存仅 +1 % (arxiv.org)。

2.3 重要 KV 对的高精度混合 (Section 3.3)

  • 默认把 Top-20 % 重击词保持 FP16；若显存仍吃紧，可降到 INT8/INT4，但过低会拖精度（INT2 时跌到 65 %）(arxiv.org)。

  • 通过调节“高-低精度比例 + 位宽”，MiKV 提供弹性的压缩-精度折衷曲线，比纯量化或纯驱逐都更优(arxiv.org, arxiv.org)。

2.4 高速混合内核 (Section 3.4)

  • 观察到注意力在位置编码后 顺序可打乱，因此可把不同位宽的 KV 按精度分组、批量解码，再拼回去，不影响语义(arxiv.org)。

  • 借助现成 weight-only GEMV 核（Park 2022；Lin 2023）解码-乘法一次完成，避免 memory-wall 瓶颈。

  • 与全 FP16 相比，生成阶段 KV mat-vec 可提速 1.6–1.8 ×，同时显存降至 20–38 % (arxiv.org)。

3 整体流程示意图

flowchart LR
    subgraph 准备阶段
        A[预填 prompt] --> B[计算通道 balancer  b_{l,h,c}]
    end
    subgraph 每步生成
        C[新 KV 对] --> D{是否重击词?}
        D -- 是 --> E1[高精度缓存<br/>FP16/INT8/INT4]
        D -- 否 --> E2[低精度缓存<br/>INT2/3 + balancer]
        E1 & E2 --> F[按精度分组 解码+RoPE+GEMV (融合核)]
        F --> G[Attention 输出]
    end

4 实验亮点

模型

Cache 占比

方法

GSM8k 1-shot

Line Retrieval

显存 (B=8,L=4 096)

Llama-2-13B

100 %

Full

28.4

100 %

32 GB

″

20 %

H2O 驱逐

15.3

4 %

6.4 GB

″

20 %

Uniform INT3

23.0

68 %

6.4 GB

″

20 %

MiKV (INT2 驱逐 + balancer)

27.7

99.8 %

6.6 GB

(arxiv.org, arxiv.org)

MiKV 在 20 % 缓存量下就能把 Llama-2-13B 的 GSM8k 精度维持在 98 % 以上，同时显存仅为原来的 ~1/5；对 Line-Retrieval 这种极度依赖上下文的任务更是完胜驱逐策略。

5 与现有方案对比

类别

方法

思路

对低比特误差处理

安全/幻觉风险

压缩-精度折衷

驱逐

H2O (2023)

只存重击词

—

⚠️ 高

中→高掉点(arxiv.org)

纯量化

KVQuant (2024)

全 2-4 bit

通道量化+稀疏

✅

精度略降(arxiv.org)

MiKV

本论文

混合精度+balancer

动态 outlier 平衡

✅

最佳综合(arxiv.org, arxiv.org)

6 局限与未来方向

  • 在线重要度判定依赖外部策略（文中默认 H2O）；若策略误判仍会影响压缩效果。

  • 动态 balancer 目前在 prefill 阶段一次性计算，对流式场景可能需增量更新。

  • 论文使用 逐 token 量化，与 KVQuant 的通道量化尚未融合，可探索两者结合。

  • 针对 CPU/边缘设备的融合核实现仍待开源。

一句话总结

MiKV 用“驱逐→低比特保留”替代“直接扔掉”，再辅以动态离群通道平衡和精度分级缓存，成功把 KV Cache 压到 20 % 大小的同时守住安全性和生成质量，是当前最可靠的 KV 缓存压缩路径之一。