SpeCache: Speculative Key-Value Caching for Efficient Generation of LLMs

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

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

Shibo Jie 1, Yehui Tang 2, Kai Han 2, Zhi-Hong Deng 1, Jing Han 3

速读

一句话结论
提出 SPECACHE 方法，通过在 GPU 中保留低比特 KV Cache 副本预测下一 token 的注意力分布，并从 CPU 异步预取极少量的关键高精度 KV Cache，在仅占用 10% 显存的情况下无损实现了长上下文大语言模型的高效推理。

要解决什么问题
大语言模型在处理长序列时，KV Cache 的体积会随序列长度线性增长，极易撑爆 GPU 显存（例如 LLaMA 2-7B 处理 2K 长度、批次 16 的序列时，KV Cache 达 8.4B，超过模型自身参数量）。为了缓解显存压力，现有做法主要有两类：一是直接压缩（如丢弃、合并或量化 KV Cache），但这会导致不可逆的信息遗忘，严重影响后续解码的准确率，因为当前步不重要的 KV 对可能在未来步骤中至关重要；二是将完整的 KV Cache 卸载到容量更大的 CPU 内存中，但这会引入极其频繁且庞大的 CPU-GPU 通信，导致推理延迟大幅增加。由于大模型的注意力机制具有极高的稀疏性（仅 0.5% 的 Key 就能覆盖 90% 的注意力分数），理论上只需将最相关的少量 KV Cache 留在显存即可，但难点在于：如何在不重新训练模型的前提下，既能精准挑出未来需要的 KV Cache，又能把数据传输的延迟隐藏掉。

怎么做的
SPECACHE 的核心思路是将完整的 16-bit KV Cache 存放在 CPU 内存中，而在 GPU 显存中仅保留一份极低精度（如 1-bit 或 2-bit）的 KV Cache 副本，并利用“投机解码”的思想提前一步预测并拉取需要的关键数据。具体而言，整个推理过程由三个关键机制构成：
首先是双 token 并行解码。在第 $t$ 步解码时，模型不仅处理当前真实的输出 token $T_t$，还会同时处理一个用于预测的“投机 token” $T'_{t+1}$。因为大模型解码是内存 IO 密集型的，GPU 计算利用率低，同时算两个 token 几乎不会增加额外延迟。
其次是异步预取机制。在计算 $T'_{t+1}$ 的注意力时，模型使用显存中的低比特 KV Cache 副本 $C'$，找出注意力分数最高的 top-$k$ 个 KV 对的索引 $\mathcal{K}_{t+1}$。一旦拿到索引，系统立刻启动后台传输，将这 $k$ 个 16-bit 的 KV 对从 CPU 内存异步拉取到 GPU 显存中。这个预取过程与模型后续层的计算完全并行，从而完美掩盖了 CPU-GPU 的通信延迟。
最后是高精度注意力计算。当进入第 $t+1$ 步解码时，所需的 top-$k$ 高精度 KV 对 $C_{\mathcal{K}_{t+1}}$ 已经就绪。此时，模型将这部分 16-bit 数据与低比特副本结合进行计算：

$$ O = \mathrm{Attn}([T_t, T'_{t+1}], C' \cup C_{\mathcal{K}_t}) $$


此外，为了将显存副本极限压缩到 1-bit，作者改进了 KIVI 库的量化算法。传统的 1-bit 量化会导致数值幅度异常，SPECACHE 假设权重在最大最小值之间服从均匀分布，重新设计了零点 $z_X$ 和缩放因子 $s_X$ 以最小化累计误差：

$$ z_X = \frac{3 \cdot \min X + \max X}{4}, \quad s_X = \frac{\max X - \min X}{2} $$
通过这种方式，所有值被平滑地映射到区间的两个中点上，大幅提升了 1-bit 量化下的生成质量。

效果如何
实验在 LLaMA-2、LLaMA-3-8B 和 Mistral-7B-Instruct-v0.2 模型上进行，测试了 LongBench（最高 32K 上下文）和 Needle-in-a-Haystack（大海捞针）基准。对比基线包括代表卸载路线的 InfLLM、代表丢弃策略的 StreamLLM 和 H2O，以及代表量化路线的 KIVI。量化结果显示，在 LongBench 任务中，当 KV Cache 压缩到仅占原大小的 10% 到 16% 时，SPECACHE 的平均表现全面超越其他基线方法。例如在 Mistral-7B 模型上，仅保留 10% 显存占用的 1-bit SPECACHE 与完整的 16-bit 基线相比，性能差距仅有 2%。在“大海捞针”测试中，纯 1-bit 量化会导致长文本检索能力严重崩塌，而引入 SPECACHE 预取机制后，检索准确率完全恢复到了 16-bit 的无损水平。在吞吐量方面，由于极大地节省了显存，SPECACHE 允许使用更大的批处理大小。在单张 NVIDIA A6000 GPU 上处理 32K 上下文时，1-bit SPECACHE 将最大可用批处理大小提升了 12 倍，整体解码吞吐量达到了原始 16-bit 方案的 4.6 倍。作者也承认了该方法的局限性：由于 CPU-GPU 交互代码目前基于 PyTorch 的多流机制和张量拷贝实现，并行度尚未达到理论最优，若定制底层算子还能进一步降低延迟；此外，预取数量 $k$ 的大小需要在模型性能和数据传输量之间做权衡。

1. 主要贡献

基于Transformer的大型语言模型（LLMs）在长文本任务上取得了显著成果，但随着序列长度的增加，有限的GPU显存（VRAM）资源难以满足呈线性增长的键值（KV）缓存（KV cache）需求，这已成为LLMs在长序列应用中的瓶颈。现有的KV cache压缩方法包括驱逐（eviction）、合并（merging）或量化（quantization）KV cache以减小其体积。然而，压缩会导致不可逆的信息遗忘，可能影响后续解码的准确性。

为了解决这一问题，本文提出了SPECACHE，其核心目标是在不进行重新训练的情况下，有效减少VRAM的使用，同时避免长序列的信息遗忘。该方法的创新点在于：
1. 充分利用大容量且易于扩展的CPU内存来卸载完整的KV cache。
2. 基于在VRAM中保留的低比特KV cache副本所衡量的重要性，在每个解码步骤中动态地将KV对抓取回VRAM。
3. 为了避免CPU-GPU通信引起的推理延迟，SPECACHE推测性地预测下一个token可能关注的KV对，从而允许在下一个解码步骤之前预取它们，实现了预取和计算的并行化。

Figure 1. SPECACHE使用低比特KV cache和推测性token来“猜测”下一个token最相关的top-k 16-bit KV对，并在下一个解码步骤之前预取它们。
Figure 1. SPECACHE使用低比特KV cache和推测性token来“猜测”下一个token最相关的top-k 16-bit KV对，并在下一个解码步骤之前预取它们。

2. 背景知识与关键观察

KV cache优化的现有方向 现有优化基于Transformer的LLMs推理过程中KV cache大小的工作可分为三个方向。首先是KV高效架构，通过修改模型结构来减小KV cache大小，例如多查询注意力（MQA）、分组查询注意力（GQA）、YOCO、多头潜在注意力（MLA）以及RWKV、RetNet和状态空间模型等。这些方法改变了模型架构，必须在预训练前应用，不适合现成LLMs的推理阶段优化。其次是训练后压缩，包括KV对的驱逐（如StreamLLM、H2O、Scissorhands、RoCo、FastGen）、合并（如CaM、D2O、MiniCache）和量化（如KIVI、KVQuant、ZipCache）。这些方法在推理阶段以贪婪的方式应用，压缩固有的信息丢失可能会丢弃对未来步骤有用的信息，从而降低LLMs的性能。最后是卸载与预取，如FlexGen、Huggingface的transformers库、InfLLM和ShadowKV，将KV cache卸载到CPU内存或磁盘。虽然这些方法在不丢失信息的情况下减少了VRAM使用，但频繁的CPU-GPU通信显著增加了推理延迟。

LLM推理的瓶颈分析 LLMs的推理过程可分为两个阶段：预填充（prefilling）阶段和解码（decoding）阶段。在预填充阶段，模型为提示（prompt）生成KV cache并产生第一个输出token，该阶段的瓶颈通常是GPU的计算速度，属于计算密集型（computation-bound）。在解码阶段，前一步的输出token作为输入生成下一个token，瓶颈在于高带宽内存（High Bandwidth Memory）和静态随机存取存储器（Static Random Access Memory）之间的数据传输速度，属于内存IO密集型（memory-IO bound）。许多技术（如批处理和推测解码）利用这一特性，在解码期间同时输入多个token，虽然增加了FLOPs，但提高了GPU利用率，缓解了单步延迟的显著增加，并最终提高了整体吞吐量。

探索注意力稀疏性 现有工作已探索LLMs中注意力的稀疏性，本文首先研究其稀疏程度及仅传输稀疏KV cache的潜在效率增益。在LLaMA-3-8B模型上使用截断至8196长度的PG19数据集进行实验，测量了命中率（hit rate），即两种稀疏注意力机制（依赖query的top-$k$注意力和贪婪驱逐）捕获的注意力分数占完整注意力分数的比例。其中，依赖query的top-$k$注意力在每次计算时仅包含得分最高的top-$k$ KV对；贪婪驱逐则类似于H2O，驱逐累积得分低的KV对。

稀疏性结论 基于Figure 2（左）得出结论：i) 注意力高度稀疏，仅$0.5\%$的keys能覆盖$90\%$的query注意力。ii) 稀疏性依赖于query。尽管两种方法强制相同的稀疏度，贪婪驱逐的命中率远低于依赖query的top-$k$注意力。这表明不同query倾向于关注不同的keys集合。贪婪驱逐虽优化当前query，但未能为后续query保留重要的KV对。因此，虽然驱逐方法能实现稀疏注意力，但无法恢复已驱逐的KV对，导致命中率较低。动态预取对于维持注意力性能至关重要。

完整卸载的开销与稀疏传输优势 如Figure 2（中）所示，简单卸载并预取整个KV cache会导致巨大的CPU-GPU传输开销。这启发我们在每步解码中仅预取少量最重要的KV对，从而减少传输的数据量。由于非连续内存传输（如主流框架PyTorch）会引入额外的时间开销，因此需要进行实验以验证传输稀疏KV cache相较于完整KV cache的效率优势。

稀疏传输的效率提升 如Figure 2（右）所示，尽管稀疏CPU-GPU传输会引入约5倍的延迟，但由于注意力机制在$1\%$稀疏度下仍保持高命中率，我们可以通过减少传输数据量来提高效率。例如，在Mistral-7B-Instruct-v0.2模型和$32\text{k}$上下文长度下，仅传输前$1\%$的KV对可将传输延迟降低$95\%$。

Figure 2. 左图：依赖query的top-k注意力和贪婪缓存驱逐的命中率。中图：GPU上单步解码延迟与从CPU加载KV cache到GPU的延迟对比。右图：连续与非连续CPU内存的CPU-GPU传输延迟。我们重点标出了在32k上下文长度下的完整和top-1% KV cache大小。在NVIDIA A6000 GPU上使用Mistral-7B-Instruct-v0.2测量。
Figure 2. 左图：依赖query的top-k注意力和贪婪缓存驱逐的命中率。中图：GPU上单步解码延迟与从CPU加载KV cache到GPU的延迟对比。右图：连续与非连续CPU内存的CPU-GPU传输延迟。我们重点标出了在32k上下文长度下的完整和top-1% KV cache大小。在NVIDIA A6000 GPU上使用Mistral-7B-Instruct-v0.2测量。

3. 方法细节

推测性预取的挑战与方案 尽管top-$k$预取的效率优势显著，但我们仍需要一种方法在加载KV cache之前预测KV对的注意力分数。为了实现预取和计算的并行化，我们需要尽早开始预取KV对。这带来了一个关键挑战：我们如何在注意力操作发生之前很久就确定哪些KV对是重要的？实际上，我们不需要预取确切的top-$k$ KV对；我们只需要预取具有高命中率的KV对，确保它们包含了绝大多数被显著关注的内容。基于此，我们提出了SPECACHE，这是一种推测性预测哪些KV对对未来的query重要并相应地预取它们的方法。

推测性预测的实现基础 推测性预测未来的注意力分数要求在VRAM中提供历史key cache和未来query的近似表示。对于前者，现有的免训练KV cache量化方法可以有效解决这个问题；例如，我们可以在VRAM中存储2-bit甚至1-bit的KV cache近似值。至于后者，我们提出在每个解码步骤中并行解码一个额外的“推测性token（speculative token）”，以近似下一个token。如Figure 3所示，整个推理过程可分为三个阶段：预填充（prefilling）、预解码（pre-decoding）和解码（decoding）。

预填充阶段（Prefilling） 在预填充阶段，我们采用逐层卸载的方法。在每个注意力层的计算完成之后，我们首先将KV cache量化为低精度，接着将原始的16-bit KV cache $C$ 完全卸载到CPU内存中，从而为下一层的KV cache腾出空间。一旦预填充阶段完成，我们就能获得一个准确的第一个输出token $T_1$。

预解码阶段（Pre-decoding） 在解码阶段开始之前，只有低比特的KV cache $C'$ 驻留在VRAM中。为了预取第一个解码步骤所需的16-bit KV对，我们引入了一个单独的解码步骤作为预解码。我们使用 $T_1$ 作为输入，生成一个初步的输出 $T'_2$（可能不完全准确），作为推测性token。同时，我们记录 $T_1$ 注意力分数计算中前 $k$ 个KV对的索引，记为 $\kappa_1$，并立即开始从CPU RAM中并行预取这些16-bit KV对，此过程与后续层的计算并行。预解码步骤结束后，紧接着开始第一个解码步骤。

解码阶段（Decoding） 在第 $t$ 个解码步骤开始之前，VRAM包含前一个输出token $T_t$、一个推测性token $T'_{t+1}$、低比特KV cache $C'$，以及 $T_t$ 所需的top-$k$ 16-bit KV对 $C_{\kappa_t}$。我们并行解码 $T_t$ 和 $T'_{t+1}$，公式如下：
$O = \text{Attn}([T_t, T'_{t+1}], C' \cup C_{\kappa_t})$
其中 $C' \cup C_{\kappa_t}$ 表示在本次计算中，用16-bit的 $C_{\kappa_t}$ 替换 $C'_{\kappa_t}$（即 $C'$ 中索引为 $\kappa_t$ 的子集）。

解码过程中的预取与卸载 在注意力计算期间，我们还会根据 $T'_{t+1}$ 的注意力分数记录top-$k$ KV对的索引 $\kappa_{t+1}$，这些KV对很可能是下一步解码 $T_{t+1}$ 所需的。一旦计算结束，我们立即开始预取这些16-bit KV对 $C_{\kappa_{t+1}}$，同时从VRAM中驱逐已经使用过的非top-$k$ 16-bit KV对，并卸载新生成的KV对。所有这些内存操作都与后续层的计算并行进行。

解码输出的生成与传递 解码后，$T_t$ 和 $T'_{t+1}$ 生成两个token：$T_{t+1}$ 和 $T'_{t+2}$。由于 $T_t$ 所需的16-bit KV对在其注意力计算之前已经预取，我们可以假设 $T_{t+1}$ 是准确的，并将其作为模型的输出。相反，$T'_{t+2}$ 作为下一个解码步骤的推测性token，因为它用于输出可能不够准确。

推理过程的延迟分析 SPECACHE的整个推理过程，与原始推理过程相比，仅增加了一个单一的预解码步骤，并在解码期间同时解码两个token。额外的预解码步骤在整个句子生成中可以忽略不计。虽然解码的token数量增加了，但两个token使用的模型权重和KV cache是相同的。由于LLMs的解码过程受限于内存IO（memory-IO bound），同时解码两个token允许共享访问模型权重和KV cache，而不会引入额外的延迟。此外，由于预取与GPU操作并行运行，整体推理延迟不会显著增加。

Figure 3. SPECACHE的图示。在预填充阶段，KV cache被逐层量化并卸载。在预解码阶段，我们使用第一个输出token计算第一个推测性token，并预取第一个解码步骤所需的16-bit KV对。在解码阶段的每一步中，我们同时解码两个token：输出token和推测性token。两者的结果作为下一步的输入。推测性token最相关的top-k 16-bit KV对在下一步之前被预取。
Figure 3. SPECACHE的图示。在预填充阶段，KV cache被逐层量化并卸载。在预解码阶段，我们使用第一个输出token计算第一个推测性token，并预取第一个解码步骤所需的16-bit KV对。在解码阶段的每一步中，我们同时解码两个token：输出token和推测性token。两者的结果作为下一步的输入。推测性token最相关的top-k 16-bit KV对在下一步之前被预取。

结合KV cache量化方法 由于SPECACHE在GPU中存储KV cache的低比特副本，它可以与任何KV cache量化方法结合。现有的训练后KV cache量化方法，如KIVI 【[15], KIVI: A tuning-free asymmetric 2bit quantization for KV cache + 2024 + ICML】，无需校准即可将KV cache量化为2-bit，而像KVQuant 【[8], Kvquant: Towards 10 million context length LLM inference with KV cache quantization + 2024 + ArXiv】 这样的方法在校准的帮助下可实现1-bit量化。我们利用KIVI来量化GPU副本，因为其简单且不需要校准。为了进一步突破KV cache压缩的极限，我们修改了KIVI以适用于1-bit量化。

原始KIVI量化方式 原始的KIVI按如下方式量化KV cache：
$Q(X) = \lfloor \frac{X - z_X}{s_X} \rceil, \quad X' = Q(X) \cdot s_X + z_X$
其中 $z_X = \min X$ 是零点（zero-point），$s_X = (\max X - \min X) / (2^B - 1)$ 是缩放因子（scaling factor），$\lfloor \cdot \rceil$ 是四舍五入操作。

改进的1-bit量化策略 然而，在1-bit量化中，$X'$ 的元素要么是 $\max X$ 要么是 $\min X$。这导致量化后的KV cache具有异常大的幅度，使得模型难以有效地执行文本生成。为了解决这个问题，我们在1-bit场景下改进了KIVI，假设权重在最小值 $\min X$ 和最大值 $\max X$ 之间服从均匀分布，并确保累积量化误差最小化。基于此，零点和缩放因子修改如下：
$z_X = \frac{3 \cdot \min X + \max X}{4}, \quad s_X = \frac{\max X - \min X}{2}$
这确保了在范围 $[\min X, (\min X + \max X) / 2)$ 内的所有值都被量化为该区间的中心点 $(3 \cdot \min X + \max X) / 4$，而在范围 $((\min X + \max X) / 2, \max X]$ 内的值被量化为该区间的中心点 $(\min X + 3 \cdot \max X) / 4$。

方法细节引用汇总

  1. 【[15], KIVI: A tuning-free asymmetric 2bit quantization for KV cache + 2024 + ICML】。在“结合KV cache量化方法”段落引用。原文描述KIVI是一种无需校准即可将KV cache量化为2-bit的方法，本文利用其来量化GPU副本。
  2. 【[8], Kvquant: Towards 10 million context length LLM inference with KV cache quantization + 2024 + ArXiv】。在“结合KV cache量化方法”段落引用。原文描述KVQuant是在校准帮助下实现1-bit量化的方法，以此说明现有量化技术的背景。

4. 实验环境

  • 数据集：

    • LongBench：用于评估长上下文理解能力，包含多个任务（Qasper, MultiFieldQA, HotpotQA, 2WikiMQA, MuSiQue, GovReport, MultiNews, PassageRetrieval, LCC, RepoBench-P等）。序列长度限制：LLaMA-2设置为4k，Mistral设置为32k，LLaMA-3设置为8k。
    • Needle-in-a-Haystack (NIAH)：合成检索任务，用于评估长上下文检索能力，本文中Mistral上下文长度为32k，LLaMA-3为8k。

  • 模型参数：LLaMA-2-7B, Mistral-7B-Instruct-v0.2, LLaMA-3-8B-Instruct。

  • 硬件配置：单张NVIDIA A6000 GPU（48GB VRAM）。
  • 软件配置：使用PyTorch实现，利用其多流机制（multi-stream mechanism）和Tensor.copy()方法进行CPU-GPU交互。基于HuggingFace transformers库。量化基线使用KIVI（保留128个残差KV对，量化组大小为32和64）。SPECACHE从CPU预取top-64的16-bit KV对，为保持总显存不变，将残差长度减小至64。

5. 实验结果

LongBench性能评估
在LongBench的多个任务上对不同LLM进行了评估（结果见Table 1和Table 2）。观察发现，随着KV压缩程度的增加（较低的位宽或较大的组大小），模型性能下降。然而，当结合SPECACHE时，大部分性能损失得以恢复。KV压缩程度越大，该方法的优势越明显。具体而言，在仅保留约10% KV cache大小在GPU中的情况下，SPECACHE在Mistral-7B-Instruct-v0.2和LLaMA-3-8B-Instruct上与16-bit基线的性能差距分别保持在2%和1%以内。与InfLLM、H2O和StreamLLM等代表性压缩方法相比，SPECACHE在更高的压缩比下实现了更好的平均性能。

Needle-in-a-Haystack基准测试
评估了模型在应用KV压缩后的长上下文检索能力（Figure 4）。在1-bit KV cache量化下，LLMs的长上下文检索能力受到严重破坏。然而，在应用SPECACHE（即在1-bit量化基础上从CPU预取16-bit KV cache）后，模型在1-bit下的性能恢复到了与完整的16-bit KV cache相当的水平。

Figure 4. 在Needle-in-a-haystack基准测试上的性能。我们使用g = 32进行量化，导致Mistral-7B-Instruct-v0.2和LLaMA-3-8B-Instruct的KV cache压缩比分别为0.13和0.14。
Figure 4. 在Needle-in-a-haystack基准测试上的性能。我们使用g = 32进行量化，导致Mistral-7B-Instruct-v0.2和LLaMA-3-8B-Instruct的KV cache压缩比分别为0.13和0.14。

效率与吞吐量评估
测试了SPECACHE在解码期间的最大吞吐量（Table 3）。在2k、8k和32k的上下文长度下，通过增加批处理大小（Batch Size）以最大化48GB VRAM的使用。结果表明，在2-bit和1-bit量化下，SPECACHE允许批处理大小分别增加最多7倍和12倍，导致整体吞吐量相比原始设置分别提高了3.4倍和4.6倍。特别是在上下文长度较大时，原始KV cache只能容纳较小的批处理大小，导致并行度较低，此时SPECACHE的加速效果更为显著。

关于 $k$ 的消融实验
对预取KV对的数量 $k$ 进行了消融研究（Figure 5）。即使使用较小的 $k$（如 $k=16$），SPECACHE也比KIVI基线（$k=0$）提供了显著的改进。随着 $k$ 的增加，模型性能提升，但传输的数据量也随之增加，实际应用中可以对 $k$ 进行权衡。

Figure 5. 关于k的消融实验。我们使用g = 32的1-bit SPECACHE。
Figure 5. 关于k的消融实验。我们使用g = 32的1-bit SPECACHE。

量化方法的消融实验
为了突出为1-bit KIVI提出的新零点和缩放因子的重要性，进行了消融实验（Table 5）。当使用原始KIVI时，即使结合SPECACHE，模型性能仍然非常差。修改量化方法后，模型性能显著提高，SPECACHE的优势变得更加明显。

与非推测性抓取的对比
将SPECACHE与一种简单策略进行了比较（Table 4）：该策略不使用推测性token，而是使用上一步的准确输出token来估计top-$k$ KV对并抓取它们，然后重新计算注意力。这种策略的抓取必须在注意力计算完成后开始，且不能并行运行。结果显示，使用准确的输出token略微提高了模型性能，但显著增加了延迟。随着批处理大小的增加，加载KV对的延迟成为主导因素，这突显了计算与预取并行的重要性。

6. 结论

本文提出了SPECACHE，这是一种低延迟、高性能且免训练的KV cache压缩方法。SPECACHE将高精度的KV cache存储在CPU RAM中，而将低比特（低至1-bit）的KV cache存储在GPU VRAM中。它利用推测性token和输出token的联合解码来预测并预取下一个解码步骤所需的top-$k$ KV对。这种方法有效地恢复了因低比特KV cache而丢失的信息，同时实现了预取和计算的并行化，为长序列LLM推理提供了高效的解决方案。未来的工作可以基于自定义底层算子进一步提升SPECACHE的执行效率。

7. 附录

预填充算法（Algorithm 1） 首先输入序列 $T \in \mathbb{R}^{L \times d}$，计算 $Q = T W_q, K = T W_k, V = T W_v$。接着，对 $K$ 和 $V$ 进行量化得到 $K' = \text{Quant}(K)$ 和 $V' = \text{Quant}(V)$。然后，定义完整的16-bit缓存 $C = \{V, K\}$ 和低比特缓存 $C' = \{V', K'\}$，并将 $C$ 卸载（$\text{Offload}(C)$）到CPU中。最后，计算注意力分数 $A = \text{MaskedSoftmax}(QK^\top)$，并计算输出 $O = AV W_o$。

Input: T
Q = T * W_q, K = T * W_k, V = T * W_v
K_prime = Quant(K), V_prime = Quant(V)
C = {V, K}, C_prime = {V_prime, K_prime}
Offload(C)
A = MaskedSoftmax(Q * K^T)
O = A * V * W_o
Output: O

预解码算法（Algorithm 2） 首先输入第一个token $T_1 \in \mathbb{R}^{1 \times d}$，计算其 $Q_1, K_1, V_1$。接着，将新的KV追加到低比特缓存中，即 $K = [K', K_1], V = [V', V_1]$。然后，计算注意力分数 $A = \text{Softmax}(Q_1 K^\top)$，并获取前 $k$ 个元素的索引 $\mathcal{K}_1 = \text{ArgTopK}(A)$。随后，启动并行预取操作 $\text{Pre-fetch}(C_{\mathcal{K}_1})$，将所需的16-bit KV对取回。最后，计算输出 $O = AV W_o$。

Input: T_1
Q_1 = T_1 * W_q, K_1 = T_1 * W_k, V_1 = T_1 * W_v
K = [K_prime, K_1], V = [V_prime, V_1]
A = Softmax(Q_1 * K^T)
K_1_idx = ArgTopK(A)
Pre-fetch(C_{K_1_idx})
O = A * V * W_o
Output: O

解码算法（Algorithm 3） 首先输入包含当前token和推测token的矩阵 $T = [T_t, T'_{t+1}] \in \mathbb{R}^{2 \times d}$，计算对应的 $Q_t, K_t, V_t$。接着，对当前步生成的 $K_t, V_t$ 进行量化得到 $K'_t, V'_t$。在构建用于注意力计算的 $K$ 和 $V$ 时，用预取的16-bit缓存替换低比特缓存中的对应部分，即 $K = [K' \cup K_{\mathcal{K}_t}, K_t]$ 和 $V = [V' \cup V_{\mathcal{K}_t}, V_t]$。然后，计算注意力分数 $A = \text{MaskedSoftmax}(Q_t K^\top)$，并根据推测token的注意力分数获取下一步所需的索引 $\mathcal{K}_{t+1} = \text{ArgTopK}(A_{1,:})$。随后，启动预取操作 $\text{Pre-fetch}(C_{\mathcal{K}_{t+1}})$。更新缓存状态，将当前步生成的16-bit缓存卸载，即 $C = [C, \text{Offload}(C_t)]$，并更新低比特缓存 $C' = [C', C'_t]$。最后，计算输出 $O = AV W_o$。

Input: T = [T_t, T_prime_{t+1}]
Q_t = T * W_q, K_t = T * W_k, V_t = T * W_v
K_prime_t = Quant(K_t), V_prime_t = Quant(V_t)
K = [K_prime U K_{K_t}, K_t], V = [V_prime U V_{K_t}, V_t]
A = MaskedSoftmax(Q_t * K^T)
K_{t+1}_idx = ArgTopK(A_{1,:})
Pre-fetch(C_{K_{t+1}_idx})
C_t = {V_t, K_t}, C_prime_t = {V_prime_t, K_prime_t}
C = [C, Offload(C_t)], C_prime = [C_prime, C_prime_t]
O = A * V * W_o
Output: O

Needle-in-a-Haystack基准测试设置
遵循Greg Kamradt 【[5], Needle in a haystack - pressure testing llms + 2023 + URL】 的设置，将句子“The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day”插入到Paul Graham的文章中，并添加问题：“What is the best thing to do in San Francisco? Here is the most relevant sentence in the context:”。之后，使用GPT-4根据以下标准对模型的响应进行评分：1分（完全无关）；3分（轻微相关但不一致）；5分（中度相关但有不准确之处）；7分（一致但有轻微遗漏）；10分（完全准确且完美一致）。

LongBench完整结果
附录中提供了正文中未列出的所有15个LongBench任务的完整结果。测试中，KIVI和SPECACHE均使用组大小 $g=64$ 的2-bit和1-bit量化。

8. 补充细节

社会影响声明
本文基于大型语言模型（LLMs），这些模型具有潜在的社会影响，包括对偏见、错误信息和可访问性等问题的担忧。尽管这些问题已广为人知，但本文的研究旨在致力于提高LLMs的运行效率。随着该领域的不断发展，持续的伦理监督和跨学科合作是必不可少的。