KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head
作者: Isaac Rehg
分类: cs.CL
发布日期: 2024-09-30 (更新: 2024-10-07)
💡 一句话要点
KV-Compress：一种基于分页KV缓存和变压缩率的注意力头压缩方法
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: KV缓存压缩 长上下文推理 大型语言模型 PagedAttention 注意力机制
📋 核心要点
- 现有长文本LLM推理中，KV缓存随上下文增长导致内存瓶颈，限制了并发请求数量。
- KV-Compress通过在PagedAttention框架内驱逐连续KV块，实现变压缩率的KV缓存压缩。
- 实验表明，KV-Compress在LongBench上达到SOTA，压缩率高达64倍，吞吐量提升高达5.18倍。
📝 摘要（中文）
近年来，大型语言模型（LLM）的上下文长度呈爆炸式增长，128k token的上下文已成为标准，百万token的上下文也已成为现实。有效支持长上下文推理仍然具有挑战性，因为生成过程中键值（KV）缓存中必须分配的内存随上下文长度而扩展，从而限制了在给定内存预算下可以同时服务的长上下文请求的数量。KV缓存压缩可以通过删除每个注意力头缓存中未充分利用的KV并减少其内存占用来缓解此问题。当移除的KV数量在不同注意力头之间变化时，可以实现更高的理论压缩率，但将这种策略应用于现有的推理框架会增加碎片，并且无法在物理内存中实现理论压缩率。我们引入了KV-Compress，这是一种新颖的压缩方法，可在PagedAttention框架内驱逐连续的KV块，从而使KV缓存的内存占用与该理论压缩率成比例地减少。我们的方法在Mistral-7B-Instruct-v0.2和Llama-3.1-8B-Instruct的LongBench上实现了最先进的性能，同时与先前的方法相比，压缩的KV总数减少了4倍。在Llama-3.1-8B-Instruct和Llama-3.1-70B-Instruct-FP8上的评估实现了高达8倍的压缩率，对性能的影响可忽略不计，并且高达64倍的压缩率，同时保留了超过90％的完整缓存性能，除了该套件的三个子集。我们对该方法与vLLM的集成进行了基准测试，通过启用更大的解码批次，总吞吐量提高了高达5.18倍。
🔬 方法详解
问题定义：大型语言模型在处理长上下文时，KV缓存会占用大量内存，限制了并发推理的规模。现有的KV缓存压缩方法无法有效利用不同注意力头之间的冗余差异，导致压缩率受限，并且容易产生内存碎片，难以在实际推理框架中应用。
核心思路：KV-Compress的核心思路是根据每个注意力头的实际需求，采用不同的压缩率，移除利用率低的KV块。通过PagedAttention框架管理KV缓存，确保移除的KV块是连续的，从而避免内存碎片，提高压缩效率。
技术框架：KV-Compress构建于PagedAttention框架之上，主要包含以下几个阶段：1. 分析每个注意力头的KV缓存利用率；2. 根据利用率确定每个注意力头的压缩率；3. 在PagedAttention框架内，驱逐连续的KV块，实现变压缩率的KV缓存压缩；4. 在推理过程中，根据需要从压缩的KV缓存中恢复KV值。
关键创新：KV-Compress的关键创新在于：1. 提出了一种变压缩率的KV缓存压缩方法，能够更有效地利用不同注意力头之间的冗余差异；2. 将变压缩率的压缩方法与PagedAttention框架相结合，避免了内存碎片问题，提高了压缩效率；3. 实现了在实际推理框架中的高效应用。
关键设计：KV-Compress的关键设计包括：1. 利用率评估方法：用于评估每个注意力头的KV缓存利用率，例如基于注意力权重的统计；2. 压缩率确定策略：根据利用率确定每个注意力头的压缩率，例如设定一个阈值，低于阈值的注意力头采用更高的压缩率；3. KV块驱逐策略：在PagedAttention框架内，选择连续的KV块进行驱逐，例如选择最不常用的KV块。
🖼️ 关键图片
📊 实验亮点
KV-Compress在LongBench测试中，针对Mistral-7B和Llama-3.1模型，实现了SOTA性能，同时压缩的KV数量减少了4倍。在Llama-3.1-8B和Llama-3.1-70B-FP8上，实现了高达8倍的压缩率，性能影响可忽略不计，最高可达64倍压缩率，同时保留了超过90%的完整缓存性能。与vLLM集成后，总吞吐量提升高达5.18倍。
🎯 应用场景
KV-Compress可广泛应用于需要处理长上下文的大型语言模型推理场景，例如机器翻译、文本摘要、问答系统等。通过降低KV缓存的内存占用，KV-Compress能够提高并发推理的规模，降低推理成本，并支持更长的上下文输入，从而提升LLM在各种应用中的性能。
📄 摘要（原文）
Context lengths of Large Language Models (LLMs) have exploded in recent years, with 128k-token context becoming a standard and million-token context becoming a reality. Efficiently supporting long-context inference remains challenging as the memory that must be allocated in key-value (KV) cache for a generation scales with its context length, limiting the number of long-context requests that can be served concurrently under a given memory budget. KV cache compression can mitigate this issue by removing under-utilized KVs from each attention head's cache and reducing its memory footprint. Higher theoretical compression rates can be achieved when the number of removed KVs varies across attention heads, but application of such a strategy within existing inference frameworks adds fragmentation and cannot realize the theoretical compression rates in physical memory. We introduce KV-Compress, a novel compression method that evicts contiguous KV blocks within a PagedAttention framework, reducing the memory footprint of the KV cache proportionally to this theoretical compression rate. Our method achieves state-of-the-art performance on LongBench for both Mistral-7B-Instruct-v0.2 and Llama-3.1-8B-Instruct while lowering the total number of compressed KVs by 4x compared with prior methods. Evaluations on Llama-3.1-8B-Instruct and Llama-3.1-70B-Instruct-FP8 achieve compression rates up to 8x with negligible impact on performance, and up to 64x while retaining over 90% of full-cache performance for all but three of the suite's subsets. We benchmark an integration of our method with vLLM that increases total throughput by up to 5.18x by enabling larger decoding batches.