CacheForge

Experiments with low-bit KV-cache quantization for Llama and Mistral inference, built on the KIVI implementation.

Setup · Quantization code · LongBench workflow · Upstream paper

Contribution scope

KIVI supplies the asymmetric quantization method, model integrations, and baseline CUDA kernels. This independently maintained derivative includes focused changes to:

• Combine min/max reductions and build quantization scales in place.
• Keep Triton outputs and launches on the input tensor's device.
• Simplify packing allocations and evaluation/answer-normalization utilities.

The current project is an experimental baseline. These code changes are visible in the commit history; independent end-to-end speed or quality gains have not been established here.

How it works

The inherited method quantizes keys per channel and values per token, retaining recent tokens at higher precision. Llama and Mistral examples support 2- and 4-bit KV-cache settings without fine-tuning.

Inherited quantization scheme: key cache per-channel and value cache per-token.

CacheForge's initial quantization flow during prefill and decoding:

Using CacheForge

Setup

To install the required packages:

conda create -n cacheforge python=3.10
conda activate cacheforge
pip install --upgrade pip  # enable PEP 660 support
pip install -e .

Then install the inherited CUDA extension:

cd quant && pip install -e .

Example

Load a model with CacheForge (for example, Llama-2-7b):

# LLaMA model with CacheForge
import torch
from models.llama_cacheforge import LlamaForCausalLM_CacheForge
from transformers import LlamaConfig, AutoTokenizer
config = LlamaConfig.from_pretrained("meta-llama/Llama-2-7b-hf")

config.k_bits = 2
config.v_bits = 2
config.group_size = 32
config.residual_length = 128  # number of recent fp16 tokens
CACHE_DIR = "./model-cache"

model = LlamaForCausalLM_CacheForge.from_pretrained(
    pretrained_model_name_or_path='meta-llama/Llama-2-7b-hf',
    config=config,
    cache_dir=CACHE_DIR,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    device_map="auto",
)

tokenizer = AutoTokenizer.from_pretrained(
    'meta-llama/Llama-2-7b-hf', 
    use_fast=False, 
    trust_remote_code=True, 
    tokenizer_type='llama')

# Inference
# e.g., model.generate(...)

GSM8K example

Use example.py to run the GSM8K example:

python example.py

Passkey retrieval example

Run CacheForge's passkey-retrieval example with long_context_example.py:

python long_context_example.py

Evaluate CacheForge on LongBench

The current baseline supports Llama and Mistral model families. See long_bench.md for the inherited reference results.

bash scripts/long_test.sh {GPU_ID} {K_BITS} {V_BITS} {GROUP_LENGTH} {RESIDUAL_LENGTH} {MODEL_NAME}
python eval_long_bench.py --model {MODEL} # MODEL is the dir name under pred/ Currently it support Llama family model and Mistral model.

Contributing

Contributions that advance CacheForge beyond its initial baseline are welcome.

License

CacheForge is distributed under the MIT License. The original copyright and license notice is preserved in LICENSE.

Credits and provenance

CacheForge began from jy-yuan/KIVI at 876b4d2. The inherited implementation and research are credited to Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. CacheForge is not an official KIVI release.

The published 2.6× lower peak memory, up to 4× larger batches, and 2.35–3.47× throughput figures are upstream KIVI results, not CacheForge benchmark results. See the original paper and inherited LongBench reference for their evaluation context. If you use the method or implementation, cite KIVI:

@InProceedings{pmlr-v235-liu24bz,
  title={{KIVI}: A Tuning-Free Asymmetric 2bit Quantization for {KV} Cache},
  author={Liu, Zirui and Yuan, Jiayi and Jin, Hongye and Zhong, Shaochen and Xu, Zhaozhuo and Braverman, Vladimir and Chen, Beidi and Hu, Xia},
  booktitle={Proceedings of the 41st International Conference on Machine Learning},
  pages={32332--32344},
  year={2024},
  volume={235},
  series={Proceedings of Machine Learning Research},
  publisher={PMLR},
  url={https://proceedings.mlr.press/v235/liu24bz.html}
}