CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving
文章标题与作者/机构
- 标题: CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving
- 作者/机构: Yuhan Liu, Hanchen Li, Yihua Cheng, Siddhant Ray, Yuyang Huang, Qizheng Zhang*, Kuntai Du, Jiayi Yao, Shan Lu†, Ganesh Ananthanarayanan†, Michael Maire, Henry Hoffmann, Ari Holtzman, Junchen Jiang University of Chicago †Microsoft *Stanford University
A1 主要贡献
-
核心问题:大型语言模型(LLM)处理长上下文时,由于需要先处理完整个上下文才能生成响应,导致响应延迟很高。虽然通过重用上下文的键值缓存(KV cache)可以减少计算延迟,但如果KV cache需要从其他机器通过网络获取,其巨大的体积(可达数十GB)会引入严重的额外网络延迟,从而影响交互式用户体验。现有工作主要关注减少KV cache在GPU内存中的大小,而忽略了传输过程中的网络瓶颈。
-
研究目标:本文旨在设计一个快速的上下文加载模块CacheGen,以减少在LLM系统中获取和处理长上下文时的网络延迟。
-
创新点:
- 定制化的KV缓存编解码器:CacheGen设计了一种专门的张量编码器,利用KV缓存的分布特性,如相邻令牌间的局部性和不同层对量化损失的敏感度不同,将KV缓存编码为更紧凑的比特流表示。这种方法在显著节省带宽的同时,解码开销极小,并通过基于GPU的实现加速解码过程。
- 自适应的KV缓存流式传输:为了应对网络带宽的变化,CacheGen将长上下文的KV缓存分割成块,并为每个块预先编码多个不同压缩级别的版本。在传输时,CacheGen会根据实时网络状况,动态地为每个块选择合适的压缩级别(或回退到传输文本进行重计算),以在满足服务等级目标(SLO)的同时,最大化生成质量。
A3 背景知识/关键Observation/设计原则
2.1 大型语言模型基础
- Transformer模型工作流程:Transformer模型是大多数LLM服务的事实标准,通过两个阶段处理输入令牌序列并生成输出。
- 预填充(Prefill)阶段:注意力神经网络接收输入令牌,其中$L$个注意力层各自产生一个键(K)张量和一个值(V)张量。这些K和V张量共同构成了KV缓存,包含了LLM后续利用上下文所需的关键信息。
- 生成(解码)阶段:在此阶段,LLM利用KV缓存计算每对令牌间的注意力分数,构成注意力矩阵,并以自回归的方式生成输出令牌。为保证性能,占用大量内存的KV缓存通常在此阶段保留在GPU内存中。
- 时间到首个令牌(TTFT):主流模型中,预填充阶段的计算开销随输入长度超线性增长。由于必须在生成第一个输出令牌前完成预填充,其持续时间被称为TTFT。本文专注于在不改变解码过程的情况下减少预填充阶段的TTFT。
22.2 LLM输入中的上下文
- 上下文的作用:为解决LLM在面对模型自身知识库之外的问题时可能产生低质量或幻觉答案的问题,许多应用会在输入中补充额外的文本,即上下文,以利用模型的上下文学习能力生成高质量响应。
- 上下文的应用场景:上下文用途广泛,包括:(i) 补充领域知识文档以生成更优答案;(ii) 从代码库检索信息以回答编程问题或生成摘要;(iii) 使用角色描述作为上下文以生成符合角色性格的对话;(iv) 在少样本学习中提供问答对作为示例;(v) 在聊天应用中将对话历史作为后续输入的上下文。
- 长上下文和重用的普遍性:实践中,上下文通常很长(如法律文件、财务报告、代码文件、聊天记录等)且经常被重用以应对不同的用户输入。例如,同一份财报可用于回答关于摘要和收入来源的两个不同问题。
- 长上下文对性能的影响:长上下文虽然能提升响应质量,但也导致更高的预填充延迟和更长的TTFT。由于上下文经常被重用,通过缓存中间结果(即KV缓存)来避免重复的预填充计算,是减少TTFT的一个有前景的方案,这一方向近期已得到探索【索引23, 58, 82】。
3 隐藏的网络瓶颈
- KV缓存的加载需求:尽管重用长上下文的KV缓存能显著减少TTFT,但这要求KV缓存必须首先位于本地GPU内存中。然而在实践中,由于GPU内存有限,无法存储大量重复上下文的KV缓存,这些缓存可能需要从其他机器获取。例如,大型金融报告的KV缓存大小可达19GB,与LLM模型本身相当。当重用请求间隔数小时,或请求未命中同一GPU时,就需要将KV缓存异地加载或在机器间移动。
- KV缓存获取的延迟问题:从其他机器获取KV缓存会产生巨大的网络延迟,这一点尚未得到足够重视。虽然一些多GPU推理系统也传输KV缓存,但它们假设使用高速互连(如NVLink,带宽数百Gbps),网络延迟可忽略不计。然而,在普通云服务器之间(带宽通常为个位数Gbps),获取KV缓存的延迟可能与没有KV缓存时的预填充延迟相当甚至更长,如图2b所示。
图2:不同上下文加载方式如何影响网络延迟(传输上下文或KV缓存)和计算延迟(在上下文上运行注意力模块)。
- 本文方法:本文专注于减少获取KV缓存的网络延迟。为此,我们通过将其编码为更紧凑的比特流表示来压缩KV缓存(如图2c所示)。这与近期减少运行时GPU内存占用的工作(如从文本上下文中丢弃词元或量化KV缓存张量)有本质区别。那些技术需要保持KV缓存的张量形态,而我们的方法旨在减少传输时的大小,因此可以不保留张量形态。这两种方法是互补的,经近期工作缩减后的KV缓存仍可通过我们的方法进一步编码以减少网络延迟。
5.1 KV缓存的经验性洞察
本节重点介绍关于KV缓存值的三个观察,这些观察通过在Llama-7B和Llama-13B模型上使用LongChat数据集【索引90,LongChat: How Long Can Open-Source LLMs Truly Promise on Context Length?, 2023】进行实证验证,证明了其普遍性。
5.1.1 令牌间的局部性
- 洞察1:相邻令牌的K/V张量值更相似。在同一层和通道内,位置相近的令牌比位置较远的令牌具有更相似的K/V张量值。
- 实验验证:通过对比K/V张量的原始值分布与连续令牌间差值(delta)的分布,我们发现差值的分布远比原始值更集中于零。如图3所示,差值的方差比原始值的方差低2.4-2.9倍。这一局部性启发CacheGen编码差值而非原始值。
- 原因解释:这种令牌间的局部性可由Transformer的自注意力机制直观解释,该机制在数学上等同于基于前一个令牌的KV张量计算当前令牌的KV张量,这意味着相邻令牌的KV张量存在内在关联。
图3:原始值分布与差值分布的对比。我们用两个Llama模型和多种长上下文(§5.1)进行建模。为清晰起见,我们显示了绝对值。
5.1.2 层间的损失敏感性
- 洞察2:LLM输出质量对浅层KV缓存值的损失比对深层更敏感。
- 实验验证:如图4所示,对KV缓存的不同层组施加相同的数据损失(通过四舍五入实现),对模型早期(浅层)层施加损失会导致响应准确率显著下降,而对深层施加相同损失的影响则小得多。这一结果在测试的不同模型中保持一致。这一异构性表明KV缓存编码器应对不同层采用不同的压缩策略。
- 原因解释:直观上,KV缓存的深层提取的是更高级别的结构和知识,而浅层嵌入的是更原始的信息【索引119, 132】。因此,移除浅层缓存的精度所造成的信息损失可能会传播并影响后续层,从而妨碍模型掌握生成高质量响应所需的高级结构。
图4:对KV缓存的不同层施加数据损失对准确率有不同影响。(工作负载与图3相同)。
5.1.3 沿层、通道和令牌的分布
- 洞察3:按通道和层对KV缓存值进行分组的信息增益显著高于按令牌位置分组。
- 直观解释:这可以粗略地理解为,同一通道(或层)内的不同KV值比属于同一令牌位置的不同KV值更相似。一个可能的解释是,不同的通道或层捕获输入中的不同特征【索引49, 92】,例如主宾关系或形容词。而根据先前工作【索引49, 92】,后期的层捕获比早期层更抽象的语义信息。另一方面,在给定的层和通道内,不同令牌的KV值更相似,这可能是由于自注意力机制,其中每个令牌的KV值都源自所有前面的令牌。
- 实验验证:我们根据层、通道或令牌位置对KV缓存中的值进行分组,并计算每组的熵。如图5所示,按令牌位置分组对熵的减少远小于按通道或层分组。
图5:使用不同分组策略时的熵(每元素比特数)。(工作负载与图3相同。)
A2 方法细节
4 CacheGen:KV缓存编码与流式传输
- KV缓存流媒体模块的角色:为减少KV缓存的传输延迟,本文引入了一个名为KV缓存流媒体(KV cache streamer)的新模块,它承担三个角色:(1) 将给定的KV缓存编码成更紧凑的比特流表示(KV比特流),这可以离线完成;(2) 通过吞吐量可变的网络连接流式传输编码后的KV比特流;(3) 将接收到的KV比特流解码回KV缓存。
- 与上下文压缩技术的区别:CacheGen与近期通过丢弃不重要令牌来压缩长上下文的方法(如【索引72, 95, 153】)存在关键区别。这些技术旨在减少KV缓存在GPU内存中的运行时大小,以适应内存限制,因此必须保持KV缓存的大型浮点张量形状。相比之下,CacheGen旨在减少KV缓存的传输时大小以降低网络延迟,因此无需维持原始张量形状,可以将其编码为更紧凑的比特流,并根据网络带宽调整其表示。此外,CacheGen在处理特定查询之前决定压缩方案,无法利用生成阶段的信息。
- CacheGen设计概述:本文提出的CacheGen是KV缓存流媒体的具体设计。首先,CacheGen使用一个定制的KV缓存编解码器,通过利用KV缓存张量的多种分布特性(§5.1)来最小化KV比特流的大小,从而直接减少TTFT。其次,在动态带宽下流式传输KV比特流时,CacheGen会动态切换不同的编码级别或按需计算KV缓存,以便在给定的延迟期限内保持高响应质量。编码/解码过程的计算开销可忽略不计,并与网络传输流水线化,以最小化对端到端延迟的影响。
5.2 KV缓存编码
- 编码流程概述:受§5.1中洞察的启发,CacheGen的KV缓存编码器包含三个高级步骤:首先,计算相邻令牌的K和V张量之间的差值张量(delta tensors),利用令牌间的局部性(§5.1.1);其次,对差值张量的不同层应用不同级别的量化,利用异构的损失敏感性(§5.1.2);最后,使用无损算术编码器将量化后的差值张量编码为比特流,该编码器根据§5.1.3的洞察,分别压缩每个层和通道中的值。
-
与视频编码的类比与区别:这些步骤与视频编码相似,视频编码也计算相邻帧之间的差值、进行量化并通过算术编码进行编码【索引126,High Throughput CABAC Entropy Coding in HEVC,2012,IEEE Transactions on Circuits and Systems for Video Technology】。但直接应用现有视频编解码器效果不佳,因为它们仅针对自然视频内容中的像素值进行了优化。CacheGen的设计是基于对LLM生成的KV缓存的领域特定洞察(§5.1)。
-
基于变化的编码 (Change-based encoding):为了利用令牌间的局部性,我们将上下文分割成每组包含十个连续令牌的令牌组。如图6所示,在每组中,我们独立地压缩第一个令牌(称为锚点令牌)的KV张量,然后压缩并记录组内其他每个令牌相对于该锚点令牌的差值张量。这种方法类似于视频编码中的图像组(groups of pictures),但不同之处在于我们让组内所有令牌都参照同一个锚点令牌,而非连续令牌间的差值,这使得压缩和解压可以并行进行以节省时间。
图6:在一个令牌组内,CacheGen计算锚点令牌的KV张量与其余令牌的KV张量之间的差值张量。
-
分层量化 (Layer-wise quantization):在将令牌分组后,CacheGen使用量化来降低KV缓存中元素(浮点数)的精度,以便用更少的比特表示它们。与先前工作中对所有元素使用统一量化位数不同【索引120,High-throughput generative inference of large language models with a single gpu,2023,arXiv】,我们受异构损失敏感性(§5.1.2)的启发,对早期层的差值张量应用更保守的量化(即使用更多比特)。具体来说,我们将Transformer层分为三个层组(前1/3,中1/3,后1/3),并对每个层组的差值张量分别应用不同的量化仓大小,量化仓大小(即量化误差)从早期层组到后期层组逐渐增大。我们采用先前工作【索引48,Llm. int8 (): 8-bit matrix multiplication for transformers at scale,2022,arXiv】中常用于量化模型权重的向量式量化方法。对于锚点令牌的KV缓存,我们仍使用相对高精度的8位量化,因为它们占所有令牌的一小部分,但其精度会影响整个块中所有差值张量的分布。
-
算术编码 (Arithmetic coding):将KV缓存量化为离散符号后,CacheGen使用算术编码(AC)【索引135,Arithmetic Coding for Data Compression,1987,Commun. ACM】将上下文的差值张量和锚点张量无损压缩成比特流。为提高效率,AC需要准确、低熵的KV缓存元素概率分布。受§5.1.3中KV值分布观察的启发,我们按通道和层对KV值进行分组以获得概率分布。具体地,我们的KV编码器离线地为LLM生成的差值张量和锚点张量的每个通道-层组合剖析一个单独的概率分布,并对该LLM生成的所有KV缓存使用相同的分布。CacheGen使用修改版的AC库【索引101,Practical Full Resolution Learned Lossless Image Compression,2019,CVPR】并结合CUDA来加速编解码过程(§6)。实验(§7.5)表明,与使用单一全局符号分布的简单方法相比,我们的方法可将比特流大小减少高达53%。
5.3 KV缓存流式传输自适应
-
带宽波动带来的挑战:由于KV缓存的传输可能耗时数百毫秒到数秒,期间可用带宽可能发生波动。以固定编码级别流式传输可能会违反给定的服务等级目标(SLO)【索引33,Site Reliability Engineering: How Google Runs Production Systems,2016】。如图7所示,带宽下降可能导致传输延迟超出SLO。
-
CacheGen的工作流程:为应对带宽变化,CacheGen离线将一个上下文分割成多个连续令牌组成的上下文块(chunks),并使用KV缓存编码器将每个块编码成多个可独立解码的不同编码(量化)级别的比特流。获取上下文时,CacheGen逐个发送这些块,并为每个块选择一种流式传输配置:可以是某个编码级别,也可以是文本格式(让LLM重新计算K和V张量)。
- 自适应逻辑:CacheGen在流式传输KV缓存时调整每个块的配置,以将传输延迟控制在SLO内。如图7所示,当带宽下降时,CacheGen切换到发送文本上下文并重新计算KV;当带宽回升时,切换回发送尺寸较小的KV比特流。通过这种自适应逻辑(具体算法见§C.1),CacheGen能够满足SLO。
图7:展示CacheGen在带宽变化下的自适应逻辑的时间序列图。
- 独立编码与解码:为实现高效自适应,CacheGen首先计算整个上下文的KV缓存,然后沿令牌维度将其K和V张量分割成子张量,每个子张量对应一个块。接着,KV编码器使用不同的编码级别对每个块的K或V子张量进行编码。由于一个令牌的KV张量编码只依赖于自身及其与所在令牌组锚点令牌的差值(§5.2),只要块长度大于一个令牌组,每个块就可以独立于其他块进行编码而不影响压缩效率。因此,以不同编码级别发送的块可以被独立解码然后拼接以重建KV缓存。若一个块以文本格式发送,LLM会基于已接收并解码的前一个块的KV张量来计算其K和V张量。
- 对生成质量的影响:如果因带宽低而以较小尺寸的编码级别发送某个块,只会对该块产生较高的压缩损失,而不会影响其他块的压缩损失。但若带宽过低导致大多数块都无法以高编码级别发送,整体质量仍会受损。
- 上下文块的长度选择:块的长度选择需权衡两个因素:1) 编码后的KV比特流不应太大,否则无法及时响应带宽变化;2) 块也不应太小,否则若选择文本格式,将无法充分利用GPU的批处理能力来计算KV张量。基于这些考虑,我们在实验中凭经验选择1.5K令牌作为默认块长度。
- 流配置决策:CacheGen通过测量前一个块的吞吐量来估计当前带宽,并假设该吞吐量在剩余块的传输中保持不变,从而计算每种流配置的预期延迟。然后,它选择在SLO内的预期延迟下压缩损失最小的配置(即文本格式或最低编码级别)来发送下一个块。对于第一个块,若有先验网络吞D吐量知识则加以利用,否则从一个默认的中等编码级别开始。
- 处理多个并发请求:当多个请求在$T$秒内同时到达时,CacheGen将它们批处理并一同流式传输,最多批处理$B$个请求(GPU服务器能同时处理的最大数量)。每个请求被划分为相同大小的块。对于每个块索引$i$,CacheGen确定包含该块的请求数$N_i$,并使用为前一个块$i-1$测量的吞吐量,通过将单个请求的延迟乘以$N_i$来计算每种配置的预期延迟。在GPU服务器上,通过填充它们的KV缓存并将它们一起处理来实现请求的批处理。
6 实现
- 代码库与集成:CacheGen的实现包含约2000行Python代码和1000行CUDA核代码,基于PyTorch v2.0和CUDA 12.0。它通过两个接口与LLM进行交互:
calculate_kv(context) -> KVCache
用于获取上下文对应的KV缓存,generate_with_kv(KVCache) -> text
用于在跳过上下文预填充的情况下生成令牌。 - HuggingFace模型集成:我们在HuggingFace的transformers库【索引64】中用约500行Python代码实现了这两个接口。
calculate_kv
通过设置max_length = 0
和return_dict_in_generate = True
来让LLM只计算KV缓存。generate_with_kv
则通过在调用generate
函数时传递past_key_values
参数来实现。类似的集成也适用于FastChat【索引155】、llama.cpp【索引98】和GGML【索引57】等其他LLM库。 - LangChain框架集成:CacheGen已集成到流行的LLM应用框架LangChain【索引83】中,在
BaseLLM
模块的_generate
函数中被激活。它会检查当前上下文的KV缓存是否存在,若存在则调用generate_with_kv
,否则先调用calculate_kv
创建KV缓存。 - KV缓存管理:CacheGen实现了两个模块来管理KV缓存:
store_kv(LLM) -> {chunk_id: encoded_KV}
负责计算KV缓存、分块、编码,并将块ID到编码比特流的映射存储在存储服务器上;get_kv(chunk_id) -> encoded_KV
负责从存储服务器获取指定块的编码KV张量并传输到推理服务器。 - 性能优化:为了加速KV缓存的编解码,我们用CUDA实现了一个基于GPU的算术编码(AC)库【索引101】,其中每个CUDA线程负责从一个令牌的比特流中编解码KV缓存。此外,我们还将上下文块$i$的传输与上下文块$i-1$的解码过程进行了流水线化处理。
A4 实验环境与结果
实验环境
- 模型:
- Mistral-7B、Llama-34B、Llama-70B 的微调版本,均支持长达32K的长上下文。
-
数据集:共包含4个数据集的662个上下文(见表2)。
- LongChat【索引90】:测试LLM对长对话历史的理解能力,上下文长度约9.2K-9.6K。
- TriviaQA【索引29】:测试LLM的阅读理解能力,基于单个文档回答问题。
- NarrativeQA【索引81】:测试LLM基于故事或剧本回答问题的能力。
- Wikitext【索引102】:基于相关文档上下文预测下一个令牌的概率。
表2:评估中使用的数据集的大小和上下文长度。 -
硬件配置:
- 一台配备4个NVIDIA A40 GPU的服务器。
- 384GB内存。
- 两颗Intel(R) Xeon(R) Gold 6130 CPU。
- 软件配置:
- 基于PyTorch v2.0和CUDA 12.0实现。
- 使用vLLM【索引82】作为“文本上下文”基线的推理引擎。
实验结果
-
总体性能提升(TTFT与KV缓存大小):
- TTFT减少:在3Gbps带宽下,与从文本加载上下文相比,CacheGen将TTFT(首个令牌生成时间)减少了3.1-4.7倍;与默认量化基线相比,减少了3.2-3.7倍。即使与近乎无损的8位量化相比,CacheGen仍能将加载延迟减少1.67-1.81倍(图8)。
- KV缓存大小缩减:在达到相似下游任务性能的前提下,与默认量化基线相比,CacheGen的KV编码器将KV缓存大小减少了3.5-4.3倍。这种有损压缩对质量的影响极小,准确率下降不超过2%,F1分数下降小于0.1%,困惑度增加小于0.1(图9)。
图8:首个令牌生成时间(TTFT):在不同模型和数据集上,CacheGen在对质量(准确率、困惑度或F1分数)影响很小的情况下减少了TTFT。
图9:减少KV缓存大小:在各种模型上,CacheGen在各种数据集上以很小的准确率下降减少了KV缓存的大小。 -
在上下文压缩基线上的增益:
- 将CacheGen应用于现有的上下文压缩方法(如H2O【索引153】和LLMlingua【索引72】)的KV缓存上,可以进一步压缩。与对这些基线压缩后的KV缓存进行量化相比,CacheGen能将KV缓存大小再减少3.3-4.2倍,且不损失质量(图10)。这表明即使在上下文被压缩后,其KV缓存仍保留了CacheGen所利用的统计特性。
图10:在H2O【153】和LLMlingu【72】基础上减少KV缓存大小:在不同模型上,与H2O缩短的KV缓存相比,CacheGen进一步减小了KV缓存的大小,而在不同数据集上的准确率下降很小。 -
敏感性分析:
- 不同带宽:在0.4Gbps至400Gbps的广泛带宽范围内,CacheGen始终优于基线。虽然在高带宽(>20Gbps)下,其相对于量化基线的绝对TTFT减少量变小,但优势依然存在(图11)。
- 并发请求数:随着并发请求数的增加,CacheGen相比基线能更显著地减少TTFT,因为它避免了大量预填充计算(图12左)。
- 上下文长度:对于长上下文,CacheGen的增益主要来自减小KV缓存大小。对于短上下文(<1K),CacheGen会自动切换到加载文本,因为此时后者TTFT更低(图12右)。
图11:CacheGen在各种不同带宽下改善了TTFT。使用Mistral-7B绘制。y轴为对数尺度。
图12:当单个GPU上有多个并发请求时,CacheGen持续减少TTFT。使用Mistral-7B绘制。 -
流式传输自适应效果:
- 在带宽随机波动的场景下(0.1-10Gbps),CacheGen的自适应逻辑显著优于无自适应的CacheGen和量化基线。在0.5s的TTFT SLO下,CacheGen在达到与量化基线相同质量的同时,SLO违规率降低了60%。在1s的SLO下,SLO违规率从81%降至8%(图13)。
图13:CacheGen相比无自适应的CacheGen和量化基线降低了SLO违规率。使用Mistral-7B模型绘制。 -
开销和微基准测试:
- 解码开销:通过GPU加速和流水线化,解码开销对端到端延迟的影响极小(图14a)。与从文本生成KV缓存相比,解码所需的计算量可以忽略不计。
- 离线编码和存储开销:尽管CacheGen需要为每个上下文存储多个压缩版本,但其总存储成本与量化基线相当,离线编码速度也很快(图14c, 14d)。
- 消融研究:逐一添加CacheGen KV编码器的各个组件(定制AC、基于变化的编码、分层量化)后,性能持续提升,证明了每个组件的有效性。尤其是,不保留张量格式而采用比特流编码,能显著减小量化后的KV缓存大小(图15)。
- 用户体验:IRB批准的用户研究表明,在270份评分中,用户一致认为CacheGen因其更短的TTFT而提供了优于其他基线的体验质量(QoE)(图16)。
图14:(a) 文本上下文、量化基线和CacheGen的TTFT分解。(b) 文本基线和CacheGen的计算开销。(c) 基线量化和CacheGen的离线延迟分解。(d) CacheGen、量化基线和未压缩KV缓存的存储成本。使用Mistral-7B绘制。
图15:KV编码器背后各个思想的贡献:基于变化的编码、分层量化和基于通道-层分组的AC。
图16:真实用户研究表明,CacheGen相比其他基线显著改善了体验质量。
A5 结论
本文提出了CacheGen,一个用于最小化LLM获取和处理上下文总延迟的上下文加载模块。CacheGen通过一个专为将KV缓存压缩成紧凑比特流而设计的编码器,减少了传输长上下文KV缓存所需的带宽。跨三种不同容量模型和四个不同上下文长度数据集的实验表明,CacheGen在保持高任务性能的同时,有效降低了总体延迟。
局限性与未来工作:
- 与其他压缩工作兼容:CacheGen与智能量化【索引62, 78, 97】等新兴技术是互补的,未来可以结合使用。
- 增量式KV缓存流式传输:未来工作包括将CacheGen扩展为增量式流传输,类似于可伸缩视频编码(SVC)。
- 真实世界应用中的上下文重用:需要更多业界数据集来验证上下文重用的普遍性。
- 硬件与模型扩展:未来将在更高端的GPU和更大规模的模型(如OPT-175B)上评估CacheGen。
- 其他系统设计:未来可将CacheGen与并发工作中关于KV缓存存储、缓存策略等方面的研究相结合。
- 其他限制:未在自由文本生成任务上进行广泛评估;网络模型未包含极高带宽情况;并非所有LLM应用(如实时搜索)都能缓存KV特征。
A6 附录
A CacheGen的文本输出示例
- 示例描述:图17展示了一个来自LongChat数据集【索引90】的例子。LLM的上下文是一段长多轮对话历史,其中第一个话题是“艺术在社会中的作用”。当被问及“我们讨论的第一个话题是什么?”时,CacheGen正确地生成了答案,而压缩后KV缓存大小相似的默认量化基线则生成了错误答案。
图17:CacheGen在使用LongChat-7b-16k模型处理LongChat数据集时的一个输出示例。
B CacheGen与更具侵入性的方法的比较
- 与修改LLM或上下文的方法对比:图18将CacheGen与改变LLM或上下文的近期方法进行了对比。
- 更小的模型:图18a显示,用更小的Llama-3B模型替换Llama-7B,即使应用了不同级别的量化,其TTFT仍高于CacheGen,因为Transformer操作减慢了其速度。
- 令牌选择:图18b使用Scissorhands【索引96】作为示例,它移除低自注意力分数的令牌。由于自注意力分数在实际生成时才可用,无法减少TTFT。我们创建了一个理想化的Scissorhands*版本,离线确定要丢弃的令牌。结果显示CacheGen能更好地缩减KV缓存。
- Gisting:图18c测试了Gisting【索引104】,它将上下文缩短为gist令牌并修改LLM以接受它们。在PIQA数据集【索引34】上,CacheGen在原始Llama-7B模型上的表现优于基于Llama-7B的预训练gisting模型,因为它能将KV特征压缩为更紧凑的比特流表示。
- 结论:CacheGen在实现相似或更好性能的同时,比这些基线更能减少TTFT或KV缓存大小。重要的是,CacheGen不依赖于特定的上下文或模型,因此可以与这些方法结合使用以可能进一步提升性能。
图18:比较CacheGen和更具侵入性的方法,包括更小的模型、令牌丢弃(左)、上下文选择(中)和gisting(右)。
C CacheGen系统设置
C.1 KV流媒体自适应逻辑
- 伪代码:算法1展示了KV流媒体模块适应带宽的逻辑伪代码。
算法1:CacheGen流式传输适配器逻辑
C.2 默认编码级别
- 默认参数:默认情况下,CacheGen编码使用以下参数:将LLM中的层划分为三个等距的组,量化仓大小分别设置为0.5、1和1.5。
D CacheGen在各种工作负载下的改进
- 性能热图:图19展示了CacheGen相对于最佳基线(量化或文本上下文)在由GPU可用周期(并发请求数$B$的倒数$1/B$)和可用带宽(对数尺度)两个维度定义的完整工作负载空间中的改进。图中较亮的单元格表示TTFT减少得更多。图11和图12可视为此图的水平/垂直横截面。
图19:热图显示了CacheGen在完整工作负载空间中相对于最佳基线的改进。更亮的单元格意味着TTFT减少更多。
E 存储KV缓存的成本
- 经济性分析:对于一个8.5K令牌的上下文,在Llama-13B上使用CacheGen压缩不同版本大约需要5GB存储空间,每月在AWS上的存储成本为$0.05【索引6】。而每次从文本重新计算KV缓存的成本至少为$0.00085(仅输入)【索引4, 5, 11, 12】。如果每月有超过150个请求重用此上下文,CacheGen还能降低推理成本。此计算仅为粗略估计,突显了CacheGen的潜在经济效益。
💬 评论讨论
欢迎在这里分享您的想法和见解!