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

研究背景

  1. LLM的广泛应用

  2. LLM使用特点

    prepend an LLM input with a long context containing thousands of tokens or more.

    • e.g.在输入中加入领域内知识，便于LLM利用以生成回应（无需LLM自身携带领域知识）。
    • e.g.用户要求 LLM 利用交互期间积累的对话历史记录来补充输入。

  3. 问题：答案生成的延时。

    Using long contexts poses a challenge to the response generation latency, as no response can be generated until the whole context is loaded and processed by the LLM.

相关工作

  • 增加处理上下文的吞吐量，即提高了处理长上下文的速度或效率，能够在单位时间内处理更多的数据或上下文。但是对于较长的上下文处理，延迟仍然长达数秒。

  • 通过存储、重用上下文的 KV cache ，在再次使用上下文时避免冗余计算。但是对应的 KV cache 不总是在本地GPU存储中，可能需要从其他机器调取，造成了额外的网络延迟。

  • 减小 GPU 内存中 KV 缓存的运行时大小，以适应内存限制或 LLM 的输入限制。

    • 删除不重要的 token 或上下文文本

    • 对 KV cache 缓存进行智能量化

      对 KV 缓存中的张量进行压缩或简化处理，将其表示为更紧凑的数据形式。

本文工作

针对研究背景中的问题部分，希望减小 KV cache 的传输时间大小以减少网络延迟。文章选择不保留其张量格式，而是将其编码为更紧凑的比特流。由此，提出了 CacheGen ，一个快速上下文加载模块，能够在抓取、处理长上下文时减少网络延迟。

CacheGen

  1. KV cache encoding and decoding

    • 自定义量化和算术编码策略：利用了 KV 缓存中的分布特性。

      KV 张量在相邻 token 之间的局部性（locality）：相邻的 token 可能有相似的 KV 值，因此可以压缩。

      对不同层量化损失的敏感度：KV 缓存的不同层对量化误差的敏感度不同，因此 CacheGen 根据不同层的特点进行差异化处理。

    • 解码和传输的优化

      • GPU 加速解码：在解码（即解压缩）阶段，CacheGen 通过 GPU 实现了加速，确保解码过程更快。

  2. KV cache streaming

    • CacheGen 以流式方式传输已经编码的 KV 缓存数据。流式传输意味着数据按顺序分块发送，而不是一次性传输整个数据包。这样可以更灵活地应对网络状况的变化。

    • 在用户查询到达之前，CacheGen 会将长上下文拆分成多个块（chunks），然后对每个块的 KV cache 进行独立编码。不同块可以使用不同的压缩级别进行编码，以便灵活适应不同的网络条件。这类似于视频流传输时，根据网络带宽动态调整视频的清晰度。

    • 发送上下文的 KV cache 时，CacheGen 会逐块传输这些编码好的数据。在传输每个块时，CacheGen 会根据当前的网络状况动态调整压缩级别。

    • 如果网络带宽太低，CacheGen 还可以选择将块以文本格式发送，而不是发送压缩后的 KV cache 。此时，LLM 会接收到文本数据并重新计算该块的 KV cache 。这种方式增加了计算开销，但可以应对极端的带宽限制。