Align quantized_kv_start defaults with the CLI

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,

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