CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving

发表时间: 2023-10 · arXiv:2310.07240

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

文章标题与作者/机构
- 标题: 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


速读

一句话结论 本文提出了一个名为 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 特征的应用场景并不适用。

A1 主要贡献


A3 背景知识/关键Observation/设计原则

2.1 大型语言模型基础

22.2 LLM输入中的上下文

3 隐藏的网络瓶颈


图2:不同上下文加载方式如何影响网络延迟(传输上下文或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 令牌间的局部性


图3:原始值分布与差值分布的对比。我们用两个Llama模型和多种长上下文(§5.1)进行建模。为清晰起见,我们显示了绝对值。

5.1.2 层间的损失敏感性


图4:对KV缓存的不同层施加数据损失对准确率有不同影响。(工作负载与图3相同)。

5.1.3 沿层、通道和令牌的分布


图5:使用不同分组策略时的熵(每元素比特数)。(工作负载与图3相同。)


A2 方法细节

4 CacheGen:KV缓存编码与流式传输

5.2 KV缓存编码


图6:在一个令牌组内,CacheGen计算锚点令牌的KV张量与其余令牌的KV张量之间的差值张量。

5.3 KV缓存流式传输自适应


图7:展示CacheGen在带宽变化下的自适应逻辑的时间序列图。

6 实现


A4 实验环境与结果

实验环境

实验结果


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:CacheGen在使用LongChat-7b-16k模型处理LongChat数据集时的一个输出示例。

B CacheGen与更具侵入性的方法的比较


图18:比较CacheGen和更具侵入性的方法,包括更小的模型、令牌丢弃(左)、上下文选择(中)和gisting(右)。

C CacheGen系统设置

C.1 KV流媒体自适应逻辑

算法1:CacheGen流式传输适配器逻辑

C.2 默认编码级别

D CacheGen在各种工作负载下的改进


图19:热图显示了CacheGen在完整工作负载空间中相对于最佳基线的改进。更亮的单元格意味着TTFT减少更多。

E 存储KV缓存的成本