Metadata-Version: 2.4
Name: tiny-turboquant
Version: 0.14.0
Summary: Tiny, drop-in HuggingFace KV cache compression via random-rotation quantization.
Author: tiny-turboquant contributors
License: MIT
Project-URL: Homepage, https://github.com/pradeepboopathy/tiny-turboquant
Project-URL: Issues, https://github.com/pradeepboopathy/tiny-turboquant/issues
Keywords: llm,kv-cache,quantization,inference,turboquant,transformers
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.2.0
Requires-Dist: transformers>=4.40.0
Requires-Dist: scipy>=1.10.0
Requires-Dist: numpy>=1.24.0
Provides-Extra: server
Requires-Dist: uvicorn; extra == "server"
Requires-Dist: fastapi; extra == "server"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: pytest-benchmark; extra == "dev"
Dynamic: license-file

# tiny-turboquant

A `transformers.DynamicCache` you can swap in for `past_key_values=` to
shrink KV-cache storage by 3–4× on any HuggingFace causal LM. No
calibration set, no fine-tuning, no model surgery — just construct and
swap.

```python
from tiny_turboquant import TinyKVCache

cache = TinyKVCache(compression="balanced")
outputs = model(**inputs, past_key_values=cache, use_cache=True)
```

That's the whole integration. Everything below is optional.

---

## Install

```bash
pip install tiny-turboquant
```

Or from source:

```bash
git clone https://github.com/pradeepboopathy/tiny-turboquant
cd tiny-turboquant && pip install -e .
```

`torch >= 2.2`, `transformers >= 4.40`, `numpy`, `scipy`. Python 3.10+.

## Knobs that actually move the needle

**Compression preset.** The default knob is a named preset, not a bit
count. Three settings cover almost every real use case:

```python
TinyKVCache(compression="safe")        # ~3× smaller storage (4-bit K, 4-bit V)
TinyKVCache(compression="balanced")    # ~4× smaller storage (4-bit K, 2-bit V, recommended)
TinyKVCache(compression="aggressive")  # ~4× smaller storage (3-bit K, 2-bit V, some quality hit)
```

Under the hood these map to `(key_bits, value_bits)` of `(4, 4)`,
`(4, 2)`, and `(3, 2)` respectively. If you want to pick the bits
yourself, pass them directly instead of `compression=`:

```python
TinyKVCache(key_bits=4, value_bits=2)
```

The two paths are mutually exclusive — mixing them raises `ValueError`.

**Pick a preset from the model itself.** Keys usually carry more dynamic
range than values in Llama/Qwen/Mistral-family models. The diagnostic runs
one forward pass with hooks on every `k_proj` / `v_proj` and tells you
whether the asymmetric preset will help:

```python
from tiny_turboquant import measure_kv_norm_ratio, recommend_bits

report = measure_kv_norm_ratio(model, tokenizer)
rec    = recommend_bits(report)        # → {'key_bits': 4, 'value_bits': 2, 'reason': ...}
cache  = TinyKVCache.for_model(model, **{k: rec[k] for k in ('key_bits', 'value_bits')})
```

**Boundary-layer protection.** Embedding-adjacent and logit-adjacent layers
are noticeably more bit-sensitive than the middle. Keeping the outer
layers at FP16 costs little memory and recovers most of the small PPL
delta:

```python
TinyKVCache.for_model(
    model, compression="balanced", protect_first=2, protect_last=2,
)
```

Use `for_model` (not the bare constructor) whenever you set `protect_last`
or pass negative indices — the cache needs to know
`model.config.num_hidden_layers`.

| arg | default | notes |
|---|---|---|
| `compression` | – | preset: `"safe"` \| `"balanced"` \| `"aggressive"` |
| `bits` | 4 | (advanced) symmetric bit width, used when no preset |
| `key_bits` / `value_bits` | – | (advanced) per-stream bit widths; override `bits` |
| `protect_first` / `protect_last` | 0 | keep this many boundary layers at FP16 |
| `protected_layers` | – | explicit layer indices (negatives allowed via `for_model`) |

## Command-line tools

Two ship with the package — both are thin wrappers; the actual work happens
in `TinyKVCache`.

```bash
# OpenAI-compatible HTTP server (stdlib http.server — no FastAPI dependency)
tiny-turboquant-server --model Qwen/Qwen2.5-3B-Instruct --bits 4

# Sliding-window WikiText-2 perplexity, writes a JSON report
tiny-turboquant-eval --model Qwen/Qwen2.5-3B-Instruct --key-bits 4 --value-bits 2
```

Server endpoints: `POST /v1/chat/completions`, `GET /v1/models`, `GET /health`.

## The algorithm in three lines

```
1. for keys:   per-channel asymmetric affine quant (min/max along token axis)
               within a block of 128 tokens — absorbs outlier channels
2. for values: per-token asymmetric affine quant (min/max along channel axis)
               within the same block — matches the granularity attention
               actually reads at
3. keep the most recent N tokens (default 128) uncompressed in an FP16
               residual window
```

This matches the KIVI-2 scheme. Earlier versions used random-rotation +
a closed-form Beta-distribution codebook from the TurboQuant paper; that
path was elegant but assumed a marginal distribution real attention
streams don't have, and degraded badly at low bits. Empirical per-channel
scaling is both simpler and substantially more accurate on real models.

## Measured quality

WikiText-2-raw test split, sliding window 2048 / stride 1024, first 8 chunks:

| model | preset | PPL | Δ vs fp16 | storage savings |
|---|---|---|---|---|
| Qwen2.5-0.5B | fp16 baseline       | 13.671 | —     | 1.0× |
| Qwen2.5-0.5B | `safe` (4K / 4V)    | 13.949 | +2.0% | 3.1× |
| Qwen2.5-0.5B | `balanced` (4K / 2V)| 17.344 | +27%  | 3.9× |
| Qwen2.5-0.5B | `aggressive` (3K / 2V)| 19.466 | +42% | 4.3× |
| Qwen2.5-3B   | fp16 baseline       |  8.156 | —     | 1.0× |
| Qwen2.5-3B   | `safe` (4K / 4V)    |  8.258 | +1.3% | 3.3× |
| Qwen2.5-3B   | `balanced` (4K / 2V)| 11.148 | +37%  | 4.1× |
| Qwen2.5-3B   | `aggressive` (3K / 2V)| 12.026 | +47% | 4.7× |

`safe` is the near-lossless setting and scales well — the PPL gap shrinks
from 2.0% on 0.5B to 1.3% on 3B. `balanced` and `aggressive` degrade
because 2-bit V only has 4 levels per token — the published KIVI-2
tradeoff. Storage savings include per-block scale/zero metadata (one
fp32 scale + zero per channel for K, per token for V).

The MSE-only variant (`TinyQuantizer`) is what the cache uses. A two-stage
MSE + QJL variant (`TinyQuantizerIP`) ships for completeness but is
deprecated for attention — softmax exponentially amplifies the
JL-projection noise.

## CUDA fast path

A fused dequant-then-attention CUDA kernel is planned for the new
per-channel-K / per-token-V layout. The kernel in `cuda/` was written
against the old rotation-based path and is not used by the current
cache; dequant runs in PyTorch on the read path until the rewrite
lands.

## Where it earns its keep, where it doesn't

Useful when KV memory is the bottleneck — long contexts on a single GPU,
many concurrent serving sessions, or running one model size larger by
buying back VRAM from the cache. Not useful for short contexts (< 1k
tokens; the cache is already small), for hybrid / recurrent architectures
that don't keep a standard KV cache (Mamba, RWKV), or for tasks that need
bit-exact reproducibility.

## References

- KIVI: [A Tuning-Free Asymmetric 2bit Quantization for KV Cache](https://arxiv.org/abs/2402.02750)
  (Liu et al., ICML 2024). The current cache uses the KIVI-2 layout:
  per-channel K, per-token V.
- TurboQuant: [Online Vector Quantization with Near-optimal Distortion Rate](https://arxiv.org/abs/2504.19874)
  (Zandieh et al., ICLR 2026). Earlier versions of this package
  implemented the rotation + Beta-codebook scheme from this paper;
  it ships as `TinyQuantizer` for standalone vector quantization but
  is no longer used by `TinyKVCache`.

Architecture deep-dive: [docs/ARCHITECTURE.md](docs/ARCHITECTURE.md).

## License

MIT — see [LICENSE](LICENSE).
