TurboQuant-GPU

5.02x KV cache compression for LLM inference. Works on any NVIDIA GPU. cuTile kernels with automatic PyTorch fallback.

pip install turboquant-gpu

Check out the PyPi package

quickstart

from transformers import AutoModelForCausalLM, AutoTokenizer
from turboquant_gpu import TurboQuantEngine
import torch

model_id = "mistralai/Mistral-7B-v0.1"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="cuda")
tok   = AutoTokenizer.from_pretrained(model_id)

engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cuda")
result = engine.generate(model, tok, "The University of Waterloo is known for ")

print(result["text"])
print(f"{result['tokens']} tokens | {result['stats']['ratio']:.2f}x compression")

how it works

Random orthogonal rotation makes each KV cache coordinate approximately Gaussian. Lloyd-Max quantization is then optimal for that distribution, giving 3 bits per element with 0.98 cosine similarity.

Keys get 2-bit MSE quantization + 1-bit QJL bias correction. Values get 3-bit MSE quantization. Both are compressed in a single fused kernel launch per attention head.

vs NVIDIA FP4

TurboQuant achieves 5.02x compression compared to 3.76x (MXFP4) and 3.56x (NVFP4) because it exploits the post-rotation Gaussian structure specific to KV caches, rather than using a general-purpose FP4 format.

install

pip install turboquant-gpu

For cuTile GPU kernel acceleration (optional, requires CUDA 13.0+ driver):

pip install cuda-tile[tileiras] --extra-index-url https://pypi.nvidia.com

If cuda-tile isn't available or your driver is older, everything still works via PyTorch.

API

engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cuda")

# one-call generation
result = engine.generate(model, tokenizer, "your prompt")

# step-by-step
compressed = engine.compress_kv_cache(out.past_key_values)
cache      = engine.build_cache(compressed)
stats      = engine.compression_stats(out.past_key_values)

# auto-tune for your GPU (benchmarks cutile vs pytorch, 2-bit vs 3-bit)
engine.auto_tune(seq_len=512)

GPU support

Uses cuTile for cross-architecture kernel portability. Falls back to PyTorch automatically.

kernels

notebooks

license

MIT