PQCache: Product Quantization-based KVCache for Long Context LLM Inference

随着大型语言模型（LLMs）领域的不断发展，推理中的上下文长度正在稳步增长。由于GPU内存有限，Key-Value Cache（KVCache）已成为LLM推理中的关键组件，现在已成为主要的内存瓶颈。目前的方法是选择适合LLMs自注意力计算的键和值来解决这个问题。然而，它们要么无法保持模型质量，要么导致高的服务延迟。我们从数据库社区使用的先进嵌入检索技术中获得启发，将KVCache的存储和搜索视为典型的嵌入检索问题。我们提出了PQCache，它采用Product Quantization（PQ）来管理KVCache，既保持模型质量，又确保低延迟。在预填充阶段，我们为每个LLM层和头部的标记键应用PQ。在自回归解码阶段，对于每个新生成的标记，我们首先使用PQ代码和质心通过最大内积搜索（MIPS）识别重要标记，然后获取相应的键值对进行自注意力计算。通过精心设计重叠和缓存，我们最小化了两个阶段中的任何额外计算和通信开销。广泛的实验表明，PQCache既有效又高效。即使只涉及到1/5的标记，它也可以保持模型质量，同时达到可接受的系统延迟。

解决在大语言模型推理中，由于GPU内存有限，导致关键-值缓存成为主要的内存瓶颈的问题。现有方法要么无法保持模型质量，要么导致服务延迟高。

使用产品量化（PQ）来管理关键-值缓存，将其视为典型的嵌入检索问题。在预填充阶段，对于每个LLM层和头，将PQ应用于令牌的键。在自回归解码阶段，通过使用PQ码和质心进行最大内积搜索（MIPS）来识别重要的令牌，然后获取相应的关键-值对进行自注意计算。

论文提出了PQCache，它在维持模型质量的同时确保低延迟。通过精心设计的重叠和缓存，最小化了两个阶段中任何额外的计算和通信开销。实验表明，PQCache即使只涉及1/5的令牌进行注意时，也能保持模型质量，并实现可接受的系统延迟。

最近的相关研究包括：1.《Key-Value Retrieval Networks for Task-Oriented Dialogue》；2.《Efficient Large-Scale Language Model Inference on GPUs with Approximated Key-Value Cache》；3.《Product Quantization for Nearest Neighbor Search》等。