Research

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

← All research
Preprint Third-party arXiv·June 5, 2024

This is third-party work — not authored by or affiliated with DarcStar Technologies.

By Amir Zandieh, Majid Daliri, Insu Han

Research

Abstract

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 https://github.com/amirzandieh/QJL.

DarcStar commentary

Our perspective on this work and how it informs our research — not part of the original paper.

QJL is where this line starts, and it starts with a genuinely sharp observation: the hidden tax in KV-cache quantization isn't the quantized numbers, it's the bookkeeping — the scale and zero-point every block has to store in full precision, which quietly adds a bit or two per value. QJL removes the tax entirely. Rotate with a Johnson–Lindenstrauss transform, keep only the sign bit, and there are no constants left to store. Zero overhead, by construction rather than by tuning.

The detail we most respect is the same one its successors inherited: it takes the inner product seriously. A quantizer tuned for low reconstruction error will still bias the dot products attention actually consumes, so QJL uses an asymmetric estimator — quantize one side, leave the other a plain JL sketch — to get an unbiased estimate of the thing that matters. That's the correctness-in-the-seams instinct, and it's the seed that later work grew into a provably near-optimal quantizer.

It also does the part we always ask for: it measures the end task, not just the distortion — three-bit KV cache, a fivefold memory cut, accuracy intact across real models. Our standing caution about proxy metrics mostly doesn't bite here, because they checked.

If we're pushing anywhere, it's that the sign-bit-after-rotation trick is elegant precisely because it's extreme, and the interesting question — how gracefully it trades bits for quality across model scales — is the one the follow-on work took up. As a primitive, though, this is the clean one: a random rotation that makes a single bit carry an honest signal.