arXiv:2412.03131cs.LGcs.DC2024-12被引 25

针对大模型缓存瓶颈,提出分层差异化压缩方案,显著提升推理效率。

DiffKV: Differentiated Memory Management for Large Language Models with Parallel KV Compaction

  • 按键值重要性、令牌价值和注意力头差异分层压缩缓存
  • 实现2.7到5.7倍压缩率,吞吐量提升1.9到5.4倍
  • 适合长推理、复杂思维链任务的高效部署场景

大语言模型(LLMs)虽具备强大能力,但其高昂的内存开销导致服务成本高,其中键值(KV)缓存是主要瓶颈。现有压缩技术如量化和剪枝对键值采用统一处理,直接丢弃不重要令牌,忽视了缓存组件间的细粒度差异。为此,我们提出DiffKV,一种新型高效的KV缓存压缩框架,利用三个层次的差异化:(1)键与值对注意力计算的影响不同;(2)不同令牌的重要性各异;(3)注意力头间存在多样的动态稀疏模式。这些差异导致不同请求和注意力头间内存使用不规则,带来显著可扩展性挑战。DiffKV设计了GPU上并行执行的内存管理器,将碎片化空闲内存列表高效合并为连续区域,有效将缓存稀疏性转化为性能收益。我们在多个主流LLM上评估了DiffKV,包括生成长思维链的新兴思考型模型。结果表明,在复杂推理与长序列生成任务中,该方法实现2.7×至5.7×的缓存压缩率,精度近乎无损,吞吐量提升1.9×至5.4×。源码已开源于https://github.com/zyqCSL/DiffKV。

原文摘要 · Abstract (English)

Large language models (LLMs) demonstrate remarkable capabilities but face substantial serving costs due to their high memory demands, with the key-value (KV) cache being a primary bottleneck. State-of-the-art KV cache compression techniques, such as quantization and pruning, apply uniform treatment to both keys and values, and discard unimportant tokens entirely, overlooking the fine-grained distinctions in the significance of individual KV cache components. To address such limitations, we introduce \textit{DiffKV}, a novel framework for efficient KV cache compression that exploits three levels of differentiation in the KV cache: (1) the differing impact of keys and values on attention computation, (2) the varying importance of tokens, and (3) the diverse dynamic sparsity patterns across attention heads. These levels of differentiation introduce irregular memory usage patterns across different requests and attention heads, posing significant scalability challenges for memory management. To address these challenges, DiffKV proposes an on-GPU memory manager that compacts fragmented free memory list into contiguous regions in parallel, effectively translating sparsity in the KV cache into performance gains. We evaluate DiffKV on several mainstream LLMs, including the emerging thinking models that generate extended chains of thought. DiffKV is able to compress the KV cache by $2.7\times$ to $5.7\times$ with near-lossless accuracy on complex workloads requiring sophisticated reasoning and long-generation capabilities, and enhances throughput by $1.9\times$ to $5.4\times$. Source codes of DiffKV are available at https://github.com/zyqCSL/DiffKV.

大模型优化缓存压缩推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。