diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index a4ef4982e..d9debea55 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -307,7 +307,7 @@ def generate_step( prefill_step_size: int = 2048, kv_bits: Optional[int] = None, kv_group_size: int = 64, - quantized_kv_start: int = 0, + quantized_kv_start: int = DEFAULT_QUANTIZED_KV_START, prompt_progress_callback: Optional[Callable[[int, int], None]] = None, input_embeddings: Optional[mx.array] = None, ) -> Generator[Tuple[mx.array, mx.array], None, None]: @@ -333,7 +333,11 @@ def generate_step( None implies no cache quantization. Default: ``None``. kv_group_size (int): Group size for KV cache quantization. Default: ``64``. quantized_kv_start (int): Step to begin using a quantized KV cache. - when ``kv_bits`` is non-None. Default: ``0``. + when ``kv_bits`` is non-None. Default: ``DEFAULT_QUANTIZED_KV_START`` + (``5000``), matching the CLI. Quantizing from the first token costs + decode throughput while saving no meaningful memory on short + contexts, so quantization is deferred until the cache is long + enough for the trade to pay off. Pass ``0`` to quantize immediately. prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the prompt tokens processed so far and the total number of prompt tokens. input_embeddings (mx.array, optional): Input embeddings to use instead of or in @@ -474,7 +478,7 @@ def speculative_generate_step( prefill_step_size: int = 512, kv_bits: Optional[int] = None, kv_group_size: int = 64, - quantized_kv_start: int = 0, + quantized_kv_start: int = DEFAULT_QUANTIZED_KV_START, ) -> Generator[Tuple[mx.array, mx.array, bool], None, None]: """ A generator producing token ids based on the given prompt from the model. @@ -499,7 +503,8 @@ def speculative_generate_step( None implies no cache quantization. Default: ``None``. kv_group_size (int): Group size for KV cache quantization. Default: ``64``. quantized_kv_start (int): Step to begin using a quantized KV cache. - when ``kv_bits`` is non-None. Default: ``0``. + when ``kv_bits`` is non-None. Default: ``DEFAULT_QUANTIZED_KV_START`` + (``5000``), see ``generate_step``. Yields: Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities, diff --git a/tests/test_generate.py b/tests/test_generate.py index 77666e4b7..18deb4019 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -845,5 +845,37 @@ def test_batch_generate_no_logprobs_by_default(self): self.assertIsNone(response.token_ids) +class TestQuantizedKVStartDefault(unittest.TestCase): + """The library defaults must not disagree with the CLI's own default.""" + + def _default(self, fn): + import inspect + + return inspect.signature(fn).parameters["quantized_kv_start"].default + + def test_generate_step_matches_cli_default(self): + from mlx_lm.generate import DEFAULT_QUANTIZED_KV_START, generate_step + + self.assertEqual(self._default(generate_step), DEFAULT_QUANTIZED_KV_START) + + def test_speculative_generate_step_matches_cli_default(self): + from mlx_lm.generate import ( + DEFAULT_QUANTIZED_KV_START, + speculative_generate_step, + ) + + self.assertEqual( + self._default(speculative_generate_step), DEFAULT_QUANTIZED_KV_START + ) + + def test_cache_prompt_shares_the_same_constant(self): + from mlx_lm.cache_prompt import ( + DEFAULT_QUANTIZED_KV_START as CACHE_PROMPT_DEFAULT, + ) + from mlx_lm.generate import DEFAULT_QUANTIZED_KV_START + + self.assertEqual(CACHE_PROMPT_DEFAULT, DEFAULT_QUANTIZED_KV_START) + + if __name__ == "__main__": unittest.main()