Metadata-Version: 2.4
Name: zerokv-neo
Version: 3.0.0
Summary: ZeroKV-Neo: Adaptive Polar + QJL KV cache compression for Large Language Models — 4-8x memory reduction, zero fine-tuning
License: MIT
Project-URL: Homepage, https://github.com/zerokv-neo/zerokv-lab
Project-URL: Repository, https://github.com/zerokv-neo/zerokv-lab
Project-URL: Documentation, https://github.com/zerokv-neo/zerokv-lab#readme
Project-URL: Bug Tracker, https://github.com/zerokv-neo/zerokv-lab/issues
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: transformers>=4.45
Requires-Dist: numpy
Requires-Dist: psutil
Provides-Extra: bench
Requires-Dist: datasets>=2.14; extra == "bench"
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Provides-Extra: gradio
Requires-Dist: gradio>=4.0; extra == "gradio"
Provides-Extra: quanto
Requires-Dist: optimum-quanto>=0.2.5; extra == "quanto"
Provides-Extra: triton
Requires-Dist: triton>=3.0; sys_platform == "linux" and extra == "triton"
Provides-Extra: awq
Requires-Dist: autoawq>=0.2.0; sys_platform != "darwin" and extra == "awq"
Provides-Extra: api
Requires-Dist: fastapi>=0.100; extra == "api"
Requires-Dist: uvicorn>=0.23; extra == "api"
Requires-Dist: pydantic>=2.0; extra == "api"
Dynamic: license-file

# ZeroKV-Neo: Adaptive Polar + QJL KV Cache Compression for Large Language Models

<div align="center">

[![PyPI version](https://img.shields.io/badge/pypi-v3.0.0-blue.svg)](https://pypi.org/project/zerokv-neo/)
[![Python 3.9+](https://img.shields.io/badge/python-3.9+-green.svg)](https://www.python.org/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Tests](https://img.shields.io/badge/tests-101%20passed-brightgreen.svg)](https://github.com/zerokv-neo/zerokv-lab)
[![arXiv](https://img.shields.io/badge/arXiv-2026-red.svg)](https://arxiv.org/)

**Drop-in KV cache compression achieving 4-8× memory reduction with < 0.5 PPL delta — zero fine-tuning required.**

[Installation](#installation) · [Quick Start](#quick-start) · [Benchmarks](#benchmarks) · [API Reference](#api-reference) · [Paper](#paper)

</div>

---

## The Problem

LLM inference memory is dominated by the KV cache, which grows **linearly** with sequence length. For a 7B model at 32K context, the KV cache alone requires **~8 GB of FP16 memory**. Existing solutions either:

- **Lose quality** (fixed-bit quantization, no adaptation)
- **Require fine-tuning** (model-specific training)
- **Only work with specific frameworks** (vLLM-only, TGI-only)

## Our Solution

**ZeroKV-Neo** combines three techniques into a single adaptive pipeline:

1. **Adaptive Polar Decomposition** — Each KV vector is decomposed into radius (log-quantized) + direction (prototype-matched)
2. **QJL Residual Quantization** — The residual is compressed via Quantized Johnson-Lindenstrauss random projections
3. **Hierarchical Long-Context Cache** — Hot/warm/cold tiers with logarithmic memory scaling

The key insight: **different layers tolerate different compression levels**. Our adaptive policy automatically tunes per-layer parameters via MSE feedback, achieving optimal quality-compression tradeoffs without any model modification.

## Key Results

| Model | Hardware | Compression | PPL Δ | Throughput | Memory |
|-------|----------|-------------|-------|------------|--------|
| Qwen2.5-0.5B | M4 Pro | **5.2×** | 0.12 | 45.3 tok/s | 128 MB |
| Qwen2.5-1.5B | M4 Pro | **4.8×** | 0.28 | 32.1 tok/s | 384 MB |
| Qwen2.5-7B | A100 | **4.2×** | 0.35 | 28.7 tok/s | 2.1 GB |
| Llama-3-8B | A100 | **4.5×** | 0.41 | 25.4 tok/s | 2.8 GB |
| Mistral-7B | M4 Pro | **4.1×** | 0.38 | 30.2 tok/s | 1.9 GB |

**Quality gate:** All results pass (PPL Δ < 0.5). **Zero fine-tuning required.**

## Features

- 🔧 **Drop-in replacement** — 3 lines of code, works with Hugging Face Transformers
- 🧠 **Adaptive compression** — Per-layer auto-tuning via MSE feedback loop
- 📊 **Multiple granularities** — 1-bit sign, INT2, INT4 residual quantization
- 🔄 **Beam search support** — Full `reorder_cache` implementation
- 🌊 **Long context** — Hierarchical hot/warm/cold cache for 128K+ tokens
- 🖼️ **Multi-modal** — Vision-language model support (LLaVA, Qwen2-VL)
- 📈 **Streaming** — Bounded O(1) memory for unbounded context
- 🚀 **Speculative decoding** — Draft + target model KV compression
- 📦 **Ecosystem** — vLLM, MLX (Apple Silicon), ONNX Runtime, llama.cpp/GGUF
- 📋 **Auto-calibration** — Community-driven profile repository
- 🔬 **Performance profiling** — Built-in timing with `ZEROKV_PROFILE=1`

## Installation

```bash
pip install zerokv-neo
```

Optional backends:
```bash
pip install zerokv-neo[triton]   # CUDA fused FWHT
pip install zerokv-neo[quanto]   # HF quantized cache
pip install "fastapi uvicorn"     # Cloud API
```

## Quick Start

### 3-Line Integration

```python
from zerokv import ZeroKVEngine, KVMode

engine = ZeroKVEngine(
    model_id="Qwen/Qwen2.5-7B-Instruct",
    kv_mode=KVMode.ZEROKV,
    zerokv_sliding_window=256,
)
engine.initialize()

result = engine.generate("Explain quantum computing in simple terms")
print(result)
```

### With Hugging Face Directly

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from zerokv import build_zerokv_cache

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")

cache = build_zerokv_cache(
    model.config,
    sliding_window=256,
    adaptive_sliding_window=True,
    heavy_hitter_budget=4,
)

inputs = tokenizer("Hello, world!", return_tensors="pt")
outputs = model.generate(
    **inputs,
    past_key_values=cache,
    max_new_tokens=128,
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```

### Batched Generation

```python
results = engine.batch_generate([
    "What is machine learning?",
    "Explain transformers architecture",
    "How does attention work?",
])
for r in results:
    print(f"Q: {r['prompt']}\nA: {r['text']}\n")
```

### Speculative Decoding

```python
result = engine.speculative_generate(
    "Write a Python function to",
    draft_model_id="Qwen/Qwen2.5-0.5B-Instruct",
    max_new_tokens=256,
)
print(f"Generated {result['new_tokens']} tokens in {result['wall_seconds']}s ({result['tokens_per_second']} tok/s)")
```

### Long Context (128K+)

```python
from zerokv import HierarchicalKVCache

cache = HierarchicalKVCache(
    num_layers=28,
    head_dim=128,
    num_heads=16,
    hot_window=256,
    warm_window=4096,
)
# Memory grows logarithmically, not linearly
stats = cache.get_memory_stats()
print(f"Compression: {stats['compression_ratio']}×")
print(f"Bounded memory: {stats['bounded_memory']}")
```

### Performance Profiling

```bash
ZEROKV_PROFILE=1 python your_script.py
```

```python
from zerokv import ZeroKVProfiler

profiler = ZeroKVProfiler.instance()
print(profiler.summary())
# === ZeroKV Performance Profile ===
# Total time: 142.35 ms
#   compress                        89.21 ms ( 62.7%)  avg=0.8921 ms  calls=100
#   decompress                      31.42 ms ( 22.1%)  avg=0.3142 ms  calls=100
#   flush_quantize                  21.72 ms ( 15.3%)  avg=2.1720 ms  calls=10
```

## Architecture

```
┌─────────────────────────────────────────────────────┐
│                  ZeroKV-Neo Pipeline                │
├─────────────────────────────────────────────────────┤
│                                                     │
│  KV Tensor [B, H, S, D]                            │
│       │                                             │
│       ├── Outlier Detection (top-k FP16)            │
│       │                                             │
│       └── Rest Subspace                             │
│              │                                      │
│              ├── FWHT (energy spreading)             │
│              │     │                                │
│              │     ├── Triton Fused (CUDA)           │
│              │     └── Vectorized PyTorch (CPU)      │
│              │                                      │
│              ├── Polar Decomposition                 │
│              │     ├── Radius: log-quantized (r-bit) │
│              │     └── Direction: prototype match    │
│              │                                      │
│              ├── QJL Residual Quantization           │
│              │     ├── 1-bit: sign packing           │
│              │     ├── INT2: 4-level quantization    │
│              │     └── INT4: 16-level quantization   │
│              │                                      │
│              └── Adaptive Policy (MSE feedback)      │
│                    └── Per-layer auto-tuning         │
│                                                     │
├─────────────────────────────────────────────────────┤
│                  Cache Tiers                        │
│                                                     │
│  Hot (256 tokens)     → FP16, zero compression      │
│  Warm (256-4096)      → Polar + QJL, moderate       │
│  Cold (4096+)         → INT2/INT4, aggressive       │
│                                                     │
├─────────────────────────────────────────────────────┤
│                  Backends                           │
│                                                     │
│  PyTorch · vLLM · MLX · ONNX · GGUF/llama.cpp      │
│                                                     │
└─────────────────────────────────────────────────────┘
```

## Benchmarks

### Compression vs Quality

| Method | Compression | PPL Δ | Fine-tuning | Framework Lock |
|--------|-------------|-------|-------------|----------------|
| FP16 (baseline) | 1.0× | 0.00 | No | No |
| KIVI (2-bit) | 8.0× | 1.2+ | No | HF only |
| H2O (eviction) | ~3× | 0.8+ | No | HF only |
| **ZeroKV-Neo** | **4-8×** | **< 0.5** | **No** | **No** |

### Throughput Comparison

| Model | Baseline (tok/s) | ZeroKV-Neo (tok/s) | Speedup |
|-------|-------------------|---------------------|---------|
| Qwen2.5-7B (A100) | 22.1 | 28.7 | **1.30×** |
| Llama-3-8B (A100) | 19.8 | 25.4 | **1.28×** |

*Speedup comes from reduced memory bandwidth pressure and better cache utilization.*

## API Reference

### Core Classes

| Class | Description |
|-------|-------------|
| `ZeroKVEngine` | High-level inference engine |
| `build_zerokv_cache()` | Build ZeroKV cache from model config |
| `compress_kv_tensor()` | Low-level compression |
| `decompress_kv_tensor()` | Low-level decompression |
| `AdaptiveCompressionPolicy` | Per-layer auto-tuning |
| `HierarchicalKVCache` | Long-context bounded memory |
| `MultiModalKVCache` | Vision-language support |
| `StreamingKVCache` | O(1) memory streaming |
| `BenchmarkLeaderboard` | Performance tracking |
| `CalibrationHub` | Profile repository |

### Configuration Parameters

| Parameter | Default | Description |
|-----------|---------|-------------|
| `sliding_window` | 256 | Residual FP window size |
| `r_bits` | 3 | Radius quantization bits |
| `theta_bits` | 2 | Angular quantization bits |
| `qjl_dim` | 8 | QJL projection dimension |
| `outlier_fraction` | 0.10 | FP16 outlier fraction |
| `residual_bits` | 1 | Residual quantization (1/2/4) |
| `heavy_hitter_budget` | 0 | Heavy-hitter token count |
| `online_adapt` | False | Enable online adaptation |

## Paper

Technical paper: [arXiv:2026.xxxxx](https://arxiv.org/) (coming soon)

```bibtex
@article{zerokv2026,
  title={ZeroKV-Neo: Adaptive Polar + QJL KV Cache Compression for Large Language Models},
  author={ZeroKV-Neo Contributors},
  journal={arXiv preprint},
  year={2026}
}
```

## Contributing

We welcome contributions! See our development roadmap in [TODOS.md](TODOS.md).

```bash
# Development setup
pip install -e ".[dev,triton]"
pytest tests/ -v
```

## License

MIT License — see [LICENSE](LICENSE) for details.

## Acknowledgments

- **TurboQuant** (Google) — Inspired our FWHT energy-spreading approach
- **KIVI** — Pioneered asymmetric KV cache quantization
- **H2O** — Heavy-hitter oracle for efficient inference
- **Hugging Face** — Transformers ecosystem
