Add support for Muse Glimmer and its tool format

mlx_engine/generate.py

@@ -46,6 +46,8 @@
 from mlx_engine.tool_runtime import (
     create_gemma4_reasoning_guard_logits_processor,
     create_gemma4_tool_context_from_prompt,
+    create_muse_glimmer_tool_context_from_prompt,
+    create_muse_glimmer_tool_logits_processor,
     create_qwen35_reasoning_guard_logits_processor,
     create_qwen35_tool_context_from_prompt,
 )
@@ -727,31 +729,44 @@ def _batched_generation(
                 )
             )
         else:
-            gemma4_tool_context = create_gemma4_tool_context_from_prompt(
+            muse_glimmer_tool_context = create_muse_glimmer_tool_context_from_prompt(
                 tokenizer=model_kit.tokenizer,
                 prompt_tokens=prompt_tokens,
                 model_type=model_kit.model_type,
             )
-            if gemma4_tool_context is not None:
+            if muse_glimmer_tool_context is not None:
                 logits_processors.append(
-                    create_gemma4_reasoning_guard_logits_processor(
+                    create_muse_glimmer_tool_logits_processor(
                         tokenizer=model_kit.tokenizer,
-                        context=gemma4_tool_context,
+                        context=muse_glimmer_tool_context,
                     )
                 )
             else:
-                qwen35_tool_context = create_qwen35_tool_context_from_prompt(
+                gemma4_tool_context = create_gemma4_tool_context_from_prompt(
                     tokenizer=model_kit.tokenizer,
                     prompt_tokens=prompt_tokens,
                     model_type=model_kit.model_type,
                 )
-                if qwen35_tool_context is not None:
+                if gemma4_tool_context is not None:
                     logits_processors.append(
-                        create_qwen35_reasoning_guard_logits_processor(
+                        create_gemma4_reasoning_guard_logits_processor(
                             tokenizer=model_kit.tokenizer,
-                            context=qwen35_tool_context,
+                            context=gemma4_tool_context,
                         )
                     )
+                else:
+                    qwen35_tool_context = create_qwen35_tool_context_from_prompt(
+                        tokenizer=model_kit.tokenizer,
+                        prompt_tokens=prompt_tokens,
+                        model_type=model_kit.model_type,
+                    )
+                    if qwen35_tool_context is not None:
+                        logits_processors.append(
+                            create_qwen35_reasoning_guard_logits_processor(
+                                tokenizer=model_kit.tokenizer,
+                                context=qwen35_tool_context,
+                            )
+                        )
 
         stream = model_kit.generate(
             prompt_tokens=input_tokens,

mlx_engine/model_kit/batched_vision/batch_generator.py

@@ -1000,6 +1000,18 @@ def generate(self, stop_criteria) -> tuple[GenerationBatch, list[Response]]:
         )
         gen_batch._rope_deltas = rope_deltas
 
+        # Keep final prefill work out of the first decode step.
+        eval_targets = [
+            first_tokens,
+            [cache.state for cache in self.prompt_cache],
+            first_logprobs,
+        ]
+        if top_idx is not None:
+            eval_targets.extend([top_idx, top_logprobs])
+        if rope_deltas is not None:
+            eval_targets.append(rope_deltas)
+        mx.eval(*eval_targets)
+
         self.prompt_cache = []
         return gen_batch, prompt_responses
 

mlx_engine/tool_protocols.py

@@ -15,6 +15,11 @@
 QWEN35_REASONING_START = "<think>"
 QWEN35_REASONING_END = "</think>"
 
+MUSE_GLIMMER_ATEM_START = "<atem:function_calls>"
+MUSE_GLIMMER_ATEM_END = "</atem:function_calls>"
+MUSE_GLIMMER_EOM = "<|eom|>"
+MUSE_GLIMMER_EOT = "<|eot|>"
+
 
 @dataclass(frozen=True)
 class Gemma4ToolContext:
@@ -28,6 +33,11 @@ class Qwen35ToolContext:
     reasoning_open: bool
 
 
+@dataclass(frozen=True)
+class MuseGlimmerToolContext:
+    tool_names: tuple[str, ...]
+
+
 def gemma4_reasoning_is_open(text: str) -> bool:
     return text.rfind(GEMMA4_REASONING_START) > text.rfind(GEMMA4_CHANNEL_END)