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

  • 作者/机构: Yucheng Li†, Huiqiang Jiang‡, Qianhui Wu, Xufang Luo, Surin Ahn, Chengruidong Zhang, Amir H. Abdi, Dongsheng Li, Jianfeng Gao, Yuqing Yang, Lili Qiu
  • 单位: Microsoft Corporation, †University of Surrey
  • 链接: https://aka.ms/SCBench

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. 在长文本生成场景中发现了注意力分布偏移问题。

本文贡献
* 提出了一个新的基准测试SCBench,用于在两种典型的KV缓存复用场景(多轮和多请求)下评估长上下文方法,提供更真实的评估。
* 设计了一套广泛的下游任务,涵盖四个长上下文能力,涉及12个不同领域的子任务。
* 从以KV缓存为中心的视角系统地对长上下文方法进行分类,并使用SCBench在八个最先进的开源长上下文LLM上评估了13种不同的长上下文方法(包括本文新提出的稀疏注意力方法Tri-shape)。全面的分析揭示了关于编码和解码中的稀疏性、任务复杂性等方面的关键见解。

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)的启发,设计了三个不同难度的任务。通过改变目标字符串的位置,评估模型利用其完整上下文窗口的能力。
* (i) Retrieve.KV: 模型需从一个包含大量键值对的大型JSON对象中,为指定的键准确检索其值。随机的KV对使得输入难以压缩,要求严格的O(n)空间存储,这对评估长上下文方法的内存模糊性特别有用。每个会话检索五个KV对,目标KV均匀分布在输入中。
* (ii) Retrieve.Prefix-Suffix: 模型需从一个包含大量可变长度字符串的列表中,检索一个与指定前缀和后缀都匹配的字符串。此任务具有挑战性,需要类似前缀树的复杂功能,计算成本为O(P $w_i^2$)。仅共享前缀或后缀的干扰项使模型无法依赖简单的查找或归纳头(induction heads)【索引70,In-context learning and induction heads,2022,arXiv】进行有效检索。
* (iii) Retrieve.MultiHop: 此任务最初在RULER【索引38,RULER: What’s the real context size of your long-context language models?,2024,CoLM】中引入,评估LLM在长输入提示中的多跳追踪能力。模型必须跟踪并回忆关键信息的变化。上下文中嵌入了五个多跳变量赋值链,每个测试轮次要求检索确切的多跳链。

语义检索。许多现实世界的长上下文应用要求模型具备语义理解能力。SCBench包含了评估此类能力的必要任务,因为有损的长上下文方法在多请求场景中通常难以抽象或理解信息。
* (i) Code.RepoQA: 模型需要根据精确的自然语言描述,从一个长源代码块中检索特定函数(包括名称、参数和完整实现)。与原始RepoQA基准【索引57,Repoqa: Evaluating long context code understanding,2024c,arXiv】不同,本文的输入扩展到64K token,目标函数在代码库中位置均匀分布。函数描述由GPT-4生成。代码库和编程语言也扩展到Python、C++、Java、PHP、Rust、Go和TypeScript。每个测试会话包含一个GitHub仓库,模型需要在5轮中每轮检索一个函数。
* (ii) http://En.QA, http://Zh.QA, En.MultiChoice: 这些任务扩展自InfiniteBench【索引105,Infinitebench: Extending long context evaluation beyond 100K tokens,2024a,ACL】,提供基于虚构小说的、由人工标注的高质量QA测试,以消除外部知识的影响。任务要求模型从长输入中定位和处理信息,通过聚合或过滤进行推理。主要问题类型有两种:1) 聚合,汇编输入中分散的信息;2) 过滤,从更大的集合中识别特定细节。在SCBench中,共享相同输入上下文的QA对被组合成共享上下文的测试会话。

全局信息处理。除了检索,一些长上下文任务需要利用和聚合全局上下文信息,如摘要、统计任务和上下文学习(ICL)。本文的基准包含三个任务来评估不同长上下文方法在多请求场景下处理全局信息的能力。
* (i) Many-shot ICL: 使用来自Big-Bench Hard【索引88,Beyond the imitation game: Quantifying and extrapolating the capabilities of language models,2023,TMLR】的数据集来评估多示例ICL能力。包括三个子任务:日期理解、显著错误翻译检测和跟踪七个打乱的对象。多示例ICL上下文在测试会话的各轮次中共享,所有子任务都以四选一的多项选择题形式呈现。
* (ii) Math.Find: 扩展了InfiniteBench【索引105,Infinitebench: Extending long context evaluation beyond 100K tokens,2024a,ACL】中的数学查找任务,从仅查找最大值扩展到多个统计值(如最小值或中位数)。这要求模型有效理解全局上下文、进行比较和统计操作。
* (iii) En.Sum: 使用来自arXiv的串联学术论文作为输入,文档长度从8K到20K token不等。基准摘要(平均654个token)由GPT-4生成,它被提示为每个文档生成简洁的单句摘要。每轮的目标文档在整个上下文长度中均匀分布。

多任务处理。在现实应用中,LLM通常在单个会话中使用共享输入上下文处理多个任务。为了反映这一点,SCBench包含了两个多任务处理任务:
* (i) Mix.Sum+NIAH: 此任务将文档摘要与大海捞针(Needle in a Haystack)【索引46,Needle in a haystack - pressure testing llms,2023,】任务结合,使用共享输入提示。一个随机的“针”被均匀地插入到En.Sum任务(串联的学术论文)的输入中。模型在每个测试会话中交替执行摘要和NIAH检索。
* (ii) Mix.RepoQA+KV: 此任务将RepoQA任务与KV检索结合,使用共享输入提示。多个KV对被均匀地插入到RepoQA的输入(长段源代码)中,包括100个KV对(4个目标KV和其余为干扰项)。模型在每个测试会话中交替执行RepoQA和KV检索。

长上下文共享模式细节

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

  • (i) 多轮模式 (Multi-turn Mode):典型的长上下文应用场景包括长上下文聊天、多步推理(如思维树【索引102,Tree of thoughts: Deliberate problem solving with large language models,2024b,NeurIPS】)和长生成链式思考(CoT)。此模式对于使用KV缓存复用的长上下文方法至关重要,因为轮次间的焦点转移可能导致KV缓存中的信息丢失。我们遵循先前工作【索引107,Judging llm-as-a-judge with mt-bench and chatbot arena,2023a,NeurIPS】【索引98,MINT: Evaluating LLMs in multi-turn interaction with tools and language feedback,2024,ICLR】的做法,使用标准答案而非模型生成的内容作为后续轮次的上下文。

  • (ii) 多请求模式 (Multi-request Mode):上下文共享跨越会话或用户,例如协作者在共享代码库上工作。模型可以编码共享上下文,并在不同请求间复用KV缓存。评估长上下文方法在这种情况下的表现至关重要,因为一些方法依赖于查询来进行稀疏编码/解码。例如,MInference和SnapKV使用输入的最后部分(通常是查询)来估计稀疏模式,本模式可以评估它们在没有查询访问权限时的泛化能力。

A4 实验环境

  • 模型:

    • 开源长上下文LLMs:Llama-3.1-8B/70B【索引21,The llama 3 herd of models,2024,arXiv】,Qwen2.5-72B/32B【索引93,Qwen2.5: A party of foundation models,2024,】,Llama-3-8B-262K【索引29,Llama-3 8b instruct gradient 4194k (v0.1),2024,】,GLM-4-9B-1M【索引28,Chatglm: A family of large language models from glm-130b to glm-4 all tools,2024,arXiv】。
    • 门控线性模型:Codestal Mamba 7B【索引92,Codestral mamba,2024,】和Jamba-1.5-Mini【索引54,Jamba: A hybrid transformer-mamba language model,2024,arXiv】。
    • 模型类型覆盖:Transformer、SSMs和SSM-Attention混合模型。
  • 硬件配置:

    • GPU:4块NVIDIA A100 GPU。对于大于7B参数的模型,使用8台A100 40GB机器或4台H100 80GB机器进行张量并行。
  • 软件配置:

    • 解码策略:所有实验使用贪心解码(greedy decoding)。
    • 精度:BFloat16。
    • 推理框架:HuggingFace或vLLM-0.52。
    • 核心库/技术:
      • 使用FlashAttention-2【索引16,Flashattention-2: Faster attention with better parallelism and work partitioning,2024,ICLR】。
      • 使用MInference【索引41,MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention,2024,NeurIPS】来减少GPU内存开销。
      • 稀疏注意力测试中使用了基于PIT【索引109,Pit: Optimization of dynamic sparse deep learning models via permutation invariant transformation,2297,SOSP '23】和Triton【索引94,Triton: an intermediate language and compiler for tiled neural network computations,2019,MAPL】实现的自定义A-shape、Tri-shape和MInference核。
      • SSMs和Mamba-Attention混合模型使用了triton版本的mamba核和causal-conv1d-1.46。
      • 提示压缩使用了LLMLingua-2的官方实现。
  • 长上下文方法:

    • 评估了8类长上下文解决方案:门控线性RNN(Codestral-Mamba)、SSM-Attention混合模型(Jamba)、稀疏注意力、KV缓存丢弃、提示压缩、KV缓存量化、检索和加载。具体方法见表1。

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模型在共享上下文模式下表现挣扎。

结果分析

  • Sub-O(n)内存几乎不可行:图6b可视化了Retr.KV任务在多轮中的注意力图。虽然重要的KV对在单轮内保持稳定,但在不同查询之间变化显著。这解释了为什么O(k)的KV缓存压缩方法在单查询测试中表现良好,但在后续查询中失败。然而,SSM-Attention混合模型Jamba和CPU-GPU协同方法(将完整的O(n)内存存储在CPU RAM中,动态加载相关KV到GPU)显示出潜力。理论研究【索引34,Compression barriers for autoregressive transformers,2025,arXiv】也证明了需要线性内存来保持注意力的表征能力。

  • 编码与解码中的稀疏性:实验表明,如果解码保持密集(O(n)内存),稀疏编码方法(如Tri-Shape和A-Shape)在多请求场景中表现良好(图3a)。这一现象首次在共享上下文环境中得到验证。相反,将稀疏模式扩展到解码阶段会严重降低性能(如StreamingLLM)。即使编码是密集的,稀疏解码方法(特别是KV缓存压缩)在共享上下文场景中也表现不佳。这可能是因为编码输出存在冗余,而解码在生成中起着至关重要的作用。动态稀疏注意力(如MInference)比静态稀疏模式(图9)能更好地逼近全注意力的性能。

  • 可压缩与不可压缩任务:对于高度可压缩的输入(如大海捞针基准测试中的重复噪声),sub-O(n)方法可以取得合理的准确率。然而,对于动态和复杂的输入(如Retr.KV任务中的随机键值对),sub-O(n)方法通常无法保留所有必要信息,导致性能不佳。

  • 无查询感知的稀疏方法:在KV缓存复用场景中,一个关键问题是依赖查询进行压缩的方法在没有查询访问权限时能否有效泛化。表5比较了三种查询感知方法在有无查询情况下的性能。结果显示,KV缓存压缩方法SnapKV和静态稀疏注意力方法Tri-Shape在没有查询时准确率下降。相比之下,动态稀疏注意力方法MInference表现出更强的泛化能力。

A5 结论

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

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

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

A6 附录

相关工作

  • 前缀缓存 (Prefix Caching),也称为KV缓存复用,用于优化LLM推理框架中的首个令牌生成时间,尤其适用于多轮对话或聊天机器人会话等共享上下文场景【索引101,Cacheblend: Fast large language model serving with cached knowledge fusion,2024a,arXiv】【索引48,Efficient memory management for large language model serving with pagedattention,2023,SOSP '23】【索引27,Prompt cache: Modular attention reuse for low-latency inference,2024,MLSys】。该技术已被LLM服务提供商广泛采用。近期的优化工作集中于提升KV缓存效率,例如PagedAttention通过分块和查找表减少内存成本;HydraGen和Cascade Inference解耦共享前缀和独立后缀的注意力计算;RadixAttention使用基数树加速KV查找;RAGCache和CacheBlend则分别针对RAG和通过部分重计算来优化缓存利用率。尽管有这些进展,但尚无现有长上下文基准评估KV缓存复用场景。
  • 对话与多轮基准:尽管多轮基准更能反映真实应用,但许多评估仍侧重于单轮。MT-Bench【索引107,Judging llm-as-a-judge with mt-bench and chatbot arena,2023a,NeurIPS】、ShareGPT【索引20,Domeccleston/sharegpt: Easily share permanent links to chatgpt conversations with your friends,2023,】、MINT【索引98,MINT: Evaluating LLMs in multi-turn interaction with tools and language feedback,2024,ICLR】、MT-Bench-101【索引6,MT-bench-101: A fine-grained benchmark for evaluating large language models in multi-turn dialogues,2024a,ACL】和MT-Eval【索引47,MT-eval: A multi-turn capabilities evaluation benchmark for large language models,2024,EMNLP】等基准评估了对话能力、指令遵循和复杂任务解决能力,但主要关注模型一致性和信息提取,而非评估长上下文输入。
  • LLM的长上下文方法:长上下文推理面临两大瓶颈:预填充阶段的计算成本和解码阶段的内存成本。
    • 预填充优化包括:状态空间模型(SSM)、线性注意力方法、基于内存的方法、稀疏注意力和提示压缩等。
    • 解码优化主要关注:1) 注意力KV复用以减少存储;2) 静态KV压缩;3) 动态KV压缩,包括缓存丢弃和卸载;4) 分层推测解码。
      大多数方法都在单轮基准上测试,并采用依赖于查询的有损技术,这可能在带有前缀缓存的多轮场景中降低性能。这一局限性促使我们设计了SCBench,一个在共享上下文设置中评估长上下文解决方案的基准。

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

我们将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 长上下文方法细节

  • 状态空间模型 (SSMs):这类模型因其线性复杂性而特别适用于长序列任务,如S4【索引35,Liquid structural state-space models,2023,ICLR】和Mamba【索引30,Mamba: Linear-time sequence modeling with selective state spaces,2024,CoLM】。然而,SSM也因其记忆能力较弱和在复制粘贴任务上的局限性而受到批评。
  • Mamba-Attention混合架构:通过交错使用Transformer和Mamba层,旨在结合两种架构的优点,即Transformer的表达能力和Mamba层的线性复杂性。Jamba【索引54,Jamba: A hybrid transformer-mamba language model,2024,arXiv】和Samba【索引79,Samba: Simple hybrid state space models for efficient unlimited context language modeling,2025,ICLR】是代表性工作。
  • 稀疏注意力:本文测试了三种方法:A-shape(每个token只关注初始token和局部token)、Tri-shape(在A-shape基础上增加了底部的密集注意力空间)和MInference【索引41,MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention,2024,NeurIPS】(一种动态稀疏注意力方法,其稀疏模式在运行时动态构建以更好地逼近全注意力)。
  • KV缓存压缩:例如,StreamingLLM【索引99,Efficient streaming language models with attention sinks,2024,ICLR】在解码步骤中使用固定大小的KV缓存,只保留初始和局部token的状态。SnapKV【索引53,SnapKV: LLM knows what you are looking for before generation,2024c,NeurIPS】引入观察窗口概念,选择被高度关注的top-K个KV,并移除其他KV。
  • KV缓存量化KIVI【索引61,KIVI: A tuning-free asymmetric 2bit quantization for KV cache,2024e,ICML】对Key张量采用逐通道量化,对Value张量采用逐token量化。本评估中使用2bit算法,组大小为32,残差长度为32。
  • KV缓存检索:大多数框架采用精确匹配算法。而CacheBlend【索引101,Cacheblend: Fast large language model serving with cached knowledge fusion,2024a,arXiv】等方法在输入与缓存中的某个请求在语义上足够相似时,就会检索KV缓存。
  • KV缓存加载:这类方法利用CPU RAM存储完整的KV缓存,每个token解码时仅将部分KV缓存加载到GPU。例如,Quest【索引91,QUEST: Queryaware sparsity for efficient long-context LLM inference,2024,ICML】以页为粒度估计Key的重要性,只加载topK重要的Key和对应的Value。Retrieval Attention【索引56,Retrievalattention: Accelerating longcontext llm inference via vector retrieval,2024b,arXiv】在CPU RAM上构建向量数据库以高效找到topK关键Key。
  • 提示压缩LLMLingua-2【索引73,LLMLingua-2: Data distillation for efficient and faithful task-agnostic prompt compression,2024,ACL】是一个监督模型,将评估单个token重要性视为一个token分类任务,可在许多任务上实现高达20倍的压缩,而性能损失很小。

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)相似的结论:密集解码方法通常优于稀疏解码方法,更鲁棒和动态的稀疏模式优于静态稀疏方法。但使用模型生成内容作为上下文确实显示出更低的总体准确率,这表明了错误传播现象,即先前查询的误导性答案会影响后续轮次的表现。

案例研究

  • En.Sum任务(表14):摘要质量与模型规模正相关。Llama-3.1-70B和Qwen2.5-72B提供了更全面、细致的摘要。在高效方法中,采用密集解码的稀疏编码方法(Tri-Shape、MInference)能更好地捕捉细节,而稀疏解码方法StreamingLLM则完全失败,输出了随机不连贯的内容。
  • Retr.Prefix-Suffix任务(表15):有趣的是,Mamba-Attention混合架构Jamba取得了最准确的性能。这很不寻常,因为该任务需要较大的时空复杂度,而Mamba层据报道在这方面表现不佳。相反,所有基于全注意力的LLM(Llama、Qwen系列)都在此任务中失败,它们通常只能记住可变长度的前缀,但无法复现整个字符串。
  • Mix.RepoQA + KV任务(表16):Llama-3.1-70B及其MInference变体都能准确检索KV值,但在复现Python函数时都引入了修改,未能精确复现基准函数,这可能归因于基础Llama模型本身的长上下文能力限制。
  • Retr.KV任务(表17):此案例比较了A-shape和Tri-shape。Tri-shape即使在第一轮也表现出色,有效保持了模型的指令遵循能力。相比之下,A-shape严重破坏了模型的指令遵循能力,导致输出不完整和错误。这凸显了Tri-shape在从一开始就保持任务结构和理解方面的优势。