tiny-turboquant 0.14.0

Tiny, drop-in HuggingFace KV cache compression via random-rotation quantization.

tiny-turboquant

A transformers.DynamicCache you can swap in for past_key_values= to shrink KV-cache storage by 3–4× on any HuggingFace causal LM. No calibration set, no fine-tuning, no model surgery — just construct and swap.

from tiny_turboquant import TinyKVCache

cache = TinyKVCache(compression="balanced")
outputs = model(**inputs, past_key_values=cache, use_cache=True)

That's the whole integration. Everything below is optional.

---

Install

pip install tiny-turboquant

Or from source:

git clone https://github.com/pradeepboopathy/tiny-turboquant
cd tiny-turboquant && pip install -e .

torch >= 2.2, transformers >= 4.40, numpy, scipy. Python 3.10+.

Knobs that actually move the needle

Compression preset. The default knob is a named preset, not a bit count. Three settings cover almost every real use case:

TinyKVCache(compression="safe")        # ~3× smaller storage (4-bit K, 4-bit V)
TinyKVCache(compression="balanced")    # ~4× smaller storage (4-bit K, 2-bit V, recommended)
TinyKVCache(compression="aggressive")  # ~4× smaller storage (3-bit K, 2-bit V, some quality hit)

Under the hood these map to (key_bits, value_bits) of (4, 4), (4, 2), and (3, 2) respectively. If you want to pick the bits yourself, pass them directly instead of compression=:

TinyKVCache(key_bits=4, value_bits=2)

The two paths are mutually exclusive — mixing them raises ValueError.

Pick a preset from the model itself. Keys usually carry more dynamic range than values in Llama/Qwen/Mistral-family models. The diagnostic runs one forward pass with hooks on every k_proj / v_proj and tells you whether the asymmetric preset will help:

from tiny_turboquant import measure_kv_norm_ratio, recommend_bits

report = measure_kv_norm_ratio(model, tokenizer)
rec    = recommend_bits(report)        # → {'key_bits': 4, 'value_bits': 2, 'reason': ...}
cache  = TinyKVCache.for_model(model, **{k: rec[k] for k in ('key_bits', 'value_bits')})

Boundary-layer protection. Embedding-adjacent and logit-adjacent layers are noticeably more bit-sensitive than the middle. Keeping the outer layers at FP16 costs little memory and recovers most of the small PPL delta:

TinyKVCache.for_model(
    model, compression="balanced", protect_first=2, protect_last=2,
)

Use for_model (not the bare constructor) whenever you set protect_last or pass negative indices — the cache needs to know model.config.num_hidden_layers.

arg	default	notes
compression	–	preset: "safe" | "balanced" | "aggressive"
bits	4	(advanced) symmetric bit width, used when no preset
key_bits / value_bits	–	(advanced) per-stream bit widths; override bits
protect_first / protect_last	0	keep this many boundary layers at FP16
protected_layers	–	explicit layer indices (negatives allowed via for_model)

Command-line tools

Two ship with the package — both are thin wrappers; the actual work happens in TinyKVCache.

# OpenAI-compatible HTTP server (stdlib http.server — no FastAPI dependency)
tiny-turboquant-server --model Qwen/Qwen2.5-3B-Instruct --bits 4

# Sliding-window WikiText-2 perplexity, writes a JSON report
tiny-turboquant-eval --model Qwen/Qwen2.5-3B-Instruct --key-bits 4 --value-bits 2

Server endpoints: POST /v1/chat/completions, GET /v1/models, GET /health.

The algorithm in three lines

1. for keys:   per-channel asymmetric affine quant (min/max along token axis)
               within a block of 128 tokens — absorbs outlier channels
2. for values: per-token asymmetric affine quant (min/max along channel axis)
               within the same block — matches the granularity attention
               actually reads at
3. keep the most recent N tokens (default 128) uncompressed in an FP16
               residual window

This matches the KIVI-2 scheme. Earlier versions used random-rotation + a closed-form Beta-distribution codebook from the TurboQuant paper; that path was elegant but assumed a marginal distribution real attention streams don't have, and degraded badly at low bits. Empirical per-channel scaling is both simpler and substantially more accurate on real models.

Measured quality

WikiText-2-raw test split, sliding window 2048 / stride 1024, first 8 chunks:

model	preset	PPL	Δ vs fp16	storage savings
Qwen2.5-0.5B	fp16 baseline	13.671	—	1.0×
Qwen2.5-0.5B	safe (4K / 4V)	13.949	+2.0%	3.1×
Qwen2.5-0.5B	balanced (4K / 2V)	17.344	+27%	3.9×
Qwen2.5-0.5B	aggressive (3K / 2V)	19.466	+42%	4.3×
Qwen2.5-3B	fp16 baseline	8.156	—	1.0×
Qwen2.5-3B	safe (4K / 4V)	8.258	+1.3%	3.3×
Qwen2.5-3B	balanced (4K / 2V)	11.148	+37%	4.1×
Qwen2.5-3B	aggressive (3K / 2V)	12.026	+47%	4.7×

safe is the near-lossless setting and scales well — the PPL gap shrinks from 2.0% on 0.5B to 1.3% on 3B. balanced and aggressive degrade because 2-bit V only has 4 levels per token — the published KIVI-2 tradeoff. Storage savings include per-block scale/zero metadata (one fp32 scale + zero per channel for K, per token for V).

The MSE-only variant (TinyQuantizer) is what the cache uses. A two-stage MSE + QJL variant (TinyQuantizerIP) ships for completeness but is deprecated for attention — softmax exponentially amplifies the JL-projection noise.

CUDA fast path

A fused dequant-then-attention CUDA kernel is planned for the new per-channel-K / per-token-V layout. The kernel in cuda/ was written against the old rotation-based path and is not used by the current cache; dequant runs in PyTorch on the read path until the rewrite lands.

Where it earns its keep, where it doesn't

Useful when KV memory is the bottleneck — long contexts on a single GPU, many concurrent serving sessions, or running one model size larger by buying back VRAM from the cache. Not useful for short contexts (< 1k tokens; the cache is already small), for hybrid / recurrent architectures that don't keep a standard KV cache (Mamba, RWKV), or for tasks that need bit-exact reproducibility.

References

• KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (Liu et al., ICML 2024). The current cache uses the KIVI-2 layout: per-channel K, per-token V.
• TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate (Zandieh et al., ICLR 2026). Earlier versions of this package implemented the rotation + Beta-codebook scheme from this paper; it ships as TinyQuantizer for standalone vector quantization but is no longer used by TinyKVCache.

Architecture deep-dive: docs/ARCHITECTURE.md.

License

MIT — see LICENSE.