谈LLM的Context Cache功能
原创 孔某人 孔某人的低维认知 2024-06-24 16:48 中国香港
用的最多的开源框架vllm也是有这个优化的(叫prefix cache),然后这个优化除了共享system prompt的cache之外,还有个是多轮对话时的cache,比如用户说A,模型回B,用户再说C,AB的kv就不用重新计算了,对于对话应用来说这个省的比system prompt多得多~
1、Context Cache 功能简介
Context Cache简单来说是一个LLM推理时候的跨请求缓存相同前缀的计算结果的方案。由于KV Cache已经在LLM推理中普遍采用,所以作为其跨请求的自然推广,Context Cache这个设计是很自然的。
但在实用场景中,Context Cache的实用意义似乎没有那么强:
一方面这需要不同请求间共享足够长的prompt前缀,但在API平台上收到的实际请求中,不同用户不同请求能满足这个要求的比例很低,一般只有system prompt与prompt模版的前部才满足。在私有化部署场景下,硬件资源更少、请求的类型也更集中,所以Context Cache类的功能会更适合一些。
另一方面缓存的Cache大小又不小,显存很宝贵,把未来不知道何时会再来的请求缓存放在显存中是一种浪费。所以把cache转移到host memory是一个自然的选择,但这成本仍然很高,内存-显存带宽、host memory都不是白来的,要给每个请求都搞一下缓存还是成本比较高的。
目前来看,支持Context Cache的LLM API都是将其当作long context的降低成本和改善首token延迟的方案来设计的(Google Gemini和Moonshot都是如此)。
2、目前各家的实现方案
现在我能看到有4个Context的实现:
-
Google Gemini 1.5
-
NVIDIA TensorRT-LLM
-
Moonshot
-
SiliconFlow未公开宣传,但我在对其测试时发现的Context Cache功能
以下分别介绍一下特点:
2.1、Gemini 1.5
https://ai.google.dev/gemini-api/docs/caching?hl=zh-cn&lang=python
Cache功能的定价:https://ai.google.dev/pricing?hl=zh-cn
Gemini是目前唯一公开发布,且广泛可用的Context Cache功能。它需要用户显示创建Cache对象,并在请求时候显式指定要使用哪个cache对象。Cache对象有存储费用,按token和时间长度线性收费。在使用时,被cache部分的token计算费用减半。存储费用为 $4.5 / M token / h 。
值得一提的是Gemini不支持32k token以下的cache对象,就是为长输入场景准备的。
2.2、NVIDIA TensorRT-LLM
https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/kv_cache_reuse.md
TensorRT-LLM只是一个LLM的推理库,并不是一个LLM API服务。但由于这是NVIDIA提供的官方实现,所以本文还是提一下。不清楚除了这个公开方案外是否还有2B的优化方案。
TensorRT-LLM提供的Context Cache叫做KV cache reuse,支持把KV cache存储到host memory。但host memory缓存部分是被pin住的,不能交换到磁盘等,即只能占用物理内存。很明显这种方式很依赖host memory的大小,以及内存到缓存间的带宽。原文的描述:
Offloading to host memory increases likelihood of kv cache reuse. Reusable blocks that are needed for higher priority tasks, like propagating an already running request, are copied to a buffer in host memory instead of being evicted. This greatly extends the amount of memory available for reuse, allowing blocks to remain reusable much longer. On the other hand, offloading of blocks (and subsequent onboarding when a block is reused) has some cost since the blocks must be copied from CPU to GPU memory and vice versa. This cost is negligible on Grace-Hopper machines, and small enough to yield a net benefit for many use cases on x86 machines with Hopper GPUs. Offloading is unlikely to yield benefits on older architectures because of the (relatively) slow link between GPU and host memory.
将数据卸载到主机内存增加了 kv 缓存重用的可能性。对于高优先级任务(如传播已运行的请求)所需的可用重用块,会被复制到主机内存的缓冲区而非被逐出。这大大扩展了可重用内存的总量,使得块能够更长时间地保持可重用状态。另一方面,块的卸载(以及块重用时的再次加载)存在一定成本,因为块必须从 CPU 复制到 GPU 内存,反之亦然。在 Grace-Hopper 机器上,这一成本可以忽略不计,而在配备 Hopper GPU 的 x86 机器上,其成本小到足以在很多使用场景中带来净收益。由于 GPU 与主机内存之间的连接相对较慢,在旧架构上卸载不太可能带来好处。
此外还有个限制是,这个缓存只能在第一次请求完整的计算完成后才能生效,所以对于同时到来的多个相同前缀的缓存无法通过该方式复用计算。
2.3、Moonshot
https://platform.moonshot.cn/docs/api/cache
Moonshot已经公开了他们的Context Cache功能,但目前仍然处在内测中。
同Google Gemini一样,Moonshot的Context Cache需要手工创建cache对象,并管理存储时间。
Moonshot的Cache对象在使用时【限时】免费,未来定价未知。Cache对象创建时也会收费。需要注意的是其存储成本相对于Gemini相当高,达到 10RMB / M token / min,换算到Gemini的小时单位是:600RMB / M token / h 。
对比Gemini之后,其高昂的存储成本很难让人接受,最可能的解释是:Moonshot的技术方案跟TensorRT-LLM一样只能把cache保存在host memory中,而导致单机能存储的cache结果不多,限制了单机能承担的请求量,所以导致其均摊成本巨高。
相对来说,Moonshot的cache方案更适合那种短时间内会连续执行的workflow,稍微长一些的中间等待时间都可能导致费用爆炸。
2.4、SiliconFlow
在这里SiliconFlow是比较与众不同的,它提供了一种隐式的context cache功能,调用方无需任何配置,也无法做任何配置,命中了缓存也不会有费用的减免。只能通过相同请求的延迟时间差异来发现这个功能的存在。(是这个功能影响了我之前的long context测试才发现的)
它的实现可能很NB,也可能很naive只适合极低负载场景,单从我的测试结果无法推测。
3、个人评价
3.1、应用侧
目前来看,Context Cache更大的作用是在有公共的长context情况下缩短请求延迟,即使是Google Gemini,费用减免也颇为有限,好在Gemini的存储成本可以接受。
Moonshot的这个功能在减少延迟上也是一样的,但在费用上就感觉比较鸡肋了,实际中缓存只能存储一个较短的时间,如几分钟,否则可能还不如不用缓存直接请求。相对来讲更适合workflow执行、大量同context请求并发或者是用户回复很快的场景。
3.2、实现侧
有了两家对比之后更能看出Google Gemini所使用的工程架构和硬件的NB。能够做到很理想的长时间缓存,我估计是做了cache的remote存储,依托于云计算厂商的高性能内网,可以实现remote存储的性能接近于本地存储,且大幅降低成本、并且还不需要太受限于请求的路由效果。
相对来说Moonshot的定价就直接体现了其实现的高成本,目前大家还没有很好的压缩memory的方案,这时候看的是硬件能力,这方面Moonshot肯定没有优势,不知道未来阿里云是否能够为其解决这个问题。
我之前提的 LLM Decoder DSL 的设想 也是考虑到较大的cache数据的存储成本和存储策略较为困难,才提出的方案。虽然从使用的角度上来说Context Cache似乎对场景更普适(能够等待外部系统或者用户返回后继续执行)和对开发者更友好一点(真的更友好么?还是需要开发者自己管理context生存周期的),大家目前还是在朝着Context Cache的方向努力。但这种方案最终能够被用户接受需要能做的像gemini那样完善才行。
3.3、关于LLM的流式增量输入场景
对于延迟非常敏感的应用场景,我提过 展望LLM的流式输入增量计算能力 V2【2024Q1】 这个设计思路。目前来看,context cache类的方案似乎相对适合用来实现这个功能,这里仅提醒一下。