Closing the Gap: Trellis-Coded Quantization for KV Cache at 2-3 Bits

Authors: buun, Claude (Anthropic)

This dataset accompanies the paper Closing the Gap: Trellis-Coded Quantization for KV Cache at 2-3 Bits. It contains the trained TCQ codebooks, codebook training scripts, and the full paper PDF.

Summary

We present the first application of trellis-coded quantization (TCQ) to KV cache compression in LLM inference. TCQ constrains quantization indices to follow a finite-state trellis, enabling a much larger effective codebook at the same bit rate. The Viterbi algorithm finds globally optimal encoding, and the bitshift trellis structure allows each element to be decoded independently in O(1) -- making it fully compatible with GPU flash attention kernels.

Key results

• 10-44% KL-divergence reduction over scalar quantization at 2-3 bits per value
• At 3.25 bpv, TCQ produces lower perplexity than FP16 KV cache due to a mild regularizing effect from norm scaling (5.802 vs 5.805 on Qwen3.5-27B)
• At 2.25 bpv, TCQ significantly improves 2-bit quantization, closing the gap with 3-bit scalar methods
• O(1) parallel decode -- no sequential Viterbi at inference time; each value decodes via a sliding bit window
• Context-adaptive norm scaling -- closed-form logarithmic formulas automatically select optimal dequantization scale per context length, tracking within 2% of per-context optima with zero speed overhead

Method

1. FWHT rotation with random sign flips converts correlated KV vectors into approximately i.i.d. Gaussian entries
2. Viterbi encoding on a 512-state (3-bit) or 256-state (2-bit) right-shift trellis finds the globally optimal codeword assignment
3. O(1) sliding-window decode reads 3 (or 2) index bits plus 8 (or 6) state bits to look up the dequantized value -- no trellis traversal needed at decode time
4. Adaptive alpha scaling applies a context-dependent norm correction at decode time to compensate for CLT averaging effects in the attention mechanism

The trained codebooks are optimized via the Generalized Lloyd Algorithm (GLA) starting from coset-spaced initialization, converging in 50-100 iterations.

Contents

tcq-kv-cache.pdf              # Full paper (17 pages)
codebooks/
  3bit/
    product_aware_iter080.bin  # Recommended (compiled default, 512 floats, 2048 bytes)
    cb_50iter_finetuned.bin    # Earlier single-objective CUDA GLA
    tcq_3bit_numpy_s*.bin      # Multi-seed numpy GLA variants
    tcq_3bit_left_*.bin        # Left-shift trellis variants
    ...                        # 16 codebooks total
  2bit/
    product_aware_iter090.bin  # Recommended (compiled default, 256 floats, 1024 bytes)
    tcq_2bit_cuda_200iter.bin  # Earlier CUDA GLA variant
    tcq_2bit_100iter_s99.bin   # numpy GLA, better at very long context
    ...                        # 7 codebooks total
training-scripts/
    tcq_train_v2.py            # Main trainer with tail-biting + left-shift support
    tcq_train_vectorized.py    # Vectorized batch Viterbi (10-50x faster)
    tcq_train_2bit.py          # Specialized 2-bit trainer
    tcq_train_cuda.cu          # CUDA-accelerated trainer
    compute_centroids.py       # Lloyd-Max centroid computation from KV data

Codebook format

Each .bin file contains raw float32 values:

• 3-bit: 512 floats (2048 bytes) -- indexed by 9-bit state (512 states x 1 output per state)
• 2-bit: 256 floats (1024 bytes) -- indexed by 8-bit state (256 states x 1 output per state)

Load in Python:

import numpy as np
cb = np.fromfile("codebooks/3bit/cb_50iter_finetuned.bin", dtype=np.float32)
# cb.shape = (512,) for 3-bit, (256,) for 2-bit

Load at runtime in llama.cpp (TurboQuant fork):

TURBO_TCQ_CB=codebooks/3bit/cb_50iter_finetuned.bin \
TURBO_TCQ_CB2=codebooks/2bit/tcq_2bit_100iter_s99.bin \
./llama-server -m model.gguf -ngl 99 -fa 1 -ctk turbo3_tcq -ctv turbo3_tcq

Codebook selection

Important: Lower MSE does not always mean better downstream quality. We found that the relationship between codebook MSE and model perplexity/KLD inverts between short and long contexts due to CLT averaging of quantization error in the attention mechanism (see Section 4.2 of the paper). Higher-iteration codebooks achieve lower MSE but can produce worse PPL at short contexts.

The codebooks compiled into the default implementation are product-aware trained codebooks from the latest CUDA GLA training campaign (100 iterations on Qwen3.5-27B FWHT-rotated KV activations). These are baked into the CUDA kernels and used automatically -- no environment variable needed. The .bin files in this dataset can override them via TURBO_TCQ_CB / TURBO_TCQ_CB2 environment variables for experimentation.

The remaining codebooks document the training progression and are included for reproducibility. The multi-seed variants (_s7, _s42, etc.) demonstrate that GLA converges to similar quality regardless of random seed, confirming the robustness of post-FWHT Gaussian inputs to the training procedure.

Implementation

The full CUDA implementation is available at: https://github.com/spiritbuun/llama-cpp-turboquant-cuda

Citation

@misc{buun2026tcqkv,
  title={Closing the Gap: Trellis-Coded Quantization for KV Cache at 2--3 Bits},
  author={buun and Claude},
  year={2026},
  url={https://huggingface.co/datasets/spiritbuun/turboquant-tcq-kv-cache}
}

License

Apache 2.0