diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index 7cdcc77fbe..6e802e86c8 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/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: diff --git a/src/exo/worker/engines/mlx/constants.py b/src/exo/worker/engines/mlx/constants.py index 86a663e424..58fb827cae 100644 --- a/src/exo/worker/engines/mlx/constants.py +++ b/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 diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 2e3d051251..4abd11fcf3 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/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, ):