Add TurboQuant KV cache compression for prefix cache

Adds --turbo-kv-bits flag (1-4) to compress stored prefix cache entries
using TurboQuant (arXiv 2504.19874). 3-bit gives 4.6x compression vs FP16,
compared to ~2x from the existing 8-bit quantization.

Integration points:
- memory_cache.py: _turbo_quantize_cache/_dequantize_cache, memory estimation,
  trim support, needs_dequantize property, config validation
- scheduler.py: turbo_kv_bits in SchedulerConfig, propagation to MemoryCacheConfig
- cli.py: --turbo-kv-bits for serve and bench commands

Requires mlx-lm with TurboQuant support (ml-explore/mlx-lm#1067).

Author: arozanov

vllm_mlx/cli.py

@@ -150,6 +150,8 @@ def serve_command(args):
             kv_cache_quantization_bits=args.kv_cache_quantization_bits,
             kv_cache_quantization_group_size=args.kv_cache_quantization_group_size,
             kv_cache_min_quantize_tokens=args.kv_cache_min_quantize_tokens,
+            # TurboQuant
+            turbo_kv_bits=args.turbo_kv_bits,
         )
 
         print("Mode: Continuous batching (for multiple concurrent users)")
@@ -169,7 +171,9 @@ def serve_command(args):
                 else f"{args.cache_memory_percent*100:.0f}% of RAM"
             )
             print(f"Memory-aware cache: {cache_info}")
-            if args.kv_cache_quantization:
+            if args.turbo_kv_bits:
+                print(f"TurboQuant: {args.turbo_kv_bits}-bit prefix cache compression")
+            elif args.kv_cache_quantization:
                 print(
                     f"KV cache quantization: {args.kv_cache_quantization_bits}-bit, "
                     f"group_size={args.kv_cache_quantization_group_size}"
@@ -248,6 +252,8 @@ async def run_benchmark():
             kv_cache_quantization_bits=args.kv_cache_quantization_bits,
             kv_cache_quantization_group_size=args.kv_cache_quantization_group_size,
             kv_cache_min_quantize_tokens=args.kv_cache_min_quantize_tokens,
+            # TurboQuant
+            turbo_kv_bits=args.turbo_kv_bits,
         )
         engine_config = EngineConfig(
             model_name=args.model,
@@ -687,6 +693,16 @@ def main():
         default=256,
         help="Minimum tokens for quantization to apply (default: 256)",
     )
+    # TurboQuant KV cache compression (arXiv 2504.19874)
+    serve_parser.add_argument(
+        "--turbo-kv-bits",
+        type=int,
+        default=None,
+        choices=[1, 2, 3, 4],
+        help="TurboQuant KV cache compression bits for prefix cache. "
+        "3-bit gives 4.6x compression vs FP16 (default: disabled). "
+        "Replaces --kv-cache-quantization when set.",
+    )
     serve_parser.add_argument(
         "--stream-interval",
         type=int,
@@ -966,6 +982,15 @@ def main():
         default=256,
         help="Minimum tokens for quantization to apply (default: 256)",
     )
+    # TurboQuant KV cache compression
+    bench_parser.add_argument(
+        "--turbo-kv-bits",
+        type=int,
+        default=None,
+        choices=[1, 2, 3, 4],
+        help="TurboQuant KV cache compression bits for prefix cache. "
+        "3-bit gives 4.6x compression vs FP16 (default: disabled).",
+    )
     # Paged cache options (experimental)
     bench_parser.add_argument(
         "--use-paged-cache",

vllm_mlx/memory_cache.py

@@ -103,7 +103,18 @@ def estimate_kv_cache_memory(cache: list[Any]) -> int:
 
     total_bytes = 0
 
+    try:
+        from mlx_lm.models.turboquant_cache import TurboQuantKVCache as _TQ
+    except ImportError:
+        _TQ = None
+
     for layer_cache in cache:
+        # Handle TurboQuantKVCache — estimate from shape+dtype (avoid lazy eval)
+        if _TQ is not None and isinstance(layer_cache, _TQ):
+            if layer_cache.k_packed is not None:
+                for arr in layer_cache.state:
+                    total_bytes += _array_memory(arr)
+            continue
         # Handle different cache object types
         # Check dict first since dicts have .keys() method that would match below
         if isinstance(layer_cache, dict) and "state" in layer_cache:
@@ -165,6 +176,7 @@ class MemoryCacheConfig:
     kv_bits: int = 8
     kv_group_size: int = 64
     kv_min_quantize_tokens: int = 256
+    turbo_kv_bits: int | None = None  # TurboQuant: 1-4 bit, 4.6x compression at 3-bit
 
     def __post_init__(self) -> None:
         if not 0.0 < self.max_memory_percent <= 1.0:
@@ -177,6 +189,15 @@ def __post_init__(self) -> None:
             raise ValueError(
                 f"kv_min_quantize_tokens must be >= 0, got {self.kv_min_quantize_tokens}"
             )
+        if self.turbo_kv_bits is not None and self.turbo_kv_bits not in (1, 2, 3, 4):
+            raise ValueError(
+                f"turbo_kv_bits must be 1-4, got {self.turbo_kv_bits}"
+            )
+
+    @property
+    def needs_dequantize(self) -> bool:
+        """Whether stored caches need dequantization on fetch."""
+        return self.kv_quantize or self.turbo_kv_bits is not None
 
     def compute_memory_limit(self) -> int:
         """
@@ -255,14 +276,14 @@ def create(cls, tokens: list[int], cache: list[Any]) -> _CacheEntry:
 
 
 def _trim_cache_offset(cache: list[Any], trim_by: int) -> list[Any]:
-    """Create shallow copies of KVCache/QuantizedKVCache layers with offset reduced.
+    """Create shallow copies of KVCache/QuantizedKVCache/TurboQuantKVCache layers
+    with offset reduced.
 
     This is used when returning a cached KV state to the scheduler so that
     the last N positions are "freed" and the model will recompute them on the
     next forward pass (preventing duplicate KV entries).
 
-    Supports both KVCache (keys/values are arrays) and QuantizedKVCache
-    (keys/values are 3-tuples of arrays).
+    Supports KVCache, QuantizedKVCache, and TurboQuantKVCache.
     """
     from mlx_lm.models.cache import KVCache
 
@@ -271,6 +292,11 @@ def _trim_cache_offset(cache: list[Any], trim_by: int) -> list[Any]:
     except ImportError:
         QuantizedKVCache = None  # noqa: N806
 
+    try:
+        from mlx_lm.models.turboquant_cache import TurboQuantKVCache
+    except ImportError:
+        TurboQuantKVCache = None  # noqa: N806
+
     trimmed: list[Any] = []
     for layer_cache in cache:
         if QuantizedKVCache is not None and isinstance(layer_cache, QuantizedKVCache):
@@ -281,6 +307,18 @@ def _trim_cache_offset(cache: list[Any], trim_by: int) -> list[Any]:
             tc.group_size = layer_cache.group_size
             tc.bits = layer_cache.bits
             trimmed.append(tc)
+        elif TurboQuantKVCache is not None and isinstance(
+            layer_cache, TurboQuantKVCache
+        ):
+            # Shallow copy with adjusted offset (do NOT mutate original)
+            tc = TurboQuantKVCache.__new__(TurboQuantKVCache)
+            tc.__dict__.update(layer_cache.__dict__)
+            tc.offset = max(layer_cache.offset - trim_by, 0)
+            tc._k_deq_buf = None  # invalidate decode buffer
+            tc._v_deq_buf = None
+            tc._deq_offset = 0
+            tc._deq_alloc = 0
+            trimmed.append(tc)
         elif (
             hasattr(layer_cache, "offset")
             and hasattr(layer_cache, "keys")
@@ -366,11 +404,33 @@ def _quantize_cache(cache: list[Any], bits: int = 8, group_size: int = 64) -> li
     return quantized
 
 
+def _turbo_quantize_cache(cache: list[Any], bits: int = 3) -> list[Any]:
+    """Compress KVCache layers with TurboQuant (4.6x at 3-bit).
+
+    Uses PolarQuant: randomized Hadamard rotation + Lloyd-Max codebook
+    quantization with fused Metal kernels. See arXiv 2504.19874.
+    """
+    from mlx_lm.models.cache import KVCache
+
+    compressed = []
+    for layer in cache:
+        if isinstance(layer, KVCache) and layer.keys is not None:
+            compressed.append(layer.to_turbo_quantized(bits=bits))
+        else:
+            compressed.append(layer)
+    return compressed
+
+
 def _dequantize_cache(cache: list[Any]) -> list[Any]:
-    """Dequantize QuantizedKVCache layers back to regular KVCache."""
+    """Dequantize QuantizedKVCache or TurboQuantKVCache layers back to KVCache."""
     import mlx.core as mx
     from mlx_lm.models.cache import KVCache, QuantizedKVCache
 
+    try:
+        from mlx_lm.models.turboquant_cache import TurboQuantKVCache
+    except ImportError:
+        TurboQuantKVCache = None  # noqa: N806
+
     result = []
     for layer in cache:
         if isinstance(layer, QuantizedKVCache) and layer.keys is not None:
@@ -383,6 +443,23 @@ def _dequantize_cache(cache: list[Any]) -> list[Any]:
             )
             kv.offset = layer.offset
             result.append(kv)
+        elif TurboQuantKVCache is not None and isinstance(layer, TurboQuantKVCache) and not layer.empty():
+            # Ensure quantizer is initialized (needed after from_state)
+            if layer._k_q is None:
+                layer._ensure_quantizer(layer._k_dim, layer._v_dim)
+            B, H = layer.k_packed.shape[:2]
+            dtype = layer._dtype if layer._dtype is not None else mx.float16
+            k_all = layer._full_dequant(
+                layer.k_packed, layer.k_norms, layer._k_q,
+                layer._k_dim, B, H, layer.offset, dtype,
+            )
+            v_all = layer._full_dequant(
+                layer.v_packed, layer.v_norms, layer._v_q,
+                layer._v_dim, B, H, layer.offset, dtype,
+            )
+            kv = KVCache()
+            kv.update_and_fetch(k_all, v_all)
+            result.append(kv)
         else:
             result.append(layer)
     return result
@@ -481,7 +558,7 @@ def fetch(self, tokens: list[int]) -> tuple[list[Any] | None, list[int]]:
             self._last_match_type = "exact"
             cache_out = (
                 _dequantize_cache(entry.cache)
-                if self._config.kv_quantize
+                if self._config.needs_dequantize
                 else entry.cache
             )
             return cache_out, []
@@ -539,7 +616,7 @@ def fetch(self, tokens: list[int]) -> tuple[list[Any] | None, list[int]]:
             excess = n_cached - n_requested
 
             has_non_trimmable = any(
-                not (hasattr(lc, "offset") and hasattr(lc, "keys"))
+                not (hasattr(lc, "is_trimmable") and lc.is_trimmable())
                 for lc in best_super.cache
             )
 
@@ -556,7 +633,7 @@ def fetch(self, tokens: list[int]) -> tuple[list[Any] | None, list[int]]:
                 self._last_match_type = "supersequence"
                 trimmed_cache = (
                     _dequantize_cache(trimmed_cache)
-                    if self._config.kv_quantize
+                    if self._config.needs_dequantize
                     else trimmed_cache
                 )
                 return trimmed_cache, []
@@ -567,7 +644,7 @@ def fetch(self, tokens: list[int]) -> tuple[list[Any] | None, list[int]]:
                 self._last_match_type = "supersequence"
                 cache_out = (
                     _dequantize_cache(best_super.cache)
-                    if self._config.kv_quantize
+                    if self._config.needs_dequantize
                     else best_super.cache
                 )
                 return cache_out, []
@@ -581,7 +658,7 @@ def fetch(self, tokens: list[int]) -> tuple[list[Any] | None, list[int]]:
             self._last_match_type = "prefix"
             cache_out = (
                 _dequantize_cache(best_match.cache)
-                if self._config.kv_quantize
+                if self._config.needs_dequantize
                 else best_match.cache
             )
             return cache_out, remaining
@@ -624,7 +701,7 @@ def fetch(self, tokens: list[int]) -> tuple[list[Any] | None, list[int]]:
             excess = len(best_lcp_entry.tokens) - best_lcp_length
 
             has_non_trimmable = any(
-                not (hasattr(lc, "offset") and hasattr(lc, "keys"))
+                not (hasattr(lc, "is_trimmable") and lc.is_trimmable())
                 for lc in best_lcp_entry.cache
             )
             logger.debug(
@@ -648,7 +725,7 @@ def fetch(self, tokens: list[int]) -> tuple[list[Any] | None, list[int]]:
                 self._last_match_type = "lcp"
                 trimmed_cache = (
                     _dequantize_cache(trimmed_cache)
-                    if self._config.kv_quantize
+                    if self._config.needs_dequantize
                     else trimmed_cache
                 )
                 return trimmed_cache, remaining
@@ -693,8 +770,13 @@ def store(
         # Trim oversized KV arrays to actual used size
         cache = _trim_to_offset(cache)
 
-        # Quantize if enabled and sequence is long enough
+        # Compress KV cache for storage: TurboQuant (4.6x) or standard quantization (2x)
         if (
+            self._config.turbo_kv_bits is not None
+            and len(tokens) >= self._config.kv_min_quantize_tokens
+        ):
+            cache = _turbo_quantize_cache(cache, self._config.turbo_kv_bits)
+        elif (
             self._config.kv_quantize
             and len(tokens) >= self._config.kv_min_quantize_tokens
         ):

vllm_mlx/scheduler.py

@@ -77,6 +77,9 @@ class SchedulerConfig:
     kv_cache_quantization_group_size: int = 64
     kv_cache_min_quantize_tokens: int = 256
 
+    # TurboQuant KV cache compression (4.6x at 3-bit, replaces standard quantization)
+    turbo_kv_bits: Optional[int] = None  # 1-4 bit; None = disabled
+
     # Paged cache settings (experimental - for memory efficiency)
     use_paged_cache: bool = (
         False  # Use BlockAwarePrefixCache instead of PrefixCacheManager
@@ -1020,6 +1023,7 @@ def __init__(
                     kv_bits=self.config.kv_cache_quantization_bits,
                     kv_group_size=self.config.kv_cache_quantization_group_size,
                     kv_min_quantize_tokens=self.config.kv_cache_min_quantize_tokens,
+                    turbo_kv_bits=self.config.turbo_kv_bits,
                 )
                 self.memory_aware_cache = MemoryAwarePrefixCache(
                     model=model,