tiny-turboquant 0.11.3.1

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

tiny-turboquant

tiny-turboquant is a PyTorch/Triton research toolkit for studying low-bit vector compression, compressed RAG retrieval, KV-cache memory reduction, and experimental compressed decode-attention kernels.

The package is built for measurement, learning, and reproducible research. It is not a drop-in serving engine, not a vLLM replacement, and not a production LLM inference accelerator.

What it supports

• Low-bit vector quantization and packed storage
• Compressed vector indexes with optional reranking
• Compressed RAG index experiments
• Retrieval metrics such as recall, precision, MRR, nDCG, overlap, and label match
• KV-cache memory estimation
• Hybrid KV-cache research objects for Hugging Face-style experimentation
• Page-wise attention diagnostics
• Compressed KV page layouts
• Rotate-Q / rotated-key compressed attention diagnostics
• Residual-corrected compressed KV layout experiments
• Triton compressed decode-attention research kernels
• Split-K / sequence-parallel compressed decode-attention diagnostics
• Fixed-cache repeated decode-loop benchmarks with setup amortization reporting
• Serving-capacity simulation for KV-memory planning

What it does not claim

Do not use this package to claim:

• production LLM inference acceleration
• full model-serving acceleration
• training or fine-tuning acceleration
• exact nearest-neighbor search
• drop-in replacement for FAISS, vLLM, TensorRT-LLM, or FlashAttention
• legal, medical, or safety-critical retrieval correctness
• universal speedups across hardware or workloads

The fastest paths are research microbenchmarks. Results must be interpreted with their reported boundaries: setup cost, payload preparation, CUDA Graph use, synthetic data, fixed-cache assumptions, and hardware-specific behavior.

Install

pip install tiny-turboquant

Optional demo dependencies:

pip install "tiny-turboquant[demos]"

Check the installed package:

tiny-tq version

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

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

index = RAGCompressedIndex.from_documents(
    texts,
    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))

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

Compressed vector index

import torch
from tiny_turboquant import CompressedVectorIndex

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

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

KV-cache memory estimation

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

Real-model KV-cache benchmark

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

This benchmark is for memory and quality diagnostics. It should not be treated as production serving performance.

Real-model KV attention benchmark

tiny-tq real-model-kv \
  --model Qwen/Qwen2.5-0.5B-Instruct \
  --prompt "Explain why KV-cache memory grows during long-context decoding." \
  --layer-index 0 \
  --query-source random-normalized \
  --max-prompt-tokens 2048 \
  --decode-steps 32 \
  --amortization-steps 32 128 256 512 1024 \
  --preset safe-layout \
  --page-size 256 \
  --device cuda \
  --dtype float16 \
  --tune-kernel \
  --tune-block-m-values 8 16 32 64 128 \
  --tune-num-warps \
  --tune-num-warps-values 1 2 4 8 \
  --split-k-slabs 2 4 8 16 \
  --cuda-graph \
  --graph-replays 100 \
  --report-json real_model_kv_report.json \
  --report-md real_model_kv_report.md

This extracts real Hugging Face past_key_values from a model forward pass, compresses one selected layer's K/V tensors, and benchmarks dense attention, SDPA, single-pass compressed decode attention, split-K compressed decode attention, and CUDA Graph replay over those real KV tensors.

Use --query-source random-normalized for stable quality diagnostics. --query-source last-key is useful for stress testing but can create degenerate quality metrics on some models. Use --kv-head-repeat N only for workload stress testing; it repeats extracted KV heads and is not a literal model shape when N > 1.

Boundary: the KV tensors are real, but the decode query is a model-shaped proxy. This is not full autoregressive model serving.

Layer-aware real-KV scan

tiny-tq real-model-kv-scan \
  --model Qwen/Qwen2.5-0.5B-Instruct \
  --prompt "Explain why KV-cache memory grows during long-context decoding." \
  --layers all \
  --max-prompt-tokens 2048 \
  --page-size 256 \
  --candidates 8,8 8,6 8,4 6,4 4,4,residual-affine \
  --boundary-layer-count 1 \
  --query-source random-normalized \
  --device cuda \
  --dtype float16 \
  --report-json real_kv_scan_report.json \
  --report-md real_kv_scan_report.md

Aliases are also registered:

tiny-tq real-kv-scan --synthetic
tiny-tq scan-kv --synthetic

This command scans selected KV-cache layers independently and reports k_mse, v_mse, k_cosine, v_cosine, attention_cosine, score_cosine, softmax_kl_mean, recommended_k_bits, recommended_v_bits, protect_layer, and first_divergence_risk. Boundary layers are protected by default because uniform compression across all layers is not a stable assumption.

Boundary: this is a layer-aware diagnostic over tested prompts, query proxies, thresholds, and candidate layouts. Validate with task metrics before claiming model-serving stability.

Hybrid KV-cache object

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

Page-wise attention diagnostic

tiny-tq page-attn-bench \
  --seq-len 1024 \
  --page-size 128 \
  --device cuda \
  --dtype float16

This measures dense attention, SDPA, and page-wise compressed attention behavior. It is a diagnostic, not a serving benchmark.

Layout quality sweep

tiny-tq layout-sweep \
  --seq-len 2048 \
  --page-size 256 \
  --heads 8 \
  --head-dim 64 \
  --device cuda \
  --dtype float16 \
  --bit-pairs 8,8 8,6 8,4 6,4 \
  --report-json layout_sweep.json \
  --report-md layout_sweep.md

Use this to compare memory savings and attention-quality impact across compressed layouts.

Residual-correction sweep

tiny-tq residual-sweep \
  --device cuda \
  --dtype float16 \
  --seq-len 2048 \
  --heads 8 \
  --head-dim 64 \
  --page-size 256 \
  --bit-pairs 8,6 6,4 4,4 \
  --report-json residual_sweep.json \
  --report-md residual_sweep.md

This tests affine vs residual-affine compressed KV layouts and reports quality metrics such as attention relative error, cosine similarity, QK score error, softmax KL, and inner-product bias.

Fused compressed decode benchmark

tiny-tq fused-decode-bench \
  --preset safe-layout \
  --seq-len 8192 \
  --page-size 256 \
  --heads 8 \
  --head-dim 64 \
  --device cuda \
  --dtype float16 \
  --tune-kernel \
  --tune-block-m-values 8 16 32 64 128 \
  --tune-num-warps \
  --tune-num-warps-values 1 2 4 8 \
  --cuda-graph \
  --graph-replays 100 \
  --report-json fused_report.json \
  --report-md fused_report.md

This measures the experimental compressed decode-attention path and compares it with dense attention and SDPA.

Split-K compressed decode benchmark

tiny-tq split-k-compare \
  --seq-lens 16384 32768 \
  --preset safe-layout \
  --page-size 256 \
  --heads 8 \
  --head-dim 64 \
  --device cuda \
  --dtype float16 \
  --tune-kernel \
  --tune-block-m-values 8 16 32 64 128 \
  --tune-num-warps \
  --tune-num-warps-values 1 2 4 8 \
  --cuda-graph \
  --graph-replays 100 \
  --split-k-slabs 2 4 8 16 \
  --report-json split_k_report.json \
  --report-md split_k_report.md

This compares single-pass compressed decode attention, projected split-K, Torch reference split-K, and measured Triton split-K paths. Treat this as a kernel research benchmark.

End-to-end fixed-cache decode benchmark

tiny-tq end-to-end-decode \
  --prompt-lens 16384 32768 \
  --decode-steps 32 \
  --amortization-steps 32 128 256 512 1024 \
  --preset safe-layout \
  --page-size 256 \
  --heads 8 \
  --head-dim 64 \
  --device cuda \
  --dtype float16 \
  --tune-kernel \
  --tune-block-m-values 8 16 32 64 128 \
  --tune-num-warps \
  --tune-num-warps-values 1 2 4 8 \
  --split-k-slabs 2 4 8 16 \
  --cuda-graph \
  --graph-replays 100 \
  --report-json e2e_decode_report.json \
  --report-md e2e_decode_report.md

This benchmark reports:

• setup cost
• payload preparation cost
• repeated decode-loop timing
• per-token timing
• token throughput estimate
• SDPA comparison
• quality vs dense attention
• amortized setup cost
• break-even decode length
• cached-layout reuse scenarios

This is still a fixed-cache synthetic diagnostic, not full model serving.

Serving memory simulation

tiny-tq serving-capacity \
  --gpu-memory-gb 24 \
  --model-weight-gb 8 \
  --layers 24 \
  --kv-heads 8 \
  --head-dim 128 \
  --avg-prompt-tokens 4096 \
  --avg-decode-tokens 256 \
  --preset balanced

Use this to estimate how compressed KV layouts may affect memory capacity for long-context serving scenarios.

How to read benchmark output

Prefer these fields:

• effective_memory_saved_pct
• cosine_similarity
• relative_error
• seconds_per_token
• tokens_per_second
• over_sdpa_per_token
• speedup_vs_sdpa_per_token
• total_setup_seconds
• break_even_decode_steps_full_setup
• break_even_decode_steps_payload_prepare_only
• boundary

A kernel-only win is not the same as production acceleration. Check whether setup cost, allocation, cache preparation, and full model execution are included.

Recommended project claim

Use this wording:

tiny-turboquant is a PyTorch/Triton research toolkit for compressed RAG retrieval, KV-cache memory analysis, and experimental compressed decode-attention benchmarking.

Avoid this wording:

tiny-turboquant accelerates production LLM inference.

Stabilization plan

The path to a stable research release is:

1. Keep the real-model KV benchmark stable across current Hugging Face cache formats.
2. Make CUDA Graph results first-class in the verdict fields, because normal Python decode loops can hide kernel behavior.
3. Keep benchmark claims separated into kernel-only, CUDA Graph, setup-amortized, and full-serving categories.
4. Add model-shape stress knobs such as KV-head repetition, while clearly labeling them as synthetic workload scaling.
5. Validate quality with finite-output checks before accepting any speed result.
6. Keep asymmetric K/V and layer-aware protection first-class: preserve key precision more aggressively than value precision, and protect sensitive boundary layers when moving toward lower-bit layouts. v0.11.3.1 adds real-model-kv-scan for this diagnostic loop.

Research ideas to investigate next

• Asymmetric K/V layouts: keep K at safer precision while compressing V more aggressively.
• Boundary-layer protection: keep first and last layers at higher precision when full-model integration arrives.
• Value-only compression studies: measure whether V can be compressed more aggressively without hurting attention quality for the target model family.
• Real generation hook: move from one-layer real KV attention diagnostics to full autoregressive model integration.
• Serving capacity study: translate KV memory savings into concurrent-user capacity under fixed GPU memory budgets.

Packaging note

The wheel installs the core tiny_turboquant package. Demo, benchmark, example, and test files are kept in the source distribution or repository for reference and are not installed as top-level Python packages.