Crate tq_kv 

tq-kv: Extreme KV Cache Compression for LLMs

Pure Rust implementation of Google’s TurboQuant algorithm (ICLR 2026). Compresses KV cache keys to 2-4 bits with up to 15x compression and 0.997 cosine similarity. Zero C/C++ dependencies.

Algorithm

1. Randomized Hadamard Transform — decorrelates outliers, O(d log d)
2. Lloyd-Max Codebook Quantization — optimal centroids for Gaussian, O(d)
3. Fused Attention — pre-rotate query, centroid table lookup (no decompress)

Quick Start

use tq_kv::{TurboQuantConfig, compress_keys, decompress_keys};

let config = TurboQuantConfig::extreme(); // 2-bit, ~15x compression
let head_dim = 128;
let kv_data: Vec<f32> = vec![0.1; head_dim]; // one vector

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

let restored = decompress_keys(&compressed, &config);

Incremental KV Cache

use tq_kv::*;

let config = TurboQuantConfig::extreme();
let dim = 128;
let signs = hadamard::generate_signs(dim, config.rotation_seed);

let mut cache = CompressedKeys::new_empty(config.bits, dim, config.rotation_seed);
let key = vec![0.1f32; dim];
let (packed, norm) = compress_single_key_with_signs(&key, dim, &config, &signs);
cache.append_raw(&packed, norm);

Modules

codebook

Lloyd-Max codebook quantization (2/3/4-bit optimal centroids). Lloyd-Max Optimal Codebook for Gaussian-distributed coordinates.

compaction

KV cache compaction — reduce token count via attention matching. KV Cache Compaction — reduce token count while preserving attention behavior.

hadamard

Fast Walsh-Hadamard Transform for decorrelation. Fast Walsh-Hadamard Transform (WHT)

Structs

CompressedKeys

Paper-faithful compressed key cache. Only keys are compressed; values stay in fp16.

CompressedValues

Compressed value cache using per-vector absmax quantization.

CompressedValues4Bit

Compressed value cache using per-group 4-bit absmax quantization.

CompressedVectors

Compressed vector collection.

CompressionStats

Compression quality statistics.

DecayTier

A single decay tier: tokens older than age_threshold get compressed to bits.

SparseVStats

Statistics from a sparse V multiply: how many positions were active vs skipped.

TemporalDecayConfig

Temporal decay configuration.

TurboQuantConfig

TurboQuant configuration.

Enums

QjlMode

QJL activation mode.

Functions

calibrate_channel_scales

Calibrate per-channel scaling factors from a batch of key vectors.

calibrate_codebook

Calibrate codebook from a batch of key vectors.

calibrate_codebook_with_rotation

Calibrate codebook with optional custom rotation matrix. If rotation_matrix is Some, uses that instead of randomized Hadamard. This ensures the codebook is fitted to the same rotation used at runtime.

calibrate_rotation

Calibrate optimal rotation matrix from key vectors (SpinQuant PCA approach).

compress_keys

Compress key vectors using paper-faithful Lloyd-Max codebook.

compress_single_key

Compress a single key vector. For incremental KV cache. Returns: (packed_indices, corrected_norm)

compress_single_key_grouped

Compress a single key with per-group quantization.

compress_single_key_with_signs

Compress a single key vector with pre-computed signs. Saves signs allocation in the hot loop.

compress_vectorsDeprecated

V1 API: Compress vectors using PolarQuant pipeline. Deprecated — use compress_keys (V2 Lloyd-Max) instead for better compression ratio and speed.

decompress_keys

Decompress keys back to f32.

decompress_keys_grouped

Decompress keys with per-group norms.

decompress_vectorsDeprecated

V1 API: Decompress PolarQuant data back to f32 vectors.

evaluateDeprecated

V1 API: Evaluate PolarQuant compression quality.

evaluate_keys

Evaluate V2 compression quality.

fused_attention_scores

Batch fused attention scores: compute all attention scores between a pre-rotated query and all keys in a compressed cache.

fused_dot_product

Compute attention score between pre-rotated query and compressed key.

fused_dot_product_with_centroids

Fused dot product with pre-computed centroid table. Eliminates Codebook construction overhead in the hot loop.

pre_rotate_query

Pre-rotate a query vector for fused attention.

pre_rotate_query_with_matrix

Pre-rotate query with a custom rotation matrix (SpinQuant/PCA).

pre_rotate_query_with_signs

Pre-rotate query with pre-computed signs (alloc-free hot path).

softmax_bias_correction

Softmax bias correction (Bondarenko, arXiv:2309.01729).

sparse_attn_v_mul

Sparse attention-value multiply: only accumulate V rows where attention weight > threshold.

sparse_attn_v_mul_compressed_4bit

Fused sparse attention-value multiply on 4-bit compressed values.

sparse_attn_v_mul_compressed_8bit

Fused sparse attention-value multiply on 8-bit compressed values.

sparse_v_stats

Count how many positions would be active for a given threshold.