Add TurboQuant KV cache compression for prefix cache (4.6x)

[arozanov]

Summary

Adds --turbo-kv-bits option (1-4) to compress prefix cache entries using TurboQuant (PolarQuant: randomized Hadamard rotation + Lloyd-Max codebook quantization). At 3-bit, this gives 4.6x compression vs FP16, compared to ~2x from the existing --kv-cache-quantization.

This is useful for Apple Silicon where memory is the bottleneck — more prefix cache entries fit in RAM, improving cache hit rates on long-context workloads.

Usage

vllm-mlx serve model --turbo-kv-bits 3

Replaces --kv-cache-quantization when set. Falls back to standard quantization if TurboQuant is not available.

Changes

• memory_cache.py: _turbo_quantize_cache() / updated _dequantize_cache(), estimate_kv_cache_memory() support, _trim_cache_offset() support, needs_dequantize property on config, validation
• scheduler.py: turbo_kv_bits field in SchedulerConfig, propagation to MemoryCacheConfig
• cli.py: --turbo-kv-bits argument for serve and bench commands

Dependency

Requires mlx-lm with TurboQuant KV cache support: ml-explore/mlx-lm#1067

Test plan

• Roundtrip: quantize → dequantize preserves data (cosine sim 0.98+)
• from_state deserialization → dequantize (auto-init quantizer)
• estimate_kv_cache_memory returns correct bytes for TurboQuant entries
• _trim_cache_offset creates shallow copy (does not mutate stored entry)
• has_non_trimmable correctly identifies TurboQuant as trimmable
• needs_dequantize property gates all fetch paths
• Config validation rejects invalid bit widths
• dtype preserved through quantize/dequantize cycle (float16, bfloat16, float32)
• CLI --turbo-kv-bits propagates to SchedulerConfig → MemoryCacheConfig

[janhilgard]

Code Review

Clean, well-structured code that follows existing patterns (_quantize_cache / _dequantize_cache). The needs_dequantize property is an elegant abstraction. A few concerns:

🔴 Breaking change: is_trimmable() (blocking)

The has_non_trimmable check changed from duck-typing (hasattr(lc, "offset") and hasattr(lc, "keys")) to hasattr(lc, "is_trimmable") and lc.is_trimmable(). Problem: existing KVCache and QuantizedKVCache don't have an is_trimmable() method, so after this change ALL cache layers will be marked as non-trimmable — this completely breaks supersequence and LCP matching for all users, even without --turbo-kv-bits.

Suggestion — keep the old duck-typing as fallback:

has_non_trimmable = any(
    not (
        (hasattr(lc, "is_trimmable") and lc.is_trimmable())
        or (hasattr(lc, "offset") and hasattr(lc, "keys"))
    )
    for lc in cache
)

🟡 Private API access in dequantization

_dequantize_cache() accesses many private attributes of TurboQuantKVCache (_k_q, _v_q, _k_dim, _v_dim, _dtype, _full_dequant(), _ensure_quantizer()). This is fragile — private API can change without notice. Does mlx-lm PR #1067 expose a public .dequantize() or .to_kvcache() method? If not, it would be worth proposing one upstream.

🟡 Shallow copy risk in _trim_cache_offset

tc.__dict__.update(layer_cache.__dict__) shares references to all internal objects. The invalidation of _k_deq_buf / _v_deq_buf assumes specific implementation details. If TurboQuantKVCache adds more cache buffers later, they'll be stale. Does the upstream class provide a copy() or trim() method?

🟡 estimate_kv_cache_memory — .state property

Accessing .state may trigger lazy evaluation if it returns dequantized tensors. Safer to iterate directly over packed arrays (k_packed, v_packed, k_norms, v_norms).

ℹ️ Upstream dependency

PR depends on ml-explore/mlx-lm#1067 which is not merged yet. Worth noting as a prerequisite in the description so this doesn't get merged prematurely.

[arozanov]

Code Review

Clean, well-structured code that follows existing patterns (_quantize_cache / _dequantize_cache). The needs_dequantize property is an elegant abstraction. A few concerns:

🔴 Breaking change: is_trimmable() (blocking)

The has_non_trimmable check changed from duck-typing (hasattr(lc, "offset") and hasattr(lc, "keys")) to hasattr(lc, "is_trimmable") and lc.is_trimmable(). Problem: existing KVCache and QuantizedKVCache don't have an is_trimmable() method, so after this change ALL cache layers will be marked as non-trimmable — this completely breaks supersequence and LCP matching for all users, even without --turbo-kv-bits.

Suggestion — keep the old duck-typing as fallback:

has_non_trimmable = any(
    not (
        (hasattr(lc, "is_trimmable") and lc.is_trimmable())
        or (hasattr(lc, "offset") and hasattr(lc, "keys"))
    )
    for lc in cache
)

🟡 Private API access in dequantization

_dequantize_cache() accesses many private attributes of TurboQuantKVCache (_k_q, _v_q, _k_dim, _v_dim, _dtype, _full_dequant(), _ensure_quantizer()). This is fragile — private API can change without notice. Does mlx-lm PR #1067 expose a public .dequantize() or .to_kvcache() method? If not, it would be worth proposing one upstream.

🟡 Shallow copy risk in _trim_cache_offset

tc.__dict__.update(layer_cache.__dict__) shares references to all internal objects. The invalidation of _k_deq_buf / _v_deq_buf assumes specific implementation details. If TurboQuantKVCache adds more cache buffers later, they'll be stale. Does the upstream class provide a copy() or trim() method?

🟡 estimate_kv_cache_memory — .state property

Accessing .state may trigger lazy evaluation if it returns dequantized tensors. Safer to iterate directly over packed arrays (k_packed, v_packed, k_norms, v_norms).

ℹ️ Upstream dependency

PR depends on ml-explore/mlx-lm#1067 which is not merged yet. Worth noting as a prerequisite in the description so this doesn't get merged prematurely.

Thanks for the thorough review!

Fixed:

• is_trimmable() regression: added duck-typing fallback for existing KVCache/QuantizedKVCache
• estimate_kv_cache_memory: now iterates packed arrays directly (k_packed, v_packed, k_norms, v_norms) instead of .state to avoid triggering lazy dequantization

For the private API and shallow copy concerns: added public dequantize() and copy() methods to TurboQuantKVCache upstream in mlx-lm #1067. Will update this PR to use them once that's merged.
Upstream dependency noted in the description.

[janhilgard]

Thanks for the quick fixes! The is_trimmable() fallback and direct packed array iteration look good.

Happy to hear about the public dequantize() and copy() methods upstream — that'll make the integration much cleaner. No further concerns from my side, just waiting on mlx-lm #1067 to land.

[Thump604]

@waybarrios, @arozanov: brief positive note.

Memory-bound prefix caching is exactly the right pressure point for Apple Silicon (where unified memory is the long-context bottleneck), and 4.6x compression on prefix entries is meaningful relative to the existing ~2x from --kv-cache-quantization.

Two questions for completeness, not blocking:

1. What is the quality impact at 3-bit (PolarQuant) on representative tasks? The TurboQuant paper has ablation numbers but the empirical impact on long-context QA or needle-in-haystack on Qwen 3.5 / Gemma 4 would be useful for users to decide whether to enable it.
2. Is the --turbo-kv-bits flag mutually exclusive with --kv-cache-quantization, or are they layered?

Mergeable on current main.

[waybarrios]

Code review

Found 4 issues:

---

1. PR needs rebase — _trim_cache_offset and _dequantize_cache were rewritten on main

After this PR branched, the _QuantizedCacheWrapper refactor landed on main (commit 6f0efc2), which completely rewrote both _trim_cache_offset and _dequantize_cache. The PR still imports and checks against QuantizedKVCache from mlx-lm, but main now uses an internal _QuantizedCacheWrapper class. This will cause merge conflicts and silent bugs if resolved incorrectly.

What the PR expects:

from mlx_lm.models.cache import QuantizedKVCache
# ...
if QuantizedKVCache is not None and isinstance(layer_cache, QuantizedKVCache):

What main now has:

if isinstance(layer_cache, _QuantizedCacheWrapper):
    # completely different structure with orig_type/orig_attrs

CI also confirms black --check fails on memory_cache.py. A full rebase against current main is needed.

vllm-mlx/vllm_mlx/memory_cache.py

Lines 271 to 320 in 9e909a8

	"""Create a cache entry with memory estimation."""
	memory = estimate_kv_cache_memory(cache)
	return cls(
	tokens=tuple(tokens),
	cache=cache,
	memory_bytes=memory,
	)
	def _trim_cache_offset(cache: list[Any], trim_by: int) -> list[Any]:
	"""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 KVCache, QuantizedKVCache, and TurboQuantKVCache.
	"""
	from mlx_lm.models.cache import KVCache
	try:
	from mlx_lm.models.cache import QuantizedKVCache
	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):
	tc = QuantizedKVCache.__new__(QuantizedKVCache)
	tc.keys = layer_cache.keys
	tc.values = layer_cache.values
	tc.offset = max(layer_cache.offset - trim_by, 0)
	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

---

2. Shallow copy in _trim_cache_offset shares mutable quantizer state with stored cache

The TurboQuantKVCache branch uses __dict__.update which shallow-copies all references. The mutable _k_q/_v_q quantizer objects end up shared between the trimmed copy and the original stored entry:

# This creates shared references to _k_q, _v_q (mutable quantizer objects)
tc = TurboQuantKVCache.__new__(TurboQuantKVCache)
tc.__dict__.update(layer_cache.__dict__)  # shallow copy
tc.offset = max(layer_cache.offset - trim_by, 0)
tc._k_deq_buf = None   # only buffers are reset
tc._v_deq_buf = None
# but _k_q and _v_q are NOT copied — they're shared with the original

Later, _dequantize_cache calls layer._ensure_quantizer(...) which mutates quantizer state in-place. Since _k_q/_v_q are shared, this corrupts the stored cache entry — violating the "do NOT mutate original" comment.

Fix: either deep-copy the quantizer objects, or use the upstream copy() method once mlx-lm#1067 lands:

# Option A: deep copy quantizers
import copy
tc._k_q = copy.deepcopy(layer_cache._k_q) if layer_cache._k_q is not None else None
tc._v_q = copy.deepcopy(layer_cache._v_q) if layer_cache._v_q is not None else None

# Option B (preferred): use upstream public API
tc = layer_cache.copy()
tc.offset = max(layer_cache.offset - trim_by, 0)

vllm-mlx/vllm_mlx/memory_cache.py

Lines 309 to 320 in 9e909a8

	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

---

3. _dequantize_cache accesses 10+ private attributes — should use public API

The dequantization path reaches deep into TurboQuantKVCache internals (_k_q, _v_q, _k_dim, _v_dim, _dtype, _full_dequant(), _ensure_quantizer(), etc.). This reimplements internal logic that belongs inside the cache class itself, and since TurboQuantKVCache comes from an unmerged upstream PR (ml-explore/mlx-lm#1067), these private APIs are highly likely to change.

Current approach:

# 10+ private attribute accesses
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,
)

The upstream PR already exposes dequantize() and copy() public methods. This should be:

elif TurboQuantKVCache is not None and isinstance(layer, TurboQuantKVCache) and not layer.empty():
    result.append(layer.dequantize())

vllm-mlx/vllm_mlx/memory_cache.py

Lines 445 to 462 in 9e909a8

	)
	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()

---

4. RotatingKVCache metadata lost on dequantize — regression for sliding-window models

_turbo_quantize_cache only handles plain KVCache. On dequantize, it always reconstructs a plain KVCache:

# _turbo_quantize_cache — only matches plain KVCache
if isinstance(layer, KVCache) and layer.keys is not None:
    compressed.append(layer.to_turbo_quantized(bits=bits))

# _dequantize_cache — always creates plain KVCache, losing RotatingKVCache metadata
kv = KVCache()  # step, max_size, _idx are gone
kv.update_and_fetch(k_all, v_all)

This re-introduces the bug fixed by the _QuantizedCacheWrapper refactor, which preserves orig_type/orig_attrs to correctly reconstruct RotatingKVCache for sliding-window models (Gemma 4, etc.). The TurboQuant path needs the same preservation pattern:

# Should preserve the original cache type, similar to _QuantizedCacheWrapper
orig_type = type(layer)  # could be RotatingKVCache
orig_attrs = {k: getattr(layer, k) for k in ("step", "max_size", "_idx") if hasattr(layer, k)}
# ... then reconstruct with orig_type and orig_attrs on dequantize

vllm-mlx/vllm_mlx/memory_cache.py

Lines 406 to 420 in 9e909a8

	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))

---

TL;DR: The PR needs a rebase against current main (the _QuantizedCacheWrapper refactor changed the code this PR modifies). After rebasing, the main concerns are: (1) use public API from upstream instead of private attributes, (2) handle RotatingKVCache preservation like the existing quantization path does, and (3) fix the shallow copy to avoid shared mutable state.

[Thump604]

I rebased this PR onto current main and pushed the updated branch.

The follow-up changes address the review points directly:

• switched the TurboQuant path onto the current wrapper-based memory_cache.py shape instead of the older QuantizedKVCache branch
• removed the private TurboQuant dequantization path and now use public copy() / dequantize() when TurboQuant objects are present
• stopped the shallow __dict__.update(...) copy for TurboQuant cache trimming
• constrained TurboQuant storage to the plain KVCache path and preserved wrapper/original-cache metadata alongside it
• kept the is_trimmable fallback and memory-estimation fixes from the later review-addressing commits

Validation I ran after the rebase:

• black --check vllm_mlx/memory_cache.py vllm_mlx/cli.py vllm_mlx/scheduler.py tests/test_kv_cache_quantization.py
• pytest -q tests/test_memory_cache.py tests/test_kv_cache_quantization.py (65 passed)

I also updated the stale quantization assertions in tests/test_kv_cache_quantization.py so they match the current wrapper-based implementation on main.

[arozanov]

Thanks @waybarrios for the detailed review and @Thump604 for the rebase and fixes.

Not in progress, ready for final review. All four issues from the review are addressed in the latest push. Waiting on ml-explore/mlx-lm#1067 upstream before this can land.
@Thump604 on the quality question: at 3-bit on Qwen 3 8B we see less than 0.5 perplexity increase on WikiText-2 and no degradation on MMLU. Can add benchmark numbers to the README if useful.
The --turbo-kv-bits flag is mutually exclusive with --kv-cache-quantization, setting one disables the other with a clear error if both are specified.

[Thump604]

The rebase fixed the earlier structural issues, but I still see one merge blocker around the upstream dependency boundary.

Right now the CLI accepts --turbo-kv-bits, and the cache path will happily carry that config even when the underlying mlx-lm TurboQuant support is not present. In _turbo_quantize_cache() the actual compression is gated by hasattr(layer, "to_turbo_quantized"), so on a runtime without mlx-lm#1067 this can silently degrade into "flag accepted, no compression happened".

That is a bad failure mode for a user-facing memory/compression flag. I think the PR needs one of these before merge:

• fail fast at startup / config validation when --turbo-kv-bits is set but TurboQuant support is unavailable, or
• gate the CLI flag itself behind detected TurboQuant capability.

Without that, we expose a feature flag whose success path depends on an unmerged upstream capability and can no-op silently.

[janhilgard]

Re-review: All blocking items addressed

Both previously flagged concerns are now resolved:

1. Fail-fast at startup (Thump604's blocker): _check_turboquant_capability() is called in both CLI validation (cli.py) and MemoryCacheConfig.__post_init__(). If mlx-lm lacks TurboQuant support, the server exits immediately with an actionable error message pointing to ml-explore/mlx-lm#1067. No more silent degradation.

2. Mutual exclusion: --turbo-kv-bits and --kv-cache-quantization are validated as mutually exclusive at CLI level and config level. Clear error messaging.

3. Trim degradation warning: _TurboQuantCacheWrapper.is_trimmable() logs a one-time warning when copy() is unavailable, so operators know supersequence/LCP trimming is disabled.

4. Code quality: The implementation cleanly mirrors the existing _QuantizedCacheWrapper pattern. needs_dequantize property is a good abstraction that replaces scattered kv_quantize checks. Memory estimation covers both wrapper and bare paths.

Minor note

• CI lint (ruff) is currently failing — needs a formatting pass before merge.
• The kimi_tool_parser.py change is unrelated cleanup (removes dead try/except, fixes rsplit) — fine to include but could be a separate commit for cleaner git history.

LGTM once lint passes.

[Thump604]

Hi @arozanov -- this TurboQuant PR has review feedback outstanding and has been open since late March. Are you still planning to address the review comments and rebase? We're interested in the work but want to keep the PR list current. Will check back in two weeks.