fix(mlx_lm.server): fail fast when --draft-model set with non-trimmable cache

mlx_lm/server.py

@@ -40,12 +40,22 @@
 )
 from .models.cache import (
     LRUPromptCache,
+    can_trim_prompt_cache,
     make_prompt_cache,
 )
 from .sample_utils import make_logits_processors, make_sampler
 from .utils import _parse_size, load, sharded_load
 
 
+def _ensure_speculation_supported(prompt_cache, model_path):
+    if not can_trim_prompt_cache(prompt_cache):
+        non_trimmable = {type(c).__name__ for c in prompt_cache if not c.is_trimmable()}
+        raise ValueError(
+            f"--draft-model was set, but speculative decoding requires a trimmable "
+            f"prompt cache; '{model_path}' uses non-trimmable cache(s): {non_trimmable}"
+        )
+
+
 def get_system_fingerprint():
     gpu_arch = mx.device_info()["architecture"]
     return f"{__version__}-{mx.__version__}-{platform.platform()}-{gpu_arch}"
@@ -357,6 +367,11 @@ def _load(self, model_path, adapter_path=None, draft_model_path=None):
             if tokenizer.chat_template is None:
                 tokenizer.chat_template = tokenizer.default_chat_template
 
+        # Speculative decoding needs a trimmable target cache, fail fast if not
+        prompt_cache = make_prompt_cache(model)
+        if draft_model_path is not None:
+            _ensure_speculation_supported(prompt_cache, model_path)
+
         # Load the draft model for speculative decoding
         draft_model = None
         if draft_model_path is not None:
@@ -369,9 +384,7 @@ def _load(self, model_path, adapter_path=None, draft_model_path=None):
 
         # Compute batchability
         is_batchable = draft_model is None
-        is_batchable = is_batchable and all(
-            hasattr(c, "merge") for c in make_prompt_cache(model)
-        )
+        is_batchable = is_batchable and all(hasattr(c, "merge") for c in prompt_cache)
 
         # Update the member variables
         self.model_key = (model_path, adapter_path, draft_model_path)

tests/test_server.py

@@ -16,6 +16,7 @@
     LRUPromptCache,
     Response,
     ResponseGenerator,
+    _ensure_speculation_supported,
     _process_control_tokens,
 )
 from mlx_lm.utils import load
@@ -157,6 +158,26 @@ def test_multi_token_match_preserves_order(self):
         )
 
 
+class TestSpeculationValidation(unittest.TestCase):
+    def test_rejects_non_trimmable_cache(self):
+        with self.assertRaises(ValueError):
+            _ensure_speculation_supported(
+                [MockCache("aaa", is_trimmable=False)], "test/model"
+            )
+
+    def test_allows_trimmable_cache(self):
+        # Should not raise
+        _ensure_speculation_supported(
+            [MockCache("aaa", is_trimmable=True)], "test/model"
+        )
+
+    def test_error_message_names_bad_cache_type(self):
+        with self.assertRaisesRegex(ValueError, "MockCache"):
+            _ensure_speculation_supported(
+                [MockCache("aaa", is_trimmable=False)], "test/model"
+            )
+
+
 class TestServer(unittest.TestCase):
     @classmethod
     def setUpClass(cls):