[Feature]: Add TurboQuant Support for KV Cache Quantization

🚀 The feature, motivation and pitch

TurboQuant is a novel online vector quantization method that achieves near-optimal distortion rates for both MSE and inner product preservation, specifically designed for KV cache compression in LLMs. As described in the paper TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate, this method provides:

• Provably optimal distortion bounds within a factor of ~2.7 of information-theoretic limits
• Unbiased inner product estimation - critical for attention mechanisms
• Online application - no data-dependent preprocessing required
• Accelerator-friendly design - suitable for real-time AI workloads
• Significant memory savings - 4-5× compression while maintaining model performance

The paper demonstrates TurboQuant achieving perfect recall on Needle-in-a-Haystack tests at 4× compression and competitive performance on LongBench with just 2.5-3.5 bits per dimension.

This would be a valuable addition to vLLM's quantization portfolio, complementing existing scalar methods (FP8, INT4, etc.) with a vector quantization approach specifically optimized for attention KV caches.

Alternatives

Current vLLM KV cache quantization options include:

• FP8 quantization (fp8_e4m3, fp8_e5m2) - scalar quantization with element-wise MSE optimization
• Compressed-tensors KV cache - supports various schemes but primarily scalar-based
• ModelOpt FP8 - NVIDIA's FP8 implementation

TurboQuant differs by:

1. Using vector quantization with learned codebooks instead of scalar quantization
2. Optimizing for inner product preservation rather than just MSE
3. Providing theoretical guarantees on distortion bounds
4. Supporting variable bit-widths (2.5-bit, 3.5-bit) through outlier handling

Additional context

Technical Integration Approach

Based on vLLM's existing quantization framework, TurboQuant integration would require:

1. Extend Cache Configuration

Add "turboquant" to CacheDType literal in vllm/config/cache.py and update the dtype mapping in vllm/utils/torch_utils.py to use integer storage types for codebook indices.

2. Create TurboQuantConfig Class

Implement using the registration decorator pattern:

@register_quantization_config("turboquant")
class TurboQuantConfig(QuantizationConfig):
    def get_name(self) -> str:
        return "turboquant"
    
    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
        if isinstance(layer, Attention):
            return TurboQuantKVCacheMethod(self)
        return None

This follows the established pattern from vllm/model_executor/layers/quantization/__init__.py .

3. Implement KV Cache Method

Create TurboQuantKVCacheMethod extending BaseKVCacheMethod to:

• Register codebook parameters instead of scalar scales
• Handle both MSE-optimized and inner-product-optimized variants
• Support per-head quantization strategies

4. Update Quantization Detection

Modify is_quantized_kv_cache() to recognize TurboQuant as a quantized format.

5. Implement CUDA/Triton Kernels

Develop two key operations:

• Encode kernel: Quantize K/V tensors to codebook indices for cache storage
• Decode kernel: Reconstruct K/V tensors from indices before attention computation

These would integrate with the existing attention backend system.

6. Memory Management

Update KVCacheSpec calculations to account for:

• Reduced storage from vector quantization
• Additional codebook memory overhead
• Variable compression ratios based on bit-width

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.