RotorQuant KV Cache Compression

The QJL (two-term unbiased estimator) correction is used to improve the accuracy of quantized attention scores. Instead of a simple MSE-based quantization, the QJL approach uses a two-term score calculation to account for outliers and quantization errors.

The scoring logic follows these steps:

1. Storage: Outliers are stored in a compressed cache per key.
2. Key Compression: The QJL estimator is computed after MSE quantization.
3. Query Pre-processing: The query is projected through a random projection matrix $S$.
4. Fused Kernel Calculation: The final score is a combination of the quantized key inner product and the outlier inner product, scaled by their respective norms and sketch dimensions: score = scl * norm_k * shared_innprod + scl_otlr * norm_otlr * shared_outlier_innprod where scl and scl_otlr are scaling factors derived from the sketch dimensions.

The CompressedKVCache (extending transformers.DynamicCache) is designed to store keys in their compressed uint8 format. This prevents the need to materialize full float16 keys in memory during the KV cache lifecycle.

Key Behaviors:

• Storage: Stores uint8 indices and float16 norms per layer.
• Quantization: Keys are quantized on insertion using a provided TurboQuantMSE quantizer.
• Values: Values are stored in standard fp16 as they do not benefit as much from the fused kernel approach used for keys.

from transformers import DynamicCache
from turboquant_core import TurboQuantMSE

class CompressedKVCache(DynamicCache):
    def __init__(self, quantizer: TurboQuantMSE):
        super().__init__()
        self.tq = quantizer
        self._compressed_keys: list[dict | None] = []

    def store_compressed_key(self, key_states: torch.Tensor, layer_idx: int):
        # Quantize and store key states
        ...

The benchmark_perplexity.py script measures the language modeling quality degradation (perplexity) caused by KV cache quantization compared to an FP16 baseline using the wikitext-2 dataset.

Usage: Run the benchmark via the module interface. You can specify the model and the target bit-widths for quantization.

# Basic usage with default model (Qwen/Qwen2.5-3B-Instruct)
python -m turboquant.benchmark_perplexity

# Specify a model and multiple bit-widths (e.g., 2, 3, and 4 bits)
python -m turboquant.benchmark_perplexity --model Qwen/Qwen2.5-7B-Instruct --bits 2 3 4

Arguments:

• --model: The Hugging Face model ID (default: Qwen/Qwen2.5-3B-Instruct).
• --bits: One or more integers specifying the quantization bit-widths to test.
• --max-length: Maximum sequence length for the sliding window (default: 2048).
• --stride: Overlap stride for the sliding window (default: 512).
• --max-tokens: Limit the number of tokens processed from the dataset (0 for full set).

Once built, you can use llama-server to serve models with compressed KV caches. Use the --cache-type-k and --cache-type-v flags to specify the quantization method.

Symmetric 3-bit (Best quality per bit): Uses iso3 or planar3 for both K and V caches.

./build/bin/llama-server -m model.gguf --jinja -ngl 99 \
    --cache-type-k iso3 --cache-type-v iso3 --host 0.0.0.0 --port 8080

K-only compression (Zero PPL loss, 5x compression): Compresses only the K cache while keeping V in FP16.

./build/bin/llama-server -m model.gguf --jinja -ngl 99 \
    --cache-type-k planar3 --cache-type-v f16 --host 0.0.0.0 --port 8080

Available Cache Types:

• planar3, iso3, planar4, iso4 (RotorQuant implementations)
• turbo3, turbo4 (WHT butterfly implementations)

./build/bin/llama-server -m model.gguf --jinja -ngl 99 --cache-type-k iso3 --cache-type-v iso3 --host 0.0.0.0 --port 8080

To prevent the perplexity degradation caused by MSE-only attention scores in 3-bit compressed KV caches, implement the QJL (two-term unbiased estimator) correction. This correction compensates for the bias in MSE-reconstructed keys by adding a second term based on a random Gaussian projection.

The QJL Estimator Formula

$$\langle q, k \rangle \approx \underbrace{\langle q, k_{mse} \rangle}{term1} + \underbrace{\frac{|residual| \cdot \sqrt{\pi/2}}{m} \cdot \langle S@q, \text{sign}(S@residual) \rangle}{term2 (QJL\ correction)}$$

Implementation Requirements

1. Storage Updates

Add the following to the compressed cache for each key:

• qjl_signs: Packed bits representing sign(S @ residual.T) as int8 {+1, -1} with shape [batch, n_kv_heads, kv_len, head_dim].
• residual_norms: fp16 values representing ||residual|| with shape [batch, n_kv_heads, kv_len].
• S: A random Gaussian projection matrix of shape [head_dim, head_dim] (shared or per-layer).

2. Key Compression Workflow

After MSE quantization, compute the QJL components:

1. Dequantize indices to get k_mse.
2. Calculate residual = k_original - k_mse.
3. Compute residual_norm = ||residual||.
4. Compute qjl_signs = sign(S @ residual.T).

3. Query Pre-processing

Project the query through the projection matrix S to create a sketch: query_sketch = Q @ S.T (Shape: [batch, n_heads, q_len, m])

4. Fused Kernel Logic

The Triton kernel must compute the score using both terms:

• term1: The existing MSE score using norm * sum_d(Q_rot[d] * centroids[idx[s,d]]).
• term2: The QJL correction using res_norm[s] * sqrt(π/2)/m * sum_d(q_sketch[d] * signs[s,d]).
• score[s] = (term1 + term2) * scale

# Key compression logic snippet
k_mse = dequantize(indices, key_norms)
residual = k_original - k_mse
residual_norm = ||residual||
qjl_signs = sign(S @ residual.T)  # 1-bit per dim

# Query pre-processing
query_sketch = Q @ S.T  # [batch, n_heads, q_len, m]

# Fused kernel score calculation
term1 = norm * sum_d(Q_rot[d] * centroids[idx[s,d]])
term2 = res_norm[s] * sqrt(π/2)/m * sum_d(q_sketch[d] * signs[s,d])
score[s] = (term1 + term2) * scale

Use the following commands to evaluate the implementation:

Benchmark llama.cpp performance:

./build/bin/llama-bench -m model.gguf -ngl 99 -ctk planar3 -ctv planar3 -p 512 -n 128

Calculate Perplexity (PPL): First, ensure datasets is installed and download the wikitext dataset:

pip install datasets
python3 -c "from datasets import load_dataset; open('/tmp/wiki.txt','w').write('\n'.join(load_dataset('wikitext','wikitext-2-raw-v1',split='test')['text']))"

Then run the perplexity tool:

./build/bin/llama-perplexity -m model.gguf -f /tmp/wiki.txt -ngl 99 -c 2048 \
    --cache-type-k iso3 --cache-type-v iso3

Python/Triton Benchmarks:

python -m turboquant.benchmark_google_parity          # PPL (post-prefill)
python -m turboquant.benchmark_perplexity --bits 3 4   # PPL (roundtrip)
python -m turboquant.benchmark_triton                  # Triton kernel speed
python -m turboquant.poc_high_context --backend planar  # High-context generation

To use RotorQuant's fused attention, you must patch the model's attention layers by replacing their forward methods with a fused version. This process iterates through the model's named modules, identifies layers containing q_proj, k_proj, v_proj, and an output projection (o_proj or out_proj), and injects the head counts if they are missing (common in Qwen2 models).

for name, module in model.named_modules():
    has_projs = all(hasattr(module, a) for a in ['q_proj', 'k_proj', 'v_proj'])
    has_out = hasattr(module, 'o_proj') or hasattr(module, 'out_proj')
    if has_projs and has_out:
        # Inject head counts if not on module
        if not hasattr(module, 'num_heads'):
            module.num_heads = text_config.num_attention_heads
        if not hasattr(module, 'num_key_value_heads'):
            module.num_key_value_heads = getattr(
                text_config, 'num_key_value_heads', text_config.num_attention_heads)
        module.forward = make_fused_rotor_attention_forward(
            module, cache, layer_idx)
        patched += 1
        layer_idx += 1