tiny-turboquant 0.2.0

Low-bit vector and KV-cache compression research toolkit for PyTorch

Tiny TurboQuant

Tiny TurboQuant is a lightweight PyTorch research toolkit for low-bit vector compression and KV-cache compression experiments.

Version 0.2.0 adds a quality-aware hybrid KV-cache path: recent tokens can stay dense while older tokens are stored in packed low-bit form.

What it includes

• real uint8 bit-packing for low-bit indices
• MSE-style scalar/vector quantization
• product / inner-product-oriented quantization experiments
• outlier-split quantization
• compressed embedding storage demos
• TurboQuantKVCache: all-packed KV-cache prototype
• HybridTurboQuantKVCache: recent dense + older compressed KV-cache prototype
• separate Key/Value bit-width settings
• stronger outlier preservation for Keys and Values
• optional per-layer calibration
• demos for ANN search, embedding compression, and KV-cache memory measurement

Important limitation

This package demonstrates packed memory compression and memory-quality benchmarking. It is not a production compressed-attention engine. The KV-cache wrappers still dequantize tensors before attention. Real latency gains require fused CUDA/Triton kernels or integration with an inference engine such as vLLM or TensorRT-LLM.

Do not use this package to claim training acceleration, fine-tuning memory reduction, production legal/medical QA readiness, drop-in vLLM replacement, exact nearest-neighbor search, or faster LLM inference.

Install

pip install tiny-turboquant

Basic usage

import torch
from tiny_turboquant import TurboQuantMSE, TurboQuantKVCache, HybridTurboQuantKVCache

x = torch.randn(128, 64)
x = x / x.norm(dim=-1, keepdim=True)

q = TurboQuantMSE.build(d=64, bits=4)
idx = q.quant(x)
x_hat = q.dequant(idx)

# v0.1-compatible all-packed cache
cache = TurboQuantKVCache(bits=4)

# v0.2 quality-aware hybrid cache
hybrid_cache = HybridTurboQuantKVCache(
    key_bits=6,
    value_bits=4,
    recent_window=128,
    key_outlier_bits=8,
    value_outlier_bits=8,
    n_key_outliers=64,
    n_value_outliers=32,
)

KV benchmark examples

All-packed v0.1-style cache:

python -m benchmarks.bench_kv_real   --cache_type compressed   --model Qwen/Qwen2.5-0.5B-Instruct   --bits 4   --mode quick

Hybrid v0.2 cache:

python -m benchmarks.bench_kv_real   --cache_type hybrid   --model Qwen/Qwen2.5-0.5B-Instruct   --key_bits 6   --value_bits 4   --key_outlier_bits 8   --value_outlier_bits 8   --recent_window 128   --mode quick

For a short context, use a smaller recent window to force compression:

python -m benchmarks.bench_kv_real   --cache_type hybrid   --recent_window 16   --key_bits 6   --value_bits 4   --mode quick

Run tests

python -m pytest

Run demos from source checkout

python -m demos.demo1_distortion_vs_theory
python -m demos.demo2_ann_vs_pq
python -m demos.demo3_real_embeddings
python -m demos.demo4_kv_cache

Scope

Use this package for:

• compressed vector-search experiments
• RAG embedding compression experiments
• KV-cache memory/quality tradeoff experiments
• educational or research benchmarking
• testing hybrid recent-token retention strategies
• studying separate K/V compression settings

Current focus:

• memory compression
• quality measurement
• cache-format experiments

Future direction:

• fused dequant + attention kernels
• compressed-domain attention estimators
• better per-head/per-layer calibration
• integration experiments with serving engines