Validate generation batch limits

mlx_lm/generate.py

@@ -673,6 +673,11 @@ def stream_generate(
         GenerationResponse: An instance containing the generated text segment and
             associated metadata. See :class:`GenerationResponse` for details.
     """
+    if max_tokens < 0:
+        raise ValueError("max_tokens must be non-negative")
+    if max_tokens == 0:
+        return
+
     if not isinstance(tokenizer, TokenizerWrapper):
         tokenizer = TokenizerWrapper(tokenizer)
 
@@ -1696,15 +1701,37 @@ def insert_segments(
     ):
         uids = []
 
-        max_tokens = max_tokens or [self.max_tokens] * len(segments)
-        all_tokens = all_tokens or [[] for _ in segments]
-        samplers = samplers or [None] * len(segments)
-        logits_processors = logits_processors or (
-            [self.logits_processors] * len(segments)
+        count = len(segments)
+        values = {
+            "max_tokens": max_tokens,
+            "caches": caches,
+            "all_tokens": all_tokens,
+            "samplers": samplers,
+            "logits_processors": logits_processors,
+            "stop_matchers": stop_matchers,
+        }
+        for name, value in values.items():
+            if value is not None and len(value) != count:
+                raise ValueError(
+                    f"{name} must have one entry per segment: "
+                    f"expected {count}, got {len(value)}"
+                )
+
+        max_tokens = max_tokens if max_tokens is not None else [self.max_tokens] * count
+        all_tokens = all_tokens if all_tokens is not None else [[] for _ in segments]
+        samplers = samplers if samplers is not None else [None] * count
+        logits_processors = (
+            logits_processors
+            if logits_processors is not None
+            else [self.logits_processors] * count
+        )
+        stop_matchers = (
+            stop_matchers
+            if stop_matchers is not None
+            else [self._default_stop_matcher] * count
         )
-        stop_matchers = stop_matchers or ([self._default_stop_matcher] * len(segments))
 
-        caches = caches or [None] * len(segments)
+        caches = caches if caches is not None else [None] * count
         for i in range(len(segments)):
             if caches[i] is None:
                 caches[i] = self._make_new_cache()

tests/test_generate.py

@@ -63,6 +63,19 @@ def test_stream_generate_max_tokens(self):
             tokens.append(response.token)
         self.assertEqual(len(tokens), 4)
 
+    def test_stream_generate_zero_max_tokens_is_empty(self):
+        self.assertEqual(
+            list(stream_generate(self.model, self.tokenizer, "hello", max_tokens=0)),
+            [],
+        )
+
+    def test_batch_insert_rejects_short_per_prompt_options(self):
+        batch_gen = BatchGenerator(self.model, max_tokens=1)
+        prompts = [[1], [2]]
+
+        with self.assertRaisesRegex(ValueError, "max_tokens.*expected 2, got 1"):
+            batch_gen.insert(prompts, max_tokens=[1])
+
     def test_generate_with_processor(self):
         init_toks = self.tokenizer.encode("hello")