SYNTOLOGY HomeExplorerAtlasCodeMethodologyAboutDevelopersFeedPricing
Paper · 2406.03482 · 2024

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

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 7 functions out of this paper's own repositories and ran 5 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

RepositoryRoleRan
amirzandieh/qjl canonical 5 of 7
FunctionStatusWhere it lives
count_score Ran amirzandieh/qjl/metrics.py
code served (permissive licence) · get_code("b349b79d9cc2934b")
normalize_answer Ran amirzandieh/qjl/metrics.py
code served (permissive licence) · get_code("e7e75981cb464788")
normalize_zh_answer Ran amirzandieh/qjl/metrics.py
code served (permissive licence) · get_code("8c5c581f9264c810")
parse_args Ran amirzandieh/qjl/eval_long_bench.py
code served (permissive licence) · get_code("ca066a4c088c8cf5")
repeat_kv_quant Ran amirzandieh/qjl/models/llama3_utils_qjl.py
code served (permissive licence) · get_code("19035055b71c9787")
greedy_generate Not yet run amirzandieh/qjl/plot_memory_peak.py
code served (permissive licence) · get_code("8bf52d5d7d55d38f")
greedy_generate Not yet run amirzandieh/qjl/plot_runtime.py
code served (permissive licence) · get_code("c541ae13a407a3ce")

Repositories linked to this paper

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

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

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2406.03482")
get_code_for_paper("2406.03482")
have("2406.03482")

Connect an agent — have() is free.