KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache

📄 arXiv: 2402.02750v2 📥 PDF

作者: Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, Xia Hu

分类: cs.CL, cs.LG, cs.PF

发布日期: 2024-02-05 (更新: 2024-07-25)

备注: ICML2024

DOI: 10.13140/RG.2.2.28167.37282

🔗 代码/项目: GITHUB

💡 一句话要点

提出KIVI以解决KV缓存量化问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: KV缓存 量化 大型语言模型 推理优化 内存管理 算法设计 性能提升

📋 核心要点

  1. 现有方法在处理大型语言模型时，KV缓存的内存需求显著增加，成为推理速度的瓶颈。
  2. 本文提出了一种无调优的2位KV缓存量化算法KIVI，针对键缓存和值缓存采用不同的量化策略。
  3. 实验结果表明，KIVI在保持模型质量的同时，减少了2.6倍的内存使用，并提升了2.35至3.47倍的推理吞吐量。

📝 摘要（中文）

高效服务大型语言模型（LLMs）需要批量处理多个请求，以降低每个请求的成本。然而，随着批量大小和上下文长度的增加，存储注意力键和值的键值（KV）缓存显著增加了内存需求，成为速度和内存使用的新瓶颈。为了解决这一问题，本文进行了KV缓存元素分布的深入研究，提出了一种无调优的2位KV缓存量化算法KIVI。KIVI能够在保持几乎相同质量的同时，减少2.6倍的峰值内存使用，支持更大的批量大小，显著提升推理吞吐量。

🔬 方法详解

问题定义：本文旨在解决大型语言模型推理过程中KV缓存的内存需求过高的问题。现有方法未能深入研究KV缓存的元素分布，导致量化效果不佳。

核心思路：通过对KV缓存元素分布的分析，提出针对键缓存采用通道量化、值缓存采用按标记量化的策略，从而优化内存使用。

技术框架：KIVI算法的整体架构包括对KV缓存的元素分布分析、量化策略的设计以及硬件友好的实现。主要模块包括数据收集、量化处理和性能评估。

关键创新：KIVI的创新在于提出了无调优的2位量化方法，针对不同类型的缓存采用不同的量化策略，显著提高了内存利用率和推理速度。

关键设计：在KIVI中，键缓存按通道量化，值缓存按标记量化，确保了量化后的数据能够有效保留模型性能，同时在实现上兼顾了硬件的友好性。具体参数设置和损失函数设计未详细披露，属于未知领域。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

KIVI在实验中表现出色，能够将峰值内存使用减少2.6倍，同时支持最多4倍的批量大小，推理吞吐量提升2.35至3.47倍。这些结果表明KIVI在实际应用中的显著优势。

🎯 应用场景

KIVI算法在大型语言模型的推理过程中具有广泛的应用潜力，尤其是在需要处理大量请求的场景中。通过降低内存需求和提升推理速度，该算法可以有效支持更大规模的模型和更高的并发请求，为自然语言处理、对话系统等领域带来实际价值。

📄 摘要（原文）

Efficiently serving large language models (LLMs) requires batching of many requests to reduce the cost per request. Yet, with larger batch sizes and longer context lengths, the key-value (KV) cache, which stores attention keys and values to avoid re-computations, significantly increases memory demands and becomes the new bottleneck in speed and memory usage. Additionally, the loading of the KV cache causes the computational core to be idle, which limits the inference speed. A straightforward and effective solution to reduce KV cache size is quantization, which decreases the total bytes taken by KV cache. However, there is a lack of in-depth studies that explore the element distribution of KV cache to understand the hardness and limitation of KV cache quantization. To fill the gap, we conducted a comprehensive study on the element distribution in KV cache of popular LLMs. Our findings indicate that the key cache should be quantized per-channel, i.e., group elements along the channel dimension and quantize them together. In contrast, the value cache should be quantized per-token. From this analysis, we developed a tuning-free 2bit KV cache quantization algorithm named KIVI. With hardware-friendly implementation, KIVI can enable Llama, Falcon, and Mistral models to maintain almost the same quality while using $\mathbf{2.6\times}$ less peak memory (including model weight). This reduction in memory usage enables up to $\mathbf{4\times}$ larger batch size, bringing $\mathbf{2.35\times \sim 3.47\times}$ throughput on real LLM inference workload. The source code is available at https://github.com/jy-yuan/KIVI.