RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression

发表时间: 2025-07 · arXiv:2502.14051 (ICML 2025)

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

作者/机构：Payman Behnam 1 2, Yaosheng Fu 1, Ritchie Zhao 1, Po-An Tsai 1, Zhiding Vu 1, Alexey Tumanov 2

速读

一句话结论
提出了一种免训练的两阶段 KV Cache 压缩策略 RocketKV，通过结合粗粒度永久驱逐与细粒度动态稀疏注意力，在几乎不掉精度的前提下实现了最高 400 倍的压缩率、32.6% 的显存峰值降低以及 3.7 倍的解码端到端加速。

要解决什么问题
在基于 Transformer 的大语言模型推理中，解码阶段的 KV Cache 大小会随序列长度和批次大小线性增长，导致严重的显存容量和访存带宽瓶颈。例如，在批次大小为 32、上下文长度为 32K 时，Llama3.1-70B 模型需要约 320GB 的半精度 KV Cache，即使是最先进的硬件也难以承受。现有的免训练压缩方法主要分为两派：一派是永久性驱逐部分 Token，这能同时节省显存和带宽，但在极低 Token 预算下，如果后续解码需要用到早期被丢弃的 Token，会导致精度严重崩塌；另一派是动态选择 Token，即把所有 KV 保留在显存中，每次动态挑出一部分参与计算，这虽然缓解了精度问题，但无法节省显存容量，且在极低预算下难以准确预测真正的 Top-k 目标。作者通过对 LongBench 中的 qasper 任务进行分析发现，当 Token 预算低于 1024 时，现有实际方法的精度均大幅落后于理想的精确 Top-k 基线，说明它们在极高压缩率下均无法准确预测关键 Token。

怎么做的
核心思路是将永久驱逐与动态选择结合，先用较大的预算过滤掉明显无关的 Token，再在剩下的较小集合中做精准的动态 Top-k 选择。这种融合既剔除了无用信息，又降低了动态选择的预测难度。方法具体由以下几个关键部件构成：
第一阶段（粗粒度永久驱逐）：直接采用现有的 SnapKV 方法，利用提示词末尾的观察窗口计算注意力分数，永久丢弃不重要的 KV Token。为了保证信息的完整性，该阶段在序列维度使用了较大核尺寸（大小为 63）的池化操作，以保留关键 Token 的相邻上下文。这一步同时降低了显存占用和访存带宽。
第二阶段（细粒度动态选择）：提出了一种混合稀疏注意力机制（HSA）。它在序列维度和注意力头维度同时做降维近似，分三步执行：首先将 Key 张量按序列维度分块，预先提取并保存每个块在各个头上的最大值 $K_{\max}$ 和最小值 $K_{\min}$；接着在每次生成时，根据当前 Query 向量 $q$ 的绝对值大小，在头维度挑出最大的 $k_1$ 个头，并根据 $q$ 的符号去取对应的 $K_{\max}$ 或 $K_{\min}$，近似计算出每个块的注意力得分，从而在序列维度选出得分最高的 $k_2$ 个块；最后只拉取这 $k_2$ 个块的完整 KV 数据做精确注意力计算。
自适应压缩分解：为了在两阶段间合理分配总压缩率 $c$，设计了一个自适应切分因子 $r$：

$$r = \min(0.2 + 0.06 \times \log_2(c), 0.8)$$

第一阶段承担 $c^r$ 的压缩率，第二阶段承担 $c^{(1-r)}$ 的压缩率。压缩率越小，越倾向于少做永久驱逐以防信息丢失；压缩率越大，则让第一阶段多承担一些，因为基于精确分数的 SnapKV 比基于近似分数的 HSA 预测更准。
多轮对话变体（RocketKV-MT）：在多轮对话中，早期不重要的 Token 后来可能变得重要。因此该变体在第一阶段不真正从显存中删除 Token，而是保留全量历史，但限制第二阶段只能在第一阶段筛选出的子集里做动态选择。

效果如何
实验在 Llama3.1-8B-Instruct、Mistral-7B-Instruct-v0.2 和 LongChat-7B-v1.5-32k 三个模型上展开，硬件使用 NVIDIA A100 和 H100 GPU。对比基线包括：全量 KV 缓存（Full-KV）、理想的精确选择（Exact-TopK）、代表混合注意力架构的 DuoAttention、代表永久驱逐路线的 SnapKV，以及代表单维度动态选择路线的 Quest 和 SparQ。
在单轮长文本任务中，RocketKV 在极低预算下显著优于其他基线。在 Needle-in-a-Haystack 测试中，Llama3.1-8B-Instruct 在 256 的极低 Token 预算下（对应超过 400 倍压缩率、109K 上下文），RocketKV 依然保持了 100% 的准确率，而其他方法均出现严重掉点。在 RULER 基准测试中，RocketKV 的精度损失在短序列下微乎其微，仅在超出模型有效上下文长度时才逐渐显现。
在效率方面，RocketKV 在解码阶段实现了最高 32.6% 的峰值显存节省。由于 A100 的访存计算比低于 H100，推理过程更受访存带宽限制，因此 RocketKV 在 A100 上获得了 3.7 倍的端到端加速，略高于 H100 上的 3.3 倍。
局限与代价：在多轮对话基准 SCBench 中，基础版 RocketKV 在高预算下反而不如 SparQ，作者承认这是因为永久驱逐导致了跨轮次的信息丢失；但换用 RocketKV-MT 变体后，精度回升至与 Exact-TopK 相当，代价是该变体无法节省显存容量。此外，HSA 机制本身需要引入额外的辅助张量存储开销。

主要贡献

在基于Transformer的大语言模型（LLM）推理中，解码阶段的键值缓存（KV cache）大小会随着序列长度和批处理大小呈线性增长，这给内存带宽和容量带来了巨大的负担（例如，批处理大小为32、上下文长度为32K的Llama3.1-70B模型需要约320GB的KV缓存）。现有的KV缓存压缩方法主要分为两类：永久性KV token驱逐（节省带宽和存储，但可能导致准确率下降）和动态KV token选择（保留所有token动态选择，无准确率损失但仅节省带宽且需额外辅助存储）。如图1所示，现有方法在极低的token预算（低于1024）下，无法准确预测top-$k$的KV tokens，导致准确率显著下降。

为了解决这一挑战并提高预测准确率，本文提出了RocketKV，这是一种免训练的两阶段KV缓存压缩策略。其主要创新点如下：
1. 两阶段KV缓存压缩框架：结合了永久性KV token驱逐与动态KV token选择，第一阶段执行粗粒度的永久性缓存驱逐，第二阶段执行细粒度的动态top-$k$稀疏注意力。
2. 混合稀疏注意力（HSA）：在第二阶段提出了一种混合稀疏注意力机制，通过结合头部维度和序列维度的二维缩减来近似计算注意力分数，以进行高精度的动态选择。
3. 自适应压缩分解机制：引入了一种自适应机制，能够智能地将给定的目标压缩率在两个阶段之间进行分配平衡。
4. 针对多轮对话的RocketKV-MT变体：为了解决多轮对话中永久驱逐导致的上下文丢失问题，提出了RocketKV-MT，保留所有历史KV tokens，仅在每轮解码时对过滤后的子集进行动态选择。

背景知识与关键Observation

现有方法的局限性分析：如图1所示，无论是永久性KV token驱逐还是动态KV token选择，在低token预算下都无法达到oracle top-$k$注意力（Exact-TopK）的准确率。

准确率差距的根本原因探究：为了理解导致准确率不匹配的原因，作者分析了Mistral-7B-Instruct-v0.2中一个随机注意力头（第31层，第0头），并展示了在qasper基准测试的所有解码步骤中，最大序列长度和Exact-TopK（$k=256$）选择的唯一KV索引数量的累积分布函数（CDF）。进行此分析的原因在于，为了匹配Exact-TopK的准确率，需要保留在所有解码步骤中至少被一次top-$k$注意力操作选中的所有重要KV tokens。

重要发现与理论预算：如图2所示，尽管最大序列长度可以高达25000，但唯一的top-$k$索引数量仅达到1200。这意味着，理想情况下，一种永久性KV token驱逐方法应该能够在1200的token预算下消除准确率差距。

两阶段设计的动机：为了进一步降低token预算，作者意识到可以在经过永久性KV token驱逐后过滤出的KV token集合上应用动态KV token选择。由于过滤后的集合比原始的完整KV缓存小得多，准确预测top-$k$的难度大大降低。因此，一个理想的解决方案是首先使用较大的token预算执行永久性KV缓存驱逐，然后对剩余的KV tokens进行动态KV token选择。这种融合既驱逐了不重要的tokens，又使得动态选择更加准确，从而激发了RocketKV的设计。

方法细节

两阶段KV缓存压缩方法概述：基于上述观察，作者提出了RocketKV，这是一种用于加速解码阶段的两阶段KV缓存压缩方法。如图3所示，RocketKV在第一阶段执行粗粒度的KV缓存驱逐。此阶段的目的是移除重要性较低的KV tokens，同时保留大部分重要的tokens。

第二阶段与框架通用性：在第二阶段，它对剩余的KV tokens进行细粒度的动态KV token选择，随后执行top-$k$稀疏注意力。RocketKV框架是通用的，许多现有的KV缓存压缩方法都可以适用于相应的阶段。例如，SnapKV【Li等人，SnapKV: LLM knows what you are looking for before generation + 2024 + NeurIPS】或Ada-KV【Feng等人，Adakv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference + 2024 + arXiv】可用于第一阶段，而Quest【Tang等人，Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference + 2024 + ICML】或SparQ【Ribar等人，SparQ Attention: BandwidthEfficient LLM Inference + 2024 + ICML】可应用于第二阶段。为了在每个阶段达到最佳性能，作者在第一阶段直接采用了SnapKV，而在第二阶段提出了一种混合稀疏注意力（HSA）方法。

SnapKV的直接应用与核心思想：在第一阶段，作者直接采用SnapKV对输入KV tokens进行永久性KV token驱逐。SnapKV的核心思想是依赖于输入上下文和末尾观察窗口之间的聚合注意力分数，来选择输入提示中最相关的tokens并保留。

针对GQA的组级别token选择优化：原始的SnapKV方法是基于每个注意力头来选择关键KV tokens的。在分组查询注意力（GQA）的情况下，注意力组内的每个注意力头都保留一组独立的KV缓存tokens，这可能会导致同一KV token的冗余存储。为了减少GQA的KV token存储，作者遵循Ada-KV的工作，根据聚合的组级别注意力分数在每个组的基础上执行token选择，从而使选定的KV tokens在整个注意力组中共享。

池化核大小的经验性调整：SnapKV沿序列维度使用池化，以确保关键KV tokens与其相邻tokens一起被选中。它通过池化展示了更好的准确率，因为它保留了所选信息的完整性。其使用的池化核大小非常小（例如，对于LongBench核大小为7）。由于在RocketKV的情况下，SnapKV仅用于第一阶段的粗粒度KV token驱逐，作者发现最佳的池化核大小要大得多。作者在所有实验中经验性地将核大小设置为63。

现有动态选择方法的局限性：先前关于动态KV token选择的方法通常通过沿单一维度减少计算来估计top-$k$ KV索引。例如，Quest使用逐元素的最小值和最大值来表示沿序列维度的连续页面。同时，SparQ和Loki【Singhania等人，Loki: Low-Rank Keys for Efficient Sparse Attention + 2024 + NeurIPS】利用头部维度的稀疏性来进行低秩估计。不幸的是，依赖一维稀疏性只能达到一定程度的压缩率，超过该压缩率后准确率可能会迅速下降，如前文图1所示。

混合稀疏注意力（HSA）的提出：相比之下，作者提出了混合稀疏注意力（HSA），它同时利用了序列和头部维度的二维缩减，以在具有top-$k$注意力分数的KV token索引上实现更好的估计准确率。

HSA算法实现细节：图4展示了作者提出的算法的详细实现，该算法受Quest和SparQ的启发。HSA算法可以分解为三个步骤：
* 步骤1：将键张量（key tensor）中的tokens沿序列维度分组为连续的页面，并存储每个页面的逐元素最大值（$K_{max}$）和最小值（$K_{min}$）作为辅助存储，类似于Quest。与Quest不同的是，它们以不同的布局存储，通过沿头部维度对齐，以便在步骤2中进行高效收集。每当生成新的键token时，辅助存储会相应更新。
* 步骤2：对于每个查询$q$，沿头部维度找到$k_1$个最大绝对值。然后，根据$q$在这些索引处的符号，仅在逐元素最大值或最小值张量中获取相应的索引。目标是计算每个页面的逐元素 $\max(q \times K_{max}, q \times K_{min})$，以近似页面内的最高可能注意力分数。为了进一步减少近似开销，作者仅在头部维度上$q$的幅度较大的$k_1$个部分位置上进行计算，并忽略其他位置，类似于SparQ。一旦计算出近似注意力分数，就会沿序列维度选择具有最大注意力分数的$k_2$个索引。
* 步骤3：通过从预测的$k_2$个索引中获取原始键和值向量来执行稀疏注意力。

HSA与GQA的兼容性与算法伪代码：作者的HSA算法与GQA完全兼容。为了实现这一点，作者在每个注意力组的基础上执行所有键张量选择。更多细节显示在算法1中，其中作者根据需要对组维度中的$q$或$|q|$进行求和，以确保组内的所有注意力头在每一步做出相同的选择。

# 输入: 查询向量 q, 键张量 K, 值张量 V, 逐元素最大/最小键张量 K_max / K_min
# 从组维度中 |q| 的总和获取沿头部维度的 top-k1 索引
i_1 = argtopk(sum(|q|, dim=group), k_1)
# 从组维度中 q 的总和获取 top-k1 索引的符号
g = sign(sum(q[i_1], dim=group))
# 从分页的 min 或 max 中获取相应的索引
P = K_max[i: g_i >= 0], K_min[i: g_i < 0]
# 计算近似注意力分数
s_1 = score(q[i_1], P)
# 获取沿序列维度具有 top-k2 注意力分数的索引
i_2 = argtopk(s_1, k_2)
# 执行稀疏注意力
y = attn(q, K[i_2], V[i_2])
return y

多轮对话中永久驱逐的问题：在多轮对话设置中，由于KV token的重要性在不同轮次之间可能发生显著变化，因此在早期轮次中被修剪的KV tokens可能成为回答后续轮次查询的必要条件。因此，永久移除这些tokens可能会导致后续轮次的准确率显著下降。

RocketKV-MT的核心思想：为了缓解这个问题，作者引入了其方法的多轮变体，称为RocketKV-MT。RocketKV-MT的核心思想是避免在第一阶段永久驱逐任何KV tokens；相反，所有KV tokens都在各轮次之间保留在内存中，确保不会丢失任何可能有用的上下文。

多轮场景下的动态选择与全量保留机制：同时，为了保持计算效率，第二阶段仍然在第一阶段过滤的KV tokens子集上执行动态选择，类似于原始的RocketKV。换句话说，模型使用减少的KV token集合生成响应以提高速度，同时保留完整的KV历史记录以供未来轮次使用。例如，假设RocketKV-MT的第一阶段在第一轮中仅保留了输入提示中总共$M$个KV tokens中的$N$个。RocketKV-MT仍将所有$M$个tokens保留在内存中，但限制第二阶段在解码阶段仅从这$N$个输入tokens（加上任何新生成的tokens）中动态选择。在下一轮中，先前存储的完整KV tokens集（所有$M$个输入tokens加上所有输出tokens）被添加到新的输入KV缓存中。过滤过程（在作者的案例中使用SnapKV）再次应用于整个输入KV缓存，以为本轮的解码阶段选择一个新的重要tokens子集。

RocketKV-MT的效果与代价：通过遵循这一策略，RocketKV-MT在每一轮中实现了与RocketKV相似的解码加速，同时保留了所有轮次的完整KV缓存历史。这种方法有效地消除了多轮场景中由永久性KV token驱逐引起的准确率下降，代价是不节省内存存储。

解码阶段的内存瓶颈与Token预算定义：LLM推理的解码阶段通常受限于内存带宽；因此，在注意力模块中花费的时间大致与总内存流量成正比。在这项工作中，作者使用token预算$t$来估计解码阶段每个注意力操作的内存流量（作者主要关注KV缓存流量，因为它构成了此场景中大部分的内存流量）。例如，512的token预算意味着每个注意力模块需要从内存中获取相当于512个键值对的总量。与先前的研究中token预算仅反映top-$k$注意力的内存流量（HSA中的步骤3）不同，作者将token预算定义为也包含top-$k$估计的内存流量（HSA中的步骤2）。通过这样做，token预算可以更精确地反映注意力模块中的整体内存流量。为了简单起见，在后续实验中，作者将token预算在HSA的这两个步骤以及所有其他动态KV token选择方法之间平均分配。对于具有GQA的模型，此token预算是为整个注意力组而不是每个注意力头定义的。对于给定的序列长度$S$，总压缩率$c$可以定义为$c = S / t$。

自适应压缩分解机制的设计：由于RocketKV是一个两阶段KV缓存压缩框架，第一阶段过滤的KV token集作为第二阶段动态选择的输入，因此确定过滤后的KV token集的中间token预算非常重要。对于总体压缩率$c$，作者定义了一个分割因子$r$，使得$c$被分割为第一阶段的$c^r$和第二阶段的$c^{(1-r)}$，其中$0 <= r <= 1$。作者使用以下公式自适应地确定$r$：
$r = \min(0.2 + 0.06 * \log_2(c), 0.8)$

公式背后的原理与约束：该公式背后的见解是，当$c$较小时，作者希望最小化第一阶段永久驱逐的KV tokens数量以防止信息丢失。随着$c$的增加，第二阶段中HSA引起的准确率下降变大，因此最好为第一阶段分配更高的压缩率，因为SnapKV可以使用精确的注意力分数更精确地估计重要的KV tokens。作者进一步将$r$的范围限制在0.2到0.8之间，以平衡这两个阶段之间的压缩分解。

HSA内部的压缩分解：对于HSA内的压缩分解，作者简单地将压缩率$c^{(1-r)}$在序列和头部维度之间平均分配，以便每个维度获得$c^{(1-r)/2}$的压缩率。请注意，沿序列维度的压缩率等于页面大小，因此作者需要将其向上取整为最接近的整数$\lceil c^{(1-r)/2} \rceil$，另一个维度获得$c^{(1-r)} / \lceil c^{(1-r)/2} \rceil$。

分解策略的优势与示例：因为RocketKV战略性地跨多个阶段和维度分解KV缓存压缩，它显著增强了实现高压缩率的潜力，同时保持了强大的准确率——超越了依赖单阶段、单维度方法的方法。例如，给定$64\times$的压缩率，分割因子可以计算为$r = 0.2 + 0.06 * \log_2(64) = 0.56$。因此，压缩率被分割为第一阶段的$64^{0.56} \approx 10.3\times$和第二阶段的$64^{(1-0.56)} = 6.2\times$。HSA进一步将其压缩率分割为序列维度的$3\times$（页面大小为3）和头部维度的$2.1\times$。我们可以看到，分解后每个单独的阶段和维度都分配到了小得多的压缩率。

存储与流量开销的理论分析：在RocketKV中，第一阶段导致KV缓存存储和流量都减少了$c^r$。在第二阶段，作者需要考虑近似注意力引入的额外内存存储开销。由于作者在HSA的两个维度之间平均分配压缩率$c^{(1-r)}$为$c^{(1-r)/2}$（为简单起见此处忽略向上取整操作），它引入了 $(1/c^r) \times (1/c^{(1-r)/2}) \times 2 = 2/c^{(1+r)/2}$ 的内存存储开销，其中$1/c^r$是第一阶段后的相对KV缓存存储，并且逐元素最大值和最小值张量都在其之上引入了$1/c^{(1-r)/2}$的存储开销。因此，RocketKV中的总KV缓存存储和流量分别是完整KV基线的 $1/c^r + 2/c^{(1+r)/2}$ 和 $1/c$。RocketKV-MT在第一阶段不产生存储节省，因此其相对KV缓存存储改为 $1 + 2/c^{(1+r)/2}$。表1比较了在给定压缩率$c$下，RocketKV和RocketKV-MT与其他方法的KV缓存存储和流量。我们可以看到，虽然所有方法都带来相同的KV缓存流量节省，但只有RocketKV、DuoAttention和SnapKV提供额外的KV缓存存储节省，而RocketKV-MT、Quest和SparQ需要额外的存储空间用于辅助数据。

表1：RocketKV、RocketKV-MT与其他方法之间归一化的KV缓存存储（包括辅助数据）和流量比较。

Method	Compression Ratio	Storage	Traffic
Full-KV	1	1	1
DuoAttention	c	1/c	1/c
SnapKV	c	1/c	1/c
Quest	c	1 + 1/c	1/c
SparQ	c	2	1/c
RocketKV	c	$1/c^r + 2/c^{(1+r)/2}$	1/c
RocketKV-MT	c	$1 + 2/c^{(1+r)/2}$	1/c

系统兼容性与分离式服务部署：RocketKV与FlashAttention【Dao等人，Flashattention: Fast and memory-efficient exact attention with io-awareness + 2022 + NeurIPS】完全兼容，因为它不修改预填充阶段的注意力。此外，它与张量并行【Shoeybi等人，Megatron-LM: Training multibillion parameter language models using model parallelism + 2019 + arXiv】无缝集成，因为所有操作在注意力头/组之间是对称的。值得注意的是，RocketKV和RocketKV-MT都能很好地与分离式服务系统配合使用，其中不同的GPU用于预填充和解码阶段。对于RocketKV，预填充和解码GPU上的KV缓存存储都得到了减少，它们之间的KV缓存传输流量也减少了。虽然对于RocketKV-MT需要将完整的KV缓存存储在预填充GPU中，但只有过滤后的集合需要传输并存储在解码GPU中，从而产生与RocketKV相同的通信和解码优势。

实验环境

• 评估模型：Llama3.1-8B-Instruct (GQA架构，128K上下文)、Mistral-7B-Instruct-v0.2 (GQA架构，32K上下文)、LongChat-7B-v1.5-32k (MHA架构，32K上下文)。
• 数据集：LongBench（长上下文理解）、Needle-in-a-Haystack (NIAH, 检索任务)、RULER（上下文长度真实能力测试）、SCBench（多轮对话）。
• 硬件配置：NVIDIA H100 和 A100 GPUs，批处理大小设为1。
• 软件与代码环境：使用FP16精度，基于gpt-fast（PyTorch原生的低延迟LLM框架）实现。
• 对比基线：Full-KV（全量KV缓存）、Exact-TopK（理想Top-K预测）、DuoAttention、SnapKV、Quest、SparQ。

实验结果

1. LongBench 基准测试结果
* 实验内容：在三个模型上，评估不同token预算（256至4096）下的各项任务平均准确率。
* 实验结果与分析：如图5 (a-c) 所示，RocketKV在低token预算下始终优于所有其他压缩方法。对于Llama3.1-8B-Ins，RocketKV在512及以上的预算下几乎没有准确率损失，在256预算下平均准确率仅下降1.1%。SnapKV在低预算下表现尚可，但随着预算降低，与RocketKV的差距逐渐拉大。

2. Needle-in-a-Haystack (NIAH) 检索测试
* 实验内容：在不同序列长度和插入深度下测试“大海捞针”能力。
* 实验结果与分析：如图5 (d-f) 所示，即使在256的极低token预算下（对应Llama3.1上超过400倍的压缩率和109K的序列长度），RocketKV在所有模型上都达到了接近Full-KV的准确率（在Llama3.1-8B-Ins上达到100%）。相比之下，其他方法出现了大幅度的准确率下降。

3. RULER 基准测试
* 实验内容：在不同序列长度下评估各种方法的鲁棒性。
* 实验结果与分析：如图5 (g-r) 所示，RocketKV在短序列长度下准确率损失可忽略不计，随着序列长度增加准确率损失逐渐增大。即便如此，在更长的序列长度下，RocketKV与其他方法之间的准确率优势差距变得更宽。

4. SCBench 多轮对话基准测试
* 实验内容：在多轮对话设置下，评估RocketKV与RocketKV-MT的准确率表现（token预算从1024至16384）。
* 实验结果与分析：如图6所示，在多轮场景下，原始RocketKV在较高token预算（>=8192）时表现不如SparQ，且与Exact-TopK存在明显差距，这是因为早期轮次驱逐的KV tokens在后续轮次中可能变得重要。而专门为多轮设计的RocketKV-MT修复了这一问题，在所有token预算下都取得了与Exact-TopK相当的显著准确率提升。

5. 端到端推理效率与内存节省
* 实验内容：在NVIDIA A100和H100 GPU上，测试Llama3.1-8B-Ins解码阶段的端到端加速比和峰值内存节省百分比。
* 实验结果与分析：如图7所示，相比Full-KV，RocketKV在A100和H100上分别实现了高达3.7倍和3.3倍的端到端加速，以及高达32.6%的峰值内存节省。A100的加速比更高，是因为其内存带宽与计算能力的比值较低，LLM推理受内存限制更严重，从而从RocketKV的内存流量节省中获益更多。

结论

RocketKV 提出了一种新颖的、免训练的KV缓存压缩方法，有效解决了LLM推理解码阶段对内存带宽和容量的巨大需求挑战。通过将SnapKV的粗粒度永久性驱逐与混合稀疏注意力（HSA）的细粒度动态选择相结合，RocketKV在实现高达400倍压缩率的同时，保持了与完整KV缓存相当的准确率。此外，它在解码阶段实现了高达3.7倍的端到端加速和32.6%的峰值内存减少。针对多轮对话场景，本文进一步提出的RocketKV-MT变体完美解决了历史上下文丢失的问题，展现了该框架在优化LLM性能方面的巨大潜力和广泛应用前景。

补充细节与附录

A.1 模型处理细节：对于LongChat-7B-v1.5，先前的研究（Tang等人，2024；Chen等人，2025）通常在KV缓存压缩期间跳过前两层以保持准确率。作者发现这仅对LongChat-7B-v1.5是必要的；因此，作者对Llama3.1-8B-Ins和Mistral-7B-Ins-v0.2的所有注意力层进行KV缓存压缩，但跳过LongChat-7B-v1.5的前两层。

A.2 基准测试参数设置：
* LongBench：Llama3.1-8B-Ins的最大提示长度设置为127,500，其他两个模型设置为31,500。超出长度的提示通过保留前后一半的tokens进行中间截断。
* Needle-in-a-Haystack：对Llama3.1-8B-Ins使用10个不同的输入序列长度（从2048到81,920个单词，约转换为109K tokens）和10个不同深度进行评估。其他两个模型的长度范围从512到20480个单词。
* RULER：基本遵循原始配置，但为加快评估，将每个任务的示例数量从500减少到50。
* SCBench：遵循多轮设置，Llama3.1-8B-Ins的最大提示长度设置为127,500。

A.3 基线方法公平性设置：因为RocketKV主要侧重于加速LLM推理的解码阶段，为了公平比较，作者不在所有对比基线（包括DuoAttention）的预填充阶段执行任何KV缓存压缩或稀疏注意力机制。对于Quest和SparQ，作者将token预算在识别top-$k$索引的近似注意力和top-$k$稀疏注意力之间平均分配，以准确反映注意力模块所需的总内存获取带宽。原始的Quest方法不兼容GQA，作者对其进行了修改，使其根据每个注意力组而不是每个注意力头选择top-$k$累积注意力分数的索引。

B.1.1 混合稀疏注意力（HSA）的消融实验：
* HSA与Quest、SparQ的独立对比：为了说明混合稀疏注意力（HSA）的有效性，作者将独立的HSA机制的准确率与Quest和SparQ进行了比较。图8展示了Llama3.1-8B-Ins在多个不同基准测试中的结果。在所有情况下，HSA始终优于Quest和SparQ，尤其是在低token预算下。这清楚地证明了HSA的优势，与Quest和SparQ等单一维度近似方法相比，它智能地利用了序列和头部维度的近似。

B.1.2 分割因子（Split Factor）的消融实验：
* 自适应与静态分割因子的比较：在这项研究中，作者将自适应压缩分解方法与范围从0.3到0.7的静态确定分割因子$r$进行了比较。0.5的分割因子表示RocketKV的第一阶段和第二阶段之间平均分配。如图9所示，最佳静态分割因子随不同的序列长度和token预算而变化，而自适应压缩分解在大多数情况下提供了与最佳静态分割因子相当的准确率。

B.2 Needle-in-a-Haystack 可视化：
* 大海捞针任务热力图：Needle-in-a-haystack (NIAH) 是一种合成挑战，旨在测试LLM在大量文本中检索特定信息的效率（Kamradt，2023）。在图10、11、12中，x轴显示文档长度（即“草垛”），y轴标记“针”（即一个短句）在文本中的相对位置。结果表明，即使token预算低至256，RocketKV也能够在三个模型的所有文本中以几乎与Full-KV相同的准确率检索到针。