通过可变压缩率优化注意力头的键值缓存,显著提升长文本生成效率。
KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head
- 在分页注意力框架中按注意力头差异压缩键值缓存块,减少内存占用。
- 在Llama-3.1模型上实现最高8倍压缩率,性能损失极小;64倍压缩仍保持90%以上性能。
- 适合需要高并发长文本推理的场景,尤其适用于vLLM等主流推理引擎。
近年来大语言模型上下文长度急剧增长,128k token已成为标准,百万token上下文也已实现。高效支持长上下文推理仍具挑战,因生成时需为每个请求分配与上下文长度成比例的键值(KV)缓存内存,限制了在固定内存预算下可并发处理的长上下文请求数量。通过移除各注意力头中低利用率的键值(KVs)来压缩KV缓存,可缓解此问题。若不同注意力头采用不同压缩率,理论上可获得更高压缩率,但现有推理框架中应用此类策略会引入碎片化,无法在物理内存中实现理论压缩率。我们提出KV-Compress,一种新型压缩方法,在分页注意力(PagedAttention)框架中连续地剔除各个注意力头的键值块,使内存占用与理论压缩率成比例降低。该方法在LongBench评测上对Mistral-7B-Instruct-v0.2和Llama-3.1-8B-Instruct均达到业界最优表现,且相比之前方法压缩的键值数量减少4倍。在Llama-3.1-8B-Instruct和Llama-3.1-70B-Instruct-FP8上,压缩率可达8倍,性能影响可忽略;在绝大多数数据子集上,压缩率达64倍时仍保持超过90%的全缓存性能。将该方法集成至vLLM后,最大吞吐量提升5.18倍,得益于支持更大的解码批次。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。