tiny-turboquant 0.18.0

Tiny, drop-in Hugging Face KV cache compression using KIVI-style low-bit KV storage.

tiny-turboquant

A transformers.DynamicCache you can swap in for past_key_values= to shrink KV-cache storage by 3–4× on Hugging Face decoder-only models that support DynamicCache / past_key_values. No calibration set, no fine-tuning, no model surgery — just construct and swap.

from tiny_turboquant import TinyKVCache

cache = TinyKVCache(compression="safe")
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.49. Python 3.10+. numpy / scipy are only required for the standalone TinyQuantizer (pip install "tiny-turboquant[vector]"); the fused decode-attention path requires Triton (pip install "tiny-turboquant[fused]").

Quick Start

Drop into any Hugging Face decoder-only model

from transformers import AutoModelForCausalLM, AutoTokenizer
from tiny_turboquant import TinyKVCache
import torch

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B-Instruct", dtype=torch.float16, device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B-Instruct")

# Create a compressed cache
cache = TinyKVCache(compression="safe")   # ~3× smaller, quality-first preset

# Use it like a normal past_key_values
inputs = tokenizer("Hello, how are you?", return_tensors="pt").to(model.device)
outputs = model(**inputs, past_key_values=cache, use_cache=True)

Asymmetric K/V (more savings, similar quality)

cache = TinyKVCache(key_bits=4, value_bits=2)   # ~4× smaller

Protect boundary layers (best quality at the same bits)

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

Let the library pick bits for you

cache = TinyKVCache.from_pretrained_model(model, tokenizer)

Runs one short forward pass to measure K/V dynamic range, then constructs the cache with the recommended key_bits/value_bits. Pass return_report=True to also get the KVNormReport back.

Run the inference server

Ships with an OpenAI-compatible HTTP server (FastAPI, SSE streaming). Point any OpenAI client at it:

pip install "tiny-turboquant[server]"
tiny-turboquant-server --model Qwen/Qwen2.5-3B-Instruct --compression safe --port 8000

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"Hello!"}],"max_tokens":100}'

Set "stream": true in the request body to get token-by-token SSE chunks.

Knobs that actually move the needle

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

TinyKVCache(compression="safe")        # K4/V4   default, best quality (~3.3x)
TinyKVCache(compression="balanced")    # K4/V3   recommended tradeoff
TinyKVCache(compression="compact")     # K3/V3   storage-focused
TinyKVCache(compression="aggressive")  # K3/V2 group=32   experimental

Under the hood these map to (key_bits, value_bits) of (4, 4), (4, 3), (3, 3), 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)

# Quality-first default: safest public recommendation.
rec = recommend_bits(report)  # usually K4/V4

# Tradeoff mode: more compression, validate on your task.
# rec = recommend_bits(report, mode="balanced")  # often K4/V3

cache = TinyKVCache.for_model(
    model,
    key_bits=rec["key_bits"],
    value_bits=rec["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" | "compact" | "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

Three ship with the package — all thin wrappers; the actual work happens in TinyKVCache.

# OpenAI-compatible HTTP server (FastAPI + SSE streaming)
pip install "tiny-turboquant[server]"
tiny-turboquant-server --model Qwen/Qwen2.5-3B-Instruct --compression safe

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

# Token-by-token agreement A/B against stock DynamicCache
tiny-turboquant-validate --model Qwen/Qwen2.5-3B-Instruct --compression balanced

Server endpoints: POST /v1/chat/completions (supports "stream": true SSE), 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, prompt 5230 tokens. Strictly monotone Pareto frontier across both model sizes:

preset	bits	3B ratio / ΔPPL	0.5B ratio / ΔPPL	status
safe	K4/V4	3.34x / +0.07	3.26x / +0.23	default
balanced	K4/V3	3.71x / +0.24	3.61x / +0.59	recommended tradeoff
compact	K3/V3	4.18x / +0.45	4.05x / +1.41	storage-focused
aggressive	K3/V2 gs=32	4.31x / +0.81	4.31x / +2.80	experimental

aggressive applies grouped V quantization (value_group_size=32): each token's value vector is split into head_dim / 32 groups and each group gets its own scale/zero pair. This is what lets 2-bit V stay usable. To opt out, pass value_group_size=None.

safe is the lowest-drift preset in our Qwen2.5 tests and scales with model size — the PPL gap shrinks from +0.23 on 0.5B to +0.07 on 3B. Reach for balanced when you want more compression with mild quality cost; compact when storage matters more than a small PPL bump; aggressive only after validating on your task (the 0.5B numbers warn that quality cliffs exist on small models).

TinyKVCache uses KIVI-style block affine quantization for live KV-cache compression. TinyQuantizer is kept as a standalone experimental vector quantizer and is not used by the cache path. TinyQuantizerIP is deprecated for attention experiments because JL residual noise can be amplified by softmax.

Performance

Decode-loop overhead vs. stock DynamicCache on Qwen2.5-3B-Instruct, T4 GPU, 4096-token context, 64 new tokens, balanced preset:

cache	decode ms/tok	overhead
DynamicCache fp16	73.0 ms	1.00×
TinyKVCache (PyTorch)	87.9 ms	1.20×
TinyKVCache (Triton)	80.6 ms	1.10×

Overhead is essentially all dequant cost — attention compute is unchanged. The dequant Triton path is opt-in (use_triton=True) while it accumulates generation/PPL coverage; it replaces the PyTorch dequant with a fused per-(batch, head, token) kernel and is bit-exact vs. the PyTorch path (max |diff| = 0 across all 4 presets on T4). Per-layer dequant cost at (B=1, H=16, D=128, ctx=4096) on T4:

preset	PyTorch	Triton	speedup
safe	2.9 ms	1.3 ms	2.32×
balanced	4.2 ms	2.6 ms	1.63×
aggressive	4.3 ms	2.7 ms	1.59×
compact	5.5 ms	3.9 ms	1.41×

Two non-obvious notes:

• compact is slower than aggressive, despite using more bits per value. Per-token V scales pay an allocation cost that aggressive's grouped V (one scale per 32 channels) avoids. Pick compact for quality reasons, not perf reasons.
• T4 is memory-bandwidth-limited; the overhead ratio is smaller on A100/H100. Per-token throughput scales accordingly.

Fused decode-attention path (experimental, 0.18.x)

A second Triton path fuses dequant + attention into two kernel launches per decode step, eliminating the dense SDPA call on the dequantized K/V. End-to-end Kaggle T4×2 numbers on Llama-3.2-1B-Instruct, balanced preset:

ctx	fp16 dense	TinyKV stock	TinyKV fused	vs dense	vs stock
1024	29.5 ms	43.9 ms	28.9 ms	1.02×	1.52×
4096	32.6 ms	46.9 ms	28.6 ms	1.14×	1.64×
8192	55.1 ms	85.3 ms	28.7 ms	1.92×	2.97×

WikiText-2 PPL at the same setup (12,288 scored tokens):

config	PPL	Δ vs fp16
fp16 dense	10.7858	—
TinyKV stock	11.0425	+0.2567
TinyKV fused	11.0431	+0.2573

Fused-vs-stock PPL gap is 6e-4 — fp16 reduction-order noise only. Wire it into your model with:

from tiny_turboquant import TinyKVCache
from tiny_turboquant.hf_attention import patch_model

patch_model(model)                          # opt-in monkey-patch
cache = TinyKVCache(compression="balanced")
out = model(**inputs, past_key_values=cache, use_cache=True)

Alpha limitations (the fallback path runs automatically when any apply):

• fp16 only (bf16 silently routes through the PyTorch reference).
• single-request decode only: attention_mask must be None; no padding, no sliding window, no custom mask. model.generate() builds an internal mask even when you pass None — this falls back to stock SDPA silently. Watch INFO logs from tiny_turboquant.hf_attention / tiny_turboquant.triton_fused_attn to see the exact fallback reason (logged once per process per reason).
• outlier_channels=0 presets only.
• prefill and protected layers always use stock attention.

Tested configurations

What we've actually run end-to-end on the numbers above:

Axis	Tested	Untested (kernel is portable, expected to work)
GPU	Tesla T4 (sm75) ×1, ×2 via device_map="auto"	A100 / H100, ≥4 GPU
Model	Llama-3.2-1B-Instruct (fp16)	Larger Llama, Qwen2 *, Mistral *
Dtype	fp16	bf16 (falls back to reference path)
Batch	B=1	B>1
Mask	attention_mask=None	Any non-trivial mask (falls back to stock)

* Qwen2 portability bench published at kaggle_bench_qwen2.py; not yet verified by maintainers, runnable by users.

CUDA fast path

TinyKVCache(..., use_triton=...) controls the read-path dequant backend. Two modes:

• False (default) — always use the PyTorch dequant path.
• True — use the fused Triton kernel; raises on CUDA tensors if Triton or CUDA is missing.

The dequant Triton path is bit-exact vs. PyTorch on the four shipped presets (max |diff| = 0 across safe / balanced / compact / aggressive on T4, fp16) and is opt-in until it has broader generation/PPL coverage. It will become the default in a future release. The dequant Triton kernels live in tiny_turboquant.triton_dequant; the experimental fused decode-attention kernels live in tiny_turboquant.triton_fused_attn and the HF monkey-patch in tiny_turboquant.hf_attention.

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.