Skip to main content

Low-bit vector and KV-cache compression research toolkit for PyTorch

Project description

Tiny TurboQuant

v0.5.0: real LLM KV-cache benchmarking

This release adds repeatable KV-cache benchmark utilities for real Hugging Face causal language models.

New capabilities:

  • tiny-tq kv-bench for DynamicCache vs HybridTurboQuantKVCache comparison
  • prompt modes: short, medium, long, and stress
  • safe / balanced / aggressive / quality-headwise KV-cache presets
  • automatic outlier count selection with "auto"
  • KV-cache memory estimator via tiny-tq kv-estimate
  • JSON and Markdown benchmark reports
  • baseline vs compressed output comparison
  • generation-drift diagnostics such as first divergence and optional logit KL

This release continues to focus on memory compression and quality measurement. It does not claim production inference acceleration.

Example:

tiny-tq kv-estimate --layers 24 --kv-heads 8 --head-dim 128 --seq-len 4096 --batch-size 4

tiny-tq kv-bench \
  --model Qwen/Qwen2.5-0.5B-Instruct \
  --preset safe \
  --prompt-mode short \
  --max-new-tokens 8 \
  --report-json kv_report.json \
  --report-md kv_report.md

Tiny TurboQuant is a lightweight PyTorch research toolkit for low-bit vector compression, compressed RAG retrieval, and KV-cache compression experiments.

Version 0.4.0 focuses on the RAG/vector-index roadmap:

  • RAGCompressedIndex for compressed retrieval experiments
  • document chunking helpers
  • JSONL / folder text ingestion utilities
  • compressed vector index save/load
  • RAG index save/load
  • retrieval metrics: recall@k, precision@k, MRR, nDCG, overlap, label-match ratio
  • CLI entry point: tiny-tq rag-bench
  • existing hybrid KV-cache and paged-attention research utilities from v0.3.x

Important limitation

This package demonstrates packed memory compression and memory-quality benchmarking. It is not a production compressed-attention engine. Hugging Face generation still receives dense K/V tensors. The paged attention utility dequantizes page-by-page and avoids one full dense cache tensor, but it is not a fused CUDA/Triton kernel. Real latency gains require fused kernels or serving-engine integration.

Do not use this package to claim training acceleration, fine-tuning memory reduction, production legal/medical QA readiness, drop-in vLLM replacement, exact nearest-neighbor search, or faster LLM inference.

Install

pip install tiny-turboquant

Optional demo dependencies:

pip install "tiny-turboquant[demos]"

Compressed RAG index from precomputed embeddings

import torch
from tiny_turboquant import RAGCompressedIndex

texts = [
    "Embedding compression can reduce vector-store memory in RAG systems.",
    "KV cache stores Key and Value tensors during LLM generation.",
]
embeddings = torch.randn(len(texts), 384)

index = RAGCompressedIndex.from_embeddings(
    texts,
    embeddings,
    bits=4,
    store_original_for_rerank=True,
)

results = index.search(embeddings[0], top_k=1, rerank_top_k=2)
print(index.memory_report())
print(results[0].text)

Compressed RAG index from documents

Requires sentence-transformers:

from tiny_turboquant import RAGCompressedIndex

docs = [
    "RAG systems retrieve relevant document chunks and pass them to an LLM.",
    "Compressed vector indexes reduce embedding memory usage.",
]

index = RAGCompressedIndex.from_documents(
    docs,
    embedding_model="sentence-transformers/all-MiniLM-L6-v2",
    bits=4,
    chunk_size=500,
    overlap=50,
)

results = index.search("How do we reduce vector-store memory?", top_k=3)

Save and load

index.save("rag_index.ttq")
loaded = RAGCompressedIndex.load("rag_index.ttq")

Retrieval metrics

from tiny_turboquant import recall_at_k, mrr_at_k, ndcg_at_k

retrieved = ["doc-1", "doc-2", "doc-3"]
relevant = {"doc-2", "doc-5"}

print(recall_at_k(retrieved, relevant, k=3))
print(mrr_at_k(retrieved, relevant, k=3))
print(ndcg_at_k(retrieved, relevant, k=3))

CLI

tiny-tq version

Synthetic RAG benchmark:

tiny-tq rag-bench --synthetic --bits 4 --top-k 10 --rerank-top-k 50

JSONL benchmark:

tiny-tq rag-bench \
  --input-jsonl docs.jsonl \
  --text-field text \
  --query "How can we reduce vector-store memory?" \
  --bits 4 \
  --top-k 10 \
  --rerank-top-k 50

Hybrid KV-cache usage

from tiny_turboquant import HybridTurboQuantKVCache

cache = HybridTurboQuantKVCache(
    key_bits=6,
    value_bits=4,
    key_outlier_bits=8,
    value_outlier_bits=8,
    n_key_outliers=32,
    n_value_outliers=16,
    key_recent_window=128,
    value_recent_window=64,
    per_layer_calibration=True,
    per_head_calibration=True,
)

Compressed vector index usage

import torch
from tiny_turboquant import CompressedVectorIndex

emb = torch.randn(10_000, 384)
index = CompressedVectorIndex(bits=4, store_original_for_rerank=True).add(emb)
results = index.search(emb[0], top_k=5, rerank_top_k=100)

print(index.compression_ratio())
print(results[0])

Project position

Current focus:

  • memory compression
  • retrieval quality measurement
  • RAG/vector-index experiments
  • KV-cache compression research

Future direction:

  • real workload RAG benchmarks
  • FAISS/vector database integration
  • long-context real-model KV-cache benchmarks
  • fused dequant + attention kernels
  • serving-engine integration experiments

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

tiny_turboquant-0.5.0.tar.gz (58.5 kB view details)

Uploaded Source

Built Distribution

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

tiny_turboquant-0.5.0-py3-none-any.whl (63.1 kB view details)

Uploaded Python 3

File details

Details for the file tiny_turboquant-0.5.0.tar.gz.

File metadata

  • Download URL: tiny_turboquant-0.5.0.tar.gz
  • Upload date:
  • Size: 58.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.0

File hashes

Hashes for tiny_turboquant-0.5.0.tar.gz
Algorithm Hash digest
SHA256 47cf5e940785fe228ee6d5e21b98cf016a188076c784b790d332dbc7e2d30c89
MD5 dc9c992221ab577c7390a65063316531
BLAKE2b-256 75b415476a5a1b6d4df6167cf0258de1e6885513e9ac720fb296e7bb52450530

See more details on using hashes here.

File details

Details for the file tiny_turboquant-0.5.0-py3-none-any.whl.

File metadata

File hashes

Hashes for tiny_turboquant-0.5.0-py3-none-any.whl
Algorithm Hash digest
SHA256 170b9ff11ed5e9a5e611b9258dd80c7bc39f11d45c61e3ec79d23b076592c709
MD5 023f6bd7b54771f0133f0777cb895ab6
BLAKE2b-256 4c559bfc1acb5848eb4ecdb1f90c91e317dd6f172f3be240d6faac1de105e07d

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