tq-kv

Pure Rust TurboQuant KV cache compression. CUDA. AVX2 SIMD. C FFI. crates.io.

Implementation of Google's TurboQuant (ICLR 2026) with the 3-Fix framework that enables aggressive key compression (4-bit, 7.5x) on GGUF quantized models -- where symmetric K compression produces catastrophic output.

Now with Pre-RoPE key quantization (34-59% less PPL gap), KV Compaction (up to 25x token reduction), TriAttention eviction for constant-memory KV cache, per-token mean removal (+40% attention quality at 2-bit), and multi-arch CUDA (Turing → Hopper).

---

Current performance

RTX 3080 10 GB, Qwen2.5-7B-Instruct Q4_K_M, own CUDA kernels, no external inference dependency. Three-run mean at 2026-04-19, warmup discarded, TTFT separated.

• llama.cpp head-to-head gap narrowed from 4.0× at v0.6.0 to 1.83× at v0.7.0.
• A pre-push hook enforces PPL thresholds (Std ≤ 12.66, TQ ≤ 13.72) on every kernel change.
• Known limit on 10 GB cards: prefill attention is O(N²), so the CLI OOMs past ~2 K prompt length regardless of KV compression. See scripts/bench-long-context.sh for the forcing measurement.

Full matrix (Llama 3.1 8B, Mistral 7B, multi-context, reproducible CLIs): BENCHMARKS.md.

What's next (v0.8.0)

Everything below is gated and unscoped; priorities in order:

• Megakernel wiring. v0.7.0 ships the 6-phase persistent kernel with all stubs promoted to real bodies and the >48 KB shmem opt-in; the last mile is routing TQ_MEGAKERNEL=1 through model.rs::forward. Architect projection +~10 %; will be measured, not claimed, at merge.
• Flash-attention for PREFILL. We only have flash_decode (single-token). Prefill attention's N×N score tensor is the real long-context OOM cliff; a flash-prefill port unlocks ≥ 4 K contexts on 10 GB cards.
• EAGLE acceptance debug. Sprint 1–3 shipped the draft loader, tree attention, and acceptance probe in v0.7.0. The draft forward is bit-correct against a numpy reference but acceptance sits at ~6.7 % — a torch-based reference covering non-zero inputs is the likely unlock. Week-scale.
• Metal backend kickoff (blocked on a collaborator with Apple Silicon).

Full plan: see the status:planned nodes in the project knowledge graph.

---

The Compound Error Problem

GGUF quantized models (Q4_K_M) already have weight quantization noise. Compressing KV cache on top introduces compound error through softmax. tq-kv solves this with a multi-stage pipeline:

Key innovations for GGUF compound error:

• 3-Fix framework: sink tokens (FP16), current token (lossless POQ), cache reset
• Pre-RoPE quantization: position-independent per-channel stats, better codebook fit
• Per-token mean removal: softmax shift-invariance exploited, +40% attention quality at 2-bit
• Auto-calibration: zero-config — calibrates on first use (256 samples, 2-3 seconds)

---

Measured Quality (Honest Numbers)

All PPL measured with automated tq perplexity / tq ablate.

GGUF Q4_K_M Models (compound error)

Qwen 2.5 7B Q4_K_M, modern English, skip=4, sink=4

Pre-RoPE Improvement Across Context Lengths

Value Compression is Nearly Free

FP16 Models (no compound error)

Qwen 2.5 0.5B FP16 safetensors, wikitext-2

Compression Quality (per-layer)

NIAH (Needle-In-A-Haystack)

---

Quick Start

Install

[dependencies]
tq-kv = "0.6"

CLI (tq-engine)

# Pull a model
tq pull qwen2:7b

# Chat with TurboQuant compression
tq chat qwen2:7b --turbo-quant

# Chat with Pre-RoPE (best quality)
TQ_PRE_ROPE=1 tq chat qwen2:7b --turbo-quant

# Start OpenAI-compatible API server
tq serve --model qwen2:7b --turbo-quant --port 11435

# Evaluate perplexity
tq perplexity --model qwen2:7b eval.txt --turbo-quant

# Calibrate (optimal codebook + rotation from real activations)
tq calibrate qwen2:7b --text calibration_data.txt

# Run ablation study
tq ablate qwen2:7b --file eval.txt --quick --output results.csv

Library API

use tq_kv::*;

let config = TurboQuantConfig::balanced(); // 4-bit, 7.5x compression
let dim = 128;

// Batch compress
let compressed = compress_keys(&kv_data, dim, &config);
println!("Ratio: {:.1}x", compressed.compression_ratio());

// Fused attention -- no decompression, AVX2+FMA SIMD
let signs = hadamard::generate_signs(dim, config.rotation_seed);
let centroids = codebook::get_centroids(config.bits);
let rotated_q = pre_rotate_query_with_signs(&query, &signs);
let scores = fused_attention_scores(&rotated_q, &compressed, centroids, scale);

// KV Compaction -- reduce token count
let compacted = compaction::compact_head(&keys, &values, &queries,
    seq_len, n_queries, dim, target_size);
// compacted.keys, compacted.beta, compacted.values

---

How It Works

Input KV vector (from GGUF Q4_K_M model)
    |
[0] Pre-RoPE capture (optional)            O(1)
    |  Compress BEFORE RoPE for position-independent stats
    |
[1] Channel bias subtraction (calibrated)   O(d)
    |  Remove weight-quantization artifacts
    |
[2] Per-token mean removal                   O(d)
    |  Softmax shift-invariant: attention ignores mean
    |  Frees codebook precision (+40% cosine sim @ 2-bit)
    |
[3] Randomized Hadamard Transform            O(d log d)
    |  Decorrelates outliers -> coordinates ~ Gaussian
    |
[4] Per-channel adaptive sigma (KIVI-style)  O(d)
    |  Per-dimension variance from calibration
    |  Lloyd-Max codebook + norm correction
    |
[5] SRHT QJL Error Correction (optional)     O(d log d)
    |  Structured Hadamard projection (+4.5 dB SNR)
    |
Output: packed indices + corrected norm + mean
        7.5x compression at 4-bit (keys only)

3-Fix for GGUF Models

Fix 1: Sink tokens (first 4) stay FP16     -> -81% attention error
Fix 2: Current token = lossless (POQ)       -> highest-impact position protected
Fix 3: Cache reset per conversation         -> prevents cross-contamination

6-Segment Attention

[sink FP16] [cold decayed] [compacted + beta] [hot compressed] [current FP16]
     |            |                |                   |                |
  Always       Temporal        Attention-           Per-head          POQ
  lossless     decay           matching             adaptive         lossless
  (Fix 1)                      reduction            bitwidth         (Fix 2)
                                    ^
                              TriAttention eviction scores all segments,
                              keeps top-B tokens, maintains fixed budget

TriAttention: Constant-Memory KV Cache

Based on TriAttention (Mao et al., 2026). Pre-RoPE Q/K vectors concentrate around fixed centers, enabling cheap trigonometric importance scoring without full attention computation.

Orthogonal to TurboQuant: TriAttention decides which tokens to keep (eviction), TurboQuant decides how to compress them (quantization). Combined: fixed-size KV cache at any context length.

# TQ+TriAttention is ON by default with --turbo-quant (requires calibration)
tq chat qwen2:7b --turbo-quant

# Disable TriAttention (TQ-only mode)
TQ_TRIATTN=0 tq chat qwen2:7b --turbo-quant

# Custom budget
TQ_TRIATTN_BUDGET=256 tq chat qwen2:7b --turbo-quant

Memory Projection (Qwen2.5-7B, TQ 4-bit + TriAttention, budget=128)

Speed vs Budget (RTX 3080, Qwen2.5-7B, 100 tokens)

What Fits in 10GB VRAM (5GB available for KV cache)

---

Configuration Guide

Recommended Configs by Use Case

Environment Variables

---

VRAM Savings

With KV Compaction: effective compression reaches 100-400x. With TriAttention: constant-memory KV cache -- up to 1,560x at 128K context.

---

SRHT QJL Performance (32K vectors, d=128, release)

---

Full Product: tq-engine

tq-kv powers tq-engine -- "Rust's Ollama" with TurboQuant compression:

tq pull qwen2:7b          # download from HuggingFace
tq serve --turbo-quant     # OpenAI-compatible API (SSE streaming)
tq chat qwen2:7b           # terminal chat

Web UI at localhost:11435. Works with ChatBox and Open WebUI.

4 validated models: Qwen2.5 7B/0.5B, Llama 3.1 8B, Mistral 7B. Auto-detected from GGUF metadata.

3-Way Benchmark

tq bench qwen2:7b                    # Standard vs TQ vs TQ+TriAttention
tq perplexity -m qwen2:7b --compare eval.txt  # 3-way PPL comparison
scripts/ppl-check.sh                  # Regression CI (9 checks, tight thresholds)

---

GPU Inference Performance

RTX 3080 10GB, Qwen 2.5 7B Q4_K_M, own CUDA kernels (no candle dependency). Full reproducible matrix in BENCHMARKS.md.

4 validated models: Qwen2.5 7B/0.5B, Llama 3.1 8B, Mistral 7B. Auto-calibration on first use.

Custom CUDA Kernel Stack (78 kernels, 16 files, 4.4K lines)

• Multi-arch: sm_75 (Turing) → sm_90 (Hopper), runtime GPU detect, per-arch PTX
• Butterfly reduce: __shfl_xor_sync broadcast to all threads (no broadcast hacks)
• cp.async pipeline: double-buffered qmatmul on Ampere+ (sm_80+), sync fallback on Turing
• Fused layer kernels: RmsNorm + Q4K QKV + bias, gateup + SiLU, down + residual
• Q4K/Q6K/Q8_0 fused matvec: __ldg read-only cache, dequant + dot in single kernel
• Flash decode v2: split-KV parallelism for long context (>256 tokens), online softmax
• TQ fused attention: compressed KV → attention score without decompression (moat kernel)
• CUDA Graph replay: short context (<256) as single GPU operation, auto-fallback to eager mode

Optimization History

---

Benchmark

cargo run --release -p tq-kv --bin tq-kv-bench

Full results: BENCHMARK.md

Paper

Our work:

• "TurboQuant on Quantized Models: Solving Compound Quantization Error with Pre-RoPE Compression and KV Compaction" -- BHI Research (2026)
• 3-Fix framework, SRHT QJL (115x speedup), Pre-RoPE quantization, KV Compaction, Adaptive QJL

Original:

• Zandieh, Daliri, Hadian, Mirrokni. "TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate." ICLR 2026. arXiv:2504.19874
• Mao, Lin, Huang et al. "TriAttention: Efficient Long Reasoning with Trigonometric KV Compression." 2026. arXiv:2604.04921

License

MIT OR Apache-2.0