Skip to main content

ZeroKV-Neo: Adaptive Polar + QJL KV cache compression for Large Language Models — 4-8x memory reduction, zero fine-tuning

Project description

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

PyPI version Python 3.9+ License: MIT Tests arXiv

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

Installation · Quick Start · Benchmarks · API Reference · Paper


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

pip install zerokv-neo

Optional backends:

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

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

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

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

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+)

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

ZEROKV_PROFILE=1 python your_script.py
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 (coming soon)

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

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

License

MIT License — see 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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

zerokv_neo-3.0.2.tar.gz (80.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

zerokv_neo-3.0.2-py3-none-any.whl (75.7 kB view details)

Uploaded Python 3

File details

Details for the file zerokv_neo-3.0.2.tar.gz.

File metadata

  • Download URL: zerokv_neo-3.0.2.tar.gz
  • Upload date:
  • Size: 80.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.9.6

File hashes

Hashes for zerokv_neo-3.0.2.tar.gz
Algorithm Hash digest
SHA256 d04e0d3662b46db3870b79de9712dffc73c2c40baedcf4d04ed94831694f3efb
MD5 082296ee9420e63ef4c6acfe47fa23ce
BLAKE2b-256 72a516cd6a735cfa9068f69fd2fa197e0492f6ef842690bca6195ec533e38d4b

See more details on using hashes here.

File details

Details for the file zerokv_neo-3.0.2-py3-none-any.whl.

File metadata

  • Download URL: zerokv_neo-3.0.2-py3-none-any.whl
  • Upload date:
  • Size: 75.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.9.6

File hashes

Hashes for zerokv_neo-3.0.2-py3-none-any.whl
Algorithm Hash digest
SHA256 1ba5e84589fa8bdb231cf50203517ab1ae429cc044819ce32dd50a4c6f904418
MD5 c63c5c408cb403ce08b3f1a6e7f713f5
BLAKE2b-256 9466e33a9cf7efafdbd076f0eb8fb1b99e7fd698289ad6f918ea4c8f507e9d4e

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page