TailorKV: A Hybrid Framework for Long-Context Inference via Tailored KV Cache Optimization

发表时间: 2025-07 · ACL 2025 Findings

原文: https://aclanthology.org/2025.findings-acl.1043

作者/机构：
Dingyu Yao$^{1,2}$, Bowen Shen$^{1,2}$, Zheng Lin$^{1,2}$, Wei Liu$^3$, Jian Luan$^3$, Bin Wang$^3$, Weiping Wang$^1$

  1. Institute of Information Engineering, Chinese Academy of Sciences, Beijing, China
  2. School of Cyber Security, University of Chinese Academy of Sciences, Beijing, China
  3. MiLM Plus, Xiaomi Inc, Beijing, China
速读

一句话结论
提出了一种名为 TailorKV 的混合 KV Cache 压缩框架，通过为不同网络层分别定制 1-bit 量化或动态稀疏检索策略，在单张 RTX 3090 上实现了 Llama-3.1-8B 模型 128k 长上下文的近乎无损推理。

要解决什么问题
在生成式大语言模型的长上下文推理中，KV Cache 的显存占用会随序列长度线性增长，成为核心卡点。现有的缓解方案主要分为两类，但都存在致命缺陷：第一类是不可逆的驱逐策略（如 StreamingLLM、SnapKV），它们会永久丢弃部分 token，导致模型在多轮对话或需要全局信息的任务中出现严重的精度下降；第二类是可召回的选择策略（如 Quest、PQCache），它们将完整的 KV Cache 卸载到廉价的 CPU 内存中，但在解码阶段需要频繁将大量 token 召回 GPU，受限于 CPU-GPU 之间极低的 PCIe 带宽（例如传输单层缓存需 2 秒，而计算仅需 10 毫秒），会引入极高的 I/O 延迟。此外，现有研究通常对所有网络层采用统一的压缩或稀疏比例，但作者观察到，模型较浅层倾向于密集注意力以维持全局信息，对量化鲁棒但对稀疏敏感；而其他层则聚焦于少数主导 token，存在大量冗余，对稀疏鲁棒但对量化敏感。强行使用统一策略不可避免地会导致性能崩塌。

怎么做的
TailorKV 的核心思路是“因层制宜”，将量化与卸载检索无缝结合。整个方法由离线层类型识别、静态量化和动态检索三个关键部件构成。首先，在不干扰标准推理的前提下，通过离线计算密集偏好得分 $\mathcal{P}$ 来对 Transformer 层进行分类。给定预填充阶段的注意力分数矩阵 $\hat{\mathbf{A}}$，选取前 $k$ 个最大值所在的索引集合 $\hat{\mathcal{I}}$，得分定义为：

$$ \mathcal{P} = n_q - \sum_{(i,j) \in \hat{\mathcal{I}}} \hat{\mathbf{A}}_{i,j} $$


当某层的 $\mathcal{P}$ 大于预设阈值时，将其判定为“量化友好层”（通常是浅层），否则判定为“稀疏友好层”。其次，对于量化友好层，系统采用激进的静态 1-bit 量化。由于异常值主要存在于 Key 缓存的通道维度，而 Value 缓存没有异常值，因此对 Value 采用逐 token 量化，对 Key 采用逐通道量化，并配合专门的 FP16×INT1 算子以保证硬件执行效率。最后，对于稀疏友好层，系统在预填充阶段将 KV Cache 逐层卸载到 CPU，并在解码阶段动态检索 Top-K 的 token。为了绕开 PCIe 带宽瓶颈，作者设计了异步的两阶段检索流水线。由于注意力分数高度依赖 Query 和 Key 中的动态异常通道，且相邻层的隐藏状态具有高度相似性，系统在第 $l-1$ 层时，就利用当前隐藏状态提前估算第 $l$ 层的 Query $\hat{\mathbf{q}}$，并计算第 $i$ 个通道的贡献度：

$$ s_i = |\hat{\mathbf{q}}_i| \cdot \max(|\mathbf{K}_i|) $$
基于该贡献度，系统提前从 CPU 预取第 $l$ 层的关键 Key 缓存（即异常通道），随后在 GPU 上近似计算注意力分数并筛选出 Top-K 的 token 索引，最后仅将这极少量的完整 token 提取到 GPU。这种 CPU-GPU 异步协同设计将通信与计算重叠，极大掩盖了数据传输延迟。

效果如何
实验评估使用了 Llama-3.1-8B、Yi-6B 和 Yi-9B 模型，上下文长度覆盖 128k 到 200k。硬件平台包括 PCIe 1.0 环境的 RTX 3090 和 PCIe 4.0 环境的 A100。对比基线涵盖了代表不可逆驱逐路线的 StreamingLLM 和 SnapKV，以及代表可召回选择路线的 Quest 和 PQCache。在 LongBench、InfiniteBench 和 RULER 三个长上下文基准测试中，TailorKV 均展现出极强的压制力。在 Llama-3.1-8B 上，TailorKV 仅需对 1 到 2 层进行 1-bit 量化，其余层仅加载 1% 到 3% 的 token，即可实现近乎无损的精度，平均准确率比最强基线高出 2.32%。在资源受限的单张 RTX 3090 上，运行 128k 上下文的 Llama-3.1-8B 时，TailorKV 将显存峰值降低了 53.7%，解码延迟控制在每 token 82 毫秒；而在 A100 上测试 64k 上下文时，其延迟比 PQCache 快 1.2 到 2.0 倍，几乎与全量缓存的计算速度持平。作者也坦诚了该方法的局限性：当前设计主要优化了解码阶段的效率，预填充阶段的卸载延迟仍难以完全被计算重叠掩盖。

主要贡献

生成式大型语言模型（LLMs）中的键值（KV）缓存引入了巨大的内存开销。现有的工作通过卸载或压缩KV缓存来减轻这一负担。然而，由于CPU-GPU通信中的PCIe带宽瓶颈，加载整个缓存会引发显著的延迟，而激进的压缩则会导致明显的性能下降。本文的核心研究目标是同时优化长上下文LLM推理中的准确性、内存和延迟。

本文的主要创新点如下：

  1. 确定了特定层的压缩偏好，并开发了一种识别指标来确定模型中不同层的最佳压缩策略。浅层具有密集的注意力模式并强调全局信息，更适合统一压缩（如量化）；而其他层主要关注少数具有主导激活的tokens，存在大量冗余信息，适合稀疏性（如选择性加载）。
  2. 提出了一种名为 TailorKV 的混合KV缓存压缩框架，该框架通过算法-系统协同设计，将量化和卸载技术无缝集成，在保留模型精度的同时保持了执行效率。
  3. 在长上下文基准测试上的广泛实验表明，TailorKV在激进的压缩设置下实现了几乎无损的性能，超越了现有最先进的方法。特别是，具有128k上下文的 Llama-3.1-8B 可以在单块 RTX 3090 GPU 上提供服务，解码期间达到每个token $82\mathrm{ms}$ 的速度。

背景知识与关键观察

LLM推理的两阶段过程。 LLM推理包含prefill和decode两个阶段。在prefill阶段，使用整个提示（prompt）生成第一个token。给定提示嵌入 $\mathbf{X} \in \mathbb{R}^{n \times d}$ 以及权重矩阵 $\mathbf{W}_i^q, \mathbf{W}_i^k, \mathbf{W}_i^v \in \mathbb{R}^{d \times d_h}$（对于头 $i \in [1, h]$），其中 $n$ 是序列长度，$d$ 是隐藏维度，$d_h$ 是头维度。计算并缓存头 $i$ 的键和值：$\mathbf{K}_i = \mathbf{X} \mathbf{W}_i^k, \mathbf{V}_i = \mathbf{X} \mathbf{W}_i^v$。在decode阶段，迭代计算新的token嵌入 $\mathbf{x} \in \mathbb{R}^{1 \times d}$ 以产生查询、键和值向量。缓存被更新，每个注意力头的输出 $\mathbf{o}_i$ 计算如下：$\mathbf{K}_i \gets \mathrm{Cat}[\mathbf{K}_i, \mathbf{xW}_i^k], \mathbf{V}_i \gets \mathrm{Cat}[\mathbf{V}_i, \mathbf{xW}_i^v]$，$\mathbf{a}_i = \mathrm{Softmax}\big(\mathbf{q}_i \mathbf{K}_i^\top / \sqrt{d_h}\big), \mathbf{o}_i = \mathbf{a}_i \mathbf{V}_i$，其中 $\mathbf{q}_i = \mathbf{x} \mathbf{W}_i^q$，所有头的注意力输出被拼接并发送到FFN。

KV缓存的量化过程。 量化将连续或高精度值转换为低精度离散表示。给定高精度张量 $\mathbf{X}$，典型的均匀量化过程可表示为：$\mathbf{X}_Q = \mathrm{Quant}_b(\mathbf{X}, s, z) = \mathrm{clamp}(\lfloor \frac{\mathbf{X} - z}{s} \rceil, 0, 2^b - 1)$，其中 $\mathbf{X}_Q$ 表示 $b$-bit精度的量化张量，$z = \min \mathbf{X}$ 作为零点，$s = \max \mathbf{X} - \min \mathbf{X}b$ 作为缩放因子。clamp函数将值限制在 $b$-bit 整数范围内，$\lfloor \cdot \rceil$ 表示舍入函数。

GPU-CPU协同执行的瓶颈。 随着序列长度增加，KV缓存大小增长，显著提高了对GPU资源的需求（例如，在512k序列长度下，Llama-2-7B需要高达256GB的内存）。当前的LLM服务系统【Kwon等，Efficient memory management for large language model serving with pagedattention+2023+SOSP】【Qin等，Mooncake: Trading more storage for less computation...+2025+FAST】采用卸载策略，将KV缓存存储在经济高效的CPU内存中，并在推理时将其加载到GPU。然而，由于低带宽的PCIe接口，I/O传输延迟成为推理瓶颈【Zhang等，Dovetail: A cpu/gpu heterogeneous speculative decoding...+2024】。例如，通过PCIe 1.0链路（4GB/s）将单层KV缓存（约8GB）从CPU内存传输到RTX 3090 GPU需要约2秒，而在RTX 3090 GPU上单层的注意力计算仅需约10毫秒。因此，按需获取是目前减少GPU空闲时间最常用的方法。

不同层具有压缩偏好。 与之前认为所有层都适合稀疏性的观点【Li等，SnapKV: LLM knows what you are looking for before generation+2024+NeurIPS】【Zhang等，H2O: heavy-hitter oracle for efficient generative inference...+2023+NeurIPS】不同，作者提出并非所有层都适合稀疏化。为了量化解码期间的稀疏挑战，定义了稀疏误差 $\mathcal{E}$。设 $\mathbf{a} \in \mathbb{R}^{1 \times n}$ 代表注意力权重，$\mathcal{M} \in \{0, 1\}^n$ 表示选择 $\mathbf{a}$ 中前 $k$ 个元素的二值掩码。每个头的稀疏误差 $\mathcal{E}$ 定义为：$\hat{\mathbf{a}} = \mathbf{a} \odot \mathcal{M}, \mathcal{E} = 1 - \sum_{i=1}^n \hat{\mathbf{a}}_i$。如图1a所示，具有密集注意力分布的层比具有稀疏分布的层表现出更高的稀疏误差。此外，观察到跨模型和数据集的稀疏误差模式。图1b显示稀疏误差在不同模型间相似，较浅的层（如0、1层）具有较高的稀疏误差。图1c显示，对于同一模型，不同数据集上的稀疏误差分布保持一致。
图1：关于注意力的观察。(a) Llama-2-7B-32K-Instruct上的注意力权重。(b) 仅保留前5%注意力分数时，不同模型在2WikiMQA数据集上的稀疏误差。(c) 仅保留前5%注意力分数时，在不同数据集上的稀疏误差。

部分层不适合量化。 同样，并非所有层都适合量化。如表1所示，将KV缓存量化为1-bit会导致显著的性能下降。这种下降主要由具有稀疏分布的层引起，它们对量化更敏感。相反，量化密集层（如第0层）不会产生性能损失。这些发现强调了定制KV缓存压缩策略的需求。作者将具有密集分布的层视为量化友好型（侧重于全局信息），将具有稀疏分布的层视为稀疏友好型（优先考虑关键信息）。

注意力分数与异常值的相关性。 键和查询中的每个通道通过它们的点积（公式为 $\mathbf{qK}^\top$）对注意力分数做出贡献。图2（上）说明查询和键中的某些通道具有很大的幅度。从点积公式可以得出，注意力分数与这些异常值相关。最近的方法【Yang等，Post-training sparse attention with double sparsity+2024】关注静态通道稀疏性，利用离线校准技术识别大幅度通道。然而，作者发现查询和键通道的稀疏性是动态的而非静态的。如图2（下）所示，查询和键中的异常值并不总是出现在固定位置；相反，它们可能出现在任何位置。此外，与使用静态离线策略相比，动态选择大幅度通道提高了主导tokens的召回率。
图2：(上) Llama-3.1-8B-Instruct中的查询和键在某些通道中显示出异常值模式，而值没有异常值。(下) 达到Top-8的次数。异常值可能出现在任何位置。

方法细节

量化与稀疏偏好的离线识别指标。 经验观察表明，某些层更受益于量化，而其他层更适合稀疏性。为了避免破坏标准推理，应用离线策略来识别每层的压缩偏好。在这个阶段，引入了一个指标——密集偏好分数 $\mathcal{P}$——来评估每个注意力层是倾向于量化还是稀疏。给定提示长度 $n$，首先在prefill阶段使用最近的 $n_q$ 个查询向量 $\mathbf{Q}_{\mathrm{last\_q}} \in \mathbb{R}^{n_q \times d_h}$ 和键向量 $\mathbf{K} \in \mathbb{R}^{n \times d_h}$ 计算每个头的注意力分数矩阵 $\hat{\mathbf{A}}$：$\hat{\mathbf{A}} = \mathrm{Softmax}\big(\mathbf{Q}_{\mathrm{last\_q}} \mathbf{K}^\top / \sqrt{d_h}\big)$。接下来，从 $\hat{\mathbf{A}}$ 中选择前 $k$ 个索引，并对前 $k$ 个元素求和以计算密集偏好分数 $\mathcal{P}$：$\hat{\mathcal{I}} = \left\{ (i, j) \mid \mathrm{Top}_k(\hat{\mathbf{A}}_{i, :}, k) \right\}_{i=1}^{n_q}$，$\mathcal{P} = n_q - \sum_{(i, j) \in \hat{\mathcal{I}}} \hat{\mathbf{A}}_{i, j}$。

层级别的分类与工作流。 如果第 $l$ 层的密集偏好分数 $\mathcal{P}_l$ 超过阈值 $\tau$，则该层被视为量化友好型；否则，被认为是稀疏友好型。这可以形式化为：$C(l) = \mathsf{Quantization\text{-}Friendly}$ 如果 $\mathcal{P}_l > \tau$，否则为 $\mathsf{Sparsity\text{-}Friendly}$。阈值 $\tau$ 是一个预定义的超参数，其最佳值通过在合成的Longbench任务上进行实验确定。指标 $\mathcal{P}$ 在不同数据集上对同一模型的评估保持一致。在层级识别之后，对稀疏友好层应用动态检索，对量化友好层应用静态量化。整体工作流如图3所示。
图3：TailorKV的系统概述。离线识别将层分类为量化友好型和稀疏友好型。对于量化友好型层，采用激进的静态量化。对于稀疏友好型层，动态检索Top-K tokens。关键当前查询和关键键缓存分别代表查询和键缓存中的异常值。

稀疏友好层的动态检索与异步系统设计。 对于稀疏友好层，提出了一种带有异步系统设计的动态检索算法。为了在内存受限的设备上促进LLM推理，在prefill期间逐层将KV缓存卸载到成本较低的CPU内存中。随后，在decode期间按需检索Top-K tokens，从而最小化通信开销。
图4：两阶段动态检索过程：阶段1在第l-1层估计关键通道，并为第l层预取关键键缓存。阶段2在第l层近似注意力分数并选择Top-K tokens。

异常值识别与关键通道估计。 注意力分数与查询和键中的异常值相关。为了更准确地评估token重要性，基于此洞察在原始操作之前近似注意力分数。首先估计关键通道以识别查询和键缓存中的异常值，称为关键当前查询和关键键缓存。由于关键键缓存驻留在CPU中，采用预取（prefetching）提前加载它。利用层间相似性提前预测关键通道。相邻层之间的相似性源于残差连接，这在先前的研究中得到了验证【Lee等，Infinigen: Efficient generative inference...+2024+OSDI】。在第 $l - 1$ 层，使用第 $l$ 层的权重矩阵和第 $l - 1$ 层的隐藏状态估计第 $l$ 层的查询 $\hat{\mathbf{q}}$。第 $i$ 个通道对注意力分数的贡献通过 $\hat{\mathbf{q}}$ 和 $\mathbf{K}$ 的逐元素乘法计算：$\mathbf{s}_i = |\hat{\mathbf{q}}_i| \cdot \max(|\mathbf{K}_i|), i = 1, 2, ..., d_h$。

双缓冲预取与Top-K检索通信。 接下来，基于 $\mathbf{s}$ 预取第 $l$ 层的关键键缓存，使用双缓冲——一个缓冲区用于写入，另一个用于读取——以实现并发执行。然后，基于关键当前查询和关键键缓存近似第 $l$ 层的注意力分数来检索Top-K tokens，随后获取Top-K tokens。图5概述了解码期间的计算和通信。唯一不可重叠的操作是获取Top-K tokens，因为它依赖于当前层的查询。TailorKV展示了异构设计如何通过利用CPU-GPU协同执行来克服资源限制。
图5：动态检索的时间线。蓝色表示计算，粉色表示通信。

量化友好层的静态量化策略。 与传统的量化方法【Liu等，Kivi: a tuning-free asymmetric 2bit quantization...+2024+ICML】【Yang等，No token left behind...+2024】【He等，Zipcache: Accurate and efficient KV cache quantization...+2024+NeurIPS】不同，TailorKV专注于确保每一层“发挥其作用”，从而实现更激进的压缩方案，如1-bit量化。异常值存在于键缓存的通道维度中，而值缓存不包含异常值。对于量化友好层，对值缓存应用静态per-token量化，对键缓存应用静态per-channel量化【Liu等，Kivi...+2024+ICML】。引入了1-bit量化内核，并实现了 $\mathrm{FP}16 \times \mathrm{INT}1$ GEMV，以提高激进压缩下的硬件性能。

内存占用分析。 设层数为 $L$，头数为 $h$，序列长度为 $n$，头维度为 $d_h$。所有输入tokens均以FP16表示。TailorKV主要在量化友好层管理量化KV缓存缓冲区，在稀疏友好层管理关键键缓冲区。量化零点和缩放因子以FP16格式存储。

实验环境

  • 数据集：

    • LongBench：包含13个子任务，聚合为6类（单文档QA、多文档QA、摘要、少样本学习、代码补全、合成任务）。
    • InfiniteBench：包含9个子任务，聚合为5类（检索、对话、小说、数学、代码），用于评估极长上下文（200k+ tokens）处理能力。
    • RULER：包含13个代表性任务（问答、检索、聚合、多跳追踪等），序列长度从4K到128K不等。

  • 模型架构关键参数：评估了三种广泛使用的模型：Llama-3.1-8B-Instruct (GQA, 128k上下文, 第0层为量化友好层)、Yi-6B-200K (GQA, 200k上下文, 第0,1层为量化友好层)、Yi-9B-200K (GQA, 200k上下文, 第0,1层为量化友好层) 以及 Llama-2-7B-32K-Instruct (MHA, 32k上下文, 第0,1层为量化友好层)。

  • 硬件配置：

    • 配置一：NVIDIA RTX 3090 GPU (24GB) + Intel Xeon Gold 6240 CPU，通过PCIe 1.0 x16 (4GB/s) 互连。
    • 配置二：NVIDIA A100 GPU (80GB) + Intel Xeon Platinum 8369B CPU，通过PCIe 4.0 x16 (32GB/s) 互连。

  • 软件配置：基线方法包括 StreamingLLM、SnapKV、Quest、PQCache、Full Cache (通过 FlashAttention-2 实现) 和 OffloadCache。TailorKV 的 $\tau$ 设为0.2。TailorKV-1 和 TailorKV-2 分别表示在量化友好层使用 1-bit 和 2-bit 精度，组大小为64。稀疏友好层中，LongBench的关键通道数设为8，InfiniteBench和RULER设为12。

实验结果

  • 长上下文任务准确性 (LongBench & InfiniteBench)：

    • 实验内容：在 LongBench 和 InfiniteBench 上对比不同方法的性能。
    • 结果与分析：由于丢失了关键信息，SnapKV和StreamingLLM性能下降。Quest和PQCache在预算受限时面临局限性。TailorKV通过在量化友好层保留1-bit KV缓存，在稀疏友好层选择192个tokens，在 Llama-3.1-8B、Yi-9B 和 Yi-6B 上分别比最佳基线方法高出 2.32%、5.42% 和 3.66% (表3)。在 InfiniteBench 上，PQCache 由于 K-Means 聚类开销随上下文长度增加而增大，限制为单次迭代导致准确率下降。相比之下，TailorKV 的混合策略优于单一策略，与全缓存相比平均性能损失不到 1.5%，在对话、小说和数学任务中表现尤为出色。

  • RULER基准测试准确性：

    • 实验内容：评估序列长度从 4K 到 128K 的长上下文建模能力。
    • 结果与分析：TailorKV 从冗余上下文中捕获了关键信息，在大多数任务（如大海捞针、问答和变量追踪）上取得了卓越的性能。
      图6：不同方法在RULER上的平均准确率。TailorKV中的稀疏友好层使用128+(896)个tokens，而其他方法使用1024个tokens。

  • 峰值内存使用量：

    • 实验内容：在 A100 GPU 上评估不同方法的峰值内存占用。
    • 结果与分析：TailorKV 实现了卓越的内存效率。具体而言，与全缓存相比，TailorKV 在 Llama-2-7B 序列长度为 128k 时，GPU内存使用量减少了约 73.8%，从而使其能够部署在较低端的 GPU（如 RTX 3090）上。
      图7：A100 (80GB) 上的峰值内存使用量。

  • 端到端延迟与延迟分解：

    • 实验内容：在 RTX 3090 和 A100 上评估解码延迟。
    • 结果与分析：在 RTX 3090 上，序列长度增加导致基线方法 OOM，而 TailorKV 能正常运行。在 A100 (64k上下文) 上，TailorKV 的延迟显著低于 OffloadCache 和 PQCache：对于 MHA 模型分别快 $18.0\times$ 和 $1.2\times$，对于 GQA 模型分别快 $8.1\times$ 和 $2.0\times$。TailorKV 的延迟与全注意力相当，这归功于使用多线程执行异步任务，实现了计算和 CPU-GPU 通信的重叠。延迟分解显示，与 PQCache 相比，TailorKV 将检索延迟降低了 27.8% (GQA) 和 40.5% (MHA)，将数据传输延迟降低了 83.5% 和 82.2%（主要归功于使用 DGL 直接将行从 CPU 张量传输到 GPU）。
      图8：不同方法下的延迟分解 (ms)。Q和S分别表示量化友好层和稀疏友好层。

  • 消融实验：

    • 实验内容：在 LongBench 上使用 Llama-3.1-8B-Instruct 评估定制策略、动态通道和关键通道数量的影响。
    • 结果与分析：(1) 定制策略的影响：仅量化第0层（量化友好层）产生最佳性能，而量化稀疏友好层会降低性能，突出了定制压缩策略的必要性。(2) 动态通道的影响：动态检索方法的性能优于静态选择大幅度通道的方法。(3) 关键通道数量的影响：选择8个关键通道在性能和延迟之间取得了良好的平衡，将数量减少到2或4会显著降低性能。
      图9：消融实验。(a) 将不同层量化为1-bit的性能比较。(b) 具有动态或静态通道的TailorKV性能。(c) 不同关键通道数量的性能比较。

补充细节

现有KV缓存压缩方法。 现有的方法主要包括驱逐（eviction）、选择（selection）和量化（quantization）。

  • 驱逐方法：通过在推理期间驱逐大多数tokens来减小KV缓存大小。StreamingLLM【Xiao等，Efficient streaming language models with attention sinks+2024+ICLR】通过保留初始和最近的tokens来识别“注意力汇（Attention Sinks）”。H2O【Zhang等，2023+NeurIPS】、SnapKV【Li等，2024+NeurIPS】和Scissorhands【Liu等，Scissorhands: Exploiting the persistence of importance hypothesis...+2023+NeurIPS】基于历史注意力分数估计token重要性。然而，驱逐主要tokens可能会降低在“大海捞针”和多轮对话任务中的准确性。
  • 选择方法：常用于稀疏注意力场景。Quest【Tang等，Quest: queryaware sparsity...+2024+ICML】保留KV缓存并利用分页键检索tokens，但它未能减少内存使用且召回率较低。KV缓存卸载方法如PQCache【Zhang等，Pqcache: Product quantization-based kvcache...+2024】和InfiniGen【Lee等，2024+OSDI】近似注意力分数，以识别并从CPU加载关键tokens到GPU，但由于KV缓存负载大，它们面临平衡计算和通信的挑战。部分方法【Chen等，MagicPIG: LSH sampling for efficient LLM generation+2025+ICLR】【Liu等，Retrievalattention: Accelerating long-context llm inference via vector retrieval+2024】使用LSH和KNN检索关键tokens，在CPU上处理后随后与GPU输出合并；然而，不平衡的计算时间可能导致GPU空闲时间。
  • 量化方法：将高精度浮点数转换为低精度整数。KVQuant【Hooper等，Kvquant: Towards 10 million context length llm inference...+2024+NeurIPS】隔离异常值以实现混合精度，GEAR【Kang等，Gear: An efficient kv cache compression recipe...+2024】利用SVD恢复残差，KIVI【Liu等，2024+ICML】对键进行每通道量化，对值进行每token量化。FlexGen【Sheng等，Flexgen: High-throughput generative inference...+2023+ICML】将KV缓存量化为4-bits以减少I/O传输延迟。然而，这些方法都没有将KV缓存减少到1-bit。相比之下，本文专注于探索层特征并选择最合适的压缩策略。

结论

本文提出了 TailorKV，一个用于 LLM 中 KV 缓存管理的有效框架。通过观察到不同层表现出不同的压缩偏好，将其分类为量化友好型和稀疏友好型，并分别采用定制策略。量化友好型层激进地将 KV 缓存量化为 1-bit；稀疏友好型层基于查询和键通道中的大幅度异常值动态检索主导 tokens，并集成了 CPU-GPU 协同设计。长上下文基准测试表明，TailorKV 在保持模型性能的同时，有效最小化了 KV 缓存的使用，且延迟代价可接受。该混合框架展示了在资源受限的 GPU 上部署 LLM 的潜力。

局限性与未来工作：(1) TailorKV 主要侧重于通过异步传输提高 decode 阶段的效率。然而，在 prefill 阶段完全重叠卸载延迟具有挑战性。该工作与用于 prefilling 加速的方法是兼容且互补的。(2) 目前为不同层设计了定制策略，未来研究计划进一步探索在 Head 级别（head-wise）进行自适应调整。

附录细节

与其他方法的比较。 图10比较了TailorKV与其他方法：(a) Full cache保留整个KV缓存。(b) 驱逐（eviction）方法从每层永久驱逐特定tokens，导致不可逆的信息丢失，因为被驱逐的tokens稍后可能变得重要。(c) 选择（selection）方法将整个KV缓存卸载到CPU，允许召回tokens，但由于涉及大量tokens，会产生显著的通信开销。(d) 本文的方法采用特定层的压缩策略，促进更激进的压缩。
图10：TailorKV与其他方法在管理跨层KV缓存预算方面的比较。

层间相似性计算。 设 $\mathbf{h}^{(l)}$ 表示第 $l$ 层的隐藏状态。为了量化相邻层隐藏状态之间的相似性，采用余弦相似度：$\mathrm{sim}(\mathbf{h}^{(l-1)}, \mathbf{h}^{(l)}) = \frac{\mathbf{h}^{(l-1)} \cdot \mathbf{h}^{(l)}}{\|\mathbf{h}^{(l-1)}\| \|\mathbf{h}^{(l)}\|}$。定义第 $l$ 层的查询权重为 $\mathbf{W}_q^{(l)}$，第 $l$ 层的查询向量计算为：$\mathbf{q}^{(l)} = \mathbf{W}_q^{(l)}(\mathbf{h}^{(l)})$。如图11所示，$\mathbf{h}^{(l)}$ 和 $\mathbf{h}^{(l-1)}$ 非常相似，允许基于第 $l - 1$ 层的隐藏状态近似第 $l$ 层的查询：$\hat{\mathbf{q}}^{(l)} = \mathbf{W}_q^{(l)}(\mathbf{h}^{(l-1)})$。现有研究【Liu等，Minicache: Kv cache compression in depth dimension...+2024+NeurIPS】阐明KV缓存在相邻层表现出相似性。然而，如图11所示，$\hat{\mathbf{q}}^{(l)}$ 和 $\mathbf{q}^{(l)}$ 之间的相似度超过了 $\mathbf{q}^{(l-1)}$ 和 $\mathbf{q}^{(l)}$ 之间的相似度，表明使用前一层的隐藏状态能提高预测准确性。
图11：相邻层之间的余弦相似度。

不同数据集上的离线识别。 如图12所示，曲线代表不同数据集。同一模型在不同数据集上的 $\mathcal{P}$ 分布是一致的，表明指标 $\mathcal{P}$ 有效捕获了不同层的特征，从而能够采用适当的压缩策略。
图12：不同离线数据集各层的密集偏好分数P。

与混合方法的比较。 为了验证量化-稀疏混合框架的有效性，将其与类似的混合方法 SimLayerKV【Zhang等，Simlayerkv: A simple framework for layer-level kv cache reduction+2024】进行比较。SimLayerKV 假设 LLM 中某些层表现出“懒惰”行为（仅保留初始和最近的tokens），而“非懒惰”层需要全精度以保留所有tokens。表7展示了在 LongBench 上的实验结果。结果表明，在平均压缩率 $34.2\times$ 下，TailorKV 的性能与 SimLayerKV 在 $1.53\times$ 压缩率下的性能相当。TailorKV 以最小的内存开销实现了最佳性能，为这种量化-稀疏混合架构的实用性提供了有力证据。相比之下，SimLayerKV 需要基于历史注意力分数实时识别层类型，这使其与 FlashAttention 不兼容。这引入了额外的计算和内存开销，从而增加了延迟并可能导致 OOM 问题。

动态检索的有效性。 表8比较了稀疏友好层与其他方法的检索准确性（在 Llama-3.1-8B-Instruct 模型上，LongBench 基准）。具体而言，在 StreamLLM、SnapKV 和 Quest 的第0层保留全精度 KV 缓存，从而保存第0层的全局信息。TailorKV-1 和 TailorKV-2 分别表示将第0层的 KV 缓存量化为 1-bit 和 2-bit 精度。实验结果表明，当第0层保留全局信息时，TailorKV 的检索方法优于其他稀疏方法。具体来说，TailorKV 对第0层应用量化，而其他方法使用全精度（16-bit），并且从第1层到第31层的注意力计算使用相同的 tokens。这一显著的性能优势凸显了该检索方法能有效识别最重要的 tokens，从而最小化关键信息的丢失。