Turboquant Pytorch

作者 reason-machines2384a003145a無授權條款83 個星標收錄於 2026年10月8日更新於 2026年10月8日儲存庫3 個月前更新

PyTorch implementation of TurboQuant for LLM KV cache compression using two-stage vector quantization (random rotation + Lloyd-Max + QJL residual correction).

僅含說明Software DevelopmentData & Analytics
AI 產生的概覽

在 PyTorch 中實作 TurboQuant 兩階段向量量化，將 LLM KV 快取壓縮至每座標 2–4 位元。

功能
提供從零實作的 PyTorch 版 TurboQuant，用於 LLM KV 快取壓縮，結合隨機正交旋轉、Lloyd-Max 純量量化與 QJL 1 位元殘差修正。它包含碼本建構、MSE 量化、無偏內積估計、單注意力頭 KV 快取封裝，以及針對真實模型張量的生產級壓縮器。文件也提供注意力保真度、大海撈針檢索與記憶體節省估算的驗證模式。
適用情境
當需要在 LLM 推論中降低 KV 快取記憶體用量，或在 PyTorch 中實作 TurboQuant 風格量化時使用。它適合測量 2–4 位元下的注意力保真度、檢索準確率或壓縮比的實驗。
執行需求
需要 Python、PyTorch 2.0 或更新版本，以及用於 Lloyd-Max 碼本計算的 scipy；transformers、accelerate 與 bitsandbytes 僅在真實模型驗證時需要。真實模型驗證需要至少 6GB VRAM 的 CUDA GPU，首次執行會下載約 2GB 模型。此技能僅為說明文件，不附帶指令碼。

TurboQuant PyTorch

Skill by ara.so — Daily 2026 Skills collection.

From-scratch PyTorch implementation of Google's TurboQuant (ICLR 2026) for compressing LLM KV caches. Achieves 5x compression at 3-bit with 99.5% attention fidelity via two-stage vector quantization.

What It Does

TurboQuant compresses LLM key-value caches to 2–4 bits per coordinate:

  • Stage 1: Random orthogonal rotation + Lloyd-Max scalar quantization (MSE-optimal)
  • Stage 2: QJL residual correction — 1-bit sign projection that makes inner product estimates unbiased

Result: attention scores remain accurate even when individual vectors look quite different from originals. The algorithm preserves inner products, not vector fidelity.

Compression ratios at 8K context on Qwen2.5-3B (289 MB FP16 baseline):

  • 4-bit → 76 MB (3.8x)
  • 3-bit → 58 MB (5.0x) ← practical sweet spot
  • 2-bit → 40 MB (7.3x)

Installation

bash
git clone https://github.com/tonbistudio/turboquant-pytorchcd turboquant-pytorchpip install -r requirements.txt
# For CUDA PyTorch:pip install torch --index-url https://download.pytorch.org/whl/cu128

requirements.txt includes:

  • torch>=2.0
  • scipy (Lloyd-Max codebook computation)
  • transformers, accelerate, bitsandbytes (only for real model validation)

Project Structure

turboquant/  __init__.py           # Package exports  lloyd_max.py          # Lloyd-Max optimal scalar quantizer  turboquant.py         # Core: TurboQuantMSE, TurboQuantProd, TurboQuantKVCache  compressors.py        # Production compressors for real model tensors  test_turboquant.py    # Synthetic validation tests  validate.py           # Real model (Qwen2.5-3B) validation

Key Commands

bash
# Run synthetic algorithm validation (no GPU required, but GPU enables speed benchmark)python -m turboquant.test_turboquant
# Run real model validation on Qwen2.5-3B-Instruct# Requires CUDA GPU with ≥6GB VRAM; downloads ~2GB model on first runpython -m turboquant.validate

Core API

Lloyd-Max Codebook

python
from turboquant.lloyd_max import build_lloyd_max_codebook
# Build optimal scalar quantizer codebook for d-dimensional rotated unit vectors# Returns (boundaries, centroids) for the given bit-widthboundaries, centroids = build_lloyd_max_codebook(dim=128, bits=3)

Stage 1: MSE Quantization (TurboQuantMSE)

python
from turboquant.turboquant import TurboQuantMSE
# Initialize for head_dim=128, 3-bit quantizationtq_mse = TurboQuantMSE(dim=128, bits=3)
# Compress a batch of vectors: shape (batch, dim)keys = torch.randn(512, 128)  # 512 key vectorscodes = tq_mse.quantize(keys)       # integer codes, (512, 128)reconstructed = tq_mse.dequantize(codes)  # approximate keys, (512, 128)

Stage 2: Unbiased Inner Product Estimation (TurboQuantProd)

python
from turboquant.turboquant import TurboQuantProd
# Initialize with QJL correctiontq_prod = TurboQuantProd(dim=128, bits=3, proj_dim=64)
# Compress key vectors (stores codes + QJL residual signs)compressed = tq_prod.compress(keys)  # dict with 'codes', 'signs', 'residual_norms'
# Estimate inner products <query, key> for all keys — unbiased estimatorquery = torch.randn(128)scores = tq_prod.inner_product(query, compressed)  # shape (512,)

KV Cache Wrapper (TurboQuantKVCache)

python
from turboquant.turboquant import TurboQuantKVCache
# Wrap a KV cache for a single attention headcache = TurboQuantKVCache(dim=128, bits=3, proj_dim=64)
# Add key/value vectors as tokens are generatedcache.append_key(new_key)    # shape (dim,)cache.append_value(new_val)  # shape (dim,)
# Compute attention scores for a query against all cached keysquery = torch.randn(128)scores = cache.attention_scores(query)  # shape (seq_len,), unbiased
# Get values (MSE-reconstructed, used for weighted sum)values = cache.get_values()  # shape (seq_len, dim)

Production Compressors (for real model tensors)

python
from turboquant.compressors import TurboQuantCompressorV2, TurboQuantCompressorMSE
# Key compressor — supports asymmetric attention score computationkey_compressor = TurboQuantCompressorV2(dim=128, bits=3, proj_dim=64)
# Compress all keys in a layer: shape (num_heads, seq_len, head_dim)compressed_keys = key_compressor.compress(layer_keys)
# Compute attention scores directly from compressed keys (no decompress needed)# query shape: (num_heads, head_dim)scores = key_compressor.asymmetric_attention_scores(query, compressed_keys)# scores shape: (num_heads, seq_len)
# Value compressor — MSE reconstruction (Stage 1 only, acceptable for values)val_compressor = TurboQuantCompressorMSE(dim=128, bits=3)compressed_vals = val_compressor.compress(layer_values)reconstructed_vals = val_compressor.decompress(compressed_vals)

Common Patterns

Pattern 1: Compress a Full Model's KV Cache

python
import torchfrom turboquant.compressors import TurboQuantCompressorV2, TurboQuantCompressorMSE
def compress_kv_cache(kv_cache, head_dim=128, bits=3, proj_dim=64):    """    kv_cache: list of (keys, values) per layer              keys/values shape: (num_heads, seq_len, head_dim)    Returns list of compressed (keys, values) per layer.    """    key_comp = TurboQuantCompressorV2(dim=head_dim, bits=bits, proj_dim=proj_dim)    val_comp = TurboQuantCompressorMSE(dim=head_dim, bits=bits)
    compressed = []    for layer_keys, layer_vals in kv_cache:        c_keys = key_comp.compress(layer_keys)        c_vals = val_comp.compress(layer_vals)        compressed.append((c_keys, c_vals))
    return compressed, key_comp, val_comp
def run_attention_with_compressed_cache(query, compressed_keys, compressed_vals,                                        key_comp, val_comp):    """    query: (num_heads, head_dim)    Returns: attention output (num_heads, head_dim)    """    # Unbiased attention scores from compressed keys    scores = key_comp.asymmetric_attention_scores(query, compressed_keys)    # scores: (num_heads, seq_len)
    attn_weights = torch.softmax(scores, dim=-1)  # (num_heads, seq_len)
    # Decompress values and compute weighted sum    values = val_comp.decompress(compressed_vals)  # (num_heads, seq_len, head_dim)    output = torch.einsum('hs,hsd->hd', attn_weights, values)    return output

Pattern 2: Validate Compression Quality

python
import torchimport torch.nn.functional as Ffrom turboquant.turboquant import TurboQuantProd
def measure_attention_fidelity(keys, queries, bits=3, proj_dim=64):    """    Measure how well TurboQuant preserves attention distributions.    keys:    (seq_len, head_dim)    queries: (num_queries, head_dim)    """    dim = keys.shape[-1]    tq = TurboQuantProd(dim=dim, bits=bits, proj_dim=proj_dim)
    compressed = tq.compress(keys)
    cosine_sims = []    top1_matches = []
    for q in queries:        # True attention scores        true_scores = (keys @ q)  # (seq_len,)        true_attn = torch.softmax(true_scores, dim=0)
        # TurboQuant estimated scores        est_scores = tq.inner_product(q, compressed)  # (seq_len,)        est_attn = torch.softmax(est_scores, dim=0)
        # Cosine similarity of attention distributions        cos_sim = F.cosine_similarity(true_attn.unsqueeze(0),                                       est_attn.unsqueeze(0)).item()        cosine_sims.append(cos_sim)
        # Top-1 match        top1_matches.append(true_attn.argmax() == est_attn.argmax())
    return {        'mean_cosine_sim': sum(cosine_sims) / len(cosine_sims),        'top1_accuracy': sum(top1_matches) / len(top1_matches),    }
# Example usagekeys = torch.randn(2048, 128)keys = F.normalize(keys, dim=-1)queries = torch.randn(100, 128)queries = F.normalize(queries, dim=-1)
results = measure_attention_fidelity(keys, queries, bits=3)print(f"Cosine similarity: {results['mean_cosine_sim']:.4f}")print(f"Top-1 accuracy:    {results['top1_accuracy']:.2%}")

Pattern 3: Needle-in-Haystack Retrieval Test

python
import torchimport torch.nn.functional as Ffrom turboquant.turboquant import TurboQuantProd
def needle_in_haystack(seq_len=2048, dim=128, bits=3):    """Test whether TurboQuant preserves nearest-neighbor ordering."""    tq = TurboQuantProd(dim=dim, bits=bits, proj_dim=64)
    # Build haystack of random unit vectors    haystack = F.normalize(torch.randn(seq_len, dim), dim=-1)
    # Insert needle at random position    needle_idx = torch.randint(0, seq_len, (1,)).item()    query = F.normalize(torch.randn(dim), dim=0)    needle = query + 0.1 * torch.randn(dim)  # Similar to query    needle = F.normalize(needle, dim=0)    haystack[needle_idx] = needle
    # Compress    compressed = tq.compress(haystack)
    # True nearest neighbor    true_scores = haystack @ query    true_best = true_scores.argmax().item()
    # TurboQuant estimated nearest neighbor    est_scores = tq.inner_product(query, compressed)    est_best = est_scores.argmax().item()
    return true_best == est_best, true_best, est_best
# Run multiple trialssuccesses = sum(needle_in_haystack(seq_len=8192)[0] for _ in range(20))print(f"Retrieval accuracy: {successes}/20")

Pattern 4: Compute Memory Savings

python
def estimate_memory_savings(num_layers, num_kv_heads, seq_len, head_dim,                             bits, proj_dim=64):    """    Estimate compressed KV cache size vs FP16 baseline.    """    # FP16 baseline: 2 bytes per element    fp16_bytes = num_layers * 2 * num_kv_heads * seq_len * head_dim * 2
    # Stage 1 codes: bits per element, packed into bytes    codes_bytes = (num_layers * 2 * num_kv_heads * seq_len * head_dim * bits) // 8
    # Stage 2 signs (keys only): 1 bit per proj_dim element    signs_bytes = (num_layers * num_kv_heads * seq_len * proj_dim) // 8
    # Residual norms: 1 float16 per vector (keys only)    norms_bytes = num_layers * num_kv_heads * seq_len * 2
    total_compressed = codes_bytes + signs_bytes + norms_bytes    ratio = fp16_bytes / total_compressed
    print(f"FP16 baseline:      {fp16_bytes / 1e6:.1f} MB")    print(f"TurboQuant {bits}-bit:  {total_compressed / 1e6:.1f} MB")    print(f"Compression ratio:  {ratio:.1f}x")    return ratio
# Qwen2.5-3B: 36 layers, 2 KV heads, head_dim=128estimate_memory_savings(    num_layers=36, num_kv_heads=2,    seq_len=8192, head_dim=128, bits=3)# FP16 baseline:      289.4 MB# TurboQuant 3-bit:   57.9 MB# Compression ratio:  5.0x

Algorithm Details

Why Random Rotation?

Rotating by a random orthogonal matrix R maps unit vectors to a space where each coordinate follows N(0, 1/d). This makes coordinates nearly independent with known distribution — enabling optimal per-coordinate scalar quantization (Lloyd-Max).

Why QJL for Keys but Not Values?

  • Keys: Used in dot products with queries. Bias in inner product estimates directly corrupts attention weights. QJL correction is essential.
  • Values: Used in weighted sums after softmax. Small per-vector MSE errors average out. Stage 1 MSE quantization is sufficient.

Choosing proj_dim (QJL projection dimension)

Higher proj_dim → lower variance in inner product estimates, but more memory:

python
# Rule of thumb: proj_dim = head_dim // 2 is a good default# head_dim=128 → proj_dim=64# head_dim=64  → proj_dim=32# head_dim=256 → proj_dim=128

Bit-width Selection Guide

BitsCompressionCosine SimTop-1 MatchUse Case
43.8x0.99987%Quality-critical tasks
35.0x0.99582%Recommended default
27.3x0.98866%Extreme memory pressure

Troubleshooting

scipy import error when building codebooks:

bash
pip install scipy

CUDA out of memory during validate.py:

  • Requires ≥6GB VRAM for Qwen2.5-3B in 4-bit
  • Reduce seq_len in the validation script or use a smaller model

Inner product estimates have high variance:

  • Increase proj_dim (try head_dim instead of head_dim // 2)
  • Check that input vectors are normalized before compressing

Codebook build is slow on first run:

  • Lloyd-Max uses numerical integration (scipy) — this is expected
  • Codebooks are precomputed once per (dim, bits) combination; cache them:
python
import pickle
# Save codebookboundaries, centroids = build_lloyd_max_codebook(dim=128, bits=3)with open('codebook_128_3bit.pkl', 'wb') as f:    pickle.dump((boundaries, centroids), f)
# Load cached codebookwith open('codebook_128_3bit.pkl', 'rb') as f:    boundaries, centroids = pickle.load(f)

Attention fidelity lower than expected:

  • Ensure vectors are L2-normalized before compressing (F.normalize(x, dim=-1))
  • The compressors in compressors.py handle normalization internally; TurboQuantProd expects unit vectors

References

來源與署名

來源：reason-machines/trending-skills位於skills/turboquant-pytorch提交2384a00

授權條款： 無授權條款

內容歸原作者所有。SourceWeft 從公開儲存庫中收錄這些內容。

檢舉或申請下架

更多來自 reason-machines/trending-skills 的技能

Yourvpndead Vpn Detection

reason-machines

介紹一款 Android 應用程式，用來偵測 VPN 用戶端並探測本機 SOCKS5 與 Clash API 連接埠是否外洩。

Security833 個月前更新

X Mentor Skill Nuwa

reason-machines

提供 X（Twitter）內容策略指引：推文與長推撰寫、鉤子評論、成長策略與帳號診斷。

Marketing & Sales833 個月前更新

Wildworld Dataset

reason-machines

介紹 WildWorld 動作條件 ARPG 世界建模資料集，並提供預期的載入、篩選與基準評測程式碼。

Data & Analytics833 個月前更新

Understand Anything Knowledge Graph

reason-machines

為程式碼庫建立可互動的知識圖譜，並支援視覺化探索、搜尋與提問。

Software Development833 個月前更新

Webphysics Avbd Engine

reason-machines

指導將實驗性 WebGPU AVBD 剛體與軟體物理引擎整合到瀏覽器專案中。

Software Development833 個月前更新

Ui Ux Pro Max Skill

reason-machines

依產品描述產生 UI/UX 設計系統，包含頁面結構、風格、配色、字體搭配與檢查清單。

Design & Creative833 個月前更新

更多Software Development技能

Playground

anthropics

精選

建立自成一體的互動式 HTML 遊樂場，包含控制項、即時預覽與可複製的提示詞輸出。

Software Development37K今天更新

M5 Onboard

anthropics

精選

透過 USB 偵測 M5Stack ESP32 開發板，燒錄 UIFlow 2.0 韌體並安裝 MicroPython 應用程式套件。

Software Development37K今天更新

Cardputer Buddy

anthropics

精選

指導迭代 Cardputer-Adv 的 MicroPython 應用程式套件：新增應用程式、透過 USB 序列埠推送檔案、查看日誌並執行 REPL 指令。

Software Development37K今天更新

Web Design Guidelines

vercel-labs

精選

依據從遠端來源取得的 Web Interface Guidelines 審查 UI 程式碼，並以 file:line 形式回報問題。

Software Development32K5 週前更新

Vercel React Native Skills

vercel-labs

精選

React Native 與 Expo 最佳實務規則，涵蓋清單效能、動畫、導覽、UI 模式、狀態與 monorepo 設定。

Software Development32K5 週前更新

Vercel React Best Practices

vercel-labs

精選

Vercel 的 React 與 Next.js 效能指南，以 70 條規則的形式用於撰寫、審查與重構程式碼。

Software Development32K5 週前更新