SqueezeAttention: 2D Management of KVCache in LLM Inference via Layer-Wise Optimal Budget

文章标题：SQUEEZEATTENTION: 通过逐层最优预算实现LLM推理中KV缓存的二维管理
作者/机构：Zihao Wang, Bin Cui, Shaoduo Gan (北京大学信息科学技术学院 & 高可信软件技术教育部重点实验室), Geoming AI

A1 主要贡献

核心问题：大型语言模型（LLM）的推理成本高昂，其中一个主要瓶颈是Key-Value（KV）缓存。现有的KV缓存压缩算法大多通过稀疏化token序列来减少缓存大小，但这些方法通常平等对待所有注意力层，为每一层分配相同的KV预算。这种策略是次优的，因为某些层对输入token的敏感度较低，却获得了与其他层相同的预算。

研究目标：本文旨在解决现有KV缓存优化方法中忽略层级差异性的问题，提出了两个核心问题：1) 所有共享相同KV缓存策略的注意力层是否都必须缓存相同数量的token？2) 如果不是，如何为每一层精确分配缓存预算，以便在序列级压缩的基础上进一步减少KV缓存？

创新点：
本文提出了一种名为SQUEEZEATTENTION的二维KV缓存压缩算法，它从序列和层两个维度联合优化KV缓存。
1. 提出层级重要性的量化指标：通过计算自注意力层前后隐藏表示的余弦相似度，来量化每个注意力层的重要性。相似度越高，表示该层对嵌入（embedding）的贡献越小。
2. 动态的层级预算分配机制：SQUEEZEATTENTION在推理的预填充（prefilling）阶段，动态地测量每个层的重要性，并通过聚类将所有层分组。然后，它重新分配初始的统一缓存预算，为重要的层分配更多预算以保证模型准确性，为不重要的层分配更少预算以节省I/O成本。
3. 正交性与兼容性：该算法与现有的序列级KV缓存压缩算法（如H2O、滑动窗口注意力等）是正交的，可以平滑地与它们结合使用，在现有优化的基础上实现进一步的性能提升。
4. 显著的性能提升：实验证明，SQUEEZEATTENTION能够在多种LLM和基准测试中实现约30%至70%的内存节省和高达2.2倍的吞吐量提升。

下图展示了不同KV缓存策略的对比，说明了SQUEEZEATTENTION如何在序列级压缩的基础上，通过层级预算的自适应再分配进一步压缩KV缓存。

A3 背景知识与关键观察

2.1 LLM推理中KV缓存的剖析

推理的两个阶段。对于仅解码器的transformer模型，推理过程通常包括预填充（prefilling）和解码（decoding）两个阶段。在预填充阶段，LLM将整个提示（prompt）作为输入，计算并缓存每个注意力层中每个token的键值（key-value）嵌入。随后，解码阶段一次处理一个嵌入以迭代生成token，同时将新计算的KV嵌入连接到KV缓存中。

预填充阶段的KV缓存计算。设$p$为提示长度，$o$为输出长度，$b$为批量大小，$n_{layer}$为注意力层总数，$d_{model}$为隐藏维度。在第$i$层，注意力相关的Key和Value权重记为$w_K^i$和$w_V^i$，其中$w_K^i \in \mathbb{R}^{d_{model} \times d_{model}}$，$w_V^i \in \mathbb{R}^{d_{model} \times d_{model}}$。在预填充阶段，第$i$层的隐藏状态记为$h_{prompt}^i$，其中$h_{prompt}^i \in \mathbb{R}^{b \times p \times d_{model}}$。预填充后第$i$层的KV缓存可以表示为：

$$\mathbf{C}_{K}^{i}=\mathbf{h}_{p r o m p t}^{i} \cdot \mathbf{w}_{K}^{i} ; \mathbf{C}_{V}^{i}=\mathbf{h}_{p r o m p t}^{i} \cdot \mathbf{w}_{V}^{i}$$

解码阶段的KV缓存计算。在解码阶段，第$i$层第$j$个输出token的隐藏状态记为$h_{output_j}^i$（$1 \le j \le o$），其中$h_{output_j}^i \in \mathbb{R}^{b \times 1 \times d_{model}}$。生成第$j$个token后，第$i$层的KV缓存可以表示为：

$$\begin{aligned} \begin{aligned} \mathbf{C}_{K}^{i} & =\operatorname{CONCAT}\left(\mathbf{C}_{K}^{i}, \mathbf{h}_{output_{j}}^{i} \cdot \mathbf{w}_{K}^{i}\right) \\ \mathbf{C}_{V}^{i} & =\operatorname{CONCAT}\left(\mathbf{C}_{V}^{i}, \mathbf{h}_{output_{j}}^{i} \cdot \mathbf{w}_{V}^{i}\right) \end{aligned} \end{aligned}$$

KV缓存的总大小。随着解码过程的进行，KV缓存会不断增长，直到输出序列完全生成。因此，KV缓存中的浮点数总量最大为：

$$2 \cdot \sum_{i=1}^{n_{layer}} [b \cdot (p + o) \cdot d_{model}]$$

或者简单表示为 $2 \cdot d_{model} \cdot n_{layer} \cdot b \cdot (p + o)$。

Llama-2-7B示例。以FP16精度的Llama-2-7B为例，其中$n_{layer} = 32$，$d_{model} = 4096$。整个模型的权重消耗约14GB内存，而KV缓存每个token约消耗0.5MB。这意味着当处理超过28K个token时（例如，28个请求，每个请求上下文长度为1K），KV缓存的大小就开始超过模型权重。

2.2 现有的KV缓存优化

优化维度。如上分析，层数、批量大小和上下文长度是决定KV缓存大小的三个关键因素。因此，现有的优化研究主要从这些角度寻求机会。

序列稀疏化方法。稀疏化上下文序列是打破上下文长度与KV缓存之间线性关系的有效方法，相关研究包括【Del Corro et al., 2023】、【【17，H2O: Heavy-hitter oracle for efficient generative inference of large language models，2024，Advances in Neural Information Processing Systems】，Zhang et al., 2024】、【Anagnostidis et al., 2023】、【Sukhbaatar et al., 2019】和【Rae & Razavi, 2020】。这些算法的通用思路是找出序列中的不重要token并丢弃它们的KV缓存。例如，滑动窗口注意力（Sliding Window Attention）【【4，Longformer: The long-document transformer，2020，arXiv】，Beltagy et al., 2020】只缓存一定数量的最近token。StreamingLLM【【16，Efficient streaming language models with attention sinks，2023，arXiv】，Xiao et al., 2023】发现除了最近的token，序列开头的token对输出也至关重要。Heavy-Hitter（H2O）【【17，H2O: Heavy-hitter oracle for efficient generative inference of large language models，2024，Advances in Neural Information Processing Systems】，Zhang et al., 2024】和Scissorhands【【10，Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time，2024，Advances in Neural Information Processing Systems】，Liu et al., 2024】通过比较token的注意力分数来排序其重要性。然而，这些算法都平等对待所有注意力层，为每层设置固定的KV缓存预算。

批处理优化方法。在批处理层面优化KV缓存主要需要高效地管理不同请求的内存。例如，vLLM【【9，Efficient memory management for large language model serving with pagedattention，2023，Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles】，Kwon et al., 2023】采用按需分配小内存块的方式，而不是为每个提示分配固定的大块内存，以减少批处理中的内存碎片。RadixAttention【【18，Efficiently programming large language models using sglang，2023，arXiv】，Zheng et al., 2023】则设法在批处理中当请求具有相同前缀时共享KV缓存。

层级优化方法的探索。与前两个维度相比，如何放宽KV缓存与层数之间的线性关系在很大程度上仍未被充分探索。FastGen【【7，Model tells you what to discard: Adaptive kv cache compression for llms，2023，arXiv】，Ge et al., 2023】是一项近期工作，它发现不同位置的层可能具有不同的最优序列级KV缓存驱逐策略。该工作提出了一种算法，在推理过程中为每个注意力头从多种策略中选择最佳驱逐策略，例如：1) 局部性策略（如滑动窗口），2) 特殊Token策略（如StreamingLLM），3) 局部性与频率策略（如Scissorhands, H2O）等。然而，对于所有被分配了相同驱逐策略的注意力头，FastGen仅给予它们一个统一的、预定义的缓存预算，如序列长度的30%。因此，如何为具有相同序列级token驱逐策略的层自适应地分配KV缓存预算，仍然是一个悬而未决的问题。

3 观察

层级重要性假设与度量方法。受到一些从层级视角优化LLM推理的先前工作【Del Corro et al., 2023; 【7，Model tells you what to discard: Adaptive kv cache compression for llms，2023，arXiv】，Ge et al., 2023】的启发，我们假设不同位置的注意力层在重要性方面扮演着不同的角色，因此应该有不同的最优KV缓存预算。为了更好地理解各层对输出嵌入的贡献，我们采用余弦相似度作为量化每层在推理过程中重要性的指标。具体来说，在每个注意力层中，我们跟踪每个输入嵌入在两个位置的隐藏状态，即自注意力模块之前和之后的嵌入。记自注意力之前的隐藏状态为$A$，之后的隐藏状态为$B$。

余弦相似度计算公式。两个嵌入$A$和$B$之间的余弦相似度计算如下：

$$\text{CosineSimilarity}(\mathbf{A}, \mathbf{B}) = \frac{\mathbf{A} \cdot \mathbf{B}}{\|\mathbf{A}\| \cdot \|\mathbf{B}\|} = \frac{\sum_{i=1}^{n} A_{i} \cdot B_{i}}{\sqrt{\sum_{i=1}^{n} A_{i}^{2}} \cdot \sqrt{\sum_{i=1}^{n} B_{i}^{2}}}$$

其中$A = (A_1, A_2, ..., A_n)$和$B = (B_1, B_2, ..., B_n)$是向量，$n$是维度。

实验观察与发现。对于任何给定的输入嵌入，我们可以得到每个注意力层的余弦相似度，该值大致量化了该层自注意力计算后嵌入的变化程度。通过比较各层之间的余弦相似度，我们可以找到一种按重要性对注意力层进行排序的方法。沿着这一思路，我们选择了4个代表性的LLM进行实验，即Mistral-7B【【8，Mistral 7b，2023，arXiv】，Jiang et al., 2023】、Llama2-7B-32K【【15，Llama 2: Open foundation and fine-tuned chat models，2023，arXiv】，Touvron et al., 2023】、Llama2-70B【【15，Llama 2: Open foundation and fine-tuned chat models，2023，arXiv】，Touvron et al., 2023】和Falcon-7B【【1，The falcon series of open language models，2023，https://arxiv.org/abs/2311.16867】 ，Almazrouei et al., 2023】。每个模型输入200个提示，我们跟踪每层中每个token的余弦相似度。图2显示了对所有提示进行平均后的结果。热力图的每一行展示了给定位置的输入嵌入如何通过所有层演变。

从热力图得出的洞察。从图的亮度中，我们可以得出以下见解：1) 总体而言，前半部分的层（颜色较深）对输出嵌入的贡献大于后半部分的层（颜色较浅）；2) 最初和最后几层往往比其他层更关键，具体取决于特定的模型和任务；3) 余弦相似度可以有效地描绘层级重要性，因为它反映的趋势与先前的研究（如Early-exiting【Del Corro et al., 2023】和FastGen【【7，Model tells you what to discard: Adaptive kv cache compression for llms，2023，arXiv】，Ge et al., 2023】）相符。这一观察为我们提供了一个简单而有效的指标，来设计一个能够从上下文和层两个维度优化KV缓存的新算法。

A2 方法细节

SQUEEZEATTENTION的二维优化思想。在本节中，我们描述了受上一章观察启发的SQUEEZEATTENTION算法。该算法最显著的特点是它将KV缓存中的token视为一个二维矩阵，一个维度是序列，另一个维度是层，并且两个维度将进行联合优化。

4.1 SQUEEZEATTENTION

序列维度与层维度的协同工作。在序列维度上，有多种我们可以直接结合使用的缓存驱逐策略，如最近最少使用（LRU）方法（Sliding Window, StreamingLLM）、最不常用（LFU）方法（Scissorhands, H2O）等。我们将一个在序列维度上压缩KV缓存的策略表示为$C_{seq}$，其缓存预算为$b_{init}$。需要注意的是，默认情况下所有层都具有相同的缓存预算，这与所有这些序列级KV缓存压缩算法所做的假设相同。在层维度上，SQUEEZEATTENTION首先在预填充阶段通过收集每层的余弦相似度来跟踪给定提示的层重要性。在预填充结束时，每个层都会得到一组余弦相似度，每个相似度对应一个流经该层的token。然后，我们使用提示token的平均值来表示该层在此任务中的层级重要性。通过使用KMeans根据层级重要性将各层聚类成组，我们重新分配$b_{init}$，使得更多的预算被分配给更“重要”的层组。由于各层具有不同的缓存预算，在解码阶段，$C_{seq}$会根据每层各自的预算独立工作。详细过程在算法1中描述。

算法1：SQUEEZEATTENTION
要求: prompt: 输入序列; $C_{seq}$: 序列维度的KV缓存压缩器; $b_{init}$: 每层的初始缓存预算; $n_{layer}$: 注意力层数; $p$: 超参数 ($0 < p < 1$); $K(i)$: 第$i$层的KV缓存。
1. 将prompt送入模型进行预填充，通过公式5计算第$i$层中第$j$个token的余弦相似度$cos\_sim(i)_j$。
2. 对于$i$从1到$n_{layer}$：
3. $cos\_sim(i) = \frac{\sum_{j=1}^{len(prompt)} cos\_sim(i)_j}{len(prompt)}$
4. 结束循环
5. $G_1, G_2, G_3 \leftarrow KMeans(cos\_sim(i))$ （通过$cos\_sim(i)$将层聚类为3组，其中$G_3$的平均$cos\_sim$最大）
6. 对于$i$从1到$n_{layer}$：
7. 如果$i \in G_3$：
8. $b(i) = b_{init} \times p$
9. 否则：
10. $b(i) = \frac{n_{layer} \times b_{init} - len(G_3) \times b_{init} \times p}{len(G_1) + len(G_2)}$
11. 结束判断
12. $K(i) = KV(prompt)$ （预填充后缓存提示的KV）
13. 结束循环
14. 对于$o$从1到$len(output)$： （逐个解码输出token）
15. 对于$i$从1到$n_{layer}$：
16. 如果$len(K(i)) > b(i)$：
17. $K(i) = C_{seq}(K(i), b(i))$ （根据该层自身的预算压缩KV缓存）
18. 结束判断
19. 基于压缩后的$K(i)$完成自注意力计算。
20. 结束循环
21. 结束循环

4.2 讨论

超参数p的决定。SQUEEZEATTENTION涉及一个超参数$p$，用于控制可以从“不重要”层中移除的初始预算百分比。$p$值越小，重新分配的预算就越多。在实验中，我们发现在大多数情况下，0.3-0.4是一个合理的选择范围。为了精确理解$p$的影响，我们进行了额外的实验来展示模型准确率如何随$p$值变化，详情请参阅附录A.2。

聚类为三组的原因。根据我们尝试的7个模型的观察，我们发现它们在层重要性方面都具有一个典型的模式（3组）。具体来说，第1组由少数特殊层组成（通常是最初和最后几层），可以看作是永远不应被驱逐的特殊token的类比。然后，第2组和第3组之间没有固定的界限，但我们发现第3组对嵌入的影响明显较小，可以看作是序列级方法中“频繁”和“不频繁”token的类比。因此，我们的3组策略可以重新表述为：“我们首先识别并优先处理特殊层（第1组），然后将其余层分为重要（第2组）和不重要（第3组）两类，再根据聚类结果重新分配缓存预算。”即使我们将层聚类成超过3组，我们只是将第2组和第3组分解成许多小组，但它们最终仍需被归约为两类，即减少预算或增加预算。为了保持模型准确性，我们只从第3组减少缓存预算，该组约占总层数的50%到70%。

A4 实验环境

LLM模型：我们选择了7个具有代表性的LLM进行评估，模型大小从6.7B到70B，上下文长度从2K到32K不等。这些模型包括：GPT-NeoX-20B、OPT-6.7B、Falcon-7B、Mistral-7B、Mixtral-8×7B、LLama2-7B-32k和LLama2-70B。

数据集：实验在5个数据集上进行：CNN/Daily Mail、XSUM、TriviaQA、SAMSUM和NarrativeQA。其中，TriviaQA、SAMSUM和NarrativeQA来自LongBench【【3，Longbench: A bilingual, multitask benchmark for long context understanding，2023，arXiv】，Bai et al., 2023】，数据长度通常超过8k。CNN/Daily Mail和XSUM的平均长度约为2k。详细信息见表1。

Table 1: 实验中使用的数据集。

基线算法：我们选择了3种序列级稀疏化算法作为基线，并将其集成到SQUEEZEATTENTION中。我们为每个模型分配了效果最好的基线算法，称之为“最佳基线算法”：
* Heavy-Hitter (H2O)【【17，H2O: Heavy-hitter oracle for efficient generative inference of large language models，2024，Advances in Neural Information Processing Systems】，Zhang et al., 2024】：通过比较每个token累积的注意力分数来识别序列中的重要token。
* Sliding Window Attention【【4，Longformer: The long-document transformer，2020，arXiv】，Beltagy et al., 2020】：一种“局部”策略，仅缓存最近token的KV嵌入。该方法在Mistral和Mixtral上效果尤为出色。
* StreamingLLM【【16，Efficient streaming language models with attention sinks，2023，arXiv】，Xiao et al., 2023】：除了最近的token，StreamingLLM总是缓存序列中的前n个token，因为它们被识别为“注意力汇点（sink tokens）”。我们根据论文建议取n=4。

硬件配置：所有实验均在AWS平台（p4d.24xlarge）上进行，配备8个Nvidia A100-40GB GPU，通过NVLinks互连（600 GB/s GPU点对点带宽）。

A5 实验结果

端到端结果：图3展示了SQUEEZEATTENTION与3种基线算法在7个模型和5个数据集上的比较结果。"Full Cache"（虚线）表示所有token的KV嵌入在推理过程中被完全缓存，因此代表了相对较好的模型准确率。图中蓝色和橙色曲线展示了模型准确率如何随KV缓存预算（占总序列长度的10%到100%）变化。在每个任务中，我们选择表现最好的序列级压缩算法作为基线，并在其基础上应用SQUEEZEATTENTION。如图所示，通过在层间重新分配缓存预算，SQUEEZEATTENTION在各种KV缓存预算下都能持续提高模型准确率。换言之，SQUEEZEATTENTION可以用更少的总KV缓存实现相似的推理准确率。

内存消耗：我们评估了SQUEEZEATTENTION的内存消耗。在不降低模型准确率的前提下，我们比较了三种设置所需的GPU内存。我们选择了Mistral-7B（滑动窗口）、GPT-NeoX-20B（Heavy-Hitter）和LLama2-70B（StreamingLLM），以覆盖三种基线算法以及小、中、大三种尺寸的模型。当模型和KV缓存无法装入单个GPU时，我们使用多GPU推理。

表2显示，在所有三种设置中，为了达到相同的模型准确率，与仅从序列维度进行压缩的算法相比，SQUEEZEATTENTION消耗的KV缓存预算最少。在某些情况下，它仅需要H2O算法缓存预算的三分之一。随后，我们使用PYTORCH PROFILER评估了推理过程中生成单个token所减少的内存使用量（不包括模型权重）。图4显示，与Full Cache方法相比，SQUEEZEATTENTION可以节省70%至80%的每token内存使用量，与基线算法相比可以节省25%至66%。

Table 2: 达到最佳准确率所需的KV缓存预算比较。选择了三个模型代表小（7B）、中（20B）和大模型（70B）。对于每个任务，我们选择最佳的现有序列级稀疏化算法作为基线。

Token生成吞吐量：由于SQUEEZEATTENTION节省了推理的内存成本，我们进一步探究了这些内存减少如何转化为token吞吐量的提升。我们选择了Mistral-7B和Llama-70B两个模型，代表小尺寸和 大尺寸模型。在固定内容长度的情况下，我们将Mistral-7B的批量大小从1增加到224，Llama2的批量大小从1增加到64。对于每个任务，我们选择了最佳基线算法，即Mistral-7B使用Sliding Window，Llama2-70B使用StreamingLLM。

表3显示了在8个A100-40GB GPU上的token吞吐量。在相同批量大小下，与Full Cache相比，SQUEEZEATTENTION可以将Mistral-7B的吞吐量提高至2.2倍，将Llama2-70B的吞吐量提高至1.4倍。此外，SQUEEZEATTENTION还使得两种模型的批量大小能够分别达到224和64，而这在Full Cache方法下会导致内存不足（OOM）。

Table 3: 在八个A100 GPU上，Mistral-7B和LLama2-70B使用SQUEEZEATTENTION和Full Cache的生成吞吐量（token/秒）。为保持模型准确率，SQUEEZEATTENTION对Mistral-7B使用20%的缓存预算，对LLama2-70B使用30%的缓存预算。“OOM”表示内存不足。更多结果见A.4。

算法开销：SQUEEZEATTENTION的计算开销来自两个操作：余弦相似度计算和K-means聚类。所有这些计算仅在预填充阶段发生。因此，SQUEEZEATTENTION的成本是一次性的，比那些需要在每次token生成迭代中进行额外计算的算法（如H2O）更具成本效益。

我们对Mistral-7B在有/无SQUEEZEATTENTION情况下的整个预填充阶段进行了性能分析，以比较实际运行时间。提示长度最高达8k token。如表4所示，SQUEEZEATTENTION仅导致预填充时间增加6.3%。需要注意的是，如果将解码时间考虑在内，这个比例会小得多，具体取决于实际生成的token数量。因此，SQUEEZEATTENTION的开销基本可以忽略不计。

Table 4: SQUEEZEATTENTION的开销（在单个Nvidia A100-40GB上的预填充时间，单位：秒）

A7 补充细节

5.6 局限性与更广泛的影响

局限性。SQUEEZEATTENTION与一个序列级稀疏化策略协同工作，因此我们的假设是：对于给定的模型和数据集，存在一个序列级KV缓存驱逐策略，该策略在一定的容忍度下不会损害模型准确性。然而，这些序列级算法的普适性仍然是一个活跃的研究课题。如果这个假设不能满足，SQUEEZEATTENTION可能无法达到预期的效果。

更广泛的影响。积极的社会影响包括降低LLM推理的能源成本和碳排放。此外，它还可以优化LLM应用的用户体验。由于该算法可以加速LLM推理，潜在的负面影响是它可能会加剧LLM的不当使用，例如生成有害信息。

A6 结论

本文提出了一种名为SQUEEZEATTENTION的二维KV缓存压缩算法。通过跟踪每个注意力层的余弦相似度，我们发现不同位置的层对于输出嵌入具有不同程度的重要性。受此观察启发，SQUEEZEATTENTION在注意力层之间重新分配KV缓存预算，以进一步降低推理的内存成本。在广泛的模型和任务上的实验表明，与仅在序列维度上压缩KV缓存的最新算法相比，SQUEEZEATTENTION可以用更低的内存消耗实现更高的模型准确率。

A7 附录

A.1 SQUEEZEATTENTION的详细开销

开销分解。在本节中，我们进一步评估了SQUEEZEATTENTION引入的开销，将其分解为两个主要操作所花费的时间：余弦相似度计算和K-means聚类。实验设置与5.5节相同。该实验使用单个Nvidia A100-40GB GPU，提示长度最高为8k token。

Table 5: 详细开销

计算过程与成本。计算余弦相似度的操作涉及计算两个大小为8000×4096的数组之间的相似度，重复32次（因为Mistral有32层）。此外，K-means将32个数字聚类为3个类别。因此，总时间消耗可以计算为 $0.00068 \times 32 + 0.001 = 0.02276$ 秒。值得注意的是，这个开销只产生一次，与处理的token数量无关。

A.2 超参数p的作用

p的角色与具体示例。超参数$p$在SqueezeAttention机制中至关重要，因为它直接决定了KV缓存的最终分配。举一个具体的例子：假设一个模型有32层，其中18层被认为是重要的，其余14层被认为是不太重要的。每层最初的预算为1000个token。如果我们设置$p$为0.3，我们将从不太重要的层中拿走70%的预算，并将其平均分配给重要的层。具体来说，不太重要的层的预算将减少到$1000 \times 0.3 = 300$。同时，重要层的预算将计算如下：$(1000 \times 18 + 1000 \times 0.7 \times 14) / 18 = 1544$。这样，总预算保持不变。

p的敏感性分析。为了评估超参数$p$的敏感性，我们在Samsum数据集上测试了Mistral-7B模型，p值范围从0.1到1.0。总KV预算设置为提示长度的20%。结果如下表所示：

Table 6: 精度随p的变化

分析与结论。100%意味着我们不改变模型KV缓存的结构。随着$p$的减小，更多的KV预算将被转移到其他层，这在一定程度上可以产生更好的结果。当$p$减小到10%甚至更低时，性能下降，因为不太重要的层的KV预算变得严重不足。我们可以清楚地观察到，在保持总预算不变的情况下调整$p$可以提高准确性。这证明了$p$对模型性能的影响，突显了其在优化KV预算分配中的重要性。

A.3 不同任务下的层重要性

跨任务层重要性实验。我们进行了一项额外实验，使用两个模型和多个数据集来确定不同层的重要性是否是模型的内在属性。

Table 7: Mistral-7B

Table 8: LLama2-70B

结果分析。表7显示了Mistral模型在三个不同数据集上的重要层分布：Samsum（Few shot）、TriviaQA（单文档问答）和LCC（代码，Python/C#/Java）。表8显示了LLama2-70B模型在三个不同数据集上的重要层分布：Xsum（摘要）、Samsum（Few shot）和LCC（代码，Python/C#/Java）。从这些表格中，我们可以观察到层的分组存在一个大致的模式，但会随具体任务而波动。我们认为存在一些任务敏感的层，它们在不同任务下可能被分到不同的组。同样，也存在一些总是重要或总是不重要的层。对这一现象的详细分析可以作为这项工作的有趣扩展。然而，我们仍然推荐使用自适应的方式，因为它可以精确地捕捉层的动态重要性。

A.4 最佳基线算法与SQUEEZEATTENTION的吞吐量比较

吞吐量对比实验。我们还进行了实验，比较SQUEEZEATTENTION与最佳基线在一系列批量大小下的吞吐量。两个实验都使用了512的输入长度和1024的输出长度。我们为每种算法选择了能达到最佳模型准确率的压缩超参数。结果表明，与那些仅从序列维度压缩KV缓存的SOTA算法相比，我们的算法可以显著提高吞吐量。

Table 9: 在八个A100 GPU上，Mistral-7B和LLama2-7B使用SQUEEZEATTENTION和最佳基线算法的生成吞吐量（token/秒）。为保持模型准确率，SQUEEZEATTENTION对Mistral-7B使用20%的缓存预算，对LLama2-7B使用40%的缓存预算，而Sliding Window对Mistral-7B使用30%的缓存预算，StreamingLLM对LLama2-7B使用60%的缓存预算。“OOM”表示内存不足。