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

src/exo/worker/engines/mlx/cache.py

@@ -22,7 +22,7 @@
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.memory import Memory
-from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
+from exo.worker.engines.mlx.constants import KV_CACHE_BITS, KV_CACHE_GROUP_SIZE
 from exo.worker.engines.mlx.types import KVCacheType, Model
 from exo.worker.runner.bootstrap import logger
 
@@ -573,7 +573,7 @@ def make_kv_cache(
         else:
             logger.info("Using quantized KV cache")
             return [
-                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
+                QuantizedKVCache(group_size=KV_CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
                 for _ in model.layers
             ]
     else:

src/exo/worker/engines/mlx/constants.py

@@ -1,15 +1,27 @@
+import os
+
 # TODO: Do we want so many constants?
 #  I think we want a lot of these as parameters?
 
-KV_GROUP_SIZE: int | None = 32
-KV_BITS: int | None = None
 ATTENTION_KV_BITS: int | None = 4
 MAX_TOKENS: int = 32168
 MAX_KV_SIZE: int | None = 3200
 KEEP_KV_SIZE: int | None = 1600
 QUANTIZE_MODEL_MODE: str | None = "affine"
-CACHE_GROUP_SIZE: int = 64
-KV_CACHE_BITS: int | None = None
+
+# Number of bits to quantize the KV cache to (mlx_lm's QuantizedKVCache, e.g.
+# 4 or 8). None (the default) keeps the cache in full precision. Opt-in via
+# env var: this path is wired into every generation code path here (both
+# make_kv_cache's direct QuantizedKVCache construction and mlx_lm's
+# maybe_quantize_kv_cache during stream_generate/pipeline prefill), and
+# exo's prefix-cache trim/snapshot logic (cache.py) already handles
+# QuantizedKVCache correctly via the shared _BaseCache.trim()/.offset
+# interface. What's unverified is real-hardware behavior -- turn on only
+# after testing on an actual cluster.
+KV_CACHE_BITS: int | None = (
+    int(os.environ["EXO_KV_CACHE_BITS"]) if "EXO_KV_CACHE_BITS" in os.environ else None
+)
+KV_CACHE_GROUP_SIZE: int = int(os.environ.get("EXO_KV_CACHE_GROUP_SIZE", "64"))
 
 DEFAULT_TOP_LOGPROBS: int = 5
 

src/exo/worker/engines/mlx/generator/generate.py

@@ -51,8 +51,8 @@
 )
 from exo.worker.engines.mlx.constants import (
     DEFAULT_TOP_LOGPROBS,
-    KV_BITS,
-    KV_GROUP_SIZE,
+    KV_CACHE_BITS,
+    KV_CACHE_GROUP_SIZE,
     MAX_TOKENS,
 )
 from exo.worker.engines.mlx.generator.remote_prefill import remote_prefill
@@ -342,8 +342,8 @@ def combined_progress_callback(processed: int, total: int) -> None:
                 prompt=prompt_tokens,
                 prompt_cache=cache,
                 prefill_step_size=prefill_step_size,
-                kv_group_size=KV_GROUP_SIZE,
-                kv_bits=KV_BITS,
+                kv_group_size=KV_CACHE_GROUP_SIZE,
+                kv_bits=KV_CACHE_BITS,
                 prompt_progress_callback=progress_callback,
                 distributed_prompt_progress_callback=distributed_prompt_progress_callback,
                 group=group,
@@ -359,8 +359,8 @@ def combined_progress_callback(processed: int, total: int) -> None:
                 sampler=sampler,
                 prompt_cache=cache,
                 prefill_step_size=prefill_step_size,
-                kv_group_size=KV_GROUP_SIZE,
-                kv_bits=KV_BITS,
+                kv_group_size=KV_CACHE_GROUP_SIZE,
+                kv_bits=KV_CACHE_BITS,
                 prompt_progress_callback=combined_progress_callback,
             ):
                 break  # Stop after first iteration - cache is now filled
@@ -727,8 +727,8 @@ def mlx_generate(
             logits_processors=logits_processors,
             prompt_cache=caches,
             prefill_step_size=1,
-            kv_group_size=KV_GROUP_SIZE,
-            kv_bits=KV_BITS,
+            kv_group_size=KV_CACHE_GROUP_SIZE,
+            kv_bits=KV_CACHE_BITS,
         ),
         start=1,
     ):