Expose existing KV-cache quantization behind an env-var opt-in

[alytaphoenix]

Summary

This started as an investigation into adding Google's TurboQuant (rotation + Lloyd-Max KV-cache compression, Google Research blog, arXiv) as a memory-reduction option. Findings on that:

• It's not merged into mlx or mlx_lm upstream — there are two open feature requests (ml-explore/mlx-lm#1060, ml-explore/mlx#3404) and a handful of small third-party MLX reimplementations, which look like solo/unvetted efforts, not something to pull in as a dependency.
• Note: some sources use "TurboQuant" and "PolarQuant" (a related ICLR'26 technique) somewhat interchangeably — treat the exact technique/claims as attributed to the primary sources above, not independently verified here.

Building real TurboQuant support would mean writing correctness-critical rotation+quantization math ourselves with no canonical reference — out of scope for now.

What this PR actually does instead: exposes a KV-cache quantization capability that already exists in this codebase but is unreachable. make_kv_cache() (cache.py) already constructs QuantizedKVCache when KV_CACHE_BITS is set, and every generation call site (prefill(), mlx_generate(), pipeline_parallel_prefill() in generate.py) already passes kv_bits/kv_group_size into mlx_lm's own stream_generate/maybe_quantize_kv_cache. Both were hardcoded to None in constants.py with no config surface — pure dead code. This PR turns them into env vars:

• EXO_KV_CACHE_BITS: bits to quantize the KV cache to (e.g. 4 or 8). Unset (default) preserves current full-precision behavior — opt-in, no behavior change by default.
• EXO_KV_CACHE_GROUP_SIZE: quantization group size, default 64.

Notable side-effect: group-size consolidation

There were two separate, inconsistent group-size constants: KV_GROUP_SIZE=32 (used by the stream_generate/pipeline-prefill path) and CACHE_GROUP_SIZE=64 (used by make_kv_cache's direct construction). This PR unifies them into one KV_CACHE_GROUP_SIZE (default 64). This is harmless while quantization is off (group size is ignored when kv_bits=None), but it means once someone sets EXO_KV_CACHE_BITS, the stream_generate path now quantizes at group 64 instead of the old 32 — coarser grouping, smaller relative scale/bias overhead, likely a minor improvement, but flagging it as a behavior change beyond pure exposure.

Investigated but not resolved: prefix-cache + quantization interaction

Before shipping this, I checked whether exo's prefix-cache reuse (KVPrefixCache in cache.py) actually handles QuantizedKVCache correctly, since that seemed like the obvious risk. It does: _entry_length explicitly lists QuantizedKVCache in its type union with a comment, and its rollback goes through the shared _BaseCache.trim()/.offset interface — the same generic path used for plain KVCache — not the RotatingKVCache/ArraysCache-specific snapshot machinery. This looks like it was deliberately built to support QuantizedKVCache, just never wired to a config knob.

What's still unverified: actual behavior on real Apple Silicon hardware — this development environment has no mlx runtime, so nothing here has been exercised end-to-end.

One open interaction, not addressed here: #2240 (context-window support, not yet merged) makes make_kv_cache build a RotatingKVCache when a context cap (max_context_length) is set. Whether mlx_lm's maybe_quantize_kv_cache (invoked via kv_bits on every stream_generate/pipeline_parallel_prefill call, unconditionally of cache type) behaves correctly when handed a RotatingKVCache instead of a plain KVCache is untested. Combining a context cap with KV quantization is a combination this PR doesn't validate — flagging for whoever tests this on hardware.

Test plan

• uv run basedpyright (full repo, not just touched files) — 319 errors / 110 warnings, identical to main baseline (all pre-existing, due to mlx not being installed in this environment)
• uv run ruff check (full repo) — all checks passed
• uv run ruff format --check — all touched files already formatted
• uv run pytest src/exo (excluding mlx-dependent test dirs, which can't collect without mlx installed) — 321 passed, 3 skipped, no regressions
• Manually verified env var parsing: EXO_KV_CACHE_BITS=4 EXO_KV_CACHE_GROUP_SIZE=32 correctly sets the constants; unset correctly defaults to None/64 (current behavior preserved)
• nix fmt / nix flake check not run in this environment (ruff substituted)
• No on-device verification of actual quantized generation or its interaction with prefix-cache reuse or the context-window cap (see above)

🤖 Generated with Claude Code

https://claude.ai/code/session_011rjSfwDBTkmySmfU6NgHKF