turboquant-pytorch

TurboQuant V3 is designed to compress LLM Key-Value (KV) caches using a combination of random rotation and Lloyd-Max quantization.

Core Components:

• MSECompressor: The fundamental building block. It performs single-stage compression using bit-packed storage and focuses on Mean Squared Error (MSE) minimization.
• TurboQuantV3: The high-level orchestrator that manages multiple MSECompressor instances for different layers and handles the asymmetric bit allocation between Keys and Values.

Why V3 is preferred over V2: Unlike the original paper's Stage 2 (QJL), V3 removes the QJL residual correction. While QJL is unbiased for raw inner products, the softmax operation in attention exponentially amplifies its noise. V3's MSE-only approach results in lower variance, which is critical for maintaining text generation quality after the softmax layer.

To use TurboQuant, ensure you have Python 3.10+ and a CUDA-capable NVIDIA GPU. You can install the dependencies using the provided requirements file.

For local development, install the package in editable mode.

# Install dependencies
pip install -r requirements.txt

# For editable local development
pip install -e .

# For CUDA PyTorch (example for cu128)
pip install torch --index-url https://download.pytorch.org/whl/cu128

The Generation Test is the recommended way to verify if the model produces correct text with a compressed KV cache. This test downloads the Qwen2.5-3B-Instruct model (~2GB) and evaluates multiple configurations across different context lengths.

python -m turboquant.generation_test

TurboQuant provides several validation modules to test different aspects of the algorithm:

• Attention Validation (V3 vs V2): Compares the attention score accuracy of the V3 implementation against the V2 implementation.
• Synthetic Tests: Validates the core algorithm against theoretical bounds from the original paper without requiring a large model.
• Original V2 Validation: Runs the original attention-score comparison (without generation testing).

# Compare V3 vs V2 attention accuracy
python -m turboquant.validate_v3

# Run synthetic algorithm tests
python -m turboquant.test_turboquant

# Run original V2 validation
python -m turboquant.validate

The TurboQuantV3 class in compressors_v3.py is the primary orchestrator for V3 compression. It supports asymmetric bit-widths for Keys (K) and Values (V), as well as layer-adaptive precision to protect sensitive layers.

Key Features:

• Asymmetric K/V: Allocate more bits to Keys (which require higher precision for attention) than to Values.
• Layer-adaptive: Protect specific layers (e.g., the first or last few layers) with higher precision.
• MSE-only: Uses MSECompressor for improved reconstruction quality by removing the QJL stage.

# Conceptual usage based on documentation
from turboquant.compressors_v3 import TurboQuantV3

# Example: 4-bit keys, 2-bit values, protecting the first 4 layers
compressor = TurboQuantV3(key_bits=4, value_bits=2, protected_layers=4)

The following classes are the primary interfaces for using the V3 compression logic:

• MSECompressor (compressors_v3.py): Single-stage compressor with bit-packed storage. Used for both keys and values.
• TurboQuantV3 (compressors_v3.py): Orchestrator that creates separate key/value compressors with different bit-widths and handles layer-adaptive precision.
• TurboQuantMSE (turboquant.py): Original Stage 1 quantizer.
• TurboQuantProd (turboquant.py): Original two-stage quantizer with QJL (kept for reference).