MosaicKV动态压缩长文本模型的键值缓存,显著降低内存占用并提升推理速度。
MosaicKV: Serving Long-Context LLM with Dynamic Two-D KV Cache Compression

- 按缓存段动态选择压缩策略,区分重要与次要信息
- 在H800上实现7.3倍吞吐量提升,内存减少3倍
- 适合需要超长上下文的实时大模型服务场景
长上下文大模型服务需处理数十万至数百万个词元的提示,导致键值(KV)缓存成为主要服务开销。由于缓存随上下文长度线性增长,易耗尽GPU内存,迫使使用更小批量,降低吞吐量。现有压缩方法通常仅针对序列或通道维度,难以应对上下文窗口扩大带来的压力。同时压缩两个维度虽可大幅减内存,但直接应用会引发显著精度损失。本文提出MosaicKV,一种面向极长上下文推理的动态二维KV缓存压缩系统。该系统利用缓存元素重要性分布不均的特点,不对全局统一压缩,而是为每个KV向量识别关键元素,并以缓存段为粒度选择压缩策略。为应对细粒度稀疏与压缩管理带来的性能开销,MosaicKV引入压缩缓存管理机制,利用闲置的GPU与CPU资源维护压缩缓存并加速注意力计算。在H800 GPU上对多个LLM的评估显示,相比未压缩基线,MosaicKV实现最高16倍注意力加速、4.8倍更低解码延迟、7.3倍更高吞吐量;同时内存消耗降低3倍,且在LongBench和RULER上仅造成1.76%平均精度损失。
原文摘要 · Abstract (English)
Long-context LLM services now sustain prompts with hundreds of thousands to millions of tokens, making the key-value (KV) cache a first-order serving cost. Because the cache grows linearly with context length, it can exhaust GPU memory, force smaller batches, and reduce serving throughput. Prior KV cache compression techniques typically target only the sequence dimension or only the channel dimension, which leaves limited headroom as context windows scale. Compressing both dimensions promises higher memory reduction, but applying the two forms of compression directly leads to significant accuracy loss. This paper introduces MosaicKV, a dynamic two-D (dimensional) KV cache compression system for extremely long-context serving. MosaicKV uses dynamic two-D compression to address the accuracy challenge, exploiting the non-uniform importance distribution of elements within the KV cache. Instead of applying one compression pattern globally, MosaicKV identifies important elements for each KV vector and selects compression strategies at the granularity of KV cache segments. To address the performance challenge, where fine-grained sparsity and compression management overhead can offset the gains from compression, MosaicKV introduces compressed KV cache management. This mechanism uses underutilized GPU and CPU resources to maintain compressed KV caches and accelerate attention computation. Evaluation on an H800 GPU with multiple LLMs shows that MosaicKV delivers up to 16x attention speedup, 4.8x lower decode latency, and 7.3x higher throughput than the uncompressed baseline. At the same time, it reduces memory usage by 3x and incurs only 1.76% average accuracy loss on LongBench and RULER.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。