KV-Compress:每个注意力头采用可变压缩率的分页KV缓存压缩方法
KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head
AI总结:
针对大语言模型长上下文推理中KV缓存内存占用高的问题,提出KV-Compress分页KV缓存压缩方法,按注意力头可变压缩率驱逐连续KV块,实现高压缩率的同时保障性能,提升推理吞吐量。
摘要:
近年来,大语言模型(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倍的压缩率且对性能影响可忽略不计;在除3个子集外的所有测试子集中,压缩率最高可达64倍,同时保留了全缓存性能的90%以上。我们将该方法与vLLM集成后进行基准测试,结果显示通过支持更大的解码批次,总吞吐量最高提升了5.18倍。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。