Fix server batching condition for SSMs

mlx_lm/server.py

@@ -34,7 +34,13 @@
 
 from ._version import __version__
 from .generate import BatchGenerator, stream_generate
-from .models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
+from .models.cache import (
+    KVCache,
+    RotatingKVCache,
+    can_trim_prompt_cache,
+    make_prompt_cache,
+    trim_prompt_cache,
+)
 from .sample_utils import make_logits_processors, make_sampler
 from .utils import load
 
@@ -378,6 +384,7 @@ def __init__(self, cli_args: argparse.Namespace):
         self.model = None
         self.tokenizer = None
         self.draft_model = None
+        self.cache_types = set()
 
         # Preload the default model if it is provided
         self.default_model_map = {}
@@ -448,6 +455,15 @@ def validate_draft_tokenizer(draft_tokenizer):
         elif draft_model_path is not None and draft_model_path != "default_model":
             self.draft_model, draft_tokenizer = load(draft_model_path)
             validate_draft_tokenizer(draft_tokenizer)
+
+        # Figure out the cache types and save them in a set for anybody that
+        # wants to make a decision based on those.
+        for c in make_prompt_cache(self.model):
+            self.cache_types.add(type(c))
+        if self.draft_model is not None:
+            for c in make_prompt_cache(self.draft_model):
+                self.cache_types.add(type(c))
+
         return self.model, self.tokenizer
 
 
@@ -477,6 +493,7 @@ def _tokenize(self, tokenizer, request):
                     messages,
                     tools,
                     add_generation_prompt=True,
+                    tokenize=True,
                     **self.model_provider.cli_args.chat_template_args,
                 )
             else:
@@ -490,6 +507,9 @@ def _is_batchable(self, args):
             or self.model_provider.cli_args.draft_model is not None
         ):
             return False
+        for c in self.model_provider.cache_types:
+            if c not in (KVCache, RotatingKVCache):
+                return False
         if args.logits.logit_bias is not None:
             return False
         if args.logits.repetition_penalty != 0: