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 3-bit scalar quantization against N(0, 1/d), shared codebook for K and V

For HuggingFace integration, keys and values are both compressed and decompressed via fused kernels — a single kernel launch compresses both K and V, and a single launch decompresses both. The reconstructed FP16 tensors are packed into a standard DynamicCache that HuggingFace's attention uses directly. No model changes needed.

The package also ships fused attention kernels with QJL bias correction (2-bit Lloyd-Max keys + 1-bit sign sketch of the quantization residual). These perform scoring, online softmax, and V accumulation in one kernel with on-chip V decompression. They're fully implemented but not yet wired into the HuggingFace path — integrating them requires replacing the model's internal attention, which is model-specific. This is a candidate for a cuTile Gym contribution.

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.

kernels

HuggingFace path (used by default):

Fused attention path (included, not in HuggingFace API):

license

MIT