Metadata-Version: 2.4
Name: turboquant-gpu
Version: 0.1.8
Summary: TurboQuant KV cache compression for LLM inference — cuTile GPU kernels
Author: Anirudh Bharadwaj Vangara
License-Expression: MIT
Keywords: quantization,kv-cache,llm,inference,cutile,cuda,gpu,attention,blackwell,hopper,h100,b200
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Requires-Dist: scipy
Provides-Extra: gpu
Requires-Dist: cuda-tile; extra == "gpu"
Dynamic: license-file

# turboquant-gpu

![TurboQuant-GPU](https://i.ibb.co/hxTFsTLn/thumbnail.png)

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

```python
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

```bash
pip install turboquant-gpu
```

For cuTile acceleration (optional, requires CUDA 13.0+ driver):
```bash
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](https://arxiv.org/abs/2501.09747) 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

```python
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](https://docs.nvidia.com/cuda/cutile-python/) 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
