SCBENCH: A KV Cache-Centric Analysis of Long-Context Methods

发表时间: 2024-12 · arXiv:2412.10319

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

速读

一句话结论 本文提出了一个以KV缓存为中心的共享上下文长文本基准测试 SCBench,通过在多轮交互和多请求场景下评估多条技术路线,揭示了压缩KV缓存的方法在多轮查询中会因注意力偏移而失效,并据此提出了一种兼顾首轮与多轮性能的免训练稀疏注意力方法 Tri-shape。

要解决什么问题 现有的长上下文大型语言模型评估通常局限于单次请求场景,忽略了真实世界应用中广泛存在的KV缓存复用机制。在多轮对话、多步推理或多个用户共享代码库等真实场景中,上下文会在不同请求间共享。为了降低长文本推理的显存和计算开销,业界提出了大量优化方法,但这些方法在共享上下文场景下的表现却是一个盲区。具体而言,许多方法试图通过压缩KV缓存来突破显存瓶颈,将显存复杂度降低到 $O(n)$ 以下。然而,在多轮查询中,模型对上下文中关键信息的注意力分布会随着查询的不同而发生剧烈偏移。如果在第一轮查询时就丢弃了部分KV缓存,后续查询一旦需要用到被丢弃的信息,就会导致准确率断崖式下降。这种因注意力分布偏移导致的上下文信息丢失,是现有单次请求基准无法测出的核心卡点。

怎么做的 本文首先建立了一个以KV缓存为中心的分析框架,将长上下文方法系统地划分为四个生命周期阶段:KV缓存生成(如稀疏注意力、状态空间模型、提示词压缩)、KV缓存压缩(如丢弃部分缓存、量化)、KV缓存检索(根据前缀检索相关块)和KV缓存加载(从内存动态加载到显存)。基于此框架,本文构建了 SCBench 基准测试,包含12个涵盖字符串检索、语义检索、全局信息处理和多任务处理的子任务,并强制在多轮模式(使用标准答案作为后续上下文)和多请求模式(跨用户共享上下文且无权访问彼此查询)下进行测试。在分析现有方法缺陷的基础上,本文发现如果解码阶段保持密集的 $O(n)$ 显存,仅在编码(预填充)阶段引入稀疏性,就能在多请求场景中保持稳健。为此,本文提出了一种名为 Tri-shape 的新型免训练稀疏注意力方法。传统的 A-shape 稀疏注意力仅保留初始的注意力下沉词(sink token)和局部窗口,这会破坏模型的指令遵循能力,导致输出混乱。Tri-shape 的核心思路是在 A-shape 的基础上,在预填充阶段为最后一个窗口的查询区域额外保留一个密集的注意力空间,从而在注意力矩阵底部形成一个三角形的稀疏模式。这一设计既利用了稀疏编码降低了预填充阶段的计算复杂度,又通过保留完整的底部查询区域增强了首轮生成的指令遵循能力,同时在后续多轮请求中凭借密集的解码阶段维持了接近全注意力机制的准确率。

效果如何 实验在4块或8块 A100 以及4块 H100 GPU 上进行,测试了包括 Llama-3.1-8B/70B、Qwen2.5-72B/32B、Llama-3-8B-262K 和 GLM-4-9B 在内的6个基于 Transformer 的长上下文大模型。对比基线涵盖了8条技术路线的13种方法,代表性基线包括:门控线性RNN路线的 Codestral-Mamba,混合架构路线的 Jamba-1.5-Mini,动态稀疏注意力路线的 MInference,静态稀疏注意力路线的 A-shape,KV缓存丢弃路线的 StreamingLLM 和 SnapKV,KV缓存量化路线的 KIVI,KV缓存加载路线的 Quest 和 RetrievalAttention,以及提示词压缩路线的 LLMLingua-2。量化结果表明,显存复杂度低于 $O(n)$ 的 KV缓存压缩方法(如 StreamingLLM 和 SnapKV)在多轮精确检索任务中几乎完全失效,在1/4计算预算下准确率分别暴跌26个和19个百分点。依赖查询感知的压缩方法(如 SnapKV)在无法获取查询的多请求模式下泛化能力极差。相比之下,保持 $O(n)$ 显存的稀疏编码方法表现优异,其中动态稀疏注意力 MInference 在所有任务中表现最佳,在1/32预算下即可达到其他静态方法1/4预算的性能;本文提出的 Tri-shape 则在多请求模式下显著提升了首轮准确率。实验也暴露了各路线的局限性:提示词压缩方法虽然在多示例上下文学习等全局任务中有效,但会严重损害检索性能;状态空间模型及其混合架构虽然在单轮交互中表现良好,但在多轮代码检索和数学任务中准确率会出现明显下降。

A1 主要贡献

长上下文大型语言模型(LLMs)虽然推动了众多下游应用,但也带来了巨大的计算和内存效率挑战。为应对这些挑战,研究者们开发了以KV缓存为中心的优化方法。然而,现有基准测试通常在单次请求场景下进行评估,忽略了真实世界应用中KV缓存的完整生命周期。由于KV缓存复用已在vLLM、SGLang等推理框架以及OpenAI、Microsoft、Google和Anthropic等LLM服务提供商中广泛采用,这一疏忽尤为关键。

为填补这一空白,本文提出了SCBENCH (Shared-ContextBENCH),一个从KV缓存中心视角全面评估长上下文方法的基准测试。该基准测试关注KV缓存的四个阶段:1) KV缓存生成2) KV缓存压缩3) KV缓存检索,和4) KV缓存加载。SCBench使用带有共享上下文的测试样本,涵盖12项任务和两种共享上下文模式,覆盖了四类长上下文能力:字符串检索、语义检索、全局信息处理和多任务处理。

借助SCBench,本文对八类长上下文解决方案进行了广泛的以KV缓存为中心的分析,这些方案包括门控线性RNN(如Codestal-Mamba)、Mamba-Attention混合模型(如Jamba-1.5-Mini),以及稀疏注意力、KV缓存丢弃、量化、检索、加载和提示压缩等高效方法。评估在六个基于Transformer的长上下文LLM上进行:Llama-3.1-8B/70B、Qwen2.5-72B/32B、Llama-3-8B-262K和GLM-4-9B。

核心发现

  1. 如图3所示,内存复杂度低于O(n)的方法在多轮场景中性能不佳。这类方法在首次查询时表现良好,但在后续请求中准确率下降。
  2. 具有O(n)内存和低于O(n²)预填充计算的稀疏编码方法表现稳健,可以在多个查询中逼近全注意力机制的准确率。
  3. 动态稀疏性比静态模式能产生更具表达力的KV缓存。
  4. 混合架构中的层级稀疏性可以在保持强劲性能的同时减少内存使用。
  5. 在长文本生成场景中发现了注意力分布偏移问题。

本文贡献

A3 背景知识与设计原则

以KV缓存为中心的视角看长上下文方法

近期,一系列工作探索了多种策略来降低长上下文LLM的推理成本,使其能够以更低的计算开销应用于下游任务。在长上下文LLM推理中,KV缓存通过有效减少解码阶段的计算开销扮演了关键角色。这导致了许多专注于KV缓存管理和调度的系统级优化。

本文提出了一个新颖的视角:这些长上下文方法可以被看作是在不同阶段围绕KV缓存进行的优化。具体来说,本文引入了一个以KV缓存为中心的框架,系统地将长上下文方法分为四个阶段:KV缓存生成、压缩、检索和加载,如图1所示。

该框架的四个阶段定义如下:

  1. KV缓存生成 (KV Cache Generation):此阶段优化推理过程中KV缓存的高效生成。技术包括稀疏注意力(如A-shape、Tri-shape、MInference【索引41,MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention,2024,NeurIPS】、NSA【索引103,Native sparse attention: Hardware-aligned and natively trainable sparse attention,2025,arXiv】、MoBA【索引62,Moba: Mixture of block attention for long-context llms,2025,arXiv】)、状态空间模型(SSM)或混合方法(如Mamba【索引30,Mamba: Linear-time sequence modeling with selective state spaces,2024,CoLM】、Jamba【索引54,Jamba: A hybrid transformer-mamba language model,2024,arXiv】)以及提示压缩(如LLMLingua-2【索引73,LLMLingua-2: Data distillation for efficient and faithful task-agnostic prompt compression,2024,ACL】)。
  2. KV缓存压缩 (KV Cache Compression):生成后,KV缓存在存储前被压缩。方法包括KV缓存丢弃(如StreamingLLM【索引99,Efficient streaming language models with attention sinks,2024,ICLR】、SnapKV【索引53,SnapKV: LLM knows what you are looking for before generation,2024c,NeurIPS】)和KV缓存量化(如KIVI【索引61,KIVI: A tuning-free asymmetric 2bit quantization for KV cache,2024e,ICML】)。
  3. KV缓存检索 (KV Cache Retrieval):根据请求的前缀从存储池中检索相关的KV缓存块,以减少首个令牌生成时间(TTFT)。方法包括语义检索方法,如CacheBlend【索引101,Cacheblend: Fast large language model serving with cached knowledge fusion,2024a,arXiv】。
  4. KV缓存加载 (KV Cache Loading):此阶段将KV缓存从存储(如VRAM、DRAM、SSD或RDMA)动态加载到GPU片上SRAM并计算稀疏注意力。方法包括Quest【索引91,QUEST: Queryaware sparsity for efficient long-context LLM inference,2024,ICML】、RetrievalAttention【索引56,Retrievalattention: Accelerating longcontext llm inference via vector retrieval,2024b,arXiv】和MagicPIG【索引12,MagicPIG: LSH sampling for efficient LLM generation,2025,ICLR】。

本文评估了表1中列出的13种长上下文方法的所有四个阶段。此外,还列出了每种方法的KV缓存大小、预填充阶段复杂度和解码阶段复杂度,以及是否在预填充和解码阶段执行了高效操作。

Tri-shape稀疏注意力。本文引入了一种名为Tri-shape的新型免训练稀疏注意力方法,它能提高首轮准确率(图5)。与仅保留初始令牌(sink token)和局部窗口的A-shape不同,Tri-shape还保留了最后一个窗口的查询区域,从而在预填充阶段形成一个三角形的稀疏注意力模式。这一设计的动机源于本文在SCBench上的发现:采用密集解码的A-shape在多次请求后性能有所提升。因此,Tri-shape旨在同时增强首次(turn-0)和多次请求的性能,同时保持LLM的指令遵循能力。值得注意的是,近期有并发工作【索引1,Star attention: Efficient llm inference over long sequences,2024,arXiv】也探索了类似的模式来加速长上下文预填充。

A2 方法细节

基准构建

SCBench包含12个任务,用于评估四种长上下文能力:字符串检索、语义检索、全局信息处理和多任务处理。这些任务在多轮多请求两种共享上下文模式下进行。任务涵盖了代码、检索、问答、摘要、上下文学习和多跳追踪等多个领域(图2b)。SCBench共包含931个多轮会话,总计4,853个查询,平均每个会话5轮。任务的统计数据见表2,示例和配置见表3。

长上下文任务细节

字符串检索。长上下文LLM的核心要求是从冗长且可能包含噪声的输入中检索相关信息。本文受算法问题解决(如LeetCode)的启发,设计了三个不同难度的任务。通过改变目标字符串的位置,评估模型利用其完整上下文窗口的能力。

语义检索。许多现实世界的长上下文应用要求模型具备语义理解能力。SCBench包含了评估此类能力的必要任务,因为有损的长上下文方法在多请求场景中通常难以抽象或理解信息。

全局信息处理。除了检索,一些长上下文任务需要利用和聚合全局上下文信息,如摘要、统计任务和上下文学习(ICL)。本文的基准包含三个任务来评估不同长上下文方法在多请求场景下处理全局信息的能力。

多任务处理。在现实应用中,LLM通常在单个会话中使用共享输入上下文处理多个任务。为了反映这一点,SCBench包含了两个多任务处理任务:

长上下文共享模式细节

除了精心设计的长上下文任务外,我们还包括了多轮多请求两种共享上下文模式,以更好地反映真实世界的应用,如图2b所示。

A4 实验环境

A4 实验结果

主要结果
表4、表10和图9展示了各种长上下文方法在不同基础LLM上,跨任务和共享上下文模式的性能。

  1. 检索任务性能普遍不佳:在检索任务中,除了MInference,大多数方法表现不佳,尤其是在字符串匹配等精确检索任务中。
  2. 稀疏注意力优于稀疏解码:随着请求轮数的增加,稀疏注意力方法的性能优于稀疏解码方法。其中,A-shape的性能提升最为显著。Tri-shape通过在A-shape的基础上增加密集的底部查询token,提升了首轮性能,但对后续轮次影响不大。它在各种任务中泛化良好,性能仅次于MInference。分析表明,Tri-shape改善了首轮的指令遵循能力,而A-shape会破坏指令信息导致随机输出(见表17)。
  3. KV缓存压缩方法性能不佳:在共享上下文场景中,KV缓存压缩方法通常表现不佳,仅在第一轮提供微小的好处。
  4. 提示压缩的权衡:提示压缩改善了需要全局信息的任务(如多示例ICL)的性能,但显著降低了与检索相关的性能。
  5. SSM相关模型表现:SSM-Attention混合模型在单轮交互中表现良好,但在多轮任务(特别是RepoQA和Math)中准确率下降。门控线性RNN模型在共享上下文模式下表现挣扎。

结果分析

A5 结论

本文解决了长上下文方法评估中的一个关键空白,即现有评估传统上侧重于单轮交互,而忽略了在真实世界LLM应用中常见的共享长上下文场景。为了弥补这一不足,我们引入了SCBench,这是一个综合性基准,用于评估在两种共享上下文模式下,涉及KV缓存复用的长上下文方法。该基准涵盖12项任务,包括字符串检索、语义检索、全局信息处理和多任务处理。

基于此基准,我们将长上下文方法分为四个以KV缓存为中心的阶段:生成、压缩、检索和加载。我们评估了八类方法(如门控线性RNN、混合模型、稀疏注意力、KV缓存丢弃、量化、检索、加载和提示压缩)在八个最先进的LLM上的表现。

我们的研究结果揭示了一个清晰的KV缓存管理权衡:O(n)内存方法在多请求场景中表现出色,而sub-O(n)方法虽然在单轮交互中表现良好,但在处理复杂交互时则举步维艰。这些发现强调了在共享上下文、多轮场景中评估长上下文方法的必要性,为改进未来的长上下文模型和架构提供了更现实的基准和宝贵的见解。

A6 附录

相关工作

与现有长上下文基准的比较

我们将SCBench与现有的长上下文基准在评估的长上下文能力、考虑的请求类型和采用的实现方式上进行了比较,如表6所示。

我们还直接比较了长上下文方法在先前基准和SCBench上的测试结果,以展示我们基准提供的独特见解。主要比较了两种常见的长上下文能力:摘要(如表7所示)和检索(如表8所示)。我们发现SCBench能更好地识别长上下文方法在KV缓存复用场景下的弱点,例如KV缓存压缩方法在多请求模式和多轮模式的后续查询中普遍表现不佳,以及稀疏注意力在多轮模式下准确率的提升。

高效长上下文方法的超参数

我们对所涵盖的高效长上下文方法在不同计算预算下进行了广泛实验。结果分别如图7(多轮模式)和图8(多请求模式)所示。

从结果中可以得出以下见解:

  1. 大多数方法在1/2预算下性能下降很小(例如,A-shape和Tri-shape下降5-6个点,SnapKV下降11个点)。然而,随着稀疏度的增加,性能显著下降。例如,在1/4预算下,StreamingLLM和SnapKV分别下降了26和19个点。
  2. 更精确的稀疏方法即使在更高稀疏度下也能保持性能。例如,MInference在1/32预算下实现的性能与A-shape和Tri-shape在1/4预算下的性能相当。
  3. 不同模式下性能差异显著。虽然某些方法在单轮场景中表现相似,但它们在多轮和多请求场景中表现差异巨大。例如,SnapKV在第1轮中优于StreamingLLM,但在第2轮中表现明显更差。在某些任务中(如长文档QA和摘要),改变预算对第1轮性能影响不大,但对第2轮及后续轮次有显著影响。

实验细节

D.1 长上下文方法细节

D.2 额外的实现细节

表9报告了实验中使用的长上下文方法的配置。对于Mamba-Codestral-7B-v0.1模型和AI21-Jamba-1.5-Large,报告了其架构细节。对于稀疏注意力方法,报告了Tri-shape和A-shape的局部大小和初始大小。对于MInference,报告了用于搜索稀疏模式的任务和搜索空间。对于KV缓存压缩、量化、检索和加载,使用了其原始实现中的默认超参数。

额外的实验结果

Llama-3.1-70B、Qwen2.5-32B和Llama-3-8B-262K的结果如表10所示,其结论与主要结果一致。MInference在所有任务中持续优于其他方法。Tri-shape由于整合了底部查询token,在多请求模式下表现出色,提升了首轮性能。KV缓存压缩方法在共享上下文中表现不佳。提示压缩方法在需要全局上下文的任务中表现良好,但在检索任务中表现不佳。

图9和图10展示了不同方法在各类任务和轮次中的详细性能。在字符串和语义检索任务中,MInference和FullAttention表现最佳,而StreamingLLM和SnapKV几乎完全失败。在全局信息任务中,大多数方法性能稳定,但在多任务场景下,性能下降明显。

表11和表12提供了所有子任务在多轮和多请求模式下的详细分解结果。总体而言,GLM-4-1M和MInference在大多数任务中表现一致优越,尤其是在检索、QA和ICL方面。Qwen2模型在自然语言处理和ICL方面也很出色。稀疏注意力方法A-shape和Tri-shape在特定领域表现尚可,而StreamingLLM和SnapKV在所有任务中持续表现不佳。多请求模式下,MInference在数学任务上的性能显著提升,显示出其在重复查询下的适应能力。

使用生成内容作为上下文的错误传播

在多轮测试中,我们遵循先前工作,使用标准答案而非模型生成内容作为下一轮查询的上下文。此方法可防止误导性生成内容对后续轮次的干扰。这里我们分析了禁用此设置(即使用模型生成内容作为上下文)的效果,以观察我们的发现是否仍然成立。

如表13所示,当使用模型生成内容作为上下文时,我们得到了与主结果(表4)相似的结论:密集解码方法通常优于稀疏解码方法,更鲁棒和动态的稀疏模式优于静态稀疏方法。但使用模型生成内容作为上下文确实显示出更低的总体准确率,这表明了错误传播现象,即先前查询的误导性答案会影响后续轮次的表现。

案例研究