turboquant-gpu 0.1.4

TurboQuant KV cache compression for LLM inference — cuTile GPU kernels

turboquant-gpu

5.02x KV cache compression for LLM inference — cuTile kernels with automatic PyTorch fallback.

pip install turboquant-gpu

Works on any NVIDIA GPU. Uses cuTile kernels when available, otherwise falls back to PyTorch automatically — no driver upgrades or manual config needed.

quick start

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")

install

pip install turboquant-gpu

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

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

If you skip cuda-tile or your driver is older, everything still works via PyTorch.

how it works

Implements the TurboQuant algorithm:

1. normalize + rotate — random orthogonal rotation (Pi) makes coordinates near-Gaussian
2. Lloyd-Max quantize — optimal scalar quantization against N(0, 1/d)

Both keys and values are compressed to 3-bit via MSE-optimal Lloyd-Max quantization, then reconstructed for standard attention. The package also includes fused attention kernels with QJL bias correction (1-bit sign sketch of the quantization residual), but these are not yet exposed in the high-level API.

step-by-step api

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

# after model prefill:
compressed = engine.compress_kv_cache(out.past_key_values)
cache      = engine.build_cache(compressed)
stats      = engine.compression_stats(out.past_key_values)

# or just do it all in one call:
result = engine.generate(model, tokenizer, "your prompt here")

# auto-tune for your specific GPU:
engine.auto_tune(seq_len=512)

gpu support

Written in cuTile for cross-architecture portability. Falls back to PyTorch if cuTile or a compatible driver isn't available.

GPU	cuTile kernels	PyTorch fallback
A100 (Ampere, sm_80)	CUDA 13.2+ driver	always works
H100 (Hopper, sm_90)	not yet supported by tileiras	always works
RTX 4090 (Ada, sm_89)	CUDA 13.2+ driver	always works
B200/B300 (Blackwell, sm_100)	CUDA 13.0+ driver	always works
Any other CUDA GPU	depends on tileiras	always works

kernels

kernel	what it does	status
compress_kv_3bit	fused K+V compression in a single kernel launch	used (default)
compress_keys	key-only: normalize → rotate → Lloyd-Max → QJL signs	fallback
compress_values	value-only: normalize → rotate → Lloyd-Max	fallback
decompress_values	dequantize → un-rotate(Pi) → scale by norms	used
attention_scores	asymmetric dot product with QJL correction	included, not in API
fused_attention	scores + online softmax + V accumulation	included, not in API

license

MIT