turboquant-gpu 0.1.6

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 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.

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

HuggingFace path (used by default):

kernel	what it does
compress_kv_3bit	fused K+V compression, 3-bit shared codebook, single launch
decompress_kv_3bit	fused K+V decompression, single launch
compress_values_3bit / 2bit	separate fallback for K or V individually
decompress_3bit / 2bit	separate fallback decompression

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

kernel	what it does
compress_keys_2bit_qjl	2-bit Lloyd-Max + 1-bit QJL signs for keys
fused_attention	QJL-corrected scores + online softmax + V accumulation
fused_attention_vfused_3bit	same + on-chip V decompression from compressed indices
attention_scores	score-only (no softmax), for debugging

license

MIT