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】)。
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】)。
KV缓存检索 (KV Cache Retrieval):根据请求的前缀从存储池中检索相关的KV缓存块,以减少首个令牌生成时间(TTFT)。方法包括语义检索方法,如CacheBlend【索引101,Cacheblend: Fast large language model serving with cached knowledge fusion,2024a,arXiv】。
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】。
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】也探索了类似的模式来加速长上下文预填充。
(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在长输入提示中的多跳追踪能力。模型必须跟踪并回忆关键信息的变化。上下文中嵌入了五个多跳变量赋值链,每个测试轮次要求检索确切的多跳链。
(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】中的数学查找任务,从仅查找最大值扩展到多个统计值(如最小值或中位数)。这要求模型有效理解全局上下文、进行比较和统计操作。
(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检索。
(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】。
使用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核。
前缀缓存 (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】等基准评估了对话能力、指令遵循和复杂任务解决能力,但主要关注模型一致性和信息提取,而非评估长上下文输入。
状态空间模型 (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倍的压缩,而性能损失很小。