Boost LLM Speed: How KV Cache Quantization Cuts Memory While Preserving Quality

Hugging Face added KV‑cache quantization to accelerate large language model (LLM) generation by compressing the key‑value (KV) cache, enabling longer contexts with lower memory while preserving generation quality.

What is KV cache quantization?

The KV cache stores attention keys and values of previously generated tokens so they can be reused without recomputation. Quantization reduces the numerical precision of these tensors (e.g., to 2‑ or 4‑bit formats), saving memory with only a small accuracy loss when applied carefully.

Memory‑savings example

For a 7 B Llama‑2 model processing 10 000 tokens, the KV cache size is roughly: 2 * 2 * 32 * 32 * 128 * 10000 ≈ 5 GB This is about one‑third of the memory required for the model parameters in fp16.

Implementation details

The method follows the KIVI paper (asymmetric 2‑bit quantization). Keys are quantized per channel, values per token, and a residual cache of recent entries is kept in full precision. When the residual cache reaches its limit (default length = 128), its contents are quantized and cleared. The affine quantization formula is: X_Q = round(X / S) - Z where S = (maxX - minX) / (maxVal - minVal) (scale) Z = round(-minX / S) (zero‑point)

Supported backends

quanto

– supports int2, int4, int8 HQQ – supports int2, int4, int8

Repositories:

https://github.com/huggingface/quanto

https://github.com/mobiusml/hqq/tree/master

Performance comparison

On the PG‑19 dataset with Llama‑2‑7B‑chat, int4 quantization (nbits=4, groupsize=64, residualLength=128, perToken=True) yields perplexity comparable to fp16, while int2 shows noticeable degradation. Reproduction scripts are available at:

https://gist.github.com/zucchini-nlp/a7b19ec32f8c402761d48f3736eac808

LongBench benchmarks confirm that Quanto int4 matches or slightly exceeds fp16 on tasks such as TREC and SAMSum, whereas int2 lags behind. Example scores (higher is better):

Dataset   fp16   Quanto int4   Quanto int2
TREC      63.0   63.0          55.0
SAMSum    41.12  41.3          14.04
TriviaQA  84.28  84.76         63.64
HotPotQA  30.08  30.04         17.3

Memory‑speed trade‑off

Quantizing to int4 reduces KV‑cache memory by ~2.5× but can slow generation as batch size grows. On an 80 GB A100 with FlashAttention, int4 enables contexts up to ~128 k tokens, compared to ~40 k tokens with fp16.

Enabling KV‑cache quantization in 🤗 Transformers

Install the backend and activate quantization: pip install quanto Example generation call:

model.generate(
    **inputs,
    cache_implementation="quantized",
    cache_config={"backend": "quanto", "nbits": 4}
)

The quanto backend works on CPU, GPU, and Apple MPS devices.

Additional optimizations

Combining KV‑cache quantization with weight quantization can further reduce memory but may increase latency (up to ~3×). Other techniques that alleviate the pre‑fill memory bottleneck include local‑window attention and FlashAttention. Relevant references: https://arxiv.org/abs/2004.05150 (local‑window attention) https://arxiv.org/abs/2307.08691 (FlashAttention)

https://hf.co/docs/transformers/main/en/perfinfergpuone#flashattention-2

Key takeaways

Memory‑speed balance : KV‑cache quantization dramatically cuts memory, enabling longer contexts, but may slow generation for large batches.

Accuracy retention : int4 quantization keeps model quality close to fp16; int2 can degrade performance.

Flexibility : Users can choose precision (int2, int4, int8) and configure residual cache length to suit workloads.

Future potential : KV‑cache quantization can be combined with other optimizations (weight quantization, advanced attention) for greater efficiency.