Skip to main content

KIVI

KIVI is VeloxQuant-MLX's re-implementation of the field's most widely-cited KV-cache quantization baseline: "KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache" (Liu, Yuan et al., ICML 2024). It is included so every other algorithm in this library can be measured against a recognized reference point.

:::info Why a baseline? KIVI is not the highest-compression method here (VecInfer-2bit reaches 8× key compression vs KIVI-2bit's ~5.8×). Its value is being calibration-free, deterministic, and the number reviewers expect to see compared against. :::

How it works​

KIVI's insight is asymmetry between keys and values:

  1. Keys are quantized per channel — the quantization group runs along the token axis, with one (scale, zero) pair per channel-group. Key distributions have a few high-variance channels, so per-channel scales keep them accurate.
  2. Values are quantized per token — the group runs along the channel axis, one (scale, zero) per token-group.
  3. fp16 residual window — the most recent residual_length tokens are kept at full precision. Newly generated tokens dominate attention and are cheap to keep exact; they are quantized only once they age out of the window.

Each group uses asymmetric min/max quantization:

zero  = min(group)
scale = (max(group) - min(group)) / (2**b - 1)
q     = round((group - zero) / scale)      # uint, [0, 2**b - 1]
recon = q * scale + zero

KIVI is fully deterministic — no codebook training, no rotation, no RNG — so it adds no run-to-run variance.

Usage​

import mlx_lm
from veloxquant_mlx import KVCacheConfig, KVCacheBuilder

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit")

config = KVCacheConfig(
    method="kivi",
    bit_width_inlier=2,      # KIVI default
    kivi_group_size=32,      # min/max group size
    residual_length=32,      # recent tokens kept in fp16
)
caches = KVCacheBuilder.for_model(model, config)
model.make_cache = lambda *_a, **_k: caches

response = mlx_lm.generate(model, tokenizer, prompt="...", max_tokens=120)

Measured results​

Apple M4, max_tokens≈120, residual_length=32, long-context prompt. Source: figures/kivi/<model>/results.json.

ModelKIVI-2bit key comp.full-KV comp.throughput vs fp16
Llama-3.2-3B-4bit5.79×3.98×102%
Qwen2.5-7B-4bit5.78×3.98×100%
Mistral-7B-4bit5.76×4.03×106%

Full-KV compression includes the fp16 residual window, so it is not inflated.

Honest scope​

:::warning Memory, not raw speed; storage, not peak

  • KIVI's published speedup comes from a CUDA kernel that does not port to Metal. On Apple Silicon the win is memory; throughput is at-or-near fp16 because the min/max arithmetic is cheap on a memory-bound decode path.
  • Compression only manifests once context exceeds the residual window — at short prompts the whole prefill stays fp16 and the realized ratio is 1.0× (correct behavior, not a bug).
  • Peak runtime memory is not reduced (sometimes marginally higher): keys are dequantized to fp16 before SDPA, so the compression is in cache-storage accounting, not the peak fp16 working set.
  • At 2 bits, raw-key reconstruction cosine on synthetic unit-norm Gaussian keys is ~0.93 — KIVI 2-bit is genuinely lossy, which is exactly why the fp16 residual window exists. :::

See figures/kivi/fig4_vs_existing.png for the KIVI-vs-VecInfer comparison.

See also: NSNQuant — the other residual-window wrapper; it differs because it adapts the data to a fixed universal codebook (NSN + Hadamard Gaussianization) rather than fitting scalar min/max scales to the data.