Match mlx-vlm Gemma 4 prefill logits

mlx_engine/model_kit/batched_vision/batch_generator.py

@@ -135,6 +135,12 @@ def _clear_qwen3_5_text_rope_state(model: nn.Module, prompt_kwargs: dict) -> Non
     language_model._rope_deltas = None
 
 
+def _with_logits_to_keep(model: nn.Module, kwargs: dict) -> dict:
+    if getattr(model, "supports_logits_to_keep", False):
+        return {**kwargs, "logits_to_keep": 1}
+    return kwargs
+
+
 def _extend_cache(cache_a, cache_b):
     if not cache_a:
         return cache_b
@@ -446,6 +452,7 @@ def _step(self):
             # external/src/mlx-vlm/mlx_vlm/models/qwen3_5/language.py.
             fwd_kwargs["rope_deltas"] = self._rope_deltas
             _sync_scalar_rope_deltas(self.model, self.prompt_cache, self._rope_deltas)
+        fwd_kwargs = _with_logits_to_keep(self.model, fwd_kwargs)
 
         output = self.model(inputs[:, None], cache=self.prompt_cache, **fwd_kwargs)
         logits = output.logits if hasattr(output, "logits") else output
@@ -809,7 +816,9 @@ def needs_processing(self):
 
     def prompt_step(self) -> int:
         n = self._next_prompt_step_size()
-        prompt_kwargs = self._prompt_kwargs_for_next(n)
+        prompt_kwargs = _with_logits_to_keep(
+            self.model, self._prompt_kwargs_for_next(n)
+        )
         # Prompt kwargs with explicit MRoPE state belong to an image prompt; otherwise
         # this text-only chunk must not inherit state from the active decode batch.
         _clear_qwen3_5_text_rope_state(self.model, prompt_kwargs)
@@ -935,7 +944,9 @@ def _emit_cache_save_snapshots(self) -> None:
 
     def generate(self, stop_criteria) -> tuple[GenerationBatch, list[Response]]:
         # This final prompt pass runs after active batched decode in the same tick.
-        prompt_kwargs = self._prompt_kwargs_for_final()
+        prompt_kwargs = _with_logits_to_keep(
+            self.model, self._prompt_kwargs_for_final()
+        )
         _clear_qwen3_5_text_rope_state(self.model, prompt_kwargs)
         try:
             output = self.model(

tests/test_batched_vision_batch_generator.py

@@ -125,6 +125,7 @@ def __call__(self, input_ids, cache=None, inputs_embeds=None, **kwargs):
                     None if inputs_embeds is None else inputs_embeds.shape
                 ),
                 "n_to_process": kwargs.get("n_to_process"),
+                "logits_to_keep": kwargs.get("logits_to_keep"),
                 "position_ids": (
                     None
                     if kwargs.get("position_ids") is None
@@ -176,6 +177,35 @@ def test_batch_generator_uses_vlm_prompt_cache_factory():
     assert type(prompt_cache[0]) is KVCache
 
 
+def test_prefill_and_decode_honor_model_logits_to_keep(monkeypatch):
+    monkeypatch.setattr(
+        batcher,
+        "make_prompt_cache",
+        lambda _model: [_FakeBatchCache()],
+    )
+    model = _FakeModel()
+    model.supports_logits_to_keep = True
+    prompt_prefill = batcher._PromptPrefill(
+        model=model,
+        uid=1,
+        input_ids=[1, 2, 3],
+        max_tokens=1,
+        top_logprobs=0,
+        sampler=_argmax_sampler,
+        logits_processors=[],
+        inputs_embeds=mx.zeros((1, 3, 2), dtype=mx.float32),
+        prompt_kwargs={},
+        prefix_cache_save_state=_prefix_cache_save_states(1)[0],
+        prefill_step_size=2,
+    )
+
+    assert prompt_prefill.prompt_step() == 2
+    generation_batch, _ = prompt_prefill.generate(lambda _token: False)
+    generation_batch.next()
+
+    assert [call["logits_to_keep"] for call in model.calls] == [1, 1, 1]
+
+
 def test_generation_batch_applies_per_sequence_processors_and_top_logprobs():
     """Processors are per-row, and sampled token metadata follows decode-ahead."""
     model = _FakeModel()