发表时间: 2023-10 · arXiv:2310.07240
原文: https://arxiv.org/abs/2310.07240
文章标题与作者/机构
一句话结论 本文提出了一个名为 CacheGen 的上下文加载模块,通过将大语言模型的 KV cache 压缩为紧凑的比特流并根据网络带宽动态流式传输,在几乎不损失生成质量的前提下,将长上下文的加载延迟和首字生成时间缩短了 3.1 至 4.7 倍。
要解决什么问题 现有的大语言模型在处理长上下文时,预填充阶段的计算开销极大,导致首字生成时间极长。为了加速这一过程,业界通常会复用历史请求的 KV cache。然而,由于单机 GPU 显存有限,系统往往需要跨机器拉取这些体积庞大(例如一份财报的 KV cache 可达 19GB)的缓存数据。在普通云服务器个位数 Gbps 的网络带宽下,传输几十 GB 张量所带来的网络延迟,甚至会超过直接从头计算纯文本的延迟。现有的上下文压缩技术大多只关注如何缩减 KV cache 在 GPU 显存中的运行时占用,要求必须保持浮点张量的形态,完全忽略了跨机传输时的网络带宽瓶颈。因此,如何打破网络传输延迟的卡点,让异地加载 KV cache 的总耗时真正降下来,是本文要解决的核心问题。
怎么做的 为了绕开网络带宽瓶颈,CacheGen 放弃了在传输阶段保持 KV cache 原始浮点张量形态的传统做法,转而将其编码为体积更小的比特流,核心由定制化编解码器和自适应流式传输两部分构成。首先,定制化编解码器利用了 KV cache 的三个统计特性来极致压缩数据。第一是相邻令牌间的局部性,即位置相近的令牌其 K 和 V 张量值高度相似。系统将上下文按每 10 个连续令牌分为一组,保留首个令牌作为锚点,其余令牌只计算并记录相对于锚点的差值张量(即 $\Delta = KV_{current} - KV_{anchor}$)。第二是层间的损失敏感性差异,即模型浅层对精度损失比深层更敏感。系统将 Transformer 层分为前、中、后三组,对差值张量实施分层量化,浅层使用较小的量化误差并保留更多比特,深层逐渐增大误差,而锚点令牌统一保留 8 位精度。第三是通道和层维度的分布聚集性,系统按通道和层对量化后的值进行分组,利用离线生成的概率分布,通过算术编码将差值和锚点无损压缩成比特流。其次,为了应对网络带宽的实时波动,CacheGen 采用了自适应流式传输机制。它在离线阶段将长上下文切分为多个长度约 1.5K 令牌的数据块,并为每个块预先生成多个不同压缩级别的比特流版本。在实际传输时,系统会根据上一个数据块的传输吞吐量来估算当前带宽,动态为下一个块选择能满足延迟服务等级目标且压缩损失最小的编码级别;如果网络极差,系统会直接回退到传输纯文本,让目标 GPU 利用已接收的前置块重新计算当前块的 KV cache。这种设计使得各个数据块的编解码相互独立,既能极限压缩传输体积,又能灵活兜底,从而彻底打破了跨机加载的网络卡点。
效果如何 实验在配备 4 张 NVIDIA A40 GPU 的服务器上进行,测试了支持 32K 上下文的 Mistral-7B、Llama-34B 和 Llama-70B 模型,数据集涵盖长对话、阅读理解、故事问答和语言建模。对比基线主要包括代表重新计算路线的纯文本上下文基线(基于 vLLM 引擎)、代表常规压缩路线的默认均匀量化基线,以及代表上下文缩减路线的 H2O 和 LLMlingua 方法。量化结果显示,在 3Gbps 带宽下,CacheGen 将首字生成时间相比纯文本基线缩短了 3.1 至 4.7 倍,相比默认量化基线缩短了 3.2 至 3.7 倍。在保持下游任务准确率下降不超过 2%、F1 分数下降小于 0.1% 的前提下,它将 KV cache 的传输体积相比默认量化基线进一步缩小了 3.5 至 4.3 倍。即使叠加在 H2O 等上下文缩减方法之上,它依然能带来 3.3 至 4.2 倍的体积缩减。在 0.1Gbps 至 10Gbps 的随机波动带宽场景下,以 1 秒为延迟达标线,CacheGen 的自适应机制将超时违规率从默认量化基线的 81% 大幅降至 8%。该方法的代价是需要离线预计算并存储多个压缩版本的缓存,占用一定的额外存储空间。作者也承认了其局限性:目前尚未在自由文本生成任务上进行广泛评估,测试环境未覆盖极高带宽的网络场景,且对于实时搜索等无法预先缓存 KV 特征的应用场景并不适用。
核心问题:大型语言模型(LLM)处理长上下文时,由于需要先处理完整个上下文才能生成响应,导致响应延迟很高。虽然通过重用上下文的键值缓存(KV cache)可以减少计算延迟,但如果KV cache需要从其他机器通过网络获取,其巨大的体积(可达数十GB)会引入严重的额外网络延迟,从而影响交互式用户体验。现有工作主要关注减少KV cache在GPU内存中的大小,而忽略了传输过程中的网络瓶颈。
研究目标:本文旨在设计一个快速的上下文加载模块CacheGen,以减少在LLM系统中获取和处理长上下文时的网络延迟。
创新点:
图2:不同上下文加载方式如何影响网络延迟(传输上下文或KV缓存)和计算延迟(在上下文上运行注意力模块)。
本节重点介绍关于KV缓存值的三个观察,这些观察通过在Llama-7B和Llama-13B模型上使用LongChat数据集【索引90,LongChat: How Long Can Open-Source LLMs Truly Promise on Context Length?, 2023】进行实证验证,证明了其普遍性。
图3:原始值分布与差值分布的对比。我们用两个Llama模型和多种长上下文(§5.1)进行建模。为清晰起见,我们显示了绝对值。
图4:对KV缓存的不同层施加数据损失对准确率有不同影响。(工作负载与图3相同)。
图5:使用不同分组策略时的熵(每元素比特数)。(工作负载与图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%。
带宽波动带来的挑战:由于KV缓存的传输可能耗时数百毫秒到数秒,期间可用带宽可能发生波动。以固定编码级别流式传输可能会违反给定的服务等级目标(SLO)【索引33,Site Reliability Engineering: How Google Runs Production Systems,2016】。如图7所示,带宽下降可能导致传输延迟超出SLO。
CacheGen的工作流程:为应对带宽变化,CacheGen离线将一个上下文分割成多个连续令牌组成的上下文块(chunks),并使用KV缓存编码器将每个块编码成多个可独立解码的不同编码(量化)级别的比特流。获取上下文时,CacheGen逐个发送这些块,并为每个块选择一种流式传输配置:可以是某个编码级别,也可以是文本格式(让LLM重新计算K和V张量)。
图7:展示CacheGen在带宽变化下的自适应逻辑的时间序列图。
calculate_kv(context) -> KVCache用于获取上下文对应的KV缓存,generate_with_kv(KVCache) -> text用于在跳过上下文预填充的情况下生成令牌。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库。BaseLLM模块的_generate函数中被激活。它会检查当前上下文的KV缓存是否存在,若存在则调用generate_with_kv,否则先调用calculate_kv创建KV缓存。store_kv(LLM) -> {chunk_id: encoded_KV}负责计算KV缓存、分块、编码,并将块ID到编码比特流的映射存储在存储服务器上;get_kv(chunk_id) -> encoded_KV负责从存储服务器获取指定块的编码KV张量并传输到推理服务器。数据集:共包含4个数据集的662个上下文(见表2)。
表2:评估中使用的数据集的大小和上下文长度。
硬件配置:
总体性能提升(TTFT与KV缓存大小):
图8:首个令牌生成时间(TTFT):在不同模型和数据集上,CacheGen在对质量(准确率、困惑度或F1分数)影响很小的情况下减少了TTFT。
图9:减少KV缓存大小:在各种模型上,CacheGen在各种数据集上以很小的准确率下降减少了KV缓存的大小。
在上下文压缩基线上的增益:
图10:在H2O【153】和LLMlingu【72】基础上减少KV缓存大小:在不同模型上,与H2O缩短的KV缓存相比,CacheGen进一步减小了KV缓存的大小,而在不同数据集上的准确率下降很小。
敏感性分析:
图11:CacheGen在各种不同带宽下改善了TTFT。使用Mistral-7B绘制。y轴为对数尺度。
图12:当单个GPU上有多个并发请求时,CacheGen持续减少TTFT。使用Mistral-7B绘制。
流式传输自适应效果:
图13:CacheGen相比无自适应的CacheGen和量化基线降低了SLO违规率。使用Mistral-7B模型绘制。
开销和微基准测试:
图14:(a) 文本上下文、量化基线和CacheGen的TTFT分解。(b) 文本基线和CacheGen的计算开销。(c) 基线量化和CacheGen的离线延迟分解。(d) CacheGen、量化基线和未压缩KV缓存的存储成本。使用Mistral-7B绘制。
图15:KV编码器背后各个思想的贡献:基于变化的编码、分层量化和基于通道-层分组的AC。
图16:真实用户研究表明,CacheGen相比其他基线显著改善了体验质量。
本文提出了CacheGen,一个用于最小化LLM获取和处理上下文总延迟的上下文加载模块。CacheGen通过一个专为将KV缓存压缩成紧凑比特流而设计的编码器,减少了传输长上下文KV缓存所需的带宽。跨三种不同容量模型和四个不同上下文长度数据集的实验表明,CacheGen在保持高任务性能的同时,有效降低了总体延迟。
局限性与未来工作:
图17:CacheGen在使用LongChat-7b-16k模型处理LongChat数据集时的一个输出示例。
图18:比较CacheGen和更具侵入性的方法,包括更小的模型、令牌丢弃(左)、上下文选择(中)和gisting(右)。
算法1:CacheGen流式传输适配器逻辑
图19:热图显示了CacheGen在完整工作负载空间中相对于最佳基线的改进。更亮的单元格意味着TTFT减少更多。