arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2406.03482cs.LGcs.AIcs.CLcs.PF

QJL:零开销的1比特量化JL变换用于KV缓存量化

QJL: 1-Bit Quantized JL Transform for KV Cache Quantization with Zero Overhead

New York University(纽约大学);Yale University(耶鲁大学);Adobe Research(奥多比研究机构)

Amir Zandieh, Majid Daliri, Insu Han

AI总结:

QJL通过JL变换加符号位量化消除量化常数存储开销,实现零开销KV缓存量化,在3比特下减少五倍内存且不损精度。

摘要:

由于KV缓存中键值(KV)嵌入的存储需求随序列长度增长,服务大型语言模型(LLM)需要大量内存。压缩KV缓存的有效方法是量化。然而,传统量化方法因需在每个数据块中以全精度存储量化常数(至少一个零点和缩放因子)而面临显著的内存开销。根据块大小,这种开销每个量化数值会增加1或2比特。我们提出QJL,一种新的量化方法,由Johnson-Lindenstrauss(JL)变换后接符号位量化组成。与现有方法不同,QJL通过消除存储量化常数的需求来消除内存开销。我们提出了一种非对称估计器用于两个向量的内积,并证明将一个向量应用QJL,另一个向量应用无量化的标准JL变换,可提供具有最小失真的无偏估计器。我们开发了QJL草图及其对应内积估计器的高效实现,并集成了轻量级CUDA内核以优化计算。当应用于各种LLM和NLP任务,将KV缓存量化至仅3比特时,QJL在不牺牲准确性的情况下,实现了KV缓存内存使用减少超过五倍,同时运行时间更快。代码可在https://github.com/amirzandieh/QJL获取。

英文摘要:

Serving LLMs requires substantial memory due to the storage requirements of Key-Value (KV) embeddings in the KV cache, which grows with sequence length. An effective approach to compress KV cache is quantization. However, traditional quantization methods face significant memory overhead due to the need to store quantization constants (at least a zero point and a scale) in full precision per data block. Depending on the block size, this overhead can add 1 or 2 bits per quantized number. We introduce QJL, a new quantization approach that consists of a Johnson-Lindenstrauss (JL) transform followed by sign-bit quantization. In contrast to existing methods, QJL eliminates memory overheads by removing the need for storing quantization constants. We propose an asymmetric estimator for the inner product of two vectors and demonstrate that applying QJL to one vector and a standard JL transform without quantization to the other provides an unbiased estimator with minimal distortion. We have developed an efficient implementation of the QJL sketch and its corresponding inner product estimator, incorporating a lightweight CUDA kernel for optimized computation. When applied across various LLMs and NLP tasks to quantize the KV cache to only 3 bits, QJL demonstrates a more than fivefold reduction in KV cache memory usage without compromising accuracy, all while achieving faster runtime. Codes are available at \url{https://github.com/amirzandieh/QJL}.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑