VL-Cache: Sparsity and Modality-Aware KV Cache Compression for Vision-Language Model Inference Acceleration

作者/机构: Dezhan Tu (University of California, Los Angeles), Danylo Vashchilenko (AWS AI, Amazon), Yuzhe Lu (AWS AI, Amazon), Panpan Xu (AWS AI, Amazon)

A1 主要贡献

本文提出了一种名为 VL-Cache 的新型 KV 缓存压缩方法,专为加速视觉语言模型(VLM)推理而设计。

  • 核心问题:随着视觉语言模型(VLM)处理长视觉上下文(如高分辨率图像、多图像和视频)能力的提升,视觉 token 的数量急剧增加。这导致存储和访问大型键值(KV)缓存成为一个重大挑战,不仅消耗大量 GPU 内存,还因数据移动增加了推理延迟。现有的为大型语言模型(LLM)设计的 KV 缓存压缩方法直接应用于 VLM 时,效果并不理想。

  • 研究目标:旨在通过一种针对 VLM 定制的 KV 缓存压缩方案,有效减少 KV 缓存大小,从而降低 GPU 内存占用和推理延迟,同时最大限度地减少对模型准确率的影响。

  • 创新点与主要贡献

    1. VLM 注意力稀疏性分析:首次深入研究了 VLM 在预填充(prefill)和解码(decoding)阶段独特的注意力稀疏性模式,揭示了其与 LLM 的显著差异。特别地,VLM 的注意力在视觉和文本 token 之间存在明显的模态边界,解码阶段的语言 token 主要关注提示中视觉 token 之后的语言 token(后视觉注意力),而非视觉 token 本身。


      图 1:LLaVA v1.6 Mistral 7B 的注意力分数矩阵,(a) 仅语言上下文,(b) 语言和视觉上下文。红色越深表示注意力分数越高。两个矩阵都表明解码阶段使用的关键 token 主要与预填充上下文中的 token 一致。关键区别在于,在 (b) VLM 的注意力中,沿查询维度出现了一个清晰的模态边界。

    2. 层级自适应的稀疏性感知缓存预算分配:提出了一种在推理时根据每层注意力稀疏性动态分配 KV 缓存预算的方法。该方法能够将有限的缓存预算有效地分配到不同层,避免了对信息密集层分配不足和对稀疏层分配浪费的问题,从而在不牺牲准确率的情况下进一步减小 KV 缓存大小。

    3. 模态感知的 Token 评分策略:开发了一种新的 token 评分策略,该策略利用“后视觉注意力”(即提示中视觉 token 之后的语言 token 所产生的注意力)来评估所有 token 的重要性。研究发现,这种策略能更准确地识别并保留对生成任务至关重要的视觉和语言 token。

  • 核心成果:实验结果表明,仅保留 10% 的 KV 缓存,VL-Cache 就能在多个基准数据集上达到与使用完整缓存相当的准确率。在速度测试中,该方法可将生成 100 个 token 的端到端延迟最多加速 2.33 倍,解码速度最多提升 7.08 倍,同时将 GPU 中的 KV 缓存内存占用减少 90%。

A3 背景知识

2.1 VLM 推理

预填充(Prefill)阶段。VLM 的输入包括图像和语言。图像由视觉编码器处理生成视觉 token,随后通过投影层映射到统一的嵌入空间。同时,语言提示被送入分词器和嵌入层以创建语言 token 的初始隐藏状态。为简化表示,我们将包含视觉和语言 token 在内的 m 个提示 token 序列表示为 {$x_1, ..., x_m$}。这些 token 由语言模型并行处理,以计算第一个解码 token 的概率 $P_\theta(x_{m+1}|x_1, ..., x_m)$。同时,每个 Transformer 层 l 的键向量 {$k_1^{(l)}, ..., k_m^{(l)}$} 和值向量 {$v_1^{(l)}, ..., v_m^{(l)}$} 被缓存在 GPU 内存中,以避免在后续解码步骤中重新计算。

解码(Decoding)阶段。解码开始后,VLM 中的语言模型以自回归循环的方式每步生成一个 token。在第 i 步,模型接收 token $x_{m+i}$ 作为输入,并计算概率 $P_\theta(x_{m+i+1}|x_1, ..., x_{m+i})$。每个解码步骤都涉及生成新的键向量 $k_{m+i}^{(l)}$ 和值向量 $v_{m+i}^{(l)}$,这些向量被附加到先前缓存的每层键值对中,形成 {$k_1^{(l)}, ..., k_{m+i}^{(l)}$} 和 {$v_1^{(l)}, ..., v_{m+i}^{(l)}$}。在长上下文(如多张或高分辨率图像)的情况下,键值缓存可能会变得比模型参数和其他中间张量大得多,使得内存容量和带宽成为主要的性能瓶颈。

2.2 KV 缓存压缩

设计维度。为了解决解码过程中存储和访问大型 KV 缓存的瓶颈,许多研究人员专注于 KV 缓存压缩,旨在仅维护完整 KV 缓存的一个子集,以实现更高效的解码,同时最小化准确率损失。此类算法有两个主要的设计维度:每个层应保留多少缓存 token,以及在压缩过程中应驱逐哪些 token。

预算分配。由于 Transformer 架构由多个相同的层组成,一个直接的策略是为每层分配相等的 KV 缓存槽预算【索引26, Xiao et al., Efficient streaming language models with attention sinks, 2023; 索引42, Zhang et al., H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2024; 索引13, He et al., Zipcache: Accurate and efficient kv cache quantization with salient token identification, 2024】。最近,受观察到在不同层移除缓存 token 会导致不同程度性能损失的启发,PyramidKV【索引41, Zhang et al., Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024】和 PyramidInfer【索引36, Yang et al., Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference, 2024】提出了一种衰减调度方案,为较浅的层分配更多的缓存预算,并观察到相对于每层预算相等的基线,准确率有所提高。

Token 评分策略。对于给定的缓存 token 预算,需要一个 token 评分策略 $\psi$ 来对 KV 缓存 token 的重要性进行排序,并决定保留哪些 token。设 n 为 l 层的缓存 token 数量,$S = \{0, 1, ..., n\}$ 为这些缓存 token 的索引。我们定义 $\psi : S \rightarrow R^n$,其输出是当前缓存 token 的分数。根据这些分数,选择得分最高的 k 个索引,$S_\psi := \{i_1, i_2, ..., i_k : \psi(S)_{i_j} \ge \psi(S)_{x \in [n]\setminus\{i_1,i_2,...,i_k\}}\}$,其中 $k \in [1, n)$。最近的研究表明,注意力分数是设计此类策略的有效来源。StreamingLLM【索引35, Xiao et al., Efficient streaming language models with attention sinks, 2023】发现高注意力分数经常分配给初始 token,并通过仅保留初始和最近的缓存 token 同时移除中间 token 来实现长度泛化。H2O【索引42, Zhang et al., H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2024】使用累积的注意力分数来识别要保留的关键 token。

基于稀疏性的统一指导。我们观察到,从根本上说,KV 缓存压缩方法之所以有效,是因为使用 Transformer 层的推理过程是稀疏的。因此,在这项工作中,我们利用注意力稀疏性作为统一的指导,为 VLM 中的 KV 缓存压缩设计缓存预算分配机制和 token 驱逐策略。

3.1 测量注意力稀疏性

稀疏性测量方法。在本节中,我们测量了预填充和解码阶段不同 Transformer 层中注意力分数矩阵的稀疏性。首先,我们对注意力分数矩阵 A 应用一个带有相对阈值 p 的过滤器:

$$\begin{aligned} \mathbf{T h r e s h o l d F i l t e r}(A, p)_{i j}= \begin{cases}A_{i j} & \text { if } A_{i j} \geq p \cdot \max _{j}\left(A_{i j}\right) \\ 0 & \text { otherwise }\end{cases} \end{aligned}$$

其中阈值 $p \in (0, 1)$ 控制了引入稀疏化的强度,遵循 Zhang 等人的工作【索引42, Zhang et al., H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2024】。我们凭经验设置 p = 1%,这样被过滤掉的分数对 Transformer 层的输出影响很小。过滤后,我们计算层 l 的稀疏度 $\gamma^{(l)} \in [0, 1]$,即零条目的数量,并用注意力分数矩阵的下三角部分的大小进行归一化:

$$\gamma^{(l)} := \frac{\sum_{i \ge j} \mathbb{1}[\text{ThresholdFilter}(A^{(l)}, p)_{ij} = 0]}{|\{A^{(l)}_{ij} : i \ge j\}|}$$

稀疏性观察。我们计算了评估数据集中样本的平均稀疏度,并在图 2 中绘制了每个层在不同注意力头下的稀疏度范围。我们观察到,在预填充阶段(图 2a),前两层的稀疏度显著较低,密度较高。此外,中间的一些层也比其相邻层表现出更高的密度。在解码阶段(图 2b),也观察到了类似的趋势。基于这些观察,我们预期预填充期间的聚合注意力稀疏性可以有效地预测稳健解码所需的 KV 缓存大小。


图 2:预填充和解码阶段的层级注意力稀疏性。不同层表现出不同程度的稀疏性;解码阶段的层级稀疏性趋势与预填充阶段相似。


图 3:平均缓存命中率。累积的后视觉注意力(我们的方法)比其他两种 token 评分策略表现出更高的缓存命中率。

对现有方法的反思。我们注意到,先前的方法,如 H2O【索引42, Zhang et al., H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2024】和 Keyformer【索引1, Adnan et al., Keyformer: Kv cache reduction through key tokens selection for efficient generative inference, 2024】,在所有层中分配相同的缓存大小,导致对信息量大的层分配不足,而对信息量少的层造成浪费。更新的工作包括 PyramidKV【索引41, Zhang et al., Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024】和 PyramidInfer【索引36, Yang et al., Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference, 2024】,它们随层深度的增加单调地减小 KV 缓存大小,我们发现这种做法也是次优的。我们的观察揭示了一种更微妙的稀疏性模式,其中不同层需要非单调变化的缓存大小来在解码期间保留上下文信息。

3.2 测量缓存命中率

后视觉注意力的假设。在第 1 节的图 1b 中,我们注意到 VLM 注意力在提示中的视觉和语言 token 之间存在明显的边界。基于这些观察,我们进一步假设,基于后视觉预填充注意力(post-vision prefill attention)而非全提示预填充注意力来保留 token,将能以更高的召回率保留重要的缓存 token。为了验证这一假设,我们定义了 CacheHitRate 来衡量在应用不同评分策略后,重要 token 在被驱逐后仍被保留的比例。设 m 为提示 token 的数量,$Q_{:m}, K_{:m} \in R^{m \times d}$ 为相应的查询和键矩阵,$Q_{m+1} \in R^d$ 为第一个解码 token 的查询向量。同时,回顾第 2.2 节,S 表示提示缓存 token 的索引,$\psi$ 表示一个将每个 token 索引映射为实数作为其重要性的评分策略。

CacheHitRate 定义。给定 $A_{m+1} := softmax(\frac{Q_{m+1}K_{:m}^T}{\sqrt{d}}) \in R^m$,我们定义:
- $\psi^* : S \rightarrow A_{m+1}$ 为最优评分函数,因为 $A_{m+1}$ 是解码中的真实注意力分数。
- $S_{\psi^*,k}$ 为根据 $\psi^*$ 排名选出的 top-k token,我们将其视为所有 k 值的基准真相。
- CacheHitRate 定义为 $\frac{|S_{\psi,k} \cap S_{\psi^*,k}|}{|S_{\psi^*,k}|}$,即被任意策略 $\psi$ 保留的 top-k token $S_{\psi,k}$ 中,也属于真实 top-k token $S_{\psi^*,k}$ 的百分比。

现有评分策略。我们注意到,在推理过程中,$A_{m+1}$ 是不可用的,因为我们的目标是在解码过程开始前压缩 KV 缓存,以解除内存瓶颈并提高解码吞吐量。因此,最近的研究提出了多种 $\psi$ 来近似 $\psi^*$。为简洁起见,我们假设以下 softmax 的输出都应用了必要的因果掩码作为注意力分数。
- 累积注意力(先前工作):使用沿查询维度的累积注意力分数,$\psi(S) := \sum_i softmax(\frac{Q_{:m}K_{:m}^T}{\sqrt{d}})_i$ 来对缓存 token 的重要性进行排序。
- 归一化注意力(先前工作):该方法计算沿查询维度的平均注意力分数,基于的观察是累积分数不成比例地偏向于较早的 token。该策略可写作 $\psi(S) := \frac{1}{n} \odot \sum_i softmax(\frac{Q_{:m}K_{:m}^T}{\sqrt{d}})_i$,其中向量 n 表示每列中未被掩码的元素数量,$\odot$ 表示逐元素乘法。
- 滑动窗口注意力(先前工作):同样计算沿查询维度的累积注意力分数,但只在一个最近的窗口内。该策略定义为 $\psi(S) := \sum softmax(\frac{Q_{m-w:m}K_{m-w:m}^T}{\sqrt{d}})_i$,其中 w 是一个固定的窗口大小。
- 后视觉注意力(我们的方法):该方法基于图 1b 中观察到的后视觉提示 token 的注意力与解码 token 的注意力之间的相似性。我们使用术语“后视觉”来区分跟随视觉 token 的语言 token 与在提示中位于图像之前的指令。形式上,我们使用注意力分数的一部分 $Q_{m-\tau:m}$,其中 $\tau$ 是提示中跟随视觉 token 的语言 token 的数量。我们的策略由 $\psi(S) := \sum_i softmax(\frac{Q_{m-\tau:m}K_{m-\tau:m}^T}{\sqrt{d}})_i$ 定义。在预填充期间,我们的策略可以解释为一个动态的、依赖于提示的滑动窗口,其窗口大小设置为后视觉语言提示的长度,而不是一个静态的、与提示无关的值。从图 3 中,我们发现我们基于后视觉注意力的评分策略在各层中始终比其他策略带来更高的缓存命中率。

A2 方法细节

4. VL-CACHE 方法

方法概述。受初步实验观察的启发,我们引入了 VL-Cache 方法,它策略性地结合了稀疏性感知的缓存预算分配和模态感知的 token 评分策略,以在有限的 KV 缓存预算下提升 VLM 在准确率和效率方面的性能。具体来说,我们使用后视觉注意力(Post-vision Attention)来计算层间稀疏度和层内 token 的重要性。前者指导每层应分配多少缓存 token,而后者决定了在一层内应保留哪些 k 个 token。我们方法的高层描述如图 4 所示。为简洁起见,本节中我们将用 A' 表示后视觉注意力矩阵。


图 4:VL-Cache 概览。在预填充阶段,利用后视觉注意力矩阵计算由稀疏性驱动的层级预算并对缓存 token 的重要性进行排序。然后驱逐不重要的 token,从而以减少的 KV 缓存实现更低的内存使用和加速的解码。

后视觉注意力的优势。我们认为,使用后视觉注意力有两个重要优势。首先,与使用完整注意力矩阵 $O(m^2)$ 的方法【索引42, Zhang et al., H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2024; 索引13, He et al., Zipcache: Accurate and efficient kv cache quantization with salient token identification, 2024】相比,使用后视觉注意力计算层级稀疏性的内存和延迟开销要低得多,为 $O(\tau m)$,因为在当前的视觉语言任务中,视觉 token 占据了提示长度的主导地位,$\tau \ll m$。其次,在固定的缓存预算下,使用后视觉注意力能更好地保留重要 token,这由图 3 中的高 CacheHitRate 所衡量。

4.1 稀疏性感知的 KV 缓存预算分配

预算分配算法。在确定要驱逐的具体 token 之前,我们需要分配 KV 缓存预算,即每层要保留的 KV 缓存的百分比。基于我们从第 3.1 节的观察,我们在预填充阶段实现了一个分为两步的稀疏性感知层级 KV 缓存分配方法,详见算法 1。首先,我们对后视觉注意力分数应用 ThresholdFilter 剪枝(p = 1%)并计算层级稀疏度(第 4-8 行和 13-18 行)。其次,给定整个模型的目标 KV 缓存预算,我们根据每层的稀疏度比率将此预算分配到各层(第 9-12 行)。这种方法优化了有限内存的使用,以在每层存储适量的上下文信息。由于分配仅在预填充后发生一次,其延迟开销被分摊到多个解码步骤中。

算法 1:稀疏性感知的缓存预算分配

1: 输入: 查询和键 Q, K ∈ R^(L×H×m×d), 层数 L, 头数 H, 后视觉提示长度 τ, 整体缓存预算 α.
2: procedure SKEWEDCACHEBUDGETALLOCATION(Q, K, α, L, H)     ▷ 计算层级缓存预算的主方法.
3:   Γ[L][H] ← 0                                           ▷ Γ 存储层和头的稀疏度
4:   for l = 1 → L do
5:     for h = 1 → H do
6:       Γ(l)h ← ComputePostVisionSparsity(Q(l)h, K(l)h)     ▷ 调用下面的辅助方法
7:   γ ← Γ.mean(1)                                         ▷ γ 是每层的头平均稀疏度
8:
9:   Z ← Σ_l(1 − γ(l))                                     ▷ Z 是所有层非稀疏比率之和,作为归一化因子
10:  β[L] ← 0                                             ▷ β 存储每层的缓存 token 预算
11:  for l = 1 → L do
12:    β(l) ← clip((1.0 − γ(l)) / Z * αL, 0.01, 1)        ▷ 按归一化的非稀疏比率分配预算
13:
14: procedure COMPUTEPOSTVISIONSPARSITY(Q, K)
15:   ▷ 计算后视觉注意力头稀疏度的辅助方法.
16:   Q' ← Q_{m−τ:m}                                         ▷ Q' 是后视觉查询
17:   A' ← softmax(Q'K^T / d)                               ▷ A' 是后视觉注意力
18:   γ' ← Σ_{i+m−τ ≥ j} 1[ThresholdFilter(A', p)_{ij}=0] / |{A'_{ij} : i+m−τ ≥ j}| ▷ γ' 是未掩码条目中零的比例
19:   return γ'
20:
21: 输出: 层级预算 β

与 PyramidKV 的关键区别。在算法 1 中,我们展示了我们的 KV 缓存预算分配算法,其中 α 是一个超参数,编码了整个模型期望的 KV 缓存预算。当准确率不理想时,可以使用更高的 α 来保留更多的 KV 缓存。最后,我们想指出我们的算法与 PyramidKV【索引41, Zhang et al., Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024】的关键区别:我们的分配预算是根据每个提示的稀疏性模式为其量身定制的,而不是对所有提示使用固定的层级预算,从而提供了额外的灵活性。

4.2 模态感知的 TOKEN 评分策略

策略选择。在我们确定了层 l 的缓存预算后,我们从完整的缓存中选择一个包含 $k^{(l)}$ 个缓存 token 的子集。先前的工作(主要针对 LLM 的 KV 缓存压缩)已经探索了多种评分策略,如 H2O 中的累积注意力【索引42, Zhang et al., H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2024】、ZipCache 中的归一化注意力【索引13, He et al., Zipcache: Accurate and efficient kv cache quantization with salient token identification, 2024】以及其他工作中的滑动窗口注意力【索引21, Li et al., Snapkv: Llm knows what you are looking for before generation, 2024; 索引41, Zhang et al., Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024; 索引36, Yang et al., Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference, 2024】,正如我们在 3.2 节中讨论的。本节我们将深入解释为什么后视觉注意力为 VLM 提供了比以往策略更好的评分策略。

后视觉注意力的优越性。虽然累积注意力作为一个简单而有效的基线,但一个关键问题是,对整个查询维度进行求和不可避免地会给早期的 token 分配高分。观察到这种长度偏差,滑动窗口注意力只对最近的窗口内的注意力分数求和。我们确认这确实导致了更高的 CacheHitRate,如图 3 所示。然而,正如我们在图 1a 和 1b 中所见,VLM 的注意力模式与 LLM 有显著不同。沿着查询维度,我们观察到一个清晰的模态边界:视觉 token 对其他视觉 token 的注意力相当均匀,而语言 token 只高度集中地关注少数几个视觉 token。如果我们应用累积或归一化注意力作为评分策略,当我们对所有前面的视觉 token 的分数进行求和或平均时,关键缓存 token 的信号将被掩盖。使用最近的窗口有所帮助,但固定的窗口大小很容易变得过大或过小,因为后视觉提示的长度因情况而异。因此,我们通过实现一个动态的、特定于提示的窗口大小,引入后视觉注意力作为 VLM 的优化评分策略。我们评分策略的优势在其图 3 中的高 CacheHitRate 中显而易见,并将在第 5 节的准确率比较中进一步得到证明。

A4 实验环境

  • 模型架构:

    • LLaVA-Mistral-7B【索引23, Liu et al., Visual instruction tuning, 2023】
    • LLaVA-1.6-34B【索引22, Liu et al., Improved baselines with visual instruction tuning, 2023】
    • 共享相同的视觉模型:openai/clip-vit-large-patch14-336【索引29, Radford et al., Learning transferable visual models from natural language supervision, 2021】
    • 基于不同的语言模型主干:Mistral【索引17, Jiang et al., Mistral 7b, 2023】和 Nous-Hermes-2-Yi-34B【索引30, Research, Nous hermes 2 on yi 34b, 2023】。
    • LLaVA-Mistral-7B 使用分组查询注意力(GQA),LLaVA-1.6-34B 使用多头注意力(MHA)。
  • 硬件配置:

    • AWS EC2 P4 实例,配备 8 块 A100 40GB GPU。
  • 数据集与任务:

    • 从 lmms-eval【索引6, Bo Li* & Liu, Lmms-eval: Accelerating the development of large multimoal models, 2024】中采样了三个任务,涵盖 OCR、推理和世界知识领域:

      • Coco-Caption【索引10, Chen et al., Microsoft coco captions: Data collection and evaluation server, 2015】
      • DocVQA【索引28, Mathew et al., Docvqa: A dataset for vqa on document images, 2021】
      • MathVista【索引26, Lu et al., Mathvista: Evaluating mathematical reasoning of foundation models in visual contexts, 2023】
    • 附录中还包括 TextVQA【索引31, Singh et al., Towards vqa models that can read, 2019】和 ChartQA【索引27, Masry et al., Chartqa: A benchmark for question answering about charts with visual and logical reasoning, 2022】。

  • 软件与基线配置:

    • 基线方法:

      • StreamingLLM【索引34, Xiao et al., Efficient streaming language models with attention sinks, 2023】
      • H2O【索引42, Zhang et al., H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2024】
      • ZipCache【索引13, He et al., Zipcache: Accurate and efficient kv cache quantization with salient token identification, 2024】
      • PyramidKV【索引41, Zhang et al., Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024】
      • FastV【索引8, Chen et al., An image is worth 1/2 tokens after layer 2: Plug-and-play inference acceleration for large visionlanguage models, 2025】(附录)
      • HiRED【索引4, Arif et al., Hired: Attention-guided token dropping for efficient inference of high-resolution vision-language models in resource-constrained environments, 2024】(附录)
    • 实现细节:

      • 所有基线均使用默认设置,但 KV 缓存预算与提示长度成比例缩放,最近 token 窗口大小固定为该预算的 10%,以进行公平比较。
      • 速度基准测试使用 HuggingFace 实现的默认设置,包括 CUDA-based FlashAttention-v2。
      • VL-Cache 性能优化使用了基于 Triton 的自注意力前向传播、层级稀疏度评估和模态感知 token 评分的解决方案。

A4 实验结果

5.1 准确率评估

实验结果如图 5 和表 1 所示,展示了在 KV 缓存预算从 1% 到 100% 变化时的平均准确率得分。总体而言,VL-Cache 在不同 KV 缓存预算和语言模型主干上均优于其他基线方法。

  • Coco-Caption 数据集:

    • 结果: 当 KV 缓存预算超过 40% 时,所有基线都能保持高准确率。但预算进一步减少时,准确率显著下降。值得注意的是,VL-Cache 即使在只有 5% 到 10% 的 KV 缓存预算下,也能达到与完整 KV 缓存相当的准确率(见图 5 (a), (d) 和表 1)。
    • 分析: H2O 和 StreamingLLM 对所有层使用相同的缓存预算,导致信息密集层丢失重要 token,而稀疏层则包含冗余 token。PyramidKV 静态地以单调递减方式分配缓存,对所有输入查询效果不佳。
  • DocVQA 数据集:

    • 结果: PyramidKV 表现出色,但 VL-Cache 始终优于所有其他基线。H2O 和 StreamingLLM 的鲁棒性较差,随着 KV 缓存的进一步压缩,性能显著下降(见图 5 (b), (e) 和表 1)。
  • MathVista 数据集:

    • 结果: VL-Cache 在大多数情况下再次优于其他方法。同时,多种方法在 1% 的 KV 缓存下实现了接近完整缓存的性能(见图 5 (c), (f) 和表 1)。
    • 分析: 这可能是因为该数据集中 54% 是选择题,挑战性低于自由形式的问题。


图 5:在不同数据集上使用不同缓存预算的评估结果。VL-Cache 在有限的 KV 缓存预算下,实现了与完整缓存相当的准确率,并优于多个基线。有趣的是,我们发现 VLM 在使用部分 KV 缓存时性能偶尔会略好,我们将其归因于 KV 缓存压缩的正则化效应。


表 1:在 10% KV 缓存下,VL-Cache 在大多数数据集和模型上达到了接近 100% 的完整缓存性能,并显著优于其他方法。由于所有 KV 缓存压缩方法在缓存预算增加时都会收敛到完整缓存性能,因此在低缓存预算下的强劲表现证明了 VL-Cache 在保留重要 KV token 方面的有效性。

5.2 速度基准测试

为了展示 VL-Cache 的速度优势,我们使用合成提示测量了预填充和解码前向传播的 GPU 内核延迟。

  • 实验设置:

    • 提示大小从 1K 到 128K token 不等。
    • 批处理大小从 1 到 64 不等,且为静态批处理。
    • 假设提示模板格式与准确率基准测试相似,使用提示的最后 50 个 token 来决定驱逐哪些 KV 缓存 token。
    • 将 KV 缓存压缩至 10% 进行此基准测试。
  • 结果与分析:

    • 预填充开销: VL-Cache 方法的开销仅为预填充延迟的 1-4%,非常小(见表 2)。
    • 解码加速: 由于 KV 缓存大小减少,解码速度最多可提升 7.08 倍。
    • 端到端加速: 端到端加速受预填充延迟的限制。例如,当提示长度为 128K 且批大小为 1 时,7.08 倍的解码加速被占端到端延迟 53% 的预填充所稀释,最终导致 1.66 倍的端到端加速。随着输出 token 数量的增加,端到端加速会单调接近解码加速。这表明 KV 缓存稀疏性在长输出任务(如图像描述、视频描述、多模态思维链推理等)中将提供最大的速度优势。
    • 吞吐量与延迟: 如图 6 所示,VL-Cache 在任何期望的服务器级吞吐量下,都能提供更高的峰值吞吐量和更低的延迟。

表 2:生成 100 个输出 token 时,按批大小和提示长度划分的性能指标。


图 6:服务器级吞吐量 vs. 请求级延迟曲线(提示长度=2K)。标记点表示批大小。

A5 结论

本文提出了 VL-Cache,一种专为 VLM 优化的新型 KV 缓存压缩方法。我们揭示了在预填充和解码阶段视觉和语言 token 独特的稀疏性模式。基于这些观察,我们引入了模态感知的 token 评分策略和稀疏性感知的缓存预算分配,以在不损失准确率的情况下减小 KV 缓存大小。在多个基准数据集上的实验结果表明,仅保留 10% 的 KV 缓存,我们的方法就能达到与完整 KV 缓存相当的准确率,并优于所有现有方法。在速度基准测试中,我们的方法将生成 100 个 token 的端到端延迟最多加速 2.33 倍。

未来工作:

  1. 解码阶段压缩: 目前的方法主要关注压缩预填充阶段的 KV 缓存。对于长输出任务,可以开发对解码生成的 token 进行周期性压缩的方法,以在延迟开销和内存节省之间取得平衡。
  2. 扩展到视频模型: 将该方法扩展到视频模型会很有趣,因为视频输入会因 KV 缓存而带来更大的内存压力。

A6 附录

A.1 视觉-语言提示模板构建

构建提示模板。在基于图像的对话中构建提示模板是 VLM 的常见做法【索引20, Li et al., Llava-onevision: Easy visual task transfer, 2024; 索引33, Team et al., Gemini: a family of highly capable multimodal models, 2023; 索引16, Islam & Moushi, Gpt-4o: The cutting-edge advancement in multimodal llm, 2024; 索引5, Bai et al., Qwen-vl: A versatile vision-language model for understanding, localization, text reading, and beyond, 2023; 索引3, Anthropic, Introducing the next generation of claude, 2023】,因为它能指导语言模型生成更准确的响应。例如,如图 7 所示,输入图像通过视觉编码器和投影层处理,生成由<image>表示的图像嵌入。对于语言输入,除了用户输入外,还使用了一个提示模板。通过适当的提示模板设计,无论用户输入的原始图像顺序如何(在语言输入之前或之后),总会有一个基于语言的指令或问题位于后视觉位置,为我们的 VL-Cache 提供了驱逐不重要视觉 token 的强信号。


图 7:视觉-语言提示模板示例。

A.2 关于阈值 p 的消融研究

对超参数p的鲁棒性分析。在 VL-Cache 中,我们利用公式 1 中的 ThresholdFilter 来稀疏化注意力矩阵。它使用一个相对阈值 p,我们在实验中凭经验设置为 0.01。为了解我们的方法对这个超参数的鲁棒性,我们使用 LLaVA-V1.6-Mistral-7B 和 Coco-Caption 进行了消融研究。我们将 p 从 0.0001 变化到 0.1,并在下图中展示了相应的性能。我们观察到:1) 总体而言,VL-Cache 对 p 的选择是鲁棒的,因为其性能变化相对较小;2) 更高的 token 预算似乎对 p 更具鲁棒性;3) 没有一个单一的阈值 p 在所有缓存预算下都表现得始终更好,但 p = 0.01 似乎是一个不错的选择。


图 8:LLaVA-V1.6-Mistral-7B 在 Coco-Caption 上使用 VL-Cache 在不同缓存预算和阈值下的性能。

A.3 完整的实验结果

表 3:不同压缩方法和数据集下 VLM 的性能。在 10% KV 缓存下,VL-Cache 在大多数数据集和模型上接近 100% 的完整缓存性能,并显著优于其他方法。


图 9:在不同数据集上使用不同缓存预算的评估结果。VL-Cache 在有限的 KV 缓存预算下,实现了与完整缓存相当的准确率,并优于多个基线。

A.4 扩展的相关工作

KV缓存压缩技术。KV 缓存虽然对基于 Transformer 的 LLM 和 VLM 家族至关重要,但需要大量的计算资源和内存,限制了推理速度。为应对这些挑战,研究人员探索了各种 KV 缓存压缩技术,如 KV 缓存稀疏化、量化或两者的结合。

KV 缓存稀疏化。Heavy-Hitters (H2O)【索引42, Zhang et al., H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2024】采用累积注意力分数来贪婪地驱逐不重要的 token。然而,这种方法倾向于在初始 token 上累积更多注意力,引入偏差并对解码过程中关键 token 的识别产生负面影响。ZipCache【索引13, He et al., Zipcache: Accurate and efficient kv cache quantization with salient token identification, 2024】进一步对累积注意力分数进行归一化,从而实现更精确的预测。Keyformer【索引1, Adnan et al., Keyformer: Kv cache reduction through key tokens selection for efficient generative inference, 2024】提出了一个新颖的分数函数来预测每个 token 的重要性,并只在 KV 缓存中保留窗口 token 和动态关键 token。PyramidKV【索引36, Yang et al., Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference, 2024】随着层数加深,分配逐渐减小的 KV 缓存大小。这种方法可以在特定任务上实现近 20.5 的绝对准确率提升。


图 9:在不同数据集和不同缓存预算下的额外评估结果,展示了 VL-Cache 在广泛设置下的一致性能。

KV 缓存量化。该方法通过利用键值对的低比特表示来减小 KV 缓存的大小。特别是,KV 缓存的 FP8 和 INT8 量化常用于 LLM 推理。将 KV 缓存进一步压缩到更低比特也正在被探索,例如 KIVI【索引25, Liu et al., Kivi: A tuning-free asymmetric 2bit quantization for kv cache, 2024】将 KV 缓存压缩到 2 比特表示。

量化与稀疏化的结合。量化和稀疏化【索引37, Yang et al., No token left behind: Reliable kv cache compression via importance-aware mixed precision quantization, 2024】在技术上可以有正交的实现,可以同时启用以获得进一步的提升。然而,由于两者都是有损方法,可能影响下游任务的准确率,因此需要对输出质量进行仔细的准确率评估。

视觉 Token 缩减。在 VLM 领域,大多数工作都集中在剪枝嵌入的图像特征,无论是在它们进入 LLM 之前还是在 LLM 的隐藏状态内。HiRED【索引4, Arif et al., Hired: Attention-guided token dropping for efficient inference of high-resolution vision-language models in resource-constrained environments, 2024】利用视觉 Transformer(如 CLIP-ViT)中的稀疏注意力模式来丢弃视觉 token,有效地从图像中提取关键对象,并减少后续 LLM 生成所需的计算资源。然而,它忽略了文本提示,可能意外地丢弃了必要的 token。另一方面,FastV【索引8, Chen et al., An image is worth 1/2 tokens after layer 2: Plug-and-play inference acceleration for large visionlanguage models, 2025】在 LLM 的特定层剪枝隐藏状态,显著加速了后续层的计算。然而,它对所有后续层应用统一的 token 预算,导致在 token 预算高度受限时准确率不佳。本文提出的 VL-Cache 分析并利用 VLM 中独特的稀疏性模式,为 VLM 推理带来了更好的准确率。

A.5 高效实现

高效统计计算。为了在预填充期间丢弃注意力较低的 token,我们需要从注意力分数矩阵中计算 2 个统计量:(1)键维度中每个 token 的平均注意力分数;(2)小于查询 Q 维度内最大值 p% 的注意力分数数量。我们算法中采用的 ThresholdFilter 在渐近性上比截断分布的其他方法(如 top-p 和 top-k)更快,因为它不需要对注意力分数进行排序。

自定义内核的需求。常规的注意力内核,如 FlashAttention【索引11, Dao, Flashattention-2: Faster attention with better parallelism and work partitioning, 2023】和 PagedAttention【索引18, Kwon et al., Efficient memory management for large language model serving with pagedattention, 2023】,不会在 HBM 中物化注意力分数,因此我们需要一个新的内存高效内核来计算注意力统计量。总体目标是调度这些操作以确保:(1)注意力掩码不写入 HBM;(2)QK 乘积不写入 HBM;(3)注意力分数不写入 HBM;(4)类似并行度的操作(如 div 和 add)被融合以避免在 HBM 中产生中间张量。

实现方案。初次尝试使用带 Triton 后端的 TorchCompiler,满足了要求(1)、(3)和(4),但未满足(2)。FlashAttention 论文中描述了 matmul 操作与 softmax 的高效融合,并可在 Triton 中实现,但 TorchCompiler 无法应用此优化。我们现在描述满足所有性能要求的最优算法。其关键在于沿 Q 维度划分工作,将 Q 张量的单个瓦片(tile)保留在 SRAM 中,同时连续加载另一个张量的瓦片并计算归约,而无需将 QK 乘积写回 HBM。

三核解决方案。最终,优化的计算由 3 个串行执行的内核组成:
1. Softmax + 行式统计(沿 K 维度归约)。读取 Q 和 K,计算并写入 HBM:(a) 指数化前的 K 维度最大值(用于数值稳定性),(b) 指数化后的 K 维度总和(用作归一化因子)。
2. Softmax + 列式统计(沿 Q 维度归约)。读取 Q、K 以及第一个内核的输出。计算并写入 HBM:(a) softmax 后的 Q 维度总和,(b) 低于阈值的注意力分数计数。
3. 部分计数的求和归约。对第二个内核输出的部分计数值进行求和。

我们在 Triton 中实现了这些内核,并观察到每个内核的延迟与 FlashAttention 内核大致相同,这意味着这种划分是最佳的。然而,即使采用这种方法,我们仍希望在预填充期间仅使用几十个 token 来查询统计信息。

A.6 速度基准测试结果

详细指标。本节我们展示速度基准测试的详细指标,包括预填充和解码阶段的精确时间,并概述了基准测试方法。我们使用 Torch Profiler 测量 GPU 内核延迟,并将其汇总以计算每个操作的总延迟。

表 4:生成 100 个输出 token 的详细性能指标(按批大小和提示长度)

预填充开销。下表展示了与从 50 个查询 token 计算注意力统计数据、为所有提示 token 计算驱逐分数,以及将驱逐 90% KV 缓存后的 KV 缓存张量复制到连续内存相关的所有 GPU 内核的累积延迟。此预填充开销以毫秒为单位:

预填充加速比。我们将预填充延迟定义为生成 1 个输出 token 的延迟。下表展示了预填充加速比 (prefill+overhead)。我们观察到开销从未超过 6%,并且随着提示变长而变小:

解码加速比。我们将解码延迟计算为端到端延迟与预填充延迟之差。下表展示了解码加速比 (完整 KV 缓存的解码延迟 / 10% KV 缓存的解码延迟)

端到端加速比。端到端加速比计算为 基线延迟 / VLCache 延迟。VLCache 延迟包括预填充延迟、VLCache 开销和 99 次解码过程,共生成 100 个输出 token:

A.7 测量对视觉和语言 Token 的注意力

进一步分析。正如我们在 3.1 节中确定的,解码期间的注意力稀疏性可以从预填充阶段的注意力分数中预测。为了进一步理解提示中视觉 token 和语言 token 的重要性,我们检查了解码 token 对这两种模态提示 token 的注意力分配情况。

Contribution 指标。我们提出 Contribution 用于对层级模态特定注意力模式进行定量分析和可视化。设 T 为当前序列长度,t 为第一个解码 token 的索引,$A^{(l)} \in R^{T \times T}$ 为特定头在 l 层的注意力矩阵。

$$\mathbf{Contribution}_{mod}^{(l)} := \frac{1}{T - t + 1} \sum_{i \ge t} \frac{\sum_{j \in J_{mod}} \mathbf{ThresholdFilter}(A^{(l)}, p)_{ij}}{\sum_{j \in J_{all}} \mathbf{ThresholdFilter}(A^{(l)}, p)_{ij}}$$


这里,$J_{mod}$ 和 $J_{all}$ 分别表示输入提示中选定的一种模态(语言或视觉)和所有模态的 KV 缓存索引。如图 10 所示,VLM 在第一层将主要注意力分配给视觉 token,而在第二层,视觉和语言 token 的贡献相当。从第三层开始,语言 token 占主导地位,中间层的视觉 token 贡献略有增加。

Coverage 指标。此外,我们还提出了 Coverage 来分析特定模态的 token 数量比例,定义如下:

$$\mathbf{Coverage}_{mod}^{(l)} := \frac{1}{T - t + 1} \sum_{i \ge t} \frac{\sum_{j \in J_{mod}} \mathbf{TopK}(A^{(l)}, k)_{ij}}{\sum_{j \in J_{all}} \mathbf{TopK}(A^{(l)}, k)_{ij}}$$


其中 TopK 定义为:

$$\mathbf{TopK}(A, k)_{ij} = \mathbb{1} [A_{ij} \in \{ A_i^{(r-k+1)}, A_i^{(r-k+2)}, ..., A_i^{(r)} \}]$$
这里,我们定义 $k$ 为 $\lfloor\alpha \cdot T\rfloor$,$r$ 为矩阵 A 的列数,并重载符号以使用 $A_i^{(n)}$作为行向量的第 n 个顺序统计量。此外,$\alpha \in (0, 1)$ 表示 token 预算阈值。在用 TopK 选择进行 KV 缓存压缩后,我们评估了剩余 token 中视觉注意力与语言注意力的比例。在图 11 中,当应用 $\alpha = 10\%$ 时,观察到与图 10 相比,覆盖率有类似的趋势。值得注意的是,在中间层,视觉 token 构成的比例大于它们的贡献,这表明每个 token 的贡献较小。


图 10:解码中的平均贡献


图 11:解码中的平均覆盖率