Skip to content

Intelligent-Computing-Lab-Panda/KronQ

Repository files navigation

KronQ: LLM Quantization via a Kronecker-Factored Hessian

Donghyun Lee · Yuhang Li · Ruokai Yin · Priyadarshini Panda

COLM 2026

arXiv HF Models License Llama

KronQ results: perplexity vs bit-width, peak VRAM, and decode latency

KronQ extends GPTQ/GPTAQ with an output-side curvature term $H_G$ alongside the usual input-side Hessian $H_X$, under a K-FAC factorization $H \approx H_X \otimes H_G$. This output curvature drives two mechanisms that improve low-bit (W2–W4) weight-only quantization with no fine-tuning:

  • Bidirectional incoherence processing (BiIP) — incoherence rotation + rescaling on both the input and output sides.
  • Inter-layer mixed precision — per-sublayer bit allocation guided by $\mathrm{tr}(H_G)\cdot\mathrm{tr}(H_X)$ sensitivity.

Quantized weights are stored packed int4/int2; a fused dequant + BiIP CUDA matvec unpacks them on the fly (int-weight × fp16-activation), CUDA-graph wrapped for single-token decode. On Llama-2-7B (A100, batch 1) the deploy path runs at 6.30 ms/token for both W2 and W4 (vs 11.6 ms fp16), and WikiText-2 PPL matches the research pipeline exactly (bf16 5.47 / W4 5.56 / W2 8.23).

Contents

Installation

# Turnkey (brings its own nvcc 12.1 so the CUDA kernels build anywhere):
conda env create -f environment.yml
conda activate kronq

# OR with an existing CUDA-12.1 toolkit on PATH (e.g. `module load CUDA/12.1.1`):
pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu121

# Build the CUDA extensions (needs nvcc):
pip install git+https://github.com/Dao-AILab/fast-hadamard-transform  # required by all paths
pip install ./runtime/kronq_kernels   # only for the packed-int deploy path

Quick start — run a released model

Evaluate a pre-quantized KronQ checkpoint straight from the 🤗 Hub — WikiText-2 PPL and/or zero-shot. No calibration code involved (uses only common/ + runtime/).

python eval_pretrained.py meta-llama/Llama-2-7b-hf donghyunli/Llama-2-7b-KronQ-W4A16 --ppl --zs

The first argument supplies the architecture + tokenizer; the second is the packed KronQ repo (a Hub id or a local packed dir). For the 70B models add --distribute to shard across all visible GPUs.

Model zoo

All checkpoints live under 🤗 donghyunli. W4 / W2 are real packed int4 / int2; W3 is shipped as fp16 fake-quant (3-bit is not bit-packable). Suffix g128 = group-size 128; no suffix = per-channel.

Base model per-channel group-128
Llama-2-7B W4 · W3 · W2 W4 · W3 · W2
Llama-2-13B W4 · W3 · W2 W4 · W3 · W2
Llama-2-70B W4 · W3 · W2 W4 · W3 · W2
Llama-3-8B W4 · W3 · W2 W4 · W3 · W2
Llama-3-70B W4 · W3 · W2 W4 · W3 · W2

Reproduce from scratch

# 1) Get the H_G (output-curvature) cache. Either DOWNLOAD the companion dataset...
huggingface-cli download donghyunli/Llama-2-7b-KronQ-HG --repo-type dataset \
    --local-dir grad_cache/llama-2-7b
#    ...or precompute it yourself (one-time per model):
# python calib/precompute_gradients.py --model meta-llama/Llama-2-7b-hf \
#     --output grad_cache/llama-2-7b --nsamples 128

# 2) Quantize (fake-quant) + WikiText-2 PPL — kronq_kernels not needed here.
python main.py --model meta-llama/Llama-2-7b-hf \
    --w_bits 4 --w_groupsize -1 --w_clip --w_asym --a_bits 16 --act_order \
    --bi_calibration --use_gptaq --incoh_rotate --incoh_kernel had --incoh_mode full \
    --alpha 0.25 --grad_dir grad_cache/llama-2-7b

# 3) Save the packed Hub artifact directly (~bits/16 the size of fp16) by adding to step 2:
#    --save_unrotated_kronq --save_packed packed_dir

Companion $H_G$ caches (one per base model — the --grad_dir for step 2; layer-wise layer_N/<sublayer>_G.pt, drop-in for any bit-width): Llama-2-7B · Llama-2-13B · Llama-2-70B · Llama-3-8B · Llama-3-70B

A full one-command demo (precompute → quantize → deploy benchmark + generation) is in scripts/run_example.sh.

Mapping paper results to flags (all on main.py):

Result Flags
Weight-only PPL (per-channel) base recipe above, --alpha 0.25
Group-128 weight-only add --w_groupsize 128
Inter-layer mixed precision add --inter_layer_mp N (upgrade top-N sensitive sublayer types by +1 bit)
Component ablation vary --incoh_mode {full,no_rescale,su_only,sv_only} and toggle --use_gptaq
Weight-and-activation add --rotate --enable_aq_calibration --a_bits 4 --a_asym --a_clip_ratio 0.9, and use --alpha 0.5
Baselines --asym_calibrate (GPTAQ), --w_rtn (RTN), or none (GPTQ)
Deploy latency --save_packed (step 3) → runtime/bench_decode.py

Note — weight-only reproduces the paper with --alpha 0.25; weight-and-activation uses --alpha 0.5.

Repository layout

KronQ/
├── eval_pretrained.py             # ① USE: load a packed model (Hub/local) → PPL / zero-shot
├── main.py                        # ② REPRODUCE: quantize + flag routing
│
├── common/                        # shared infra (both ① and ②)
│   └── model_utils · quant_utils · hadamard_utils · data_utils · utils · monkeypatch
│
├── calib/                         # quantization algorithm (only ② needs this)
│   ├── kronq_utils.py             #   KronQ calibration (--bi_calibration)
│   ├── incoherence.py             #   bidirectional incoherence (BiIP)
│   ├── inter_layer_mp.py          #   mixed-precision bit allocator
│   ├── precompute_gradients.py    #   offline H_G cache
│   └── gptq_utils · gptaq_utils · rotation_utils    #   baselines + rotation
│
├── runtime/                       # inference & deploy (calibration-free)
│   ├── biip_linear.py             #   online BiIP forward
│   ├── packed_io.py               #   packed int4/int2 save/load
│   ├── lowbit_utils.py            #   pack / unpack / compress
│   ├── kronq_kernels/             #   fused dequant + BiIP CUDA matvec
│   └── eval_utils · generate · bench_decode · graph_wrapper · lowbit_triton
│
└── scripts/                       # run_example.sh (demo) · zs_distribute.py (multi-GPU eval)

common/, runtime/, calib/ are added to sys.path by a small bootstrap at the top of each entry point, so the flat import xxx style works from any directory.

Citation

@article{lee2026kronq,
  title={KronQ: LLM Quantization via Kronecker-Factored Hessian},
  author={Lee, Donghyun and Li, Yuhang and Yin, Ruokai and Panda, Priyadarshini},
  journal={arXiv preprint arXiv:2607.07964},
  year={2026}
}

Contact

Donghyun Lee — donghyun.lee.1@usc.edu

Acknowledgements & License

Apache License 2.0 — see LICENSE. Builds on GPTAQ and QuaRot; the incoherence machinery follows QuIP#. See NOTICE for attributions.

About

PyTorch Implementation of KronQ: LLM Quantization via Kronecker-Factored Hessian (COLM 2026)

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages