From d6a85130ef46a959ac98cc838440f1cec6b22dda Mon Sep 17 00:00:00 2001 From: Monroe Williams Date: Mon, 8 Jun 2026 17:38:00 -0700 Subject: [PATCH] fix: normalise logits_processors None to [] to avoid mlx-lm crash (#934) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mlx-lm's GenerationBatch._step does any(self.logits_processors) and for p in self.logits_processors[e], both of which crash with 'NoneType' object is not iterable when logits_processors is None. Apply three layers of defense: 1. _create_batch_generator: pass [] instead of None to BatchGenerator 2. _patched_generation_batch_step: normalise None → [] before calling the original mlx-lm _step (which crashes at any(None)) 3. _patched_ppb_split: normalise None → [] when splitting batches --- omlx/scheduler.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/omlx/scheduler.py b/omlx/scheduler.py index f1e61abee1..a51c8c4b95 100644 --- a/omlx/scheduler.py +++ b/omlx/scheduler.py @@ -385,6 +385,13 @@ def _patched_generation_batch_step(self): deltas = [model._uid_rope_deltas.get(uid, 0.0) for uid in self.uids] model.set_batch_rope_deltas(mx.array(deltas)) + # Defensive: mlx-lm's GenerationBatch._step does `any(self.logits_processors)` + # and `for p in self.logits_processors[e]`, both of which crash when + # logits_processors is None. Normalise to [] so the original code path + # works without modification. See #934. + if self.logits_processors is None: + self.logits_processors = [] + result = _original_generation_batch_step(self) # self._next_tokens contains the just-sampled tokens (async eval pending). @@ -571,7 +578,9 @@ def _patched_ppb_split(self, indices): new_batch.prefill_step_size = self.prefill_step_size new_batch.samplers = self.samplers new_batch.fallback_sampler = self.fallback_sampler - new_batch.logits_processors = self.logits_processors + # Defensive: normalise None → [] to avoid mlx-lm crash in _step + lps = self.logits_processors if self.logits_processors is not None else [] + new_batch.logits_processors = lps new_batch.state_machines = self.state_machines new_batch.max_tokens = self.max_tokens @@ -2030,7 +2039,7 @@ def _create_batch_generator( max_tokens=sampling_params.max_tokens, stop_tokens=stop_tokens_seq, sampler=sampler, - logits_processors=logits_processors if logits_processors else None, + logits_processors=logits_processors if logits_processors else [], prefill_batch_size=1, completion_batch_size=self.config.completion_batch_size, prefill_step_size=self.config.prefill_step_size,