Skip to content

DHdroid/NSNQuant

Repository files navigation

NSNQuant: A Double Normalization Approach for Calibration-Free Low-Bit Vector Quantization of KV Cache [Paper]

This is the official implementation of the paper NSNQuant: A Double Normalization Approach for Calibration-Free Low-Bit Vector Quantization of KV Cache (NeurIPS 2025). Some parts of the code are inspired by and adapted from Palu, KIVI, and KVQuant.

We also include pseudo-quantized implementations of three baseline methods: KIVI, KVQuant, and CQ (Coupled Quantization). The artifacts (NUQ tables and codebooks) needed for KVQuant and CQ will be released soon.

Overview

NSN Overview NSNQuant is a calibration-free vector quantization (VQ) algorithm designed for low-bit (<= 2-bit) KV cache quantization. By applying the Hadamard transform and a 3-step normalization algorithm (NSN), it aligns the KV distribution with the standard normal distribution. Leveraging this alignment, a single reusable VQ codebook can be pre-built using only synthetic standard normal data (e.g. torch.randn).

NSN_effect The main components of NSNQuant are two transforms: the Hadamard transform and Normalize-Shift-Normalize (NSN). These two transforms work in a complementary manner to align the KV distribution with the standard normal distribution, enabling us to build an optimized codebook without relying on external data (i.e., a calibration dataset).

  1. Hadamard transform yields a Gaussian-like KV distribution, supported by the central limit theorem.
  2. NSN standardizes the output distribution when used with the Hadamard transform.

Example commands

Here, we provide sample commands to reproduce the results from our paper. For more detailed options, please refer to the config files located in src/conf.

Setup

You need two separate environments for lm_eval and the other experiments because their hydra-core requirements are incompatible.

# General setup
conda create -n NSNQuant python=3.10
conda activate NSNQuant
pip install -r requirements.txt
python setup.py install
pip install -e 3rdparty/fast-hadamard-transform

# Setup for lm_eval
conda create -n lm_eval python=3.10
conda activate lm_eval
pip install -r requirements-lm-eval.txt
python setup.py install
pip install -e 3rdparty/fast-hadamard-transform
pip install -e 3rdparty/lm-evaluation-harness
pip install -e '3rdparty/lm-evaluation-harness[math]'

LongBench evaluation

# Generate predictions
python run_longbench_pred.py model=your_model_name quantizer=your_quantizer

# Run evaluation script to compute metrics based on predictions
python run_longbench_eval.py --model your_model_name_your_quantizer_postfix

LM-Evaluation-Harness

python run_lm_eval.py --task your_task --model_name_or_path your_model_name_or_path --quantizer your_quantizer

Codebook tuning

# 2-bit codebook
python generate_codebook.py --method learned --output_path /path/to/your/custom/codebook --abs --save

# 1-bit codebook
python generate_codebook.py --method learned --output_path /path/to/your/custom/codebook --save

Note

  • For generation-style evaluations such as LongBench and LM-Evaluation-Harness, we recommend batch_size=1 because padding can change residual cache split boundaries.
  • The numbers may differ slightly from those in the paper depending on the environment, but they should not differ significantly. If you observe a severe error, please report it via an issue.
  • Except for NSNQuant, the included baselines are implemented with pseudo quantization for quality comparison, not as full memory-saving deployment kernels.

Citation

@article{son2026nsnquant,
  title={NSNQuant: A double normalization approach for calibration-free low-bit vector quantization of KV cache},
  author={Son, Donghyun and Choi, Euntae and Yoo, Sungjoo},
  journal={Advances in Neural Information Processing Systems},
  volume={38},
  pages={43124--43159},
  year={2026}
}

About

[NeurIPS 2025] NSNQuant: A Double Normalization Approach for Calibration-Free Low-Bit Vector Quantization of KV Cache

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors