Match mlx-vlm Gemma 4 prefill logits

[samir-lms]

Summary

• Align mlx-engine’s batched VLM generation path with the opt-in logits_to_keep=1 contract introduced by mlx-vlm 0.6.11.
• Forward the option during chunked prefill, final prefill, and decode whenever the model advertises supports_logits_to_keep.
• For Gemma 4, this limits the vocabulary projection to the final token instead of computing and then discarding logits for the entire prompt.
• This preserves generation parity with mlx-vlm while reducing prefill memory use and unnecessary LM-head computation.
• Add focused coverage verifying that all three batched generation stages forward the option.

Validation

• Focused batched vision/routing suite: 125 passed
• Pre-commit: ruff and ruff-format passed

Performance

Using gemma-4-E2B-it-MLX-4bit on an M5 Max:

• 512-token prefill: 199 MiB lower peak memory and 15.6% faster.
• 2,048-token unchunked prefill: 811 MiB lower peak memory and 10.5% faster.
• Logits allocation dropped from 256 MiB / 1 GiB to 0.5 MiB.
• Decode performance remained unchanged at approximately 116 tokens/sec.
• Results are medians from five warmed prefill runs; decode used four alternating 250-token runs per configuration.

[github-actions]

Thank you for your submission, we really appreciate it. Like many open-source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution. You can sign the CLA by just posting a Pull Request Comment same as the below format.

---

I have read the CLA Document and I hereby sign the CLA

---

_{You can retrigger this bot by commenting recheck in this Pull Request. Posted by the CLA Assistant Lite bot.}

[samir-lms]